[Mlir-commits] [mlir] 3b8cd6c - [mlir][affine] Fix crash in affine-super-vectorize for index constants inside loops (#184614)

[llvmlistbot at llvm.org]

Author: Mehdi Amini
Date: 2026-03-11T13:36:45+01:00
New Revision: 3b8cd6c5287ebf17397f25e0e5e8513ebdea7bd1

URL: https://github.com/llvm/llvm-project/commit/3b8cd6c5287ebf17397f25e0e5e8513ebdea7bd1
DIFF: https://github.com/llvm/llvm-project/commit/3b8cd6c5287ebf17397f25e0e5e8513ebdea7bd1.diff

LOG: [mlir][affine] Fix crash in affine-super-vectorize for index constants inside loops (#184614)

When an arith.constant of index type is defined inside the loop body
being vectorized, vectorizeConstant creates a vector<Nxindex> constant
and registers it as the vector replacement. However,
getScalarValueReplacementsFor (used by vectorizeAffineStore to compute
indices for vector.transfer_write) looks only in the scalar replacement
map. With no scalar replacement registered for the index constant, it
falls back to the original scalar value, which is erased when the scalar
loop is cleaned up. This results in a "operation destroyed but still has
uses" crash.

Fix: when vectorizeConstant processes an index-typed constant, also
create a new scalar constant in the vector loop body and register it as
the scalar replacement. This ensures that memory operation index
computation can find a live value in the vectorized IR.

Fixes #122213

Assisted-by: Claude Code

Added: 
    

Modified: 
    mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp
    mlir/test/Dialect/Affine/SuperVectorize/vectorize_1d.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp b/mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp
index b4653cdd2a174..ef5fc64857c81 100644
--- a/mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp
+++ b/mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp
@@ -965,6 +965,19 @@ static arith::ConstantOp vectorizeConstant(arith::ConstantOp constOp,
 
   // Register vector replacement for future uses in the scope.
   state.registerOpVectorReplacement(constOp, newConstOp);
+
+  // Index-typed constants are also used as scalar indices in vectorized memory
+  // operations (e.g., as operands to vector.transfer_read/write). Register a
+  // scalar replacement so that getScalarValueReplacementsFor can find a live
+  // value in the vectorized loop body instead of falling back to the original
+  // constant, which will be erased along with the scalar loop.
+  if (isa<IndexType>(scalarTy)) {
+    auto scalarConstOp = arith::ConstantOp::create(
+        state.builder, constOp.getLoc(), constOp.getValue());
+    state.registerValueScalarReplacement(constOp.getResult(),
+                                         scalarConstOp.getResult());
+  }
+
   return newConstOp;
 }
 

diff  --git a/mlir/test/Dialect/Affine/SuperVectorize/vectorize_1d.mlir b/mlir/test/Dialect/Affine/SuperVectorize/vectorize_1d.mlir
index 72ced5b53879b..f9593221e1843 100644
--- a/mlir/test/Dialect/Affine/SuperVectorize/vectorize_1d.mlir
+++ b/mlir/test/Dialect/Affine/SuperVectorize/vectorize_1d.mlir
@@ -700,3 +700,25 @@ func.func @vec_non_scalar_type() {
 
 // CHECK-LABEL: @vec_non_scalar_type
 // CHECK-NOT: vector
+
+// -----
+
+// Regression test: an arith.constant of index type defined inside the loop
+// body must have a scalar replacement registered so that it can be used as
+// a scalar index in the vectorized memory transfer op. Without the fix,
+// affine-super-vectorize crashes with "operation destroyed but still has uses".
+
+// CHECK-LABEL: @index_const_inside_loop
+// CHECK:       affine.for %[[IV:.*]] = 0 to 8 step 128 {
+// CHECK-DAG:     %[[CST_VEC:.*]] = arith.constant dense<0> : vector<128xi32>
+// CHECK-DAG:     %[[C0:.*]] = arith.constant 0 : index
+// CHECK:         vector.transfer_write %[[CST_VEC]], %{{.*}}[%[[C0]], %[[IV]]] : vector<128xi32>, memref<1x8xi32>
+// CHECK:       }
+func.func @index_const_inside_loop(%mem: memref<1x8xi32>) {
+  affine.for %arg0 = 0 to 8 {
+    %c0_i32 = arith.constant 0 : i32
+    %c0 = arith.constant 0 : index
+    affine.store %c0_i32, %mem[%c0, %arg0] : memref<1x8xi32>
+  }
+  return
+}