[Mlir-commits] [mlir] [mlir][Transforms] Allow RemoveDeadValues to process a function which the last block is not the exit. (PR #156123)

[xin liu]

https://github.com/navyxliu created https://github.com/llvm/llvm-project/pull/156123

'processFuncOp' queries the number of returned values of a function using the terminator of the last block's getNumOperands(). It presumes the last block is the exit. It is not always the case. We encounter a function that the last block jumps backward and then exit. append it as a lit test.

If processFuncOp gets the wrong number, it can't delete a dead returned value. here is the stacktrace.
```
```

This patch fixes the bug by querying from FunctionInterfaceOp directly.

>From 10c2506630f1d0e1a4c6f1a61844219b572a11cf Mon Sep 17 00:00:00 2001
From: Xin Liu <xxinliu at meta.com>
Date: Fri, 29 Aug 2025 15:42:20 -0700
Subject: [PATCH] [mlir][Transforms] Allow RemoveDeadValues to process a
 function which the last block is not the exit.

'processFuncOp' queries the number of returned values of a function using the terminator of the last block's getNumOperands().
It presumes the last block is the exit. It is not always the case. We encounter a function that the last block jumps backward and then exit. append it as a lit test.

This patch fixes the bug by querying from FunctionInterfaceOp directly.
If processFuncOp gets the wrong number, it can't delete a dead returned value. here is the stacktrace.
```
 Assertion failed: (op->getNumResults() == toErase.size() && "expected the number of results in `op` and the size of `toErase` to " "be the same"), function dropUsesAndEraseResults, file RemoveDeadValues.cpp, line 200.
 PLEASE submit a bug report to https://github.com/llvm/llvm-project/issues/ and include the crash backtrace.
 Stack dump:
 0.	Program arguments: /Users/xxinliu/Devel/llvm-project/build/bin/mlir-opt /Users/xxinliu/Devel/llvm-project/mlir/test/Transforms/remove-dead-values.mlir -remove-dead-values -split-input-file -verify-diagnostics
 Stack dump without symbol names (ensure you have llvm-symbolizer in your PATH or set the environment var `LLVM_SYMBOLIZER_PATH` to point to it):
 0  mlir-opt                 0x00000001044b4f68 llvm::sys::PrintStackTrace(llvm::raw_ostream&, int) + 56
 1  mlir-opt                 0x00000001044b2bfc llvm::sys::RunSignalHandlers() + 172
 2  mlir-opt                 0x00000001044b5a60 SignalHandler(int, __siginfo*, void*) + 360
 3  libsystem_platform.dylib 0x000000018f0fd6a4 _sigtramp + 56
 4  libsystem_pthread.dylib  0x000000018f0c388c pthread_kill + 296
 5  libsystem_c.dylib        0x000000018efcca3c abort + 124
 6  libsystem_c.dylib        0x000000018efcbc70 err + 0
 7  mlir-opt                 0x000000010c372148 (anonymous namespace)::cleanUpDeadVals((anonymous namespace)::RDVFinalCleanupList&) (.cold.3) + 0
 8  mlir-opt                 0x000000010c36a484 (anonymous namespace)::cleanUpDeadVals((anonymous namespace)::RDVFinalCleanupList&) + 14920
 9  mlir-opt                 0x000000010c36675c (anonymous namespace)::RemoveDeadValues::runOnOperation() + 292
```
---
 mlir/lib/Transforms/RemoveDeadValues.cpp     |  3 +--
 mlir/test/Transforms/remove-dead-values.mlir | 23 ++++++++++++++++++++
 2 files changed, 24 insertions(+), 2 deletions(-)

diff --git a/mlir/lib/Transforms/RemoveDeadValues.cpp b/mlir/lib/Transforms/RemoveDeadValues.cpp
index 02dad69e49614..0e84b6dd17f29 100644
--- a/mlir/lib/Transforms/RemoveDeadValues.cpp
+++ b/mlir/lib/Transforms/RemoveDeadValues.cpp
@@ -344,8 +344,7 @@ static void processFuncOp(FunctionOpInterface funcOp, Operation *module,
   // being returned, in order to optimize our IR. So, this demonstrates how we
   // can make our optimization strong by even removing a live return value (%0),
   // since it forwards only to non-live value(s) (%1#1).
-  Operation *lastReturnOp = funcOp.back().getTerminator();
-  size_t numReturns = lastReturnOp->getNumOperands();
+  size_t numReturns = funcOp.getNumResults();
   BitVector nonLiveRets(numReturns, true);
   for (SymbolTable::SymbolUse use : uses) {
     Operation *callOp = use.getUser();
diff --git a/mlir/test/Transforms/remove-dead-values.mlir b/mlir/test/Transforms/remove-dead-values.mlir
index 0f8d757086e87..fa2c145bd3701 100644
--- a/mlir/test/Transforms/remove-dead-values.mlir
+++ b/mlir/test/Transforms/remove-dead-values.mlir
@@ -592,3 +592,26 @@ module @dynamically_unreachable {
     return
   }
 }
+
+// CHECK-LABEL: module @last_block_not_exit
+module @last_block_not_exit {
+  // return value can be removed because it's private.
+  func.func private @terminated_with_condbr(%arg0: i1, %arg1: i1) -> i1 {
+    %true = arith.constant true
+    %false = arith.constant false
+    cf.cond_br %arg0, ^bb1(%false : i1), ^bb2
+  ^bb1(%1: i1):  // 2 preds: ^bb0, ^bb2
+    return %1 : i1
+  ^bb2:  // pred: ^bb3
+    cf.cond_br %arg1, ^bb1(%false : i1), ^bb1(%true : i1)
+  }
+
+  func.func public @call_private_but_not_use() {
+    %i0 = arith.constant 0: i1
+    %i1 = arith.constant 1: i1
+    call @terminated_with_condbr(%i0, %i1) : (i1, i1) -> i1
+    func.return
+  }
+  // CHECK-LABEL: @call_private_but_not_use
+  // CHECK: call @terminated_with_condbr(%false, %true) : (i1, i1)
+}