[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

2384–2390

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

2384–2390

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

3406

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.

3412

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.

3419

Ditto.

lib/CodeGen/RegAllocGreedy.cpp

2398

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

10787

What do user blocks mean in this context?

[junbuml]

include/llvm/Target/TargetLowering.h

3412

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

2398

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

10787

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;
   }
 
+  /// The target can specify whether a callee-saved register should be used
+  /// rather than spliting the live range. Default behaviour is yes.
+  virtual bool useCSRInsteadOfSplit(const LiveInterval &LI) const {
+    return true;
+  }
+
+  /// The number of splits in user blocks which could be allowed to be traded
+  /// for the spill of the CSR in the entry block when detering the first use
+  /// of CSR is prefered.
+  virtual unsigned getNumberOfTradableSplitsAgainstCSR() const {
+    return 0;
+  }
+
+  /// The number of spills in user blocks which could be allowed to be traded
+  /// for the spill of the CSR in the entry block when detering the first use
+  /// of CSR is prefered.
+  virtual unsigned getNumberOfTradableSpillsAgainstCSR() 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.

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;
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>&);
while (unsigned PhysReg = Order.next()) {
 
     // No live bundles, defer to splitSingleBlocks().
     if (!Cand.LiveBundles.any()) {
       DEBUG(dbgs() << " no bundles.\n");
       continue;
     }
 
     Cost += calcGlobalSplitCost(Cand);
+
     DEBUG({
       dbgs() << ", total = "; MBFI->printBlockFreq(dbgs(), Cost)
                                 << " with bundles";
       for (int i : Cand.LiveBundles.set_bits())
         dbgs() << " EB#" << i;
       dbgs() << ".\n";
     });
     if (Cost < BestCost) {

 /// 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);
 }
 
+// Increase the cost for the first use of a callee-saved register if the live
+// range of the value spans basic blocks in which we'd prefer not to use one.
+// This will often defer use of a CSR and give shrink-wrapping an opportunity
+// to sink/hoist the save/restore from entry/exit blocks respectively.
+BlockFrequency RAGreedy::getFirstTimeCSRCost(LiveInterval &VirtReg) {
+  BlockFrequency BestCost = CSRCost;
+  if (TLI->useCSRInsteadOfSplit(VirtReg))
+    return BestCost;
+
+  // Conservatively, we try to increase the CSR cost only when all blocks in
+  // the live range have no 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 BestCost;
+
+  // Now, we prefer to defering the first use of CSR. Try to increase the CSR
+  // cost by multipling the frequency of entry block with the number of tradable
+  // splits or spills.
+  uint64_t EntryFreq = MBFI->getEntryFreq();
+  if (getStage(VirtReg) == RS_Spill && VirtReg.isSpillable())
+    return std::max(BestCost.getFrequency(),
+                    EntryFreq * TLI->getNumberOfTradableSpillsAgainstCSR());
+  else if (getStage(VirtReg) < RS_Split)
+    return std::max(BestCost.getFrequency(),
+                    EntryFreq * TLI->getNumberOfTradableSplitsAgainstCSR());
+  else
+    llvm_unreachable ("Unexpected stage to find the first time CSR cost.");
+}
+
 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;
 
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() || !TLI->useCSRInsteadOfSplit(VirtReg)) &&
+        isUnusedCalleeSavedReg(PhysReg) && NewVRegs.empty()) {
       unsigned CSRReg = tryAssignCSRFirstTime(VirtReg, Order, PhysReg,
                                               CostPerUseLimit, NewVRegs);
       if (CSRReg || !NewVRegs.empty())
         // Return now if we decide to use a CSR or create new vregs due to
         // pre-splitting.
         return CSRReg;
     } else
       return PhysReg;

 
 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");
 
   RegAllocBase::init(getAnalysis<VirtRegMap>(),
                      getAnalysis<LiveIntervals>(),
                      getAnalysis<LiveRegMatrix>());
   Indexes = &getAnalysis<SlotIndexes>();
   MBFI = &getAnalysis<MachineBlockFrequencyInfo>();
   DomTree = &getAnalysis<MachineDominatorTree>();
   ORE = &getAnalysis<MachineOptimizationRemarkEmitterPass>().getORE();
   SpillerInstance.reset(createInlineSpiller(*this, *MF, *VRM));
   Loops = &getAnalysis<MachineLoopInfo>();
   Bundles = &getAnalysis<EdgeBundles>();
   SpillPlacer = &getAnalysis<SpillPlacement>();
   DebugVars = &getAnalysis<LiveDebugVariables>();
   AA = &getAnalysis<AAResultsWrapperPass>().getAAResults();
 
   initializeCSRCost();
