[BranchFolding] Update UnavoidableBlocks for OutlineOptionalBranches

[haicheng]

Summary

The experimental code OutlineOptionalBranches of machine block placement requires updated UnavoidableBlocks if tail merging is called during the placement.

OutlineOptionalBranches partitions blocks into two groups, unavoidable blocks which dominates all the return blocks and the rest. During the layout, blocks in each group are placed consecutively. Unavoidable blocks are placed at the top and the rest blocks are at the bottom.

This patch caches and maintains the set of UnavoidableBlocks during tail merging so that we don't need to recalculate it when the layout is called again.

[haicheng]

Add a missing check.

[iteratee]

Updating just unavoidable blocks is insufficient.

See D20505 for an example of why.

[iteratee]

I'm also looking at this, and it doesn't even correctly update the list of unavoidable blocks. Consider merging 2 identical blocks, neither of which is unavoidable, but the merged block is. There are other examples, but that's probably the simplest one to describe in a comment.

[haicheng]

In D21674#466591, @iteratee wrote:

I'm also looking at this, and it doesn't even correctly update the list of unavoidable blocks. Consider merging 2 identical blocks, neither of which is unavoidable, but the merged block is. There are other examples, but that's probably the simplest one to describe in a comment.

Tail merging merges predecessors so that it is possible that an avoidable blocks can be promoted to be unavoidable, but I don't think an unavoidable block can become avoidable.

It does not make things worse if the block can be promoted remains in the avoidable area. Tail merging always make things better although maybe not best. I think you tail duplication patch can at least live with this patch and move on. If you want precise update, I can do it for you in the next patch.

[mcrosier]

@haicheng, are you still interested in trying to get this landed or is this something that can be abandon?

[haicheng]

Abandon it since it is not used.

lib/CodeGen/BranchFolding.h

namespace llvm {
   class LLVM_LIBRARY_VISIBILITY BranchFolder {
   public:
     class MBFIWrapper;
 
