[CodeMoverUtils] Enhance isSafeToMoveBefore() when moving BBs

[congzhe]

Summary

When moving an entire basic block BB before InsertPoint, currently we check for all instructions whether the operands dominates InsertPoint, however, this can be improved such that even an operand does not dominate InsertPoint, as long as it appears as a previous instruction in the same BB, it is safe to move.

[Whitney]

llvm/lib/Transforms/Utils/CodeMoverUtils.cpp

348

Do you think this is more clear?
if (CheckForEntireBlock && I.getParent() == OpInst->getParent() && DT.dominates(OpInst, &I))

[congzhe]

llvm/lib/Transforms/Utils/CodeMoverUtils.cpp

348

Thanks for the suggestion! It's definitely more clear. I updated the patch accordingly.

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

 bool isControlFlowEquivalent(const BasicBlock &BB0, const BasicBlock &BB1,
                              const DominatorTree &DT,
                              const PostDominatorTree &PDT);
 
 /// Return true if \p I can be safely moved before \p InsertPoint.
 bool isSafeToMoveBefore(Instruction &I, Instruction &InsertPoint,
                         DominatorTree &DT,
                         const PostDominatorTree *PDT = nullptr,
-                        DependenceInfo *DI = nullptr);
+                        DependenceInfo *DI = nullptr,
+                        bool CheckForEntireBlock = false);
 
 /// Return true if all instructions (except the terminator) in \p BB can be
 /// safely moved before \p InsertPoint.
 bool isSafeToMoveBefore(BasicBlock &BB, Instruction &InsertPoint,
                         DominatorTree &DT,
                         const PostDominatorTree *PDT = nullptr,
                         DependenceInfo *DI = nullptr);
 

llvm/lib/Transforms/Utils/CodeMoverUtils.cpp

if (!InBetweenInsts.insert(CurInst).second)
       continue;
 
     getNextInsts(*CurInst, WorkList);
   }
 }
 
 bool llvm::isSafeToMoveBefore(Instruction &I, Instruction &InsertPoint,
                               DominatorTree &DT, const PostDominatorTree *PDT,
-                              DependenceInfo *DI) {
+                              DependenceInfo *DI, bool CheckForEntireBlock) {
   // Skip tests when we don't have PDT or DI
   if (!PDT || !DI)
     return false;
 
   // Cannot move itself before itself.
   if (&I == &InsertPoint)
     return false;
 
bool llvm::isSafeToMoveBefore(Instruction &I, Instruction &InsertPoint,
 
   if (!DT.dominates(&InsertPoint, &I))
     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))
     for (const Value *Op : I.operands())
-      if (auto *OpInst = dyn_cast<Instruction>(Op))
-        if (&InsertPoint == OpInst || !DT.dominates(OpInst, &InsertPoint))
+      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))
+          continue;
+        if (!DT.dominates(OpInst, &InsertPoint))
           return false;
+      }
 
   DT.updateDFSNumbers();
   const bool MoveForward = domTreeLevelBefore(&DT, &I, &InsertPoint);
   Instruction &StartInst = (MoveForward ? I : InsertPoint);
   Instruction &EndInst = (MoveForward ? InsertPoint : I);
   SmallPtrSet<Instruction *, 10> InstsToCheck;
   collectInstructionsInBetween(StartInst, EndInst, InstsToCheck);
   if (!MoveForward)

 
 bool llvm::isSafeToMoveBefore(BasicBlock &BB, Instruction &InsertPoint,
                               DominatorTree &DT, const PostDominatorTree *PDT,
                               DependenceInfo *DI) {
   return llvm::all_of(BB, [&](Instruction &I) {
     if (BB.getTerminator() == &I)
       return true;
 
-    return isSafeToMoveBefore(I, InsertPoint, DT, PDT, DI);
+    return isSafeToMoveBefore(I, InsertPoint, DT, PDT, DI,
+                              /*CheckForEntireBlock=*/true);
   });
 }
 
 void llvm::moveInstructionsToTheBeginning(BasicBlock &FromBB, BasicBlock &ToBB,
                                           DominatorTree &DT,
                                           const PostDominatorTree &PDT,
                                           DependenceInfo &DI) {
   for (auto It = ++FromBB.rbegin(); It != FromBB.rend();) {

llvm/unittests/Transforms/Utils/CodeMoverUtilsTest.cpp

std::unique_ptr<Module> M = parseIR(C, R"(define void @foo(i64 %N) {
                  entry:
                    br label %for.body
                  for.body:
                    %i = phi i64 [ 0, %entry ], [ %inc, %for.latch ]
                    %inc = add nsw i64 %i, 1
                    br label %for.latch
                  for.latch:
                    %cmp = icmp slt i64 %inc, %N
+                   %add = add i64 100, %N
+                   %add2 = add i64 %add, %N
                    br i1 %cmp, label %for.body, label %for.end
                  for.end:
                    ret void
                  })");
 
   run(*M, "foo",
       [&](Function &F, DominatorTree &DT, PostDominatorTree &PDT,
           DependenceInfo &DI) {
         Instruction *IncInst = getInstructionByName(F, "inc");
         Instruction *CmpInst = getInstructionByName(F, "cmp");
+        BasicBlock *BB0 = getBasicBlockByName(F, "for.body");
+        BasicBlock *BB1 = getBasicBlockByName(F, "for.latch");
 
         // Can move as the incoming block of %inc for %i (%for.latch) dominated
         // by %cmp.
         EXPECT_TRUE(isSafeToMoveBefore(*IncInst, *CmpInst, DT, &PDT, &DI));
+
+        // Can move as the operands of instructions in BB1 either dominate
+        // InsertPoint or appear before that instruction, e.g., %add appears
+        // before %add2 although %add does not dominate InsertPoint.
+        EXPECT_TRUE(
+            isSafeToMoveBefore(*BB1, *BB0->getTerminator(), DT, &PDT, &DI));
       });
 }
 
 TEST(CodeMoverUtils, IsSafeToMoveTest4) {
   LLVMContext C;
 
   std::unique_ptr<Module> M =
       parseIR(C, R"(define void @foo(i1 %cond, i32 %op0, i32 %op1) {