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

[srishti-pm]

Summary

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 <[e-mail]>

[srishti-pm]

I'm planning on accepting and landing this change soon because it is trivial and crucial.

[srishti-pm]

Rebase to main.

[mehdi_amini]

Did you get this reviewed? It's quite unexpected to see "self approved" patch (actually I didn't even know it was possible).

Also when you fix a GitHub issue, add Fixes #64997 to the commit message to automatically close the issue.

mlir/lib/Transforms/RemoveDeadValues.cpp

143

The temp block is a cute trick, but it clobbered the intent here, this is not the kind of idiom you will find in MLIR.

The right API here should be, I believe, newOp->getRegion(index).takeBody(region);

[mehdi_amini]

mlir/lib/Transforms/RemoveDeadValues.cpp

143

Ping here?

[srishti-pm]

mlir/lib/Transforms/RemoveDeadValues.cpp

143

I'm currently in school, busy with schoolwork. I will make this minor modification when I find time.

mlir/lib/Transforms/RemoveDeadValues.cpp

static void dropUsesAndEraseResults(Operation *op, BitVector toErase) {
   builder.setInsertionPointAfter(op);
   OperationState state(op->getLoc(), op->getName().getStringRef(),
                        op->getOperands(), newResultTypes, op->getAttrs());
   for (unsigned i = 0, e = op->getNumRegions(); i < e; ++i)
     state.addRegion();
   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;
   for (auto [index, result] : llvm::enumerate(op->getResults())) {
     assert(result && "expected result to be non-null");
     if (toErase[index]) {
       result.dropAllUses();
     } else {

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

 //  and %non_live_1 which are not live and have no other predecessors.
 //
 // Note that this cleanup cannot be done by the `canonicalize` pass.
 //
 // CHECK:       func.func @clean_region_branch_op_remove_last_2_results_last_2_arguments_and_last_operand(%[[arg2:.*]]: i1) -> i32 {
 // 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):
     %non_live_4 = arith.addi %arg7, %arg8 :i32
     scf.yield %arg5, %arg6, %non_live_4 : i32, i32, i32
   }
   return %live : i32
 }
+func.func private @identity(%arg1 : i32) -> (i32) {
+  return %arg1 : i32
+}
 
 // -----
 
 // The op isn't erased because it has memory effects but its unnecessary result
 // is removed.
 //
 // Note that this cleanup cannot be done by the `canonicalize` pass.
 //