     explicit BranchFolder(bool defaultEnableTailMerge, bool CommonHoist,
                           MBFIWrapper &MBFI,
                           const MachineBranchProbabilityInfo &MBPI);
 
-    bool OptimizeFunction(MachineFunction &MF, const TargetInstrInfo *tii,
-                          const TargetRegisterInfo *tri, MachineModuleInfo *mmi,
-                          MachineLoopInfo *mli = nullptr,
-                          bool AfterPlacement = false);
+    bool OptimizeFunction(
+        MachineFunction &MF, const TargetInstrInfo *tii,
+        SmallPtrSet<MachineBasicBlock *, 4> *unavoidables = nullptr,
 
   private:
     class MergePotentialsElt {
       unsigned Hash;
       MachineBasicBlock *Block;
     public:
       MergePotentialsElt(unsigned h, MachineBasicBlock *b)
         : Hash(h), Block(b) {}
private:
 
     bool AfterBlockPlacement;
     bool EnableTailMerge;
     bool EnableHoistCommonCode;
     const TargetInstrInfo *TII;
     const TargetRegisterInfo *TRI;
     MachineModuleInfo *MMI;
     MachineLoopInfo *MLI;
+    SmallPtrSet<MachineBasicBlock *, 4> *UnavoidableBlocks;
     RegScavenger *RS;
 
   public:
     /// \brief This class keeps track of branch frequencies of newly created
     /// blocks and tail-merged blocks.
     class MBFIWrapper {
     public:
       MBFIWrapper(const MachineBlockFrequencyInfo &I) : MBFI(I) {}

lib/CodeGen/BranchFolding.cpp

void BranchFolder::RemoveDeadBlock(MachineBasicBlock *MBB) {
   // Avoid matching if this pointer gets reused.
   TriedMerging.erase(MBB);
 
   // Remove the block.
   MF->erase(MBB);
   FuncletMembership.erase(MBB);
   if (MLI)
     MLI->removeBlock(MBB);
+  if (UnavoidableBlocks)
+    UnavoidableBlocks->erase(MBB);
 }
 
 /// OptimizeImpDefsBlock - If a basic block is just a bunch of implicit_def
 /// followed by terminators, and if the implicitly defined registers are not
 /// used by the terminators, remove those implicit_def's. e.g.
 /// BB1:
 ///   r0 = implicit_def
 ///   r1 = implicit_def
bool BranchFolder::OptimizeImpDefsBlock(MachineBasicBlock *MBB) {
   }
 
   return true;
 }
 
 /// OptimizeFunction - Perhaps branch folding, tail merging and other
 /// CFG optimizations on the given function.  Block placement changes the layout
 /// and may create new tail merging opportunities.
-bool BranchFolder::OptimizeFunction(MachineFunction &MF,
-                                    const TargetInstrInfo *tii,
-                                    const TargetRegisterInfo *tri,
-                                    MachineModuleInfo *mmi,
-                                    MachineLoopInfo *mli, bool AfterPlacement) {
+bool BranchFolder::OptimizeFunction(
+    MachineFunction &MF, const TargetInstrInfo *tii,
+    const TargetRegisterInfo *tri, MachineModuleInfo *mmi, MachineLoopInfo *mli,
+    SmallPtrSet<MachineBasicBlock *, 4> *unavoidables, bool AfterPlacement) {
   if (!tii) return false;
 
   TriedMerging.clear();
 
   AfterBlockPlacement = AfterPlacement;
   TII = tii;
   TRI = tri;
   MMI = mmi;
   MLI = mli;
+  UnavoidableBlocks = unavoidables;
   RS = nullptr;
 
   // Use a RegScavenger to help update liveness when required.
   MachineRegisterInfo &MRI = MF.getRegInfo();
   if (MRI.tracksLiveness() && TRI->trackLivenessAfterRegAlloc(MF))
     RS = new RegScavenger();
   else
     MRI.invalidateLiveness();
MachineBasicBlock *BranchFolder::SplitMBBAt(MachineBasicBlock &CurMBB,
   // Splice the code over.
   NewMBB->splice(NewMBB->end(), &CurMBB, BBI1, CurMBB.end());
 
   // NewMBB belongs to the same loop as CurMBB.
   if (MLI) 
     if (MachineLoop *ML = MLI->getLoopFor(&CurMBB))
       ML->addBasicBlockToLoop(NewMBB, MLI->getBase());
 
+  if (UnavoidableBlocks && UnavoidableBlocks->count(&CurMBB))
+    UnavoidableBlocks->insert(NewMBB);
+
   // NewMBB inherits CurMBB's block frequency.
   MBBFreqInfo.setBlockFreq(NewMBB, MBBFreqInfo.getBlockFreq(&CurMBB));
 
   // For targets that use the register scavenger, we must maintain LiveIns.
   MaintainLiveIns(&CurMBB, NewMBB);
 
   // Add the new block to the funclet.
   const auto &FuncletI = FuncletMembership.find(&CurMBB);

lib/CodeGen/MachineBlockPlacement.cpp

class MachineBlockPlacement : public MachineFunctionPass {
 
   /// \brief A handle to the post dominator tree.
   MachineDominatorTree *MDT;
 
   /// \brief A set of blocks that are unavoidably execute, i.e. they dominate
   /// all terminators of the MachineFunction.
   SmallPtrSet<MachineBasicBlock *, 4> UnavoidableBlocks;
 
+  /// \brief A flag to show if UnavoidableBlocks is calculated.
+  bool HasUnavoidableBlocks;
+
   /// \brief Allocator and owner of BlockChain structures.
   ///
   /// We build BlockChains lazily while processing the loop structure of
   /// a function. To reduce malloc traffic, we allocate them using this
   /// slab-like allocator, and destroy them after the pass completes. An
   /// important guarantee is that this allocator produces stable pointers to
   /// the chains.
   SpecificBumpPtrAllocator<BlockChain> ChainAllocator;
DEBUG({
     }
     assert(!BadLoop && "Detected problems with the placement of this loop.");
   });
 }
 
 /// When OutlineOpitonalBranches is on, this method colects BBs that
 /// dominates all terminator blocks of the function \p F.
 void MachineBlockPlacement::collectMustExecuteBBs() {
-  if (OutlineOptionalBranches) {
+  if (OutlineOptionalBranches && !HasUnavoidableBlocks) {
     // Find the nearest common dominator of all of F's terminators.
     MachineBasicBlock *Terminator = nullptr;
     for (MachineBasicBlock &MBB : *F) {
       if (MBB.succ_size() == 0) {
         if (Terminator == nullptr)
           Terminator = &MBB;
         else
           Terminator = MDT->findNearestCommonDominator(Terminator, &MBB);
       }
     }
 
     // MBBs dominating this common dominator are unavoidable.
     UnavoidableBlocks.clear();
     for (MachineBasicBlock &MBB : *F) {
       if (MDT->dominates(&MBB, Terminator)) {
         UnavoidableBlocks.insert(&MBB);
       }
     }
+    HasUnavoidableBlocks = true;
   }
 }
 
 void MachineBlockPlacement::buildCFGChains() {
   // Ensure that every BB in the function has an associated chain to simplify
   // the assumptions of the remaining algorithm.
   SmallVector<MachineOperand, 4> Cond; // For AnalyzeBranch.
   for (MachineFunction::iterator FI = F->begin(), FE = F->end(); FI != FE;
bool MachineBlockPlacement::runOnMachineFunction(MachineFunction &MF) {
   F = &MF;
   MBPI = &getAnalysis<MachineBranchProbabilityInfo>();
   MBFI = llvm::make_unique<BranchFolder::MBFIWrapper>(
       getAnalysis<MachineBlockFrequencyInfo>());
   MLI = &getAnalysis<MachineLoopInfo>();
   TII = MF.getSubtarget().getInstrInfo();
   TLI = MF.getSubtarget().getTargetLowering();
   MDT = &getAnalysis<MachineDominatorTree>();
+  HasUnavoidableBlocks = false;
   assert(BlockToChain.empty());
 
   buildCFGChains();
 
   // Changing the layout can create new tail merging opportunities.
   TargetPassConfig *PassConfig = &getAnalysis<TargetPassConfig>();
   // TailMerge can create jump into if branches that make CFG irreducible for
   // HW that requires structurized CFG.
   bool EnableTailMerge = !MF.getTarget().requiresStructuredCFG() &&
                          PassConfig->getEnableTailMerge() &&
                          BranchFoldPlacement;
   // No tail merging opportunities if the block number is less than four.
   if (MF.size() > 3 && EnableTailMerge) {
     BranchFolder BF(/*EnableTailMerge=*/true, /*CommonHoist=*/false, *MBFI,
                     *MBPI);
 
-    if (BF.OptimizeFunction(MF, TII, MF.getSubtarget().getRegisterInfo(),
-                            getAnalysisIfAvailable<MachineModuleInfo>(), MLI,
-                            /*AfterBlockPlacement=*/true)) {
+    if (BF.OptimizeFunction(
+            MF, TII, MF.getSubtarget().getRegisterInfo(),
+            (OutlineOptionalBranches ? &UnavoidableBlocks : nullptr),
       // Redo the layout if tail merging creates/removes/moves blocks.
       BlockToChain.clear();
       ChainAllocator.DestroyAll();
       buildCFGChains();
     }
   }
 
   optimizeBranches();

test/CodeGen/AArch64/tailmerging_in_mbp.ll

-; RUN: llc <%s -march=aarch64 | FileCheck %s
+; RUN: llc < %s -march=aarch64 | FileCheck %s
+; RUN: llc < %s -march=aarch64 -outline-optional-branches | FileCheck %s --check-prefix=OUTLINE
 
 ; CHECK-LABEL: test:
 ; CHECK:       LBB0_7:
 ; CHECK:         b.hi 
 ; CHECK-NEXT:    b  
 ; CHECK-NEXT:  LBB0_8:
 ; CHECK-NEXT:    mov   x8, x9
 ; CHECK-NEXT:  LBB0_9:

 
 for.end:                                          ; preds = %for.end.loopexit, %entry
   %j.0.lcssa = phi i64 [ %n, %entry ], [ %j.0.lcssa.ph, %for.end.loopexit ]
   %j.2 = add i64 %j.0.lcssa, %n
   %j.3 = mul i64 %j.2, %n
   %j.4 = add i64 %j.3, 10
   ret i64 %j.4 
 }
+; OUTLINE-LABEL: test:
+; OUTLINE:       entry
+; OUTLINE:       for.end
+; OUTLINE:       for.body.preheader
+; OUTLINE:       cond.false.i
+; OUTLINE:       cond.false6.i
+; OUTLINE:       cond.false11.i
+; OUTLINE:       cond.false12.i
+; OUTLINE:       if.end
+; OUTLINE:       LBB0_9:
+