[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 <[email protected]>

Reviewed By: srishti-pm

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

Author: srishti-pm

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;

mlir/test/Transforms/remove-dead-values.mlir

@@ -207,20 +207,24 @@ 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):
 // CHECK-NEXT:      scf.yield %[[arg5]], %[[arg6]] : i32, i32
 // 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
+}
 
 // -----