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

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Darwin Morris (darwinmorris)

<details>
<summary>Changes</summary>

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. 

---
Full diff: https://github.com/llvm/llvm-project/pull/207587.diff


2 Files Affected:

- (modified) mlir/lib/Transforms/RemoveDeadValues.cpp (+13-1) 
- (modified) mlir/test/Transforms/remove-dead-values.mlir (+19) 


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

``````````

</details>


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