[MLIR] Correct block merge bug

[wsmoses]

Summary

Block merging in MLIR will incorrectly merge blocks with operations whose values are used outside of that block. This change forbids this behavior and provides a test where it is illegal to perform such a merge.

[wsmoses]

Add test comments.

[ftynse]

LGTM, forwarding to @rriddle as I am not very familiar with this code.

mlir/test/Transforms/canonicalize-block-merge.mlir

177

What is the thing that handles them?

Also ultra-nit: terminate sentences with a dot.

232

Nit: can we sink this down to the line that has the call?

[wsmoses]

Fix documentation nits

[wsmoses]

mlir/test/Transforms/canonicalize-block-merge.mlir

177

Added full-stop. Unsure where back edges are handled. This test, however, accidentally exhibited the same bug and had to be modified as a result. It now just tests for back edges and does not in the case a value is used externally (the portion we modified and checked by new test).

[rriddle]

This bug should only manifest when one of the blocks dominates the other, but we don't have access to DominanceInfo to check that easily right now though. Can you add a TODO that we can allow this merging when we can check the dominance?

mlir/lib/Transforms/Utils/RegionUtils.cpp

522

You can remove the functionality for newOpsToReplace/opsToReplace for now as well.

[wsmoses]

Remove "ops to merge" and add TODO

mlir/lib/Transforms/Utils/RegionUtils.cpp

private:
   BlockEquivalenceData leaderData;
 
   /// The set of blocks that can be merged into the leader.
   llvm::SmallSetVector<Block *, 1> blocksToMerge;
 
   /// A set of operand+index pairs that correspond to operands that need to be
   /// replaced by arguments when the cluster gets merged.
   std::set<std::pair<int, int>> operandsToMerge;
-
-  /// A map of operations with external uses to a replacement within the leader
-  /// block.
-  DenseMap<Operation *, Operation *> opsToReplace;
 };
 } // end anonymous namespace
 
 LogicalResult BlockMergeCluster::addToCluster(BlockEquivalenceData &blockData) {
   if (leaderData.hash != blockData.hash)
     return failure();
   Block *leaderBlock = leaderData.block, *mergeBlock = blockData.block;
   if (leaderBlock->getArgumentTypes() != mergeBlock->getArgumentTypes())
     return failure();
 
   // A set of operands that mismatch between the leader and the new block.
   SmallVector<std::pair<int, int>, 8> mismatchedOperands;
-  SmallVector<std::pair<Operation *, Operation *>, 2> newOpsToReplace;
   auto lhsIt = leaderBlock->begin(), lhsE = leaderBlock->end();
   auto rhsIt = blockData.block->begin(), rhsE = blockData.block->end();
   for (int opI = 0; lhsIt != lhsE && rhsIt != rhsE; ++lhsIt, ++rhsIt, ++opI) {
     // Check that the operations are equivalent.
     if (!OperationEquivalence::isEquivalentTo(
             &*lhsIt, &*rhsIt, OperationEquivalence::Flags::IgnoreOperands))
       return failure();
 
for (int operand : llvm::seq<int>(0, lhsIt->getNumOperands())) {
       }
 
       // Otherwise, these operands must have the same logical order within the
       // parent block.
       if (leaderData.getOrderOf(lhsOperand) != blockData.getOrderOf(rhsOperand))
         return failure();
     }
 
-    // If the rhs has external uses, it will need to be replaced.
-    if (rhsIt->isUsedOutsideOfBlock(mergeBlock))
-      newOpsToReplace.emplace_back(&*rhsIt, &*lhsIt);
+    // If the lhs or rhs has external uses, the blocks cannot be merged as the
+    // merged version of this operation will not be either the lhs or rhs
+    // alone (thus semantically incorrect), but some mix dependening on which
+    // block preceeded this.
+    // TODO allow merging of operations when one block does not dominate the
+    // other
+    if (rhsIt->isUsedOutsideOfBlock(mergeBlock) ||
+        lhsIt->isUsedOutsideOfBlock(leaderBlock)) {
+      return failure();
+    }
   }
   // Make sure that the block sizes are equivalent.
   if (lhsIt != lhsE || rhsIt != rhsE)
     return failure();
 
   // If we get here, the blocks are equivalent and can be merged.
   operandsToMerge.insert(mismatchedOperands.begin(), mismatchedOperands.end());
