[Mlir-commits] [mlir] [mlir][affine] Skip nested duplicate loop matches in the super-vectorizer strategy (PR #224850)

Sat Sep 19 11:13:34 PDT 2026

llvmorg-github-actions[bot] wrote:




@llvm/pr-subscribers-mlir-affine

Author: Stanislav Bardyuk (kodlan)

<details>
<summary>Changes</summary>

`-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

---
Full diff: https://github.com/llvm/llvm-project/pull/224850.diff


2 Files Affected:

- (modified) mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp (+22-2) 
- (modified) mlir/test/Dialect/Affine/SuperVectorize/vectorize_2d.mlir (+41) 


``````````diff

diff --git a/mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp b/mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp
index 6c6e8ed979d235..f417615a0e5f9b 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 eb5120a49e3d4b..f3782d9ca9a053 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
+}

``````````

</details>


https://github.com/llvm/llvm-project/pull/224850
