[CodeMoverUtils] Enhance isSafeToMoveBefore() when control flow equivalence is satisfied

[congzhe]

Summary

With improved analysis in determining CFG equivalence that does not require strict dominance and post-dominance conditions, we now relax isSafeToMoveBefore() such that an instruction I can be moved before InsertPoint even if they do not strictly dominate each other, as long as they follow the same control flow path.

For example, we can move Instruction 0 before Instruction 1, and vice versa.

if (cond1)
   // Instruction 0: %add = add i32 1, 2
if (cond1)
   // Instruction 1: %add2 = add i32 2, 1

[Whitney]

llvm/lib/Transforms/Utils/CodeMoverUtils.cpp

440

How about adding assert(isControlFlowEquivalent(*ThisBlock, *OtherBlock, *DT, *PDT));?

446

typo: any or them

[congzhe]

llvm/lib/Transforms/Utils/CodeMoverUtils.cpp

440

Thanks, all comments addressed.

[Whitney]

One minor comment, other than that LGTM.

llvm/lib/Transforms/Utils/CodeMoverUtils.cpp

440

make sure to put some kind of error message in the assertion statement

https://llvm.org/docs/CodingStandards.html#assert-liberally

llvm/include/llvm/Transforms/Utils/CodeMoverUtils.h

void moveInstructionsToTheBeginning(BasicBlock &FromBB, BasicBlock &ToBB,
                                     DependenceInfo &DI);
 
 /// Move instructions, in an order-preserving manner, from \p FromBB to the end
 /// of \p ToBB when proven safe.
 void moveInstructionsToTheEnd(BasicBlock &FromBB, BasicBlock &ToBB,
                               DominatorTree &DT, const PostDominatorTree &PDT,
                               DependenceInfo &DI);
 
+/// In case that two BBs \p ThisBlock and \p OtherBlock are control flow
+/// equivalent but they do not strictly dominate and post-dominate each
+/// other, we determine if \p ThisBlock is reached after \p OtherBlock
+/// in the control flow.
+bool nonStrictlyPostDominate(const BasicBlock *ThisBlock,
+                             const BasicBlock *OtherBlock,
+                             const DominatorTree *DT,
+                             const PostDominatorTree *PDT);
+
+// Check if I0 is reached before I1 in the control flow.
+bool isReachedBefore(const Instruction *I0, const Instruction *I1,
+                     const DominatorTree *DT, const PostDominatorTree *PDT);
+
 } // end namespace llvm
 
 #endif // LLVM_TRANSFORMS_UTILS_CODEMOVERUTILS_H

llvm/lib/Transforms/Utils/CodeMoverUtils.cpp

