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

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-affine

Author: Mehdi Amini (joker-eph)

<details>
<summary>Changes</summary>

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

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


2 Files Affected:

- (modified) mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp (+13) 
- (modified) mlir/test/Dialect/Affine/SuperVectorize/vectorize_1d.mlir (+22) 


``````````diff
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
+}

``````````

</details>


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