[CodeMoverUtils] Use dominator tree level to decide the direction of code motion

[RithikSharma]

Summary

Currently isSafeToMoveBefore uses DFS numbering for determining the relative position of instruction and insert point which is not always correct. This PR proposes the use of Dominator Tree depth for the same. If a node is at a higher level than the insert point then it is safe to say that we want to move in the forward direction.

[Whitney]

llvm/include/llvm/Analysis/OrderedInstructions.h

48

Please add description of this new function.

[fhahn]

It looks like the only user of OrderedInstructions.h is CodeMoverUtils now. At this point I think it would be good to remove OrderedInstructions and move the new code directly to CodeMoverUtils

[nikic]

llvm/lib/Transforms/Utils/CodeMoverUtils.cpp

323–337

I'd feel a lot better about this code if this code were just

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))
        return false;

which is more "obviously correct". The outer dominance checks here are just an optimization to handle obvious cases, but

for (const Use &U : I.uses())
  if (auto *UserInst = dyn_cast<Instruction>(U.getUser()))
    if (UserInst != &InsertPoint && !DT.dominates(&InsertPoint, U))
      return false;
for (const Value *Op : I.operands())
  if (auto *OpInst = dyn_cast<Instruction>(Op))
    if (&InsertPoint == OpInst || !DT.dominates(OpInst, &InsertPoint))
      return false;

would work as well.

In this case the MoveForward variable would only be needed for the collectInstructionsInBetween() check below.

[RithikSharma]

Added function description to bfsLevelBefore and made some optimization changes in the dominance checks.

[RithikSharma]

In D80084#2040795, @fhahn wrote:

It looks like the only user of OrderedInstructions.h is CodeMoverUtils now. At this point I think it would be good to remove OrderedInstructions and move the new code directly to CodeMoverUtils

It makes sense to me, surely I will move these OrderedInstructions to CodeMoverUtils in the next patch.

[RithikSharma]

This update changes bfsLevelBefore to domTreeLevelBefore, the new name adds more details about the function and makes it more descriptive.

[Whitney]

llvm/include/llvm/Analysis/OrderedInstructions.h

52

nit: sentence end with a period.

llvm/lib/Analysis/OrderedInstructions.cpp

48

clang-format?

[RithikSharma]

This update changes OrderedInstructions.cpp to Clang-format and fixes errors in function description of domTreeLevelBefore.

[Whitney]

Knowing that you will move OrderedInstructions to CodeMoverUtils in the next patch. LGTM.

[nikic]

llvm/lib/Transforms/Utils/CodeMoverUtils.cpp

342

Side note: collectInstructionsInBetween() seems to collect much more instructions than it actually needs. If there is a control flow split and the end instruction is on one side, it will still visit the other one, and everything that is reachable from it.

[Whitney]

llvm/lib/Transforms/Utils/CodeMoverUtils.cpp

342

If there is a control flow split, then isControlFlowEquivalent check on line 317 will exit.

llvm/include/llvm/Analysis/OrderedInstructions.h

public:
   /// Return true if first instruction dominates the second.
   bool dominates(const Instruction *, const Instruction *) const;
 
   /// Return true if the first instruction comes before the second in the
   /// dominator tree DFS traversal if they are in different basic blocks,
   /// or if the first instruction comes before the second in the same basic
   /// block.
   bool dfsBefore(const Instruction *, const Instruction *) const;
+  bool bfsLevelBefore(const Instruction *, const Instruction *) const;
 };
 
 } // end namespace llvm
 
 #endif // LLVM_ANALYSIS_ORDEREDINSTRUCTIONS_H

llvm/lib/Analysis/OrderedInstructions.cpp

bool OrderedInstructions::dfsBefore(const Instruction *InstA,
   // block.
   if (InstA->getParent() == InstB->getParent())
     return localDominates(InstA, InstB);
 
   DomTreeNode *DA = DT->getNode(InstA->getParent());
   DomTreeNode *DB = DT->getNode(InstB->getParent());
   return DA->getDFSNumIn() < DB->getDFSNumIn();
 }
+
+bool OrderedInstructions::bfsLevelBefore(const Instruction *InstA,
+                                    const Instruction *InstB) const {
+  // Use ordered basic block in case the 2 instructions are in the same basic
+  // block.
+  if (InstA->getParent() == InstB->getParent())
+    return localDominates(InstA, InstB);
+
+  DomTreeNode *DA = DT->getNode(InstA->getParent());
+  DomTreeNode *DB = DT->getNode(InstB->getParent());
+  return DA->getLevel() < DB->getLevel();
+}

llvm/lib/Transforms/Utils/CodeMoverUtils.cpp

if (I.isTerminator())
     return reportInvalidCandidate(I, NotMovedTerminator);
 
   // TODO remove this limitation.
   if (!isControlFlowEquivalent(I, InsertPoint, DT, PDT))
     return reportInvalidCandidate(I, NotControlFlowEquivalent);
 
   OrderedInstructions OI(&DT);
   DT.updateDFSNumbers();
-  const bool MoveForward = OI.dfsBefore(&I, &InsertPoint);
+  const bool MoveForward = OI.bfsLevelBefore(&I, &InsertPoint);
   if (MoveForward) {
     // When I is being moved forward, we need to make sure the InsertPoint
     // dominates every users. Or else, a user may be using an undefined I.
     for (const Use &U : I.uses())
       if (auto *UserInst = dyn_cast<Instruction>(U.getUser()))
         if (UserInst != &InsertPoint && !DT.dominates(&InsertPoint, U))
           return false;
   } else {
     // When I is being moved backward, we need to make sure all its opernads
     // dominates the InsertPoint. Or else, an operand may be undefined for I.
     for (const Value *Op : I.operands())
       if (auto *OpInst = dyn_cast<Instruction>(Op))
         if (&InsertPoint == OpInst || !OI.dominates(OpInst, &InsertPoint))
           return false;
   }
 
   Instruction &StartInst = (MoveForward ? I : InsertPoint);
   Instruction &EndInst = (MoveForward ? InsertPoint : I);
   SmallPtrSet<Instruction *, 10> InstsToCheck;
   collectInstructionsInBetween(StartInst, EndInst, InstsToCheck);
   if (!MoveForward)
     InstsToCheck.insert(&InsertPoint);
 
   // Check if there exists instructions which may throw, may synchonize, or may
   // never return, from I to InsertPoint.
   if (!isSafeToSpeculativelyExecute(&I))
     if (std::any_of(InstsToCheck.begin(), InstsToCheck.end(),
                     [](Instruction *I) {

llvm/unittests/Transforms/Utils/CodeMoverUtilsTest.cpp

run(*M, "foo",
         Instruction *IncInst = getInstructionByName(F, "inc");
         Instruction *CmpInst = getInstructionByName(F, "cmp");
 
         // Can move as the incoming block of %inc for %i (%for.latch) dominated
         // by %cmp.
         EXPECT_TRUE(isSafeToMoveBefore(*IncInst, *CmpInst, 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) {
+                 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
+                   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
+                   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 *SubInst = getInstructionByName(F, "sub");
+
+        // 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));
+      });
+}