[WIP][AArch64] Increase CSR cost when defering use of CSR is preferred

[junbuml]

Summary

The RegAllocGreedy currently favours allocating a CSR over splitting a region.
With this change we try to defer allocating CSR by increasing the CSR cost when
none of user blocks in the live range contains a call. Similar change was
initially suggested by Nemanja in D27366 for PowerPC and Wei posted D32201 to
increase the CSR cost.

Comparing with D27366, this change is more conservative in splitting over allocating
CSR becuase we increase the CSR cost only when none of blocks in the live range
have call. Instead of a fixed arbitrary CSR cost, we increase the CSR cost considering
the number of tradable splits/spills against a spill of a CSR in the entry block.

I'm posting this since it seems that Nemanja is no longer pursuing D27366.
Hopefully this change open up more productive discussion and make progress
in the context-sensitive CSR cost model.

With this patch, I observed +10% performance gain in spec2006/astar in AArch64.

[wmi]

lib/CodeGen/RegAllocGreedy.cpp

2387–2393

I don't see why the logic here is better than D32201. It will miss the motivational case we are targeting. For a long live range of a function argument, which crosses call in a cold bb, we want to increase the first time CSRCost so that regalloc can choose to split the long live range before and after the call in the cold bb and use non-CSR register in the places other than the cold bb. As a result shrinkwrapping will be unblocked because after the splitting there will be less or no CSR use inside of entry/exit.

[junbuml]

lib/CodeGen/RegAllocGreedy.cpp

2387–2393

Hi Wei,

This change was intended to make a process for the context-sensitive CSR cost model proposed in D27366. The case this change is targeting is different from D32201. I believe we can do this change on top of D32201 or this change can be treated differently from D32201.

This change try to defer allocating CSR when we prefer not to use one in some context. For that I check if there is any user block which contain a call. If none of the blocks has call, then I said allocating a CSR is not preferred in that context. Then I increase CSRCost based on the frequency of the entry block.

Like D32201, this change encourage more shrink-wrapping. In my experiment with spec2006/astar, one of the hot function was shrink-wrapped as we discourage the first allocation of CSR for a LR which do not expand across any call.

[junbuml]

Changed the return value of getNumberOfTradableSpillsAgainstCSR from 1 to 0 in AArch64 as I didn't see any profitable case yet.

[qcolombet]

include/llvm/Target/TargetLowering.h

3487

I don't think we should expose. Either regalloc can always do the right call or it cannot and we shouldn't ask the target to decide if it wants to be lucky here.

3493

I don't think we can have a simple numbers as in static number of instructions. The cost of splitting depends where it happens and not taking that into account would for sure lead to cases where we get it wrong.

3500

Ditto.

lib/CodeGen/RegAllocGreedy.cpp

2401

That formula does not make sense to me.
We should compare the cost of spilling vs. the CSRCost.
EntryFreq * #tradableSpills is not the cost of spilling.

lib/Target/AArch64/AArch64ISelLowering.cpp

10841

What do user blocks mean in this context?

[junbuml]

include/llvm/Target/TargetLowering.h

3493

I rename this to getWeightFactorToTheFirstCSRAllocation().
IIUC, the cost of splitting calculated in calculateRegionSplitCost() consider where it happen. We have to compare this split cost with the CSRCost which should be the cost for a spill of the CSR in the entry block. The return value of getWeightFactorToTheFirstCSRAllocation() approximates the number of copies which can be traded against the first spill of CSR in the entry block. In AArch64, I set 30, which means the very push/pop for the first CSR in the entry/exit block is tradable against to 30 copies in other blocks. Hopefully, this might estimate the benefit of shrink wrapping opportunity when deferring the first CSR allocation.

lib/CodeGen/RegAllocGreedy.cpp

2401

IIUC, the spilt cost is the sum of frequencies of the copies at all the related blocks, which is calculated in the current cost model. Similarly, the CSR cost should be a cost for a spill of the CSR in the entry block. I think that why we decide the CSRCost based on the frequency of the entry block in initializeCSRCost().
Since we have to compare the cost of copies with the cost of push/pop in prologue|epilogue, especially when allocating the first CSR is not preferred, I gave a weight factor through getWeightFactorToTheFirstCSRAllocation(). Hopefully this make sense to you. If not, please let me know any suggestion.

