[Mlir-commits] [mlir] [mlir][vector] extend `createReadOrMaskedRead`/`createWriteOrMaskedWrite` with permutation map support (PR #202766)
Federico Bruzzone
llvmlistbot at llvm.org
Thu Jun 11 02:29:36 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/6] [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/6] 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/6] 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/6] 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/6] 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/6] 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,
More information about the Mlir-commits
mailing list