[MBP] Improve nop placement in MBP pass.

[mcrosier]

Summary

While toying with the preferred loop alignment for AArch64, I noticed the MBP pass had a bad habit of inserting nop instructions into the dynamic instruction stream when aligning blocks. This patch attempts to prevent nops from being added to the dynamic instruction stream by placing additional alignment directives on non-fall through basic blocks.

This is a WIP and I don't plan on pursuing this patch further at this time. However, I figured I would post the patch in case someone else would like to take this on.. AFAIK, PowerPC is the only target that would immediately benefit from this patch. However, if other targets start setting a preferred loop alignment this work becomes much more interesting.

Chad

include/llvm/Target/TargetInstrInfo.h

public:
   bool isTriviallyReMaterializable(const MachineInstr *MI,
                                    AliasAnalysis *AA = nullptr) const {
     return MI->getOpcode() == TargetOpcode::IMPLICIT_DEF ||
            (MI->getDesc().isRematerializable() &&
             (isReallyTriviallyReMaterializable(MI, AA) ||
              isReallyTriviallyReMaterializableGeneric(MI, AA)));
   }
 
+  /// Determines the size of the specified MachineInstr.
+  /// \returns -1 if the target was unable to determine the size of the
+  /// instruction.
+  virtual unsigned GetInstSizeInBytes(const MachineInstr *MI) const {
+    return ~0x0;
+  }
+
 protected:
   /// For instructions with opcodes for which the M_REMATERIALIZABLE flag is
   /// set, this hook lets the target specify whether the instruction is actually
   /// trivially rematerializable, taking into consideration its operands. This
   /// predicate must return false if the instruction has any side effects other
   /// than producing a value, or if it requres any address registers that are
   /// not always available.
   /// Requirements must be check as stated in isTriviallyReMaterializable() .

lib/CodeGen/MachineBlockPlacement.cpp

 #define DEBUG_TYPE "block-placement"
 
 STATISTIC(NumCondBranches, "Number of conditional branches");
 STATISTIC(NumUncondBranches, "Number of unconditional branches");
 STATISTIC(CondBranchTakenFreq,
           "Potential frequency of taking conditional branches");
 STATISTIC(UncondBranchTakenFreq,
           "Potential frequency of taking unconditional branches");