lib/Target/AArch64/AArch64ISelLowering.cpp

10841

Rename the function name and update the comment.

[junbuml]

Kindling ping. Please let me know any comment or suggestion regarding this change.

I believe the current CSRCost fixed in initializeCSRCost() is very conservative. In most case, we allocate a CSR rather than split ranges. However, in some situation (e.g, when none of user blocks of a LR has a function call), we should have more possibility to split the range because if we allocate a CSR too early, the CSR would not be accessed in some paths even after spilling it in the entry block. When we properly defer the use of CSR we might expose more shrink wrapping opportunities.

In this change, I treated the cost of spilt as the frequencies of the copies at the related spots while the cost of using the CSR is equivalent to the first spill in the entry block. Based on that, if none of user blocks of the LR has a function call, I increased the CSR cost by multiplying the frequency of the entry block with a fixed weight factor for the spill of the first CSR as we compare the cost between copies and the spill of CSRs.

[• wuzish]

Should this be rebased and continued to move on? Or it already could be addressed by https://reviews.llvm.org/D41765? @junbuml @qcolombet

include/llvm/Target/TargetLowering.h

 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/Type.h"
 #include "llvm/MC/MCRegisterInfo.h"
 #include "llvm/Support/AtomicOrdering.h"
 #include "llvm/Support/Casting.h"
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Target/TargetCallingConv.h"
 #include "llvm/Target/TargetMachine.h"
+#include "llvm/CodeGen/LiveInterval.h"
 #include <algorithm>
 #include <cassert>
 #include <climits>
 #include <cstdint>
 #include <iterator>
 #include <map>
 #include <string>
 #include <utility>
virtual void AdjustInstrPostInstrSelection(MachineInstr &MI,
                                              SDNode *Node) const;
 
   /// If this function returns true, SelectionDAGBuilder emits a
   /// LOAD_STACK_GUARD node when it is lowering Intrinsic::stackprotector.
   virtual bool useLoadStackGuardNode() const {
     return false;
   }
 
+  /// Target specific weight factor to the cost of allocating callee-saved
+  /// register for the first time in a function when defering the first
+  /// allocation of CSRs is prefered. The default value 0 means no change in
+  /// current behavior.
+  virtual unsigned getWeightFactorToTheFirstCSRAllocation() const {
+    return 0;
+  }
+
   /// Lower TLS global address SDNode for target independent emulated TLS model.
   virtual SDValue LowerToTLSEmulatedModel(const GlobalAddressSDNode *GA,
                                           SelectionDAG &DAG) const;
 
   // seteq(x, 0) -> truncate(srl(ctlz(zext(x)), log2(#bits)))
   // If we're comparing for equality to zero and isCtlzFast is true, expose the
   // fact that this can be implemented as a ctlz/srl pair, so that the dag
   // combiner can fold the new nodes.
   SDValue lowerCmpEqZeroToCtlzSrl(SDValue Op, SelectionDAG &DAG) const;
 
 private:
   SDValue simplifySetCCWithAnd(EVT VT, SDValue N0, SDValue N1,
                                ISD::CondCode Cond, DAGCombinerInfo &DCI,
                                const SDLoc &DL) const;
 };
 
 /// Given an LLVM IR type and return type attributes, compute the return value
 /// EVTs and flags, and optionally also the offsets, if the return value is
 /// being lowered to memory.
 void GetReturnInfo(Type *ReturnType, AttributeList attr,
                    SmallVectorImpl<ISD::OutputArg> &Outs,
                    const TargetLowering &TLI, const DataLayout &DL);
 
 } // end namespace llvm
 
 #endif // LLVM_TARGET_TARGETLOWERING_H

lib/CodeGen/RegAllocGreedy.cpp

 #include "llvm/Support/BlockFrequency.h"
 #include "llvm/Support/BranchProbability.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/MathExtras.h"
 #include "llvm/Support/Timer.h"
 #include "llvm/Support/raw_ostream.h"
 #include "llvm/Target/TargetInstrInfo.h"
