Offset the new arguments by the number of old arguments

Author: giuseros

mlir/lib/Transforms/Utils/RegionUtils.cpp

@@ -679,12 +679,15 @@ static bool ableToUpdatePredOperands(Block *block) {
   return true;
 }
 
-/// Prunes the redundant list of arguments. E.g., if we are passing an argument
-/// list like [x, y, z, x] this would return [x, y, z] and it would update the
-/// `block` (to whom the argument are passed to) accordingly.
+/// Prunes the redundant list of new arguments. E.g., if we are passing an
+/// argument list like [x, y, z, x] this would return [x, y, z] and it would
+/// update the `block` (to whom the argument are passed to) accordingly. The new
+/// arguments are passed as arguments at the back of the block, hence we need to
+/// know how many `numOldArguments` were before, in order to correctly replace
+/// the new arguments in the block
 static SmallVector<SmallVector<Value, 8>, 2> pruneRedundantArguments(
     const SmallVector<SmallVector<Value, 8>, 2> &newArguments,
-    RewriterBase &rewriter, Block *block) {
+    RewriterBase &rewriter, unsigned numOldArguments, Block *block) {
 
   SmallVector<SmallVector<Value, 8>, 2> newArgumentsPruned(
       newArguments.size(), SmallVector<Value, 8>());
@@ -751,10 +754,11 @@ static SmallVector<SmallVector<Value, 8>, 2> pruneRedundantArguments(
   SmallVector<unsigned> toErase;
   for (auto [idx, arg] : llvm::enumerate(block->getArguments())) {
     if (idxToReplacement.contains(idx)) {
-      Value oldArg = block->getArgument(idx);
-      Value newArg = block->getArgument(idxToReplacement[idx]);
+      Value oldArg = block->getArgument(numOldArguments + idx);
+      Value newArg =
+          block->getArgument(numOldArguments + idxToReplacement[idx]);
       rewriter.replaceAllUsesWith(oldArg, newArg);
-      toErase.push_back(idx);
+      toErase.push_back(numOldArguments + idx);
     }
   }
 
@@ -793,6 +797,7 @@ LogicalResult BlockMergeCluster::merge(RewriterBase &rewriter) {
         1 + blocksToMerge.size(),
         SmallVector<Value, 8>(operandsToMerge.size()));
     unsigned curOpIndex = 0;
+    unsigned numOldArguments = leaderBlock->getNumArguments();
     for (const auto &it : llvm::enumerate(operandsToMerge)) {
       unsigned nextOpOffset = it.value().first - curOpIndex;
       curOpIndex = it.value().first;
@@ -814,7 +819,8 @@ LogicalResult BlockMergeCluster::merge(RewriterBase &rewriter) {
     }
 
     // Prune redundant arguments and update the leader block argument list
-    newArguments = pruneRedundantArguments(newArguments, rewriter, leaderBlock);
+    newArguments = pruneRedundantArguments(newArguments, rewriter,
+                                           numOldArguments, leaderBlock);
 
     // Update the predecessors for each of the blocks.
     auto updatePredecessors = [&](Block *block, unsigned clusterIndex) {

mlir/test/Transforms/test-canonicalize-merge-large-blocks.mlir

@@ -160,3 +160,33 @@ llvm.func @redundant_args_complex(%cond : i1) {
 ^bb3:
   llvm.return
 }
+
+llvm.func @blocks_with_args() {
+  %0 = llvm.mlir.zero : !llvm.ptr
+  %1 = llvm.call @rand() : () -> i1
+  // CHECK: %[[c1:.*]] = llvm.mlir.constant(1 : i64)
+  // CHECK: %[[c0:.*]] = llvm.mlir.constant(0 : i64)
+  // CHECK: %[[cond:.*]] = llvm.call @rand
+  %3 = llvm.mlir.constant(0) : i64
+  %4 = llvm.mlir.constant(1) : i64
+  // CHECK: llvm.cond_br %[[cond]], ^bb1(%[[c1]] : i64), ^bb1(%[[c0]] : i64)
+  // CHECK: ^bb1(%{{.*}}: i64):
+  // CHECK  ^bb2:
+  // CHECK  ^bb3:
+  // CHECK  llvm.return
+  llvm.cond_br %1, ^bb7(%0 : !llvm.ptr), ^bb1(%0 : !llvm.ptr)
+^bb1(%5: !llvm.ptr):
+  llvm.store %5, %0 : !llvm.ptr, !llvm.ptr
+  llvm.cond_br %1, ^bb2(%3 : i64), ^bb4(%3 : i64)
+^bb7(%6: !llvm.ptr):
+  llvm.store %6, %0 : !llvm.ptr, !llvm.ptr
+  llvm.cond_br %1, ^bb2(%4 : i64), ^bb4(%4 : i64)
+^bb2(%7: i64):
+  llvm.call @foo(%7) : (i64) -> ()
+  llvm.br ^bb8
+^bb4(%8: i64):
+  llvm.call @foo(%8) : (i64) -> ()
+  llvm.br ^bb8
+^bb8:
+  llvm.return
+}