[Mlir-commits] [mlir] 3b63f04 - [mlir][vector] extend `createReadOrMaskedRead`/`createWriteOrMaskedWrite` with permutation map support (#202766)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Fri Jun 12 05:48:27 PDT 2026
Author: Federico Bruzzone
Date: 2026-06-12T13:48:23+01:00
New Revision: 3b63f045b9a994e275d8fd186f2de96939c38378
URL: https://github.com/llvm/llvm-project/commit/3b63f045b9a994e275d8fd186f2de96939c38378
DIFF: https://github.com/llvm/llvm-project/commit/3b63f045b9a994e275d8fd186f2de96939c38378.diff
LOG: [mlir][vector] extend `createReadOrMaskedRead`/`createWriteOrMaskedWrite` with permutation map support (#202766)
Follow-up to #201180.
Extends the existing `createReadOrMaskedRead` and
`createWriteOrMaskedWrite` utilities in `VectorUtils` with two optional
trailing parameters:
- `ArrayRef<Value> indices`
- `AffineMap permutationMap`
The affine super-vectorizer is updated to call these functions instead
of constructing `TransferReadOp`/`TransferWriteOp` directly.
@banach-space, please correct me if this wasn't what you meant in the
previous PR.
---------
Signed-off-by: Federico Bruzzone <federico.bruzzone.i at gmail.com>
Co-authored-by: Andrzej Warzyński <andrzej.warzynski at gmail.com>
Added:
Modified:
mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
mlir/include/mlir/Dialect/Vector/Utils/VectorUtils.h
mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp
mlir/lib/Dialect/Vector/IR/VectorOps.cpp
mlir/lib/Dialect/Vector/Utils/VectorUtils.cpp
Removed:
################################################################################
diff --git a/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td b/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
index 5acf2b4ab7649..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,7 +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).
+ /// 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,
@@ -1473,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,
@@ -1631,7 +1637,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/include/mlir/Dialect/Vector/Utils/VectorUtils.h b/mlir/include/mlir/Dialect/Vector/Utils/VectorUtils.h
index 773b27bc6bfff..edfbbf4ba6da0 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 `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`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,
- 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 `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/IR/VectorOps.cpp b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
index 1297f4561b6b7..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,46 +5045,45 @@ 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 an empty mask (variant without attrs).
+/// 2. Builder that sets padding to zero and 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.
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())
: 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) {
- 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);
+ // Delegate to Builder 2
+ build(builder, result, vectorType, source, indices, padding,
+ /*permutationMap=*/AffineMap(), inBounds);
}
template <typename EmitFun>
@@ -5692,11 +5692,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 +5716,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 576023dbc9de1..c38213850c6ea 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 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, 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 (!sourceType.isDynamicDim(memDim) &&
+ sourceType.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,29 +457,53 @@ Value vector::createReadOrMaskedRead(OpBuilder &builder, Location loc,
int64_t vecToReadRank = vecToReadTy.getRank();
auto vecToReadShape = vecToReadTy.getShape();
- assert(sourceShape.size() == static_cast<size_t>(vecToReadRank) &&
- "expected same ranks.");
+ // 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");
- 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) {
+ // Update the inBounds attribute.
+ // FIXME: This computation is too weak - it ignores the read indices.
+ inBoundsVal = computeInBoundsFromPermutationMap(
+ permutationMap, vecToReadTy, cast<ShapedType>(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(vecToReadRank, zero);
+ // 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;
+ customIndices.empty()
+ ? indices.assign(sourceShape.size(),
+ arith::ConstantIndexOp::create(builder, loc, 0))
+ : indices.assign(customIndices.begin(), customIndices.end());
+
+ // A null permutation map means the builder defaults to a minor identity map.
auto transferReadOp =
- vector::TransferReadOp::create(builder, loc,
- /*vectorType=*/vecToReadTy,
+ vector::TransferReadOp::create(builder, loc, /*vectorType=*/vecToReadTy,
/*source=*/source,
/*indices=*/indices,
/*padding=*/padValue,
+ /*permutationMap=*/permutationMap,
/*inBounds=*/inBoundsVal);
if (useInBoundsInsteadOfMasking)
@@ -481,7 +528,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 +542,19 @@ 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) {
+ // Update the inBounds attribute.
+ // FIXME: This computation is too weak - it ignores the write indices.
+ inBoundsVal = computeInBoundsFromPermutationMap(
+ permutationMap, vecToStoreType, cast<ShapedType>(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.
@@ -512,12 +567,15 @@ Operation *vector::createWriteOrMaskedWrite(OpBuilder &builder, Location loc,
writeIndices.assign(destRank, zero);
}
- // Generate the xfer_write Op
- Operation *write = vector::TransferWriteOp::create(builder, loc,
- /*vector=*/vecToStore,
- /*dest=*/dest,
- /*indices=*/writeIndices,
- /*inBounds=*/inBoundsVal);
+ // Generate the xfer_write Op. A null permutation map means the builder
+ // defaults to a minor identity map.
+ Operation *write =
+ vector::TransferWriteOp::create(builder, loc,
+ /*vector=*/vecToStore,
+ /*dest=*/dest,
+ /*indices=*/writeIndices,
+ /*permutationMap=*/permutationMap,
+ /*inBounds=*/inBoundsVal);
// If masking is disabled, exit.
if (useInBoundsInsteadOfMasking)
More information about the Mlir-commits
mailing list