[Mlir-commits] [mlir] 7370091 - [mlir][test-ir-visitors] Fix noSkipBlockErasure crash with block args used across blocks (#183828)

[llvmlistbot at llvm.org]

Author: Mehdi Amini
Date: 2026-02-28T11:56:36+01:00
New Revision: 7370091a43e5cbcb4cf56fb83b62e056903cb918

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

LOG: [mlir][test-ir-visitors] Fix noSkipBlockErasure crash with block args used across blocks (#183828)

The noSkipBlockErasure callback in TestVisitors.cpp dropped uses of op
results within the same region before erasing a block, but did not drop
uses of the block's own arguments (e.g. function entry block arguments).
When the block was subsequently erased its block arguments were
destroyed while their use-lists were still non-empty, triggering the
assertion in IRObjectWithUseList::~IRObjectWithUseList().

Fix this by also iterating over the block's arguments and dropping any
uses that belong to the same parent region. This mirrors the existing
logic for op result uses and makes the block-erasure walk handle IRs
where function arguments are consumed by ops in sibling blocks.

Also replace `block->front().getParentRegion()` with
`block->getParent()` for robustness (avoids UB when the block has no
ops).

Add a regression test based on the reproducer from
https://github.com/llvm/llvm-project/issues/182996.

Fixes #182996

Added: 
    

Modified: 
    mlir/test/IR/visitors.mlir
    mlir/test/lib/IR/TestVisitors.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/test/IR/visitors.mlir b/mlir/test/IR/visitors.mlir
index 0e6ac879f5b94..93380cc96ab0a 100644
--- a/mlir/test/IR/visitors.mlir
+++ b/mlir/test/IR/visitors.mlir
@@ -385,7 +385,7 @@ func.func @unordered_cfg_with_loop() {
 
 // -----
 
-// The following test should not crash while visiting the intra-op blocks (inside the top level 
+// 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
@@ -399,3 +399,16 @@ func.func @test_no_skip_block_erasure() {
 ^bb4:
   return
 }
+
+// -----
+
+// Regression test for https://github.com/llvm/llvm-project/issues/182996:
+// Block erasure must also drop uses of block arguments (e.g. function args)
+// from sibling blocks in the same region before destroying the block.
+// CHECK-LABEL: func.func @test_no_skip_block_erasure_block_args
+func.func @test_no_skip_block_erasure_block_args(%arg0: i32, %arg1: i32) -> i32 {
+  cf.br ^bb1
+^bb1:
+  %0 = arith.addi %arg0, %arg1 : i32
+  return %0 : i32
+}

diff  --git a/mlir/test/lib/IR/TestVisitors.cpp b/mlir/test/lib/IR/TestVisitors.cpp
index 2667001ee10a7..dec5140d170fc 100644
--- a/mlir/test/lib/IR/TestVisitors.cpp
+++ b/mlir/test/lib/IR/TestVisitors.cpp
@@ -192,7 +192,7 @@ static void testNoSkipErasureCallbacks(Operation *op) {
       // 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();
+      Region *blockParentRegion = block->getParent();
       for (Operation &op : *block) {
         for (OpOperand &use : llvm::make_early_inc_range(op.getUses())) {
           // Early continue if the parent regions are not same.
@@ -202,6 +202,17 @@ static void testNoSkipErasureCallbacks(Operation *op) {
         }
       }
 
+      // Also drop uses of block arguments that are consumed within the same
+      // region. This handles cases like function entry blocks whose arguments
+      // are used by ops in sibling blocks.
+      for (BlockArgument arg : block->getArguments()) {
+        for (OpOperand &use : llvm::make_early_inc_range(arg.getUses())) {
+          if (blockParentRegion != use.getOwner()->getParentRegion())
+            continue;
+          use.drop();
+        }
+      }
+
       block->erase();
     } else {
       llvm::outs() << "Cannot erase ";