[Mlir-commits] [mlir] [mlir] Avoid removing returns for mismatch call results (PR #207587)

[Darwin Morris]

https://github.com/darwinmorris updated https://github.com/llvm/llvm-project/pull/207587

>From 4f922bc244704f56c2ca3922c0f21c91a3949319 Mon Sep 17 00:00:00 2001
From: darwinmorris <darwinbmorris98 at gmail.com>
Date: Sun, 5 Jul 2026 15:31:14 +0200
Subject: [PATCH] [mlir] Avoid removing returns for mismatch call results

---
 mlir/lib/Transforms/RemoveDeadValues.cpp     | 14 +++++++++++++-
 mlir/test/Transforms/remove-dead-values.mlir | 19 +++++++++++++++++++
 2 files changed, 32 insertions(+), 1 deletion(-)

diff --git a/mlir/lib/Transforms/RemoveDeadValues.cpp b/mlir/lib/Transforms/RemoveDeadValues.cpp
index f0a210a2ededb..0786ac2a391a2 100644
--- a/mlir/lib/Transforms/RemoveDeadValues.cpp
+++ b/mlir/lib/Transforms/RemoveDeadValues.cpp
@@ -335,14 +335,26 @@ static void processFuncOp(FunctionOpInterface funcOp, Operation *module,
   // can make our optimization strong by even removing a live return value (%0),
   // since it forwards only to non-live value(s) (%1#1).
   size_t numReturns = funcOp.getNumResults();
+  bool canCleanReturnValues = true;
   BitVector nonLiveRets(numReturns, true);
+
   for (SymbolTable::SymbolUse use : uses) {
     Operation *callOp = use.getUser();
     assert(isa<CallOpInterface>(callOp) && "expected a call-like user");
+
+    if (callOp->getNumResults() != numReturns) {
+      canCleanReturnValues = false;
+      nonLiveRets.reset();
+      break;
+    }
+
     BitVector liveCallRets = markLives(callOp->getResults(), nonLiveSet, la);
     nonLiveRets &= liveCallRets.flip();
   }
 
+  if (!canCleanReturnValues)
+    nonLiveArgs.reset();
+    
   // Note that in the absence of control flow ops forcing the control to go from
   // the entry (first) block to the other blocks, the control never reaches any
   // block other than the entry block, because every block has a terminator.
@@ -358,7 +370,7 @@ static void processFuncOp(FunctionOpInterface funcOp, Operation *module,
   cl.functions.push_back({funcOp, nonLiveArgs, nonLiveRets});
 
   // Do (5) and (6).
-  if (numReturns == 0)
+  if (!canCleanReturnValues || numReturns == 0)
     return;
   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 64088ce15cd48..080743e0539e8 100644
--- a/mlir/test/Transforms/remove-dead-values.mlir
+++ b/mlir/test/Transforms/remove-dead-values.mlir
@@ -868,3 +868,22 @@ module @func_with_non_call_users {
   }
   spirv.EntryPoint "GLCompute" @callee
 }
+
+// -----
+
+// Check that function cleanup is conservative when a call-like op does not
+// expose the same number of results as the callee.
+func.func private @call_and_store_callee(%arg0: memref<f32>) -> memref<f32> {
+  // CHECK-LABEL: func.func private @call_and_store_callee
+  // CHECK-SAME: (%[[ARG0:.*]]: memref<f32>) -> memref<f32>
+  // CHECK: return %[[ARG0]] : memref<f32>
+  return %arg0 : memref<f32>
+}
+
+func.func @call_and_store_before(%arg0: memref<f32>) -> memref<f32> {
+  // CHECK-LABEL: func.func @call_and_store_before
+  // CHECK: test.call_and_store
+  // CHECK: return
+  test.call_and_store @call_and_store_callee(%arg0), %arg0 {store_before_call = true, tag_name = "call"} : (memref<f32>, memref<f32>) -> ()
+  return {tag = "return"} %arg0 : memref<f32>
+}