+#include "llvm/Target/TargetLowering.h"
 #include "llvm/Target/TargetMachine.h"
 #include "llvm/Target/TargetRegisterInfo.h"
 #include "llvm/Target/TargetSubtargetInfo.h"
 #include <algorithm>
 #include <cassert>
 #include <cstdint>
 #include <memory>
 #include <queue>
class RAGreedy : public MachineFunctionPass,
   using SmallVirtRegSet = SmallSet<unsigned, 16>;
 
   // context
   MachineFunction *MF;
 
   // Shortcuts to some useful interface.
   const TargetInstrInfo *TII;
   const TargetRegisterInfo *TRI;
+  const TargetLowering *TLI;
   RegisterClassInfo RCI;
 
   // analyses
   SlotIndexes *Indexes;
   MachineBlockFrequencyInfo *MBFI;
   MachineDominatorTree *DomTree;
   MachineLoopInfo *Loops;
   MachineOptimizationRemarkEmitter *ORE;
public:
   MachineFunctionProperties getRequiredProperties() const override {
     return MachineFunctionProperties().set(
         MachineFunctionProperties::Property::NoPHIs);
   }
 
   static char ID;
 
 private:
+  /// Track if a CSR is allocated in the current function.
+  bool IsCSRAllocated;
+
   unsigned selectOrSplitImpl(LiveInterval &, SmallVectorImpl<unsigned> &,
                              SmallVirtRegSet &, unsigned = 0);
 
   bool LRE_CanEraseVirtReg(unsigned) override;
   void LRE_WillShrinkVirtReg(unsigned) override;
   void LRE_DidCloneVirtReg(unsigned, unsigned) override;
   void enqueue(PQueue &CurQueue, LiveInterval *LI);
   LiveInterval *dequeue(PQueue &CurQueue);
unsigned doRegionSplit(LiveInterval &VirtReg, unsigned BestCand,
                          bool HasCompact,
                          SmallVectorImpl<unsigned> &NewVRegs);
   /// Check other options before using a callee-saved register for the first
   /// time.
   unsigned tryAssignCSRFirstTime(LiveInterval &VirtReg, AllocationOrder &Order,
                                  unsigned PhysReg, unsigned &CostPerUseLimit,
                                  SmallVectorImpl<unsigned> &NewVRegs);
   void initializeCSRCost();
+  BlockFrequency getFirstTimeCSRCost(LiveInterval &VirtReg);
   unsigned tryBlockSplit(LiveInterval&, AllocationOrder&,
                          SmallVectorImpl<unsigned>&);
   unsigned tryInstructionSplit(LiveInterval&, AllocationOrder&,
                                SmallVectorImpl<unsigned>&);
   unsigned tryLocalSplit(LiveInterval&, AllocationOrder&,
     SmallVectorImpl<unsigned>&);
   unsigned trySplit(LiveInterval&, AllocationOrder&,
                     SmallVectorImpl<unsigned>&);

 /// to use the CSR; otherwise return 0.
 unsigned RAGreedy::tryAssignCSRFirstTime(LiveInterval &VirtReg,
                                          AllocationOrder &Order,
                                          unsigned PhysReg,
                                          unsigned &CostPerUseLimit,
                                          SmallVectorImpl<unsigned> &NewVRegs) {
   if (getStage(VirtReg) == RS_Spill && VirtReg.isSpillable()) {
     // We choose spill over using the CSR for the first time if the spill cost
-    // is lower than CSRCost.
+    // is lower than the CSR cost.
     SA->analyze(&VirtReg);
-    if (calcSpillCost() >= CSRCost)
+    if (calcSpillCost() >= getFirstTimeCSRCost(VirtReg)) {
       return PhysReg;
+    }
 
     // We are going to spill, set CostPerUseLimit to 1 to make sure that
     // we will not use a callee-saved register in tryEvict.
     CostPerUseLimit = 1;
     return 0;
   }
   if (getStage(VirtReg) < RS_Split) {
     // We choose pre-splitting over using the CSR for the first time if
-    // the cost of splitting is lower than CSRCost.
+    // the cost of splitting is lower than the CSR cost.
     SA->analyze(&VirtReg);
     unsigned NumCands = 0;
-    BlockFrequency BestCost = CSRCost; // Don't modify CSRCost.
+    BlockFrequency BestCost = getFirstTimeCSRCost(VirtReg);
     unsigned BestCand = calculateRegionSplitCost(VirtReg, Order, BestCost,
                                                  NumCands, true /*IgnoreCSR*/);
     if (BestCand == NoCand)
       // Use the CSR if we can't find a region split below CSRCost.
       return PhysReg;
 
     // Perform the actual pre-splitting.
     doRegionSplit(VirtReg, BestCand, false/*HasCompact*/, NewVRegs);
     return 0;
   }
   return PhysReg;
 }
 
 void RAGreedy::aboutToRemoveInterval(LiveInterval &LI) {
   // Do not keep invalid information around.
   SetOfBrokenHints.remove(&LI);
 }
 
