[Mlir-commits] [mlir] [mlir][vector] Canonicalize transfer_{read, write} vector<1xT> (PR #196598)
Erick Ochoa Lopez
llvmlistbot at llvm.org
Fri May 8 11:01:28 PDT 2026
https://github.com/amd-eochoalo created https://github.com/llvm/llvm-project/pull/196598
vector.transfer_read and vector.transfer_write's permutations maps are irrelevant with vector<1xT>. This pattern unblocks lowerings to vector.load and vector.store.
>From bebc8c0fa81d4fd353d144523200567c27fae55a Mon Sep 17 00:00:00 2001
From: Erick Ochoa <erick.ochoalopez at amd.com>
Date: Fri, 8 May 2026 13:53:04 -0400
Subject: [PATCH] [mlir][vector] Canonicalize transfer_{read,write} vector<1xT>
vector.transfer_read and vector.transfer_write's permutations maps are
irrelevant with vector<1xT>. This pattern unblocks lowerings to
vector.load and vector.store.
---
mlir/lib/Dialect/Vector/IR/VectorOps.cpp | 49 ++++++++-
mlir/test/Dialect/Vector/canonicalize.mlir | 109 +++++++++++++++++++++
2 files changed, 156 insertions(+), 2 deletions(-)
diff --git a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
index 51be1e4431e70..1565e27c2b5dd 100644
--- a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
+++ b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
@@ -5625,11 +5625,54 @@ struct TransferReadAfterWriteToBroadcast
return success();
}
};
+/// Canonicalize a rank-1, single-element vector transfer so that its
+/// permutation map is the minor identity.
+///
+/// When the vector type is `vector<1xT>`, the permutation map is irrelevant
+/// to which element is accessed: the single vector lane has iteration offset 0,
+/// so the element is always at `indices` regardless of which source dimension
+/// the map points at. Replacing the map with the minor identity unblocks
+/// lowering to vector.load / vector.store.
+template <typename TransferOp>
+struct CanonicalizeSize1TransferPermutationMap final
+ : OpRewritePattern<TransferOp> {
+ using OpRewritePattern<TransferOp>::OpRewritePattern;
+
+ LogicalResult matchAndRewrite(TransferOp op,
+ PatternRewriter &rewriter) const override {
+ VectorType vecType = op.getVectorType();
+ if (vecType.getRank() != 1 || vecType.getShape()[0] != 1)
+ return failure();
+
+ AffineMap map = op.getPermutationMap();
+ if (map.isMinorIdentity())
+ return failure();
+
+ int64_t srcRank = op.getShapedType().getRank();
+ AffineMap minorIdentity =
+ AffineMap::getMinorIdentityMap(srcRank, 1, rewriter.getContext());
+
+ if constexpr (std::is_same_v<TransferOp, vector::TransferReadOp>) {
+ rewriter.replaceOpWithNewOp<vector::TransferReadOp>(
+ op, vecType, op.getBase(), op.getIndices(),
+ AffineMapAttr::get(minorIdentity), op.getPadding(), op.getMask(),
+ op.getInBoundsAttr());
+ } else {
+ rewriter.replaceOpWithNewOp<vector::TransferWriteOp>(
+ op, op.getVector(), op.getBase(), op.getIndices(),
+ AffineMapAttr::get(minorIdentity), op.getMask(),
+ op.getInBoundsAttr());
+ }
+ return success();
+ }
+};
+
} // namespace
void TransferReadOp::getCanonicalizationPatterns(RewritePatternSet &results,
MLIRContext *context) {
- results.add<TransferReadAfterWriteToBroadcast>(context);
+ results.add<TransferReadAfterWriteToBroadcast,
+ CanonicalizeSize1TransferPermutationMap<TransferReadOp>>(context);
}
FailureOr<std::optional<SmallVector<Value>>>
@@ -6127,7 +6170,9 @@ struct SwapExtractSliceOfTransferWrite
void TransferWriteOp::getCanonicalizationPatterns(RewritePatternSet &results,
MLIRContext *context) {
- results.add<FoldWaw, SwapExtractSliceOfTransferWrite>(context);
+ results.add<FoldWaw, SwapExtractSliceOfTransferWrite,
+ CanonicalizeSize1TransferPermutationMap<TransferWriteOp>>(
+ context);
}
FailureOr<std::optional<SmallVector<Value>>>
diff --git a/mlir/test/Dialect/Vector/canonicalize.mlir b/mlir/test/Dialect/Vector/canonicalize.mlir
index 6aa92ab79a0dd..5f80e63f49692 100644
--- a/mlir/test/Dialect/Vector/canonicalize.mlir
+++ b/mlir/test/Dialect/Vector/canonicalize.mlir
@@ -4407,3 +4407,112 @@ func.func @no_fold_alltrue_mask_empty_body_scalar_result(
%result = vector.mask %all_true, %passthru { vector.yield %val : i32 } : vector<1xi1> -> i32
return %result : i32
}
+
+// -----
+
+// transfer_read with vector<1xf32> and non-minor-identity map: canonicalize
+// the permutation map to minor identity.
+// (d0, d1) -> (d0) becomes (d0, d1) -> (d1)
+
+// CHECK-LABEL: func.func @canonicalize_transfer_read_size1_map
+// CHECK-SAME: (%[[M:.+]]: memref<4x3xf32>, %[[I:.+]]: index, %[[J:.+]]: index, %[[PAD:.+]]: f32)
+// CHECK: %[[V:.+]] = vector.transfer_read %[[M]][%[[I]], %[[J]]], %[[PAD]]
+// CHECK-SAME: {in_bounds = [true]}
+// CHECK-SAME: : memref<4x3xf32>, vector<1xf32>
+// CHECK-NOT: permutation_map
+// CHECK: return %[[V]]
+func.func @canonicalize_transfer_read_size1_map(%m: memref<4x3xf32>, %i: index, %j: index, %pad: f32) -> vector<1xf32> {
+ %v = vector.transfer_read %m[%i, %j], %pad {in_bounds = [true], permutation_map = affine_map<(d0, d1) -> (d0)>} : memref<4x3xf32>, vector<1xf32>
+ return %v : vector<1xf32>
+}
+
+// -----
+
+// transfer_write with vector<1xf32> and non-minor-identity map: canonicalize.
+
+// CHECK-LABEL: func.func @canonicalize_transfer_write_size1_map
+// CHECK-SAME: (%[[V:.+]]: vector<1xf32>, %[[M:.+]]: memref<4x3xf32>, %[[I:.+]]: index, %[[J:.+]]: index)
+// CHECK: vector.transfer_write %[[V]], %[[M]][%[[I]], %[[J]]]
+// CHECK-SAME: {in_bounds = [true]}
+// CHECK-SAME: : vector<1xf32>, memref<4x3xf32>
+// CHECK-NOT: permutation_map
+func.func @canonicalize_transfer_write_size1_map(%v: vector<1xf32>, %m: memref<4x3xf32>, %i: index, %j: index) {
+ vector.transfer_write %v, %m[%i, %j] {in_bounds = [true], permutation_map = affine_map<(d0, d1) -> (d0)>} : vector<1xf32>, memref<4x3xf32>
+ return
+}
+
+// -----
+
+// Rank-3 memref: (d0, d1, d2) -> (d0) becomes (d0, d1, d2) -> (d2).
+
+// CHECK-LABEL: func.func @canonicalize_transfer_read_size1_rank3
+// CHECK: vector.transfer_read
+// CHECK-NOT: permutation_map
+// CHECK-SAME: : memref<4x3x2xf32>, vector<1xf32>
+func.func @canonicalize_transfer_read_size1_rank3(%m: memref<4x3x2xf32>, %i: index, %j: index, %k: index, %pad: f32) -> vector<1xf32> {
+ %v = vector.transfer_read %m[%i, %j, %k], %pad {in_bounds = [true], permutation_map = affine_map<(d0, d1, d2) -> (d0)>} : memref<4x3x2xf32>, vector<1xf32>
+ return %v : vector<1xf32>
+}
+
+// -----
+
+// Middle dimension map: (d0, d1, d2) -> (d1) becomes (d0, d1, d2) -> (d2).
+
+// CHECK-LABEL: func.func @canonicalize_transfer_read_size1_middle_dim
+// CHECK: vector.transfer_read
+// CHECK-NOT: permutation_map
+// CHECK-SAME: : memref<4x3x2xf32>, vector<1xf32>
+func.func @canonicalize_transfer_read_size1_middle_dim(%m: memref<4x3x2xf32>, %i: index, %j: index, %k: index, %pad: f32) -> vector<1xf32> {
+ %v = vector.transfer_read %m[%i, %j, %k], %pad {in_bounds = [true], permutation_map = affine_map<(d0, d1, d2) -> (d1)>} : memref<4x3x2xf32>, vector<1xf32>
+ return %v : vector<1xf32>
+}
+
+// -----
+
+// Negative: already minor identity -- should not be modified.
+
+// CHECK-LABEL: func.func @negative_size1_already_minor_identity
+// CHECK: vector.transfer_read
+// CHECK-NOT: permutation_map
+// CHECK-SAME: : memref<4x3xf32>, vector<1xf32>
+func.func @negative_size1_already_minor_identity(%m: memref<4x3xf32>, %i: index, %j: index, %pad: f32) -> vector<1xf32> {
+ %v = vector.transfer_read %m[%i, %j], %pad {in_bounds = [true], permutation_map = affine_map<(d0, d1) -> (d1)>} : memref<4x3xf32>, vector<1xf32>
+ return %v : vector<1xf32>
+}
+
+// -----
+
+// Negative: vector<4xf32> (size != 1) -- should not be modified.
+
+// CHECK-LABEL: func.func @negative_size1_size_not_one
+// CHECK: vector.transfer_read
+// CHECK-SAME: permutation_map
+// CHECK-SAME: : memref<4x3xf32>, vector<4xf32>
+func.func @negative_size1_size_not_one(%m: memref<4x3xf32>, %i: index, %j: index, %pad: f32) -> vector<4xf32> {
+ %v = vector.transfer_read %m[%i, %j], %pad {in_bounds = [true], permutation_map = affine_map<(d0, d1) -> (d0)>} : memref<4x3xf32>, vector<4xf32>
+ return %v : vector<4xf32>
+}
+
+// -----
+
+// Negative: rank-0 vector<f32> -- should not be modified.
+
+// CHECK-LABEL: func.func @negative_size1_rank0
+// CHECK: vector.transfer_read
+// CHECK-SAME: : memref<4xf32>, vector<f32>
+func.func @negative_size1_rank0(%m: memref<4xf32>, %i: index, %pad: f32) -> vector<f32> {
+ %v = vector.transfer_read %m[%i], %pad : memref<4xf32>, vector<f32>
+ return %v : vector<f32>
+}
+
+// -----
+
+// Negative: rank-2 vector -- should not be modified.
+
+// CHECK-LABEL: func.func @negative_size1_rank2
+// CHECK: vector.transfer_read
+// CHECK-SAME: : memref<4x3xf32>, vector<2x4xf32>
+func.func @negative_size1_rank2(%m: memref<4x3xf32>, %i: index, %j: index, %pad: f32) -> vector<2x4xf32> {
+ %v = vector.transfer_read %m[%i, %j], %pad {in_bounds = [true, true]} : memref<4x3xf32>, vector<2x4xf32>
+ return %v : vector<2x4xf32>
+}
More information about the Mlir-commits
mailing list