[Mlir-commits] [mlir] b6bab6d - [MLIR][transforms] Fix `cloneInto()` error in `RemoveDeadValues` pass

[Srishti Srivastava]

Author: Srishti Srivastava
Date: 2023-08-26T19:50:24Z
New Revision: b6bab6db9bea2793634bec2e33385f3d89a5e7f3

URL: https://github.com/llvm/llvm-project/commit/b6bab6db9bea2793634bec2e33385f3d89a5e7f3
DIFF: https://github.com/llvm/llvm-project/commit/b6bab6db9bea2793634bec2e33385f3d89a5e7f3.diff

LOG: [MLIR][transforms] Fix `cloneInto()` error in `RemoveDeadValues` pass

This commit fixes an error in the `RemoveDeadValues` pass that is
associated with its incorrect usage of the `cloneInto()` function.

The `setOperands()` function that is used by the `cloneInto()` function
requires all operands to not be null. But, that is not possible in this
pass because we drop uses of dead values, thus making them null. It is
only at the end of the pass that we are assured that such null values
won't exist but during the execution of the pass, there could be null
values.

To fix this, we replace the usage of the `cloneInto()` function to copy
a region with `moveBlock()` to move each block of the region one by one.
This function does not require the presence of non-null values and is
thus the right choice here. This implementation is also more opttimized
because we are moving things instead of copying them. The goal was
always moving.

Signed-off-by: Srishti Srivastava <srishtisrivastava.ai at gmail.com>

Reviewed By: srishti-pm

Differential Revision: https://reviews.llvm.org/D158941

Added: 
    

Modified: 
    mlir/lib/Transforms/RemoveDeadValues.cpp
    mlir/test/Transforms/remove-dead-values.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Transforms/RemoveDeadValues.cpp b/mlir/lib/Transforms/RemoveDeadValues.cpp
index ad49a10b30b5ae..ce19dc667f009d 100644
--- a/mlir/lib/Transforms/RemoveDeadValues.cpp
+++ b/mlir/lib/Transforms/RemoveDeadValues.cpp
@@ -135,8 +135,12 @@ static void dropUsesAndEraseResults(Operation *op, BitVector toErase) {
   Operation *newOp = builder.create(state);
   for (const auto &[index, region] : llvm::enumerate(op->getRegions())) {
     Region &newRegion = newOp->getRegion(index);
-    IRMapping mapping;
-    region.cloneInto(&newRegion, mapping);
+    // Move all blocks of `region` into `newRegion`.
+    Block *temp = new Block();
+    newRegion.push_back(temp);
+    while (!region.empty())
+      region.front().moveBefore(temp);
+    temp->erase();
   }
 
   unsigned indexOfNextNewCallOpResultToReplace = 0;

diff  --git a/mlir/test/Transforms/remove-dead-values.mlir b/mlir/test/Transforms/remove-dead-values.mlir
index 22b66b464ac69a..600810a785b1f9 100644
--- a/mlir/test/Transforms/remove-dead-values.mlir
+++ b/mlir/test/Transforms/remove-dead-values.mlir
@@ -207,6 +207,7 @@ func.func @clean_region_branch_op_dont_remove_first_2_results_but_remove_first_o
 // CHECK-NEXT:    %[[c0:.*]] = arith.constant 0
 // CHECK-NEXT:    %[[c1:.*]] = arith.constant 1
 // CHECK-NEXT:    %[[live_and_non_live:.*]]:2 = scf.while (%[[arg3:.*]] = %[[c0]], %[[arg4:.*]] = %[[c1]]) : (i32, i32) -> (i32, i32) {
+// CHECK-NEXT:      func.call @identity() : () -> ()
 // CHECK-NEXT:      scf.condition(%[[arg2]]) %[[arg4]], %[[arg3]] : i32, i32
 // CHECK-NEXT:    } do {
 // CHECK-NEXT:    ^bb0(%[[arg5:.*]]: i32, %[[arg6:.*]]: i32):
@@ -214,13 +215,16 @@ func.func @clean_region_branch_op_dont_remove_first_2_results_but_remove_first_o
 // CHECK-NEXT:    }
 // CHECK-NEXT:    return %[[live_and_non_live]]#0 : i32
 // CHECK-NEXT:  }
+// CHECK:       func.func private @identity() {
+// CHECK-NEXT:    return
+// CHECK-NEXT:  }
 func.func @clean_region_branch_op_remove_last_2_results_last_2_arguments_and_last_operand(%arg2: i1) -> (i32) {
   %c0 = arith.constant 0 : i32
   %c1 = arith.constant 1 : i32
   %c2 = arith.constant 2 : i32
   %live, %non_live, %non_live_0, %non_live_1 = scf.while (%arg3 = %c0, %arg4 = %c1, %arg10 = %c2) : (i32, i32, i32) -> (i32, i32, i32, i32) {
     %non_live_2 = arith.addi %arg10, %arg10 : i32
-    %non_live_3 = arith.muli %arg10, %arg10 : i32
+    %non_live_3 = func.call @identity(%arg10) : (i32) -> (i32)
     scf.condition(%arg2) %arg4, %arg3, %non_live_2, %non_live_3 : i32, i32, i32, i32
   } do {
   ^bb0(%arg5: i32, %arg6: i32, %arg7: i32, %arg8: i32):
@@ -229,6 +233,9 @@ func.func @clean_region_branch_op_remove_last_2_results_last_2_arguments_and_las
   }
   return %live : i32
 }
+func.func private @identity(%arg1 : i32) -> (i32) {
+  return %arg1 : i32
+}
 
 // -----