[Mlir-commits] [mlir] Fix Bug in RemoveDeadValues Pass (PR #148437)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: None (ronigoldman22)

<details>
<summary>Changes</summary>

This patch fixes a bug in the RemoveDeadValues pass where unused function arguments were not removed from the function signature in an edge case where the function returns void.
A corresponding test was added to the MLIR LIT test suite to cover this case.

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


2 Files Affected:

- (modified) mlir/lib/Transforms/RemoveDeadValues.cpp (+2-2) 
- (modified) mlir/test/Transforms/remove-dead-values.mlir (+40) 


``````````diff
diff --git a/mlir/lib/Transforms/RemoveDeadValues.cpp b/mlir/lib/Transforms/RemoveDeadValues.cpp
index 608bdcb948176..ddd5f2ba1a7b7 100644
--- a/mlir/lib/Transforms/RemoveDeadValues.cpp
+++ b/mlir/lib/Transforms/RemoveDeadValues.cpp
@@ -345,8 +345,6 @@ static void processFuncOp(FunctionOpInterface funcOp, Operation *module,
   // since it forwards only to non-live value(s) (%1#1).
   Operation *lastReturnOp = funcOp.back().getTerminator();
   size_t numReturns = lastReturnOp->getNumOperands();
-  if (numReturns == 0)
-    return;
   BitVector nonLiveRets(numReturns, true);
   for (SymbolTable::SymbolUse use : uses) {
     Operation *callOp = use.getUser();
@@ -368,6 +366,8 @@ static void processFuncOp(FunctionOpInterface funcOp, Operation *module,
   cl.functions.push_back({funcOp, nonLiveArgs, nonLiveRets});
 
   // Do (5) and (6).
+  if (numReturns == 0)
+    return;
   for (SymbolTable::SymbolUse use : uses) {
     Operation *callOp = use.getUser();
     assert(isa<CallOpInterface>(callOp) && "expected a call-like user");
diff --git a/mlir/test/Transforms/remove-dead-values.mlir b/mlir/test/Transforms/remove-dead-values.mlir
index 3af95db3c0e24..dcaeefc3f2a6f 100644
--- a/mlir/test/Transforms/remove-dead-values.mlir
+++ b/mlir/test/Transforms/remove-dead-values.mlir
@@ -548,3 +548,43 @@ func.func @test_atomic_yield(%I: memref<10xf32>, %idx : index) {
   func.return
 }
 
+// CHECK-LABEL: module @return_void_with_unused_argument {
+// CHECK-LABEL:  func.func private @noop() {
+// CHECK-NEXT:    return
+// CHECK-NEXT:  }
+// CHECK-LABEL:  func.func private @fn_return_void_with_unused_argument(%arg0: memref<1x1xsi8, 1>, %arg1: memref<1x1xsi16, 1>) {
+// CHECK-NEXT:    %alloc = memref.alloc() {alignment = 8 : i64} : memref<1x73xsi8, 3>
+// CHECK-NEXT:    %alloc_0 = memref.alloc() {alignment = 8 : i64} : memref<1x1xsi16, 3>
+// CHECK-NEXT:    %alloc_1 = memref.alloc() {alignment = 8 : i64} : memref<1x1xsi8, 3>
+// CHECK-NEXT:    memref.copy %arg0, %alloc_1 : memref<1x1xsi8, 1> to memref<1x1xsi8, 3>
+// CHECK-NEXT:    memref.copy %arg1, %alloc_0 : memref<1x1xsi16, 1> to memref<1x1xsi16, 3>
+// CHECK-NEXT:    call @noop() : () -> ()
+// CHECK-NEXT:    return
+// CHECK-NEXT:  }
+// CHECK-LABEL:  func.func @main(%arg0: memref<1x73xsi8, 1>, %arg1: memref<1x1xsi8, 1>, %arg2: memref<1x1xsi16, 1>) -> memref<1x73xsi8, 1> {
+// CHECK-NEXT:    %alloc = memref.alloc() : memref<1x73xsi8, 1>
+// CHECK-NEXT:    call @fn_return_void_with_unused_argument(%arg1, %arg2) : (memref<1x1xsi8, 1>, memref<1x1xsi16, 1>) -> ()
+// CHECK-NEXT:    return %alloc : memref<1x73xsi8, 1>
+// CHECK-NEXT:  }
+// CHECK-NEXT:}
+module @return_void_with_unused_argument {
+  func.func private @noop(%arg0: memref<1x1xsi8, 3>, %arg1: memref<1x1xsi16, 3>, %arg2: memref<1x73xsi8, 1>) {
+    return
+  }
+  func.func private @fn_return_void_with_unused_argument(%arg0: memref<1x73xsi8, 1> , %arg1: memref<1x1xsi8, 1> , %arg2: memref<1x1xsi16, 1> , %arg3: memref<1x73xsi8, 1>) {
+    %alloc = memref.alloc() {alignment = 8 : i64} : memref<1x73xsi8, 3>
+    %alloc_0 = memref.alloc() {alignment = 8 : i64} : memref<1x1xsi16, 3>
+    %alloc_1 = memref.alloc() {alignment = 8 : i64} : memref<1x1xsi8, 3>
+    memref.copy %arg1, %alloc_1 : memref<1x1xsi8, 1> to memref<1x1xsi8, 3>
+    memref.copy %arg2, %alloc_0 : memref<1x1xsi16, 1> to memref<1x1xsi16, 3>
+    call @noop(%alloc_1, %alloc_0, %arg0) : 
+        (memref<1x1xsi8, 3>, memref<1x1xsi16, 3>, memref<1x73xsi8, 1>) -> ()
+    return
+  }
+  func.func @main(%arg0: memref<1x73xsi8, 1> , %arg1: memref<1x1xsi8, 1> , %arg2: memref<1x1xsi16, 1>) -> (memref<1x73xsi8, 1> ) {
+    %alloc = memref.alloc() : memref<1x73xsi8, 1>
+    call @fn_return_void_with_unused_argument(%arg0, %arg1, %arg2, %alloc) : (memref<1x73xsi8, 1>, memref<1x1xsi8, 1>, memref<1x1xsi16, 1>, memref<1x73xsi8, 1>) -> ()
+    return %alloc : memref<1x73xsi8, 1>
+  }
+}
+

``````````

</details>


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