bool llvm::isSafeToMoveBefore(Instruction &I, Instruction &InsertPoint,
 
   if (I.isTerminator())
     return reportInvalidCandidate(I, NotMovedTerminator);
 
   // TODO remove this limitation.
   if (!isControlFlowEquivalent(I, InsertPoint, DT, *PDT))
     return reportInvalidCandidate(I, NotControlFlowEquivalent);
 
-  if (!DT.dominates(&InsertPoint, &I))
+  if (isReachedBefore(&I, &InsertPoint, &DT, PDT))
     for (const Use &U : I.uses())
       if (auto *UserInst = dyn_cast<Instruction>(U.getUser()))
         if (UserInst != &InsertPoint && !DT.dominates(&InsertPoint, U))
           return false;
-  if (!DT.dominates(&I, &InsertPoint))
+  if (isReachedBefore(&InsertPoint, &I, &DT, PDT))
     for (const Value *Op : I.operands())
       if (auto *OpInst = dyn_cast<Instruction>(Op)) {
         if (&InsertPoint == OpInst)
           return false;
         // If OpInst is an instruction that appears earlier in the same BB as
         // I, then it is okay to move since OpInst will still be available.
         if (CheckForEntireBlock && I.getParent() == OpInst->getParent() &&
             DT.dominates(OpInst, &I))
void llvm::moveInstructionsToTheEnd(BasicBlock &FromBB, BasicBlock &ToBB,
                                     DependenceInfo &DI) {
   Instruction *MovePos = ToBB.getTerminator();
   while (FromBB.size() > 1) {
     Instruction &I = FromBB.front();
     if (isSafeToMoveBefore(I, *MovePos, DT, &PDT, &DI))
       I.moveBefore(MovePos);
   }
 }
+
+bool llvm::nonStrictlyPostDominate(const BasicBlock *ThisBlock,
+                                   const BasicBlock *OtherBlock,
+                                   const DominatorTree *DT,
+                                   const PostDominatorTree *PDT) {
+  assert(isControlFlowEquivalent(*ThisBlock, *OtherBlock, *DT, *PDT) &&
+         "ThisBlock and OtherBlock must be CFG equivalent!");
+  const BasicBlock *CommonDominator =
+      DT->findNearestCommonDominator(ThisBlock, OtherBlock);
+  if (CommonDominator == nullptr)
+    return false;
+
+  /// Recursively check the predecessors of \p ThisBlock up to
+  /// their common dominator, and see if any of them post-dominates
+  /// \p OtherBlock.
+  SmallVector<const BasicBlock *, 8> WorkList;
+  SmallPtrSet<const BasicBlock *, 8> Visited;
+  WorkList.push_back(ThisBlock);
+  while (!WorkList.empty()) {
+    const BasicBlock *CurBlock = WorkList.back();
+    WorkList.pop_back();
+    Visited.insert(CurBlock);
+    if (PDT->dominates(CurBlock, OtherBlock))
+      return true;
+
+    for (auto *Pred : predecessors(CurBlock)) {
+      if (Pred == CommonDominator || Visited.count(Pred))
+        continue;
+      WorkList.push_back(Pred);
+    }
+  }
+  return false;
+}
+
+bool llvm::isReachedBefore(const Instruction *I0, const Instruction *I1,
+                           const DominatorTree *DT,
+                           const PostDominatorTree *PDT) {
+  const BasicBlock *BB0 = I0->getParent();
+  const BasicBlock *BB1 = I1->getParent();
+  if (BB0 == BB1)
+    return DT->dominates(I0, I1);
+
+  return nonStrictlyPostDominate(BB1, BB0, DT, PDT);
+}

llvm/unittests/Transforms/Utils/CodeMoverUtilsTest.cpp

TEST(CodeMoverUtils, IsSafeToMoveTest4) {
 
   std::unique_ptr<Module> M =
       parseIR(C, R"(define void @foo(i1 %cond, i32 %op0, i32 %op1) {
                  entry:
                    br i1 %cond, label %if.end.first, label %if.then.first
                  if.then.first:
                    %add = add i32 %op0, %op1
                    %user = add i32 %add, 1
+                   %add2 = add i32 %op0, 1
                    br label %if.end.first
                  if.end.first:
                    br i1 %cond, label %if.end.second, label %if.then.second
                  if.then.second:
                    %sub_op0 = add i32 %op0, 1
                    %sub = sub i32 %sub_op0, %op1
+                   %sub2 = sub i32 %op0, 1
                    br label %if.end.second
                  if.end.second:
                    ret void
                  })");
 
   run(*M, "foo",
       [&](Function &F, DominatorTree &DT, PostDominatorTree &PDT,
           DependenceInfo &DI) {
         Instruction *AddInst = getInstructionByName(F, "add");
+        Instruction *AddInst2 = getInstructionByName(F, "add2");
         Instruction *SubInst = getInstructionByName(F, "sub");
+        Instruction *SubInst2 = getInstructionByName(F, "sub2");
 
         // Cannot move as %user uses %add and %sub doesn't dominates %user.
         EXPECT_FALSE(isSafeToMoveBefore(*AddInst, *SubInst, DT, &PDT, &DI));
 
         // Cannot move as %sub_op0 is an operand of %sub and %add doesn't
         // dominates %sub_op0.
         EXPECT_FALSE(isSafeToMoveBefore(*SubInst, *AddInst, DT, &PDT, &DI));
+
+        // Can move as %add2 and %sub2 are control flow equivalent,
+        // although %add2 does not strictly dominate %sub2.
+        EXPECT_TRUE(isSafeToMoveBefore(*AddInst2, *SubInst2, DT, &PDT, &DI));
+
+        // Can move as %add2 and %sub2 are control flow equivalent,
+        // although %add2 does not strictly dominate %sub2.
+        EXPECT_TRUE(isSafeToMoveBefore(*SubInst2, *AddInst2, DT, &PDT, &DI));
       });
 }
 
 TEST(CodeMoverUtils, IsSafeToMoveTest5) {
   LLVMContext C;
 
   std::unique_ptr<Module> M =
       parseIR(C, R"(define void @dependence(i32* noalias %A, i32* noalias %B){