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

[Darwin Morris]

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

Fixes #204588

`RemoveDeadValues` assumes that a call operation exposes the same number of results as its callee. This is not true for the `call_and_store` test operation, which consumes callee results internally instead of exposing them as operation results.

This results in an error when the pass computes return value cleanup information, as it does not handle the mismatch. Now this is handled by disabling return-value cleanup when there is a result-count mismatch.

Regression tests were added based on the reproducer. 

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