[Mlir-commits] [mlir] e3bae38 - [mlir][affine] Fix crash when rolling back vectorization of reduction loops (#226912)

[llvmlistbot at llvm.org]

Author: Mattéo Rizza Murgier
Date: 2026-09-30T09:33:51+02:00
New Revision: e3bae380eb23076b081ad7db1e8f6a36c5261eb3

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

LOG: [mlir][affine] Fix crash when rolling back vectorization of reduction loops (#226912)

Reductions introduce operations outside of the created loop, that are
users of said loop. When rolling back (e.g. when an op inside the loop
being vectorized is not vectorizable), only the created loop is deleted,
which leads to the error: “operation destroyed but still has uses”.

This PR fixes this by recursively deleting users of the newly created
loop before removing it when rolling back.

Added: 
    

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

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp b/mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp
index 6c6e8ed979d23..0bb344b0484f0 100644
--- a/mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp
+++ b/mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp
@@ -891,6 +891,16 @@ static void eraseLoopNest(AffineForOp forOp) {
   forOp.erase();
 }
 
+/// Recursively erases the users of the results of 'op'.
+static void eraseUsers(Operation *op) {
+  while (!op->use_empty()) {
+    Operation *user = *op->user_begin();
+    eraseUsers(user);
+    LLVM_DEBUG(dbgs() << "[early-vect]+++++ erasing user:\n" << *user << "\n");
+    user->erase();
+  }
+}
+
 /// Erases the scalar loop nest after its successful vectorization.
 void VectorizationState::finishVectorizationPattern(AffineForOp rootLoop) {
   LLVM_DEBUG(dbgs() << "\n[early-vect] Finalizing vectorization\n");
@@ -1670,10 +1680,13 @@ vectorizeLoopNest(std::vector<SmallVector<AffineForOp, 2>> &loops,
   if (opVecResult.wasInterrupted()) {
     LLVM_DEBUG(dbgs() << "[early-vect]+++++ failed vectorization for: "
                       << rootLoop << "\n");
-    // Erase vector loop nest if it was created.
+    // Erase vector loop nest if it was created, and all its users.
     auto vecRootLoopIt = state.opVectorReplacement.find(rootLoop);
-    if (vecRootLoopIt != state.opVectorReplacement.end())
-      eraseLoopNest(cast<AffineForOp>(vecRootLoopIt->second));
+    if (vecRootLoopIt != state.opVectorReplacement.end()) {
+      auto vecRootLoop = cast<AffineForOp>(vecRootLoopIt->second);
+      eraseUsers(vecRootLoop);
+      eraseLoopNest(vecRootLoop);
+    }
 
     return failure();
   }

diff  --git a/mlir/test/Dialect/Affine/SuperVectorize/vectorize_reduction.mlir b/mlir/test/Dialect/Affine/SuperVectorize/vectorize_reduction.mlir
index 9dacc316f1072..836e7b0c71f1f 100644
--- a/mlir/test/Dialect/Affine/SuperVectorize/vectorize_reduction.mlir
+++ b/mlir/test/Dialect/Affine/SuperVectorize/vectorize_reduction.mlir
@@ -606,6 +606,25 @@ func.func @vecdim_none_reduction_rejected(%in: memref<256x512xf32>, %out: memref
 
 // -----
 
+// The index_cast operation cannot be vectorized: vectorization is rolled back
+// and both loop nest and reduction op are removed.
+
+func.func @vecdim_reduction_rollback(%in: memref<512xf32>, %init: f32) -> f32 {
+ %final_red = affine.for %j = 0 to 512 iter_args(%red_iter = %init) -> (f32) {
+   %ld = affine.load %in[%j] : memref<512xf32>
+   %j_i32 = arith.index_cast %j : index to i32
+   %add = arith.addf %red_iter, %ld : f32
+   affine.yield %add : f32
+ }
+ return %final_red : f32
+}
+
+// CHECK-LABEL: @vecdim_reduction_rollback
+// CHECK:       affine.for {{.*}} iter_args
+// CHECK-NOT:   vector.reduction
+
+// -----
+
 // The number of iterations is not divisable by the vector size, so a mask has
 // to be applied to the last update of the accumulator.