+STATISTIC(NumAlignedBlocks, "Number of aligned blocks");
+STATISTIC(
+    NumAlignedBlocksWithoutNops,
+    "Number of aligned blocks that don't insert nops in instruction stream");
 
 static cl::opt<unsigned> AlignAllBlock("align-all-blocks",
                                        cl::desc("Force the alignment of all "
                                                 "blocks in the function."),
                                        cl::init(0), cl::Hidden);
 
 static cl::opt<unsigned> AlignAllNonFallThruBlocks(
     "align-all-nofallthru-blocks",
#endif // NDEBUG
   /// Note: This field is reinitialized multiple times - once for each loop,
   /// and then once for the function as a whole.
   unsigned UnscheduledPredecessors;
 };
 }
 
 namespace {
 class MachineBlockPlacement : public MachineFunctionPass {
+  /// \brief Did the MBP pass set the alignment of any MBBs?  If so, optimize
+  /// nop placement, if possible.
+  bool SetAlignmentOfBlocks;
+
   /// \brief A typedef for a block filter set.
   typedef SmallPtrSet<MachineBasicBlock *, 16> BlockFilterSet;
 
   /// \brief A handle to the branch probability pass.
   const MachineBranchProbabilityInfo *MBPI;
 
   /// \brief A handle to the function-wide block frequency pass.
   const MachineBlockFrequencyInfo *MBFI;
class MachineBlockPlacement : public MachineFunctionPass {
   BlockFilterSet collectLoopBlockSet(MachineFunction &F, MachineLoop &L);
   void buildLoopChains(MachineFunction &F, MachineLoop &L);
   void rotateLoop(BlockChain &LoopChain, MachineBasicBlock *ExitingBB,
                   const BlockFilterSet &LoopBlockSet);
   void rotateLoopWithProfile(BlockChain &LoopChain, MachineLoop &L,
                              const BlockFilterSet &LoopBlockSet);
   void buildCFGChains(MachineFunction &F);
 
+  void optimizeAlignmentNops(MachineFunction &F);
+
 public:
   static char ID; // Pass identification, replacement for typeid
-  MachineBlockPlacement() : MachineFunctionPass(ID) {
+  MachineBlockPlacement()
+      : MachineFunctionPass(ID), SetAlignmentOfBlocks(false) {
     initializeMachineBlockPlacementPass(*PassRegistry::getPassRegistry());
   }
 
   bool runOnMachineFunction(MachineFunction &F) override;
 
   void getAnalysisUsage(AnalysisUsage &AU) const override {
     AU.addRequired<MachineBranchProbabilityInfo>();
     AU.addRequired<MachineBlockFrequencyInfo>();
for (MachineBasicBlock *ChainBB : FunctionChain) {
     // Check for the existence of a non-layout predecessor which would benefit
     // from aligning this block.
     MachineBasicBlock *LayoutPred =
         &*std::prev(MachineFunction::iterator(ChainBB));
 
     // Force alignment if all the predecessors are jumps. We already checked
     // that the block isn't cold above.
     if (!LayoutPred->isSuccessor(ChainBB)) {
+      ++NumAlignedBlocks;
+      SetAlignmentOfBlocks = true;
       ChainBB->setAlignment(Align);
       continue;
     }
 
     // Align this block if the layout predecessor's edge into this block is
     // cold relative to the block. When this is true, other predecessors make up
     // all of the hot entries into the block and thus alignment is likely to be
     // important.
     BranchProbability LayoutProb =
         MBPI->getEdgeProbability(LayoutPred, ChainBB);
     BlockFrequency LayoutEdgeFreq = MBFI->getBlockFreq(LayoutPred) * LayoutProb;
-    if (LayoutEdgeFreq <= (Freq * ColdProb))
+    if (LayoutEdgeFreq <= (Freq * ColdProb)) {
+      ++NumAlignedBlocks;
+      SetAlignmentOfBlocks = true;
       ChainBB->setAlignment(Align);
     }
   }
+}
+
+namespace {
+
+typedef struct AlignmentRegion {
+  // Alignment.
+  unsigned Alignment;
+
+  // Total bytes in the region.
+  unsigned NumBytes;
+
+  // Track blocks that have no fall-through predecessors.  Aligning these blocks
+  // do not insert nops into the dynamic instruction stream.
+  SmallVector<std::pair<MachineBasicBlock *, unsigned>, 4> NonFallThruBlocks;
+
+  AlignmentRegion(unsigned A) : Alignment(A), NumBytes(0) {}
+} AlignmentRegion;
+
+} // namespace
+
+void MachineBlockPlacement::optimizeAlignmentNops(MachineFunction &F) {
+  bool EntryBlock = true;
+  AlignmentRegion *CurrRegion = nullptr;
+  SmallVector<AlignmentRegion, 4> Regions;
+
+  for (auto MBBI = F.begin(), MBBE = F.end(); MBBI != MBBE; ++MBBI) {
+    unsigned Alignment = MBBI->getAlignment();
+    MachineBasicBlock *MBB = &*MBBI;
+
+    // If this block is aligned (or is the entry block), create a new region.
+    if (Alignment || EntryBlock) {
+      unsigned FnAlignment = F.getAlignment();
+      if (EntryBlock && !Alignment)
+        Alignment = FnAlignment;
+      Regions.push_back(Alignment);
+      CurrRegion = &Regions.back();
+    } else {
+      assert(!EntryBlock && "Didn't expect the entry block.");
+      // Determine if this is a non fall-through block.
+      auto LayoutPred = std::prev(MBBI);
+      if (!LayoutPred->isSuccessor(MBB))
+        CurrRegion->NonFallThruBlocks.push_back(
+            std::make_pair(MBB, CurrRegion->NumBytes));
+    }
+    EntryBlock = false;
+
+    // Count the number of bytes in the block.
+    for (auto &I : *MBB) {
+      // Conservatively assume we don't know the size of inline assembly and
+      // bail.
+      if (I.isInlineAsm())
+        return;
+
+      unsigned NumBytes = TII->GetInstSizeInBytes(&I);
+      if (NumBytes == (unsigned)~0x0)
+        return;
+
+      CurrRegion->NumBytes += NumBytes;
+    }
+  }
+  // If we have a single region we don't have any aligned block in the function.
+  if (Regions.size() == 1)
+    return;
+
+  // See if we can align the function to minimize nops in the first region.
+  bool AlignedFunction = false;
+  if (F.getAlignment() < Regions[1].Alignment) {
+    unsigned PrefAlignment = Regions[1].Alignment;
+    unsigned Mask = (1 << PrefAlignment) - 1;
+    unsigned UnalignedBytes = Regions[0].NumBytes & Mask;
+    if (!UnalignedBytes) {
+      AlignedFunction = true;
+      ++NumAlignedBlocksWithoutNops;
+      F.setAlignment(PrefAlignment);
+    }
+  }
+
+  // Align non-fall through blocks to minimize nops in the dynamic instruction
+  // stream.
+  for (unsigned i = AlignedFunction ? 1 : 0, ie = Regions.size() - 1; i != ie;
+       ++i) {
+    // Get the preferred alignment for the next region.
+    unsigned PrefAlignment = Regions[i + 1].Alignment;
+    unsigned Mask = (1 << PrefAlignment) - 1;
+    unsigned RegionBytes = Regions[i].NumBytes;
+    // Look for a NFTB that can be aligned to remove all nops.
+    bool Success = false;
+    for (unsigned j = 0, je = Regions[i].NonFallThruBlocks.size(); j != je;
+         ++j) {
+      unsigned NFTBBytes = Regions[i].NonFallThruBlocks[j].second;
+      unsigned NFTBUnalignedBytes = (RegionBytes - NFTBBytes) & Mask;
+      if (!NFTBUnalignedBytes) {
+        MachineBasicBlock *MBB = Regions[i].NonFallThruBlocks[j].first;
+        assert(!MBB->getAlignment() && "Aligned blocks start new regions.");
+        ++NumAlignedBlocksWithoutNops;
+        MBB->setAlignment(PrefAlignment);
+        Success = true;
+        break;
+      }
+    }
+  }
+}
 
 bool MachineBlockPlacement::runOnMachineFunction(MachineFunction &F) {
   // Check for single-block functions and skip them.
   if (std::next(F.begin()) == F.end())
     return false;
 
   if (skipOptnoneFunction(*F.getFunction()))
     return false;
 
   MBPI = &getAnalysis<MachineBranchProbabilityInfo>();
   MBFI = &getAnalysis<MachineBlockFrequencyInfo>();
   MLI = &getAnalysis<MachineLoopInfo>();
   TII = F.getSubtarget().getInstrInfo();
   TLI = F.getSubtarget().getTargetLowering();
   MDT = &getAnalysis<MachineDominatorTree>();
   assert(BlockToChain.empty());
 
   buildCFGChains(F);
 
   BlockToChain.clear();
   ChainAllocator.DestroyAll();
 
+  // If we aligned any basic blocks, optimize the placement of the nops.
+  if (SetAlignmentOfBlocks)
+    optimizeAlignmentNops(F);
+
   if (AlignAllBlock)
     // Align all of the blocks in the function to a specific alignment.
     for (MachineBasicBlock &MBB : F)
       MBB.setAlignment(AlignAllBlock);
   else if (AlignAllNonFallThruBlocks) {
     // Align all of the blocks that have no fall-through predecessors to a
     // specific alignment.
     for (auto MBI = std::next(F.begin()), MBE = F.end(); MBI != MBE; ++MBI) {

lib/Target/AArch64/AArch64InstrInfo.h

 public:
   explicit AArch64InstrInfo(const AArch64Subtarget &STI);
 
   /// getRegisterInfo - TargetInstrInfo is a superset of MRegister info.  As
   /// such, whenever a client has an instance of instruction info, it should
   /// always be able to get register info as well (through this method).
   const AArch64RegisterInfo &getRegisterInfo() const { return RI; }
 
-  unsigned GetInstSizeInBytes(const MachineInstr *MI) const;
+  unsigned GetInstSizeInBytes(const MachineInstr *MI) const override;
 
   bool isAsCheapAsAMove(const MachineInstr *MI) const override;
 
   bool isCoalescableExtInstr(const MachineInstr &MI, unsigned &SrcReg,
                              unsigned &DstReg, unsigned &SubIdx) const override;
 
   bool
   areMemAccessesTriviallyDisjoint(MachineInstr *MIa, MachineInstr *MIb,

lib/Target/ARM/ARMBaseInstrInfo.h

public:
 
   bool DefinesPredicate(MachineInstr &MI,
                         std::vector<MachineOperand> &Pred) const override;
 
   bool isPredicable(MachineInstr &MI) const override;
 
   /// GetInstSize - Returns the size of the specified MachineInstr.
   ///
-  virtual unsigned GetInstSizeInBytes(const MachineInstr* MI) const;
+  virtual unsigned GetInstSizeInBytes(const MachineInstr *MI) const override;
 
   unsigned isLoadFromStackSlot(const MachineInstr *MI,
                                int &FrameIndex) const override;
   unsigned isStoreToStackSlot(const MachineInstr *MI,
                               int &FrameIndex) const override;
   unsigned isLoadFromStackSlotPostFE(const MachineInstr *MI,
                                      int &FrameIndex) const override;
   unsigned isStoreToStackSlotPostFE(const MachineInstr *MI,

lib/Target/MSP430/MSP430InstrInfo.h

void storeRegToStackSlot(MachineBasicBlock &MBB,
                            const TargetRegisterClass *RC,
                            const TargetRegisterInfo *TRI) const override;
   void loadRegFromStackSlot(MachineBasicBlock &MBB,
                             MachineBasicBlock::iterator MI,
                             unsigned DestReg, int FrameIdx,
                             const TargetRegisterClass *RC,
                             const TargetRegisterInfo *TRI) const override;
 
-  unsigned GetInstSizeInBytes(const MachineInstr *MI) const;
+  unsigned GetInstSizeInBytes(const MachineInstr *MI) const override;
 
   // Branch folding goodness
   bool
   ReverseBranchCondition(SmallVectorImpl<MachineOperand> &Cond) const override;
   bool isUnpredicatedTerminator(const MachineInstr &MI) const override;
   bool AnalyzeBranch(MachineBasicBlock &MBB,
                      MachineBasicBlock *&TBB, MachineBasicBlock *&FBB,
                      SmallVectorImpl<MachineOperand> &Cond,

lib/Target/Mips/MipsInstrInfo.h

public:
   /// such, whenever a client has an instance of instruction info, it should
   /// always be able to get register info as well (through this method).
   ///
   virtual const MipsRegisterInfo &getRegisterInfo() const = 0;
 
   virtual unsigned getOppositeBranchOpc(unsigned Opc) const = 0;
 
   /// Return the number of bytes of code the specified instruction may be.
-  unsigned GetInstSizeInBytes(const MachineInstr *MI) const;
+  unsigned GetInstSizeInBytes(const MachineInstr *MI) const override;
 
   void storeRegToStackSlot(MachineBasicBlock &MBB,
                            MachineBasicBlock::iterator MBBI,
                            unsigned SrcReg, bool isKill, int FrameIndex,
                            const TargetRegisterClass *RC,
                            const TargetRegisterInfo *TRI) const override {
     storeRegToStack(MBB, MBBI, SrcReg, isKill, FrameIndex, RC, TRI, 0);
   }

lib/Target/PowerPC/PPCInstrInfo.h

public:
   bool optimizeCompareInstr(MachineInstr *CmpInstr,
                             unsigned SrcReg, unsigned SrcReg2,
                             int Mask, int Value,
                             const MachineRegisterInfo *MRI) const override;
 
   /// GetInstSize - Return the number of bytes of code the specified
   /// instruction may be.  This returns the maximum number of bytes.
   ///
-  unsigned GetInstSizeInBytes(const MachineInstr *MI) const;
+  unsigned GetInstSizeInBytes(const MachineInstr *MI) const override;
 
   void getNoopForMachoTarget(MCInst &NopInst) const override;
 
   std::pair<unsigned, unsigned>
   decomposeMachineOperandsTargetFlags(unsigned TF) const override;
 
   ArrayRef<std::pair<unsigned, const char *>>
   getSerializableDirectMachineOperandTargetFlags() const override;
 
   ArrayRef<std::pair<unsigned, const char *>>
   getSerializableBitmaskMachineOperandTargetFlags() const override;
 };
 
 }
 
 #endif