+// Return the cost of the first use of a callee-saved register. For the first
+// allocation from CSRs in the current function, increase the cost if the live
+// range of the value spans basic blocks in which we'd prefer not to use one.
+// This will often defer the first allocation from CSRs and give shrink-wrapping
+// an opportunity to sink/hoist the save/restore from entry/exit blocks
+// respectively.
+BlockFrequency RAGreedy::getFirstTimeCSRCost(LiveInterval &VirtReg) {
+  if (getStage(VirtReg) == RS_Spill && VirtReg.isSpillable())
+    return CSRCost;
+
+  assert(getStage(VirtReg) < RS_Split &&
+         "Unexpected stage to find the first time CSR cost." );
+
+  // Increase the CSR cost only for the first allocation from CSRs in the
+  // current function.
+  if (IsCSRAllocated)
+     return CSRCost;
+
+  if (!VirtReg.isSpillable())
+    return CSRCost;
+
+  // Conservatively, we try to increase the CSR cost only when none of the
+  // blocks in the live range have call.
+  ArrayRef<SplitAnalysis::BlockInfo> UseBlocks = SA->getUseBlocks();
+  for (int i = 0, e = UseBlocks.size(); i < e; ++i)
+    for (auto &MI : *UseBlocks[i].MBB)
+      if (MI.isCall())
+        return CSRCost;
+
+  // Now, we prefer to trying to defer the first allocation from CSRs. We see
+  // the cost of spilt is the frequencies of the copy operation at the related
+  // spots while the cost of using the CSR is equivalent to the first spill in
+  // the entry block. Based on this, try to increase the CSR cost by multiplying
+  // the frequency of the entry block with the weight factor of the first
+  // allocation from CSRs. As we compare costs between copies and the spill of
+  // CSR, the weight factor approximates the number of copies which might be
+  // tradable against the first spill of CSR in the entry block.
+  uint64_t EntryFreq = MBFI->getEntryFreq();
+  return std::max(CSRCost.getFrequency(),
+                  EntryFreq * TLI->getWeightFactorToTheFirstCSRAllocation());
+}
+
 void RAGreedy::initializeCSRCost() {
   // We use the larger one out of the command-line option and the value report
   // by TRI.
   CSRCost = BlockFrequency(
       std::max((unsigned)CSRFirstTimeCost, TRI->getCSRFirstUseCost()));
   if (!CSRCost.getFrequency())
     return;
 
   // Raw cost is relative to Entry == 2^14; scale it appropriately.
   uint64_t ActualEntry = MBFI->getEntryFreq();
   if (!ActualEntry) {
     CSRCost = 0;
     return;
   }
   uint64_t FixedEntry = 1 << 14;
   if (ActualEntry < FixedEntry)
     CSRCost *= BranchProbability(ActualEntry, FixedEntry);
   else if (ActualEntry <= UINT32_MAX)
     // Invert the fraction and divide.
     CSRCost /= BranchProbability(FixedEntry, ActualEntry);
   else
     // Can't use BranchProbability in general, since it takes 32-bit numbers.
     CSRCost = CSRCost.getFrequency() * (ActualEntry / FixedEntry);
+
+  IsCSRAllocated = false;
 }
 
 /// \brief Collect the hint info for \p Reg.
 /// The results are stored into \p Out.
 /// \p Out is not cleared before being populated.
 void RAGreedy::collectHintInfo(unsigned Reg, HintsInfo &Out) {
   for (const MachineInstr &Instr : MRI->reg_nodbg_instructions(Reg)) {
     if (!Instr.isFullCopy())
unsigned RAGreedy::selectOrSplitImpl(LiveInterval &VirtReg,
                                      unsigned Depth) {
   unsigned CostPerUseLimit = ~0u;
   // First try assigning a free register.
   AllocationOrder Order(VirtReg.reg, *VRM, RegClassInfo, Matrix);
   if (unsigned PhysReg = tryAssign(VirtReg, Order, NewVRegs)) {
     // When NewVRegs is not empty, we may have made decisions such as evicting
     // a virtual register, go with the earlier decisions and use the physical
     // register.
-    if (CSRCost.getFrequency() && isUnusedCalleeSavedReg(PhysReg) &&
-        NewVRegs.empty()) {
+
+    if ((CSRCost.getFrequency() || (!IsCSRAllocated && VirtReg.isSpillable())) &&
+        isUnusedCalleeSavedReg(PhysReg) && NewVRegs.empty()) {
+
       unsigned CSRReg = tryAssignCSRFirstTime(VirtReg, Order, PhysReg,
                                               CostPerUseLimit, NewVRegs);
-      if (CSRReg || !NewVRegs.empty())
+      if (CSRReg || !NewVRegs.empty()) {
+        if (CSRReg)
+          IsCSRAllocated = true;
         // Return now if we decide to use a CSR or create new vregs due to
         // pre-splitting.
         return CSRReg;
+      }
     } else
       return PhysReg;
   }
 
   LiveRangeStage Stage = getStage(VirtReg);
   DEBUG(dbgs() << StageName[Stage]
                << " Cascade " << ExtraRegInfo[VirtReg.reg].Cascade << '\n');
 

 
 bool RAGreedy::runOnMachineFunction(MachineFunction &mf) {
   DEBUG(dbgs() << "********** GREEDY REGISTER ALLOCATION **********\n"
                << "********** Function: " << mf.getName() << '\n');
 
   MF = &mf;
   TRI = MF->getSubtarget().getRegisterInfo();
   TII = MF->getSubtarget().getInstrInfo();
+  TLI = MF->getSubtarget().getTargetLowering();
   RCI.runOnMachineFunction(mf);
 
   EnableLocalReassign = EnableLocalReassignment ||
                         MF->getSubtarget().enableRALocalReassignment(
                             MF->getTarget().getOptLevel());
 
   if (VerifyEnabled)
     MF->verify(this, "Before greedy register allocator");

lib/Target/AArch64/AArch64ISelLowering.h

private:
   void ReplaceNodeResults(SDNode *N, SmallVectorImpl<SDValue> &Results,
                           SelectionDAG &DAG) const override;
 
   bool functionArgumentNeedsConsecutiveRegisters(Type *Ty,
                                                  CallingConv::ID CallConv,
                                                  bool isVarArg) const override;
 
   bool shouldNormalizeToSelectSequence(LLVMContext &, EVT) const override;
+
+
+  virtual unsigned getWeightFactorToTheFirstCSRAllocation() const override;
 };
 
 namespace AArch64 {
 FastISel *createFastISel(FunctionLoweringInfo &funcInfo,
                          const TargetLibraryInfo *libInfo);
 } // end namespace AArch64
 
 } // end namespace llvm
 
 #endif

lib/Target/AArch64/AArch64ISelLowering.cpp

 
 unsigned
 AArch64TargetLowering::getVaListSizeInBits(const DataLayout &DL) const {
   if (Subtarget->isTargetDarwin() || Subtarget->isTargetWindows())
     return getPointerTy(DL).getSizeInBits();
 
   return 3 * getPointerTy(DL).getSizeInBits() + 2 * 32;
 }
+
+// Weight factor to the CSR cost applied when defering the first allocation of
+// CSR is prefered. This value approximate the number of copies which can be
+// traded against the first spill of CSR in the entry block when defering the
+// first allocation from CSRs is prefered.
+unsigned AArch64TargetLowering::getWeightFactorToTheFirstCSRAllocation() const {
+  return 30;
+}