Teach ConstantFoldTerminator to preserve DomTree

[arsenm]

Summary

The pass I'm working on only cares about the BranchInst
handling, which I know works. I don't currently have a way
to test the switch handling, so alternatively I could
just assert there is no DomTree to update for now.

[kuhar]

lib/Transforms/Utils/Local.cpp

129

Can Dest1 and Dest2 be the same BB? If so, a check for that would be needed before removing the edge. LLVM IR forms a multigraph, while the DomTree updater API treats it more like a normal graph, and we have to check for multiedges.

205

Same as above: are we sure that this is the only edge from SI->getParent() to DefaultDest?

243

Same as with the other deletions.

[arsenm]

Handle repeated successors correctly. Handle indirectbr, add unit tests.

[kuhar]

LGTM.

I would only add a comment for where you insert and look at .second highlighting the multiedge problem.

[arsenm]

It turns out this was just a waste of my time and this was already implemented differently in r322401. I don't see specific tests for ConstantFoldTerminator in that commit, so I've gone ahead and committed the test portion in r322683

include/llvm/Transforms/Utils/Local.h

 /// If a terminator instruction is predicated on a constant value, convert it
 /// into an unconditional branch to the constant destination.
 /// This is a nontrivial operation because the successors of this basic block
 /// must have their PHI nodes updated.
 /// Also calls RecursivelyDeleteTriviallyDeadInstructions() on any branch/switch
 /// conditions and indirectbr addresses this might make dead if
 /// DeleteDeadConditions is true.
 bool ConstantFoldTerminator(BasicBlock *BB, bool DeleteDeadConditions = false,
-                            const TargetLibraryInfo *TLI = nullptr);
+                            const TargetLibraryInfo *TLI = nullptr,
+                            DominatorTree *DT = nullptr);
 
 //===----------------------------------------------------------------------===//
 //  Local dead code elimination.
 //
 
 /// Return true if the result produced by the instruction is not used, and the
 /// instruction has no side effects.
 bool isInstructionTriviallyDead(Instruction *I,

lib/Transforms/Utils/Local.cpp

 /// ConstantFoldTerminator - If a terminator instruction is predicated on a
 /// constant value, convert it into an unconditional branch to the constant
 /// destination.  This is a nontrivial operation because the successors of this
 /// basic block must have their PHI nodes updated.
 /// Also calls RecursivelyDeleteTriviallyDeadInstructions() on any branch/switch
 /// conditions and indirectbr addresses this might make dead if
 /// DeleteDeadConditions is true.
 bool llvm::ConstantFoldTerminator(BasicBlock *BB, bool DeleteDeadConditions,
-                                  const TargetLibraryInfo *TLI) {
+                                  const TargetLibraryInfo *TLI,
+                                  DominatorTree *DT) {
   TerminatorInst *T = BB->getTerminator();
   IRBuilder<> Builder(T);
 
   // Branch - See if we are conditional jumping on constant
   if (auto *BI = dyn_cast<BranchInst>(T)) {
     if (BI->isUnconditional()) return false;  // Can't optimize uncond branch
     BasicBlock *Dest1 = BI->getSuccessor(0);
     BasicBlock *Dest2 = BI->getSuccessor(1);
 
     if (auto *Cond = dyn_cast<ConstantInt>(BI->getCondition())) {
       // Are we branching on constant?
       // YES.  Change to unconditional branch...
       BasicBlock *Destination = Cond->getZExtValue() ? Dest1 : Dest2;
       BasicBlock *OldDest     = Cond->getZExtValue() ? Dest2 : Dest1;
 
       // Let the basic block know that we are letting go of it.  Based on this,
       // it will adjust it's PHI nodes.
       OldDest->removePredecessor(BB);
 
       // Replace the conditional branch with an unconditional one.
       Builder.CreateBr(Destination);
       BI->eraseFromParent();
+
+      if (DT && OldDest != Destination)
+        DT->deleteEdge(BB, OldDest);
       return true;
     }
 
     if (Dest2 == Dest1) {       // Conditional branch to same location?
       // This branch matches something like this:
       //     br bool %cond, label %Dest, label %Dest
       // and changes it into:  br label %Dest
 
if (auto *SI = dyn_cast<SwitchInst>(T)) {
     BasicBlock *TheOnlyDest = DefaultDest;
 
     // If the default is unreachable, ignore it when searching for TheOnlyDest.
     if (isa<UnreachableInst>(DefaultDest->getFirstNonPHIOrDbg()) &&
         SI->getNumCases() > 0) {
       TheOnlyDest = SI->case_begin()->getCaseSuccessor();
     }
 
+    SmallVector<DominatorTree::UpdateType, 8> Updates;
+
     // Figure out which case it goes to.
     for (auto i = SI->case_begin(), e = SI->case_end(); i != e;) {
       // Found case matching a constant operand?
       if (i->getCaseValue() == CI) {
         TheOnlyDest = i->getCaseSuccessor();
         break;
       }
 
for (auto i = SI->case_begin(), e = SI->case_end(); i != e;) {
           std::swap(Weights[idx+1], Weights.back());
           Weights.pop_back();
           SI->setMetadata(LLVMContext::MD_prof,
                           MDBuilder(BB->getContext()).
                           createBranchWeights(Weights));
         }
         // Remove this entry.
         DefaultDest->removePredecessor(SI->getParent());
+
         i = SI->removeCase(i);
         e = SI->case_end();
         continue;
       }
 
       // Otherwise, check to see if the switch only branches to one destination.
       // We do this by reseting "TheOnlyDest" to null when we find two non-equal
       // destinations.
       if (i->getCaseSuccessor() != TheOnlyDest)
         TheOnlyDest = nullptr;
if (auto *SI = dyn_cast<SwitchInst>(T)) {
     }
 
     // If we found a single destination that we can fold the switch into, do so
     // now.
     if (TheOnlyDest) {
       // Insert the new branch.
       Builder.CreateBr(TheOnlyDest);
       BasicBlock *BB = SI->getParent();
+      SmallPtrSet<BasicBlock *, 4> Dests;
+      Dests.insert(TheOnlyDest);
 
       // Remove entries from PHI nodes which we no longer branch to...
       for (BasicBlock *Succ : SI->successors()) {
         // Found case matching a constant operand?
         if (Succ == TheOnlyDest)
           TheOnlyDest = nullptr; // Don't modify the first branch to TheOnlyDest
-        else
+        else {
           Succ->removePredecessor(BB);
+
+          if (DT && Dests.insert(Succ).second)
+            Updates.push_back({DominatorTree::Delete, BB, Succ});
+        }
       }
 
       // Delete the old switch.
       Value *Cond = SI->getCondition();
       SI->eraseFromParent();
+
+      if (DT)
+        DT->applyUpdates(Updates);
       if (DeleteDeadConditions)
         RecursivelyDeleteTriviallyDeadInstructions(Cond, TLI);
       return true;
     }
 
     if (SI->getNumCases() == 1) {
       // Otherwise, we can fold this switch into a conditional branch
       // instruction if it has only one non-default destination.
bool llvm::ConstantFoldTerminator(BasicBlock *BB, bool DeleteDeadConditions,
   if (auto *IBI = dyn_cast<IndirectBrInst>(T)) {
     // indirectbr blockaddress(@F, @BB) -> br label @BB
     if (auto *BA =
           dyn_cast<BlockAddress>(IBI->getAddress()->stripPointerCasts())) {
       BasicBlock *TheOnlyDest = BA->getBasicBlock();
       // Insert the new branch.
       Builder.CreateBr(TheOnlyDest);
 
+      SmallVector<DominatorTree::UpdateType, 8> Updates;
+      SmallPtrSet<BasicBlock *, 4> Dests;
+      Dests.insert(TheOnlyDest);
+
       for (unsigned i = 0, e = IBI->getNumDestinations(); i != e; ++i) {
-        if (IBI->getDestination(i) == TheOnlyDest)
+        BasicBlock *Dest = IBI->getDestination(i);
+        if (Dest == TheOnlyDest)
           TheOnlyDest = nullptr;
-        else
-          IBI->getDestination(i)->removePredecessor(IBI->getParent());
+        else {
+          Dest->removePredecessor(IBI->getParent());
+          if (DT && Dests.insert(Dest).second)
+            Updates.push_back({ DominatorTree::Delete, IBI->getParent(), Dest });
+        }
       }
       Value *Address = IBI->getAddress();
       IBI->eraseFromParent();
+
+      if (DT)
+        DT->applyUpdates(Updates);
+
       if (DeleteDeadConditions)
         RecursivelyDeleteTriviallyDeadInstructions(Address, TLI);
 
       // If we didn't find our destination in the IBI successor list, then we
       // have undefined behavior.  Replace the unconditional branch with an
       // 'unreachable' instruction.
       if (TheOnlyDest) {
         BB->getTerminator()->eraseFromParent();

unittests/Transforms/Utils/Local.cpp

runWithDomTree(
           if (!SinglePred || SinglePred == BB || BB->hasAddressTaken()) continue;
           BranchInst *Term = dyn_cast<BranchInst>(SinglePred->getTerminator());
           if (Term && !Term->isConditional())
             MergeBasicBlockIntoOnlyPred(BB, DT);
         }
         EXPECT_TRUE(DT->verify());
       });
 }
+
+TEST(Local, ConstantFoldTerminator) {
+  LLVMContext C;
+
+  std::unique_ptr<Module> M = parseIR(
+      C,
+      "define void @br_same_dest() {\n"
+      "entry:\n"
+      "  br i1 false, label %bb0, label %bb0\n"
+      "bb0:\n"
+      "  ret void\n"
+      "}\n"
+      "\n"
+      "define void @br_different_dest() {\n"
+      "entry:\n"
+      "  br i1 true, label %bb0, label %bb1\n"
+      "bb0:\n"
+      "  br label %exit\n"
+      "bb1:\n"
+      "  br label %exit\n"
+      "exit:\n"
+      "  ret void\n"
+      "}\n"
+      "\n"
+      "define void @switch_2_different_dest() {\n"
+      "entry:\n"
+      "  switch i32 0, label %default [ i32 0, label %bb0 ]\n"
+      "default:\n"
+      "  ret void\n"
+      "bb0:\n"
+      "  ret void\n"
+      "}\n"
+      "define void @switch_2_different_dest_default() {\n"
+      "entry:\n"
+      "  switch i32 1, label %default [ i32 0, label %bb0 ]\n"
+      "default:\n"
+      "  ret void\n"
+      "bb0:\n"
+      "  ret void\n"
+      "}\n"
+      "define void @switch_3_different_dest() {\n"
+      "entry:\n"
+      "  switch i32 0, label %default [ i32 0, label %bb0\n"
+      "                                 i32 1, label %bb1 ]\n"
+      "default:\n"
+      "  ret void\n"
+      "bb0:\n"
+      "  ret void\n"
+      "bb1:\n"
+      "  ret void\n"
+      "}\n"
+      "\n"
+      "define void @switch_variable_2_default_dest(i32 %arg) {\n"
+      "entry:\n"
+      "  switch i32 %arg, label %default [ i32 0, label %default ]\n"
+      "default:\n"
+      "  ret void\n"
+      "}\n"
+      "\n"
+      "define void @switch_constant_2_default_dest() {\n"
+      "entry:\n"
+      "  switch i32 1, label %default [ i32 0, label %default ]\n"
+      "default:\n"
+      "  ret void\n"
+      "}\n"
+      "\n"
+      "define void @switch_constant_3_repeated_dest() {\n"
+      "entry:\n"
+      "  switch i32 0, label %default [ i32 0, label %bb0\n"
+      "                                 i32 1, label %bb0 ]\n"
+      " bb0:\n"
+      "   ret void\n"
+      "default:\n"
+      "  ret void\n"
+      "}\n"
+      "\n"
+      "define void @indirectbr() {\n"
+      "entry:\n"
+      "  indirectbr i8* blockaddress(@indirectbr, %bb0), [label %bb0, label %bb1]\n"
+      "bb0:\n"
+      "  ret void\n"
+      "bb1:\n"
+      "  ret void\n"
+      "}\n"
+      "\n"
+      "define void @indirectbr_repeated() {\n"
+      "entry:\n"
+      "  indirectbr i8* blockaddress(@indirectbr_repeated, %bb0), [label %bb0, label %bb0]\n"
+      "bb0:\n"
+      "  ret void\n"
+      "}\n"
+      "\n"
+      "define void @indirectbr_unreachable() {\n"
+      "entry:\n"
+      "  indirectbr i8* blockaddress(@indirectbr_unreachable, %bb0), [label %bb1]\n"
+      "bb0:\n"
+      "  ret void\n"
+      "bb1:\n"
+      "  ret void\n"
+      "}\n"
+      "\n"
+    );
+
+  auto CFAllTerminators = [&](Function &F, DominatorTree *DT) {
+    for (Function::iterator I = F.begin(), E = F.end(); I != E;) {
+      BasicBlock *BB = &*I++;
+      ConstantFoldTerminator(BB, true, nullptr, DT);
+    }
+
+    EXPECT_TRUE(DT->verify());
+  };
+
+  runWithDomTree(*M, "br_same_dest", CFAllTerminators);
+  runWithDomTree(*M, "br_different_dest", CFAllTerminators);
+  runWithDomTree(*M, "switch_2_different_dest", CFAllTerminators);
+  runWithDomTree(*M, "switch_2_different_dest_default", CFAllTerminators);
+  runWithDomTree(*M, "switch_3_different_dest", CFAllTerminators);
+  runWithDomTree(*M, "switch_variable_2_default_dest", CFAllTerminators);
+  runWithDomTree(*M, "switch_constant_2_default_dest", CFAllTerminators);
+  runWithDomTree(*M, "switch_constant_3_repeated_dest", CFAllTerminators);
+  runWithDomTree(*M, "indirectbr", CFAllTerminators);
+  runWithDomTree(*M, "indirectbr_repeated", CFAllTerminators);
+  runWithDomTree(*M, "indirectbr_unreachable", CFAllTerminators);
+}