[Mlir-commits] [mlir] [mlir][affine] Fix crash in super-vectorizer when inner loop bounds depend on outer loop IV (PR #184618)

[Mehdi Amini]

https://github.com/joker-eph updated https://github.com/llvm/llvm-project/pull/184618

>From b7bcf8022847d97f0578bb9a1857251006387cf5 Mon Sep 17 00:00:00 2001
From: Mehdi Amini <joker.eph at gmail.com>
Date: Wed, 11 Mar 2026 05:54:31 -0700
Subject: [PATCH] [mlir][affine] Bail out when store permutation map has
 broadcast dimensions

When the vectorized loop is an outer loop and the store index uses an
inner loop's IV (which is invariant w.r.t. the outer loop), makePermutationMap
produces a broadcast map (e.g. (d0) -> (0)). A vector.transfer_write with
broadcast dimensions is invalid, causing a verifier error.

Fix: check for AffineConstantExpr results in the store permutation map and
bail out of vectorization, preserving the scalar loop nest.
---
 .../Affine/Transforms/SuperVectorize.cpp      |  9 +++++++
 .../SuperVectorize/vectorize_unsupported.mlir | 25 +++++++++++++++++++
 2 files changed, 34 insertions(+)

diff --git a/mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp b/mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp
index ef5fc64857c81..c90d9bd8730e6 100644
--- a/mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp
+++ b/mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp
@@ -1312,6 +1312,15 @@ static Operation *vectorizeAffineStore(AffineStoreOp storeOp,
   LLVM_DEBUG(dbgs() << "\n[early-vect]+++++ permutationMap: ");
   LLVM_DEBUG(permutationMap.print(dbgs()));
 
+  // 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(),
+                   llvm::IsaPred<AffineConstantExpr>)) {
+    LLVM_DEBUG(dbgs() << "\n[early-vect]+++++ store permutation map has "
+                         "broadcast dims, bailing out\n");
+    return nullptr;
+  }
+
   auto transfer = vector::TransferWriteOp::create(
       state.builder, storeOp.getLoc(), vectorValue, storeOp.getMemRef(),
       indices, permutationMap);
diff --git a/mlir/test/Dialect/Affine/SuperVectorize/vectorize_unsupported.mlir b/mlir/test/Dialect/Affine/SuperVectorize/vectorize_unsupported.mlir
index 8ed152a5b0e20..331252bb9fe4b 100644
--- a/mlir/test/Dialect/Affine/SuperVectorize/vectorize_unsupported.mlir
+++ b/mlir/test/Dialect/Affine/SuperVectorize/vectorize_unsupported.mlir
@@ -43,6 +43,31 @@ func.func @iv_mapped_to_multiple_indices_unsupported(%arg0: index) -> memref<2x2
 
 // -----
 
+// Regression test: when the store's permutation map has a broadcast dimension
+// (because the index is invariant w.r.t. the vectorized loop), vectorization
+// must bail out gracefully instead of emitting an invalid transfer_write.
+
+#map_id = affine_map<(d0) -> (d0)>
+#map_id_p1 = affine_map<(d0) -> (d0 + 1)>
+
+// CHECK-LABEL: func.func @store_broadcast_perm_map_unsupported
+// CHECK:         affine.for
+// CHECK:           affine.for
+// CHECK:             affine.load
+// CHECK:             affine.store
+// CHECK-NOT:     vector.transfer_write
+func.func @store_broadcast_perm_map_unsupported(%arg0: memref<4x4xf32>, %arg1: memref<4xf32>) {
+  affine.for %i = 0 to 4 {
+    affine.for %j = #map_id(%i) to #map_id_p1(%i) {
+      %0 = affine.load %arg0[%j, %j] : memref<4x4xf32>
+      affine.store %0, %arg1[%j] : memref<4xf32>
+    }
+  }
+  return
+}
+
+// -----
+
 // Regression test for https://github.com/llvm/llvm-project/issues/128334
 // The vectorizer test utility used to crash when a reduction loop with a
 // dynamic upper bound was vectorized via 'vectorizeAffineLoopNest', because