[Mlir-commits] [mlir] [mlir][vector] extend `createReadOrMaskedRead`/`createWriteOrMaskedWrite` with permutation map support (PR #202766)

Federico Bruzzone llvmlistbot at llvm.org
Thu Jun 11 03:15:15 PDT 2026


https://github.com/FedericoBruzzone updated https://github.com/llvm/llvm-project/pull/202766

>From cc0c3a69dbe4f84b921a0df14c9e169bb6f14ed9 Mon Sep 17 00:00:00 2001
From: Federico Bruzzone <federico.bruzzone.i at gmail.com>
Date: Tue, 9 Jun 2026 22:12:56 +0200
Subject: [PATCH 1/8] [mlir][vector] extend
 createReadOrMaskedRead/createWriteOrMaskedWrite with permutation map support

Signed-off-by: Federico Bruzzone <federico.bruzzone.i at gmail.com>
---
 .../mlir/Dialect/Vector/Utils/VectorUtils.h   |  16 ++-
 .../Affine/Transforms/SuperVectorize.cpp      |  43 ++------
 mlir/lib/Dialect/Vector/Utils/VectorUtils.cpp | 103 +++++++++++++-----
 3 files changed, 97 insertions(+), 65 deletions(-)

diff --git a/mlir/include/mlir/Dialect/Vector/Utils/VectorUtils.h b/mlir/include/mlir/Dialect/Vector/Utils/VectorUtils.h
index 773b27bc6bfff..fd0d640dae219 100644
--- a/mlir/include/mlir/Dialect/Vector/Utils/VectorUtils.h
+++ b/mlir/include/mlir/Dialect/Vector/Utils/VectorUtils.h
@@ -224,11 +224,17 @@ bool isLinearizableVector(VectorType type);
 /// `useInBoundsInsteadOfMasking` to `true` to use the "in_bounds" attribute
 /// instead of explicit masks.
 ///
-/// Note: all read offsets are set to 0.
+/// When \p permutationMap is provided the in_bounds attribute is inferred from
+/// it: dimension i is in-bounds when the map result is an AffineDimExpr
+/// pointing to a static memref dimension divisible by the vector size, or an
+/// AffineConstantExpr (broadcast). Custom \p indices must also be supplied in
+/// that case; if \p indices is empty, all offsets default to 0.
 Value createReadOrMaskedRead(OpBuilder &builder, Location loc, Value source,
                              const VectorType &vecToReadTy,
                              std::optional<Value> padValue = std::nullopt,
-                             bool useInBoundsInsteadOfMasking = false);
+                             bool useInBoundsInsteadOfMasking = false,
+                             ArrayRef<Value> indices = {},
+                             AffineMap permutationMap = AffineMap());
 
 Value createReadOrMaskedRead(OpBuilder &builder, Location loc, Value source,
                              ArrayRef<int64_t> inputVectorSizes,
@@ -243,11 +249,13 @@ Value createReadOrMaskedRead(OpBuilder &builder, Location loc, Value source,
 /// `useInBoundsInsteadOfMasking` to `true` to use the "in_bounds" attribute
 /// instead of explicit masks.
 /// `writeIndices` specifies the offsets to use. If empty, all indices are set
-/// to 0.
+/// to 0. When \p permutationMap is provided, the in_bounds attribute is
+/// inferred from the map instead of the destination shape.
 Operation *createWriteOrMaskedWrite(OpBuilder &builder, Location loc,
                                     Value vecToStore, Value dest,
                                     SmallVector<Value> writeIndices = {},
-                                    bool useInBoundsInsteadOfMasking = false);
+                                    bool useInBoundsInsteadOfMasking = false,
+                                    AffineMap permutationMap = AffineMap());
 
 /// Returns success if `inputVectorSizes` is a valid masking configuraion for
 /// given `shape`, i.e., it meets:
diff --git a/mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp b/mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp
index 2027b389c02d3..3158a113a7600 100644
--- a/mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp
+++ b/mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp
@@ -1220,28 +1220,6 @@ static bool isIVMappedToMultipleIndices(
   return false;
 }
 
-/// Returns an in-bounds mask for a transfer op given its permutation map and
-/// the memref being accessed. Dimension i is in-bounds when the map result is
-/// an AffineDimExpr pointing to a static memref dimension that is divisible by
-/// the vector size, or an AffineConstantExpr.
-static SmallVector<bool> computeInBoundsMask(AffineMap permutationMap,
-                                             VectorType vectorType,
-                                             MemRefType memrefType) {
-  SmallVector<bool> inBounds(vectorType.getRank(), false);
-  for (unsigned i = 0; i < vectorType.getRank(); ++i) {
-    AffineExpr expr = permutationMap.getResult(i);
-    if (auto dimExpr = dyn_cast<AffineDimExpr>(expr)) {
-      unsigned memDim = dimExpr.getPosition();
-      if (!memrefType.isDynamicDim(memDim) &&
-          memrefType.getDimSize(memDim) % vectorType.getDimSize(i) == 0)
-        inBounds[i] = true;
-    } else if (isa<AffineConstantExpr>(expr)) {
-      inBounds[i] = true;
-    }
-  }
-  return inBounds;
-}
-
 /// Vectorizes an affine load with the vectorization strategy in 'state' by
 /// generating a 'vector.transfer_read' op with the proper permutation map
 /// inferred from the indices of the load. The new 'vector.transfer_read' is
@@ -1287,12 +1265,11 @@ static Operation *vectorizeAffineLoad(AffineLoadOp loadOp,
   LLVM_DEBUG(dbgs() << "\n[early-vect]+++++ permutationMap: ");
   LLVM_DEBUG(permutationMap.print(dbgs()));
 
-  SmallVector<bool> inBounds =
-      computeInBoundsMask(permutationMap, vectorType,
-                          cast<MemRefType>(loadOp.getMemRef().getType()));
-  auto transfer = vector::TransferReadOp::create(
-      state.builder, loadOp.getLoc(), vectorType, loadOp.getMemRef(), indices,
-      /*padding=*/std::nullopt, permutationMap, ArrayRef<bool>(inBounds));
+  Value transferVal = createReadOrMaskedRead(
+      state.builder, loadOp.getLoc(), loadOp.getMemRef(), vectorType,
+      /*padValue=*/std::nullopt, /*useInBoundsInsteadOfMasking=*/true, indices,
+      permutationMap);
+  Operation *transfer = transferVal.getDefiningOp();
 
   // Register replacement for future uses in the scope.
   state.registerOpVectorReplacement(loadOp, transfer);
@@ -1346,13 +1323,11 @@ static Operation *vectorizeAffineStore(AffineStoreOp storeOp,
     return nullptr;
   }
 
-  auto vType = cast<VectorType>(vectorValue.getType());
-  SmallVector<bool> inBounds = computeInBoundsMask(
-      permutationMap, vType, cast<MemRefType>(storeOp.getMemRef().getType()));
-  auto transfer = vector::TransferWriteOp::create(
+  Operation *transfer = createWriteOrMaskedWrite(
       state.builder, storeOp.getLoc(), vectorValue, storeOp.getMemRef(),
-      indices, permutationMap, ArrayRef<bool>(inBounds));
-  LLVM_DEBUG(dbgs() << "\n[early-vect]+++++ vectorized store: " << transfer);
+      SmallVector<Value>(indices.begin(), indices.end()),
+      /*useInBoundsInsteadOfMasking=*/true, permutationMap);
+  LLVM_DEBUG(dbgs() << "\n[early-vect]+++++ vectorized store: " << *transfer);
 
   // Register replacement for future uses in the scope.
   state.registerOpVectorReplacement(storeOp, transfer);
diff --git a/mlir/lib/Dialect/Vector/Utils/VectorUtils.cpp b/mlir/lib/Dialect/Vector/Utils/VectorUtils.cpp
index 576023dbc9de1..bdcb8a3fa097d 100644
--- a/mlir/lib/Dialect/Vector/Utils/VectorUtils.cpp
+++ b/mlir/lib/Dialect/Vector/Utils/VectorUtils.cpp
@@ -420,11 +420,34 @@ Value vector::createReadOrMaskedRead(OpBuilder &builder, Location loc,
                                 useInBoundsInsteadOfMasking);
 }
 
+/// Compute the in_bounds attribute for a transfer op given its permutation map
+/// and the memref being accessed. Dimension i is in-bounds when the map result
+/// is an AffineDimExpr pointing to a static memref dimension divisible by the
+/// vector size, or an AffineConstantExpr (broadcast).
+static SmallVector<bool> computeInBoundsFromPermutationMap(
+    AffineMap permutationMap, VectorType vectorType, MemRefType memrefType) {
+  SmallVector<bool> inBounds(vectorType.getRank(), false);
+  for (unsigned i = 0; i < (unsigned)vectorType.getRank(); ++i) {
+    AffineExpr expr = permutationMap.getResult(i);
+    if (auto dimExpr = dyn_cast<AffineDimExpr>(expr)) {
+      unsigned memDim = dimExpr.getPosition();
+      if (!memrefType.isDynamicDim(memDim) &&
+          memrefType.getDimSize(memDim) % vectorType.getDimSize(i) == 0)
+        inBounds[i] = true;
+    } else if (isa<AffineConstantExpr>(expr)) {
+      inBounds[i] = true;
+    }
+  }
+  return inBounds;
+}
+
 Value vector::createReadOrMaskedRead(OpBuilder &builder, Location loc,
                                      Value source,
                                      const VectorType &vecToReadTy,
                                      std::optional<Value> padValue,
-                                     bool useInBoundsInsteadOfMasking) {
+                                     bool useInBoundsInsteadOfMasking,
+                                     ArrayRef<Value> customIndices,
+                                     AffineMap permutationMap) {
   assert(!llvm::is_contained(vecToReadTy.getScalableDims(),
                              ShapedType::kDynamic) &&
          "invalid input vector sizes");
@@ -434,30 +457,45 @@ Value vector::createReadOrMaskedRead(OpBuilder &builder, Location loc,
   int64_t vecToReadRank = vecToReadTy.getRank();
   auto vecToReadShape = vecToReadTy.getShape();
 
-  assert(sourceShape.size() == static_cast<size_t>(vecToReadRank) &&
+  assert((permutationMap ||
+          sourceShape.size() == static_cast<size_t>(vecToReadRank)) &&
          "expected same ranks.");
   assert((!padValue.has_value() ||
           padValue.value().getType() == sourceShapedType.getElementType()) &&
          "expected same pad element type to match source element type");
 
-  auto zero = arith::ConstantIndexOp::create(builder, loc, 0);
   SmallVector<bool> inBoundsVal(vecToReadRank, true);
 
   if (useInBoundsInsteadOfMasking) {
-    // Update the inBounds attribute.
-    // FIXME: This computation is too weak - it ignores the read indices.
-    for (unsigned i = 0; i < vecToReadRank; i++)
-      inBoundsVal[i] = (sourceShape[i] == vecToReadShape[i]) &&
-                       ShapedType::isStatic(sourceShape[i]);
+    if (permutationMap) {
+      inBoundsVal = computeInBoundsFromPermutationMap(
+          permutationMap, vecToReadTy, cast<MemRefType>(source.getType()));
+    } else {
+      // Update the inBounds attribute.
+      // FIXME: This computation is too weak - it ignores the read indices.
+      for (unsigned i = 0; i < vecToReadRank; i++)
+        inBoundsVal[i] = (sourceShape[i] == vecToReadShape[i]) &&
+                         ShapedType::isStatic(sourceShape[i]);
+    }
+  }
+  SmallVector<Value> indices;
+  if (customIndices.empty()) {
+    auto zero = arith::ConstantIndexOp::create(builder, loc, 0);
+    indices.assign(vecToReadRank, zero);
+  } else {
+    indices.assign(customIndices.begin(), customIndices.end());
   }
-  SmallVector<Value> indices(vecToReadRank, zero);
   auto transferReadOp =
-      vector::TransferReadOp::create(builder, loc,
-                                     /*vectorType=*/vecToReadTy,
-                                     /*source=*/source,
-                                     /*indices=*/indices,
-                                     /*padding=*/padValue,
-                                     /*inBounds=*/inBoundsVal);
+      permutationMap
+          ? vector::TransferReadOp::create(builder, loc, vecToReadTy, source,
+                                           indices, padValue, permutationMap,
+                                           inBoundsVal)
+          : vector::TransferReadOp::create(builder, loc,
+                                           /*vectorType=*/vecToReadTy,
+                                           /*source=*/source,
+                                           /*indices=*/indices,
+                                           /*padding=*/padValue,
+                                           /*inBounds=*/inBoundsVal);
 
   if (useInBoundsInsteadOfMasking)
     return transferReadOp;
@@ -481,7 +519,8 @@ Value vector::createReadOrMaskedRead(OpBuilder &builder, Location loc,
 Operation *vector::createWriteOrMaskedWrite(OpBuilder &builder, Location loc,
                                             Value vecToStore, Value dest,
                                             SmallVector<Value> writeIndices,
-                                            bool useInBoundsInsteadOfMasking) {
+                                            bool useInBoundsInsteadOfMasking,
+                                            AffineMap permutationMap) {
 
   ShapedType destType = cast<ShapedType>(dest.getType());
   int64_t destRank = destType.getRank();
@@ -494,12 +533,17 @@ Operation *vector::createWriteOrMaskedWrite(OpBuilder &builder, Location loc,
   // Compute the in_bounds attribute
   SmallVector<bool> inBoundsVal(vecToStoreRank, true);
   if (useInBoundsInsteadOfMasking) {
-    // Update the inBounds attribute.
-    // FIXME: This computation is too weak - it ignores the write indices.
-    for (unsigned i = 0; i < vecToStoreRank; i++)
-      inBoundsVal[i] =
-          (destShape[destRank - vecToStoreRank + i] >= vecToStoreShape[i]) &&
-          ShapedType::isStatic(destShape[destRank - vecToStoreRank + i]);
+    if (permutationMap) {
+      inBoundsVal = computeInBoundsFromPermutationMap(
+          permutationMap, vecToStoreType, cast<MemRefType>(dest.getType()));
+    } else {
+      // Update the inBounds attribute.
+      // FIXME: This computation is too weak - it ignores the write indices.
+      for (unsigned i = 0; i < vecToStoreRank; i++)
+        inBoundsVal[i] =
+            (destShape[destRank - vecToStoreRank + i] >= vecToStoreShape[i]) &&
+            ShapedType::isStatic(destShape[destRank - vecToStoreRank + i]);
+    }
   }
 
   // If missing, initialize the write indices to 0.
@@ -513,11 +557,16 @@ Operation *vector::createWriteOrMaskedWrite(OpBuilder &builder, Location loc,
   }
 
   // Generate the xfer_write Op
-  Operation *write = vector::TransferWriteOp::create(builder, loc,
-                                                     /*vector=*/vecToStore,
-                                                     /*dest=*/dest,
-                                                     /*indices=*/writeIndices,
-                                                     /*inBounds=*/inBoundsVal);
+  Operation *write =
+      permutationMap
+          ? vector::TransferWriteOp::create(builder, loc, vecToStore, dest,
+                                            writeIndices, permutationMap,
+                                            inBoundsVal)
+          : vector::TransferWriteOp::create(builder, loc,
+                                            /*vector=*/vecToStore,
+                                            /*dest=*/dest,
+                                            /*indices=*/writeIndices,
+                                            /*inBounds=*/inBoundsVal);
 
   // If masking is disabled, exit.
   if (useInBoundsInsteadOfMasking)

>From 1b0aa90e16c31cd871a87396a10f40db69db143f Mon Sep 17 00:00:00 2001
From: Federico Bruzzone <federico.bruzzone.i at gmail.com>
Date: Wed, 10 Jun 2026 15:18:13 +0200
Subject: [PATCH 2/8] Update
 mlir/include/mlir/Dialect/Vector/Utils/VectorUtils.h
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Co-authored-by: Andrzej Warzyński <andrzej.warzynski at gmail.com>
---
 mlir/include/mlir/Dialect/Vector/Utils/VectorUtils.h | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/mlir/include/mlir/Dialect/Vector/Utils/VectorUtils.h b/mlir/include/mlir/Dialect/Vector/Utils/VectorUtils.h
index fd0d640dae219..f56df99afd230 100644
--- a/mlir/include/mlir/Dialect/Vector/Utils/VectorUtils.h
+++ b/mlir/include/mlir/Dialect/Vector/Utils/VectorUtils.h
@@ -224,11 +224,11 @@ bool isLinearizableVector(VectorType type);
 /// `useInBoundsInsteadOfMasking` to `true` to use the "in_bounds" attribute
 /// instead of explicit masks.
 ///
-/// When \p permutationMap is provided the in_bounds attribute is inferred from
+/// When `permutationMap` is provided the in_bounds attribute is inferred from
 /// it: dimension i is in-bounds when the map result is an AffineDimExpr
 /// pointing to a static memref dimension divisible by the vector size, or an
-/// AffineConstantExpr (broadcast). Custom \p indices must also be supplied in
-/// that case; if \p indices is empty, all offsets default to 0.
+/// AffineConstantExpr (broadcast). Custom`indices` must also be supplied in
+/// that case; if `indices` is empty, all offsets default to 0.
 Value createReadOrMaskedRead(OpBuilder &builder, Location loc, Value source,
                              const VectorType &vecToReadTy,
                              std::optional<Value> padValue = std::nullopt,

>From 10c360e7d7e19ad0b1d7061e4032d1570e8e6764 Mon Sep 17 00:00:00 2001
From: Federico Bruzzone <federico.bruzzone.i at gmail.com>
Date: Wed, 10 Jun 2026 15:19:22 +0200
Subject: [PATCH 3/8] Add FIXME

Signed-off-by: Federico Bruzzone <federico.bruzzone.i at gmail.com>
---
 mlir/lib/Dialect/Vector/Utils/VectorUtils.cpp | 9 ++++++---
 1 file changed, 6 insertions(+), 3 deletions(-)

diff --git a/mlir/lib/Dialect/Vector/Utils/VectorUtils.cpp b/mlir/lib/Dialect/Vector/Utils/VectorUtils.cpp
index bdcb8a3fa097d..41490c1c0df14 100644
--- a/mlir/lib/Dialect/Vector/Utils/VectorUtils.cpp
+++ b/mlir/lib/Dialect/Vector/Utils/VectorUtils.cpp
@@ -457,9 +457,10 @@ Value vector::createReadOrMaskedRead(OpBuilder &builder, Location loc,
   int64_t vecToReadRank = vecToReadTy.getRank();
   auto vecToReadShape = vecToReadTy.getShape();
 
-  assert((permutationMap ||
-          sourceShape.size() == static_cast<size_t>(vecToReadRank)) &&
-         "expected same ranks.");
+  size_t expectedSourceRank =
+      permutationMap ? permutationMap.getNumDims() : vecToReadRank;
+  assert(sourceShape.size() == expectedSourceRank &&
+         "expected source rank to match permutation map dims or vector rank.");
   assert((!padValue.has_value() ||
           padValue.value().getType() == sourceShapedType.getElementType()) &&
          "expected same pad element type to match source element type");
@@ -468,6 +469,7 @@ Value vector::createReadOrMaskedRead(OpBuilder &builder, Location loc,
 
   if (useInBoundsInsteadOfMasking) {
     if (permutationMap) {
+      // FIXME: This computation is too weak - it ignores the read indices.
       inBoundsVal = computeInBoundsFromPermutationMap(
           permutationMap, vecToReadTy, cast<MemRefType>(source.getType()));
     } else {
@@ -534,6 +536,7 @@ Operation *vector::createWriteOrMaskedWrite(OpBuilder &builder, Location loc,
   SmallVector<bool> inBoundsVal(vecToStoreRank, true);
   if (useInBoundsInsteadOfMasking) {
     if (permutationMap) {
+      // FIXME: This computation is too weak - it ignores the write indices.
       inBoundsVal = computeInBoundsFromPermutationMap(
           permutationMap, vecToStoreType, cast<MemRefType>(dest.getType()));
     } else {

>From a3c6119096ee816b739476a8f4d4dd4f79d9d538 Mon Sep 17 00:00:00 2001
From: Federico Bruzzone <federico.bruzzone.i at gmail.com>
Date: Wed, 10 Jun 2026 18:45:34 +0200
Subject: [PATCH 4/8] Use ShapedType and update assertions

Signed-off-by: Federico Bruzzone <federico.bruzzone.i at gmail.com>
---
 mlir/lib/Dialect/Vector/Utils/VectorUtils.cpp | 34 +++++++++++++------
 1 file changed, 23 insertions(+), 11 deletions(-)

diff --git a/mlir/lib/Dialect/Vector/Utils/VectorUtils.cpp b/mlir/lib/Dialect/Vector/Utils/VectorUtils.cpp
index 41490c1c0df14..2b4fedb2dbd2a 100644
--- a/mlir/lib/Dialect/Vector/Utils/VectorUtils.cpp
+++ b/mlir/lib/Dialect/Vector/Utils/VectorUtils.cpp
@@ -421,18 +421,18 @@ Value vector::createReadOrMaskedRead(OpBuilder &builder, Location loc,
 }
 
 /// Compute the in_bounds attribute for a transfer op given its permutation map
-/// and the memref being accessed. Dimension i is in-bounds when the map result
-/// is an AffineDimExpr pointing to a static memref dimension divisible by the
+/// and the source being accessed. Dimension i is in-bounds when the map result
+/// is an AffineDimExpr pointing to a static source dimension divisible by the
 /// vector size, or an AffineConstantExpr (broadcast).
 static SmallVector<bool> computeInBoundsFromPermutationMap(
-    AffineMap permutationMap, VectorType vectorType, MemRefType memrefType) {
+    AffineMap permutationMap, VectorType vectorType, ShapedType sourceType) {
   SmallVector<bool> inBounds(vectorType.getRank(), false);
   for (unsigned i = 0; i < (unsigned)vectorType.getRank(); ++i) {
     AffineExpr expr = permutationMap.getResult(i);
     if (auto dimExpr = dyn_cast<AffineDimExpr>(expr)) {
       unsigned memDim = dimExpr.getPosition();
-      if (!memrefType.isDynamicDim(memDim) &&
-          memrefType.getDimSize(memDim) % vectorType.getDimSize(i) == 0)
+      if (!sourceType.isDynamicDim(memDim) &&
+          sourceType.getDimSize(memDim) % vectorType.getDimSize(i) == 0)
         inBounds[i] = true;
     } else if (isa<AffineConstantExpr>(expr)) {
       inBounds[i] = true;
@@ -457,10 +457,16 @@ Value vector::createReadOrMaskedRead(OpBuilder &builder, Location loc,
   int64_t vecToReadRank = vecToReadTy.getRank();
   auto vecToReadShape = vecToReadTy.getShape();
 
-  size_t expectedSourceRank =
-      permutationMap ? permutationMap.getNumDims() : vecToReadRank;
-  assert(sourceShape.size() == expectedSourceRank &&
+  // The permutation map maps the source's index space to the vector's, so its
+  // dims must match the source rank and its results the vector rank. Without a
+  // map, a minor identity is implied, requiring the two ranks to match.
+  assert(sourceShape.size() == (permutationMap
+                                    ? permutationMap.getNumDims()
+                                    : static_cast<size_t>(vecToReadRank)) &&
          "expected source rank to match permutation map dims or vector rank.");
+  assert((!permutationMap || permutationMap.getNumResults() ==
+                                 static_cast<size_t>(vecToReadRank)) &&
+         "expected permutation map results to match vector rank.");
   assert((!padValue.has_value() ||
           padValue.value().getType() == sourceShapedType.getElementType()) &&
          "expected same pad element type to match source element type");
@@ -469,9 +475,10 @@ Value vector::createReadOrMaskedRead(OpBuilder &builder, Location loc,
 
   if (useInBoundsInsteadOfMasking) {
     if (permutationMap) {
+      // Update the inBounds attribute.
       // FIXME: This computation is too weak - it ignores the read indices.
       inBoundsVal = computeInBoundsFromPermutationMap(
-          permutationMap, vecToReadTy, cast<MemRefType>(source.getType()));
+          permutationMap, vecToReadTy, cast<ShapedType>(source.getType()));
     } else {
       // Update the inBounds attribute.
       // FIXME: This computation is too weak - it ignores the read indices.
@@ -480,10 +487,14 @@ Value vector::createReadOrMaskedRead(OpBuilder &builder, Location loc,
                          ShapedType::isStatic(sourceShape[i]);
     }
   }
+  // The transfer op expects one index per source dimension.
+  assert(
+      (customIndices.empty() || customIndices.size() == sourceShape.size()) &&
+      "expected as many custom indices as source dims.");
   SmallVector<Value> indices;
   if (customIndices.empty()) {
     auto zero = arith::ConstantIndexOp::create(builder, loc, 0);
-    indices.assign(vecToReadRank, zero);
+    indices.assign(sourceShape.size(), zero);
   } else {
     indices.assign(customIndices.begin(), customIndices.end());
   }
@@ -536,9 +547,10 @@ Operation *vector::createWriteOrMaskedWrite(OpBuilder &builder, Location loc,
   SmallVector<bool> inBoundsVal(vecToStoreRank, true);
   if (useInBoundsInsteadOfMasking) {
     if (permutationMap) {
+      // Update the inBounds attribute.
       // FIXME: This computation is too weak - it ignores the write indices.
       inBoundsVal = computeInBoundsFromPermutationMap(
-          permutationMap, vecToStoreType, cast<MemRefType>(dest.getType()));
+          permutationMap, vecToStoreType, cast<ShapedType>(dest.getType()));
     } else {
       // Update the inBounds attribute.
       // FIXME: This computation is too weak - it ignores the write indices.

>From a2ff5429b3e4ca63385691fe5f62d07796522b5a Mon Sep 17 00:00:00 2001
From: Federico Bruzzone <federico.bruzzone.i at gmail.com>
Date: Wed, 10 Jun 2026 20:00:22 +0200
Subject: [PATCH 5/8] Treat null permutation map as minor identity in transfer
 builders

Signed-off-by: Federico Bruzzone <federico.bruzzone.i at gmail.com>
---
 .../mlir/Dialect/Vector/IR/VectorOps.td       |  6 ++--
 mlir/lib/Dialect/Vector/IR/VectorOps.cpp      | 31 +++++++---------
 mlir/lib/Dialect/Vector/Utils/VectorUtils.cpp | 35 ++++++++-----------
 3 files changed, 32 insertions(+), 40 deletions(-)

diff --git a/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td b/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
index 5acf2b4ab7649..97035e57eb0be 100644
--- a/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
+++ b/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
@@ -1464,7 +1464,8 @@ def Vector_TransferReadOp :
                    "AffineMapAttr":$permutationMapAttr,
                    "ArrayAttr":$inBoundsAttr)>,
     /// 2. Builder that sets padding to `padding` or poison if not provided and
-    /// an empty mask (variant without attrs).
+    /// an empty mask (variant without attrs). If `permutationMap` is null, a
+    /// minor identity map is used.
     OpBuilder<(ins "VectorType":$vectorType,
                    "Value":$source,
                    "ValueRange":$indices,
@@ -1631,7 +1632,8 @@ def Vector_TransferWriteOp :
                    "ValueRange":$indices,
                    "AffineMapAttr":$permutationMapAttr,
                    "ArrayAttr":$inBoundsAttr)>,
-    /// 3. Builder with type inference that sets an empty mask (variant without attrs).
+    /// 3. Builder with type inference that sets an empty mask (variant without
+    /// attrs). If `permutationMap` is null, a minor identity map is used.
     OpBuilder<(ins "Value":$vector,
                    "Value":$dest,
                    "ValueRange":$indices,
diff --git a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
index 1297f4561b6b7..0feefae03f2b9 100644
--- a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
+++ b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
@@ -5049,11 +5049,15 @@ void TransferReadOp::build(OpBuilder &builder, OperationState &result,
 }
 
 /// 2. Builder that sets padding to zero an empty mask (variant without attrs).
+/// If `permutationMap` is null, a minor identity map is used.
 void TransferReadOp::build(OpBuilder &builder, OperationState &result,
                            VectorType vectorType, Value source,
                            ValueRange indices, std::optional<Value> padding,
                            AffineMap permutationMap,
                            std::optional<ArrayRef<bool>> inBounds) {
+  if (!permutationMap)
+    permutationMap = getTransferMinorIdentityMap(
+        llvm::cast<ShapedType>(source.getType()), vectorType);
   auto permutationMapAttr = AffineMapAttr::get(permutationMap);
   auto inBoundsAttr = (inBounds && !inBounds.value().empty())
                           ? builder.getBoolArrayAttr(inBounds.value())
@@ -5071,19 +5075,8 @@ void TransferReadOp::build(OpBuilder &builder, OperationState &result,
                            VectorType vectorType, Value source,
                            ValueRange indices, std::optional<Value> padding,
                            std::optional<ArrayRef<bool>> inBounds) {
-  AffineMap permutationMap = getTransferMinorIdentityMap(
-      llvm::cast<ShapedType>(source.getType()), vectorType);
-  auto permutationMapAttr = AffineMapAttr::get(permutationMap);
-  auto inBoundsAttr = (inBounds && !inBounds.value().empty())
-                          ? builder.getBoolArrayAttr(inBounds.value())
-                          : builder.getBoolArrayAttr(
-                                SmallVector<bool>(vectorType.getRank(), false));
-  Type elemType = llvm::cast<ShapedType>(source.getType()).getElementType();
-  if (!padding)
-    padding = ub::PoisonOp::create(builder, result.location, elemType);
-  build(builder, result, vectorType, source, indices, permutationMapAttr,
-        *padding,
-        /*mask=*/Value(), inBoundsAttr);
+  build(builder, result, vectorType, source, indices, padding,
+        /*permutationMap=*/AffineMap(), inBounds);
 }
 
 template <typename EmitFun>
@@ -5692,11 +5685,15 @@ void TransferWriteOp::build(OpBuilder &builder, OperationState &result,
 }
 
 /// 3. Builder with type inference that sets an empty mask (variant without
-/// attrs)
+/// attrs). If `permutationMap` is null, a minor identity map is used.
 void TransferWriteOp::build(OpBuilder &builder, OperationState &result,
                             Value vector, Value dest, ValueRange indices,
                             AffineMap permutationMap,
                             std::optional<ArrayRef<bool>> inBounds) {
+  if (!permutationMap)
+    permutationMap =
+        getTransferMinorIdentityMap(llvm::cast<ShapedType>(dest.getType()),
+                                    llvm::cast<VectorType>(vector.getType()));
   auto permutationMapAttr = AffineMapAttr::get(permutationMap);
   auto inBoundsAttr =
       (inBounds && !inBounds.value().empty())
@@ -5712,10 +5709,8 @@ void TransferWriteOp::build(OpBuilder &builder, OperationState &result,
 void TransferWriteOp::build(OpBuilder &builder, OperationState &result,
                             Value vector, Value dest, ValueRange indices,
                             std::optional<ArrayRef<bool>> inBounds) {
-  auto vectorType = llvm::cast<VectorType>(vector.getType());
-  AffineMap permutationMap = getTransferMinorIdentityMap(
-      llvm::cast<ShapedType>(dest.getType()), vectorType);
-  build(builder, result, vector, dest, indices, permutationMap, inBounds);
+  build(builder, result, vector, dest, indices, /*permutationMap=*/AffineMap(),
+        inBounds);
 }
 
 ParseResult TransferWriteOp::parse(OpAsmParser &parser,
diff --git a/mlir/lib/Dialect/Vector/Utils/VectorUtils.cpp b/mlir/lib/Dialect/Vector/Utils/VectorUtils.cpp
index 2b4fedb2dbd2a..9b1933150306d 100644
--- a/mlir/lib/Dialect/Vector/Utils/VectorUtils.cpp
+++ b/mlir/lib/Dialect/Vector/Utils/VectorUtils.cpp
@@ -498,17 +498,14 @@ Value vector::createReadOrMaskedRead(OpBuilder &builder, Location loc,
   } else {
     indices.assign(customIndices.begin(), customIndices.end());
   }
+  // A null permutation map means the builder defaults to a minor identity map.
   auto transferReadOp =
-      permutationMap
-          ? vector::TransferReadOp::create(builder, loc, vecToReadTy, source,
-                                           indices, padValue, permutationMap,
-                                           inBoundsVal)
-          : vector::TransferReadOp::create(builder, loc,
-                                           /*vectorType=*/vecToReadTy,
-                                           /*source=*/source,
-                                           /*indices=*/indices,
-                                           /*padding=*/padValue,
-                                           /*inBounds=*/inBoundsVal);
+      vector::TransferReadOp::create(builder, loc, /*vectorType=*/vecToReadTy,
+                                     /*source=*/source,
+                                     /*indices=*/indices,
+                                     /*padding=*/padValue,
+                                     /*permutationMap=*/permutationMap,
+                                     /*inBounds=*/inBoundsVal);
 
   if (useInBoundsInsteadOfMasking)
     return transferReadOp;
@@ -571,17 +568,15 @@ Operation *vector::createWriteOrMaskedWrite(OpBuilder &builder, Location loc,
     writeIndices.assign(destRank, zero);
   }
 
-  // Generate the xfer_write Op
+  // Generate the xfer_write Op. A null permutation map means the builder
+  // defaults to a minor identity map.
   Operation *write =
-      permutationMap
-          ? vector::TransferWriteOp::create(builder, loc, vecToStore, dest,
-                                            writeIndices, permutationMap,
-                                            inBoundsVal)
-          : vector::TransferWriteOp::create(builder, loc,
-                                            /*vector=*/vecToStore,
-                                            /*dest=*/dest,
-                                            /*indices=*/writeIndices,
-                                            /*inBounds=*/inBoundsVal);
+      vector::TransferWriteOp::create(builder, loc,
+                                      /*vector=*/vecToStore,
+                                      /*dest=*/dest,
+                                      /*indices=*/writeIndices,
+                                      /*permutationMap=*/permutationMap,
+                                      /*inBounds=*/inBoundsVal);
 
   // If masking is disabled, exit.
   if (useInBoundsInsteadOfMasking)

>From cb32d146097b989421e9534d047ab74df1f725be Mon Sep 17 00:00:00 2001
From: Federico Bruzzone <federico.bruzzone.i at gmail.com>
Date: Thu, 11 Jun 2026 11:29:24 +0200
Subject: [PATCH 6/8] Update
 mlir/include/mlir/Dialect/Vector/Utils/VectorUtils.h
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Co-authored-by: Andrzej Warzyński <andrzej.warzynski at gmail.com>
---
 mlir/include/mlir/Dialect/Vector/Utils/VectorUtils.h | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/mlir/include/mlir/Dialect/Vector/Utils/VectorUtils.h b/mlir/include/mlir/Dialect/Vector/Utils/VectorUtils.h
index f56df99afd230..edfbbf4ba6da0 100644
--- a/mlir/include/mlir/Dialect/Vector/Utils/VectorUtils.h
+++ b/mlir/include/mlir/Dialect/Vector/Utils/VectorUtils.h
@@ -249,7 +249,7 @@ Value createReadOrMaskedRead(OpBuilder &builder, Location loc, Value source,
 /// `useInBoundsInsteadOfMasking` to `true` to use the "in_bounds" attribute
 /// instead of explicit masks.
 /// `writeIndices` specifies the offsets to use. If empty, all indices are set
-/// to 0. When \p permutationMap is provided, the in_bounds attribute is
+/// to 0. When `permutationMap` is provided, the in_bounds attribute is
 /// inferred from the map instead of the destination shape.
 Operation *createWriteOrMaskedWrite(OpBuilder &builder, Location loc,
                                     Value vecToStore, Value dest,

>From 1cf96cc560556112cf834a4abac50b752c85d5c0 Mon Sep 17 00:00:00 2001
From: Federico Bruzzone <federico.bruzzone.i at gmail.com>
Date: Thu, 11 Jun 2026 11:30:12 +0200
Subject: [PATCH 7/8] Update mlir/lib/Dialect/Vector/IR/VectorOps.cpp
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Co-authored-by: Andrzej Warzyński <andrzej.warzynski at gmail.com>
---
 mlir/lib/Dialect/Vector/IR/VectorOps.cpp | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
index 0feefae03f2b9..a24671941b961 100644
--- a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
+++ b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
@@ -5048,8 +5048,8 @@ void TransferReadOp::build(OpBuilder &builder, OperationState &result,
         *padding, /*mask=*/Value(), inBoundsAttr);
 }
 
-/// 2. Builder that sets padding to zero an empty mask (variant without attrs).
-/// If `permutationMap` is null, a minor identity map is used.
+/// 2. Builder that sets padding to zero and an empty mask (variant without
+/// attrs). If `permutationMap` is null, a minor identity map is used.
 void TransferReadOp::build(OpBuilder &builder, OperationState &result,
                            VectorType vectorType, Value source,
                            ValueRange indices, std::optional<Value> padding,

>From 43336c58d03909ca409bc9ed0440ebe8efb39f0b Mon Sep 17 00:00:00 2001
From: Federico Bruzzone <federico.bruzzone.i at gmail.com>
Date: Thu, 11 Jun 2026 12:13:24 +0200
Subject: [PATCH 8/8] Add docs for VectorOps, both .td and .cpp files

Signed-off-by: Federico Bruzzone <federico.bruzzone.i at gmail.com>
---
 .../include/mlir/Dialect/Vector/IR/VectorOps.td |  9 +++++++--
 mlir/lib/Dialect/Vector/IR/VectorOps.cpp        | 17 ++++++++++++-----
 2 files changed, 19 insertions(+), 7 deletions(-)

diff --git a/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td b/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
index 97035e57eb0be..7578ce78a0f00 100644
--- a/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
+++ b/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
@@ -1457,6 +1457,7 @@ def Vector_TransferReadOp :
   let builders = [
     /// 1. Builder that sets padding to `padding` or poison if not provided and
     /// an empty mask (variant with attrs).
+    /// If `padding` is null, a poison value is used.
     OpBuilder<(ins "VectorType":$vectorType,
                    "Value":$source,
                    "ValueRange":$indices,
@@ -1464,8 +1465,10 @@ def Vector_TransferReadOp :
                    "AffineMapAttr":$permutationMapAttr,
                    "ArrayAttr":$inBoundsAttr)>,
     /// 2. Builder that sets padding to `padding` or poison if not provided and
-    /// an empty mask (variant without attrs). If `permutationMap` is null, a
-    /// minor identity map is used.
+    /// an empty mask (variant without attrs). 
+    /// If `padding` is null, a poison value is used.
+    /// If `permutationMap` is null, a minor identity map is used.
+    /// If `inBounds` is null, an empty mask is used.
     OpBuilder<(ins "VectorType":$vectorType,
                    "Value":$source,
                    "ValueRange":$indices,
@@ -1474,6 +1477,8 @@ def Vector_TransferReadOp :
                    CArg<"std::optional<ArrayRef<bool>>", "::std::nullopt">:$inBounds)>,
     /// 3. Builder that sets padding to `padding` or poison if not provided and
     /// permutation map to 'getMinorIdentityMap'.
+    /// If `padding` is null, a poison value is used.
+    /// If `inBounds` is null, an empty mask is used.
     OpBuilder<(ins "VectorType":$vectorType,
                    "Value":$source,
                    "ValueRange":$indices,
diff --git a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
index a24671941b961..67c31730f4b65 100644
--- a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
+++ b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
@@ -5035,6 +5035,7 @@ void ExtractStridedSliceOp::getCanonicalizationPatterns(
 //===----------------------------------------------------------------------===//
 
 /// 1. Builder that sets padding to zero and an empty mask (variant with attrs).
+/// If `padding` is null, a poison value is used.
 void TransferReadOp::build(OpBuilder &builder, OperationState &result,
                            VectorType vectorType, Value source,
                            ValueRange indices, std::optional<Value> padding,
@@ -5044,12 +5045,17 @@ void TransferReadOp::build(OpBuilder &builder, OperationState &result,
   Type elemType = llvm::cast<ShapedType>(source.getType()).getElementType();
   if (!padding)
     padding = ub::PoisonOp::create(builder, result.location, elemType);
+  // Delegate to the most general builder (see
+  // `mlir/Dialect/Vector/IR/VectorOps.cpp.inc`)
   build(builder, result, vectorType, source, indices, permutationMapAttr,
         *padding, /*mask=*/Value(), inBoundsAttr);
 }
 
 /// 2. Builder that sets padding to zero and an empty mask (variant without
-/// attrs). If `permutationMap` is null, a minor identity map is used.
+/// attrs).
+/// If `padding` is null, a poison value is used.
+/// If `permutationMap` is null, a minor identity map is used.
+/// If `inBounds` is null, an empty mask is used.
 void TransferReadOp::build(OpBuilder &builder, OperationState &result,
                            VectorType vectorType, Value source,
                            ValueRange indices, std::optional<Value> padding,
@@ -5063,18 +5069,19 @@ void TransferReadOp::build(OpBuilder &builder, OperationState &result,
                           ? builder.getBoolArrayAttr(inBounds.value())
                           : builder.getBoolArrayAttr(
                                 SmallVector<bool>(vectorType.getRank(), false));
-  Type elemType = llvm::cast<ShapedType>(source.getType()).getElementType();
-  if (!padding)
-    padding = ub::PoisonOp::create(builder, result.location, elemType);
-  build(builder, result, vectorType, source, indices, *padding,
+  // Delegate to Builder 1
+  build(builder, result, vectorType, source, indices, padding,
         permutationMapAttr, inBoundsAttr);
 }
 
 /// 3. Builder that sets permutation map to 'getMinorIdentityMap'.
+/// If `padding` is null, a poison value is used.
+/// If `inBounds` is null, an empty mask is used.
 void TransferReadOp::build(OpBuilder &builder, OperationState &result,
                            VectorType vectorType, Value source,
                            ValueRange indices, std::optional<Value> padding,
                            std::optional<ArrayRef<bool>> inBounds) {
+  // Delegate to Builder 2
   build(builder, result, vectorType, source, indices, padding,
         /*permutationMap=*/AffineMap(), inBounds);
 }



More information about the Mlir-commits mailing list