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

[Prathamesh Tagore]

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

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 fact that the current block might be providing global context to other blocks in the same parent op (see the attached test for an example of this) even if no one is pointing to it. Erasing this block would then lead to invalid IR. We solve this by dropping uses of ops inside such a block which belong in the same parent region op. This way, we can safely erase the block if nothing outside the surrounding region still needs it.

Fixes https://github.com/llvm/llvm-project/issues/178099

>From 51a935d63fdc00e4dd88956093076f9fa4626e3a Mon Sep 17 00:00:00 2001
From: Prathamesh Tagore <prathameshtagore at gmail.com>
Date: Wed, 11 Feb 2026 23:10:06 +0100
Subject: [PATCH] [mlir][test-ir-visitors] Fix noSkipBlockErasure for blocks
 providing global context

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 fact that
the current block might be providing global context to other blocks in the
same parent op (see the attached test for an example of this) even if no one
is pointing to it. Erasing this block would then lead to invalid IR. We solve
this by dropping uses of ops inside such a block which belong in the same
parent region op. This way, we can safely erase the block if nothing outside
the surrounding region still needs it.
---
 mlir/test/IR/visitors.mlir        | 21 +++++++++++++++++++++
 mlir/test/lib/IR/TestVisitors.cpp | 15 +++++++++++++++
 2 files changed, 36 insertions(+)

diff --git a/mlir/test/IR/visitors.mlir b/mlir/test/IR/visitors.mlir
index ec7712a45d388..93204f6a73530 100644
--- a/mlir/test/IR/visitors.mlir
+++ b/mlir/test/IR/visitors.mlir
@@ -382,3 +382,24 @@ 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) in non-post-order fashion and dropping inside op uses appropriately.
+// CHECK-LABEL: Op post-order erasures (no skip)
+func.func @test_no_skip_block_erasure() {
+  %c0 = arith.constant 0 : index
+  %c3 = arith.constant 3 : index
+  cf.br ^bb1
+^bb1:
+  %c1 = arith.constant 1 : index
+  %c128 = arith.constant 128 : index
+  %c256 = arith.constant 256 : index
+  %cond = arith.cmpi eq, %c0, %c3 : index
+  cf.cond_br %cond, ^bb2, ^bb4
+^bb2:
+  cf.br ^bb4
+^bb4:
+  return
+}
diff --git a/mlir/test/lib/IR/TestVisitors.cpp b/mlir/test/lib/IR/TestVisitors.cpp
index 7df0ddb7ca27e..542c1bb11f2df 100644
--- a/mlir/test/lib/IR/TestVisitors.cpp
+++ b/mlir/test/lib/IR/TestVisitors.cpp
@@ -184,6 +184,21 @@ static void testNoSkipErasureCallbacks(Operation *op) {
       llvm::outs() << "Erasing ";
       printBlock(block);
       llvm::outs() << "\n";
+
+      // Only drop intra-region uses of ops inside the block (they don't have a
+      // post-order visit expectation here). We only expect to visit the block
+      // parent ops in post-order, so we can safely erase the block if nothing
+      // outside the surrounding region still needs it.
+      Operation *blockParentOp = block->getParentOp();
+      for (Operation &op : *block) {
+        for (OpOperand &use : llvm::make_early_inc_range(op.getUses())) {
+          Operation *userRegionHolder = use.getOwner()->getParentOp();
+          // Drop only when both ops live under the same region holder.
+          if (userRegionHolder == blockParentOp)
+            use.drop();
+        }
+      }
+
       block->erase();
     } else {
       llvm::outs() << "Cannot erase ";