-  opsToReplace.insert(newOpsToReplace.begin(), newOpsToReplace.end());
   blocksToMerge.insert(blockData.block);
   return success();
 }
 
 /// Returns true if the predecessor terminators of the given block can not have
 /// their operands updated.
 static bool ableToUpdatePredOperands(Block *block) {
   for (auto it = block->pred_begin(), e = block->pred_end(); it != e; ++it) {
if (!operandsToMerge.empty()) {
     // terminators of each of the blocks can have their successor operands
     // updated.
     // TODO: We could try and sub-partition this cluster if only some blocks
     // cause the mismatch.
     if (!ableToUpdatePredOperands(leaderBlock) ||
         !llvm::all_of(blocksToMerge, ableToUpdatePredOperands))
       return failure();
 
-    // Replace any necessary operations.
-    for (std::pair<Operation *, Operation *> &it : opsToReplace)
-      it.first->replaceAllUsesWith(it.second);
-
     // Collect the iterators for each of the blocks to merge. We will walk all
     // of the iterators at once to avoid operand index invalidation.
     SmallVector<Block::iterator, 2> blockIterators;
     blockIterators.reserve(blocksToMerge.size() + 1);
     blockIterators.push_back(leaderBlock->begin());
     for (Block *mergeBlock : blocksToMerge)
       blockIterators.push_back(mergeBlock->begin());
 

mlir/test/Transforms/canonicalize-block-merge.mlir

^bb1:
   return
 ^bb2:
   scf.if %cond {
     "foo.op"() : () -> ()
   }
   return
 }
 
-// Check that properly handles back edges and the case where a value from one
-// block is used in another.
+// Check that properly handles back edges.
 
 // CHECK-LABEL: func @mismatch_loop(
 // CHECK-SAME: %[[ARG:.*]]: i1, %[[ARG2:.*]]: i1
 func @mismatch_loop(%cond : i1, %cond2 : i1) {
+  // CHECK-NEXT: %[[LOOP_CARRY:.*]] = "foo.op"
   // CHECK: cond_br %{{.*}}, ^bb1(%[[ARG2]] : i1), ^bb2
 
+  %cond3 = "foo.op"() : () -> (i1)
   cond_br %cond, ^bb2, ^bb3
 
 ^bb1:
   // CHECK: ^bb1(%[[ARG3:.*]]: i1):
-  // CHECK-NEXT: %[[LOOP_CARRY:.*]] = "foo.op"
   // CHECK-NEXT: cond_br %[[ARG3]], ^bb1(%[[LOOP_CARRY]] : i1), ^bb2
 
-  %ignored = "foo.op"() : () -> (i1)
   cond_br %cond3, ^bb1, ^bb3
 
 ^bb2:
-  %cond3 = "foo.op"() : () -> (i1)
   cond_br %cond2, ^bb1, ^bb3
 
 ^bb3:
   // CHECK: ^bb2:
   // CHECK-NEXT: return
 
   return
 }
^bb2:
   store %c0_i32, %arg1[] : memref<i32>
   br ^bb1
 
 ^bb3:
   // CHECK: store %{{.*}}, %{{.*}} : memref<i1>
   store %true, %arg2[] : memref<i1>
   br ^bb1
 }
+
+// Check that it is illegal to merge blocks containing an operand
+// with an external user. Incorrectly performing the optimization
+// anyways will result in print(merged, merged) rather than
+// distinct operands.
+func private @print(%arg0: i32, %arg1: i32)
+// CHECK-LABEL: @nomerge
+func @nomerge(%arg0: i32, %i: i32) {
+  %c1_i32 = constant 1 : i32
+  %icmp = cmpi "slt", %i, %arg0 : i32
+  cond_br %icmp, ^bb2, ^bb3
+
+^bb2:  // pred: ^bb1
+  %ip1 = addi %i, %c1_i32 : i32
+  br ^bb4(%ip1 : i32)
+
+^bb7:  // pred: ^bb5
+  %jp1 = addi %j, %c1_i32 : i32
+  br ^bb4(%jp1 : i32)
+
+^bb4(%j: i32):  // 2 preds: ^bb2, ^bb7
+  %jcmp = cmpi "slt", %j, %arg0 : i32
+// CHECK-NOT:  call @print(%[[arg1:.+]], %[[arg1]])
+  call @print(%j, %ip1) : (i32, i32) -> ()
+  cond_br %jcmp, ^bb7, ^bb3
+
+^bb3:  // pred: ^bb1
+  return
+}