From 5842a0f76c31fc3cc5049e035003489cd2f730ae Mon Sep 17 00:00:00 2001
From: Mehdi Amini <joker.eph@gmail.com>
Date: Wed, 11 Mar 2026 16:59:38 +0100
Subject: [PATCH] [mlir][affine] Bail out when store permutation map has
 broadcast dimensions (#184618)

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.

Fixes #131135

Assisted-by: Claude Code
---
 .../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 ef5fc64857c8..c90d9bd8730e 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 8ed152a5b0e2..331252bb9fe4 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