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

[Mehdi Amini]

https://github.com/joker-eph created https://github.com/llvm/llvm-project/pull/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

>From baf8c7c91cdc9c566bcbc037b8095a610190ad0d Mon Sep 17 00:00:00 2001
From: Mehdi Amini <joker.eph at gmail.com>
Date: Wed, 4 Mar 2026 05:07:33 -0800
Subject: [PATCH] [mlir][affine] Fix crash in affine-super-vectorize for index
 constants inside loops

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
---
 .../Affine/Transforms/SuperVectorize.cpp      | 13 +++++++++++
 .../Affine/SuperVectorize/vectorize_1d.mlir   | 22 +++++++++++++++++++
 2 files changed, 35 insertions(+)

diff --git a/mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp b/mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp
index 4e9f10b2c525c..6117e6f3a81d4 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
+}