[Mlir-commits] [mlir] [mlir] Fix visitor block erasure (PR #205854)

[llvmlistbot at llvm.org]

https://github.com/mygitljf created https://github.com/llvm/llvm-project/pull/205854

I made the visitor test drop block-defined value uses before erasing the block and added a small regression case.
Fixes #205717 

>From baa3384fa8e47c30d4740d6a0963902c8094091c Mon Sep 17 00:00:00 2001
From: mygitljf <2410316423 at qq.com>
Date: Fri, 26 Jun 2026 00:26:57 +0000
Subject: [PATCH] [mlir][IR] Fix visitor block erasure

---
 mlir/test/IR/visitors.mlir        | 26 ++++++++++++++++++++++++++
 mlir/test/lib/IR/TestVisitors.cpp |  1 +
 2 files changed, 27 insertions(+)

diff --git a/mlir/test/IR/visitors.mlir b/mlir/test/IR/visitors.mlir
index e56c67922a662..705e353d27e6d 100644
--- a/mlir/test/IR/visitors.mlir
+++ b/mlir/test/IR/visitors.mlir
@@ -444,3 +444,29 @@ func.func @graph_region_skip(%fill: tensor<2xf32>, %output: tensor<2xf32>) {
   }
   return
 }
+
+// -----
+
+// Regression test for https://github.com/llvm/llvm-project/issues/205717:
+// Block erasure in skip callbacks must drop uses of values defined in the block
+// before erasing it.
+// CHECK-LABEL: gpu.module @issue205717
+module attributes {gpu.container_module, llvm.data_layout = ""} {
+  gpu.module @issue205717 {
+    func.func @test_skip_block_erasure_nested_uses(%arg0: i32) -> i32 {
+      %c0 = arith.constant 0 : i32
+      %c1 = arith.constant 1 : i32
+      %cond = arith.cmpi eq, %arg0, %c0 : i32
+      cf.cond_br %cond, ^bb1, ^bb2
+    ^bb1:
+      gpu.barrier
+      cf.br ^bb3(%c1 : i32)
+    ^bb2:
+      gpu.barrier
+      %plus = arith.addi %arg0, %c1 : i32
+      cf.br ^bb3(%plus : i32)
+    ^bb3(%result: i32):
+      return %result : i32
+    }
+  }
+}
diff --git a/mlir/test/lib/IR/TestVisitors.cpp b/mlir/test/lib/IR/TestVisitors.cpp
index 566833ec559fa..148a57f2a739b 100644
--- a/mlir/test/lib/IR/TestVisitors.cpp
+++ b/mlir/test/lib/IR/TestVisitors.cpp
@@ -136,6 +136,7 @@ static void testSkipErasureCallbacks(Operation *op) {
       llvm::outs() << "Erasing ";
       printBlock(block);
       llvm::outs() << "\n";
+      block->dropAllDefinedValueUses();
       block->erase();
       return WalkResult::skip();
     }