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

[Mattéo Rizza Murgier]

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

>From c43c0f6c9225337e14e8627de86e03f4a759a0cc Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Matt=C3=A9o=20Rizza=20Murgier?=
 <matteo.rizza-murgier at sipearl.com>
Date: Sun, 27 Sep 2026 15:32:57 +0200
Subject: [PATCH] [mlir][affine] Fix crash when rolling back vectorization of
 reduction loops

---
 .../Affine/Transforms/SuperVectorize.cpp      | 18 ++++++++++++++---
 .../SuperVectorize/vectorize_reduction.mlir   | 20 +++++++++++++++++++
 2 files changed, 35 insertions(+), 3 deletions(-)

diff --git a/mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp b/mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp
index 6c6e8ed979d23..c5b323591764f 100644
--- a/mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp
+++ b/mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp
@@ -891,6 +891,15 @@ 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);
+    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 +1679,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..736e40a88a746 100644
--- a/mlir/test/Dialect/Affine/SuperVectorize/vectorize_reduction.mlir
+++ b/mlir/test/Dialect/Affine/SuperVectorize/vectorize_reduction.mlir
@@ -606,6 +606,26 @@ 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>) -> f32 {
+ %cst = arith.constant 0.000000e+00 : f32
+ %final_red = affine.for %j = 0 to 512 iter_args(%red_iter = %cst) -> (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.