-
   calculateSpillWeightsAndHints(*LIS, mf, VRM, *Loops, *MBFI);
 
   DEBUG(LIS->dump());
 
   SA.reset(new SplitAnalysis(*VRM, *LIS, *Loops));
   SE.reset(new SplitEditor(*SA, *AA, *LIS, *VRM, *DomTree, *MBFI));
   ExtraRegInfo.clear();
   ExtraRegInfo.resize(MRI->getNumVirtRegs());

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 bool useCSRInsteadOfSplit(const LiveInterval &LI) const override;
+
+  virtual unsigned getNumberOfTradableSplitsAgainstCSR() const override;
+
+  virtual unsigned getNumberOfTradableSpillsAgainstCSR() 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())
     return getPointerTy(DL).getSizeInBits();
 
   return 3 * getPointerTy(DL).getSizeInBits() + 2 * 32;
 }
+
+// If the live interval can be spilled, we'd prefer to do so.
+bool AArch64TargetLowering::useCSRInsteadOfSplit(const LiveInterval &LI) const {
+  return !LI.isSpillable();
+}
+
+// The number of splits in user blocks which can be traded against the spill of
+// the CSR in the entry block when detering the first use of CSR is prefered.
+unsigned AArch64TargetLowering::getNumberOfTradableSplitsAgainstCSR() const {
+  return 32;
+}
+
+// The number of spills in user blocks which can be traded against the spill of
+// the CSR in the entry block when detering the first use of CSR is prefered.
+unsigned AArch64TargetLowering::getNumberOfTradableSpillsAgainstCSR() const {
+  return 1;
+}
+

test/CodeGen/AArch64/csr-split.ll

+; RUN: llc -mtriple=aarch64-unknown-linux-gnu  < %s | FileCheck %s
+
+target datalayout = "e-m:e-i8:8:32-i16:16:32-i64:64-i128:128-n32:64-S128"
+target triple = "aarch64-linaro-linux-gnueabi"
+
+; After splitting, instead of assigning CSR, ShrinkWrap is enabled and prologue
+; are moved from entry to BB#1.
+
+define i32 @test1(i32 %x, i32 %y, i32* nocapture %P) local_unnamed_addr {
+; CHECK-LABEL: test1:
+; CHECK-LABEL: BB#0:
+; CHECK-NOT: stp
+; CHECK-LABEL: BB#1:
+; CHECK: stp x{{[0-9]+}}, x{{[0-9]+}}, [sp, {{.*}}]
+
+entry:
+  %idxprom = sext i32 %x to i64
+  %arrayidx = getelementptr inbounds i32, i32* %P, i64 %idxprom
+  %0 = load i32, i32* %arrayidx, align 4
+  %idxprom1 = sext i32 %y to i64
+  %arrayidx2 = getelementptr inbounds i32, i32* %P, i64 %idxprom1
+  %1 = load i32, i32* %arrayidx2, align 4
+  %add = add nsw i32 %1, %0
+  %cmp = icmp eq i32 %add, 0
+  br i1 %cmp, label %cleanup, label %if.end
+
+if.end:                                           ; preds = %entry
+  store i32 %add, i32* %P, align 4
+  tail call void @func()
+  br label %cleanup
+
+cleanup:                                          ; preds = %entry, %if.end
+  %retval.0 = phi i32 [ %x, %if.end ], [ 0, %entry ]
+  ret i32 %retval.0
+}
+
+declare void @func()
+