[Mlir-commits] [mlir] [mlir][affine] Skip nested duplicate loop matches in the super-vectorizer strategy (PR #224850)
Stanislav Bardyuk
llvmlistbot at llvm.org
Sat Sep 19 11:12:55 PDT 2026
https://github.com/kodlan created https://github.com/llvm/llvm-project/pull/224850
`-affine-super-vectorize=virtual-vector-size=32,256` asserts on this input:
```mlir
func.func @m(%A: memref<8x64xf32>) {
affine.for %i = 0 to 8 {
affine.for %j = 0 to 64 {
affine.for %k = 0 to 1 {
%0 = affine.load %A[%i, %j] : memref<8x64xf32>
%1 = arith.addf %0, %0 : f32
affine.store %1, %A[%i, %j] : memref<8x64xf32>
}
}
}
return
}
```
```
VectorUtils.cpp:469: Assertion `(!permutationMap || permutationMap.getNumResults() == static_cast<size_t>(vecToReadRank)) && "expected permutation map results to match vector rank."' failed.
```
The nested pattern matcher walks the whole subtree of a matched loop, so for the 2-D pattern the match rooted at `%i` reports both `%j` and `%k` as inner loops (`%k` is parallel because it runs once, and it is not used in any index). `analyzeProfitability` then maps both of them to vector dimension 1. When the load inside `%k` is vectorized, `makePermutationMap` sees three enclosing vectorized loops and builds a map with three results for a `vector<32x256xf32>`.
This only shows up when the innermost match (`%j` with `%k`) fails first, which it does since #184618 bails out on the store whose permutation map would have a broadcast dimension (the store does not use `%k`). The driver then tries the match rooted at `%i` and hits the bad strategy.
The fix skips, at each pattern level, a matched loop that is nested inside another matched loop of the same level, so only the outermost candidates get a vector dimension. The input above now vectorizes `%i` by 32 and `%j` by 256 and leaves `%k` as a scalar loop, the same output shape `vectorize_matmul` in `vectorize_2d.mlir` has for its 3-deep nest (there the third loop is a reduction and never matches, so it did not hit this). Sibling matches are unaffected. The public `vectorizeAffineLoopNest` entry already rejects such a nesting through `verifyLoopNesting`; the pass driver never checked it, and only the strategy map matters for `vectorizeLoopNest`, so the fix goes there.
The same assertion is also reachable in the other direction, when a load or store is enclosed by fewer vectorized loops than the vector rank (memory ops directly under the outer loop of a 2-D match). `vectorizeAffineLoad` and `vectorizeAffineStore` now fail the op when the permutation map's result count does not match the vector rank, so the match fails instead of the pass aborting. Both shapes have a test.
Fixes #220566
>From 691ee5499220c482453d35536f4eb91525cb8f6d Mon Sep 17 00:00:00 2001
From: Stanislav Bardyuk <sbardyuk at google.com>
Date: Sat, 19 Sep 2026 17:57:39 +0000
Subject: [PATCH] [mlir][affine] Skip nested duplicate loop matches in the
super-vectorizer strategy
The nested pattern matcher walks the whole subtree of a matched loop, so a
2-D match rooted at an outer loop reports every deeper parallel loop as an
inner match and analyzeProfitability maps all of them to the same vector
dimension. When the innermost match fails first (store broadcast bail-out),
the outer match vectorizes three loops for a 2-D vector and
makePermutationMap builds a map with too many results, which asserts in
createReadOrMaskedRead.
Skip a matched loop that is nested inside another match at the same level so
only the outermost candidates get a vector dimension. Also fail the load or
store vectorization when the permutation map's result count differs from the
vector rank, which covers memory ops enclosed by fewer vectorized loops than
the vector rank.
Fixes #220566
---
.../Affine/Transforms/SuperVectorize.cpp | 24 ++++++++++-
.../Affine/SuperVectorize/vectorize_2d.mlir | 41 +++++++++++++++++++
2 files changed, 63 insertions(+), 2 deletions(-)
diff --git a/mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp b/mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp
index 6c6e8ed979d23..f417615a0e5f9 100644
--- a/mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp
+++ b/mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp
@@ -659,12 +659,17 @@ static LogicalResult analyzeProfitability(ArrayRef<NestedMatch> matches,
unsigned patternDepth,
VectorizationStrategy *strategy) {
for (auto m : matches) {
+ Operation *loop = m.getMatchedOperation();
+ // Skip loops nested in another match at this level (same vector dim).
+ if (llvm::any_of(matches, [loop](const NestedMatch &other) {
+ return other.getMatchedOperation()->isProperAncestor(loop);
+ }))
+ continue;
if (failed(analyzeProfitability(m.getMatchedChildren(), depthInPattern + 1,
patternDepth, strategy))) {
return failure();
}
- vectorizeLoopIfProfitable(m.getMatchedOperation(), depthInPattern,
- patternDepth, strategy);
+ vectorizeLoopIfProfitable(loop, depthInPattern, patternDepth, strategy);
}
return success();
}
@@ -1265,6 +1270,13 @@ static Operation *vectorizeAffineLoad(AffineLoadOp loadOp,
LLVM_DEBUG(dbgs() << "\n[early-vect]+++++ permutationMap: ");
LLVM_DEBUG(permutationMap.print(dbgs()));
+ // One map result per enclosing vectorized loop; must match the vector rank.
+ if (permutationMap.getNumResults() != vectorType.getRank()) {
+ LLVM_DEBUG(dbgs() << "\n[early-vect]+++++ permutationMap rank does not "
+ "match the vector rank\n");
+ return nullptr;
+ }
+
Value transferVal = createReadOrMaskedRead(
state.builder, loadOp.getLoc(), loadOp.getMemRef(), vectorType,
/*padValue=*/std::nullopt, /*useInBoundsInsteadOfMasking=*/true, indices,
@@ -1314,6 +1326,14 @@ static Operation *vectorizeAffineStore(AffineStoreOp storeOp,
LLVM_DEBUG(dbgs() << "\n[early-vect]+++++ permutationMap: ");
LLVM_DEBUG(permutationMap.print(dbgs()));
+ // Same rank check as for loads.
+ if (permutationMap.getNumResults() !=
+ cast<VectorType>(vectorValue.getType()).getRank()) {
+ LLVM_DEBUG(dbgs() << "\n[early-vect]+++++ permutationMap rank does not "
+ "match the vector rank\n");
+ return nullptr;
+ }
+
// A transfer_write with a broadcast dimension (constant expr in the
// permutation map) is invalid. Bail out to avoid producing invalid IR.
if (llvm::any_of(permutationMap.getResults(),
diff --git a/mlir/test/Dialect/Affine/SuperVectorize/vectorize_2d.mlir b/mlir/test/Dialect/Affine/SuperVectorize/vectorize_2d.mlir
index eb5120a49e3d4..f3782d9ca9a05 100644
--- a/mlir/test/Dialect/Affine/SuperVectorize/vectorize_2d.mlir
+++ b/mlir/test/Dialect/Affine/SuperVectorize/vectorize_2d.mlir
@@ -143,3 +143,44 @@ func.func @vectorize_matmul(%arg0: memref<?x?xf32>, %arg1: memref<?x?xf32>, %arg
}
return
}
+
+// The innermost loop runs once and no index uses it, so it is parallel and is
+// matched as an inner loop just like %i1. The match rooted at %i0 used to map
+// %i1 and %i2 to the same vector dimension and crash on the load.
+// VECT-LABEL: func @vectorize_nested_parallel_loops
+// VECT: affine.for %[[I0:.*]] = 0 to 8 step 4 {
+// VECT-NEXT: affine.for %[[I1:.*]] = 0 to 64 step 8 {
+// VECT-NEXT: affine.for %{{.*}} = 0 to 1 {
+// VECT: %[[LD:.*]] = vector.transfer_read %{{.*}}[%[[I0]], %[[I1]]], %{{.*}} {{.*}}: memref<8x64xf32>, vector<4x8xf32>
+// VECT-NEXT: %[[ADD:.*]] = arith.addf %[[LD]], %[[LD]] : vector<4x8xf32>
+// VECT-NEXT: vector.transfer_write %[[ADD]], %{{.*}}[%[[I0]], %[[I1]]] {{.*}}: vector<4x8xf32>, memref<8x64xf32>
+func.func @vectorize_nested_parallel_loops(%A : memref<8x64xf32>) {
+ affine.for %i0 = 0 to 8 {
+ affine.for %i1 = 0 to 64 {
+ affine.for %i2 = 0 to 1 {
+ %0 = affine.load %A[%i0, %i1] : memref<8x64xf32>
+ %1 = arith.addf %0, %0 : f32
+ affine.store %1, %A[%i0, %i1] : memref<8x64xf32>
+ }
+ }
+ }
+ return
+}
+
+// The memory ops are only nested in %i0, but the match rooted at %i0 assigns
+// the second vector dimension to %i1, so their permutation maps have a single
+// result for a 2-D vector. Vectorization must fail instead of asserting.
+// VECT-LABEL: func @vectorize_memory_ops_outside_inner_loop
+// VECT-NOT: vector.transfer_read
+// VECT-NOT: vector.transfer_write
+func.func @vectorize_memory_ops_outside_inner_loop(%A : memref<8xf32>, %B : memref<8xf32>) {
+ affine.for %i0 = 0 to 8 {
+ %0 = affine.load %A[%i0] : memref<8xf32>
+ affine.store %0, %B[%i0] : memref<8xf32>
+ affine.for %i1 = 0 to 64 {
+ %c1 = arith.constant 1 : index
+ %1 = arith.addi %i1, %c1 : index
+ }
+ }
+ return
+}
More information about the Mlir-commits
mailing list