[Mlir-commits] [mlir] [mlir][test-ir-visitors] Fix noSkipBlockErasure for blocks providing global context (PR #181320)

[Prathamesh Tagore]

https://github.com/meshtag updated https://github.com/llvm/llvm-project/pull/181320

>From 3f7bede72e1eaf8c239a23e919a0483bd0ec2763 Mon Sep 17 00:00:00 2001
From: Prathamesh Tagore <prathameshtagore at gmail.com>
Date: Tue, 17 Feb 2026 21:48:18 +0100
Subject: [PATCH] [mlir][test-ir-visitors] Fix erasure of operations with
 intact uses

The pass checks if something is pointing to the current block in
testNoSkipErasureCallbacks() inside noSkipBlockErasure() callback and erases
it if no one is pointing to it. However, it doesn't consider the case where
the current block might be providing global context to other blocks in the
same parent region even if no one is pointing to it (see the attached test
for an example of this). Erasing this block would then lead to invalid IR. We
solve this by dropping intra-block op uses from ops within the same parent
region.
---
 mlir/test/IR/visitors.mlir        | 17 +++++++++++++++++
 mlir/test/lib/IR/TestVisitors.cpp | 18 ++++++++++++++++++
 2 files changed, 35 insertions(+)

diff --git a/mlir/test/IR/visitors.mlir b/mlir/test/IR/visitors.mlir
index ec7712a45d388..31d5494fecb7a 100644
--- a/mlir/test/IR/visitors.mlir
+++ b/mlir/test/IR/visitors.mlir
@@ -382,3 +382,20 @@ func.func @unordered_cfg_with_loop() {
 // CHECK:       Cannot erase block ^bb2 from region 0 from operation 'regionOp0', still has uses
 // CHECK:       Cannot erase block ^bb3 from region 0 from operation 'regionOp0', still has uses
 // CHECK:       Cannot erase block ^bb4 from region 0 from operation 'regionOp0', still has uses
+
+// -----
+
+// The following test should not crash while visiting the intra-op blocks (inside the top level 
+// function in this case). We are testing that the intra-block ops are erased after dropping their
+// uses from ops with same parent region.
+// CHECK-LABEL: func.func @test_no_skip_block_erasure
+func.func @test_no_skip_block_erasure() {
+  %c0 = arith.constant 0 : index
+  %c3 = arith.constant 3 : index
+  cf.br ^bb1
+^bb1:
+  %cond = arith.cmpi eq, %c0, %c3 : index
+  cf.br ^bb4
+^bb4:
+  return
+}
diff --git a/mlir/test/lib/IR/TestVisitors.cpp b/mlir/test/lib/IR/TestVisitors.cpp
index 7df0ddb7ca27e..2667001ee10a7 100644
--- a/mlir/test/lib/IR/TestVisitors.cpp
+++ b/mlir/test/lib/IR/TestVisitors.cpp
@@ -184,6 +184,24 @@ static void testNoSkipErasureCallbacks(Operation *op) {
       llvm::outs() << "Erasing ";
       printBlock(block);
       llvm::outs() << "\n";
+
+      // Only drop uses from operations within the same parent region. This
+      // avoids erasing operations with their uses still intact and eliminates
+      // such crashes.
+      // Note: We do not drop the use when the parent region is different for
+      // it, because this means that the use's region holding op is a child of
+      // the region holding op containing the current block and was expected to
+      // be visited and erased first - we should correctly fail here.
+      Region *blockParentRegion = block->front().getParentRegion();
+      for (Operation &op : *block) {
+        for (OpOperand &use : llvm::make_early_inc_range(op.getUses())) {
+          // Early continue if the parent regions are not same.
+          if (blockParentRegion != use.getOwner()->getParentRegion())
+            continue;
+          use.drop();
+        }
+      }
+
       block->erase();
     } else {
       llvm::outs() << "Cannot erase ";