[Mlir-commits] [mlir] [mlir][vector] extend `createReadOrMaskedRead`/`createWriteOrMaskedWrite` with permutation map support (PR #202766)
Federico Bruzzone
llvmlistbot at llvm.org
Tue Jun 9 13:54:51 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] [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)
More information about the Mlir-commits
mailing list