[mlir][Transforms] Allow RemoveDeadValues to process a function whose the last block is not the exit. (#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. 
This patch fixes the bug by querying from FunctionInterfaceOp directly.

Author: navyxliu

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();

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)
+}