[Mlir-commits] [mlir] [mlir][vector] Check the indices when inferring in_bounds from a permutation map (PR #221268)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Fri Sep 4 09:13:12 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-vector
Author: Hung-Kuan Tseng (Tim096)
<details>
<summary>Changes</summary>
`createReadOrMaskedRead` and `createWriteOrMaskedWrite` infer `in_bounds`
from the permutation map alone: a vector dimension is in bounds when the
source dimension it maps to is static and divisible by the vector size.
That only holds when the index is a multiple of the vector size, which the
affine super-vectorizer guarantees for an induction variable that starts at
zero and for nothing else. With an offset index or an unaligned lower bound
the last vector runs past the end of the memref while still carrying
`in_bounds = [true]`:
```mlir
affine.for %i = 0 to 15 {
%v = affine.load %A[%i + 1] : memref<16xf32>
affine.store %v, %B[%i] : memref<16xf32>
}
```
`-affine-super-vectorize="virtual-vector-size=8"` today:
```mlir
affine.for %arg2 = 0 to 15 step 8 {
%0 = affine.apply affine_map<(d0) -> (d0 + 1)>(%arg2)
%2 = vector.transfer_read %arg0[%0], %1 {in_bounds = [true]} : memref<16xf32>, vector<8xf32>
```
The second iteration reads elements 9 to 16, and `-convert-vector-to-llvm`
turns the read into a plain `llvm.load`. The same happens for
`affine.for %i = 4 to 16` with `%A[%i]`.
This requires the index to be a known multiple of the vector size: a
constant, an `affine.for` induction variable whose step and lower bound
are multiples, or an `affine.apply` whose map preserves that. Loop bounds
are followed recursively, so a tiled loop whose lower bound is an outer
induction variable, as in `vectorize_2d_inbounds.mlir`, still qualifies.
The check is a syntactic walk over the defining ops, no constraint sets.
The existing `affine-super-vectorize` tests keep their `in_bounds` except
`vec_affine_apply_2`, whose index `d0 mod 16 + 1` is offset by one. New
tests cover the offset and unaligned-lower-bound cases and their aligned
counterparts.
The `FIXME`s on the no-map branches stay: every in-tree caller reaches
them with zero indices, or with `tensor.insert_slice` offsets that the
verifier already bounds.
---
Full diff: https://github.com/llvm/llvm-project/pull/221268.diff
4 Files Affected:
- (modified) mlir/include/mlir/Dialect/Vector/Utils/VectorUtils.h (+6-4)
- (modified) mlir/lib/Dialect/Vector/Utils/VectorUtils.cpp (+79-30)
- (added) mlir/test/Dialect/Affine/SuperVectorize/vectorize_1d_inbounds.mlir (+103)
- (modified) mlir/test/Dialect/Affine/SuperVectorize/vectorize_affine_apply.mlir (+1-1)
``````````diff
diff --git a/mlir/include/mlir/Dialect/Vector/Utils/VectorUtils.h b/mlir/include/mlir/Dialect/Vector/Utils/VectorUtils.h
index edfbbf4ba6da0..4080be7ec4af2 100644
--- a/mlir/include/mlir/Dialect/Vector/Utils/VectorUtils.h
+++ b/mlir/include/mlir/Dialect/Vector/Utils/VectorUtils.h
@@ -225,9 +225,10 @@ bool isLinearizableVector(VectorType type);
/// instead of explicit masks.
///
/// 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
+/// it and from `indices`: dimension i is in-bounds when the map result is an
+/// AffineDimExpr pointing to a static memref dimension divisible by the vector
+/// size whose index is a known multiple of 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,
@@ -250,7 +251,8 @@ Value createReadOrMaskedRead(OpBuilder &builder, Location loc, Value source,
/// instead of explicit masks.
/// `writeIndices` specifies the offsets to use. If empty, all indices are set
/// to 0. When `permutationMap` is provided, the in_bounds attribute is
-/// inferred from the map instead of the destination shape.
+/// inferred from the map and `writeIndices` instead of the destination shape,
+/// as for createReadOrMaskedRead.
Operation *createWriteOrMaskedWrite(OpBuilder &builder, Location loc,
Value vecToStore, Value dest,
SmallVector<Value> writeIndices = {},
diff --git a/mlir/lib/Dialect/Vector/Utils/VectorUtils.cpp b/mlir/lib/Dialect/Vector/Utils/VectorUtils.cpp
index c38213850c6ea..6a2fa308f5a45 100644
--- a/mlir/lib/Dialect/Vector/Utils/VectorUtils.cpp
+++ b/mlir/lib/Dialect/Vector/Utils/VectorUtils.cpp
@@ -420,19 +420,69 @@ 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) {
+static bool isKnownMultipleOf(Value value, int64_t factor);
+
+/// Returns true if every result of `map` applied to `operands` is known to be
+/// a multiple of `factor`. Operands that are themselves multiples of `factor`
+/// are substituted with `factor * d` so that AffineExpr::isMultipleOf can see
+/// through them.
+static bool isKnownMultipleOf(AffineMap map, ValueRange operands,
+ int64_t factor) {
+ MLIRContext *ctx = map.getContext();
+ SmallVector<AffineExpr> dimRepls, symRepls;
+ for (unsigned i = 0, e = map.getNumDims(); i < e; ++i) {
+ AffineExpr dim = getAffineDimExpr(i, ctx);
+ dimRepls.push_back(isKnownMultipleOf(operands[i], factor) ? dim * factor
+ : dim);
+ }
+ for (unsigned i = 0, e = map.getNumSymbols(); i < e; ++i) {
+ AffineExpr sym = getAffineSymbolExpr(i, ctx);
+ symRepls.push_back(isKnownMultipleOf(operands[map.getNumDims() + i], factor)
+ ? sym * factor
+ : sym);
+ }
+ return llvm::all_of(map.getResults(), [&](AffineExpr expr) {
+ return expr.replaceDimsAndSymbols(dimRepls, symRepls).isMultipleOf(factor);
+ });
+}
+
+/// Returns true if `value` is known to be a multiple of `factor`: a constant,
+/// an affine.for induction variable whose step and lower bound are multiples
+/// of `factor`, or an affine.apply whose map preserves that property. Loop
+/// bounds and apply operands are followed recursively, so an induction
+/// variable whose lower bound is the induction variable of an outer loop
+/// stepping by a multiple of `factor` is recognized.
+static bool isKnownMultipleOf(Value value, int64_t factor) {
+ if (std::optional<int64_t> cst = getConstantIntValue(value))
+ return *cst % factor == 0;
+ if (affine::AffineForOp forOp = affine::getForInductionVarOwner(value))
+ return forOp.getStepAsInt() % factor == 0 &&
+ isKnownMultipleOf(forOp.getLowerBoundMap(),
+ forOp.getLowerBoundOperands(), factor);
+ if (auto applyOp = value.getDefiningOp<affine::AffineApplyOp>())
+ return isKnownMultipleOf(applyOp.getAffineMap(), applyOp.getMapOperands(),
+ factor);
+ return false;
+}
+
+/// Compute the in_bounds attribute for a transfer op given its permutation map,
+/// the source being accessed and the indices into it. Dimension i is in-bounds
+/// when the map result is an AffineDimExpr pointing to a static source
+/// dimension divisible by the vector size whose index is a known multiple of
+/// the vector size, or an AffineConstantExpr (broadcast).
+static SmallVector<bool>
+computeInBoundsFromPermutationMap(AffineMap permutationMap,
+ VectorType vectorType, ShapedType sourceType,
+ ValueRange indices) {
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();
+ int64_t vectorSize = vectorType.getDimSize(i);
if (!sourceType.isDynamicDim(memDim) &&
- sourceType.getDimSize(memDim) % vectorType.getDimSize(i) == 0)
+ sourceType.getDimSize(memDim) % vectorSize == 0 &&
+ isKnownMultipleOf(indices[memDim], vectorSize))
inBounds[i] = true;
} else if (isa<AffineConstantExpr>(expr)) {
inBounds[i] = true;
@@ -471,14 +521,23 @@ Value vector::createReadOrMaskedRead(OpBuilder &builder, Location loc,
padValue.value().getType() == sourceShapedType.getElementType()) &&
"expected same pad element type to match source element type");
+ // 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());
+
SmallVector<bool> inBoundsVal(vecToReadRank, true);
if (useInBoundsInsteadOfMasking) {
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()));
+ permutationMap, vecToReadTy, sourceShapedType, indices);
} else {
// Update the inBounds attribute.
// FIXME: This computation is too weak - it ignores the read indices.
@@ -487,15 +546,6 @@ 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;
- 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 =
@@ -539,14 +589,23 @@ Operation *vector::createWriteOrMaskedWrite(OpBuilder &builder, Location loc,
int64_t vecToStoreRank = vecToStoreType.getRank();
auto vecToStoreShape = vecToStoreType.getShape();
+ // If missing, initialize the write indices to 0.
+ bool useDefaultWriteIdxs = writeIndices.empty();
+ assert((useDefaultWriteIdxs ||
+ writeIndices.size() == static_cast<size_t>(destRank)) &&
+ "Invalid number of write indices!");
+ if (useDefaultWriteIdxs) {
+ auto zero = arith::ConstantIndexOp::create(builder, loc, 0);
+ writeIndices.assign(destRank, zero);
+ }
+
// Compute the in_bounds attribute
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<ShapedType>(dest.getType()));
+ permutationMap, vecToStoreType, destType, writeIndices);
} else {
// Update the inBounds attribute.
// FIXME: This computation is too weak - it ignores the write indices.
@@ -557,16 +616,6 @@ Operation *vector::createWriteOrMaskedWrite(OpBuilder &builder, Location loc,
}
}
- // If missing, initialize the write indices to 0.
- bool useDefaultWriteIdxs = writeIndices.empty();
- assert((useDefaultWriteIdxs ||
- writeIndices.size() == static_cast<size_t>(destRank)) &&
- "Invalid number of write indices!");
- if (useDefaultWriteIdxs) {
- auto zero = arith::ConstantIndexOp::create(builder, loc, 0);
- writeIndices.assign(destRank, zero);
- }
-
// Generate the xfer_write Op. A null permutation map means the builder
// defaults to a minor identity map.
Operation *write =
diff --git a/mlir/test/Dialect/Affine/SuperVectorize/vectorize_1d_inbounds.mlir b/mlir/test/Dialect/Affine/SuperVectorize/vectorize_1d_inbounds.mlir
new file mode 100644
index 0000000000000..83642685a678d
--- /dev/null
+++ b/mlir/test/Dialect/Affine/SuperVectorize/vectorize_1d_inbounds.mlir
@@ -0,0 +1,103 @@
+// RUN: mlir-opt %s -affine-super-vectorize="virtual-vector-size=8" -split-input-file | FileCheck %s
+
+// The dimension is divisible by the vector size, but the index is offset by
+// one, so the last vector runs past the end of the memref.
+
+// CHECK-LABEL: func.func @offset_index
+// CHECK: affine.for %{{.*}} = 0 to 15 step 8 {
+// CHECK: vector.transfer_read %{{.*}}[%{{.*}}], %{{[0-9]+}} : memref<16xf32>, vector<8xf32>
+// CHECK: vector.transfer_write %{{.*}}, %{{.*}}[%{{[a-z0-9]+}}] : vector<8xf32>, memref<16xf32>
+func.func @offset_index(%A: memref<16xf32>, %B: memref<16xf32>) {
+ affine.for %i = 0 to 15 {
+ %v = affine.load %A[%i + 1] : memref<16xf32>
+ affine.store %v, %B[%i + 1] : memref<16xf32>
+ }
+ return
+}
+
+// -----
+
+// An offset that is a multiple of the vector size keeps the accesses aligned.
+
+// CHECK-LABEL: func.func @aligned_offset_index
+// CHECK: affine.for %{{.*}} = 0 to 8 step 8 {
+// CHECK: vector.transfer_read %{{.*}}[%{{.*}}], %{{.*}} {in_bounds = [true]} : memref<16xf32>, vector<8xf32>
+// CHECK: vector.transfer_write %{{.*}}, %{{.*}}[%{{.*}}] {in_bounds = [true]} : vector<8xf32>, memref<16xf32>
+func.func @aligned_offset_index(%A: memref<16xf32>, %B: memref<16xf32>) {
+ affine.for %i = 0 to 8 {
+ %v = affine.load %A[%i + 8] : memref<16xf32>
+ affine.store %v, %B[%i + 8] : memref<16xf32>
+ }
+ return
+}
+
+// -----
+
+// A lower bound that is not a multiple of the vector size shifts every vector
+// off alignment.
+
+// CHECK-LABEL: func.func @unaligned_lower_bound
+// CHECK: affine.for %{{.*}} = 4 to 16 step 8 {
+// CHECK: vector.transfer_read %{{.*}}[%{{.*}}], %{{[0-9]+}} : memref<16xf32>, vector<8xf32>
+// CHECK: vector.transfer_write %{{.*}}, %{{.*}}[%{{[a-z0-9]+}}] : vector<8xf32>, memref<16xf32>
+func.func @unaligned_lower_bound(%A: memref<16xf32>, %B: memref<16xf32>) {
+ affine.for %i = 4 to 16 {
+ %v = affine.load %A[%i] : memref<16xf32>
+ affine.store %v, %B[%i] : memref<16xf32>
+ }
+ return
+}
+
+// -----
+
+// CHECK-LABEL: func.func @aligned_lower_bound
+// CHECK: affine.for %{{.*}} = 8 to 16 step 8 {
+// CHECK: vector.transfer_read %{{.*}}[%{{.*}}], %{{.*}} {in_bounds = [true]} : memref<16xf32>, vector<8xf32>
+// CHECK: vector.transfer_write %{{.*}}, %{{.*}}[%{{.*}}] {in_bounds = [true]} : vector<8xf32>, memref<16xf32>
+func.func @aligned_lower_bound(%A: memref<16xf32>, %B: memref<16xf32>) {
+ affine.for %i = 8 to 16 {
+ %v = affine.load %A[%i] : memref<16xf32>
+ affine.store %v, %B[%i] : memref<16xf32>
+ }
+ return
+}
+
+// -----
+
+// The lower bound of the vectorized loop is the induction variable of a loop
+// stepping by a multiple of the vector size, as after tiling.
+
+// CHECK-LABEL: func.func @lower_bound_from_outer_loop
+// CHECK: affine.for %{{.*}} = 0 to 32 step 16 {
+// CHECK: affine.for %{{.*}} = #{{.*}}(%{{.*}}) to #{{.*}}(%{{.*}}) step 8 {
+// CHECK: vector.transfer_read %{{.*}}[%{{.*}}], %{{.*}} {in_bounds = [true]} : memref<32xf32>, vector<8xf32>
+// CHECK: vector.transfer_write %{{.*}}, %{{.*}}[%{{.*}}] {in_bounds = [true]} : vector<8xf32>, memref<32xf32>
+func.func @lower_bound_from_outer_loop(%A: memref<32xf32>, %B: memref<32xf32>) {
+ affine.for %i = 0 to 32 step 16 {
+ affine.for %ii = affine_map<(d0) -> (d0)>(%i) to affine_map<(d0) -> (d0 + 16)>(%i) {
+ %v = affine.load %A[%ii] : memref<32xf32>
+ affine.store %v, %B[%ii] : memref<32xf32>
+ }
+ }
+ return
+}
+
+// -----
+
+// Same as above, but the outer loop steps by 12, so the inner lower bound is
+// not a multiple of the vector size.
+
+// CHECK-LABEL: func.func @unaligned_lower_bound_from_outer_loop
+// CHECK: affine.for %{{.*}} = 0 to 24 step 12 {
+// CHECK: affine.for %{{.*}} = #{{.*}}(%{{.*}}) to #{{.*}}(%{{.*}}) step 8 {
+// CHECK: vector.transfer_read %{{.*}}[%{{.*}}], %{{[0-9]+}} : memref<24xf32>, vector<8xf32>
+// CHECK: vector.transfer_write %{{.*}}, %{{.*}}[%{{[a-z0-9]+}}] : vector<8xf32>, memref<24xf32>
+func.func @unaligned_lower_bound_from_outer_loop(%A: memref<24xf32>, %B: memref<24xf32>) {
+ affine.for %i = 0 to 24 step 12 {
+ affine.for %ii = affine_map<(d0) -> (d0)>(%i) to affine_map<(d0) -> (d0 + 12)>(%i) {
+ %v = affine.load %A[%ii] : memref<24xf32>
+ affine.store %v, %B[%ii] : memref<24xf32>
+ }
+ }
+ return
+}
diff --git a/mlir/test/Dialect/Affine/SuperVectorize/vectorize_affine_apply.mlir b/mlir/test/Dialect/Affine/SuperVectorize/vectorize_affine_apply.mlir
index 08ebea37b159d..2dec94451cdbd 100644
--- a/mlir/test/Dialect/Affine/SuperVectorize/vectorize_affine_apply.mlir
+++ b/mlir/test/Dialect/Affine/SuperVectorize/vectorize_affine_apply.mlir
@@ -43,7 +43,7 @@ func.func @vec_affine_apply_2(%arg0: memref<8x12x16xf32>, %arg1: memref<8x24x48x
// CHECK-NEXT: affine.for %[[ARG4:.*]] = 0 to 48 step 8 {
// CHECK-NEXT: %[[S0:.*]] = affine.apply #[[$MAP_ID2]](%[[ARG4]])
// CHECK-NEXT: %[[CST:.*]] = ub.poison : f32
-// CHECK-NEXT: %[[S1:.*]] = vector.transfer_read %[[ARG0]][%[[ARG2]], %[[ARG3]], %[[S0]]], %[[CST]] {in_bounds = [true]} : memref<8x12x16xf32>, vector<8xf32>
+// CHECK-NEXT: %[[S1:.*]] = vector.transfer_read %[[ARG0]][%[[ARG2]], %[[ARG3]], %[[S0]]], %[[CST]] : memref<8x12x16xf32>, vector<8xf32>
// CHECK-NEXT: vector.transfer_write %[[S1]], %[[ARG1]][%[[ARG2]], %[[ARG3]], %[[ARG4]]] {in_bounds = [true]} : vector<8xf32>, memref<8x24x48xf32>
// CHECK-NEXT: }
// CHECK-NEXT: }
``````````
</details>
https://github.com/llvm/llvm-project/pull/221268
More information about the Mlir-commits
mailing list