AArch64: Use SplitCSR for callee save registers used for parameters.

[MatzeB]

Summary

Callee save parameters that are also used to pass parameter values
should use virtual regs for save/restore to avoid unnecessary spills.

This is done by determining a list of callee saves (the ones handled in
FrameLowering) and SavedViaCopy (register is copied to a vreg so the
register allocator handles saving/restoring).

Related to rdar://25638054

[t.p.northover]

Looks reasonable. I hope it's worth it.

Tim.

lib/Target/AArch64/AArch64RegisterInfo.h

36

Typo on "teerminated".

include/llvm/CodeGen/MachineFunction.h

public:
 
   /// Dellocate an array of MachineOperands and recycle the memory. This is
   /// only intended for use by internal MachineInstr functions.
   /// Cap must be the same capacity that was used to allocate the array.
   void deallocateOperandArray(OperandCapacity Cap, MachineOperand *Array) {
     OperandRecycler.deallocate(Cap, Array);
   }
 
+  /// Return allocator for data with the lifetime of the MachineFunction.
+  BumpPtrAllocator &getAllocator() {
+    return Allocator;
+  }
+
   /// \brief Allocate and initialize a register mask with @p NumRegister bits.
   uint32_t *allocateRegisterMask(unsigned NumRegister) {
     unsigned Size = (NumRegister + 31) / 32;
     uint32_t *Mask = Allocator.Allocate<uint32_t>(Size);
     for (unsigned i = 0; i != Size; ++i)
       Mask[i] = 0;
     return Mask;
   }

lib/Target/AArch64/AArch64ISelLowering.cpp

if (Ins[i].isOrigArg()) {
         ValVT = MVT::i16;
     }
     CCAssignFn *AssignFn = CCAssignFnForCall(CallConv, /*IsVarArg=*/false);
     bool Res =
         AssignFn(i, ValVT, ValVT, CCValAssign::Full, Ins[i].Flags, CCInfo);
     assert(!Res && "Call operand has unhandled type");
     (void)Res;
   }
+
+  const TargetRegisterInfo &TRI = *MF.getSubtarget().getRegisterInfo();
+  const uint32_t *PreservedMask = TRI.getCallPreservedMask(MF, CallConv);
+  BitVector ParametersInCalleeSaveRegs;
+
   assert(ArgLocs.size() == Ins.size());
   SmallVector<SDValue, 16> ArgValues;
   for (unsigned i = 0, e = ArgLocs.size(); i != e; ++i) {
     CCValAssign &VA = ArgLocs[i];
 
     if (Ins[i].Flags.isByVal()) {
       // Byval is used for HFAs in the PCS, but the system should work in a
       // non-compliant manner for larger structs.
if (VA.isRegLoc()) {
         RC = &AArch64::FPR32RegClass;
       else if (RegVT == MVT::f64 || RegVT.is64BitVector())
         RC = &AArch64::FPR64RegClass;
       else if (RegVT == MVT::f128 || RegVT.is128BitVector())
         RC = &AArch64::FPR128RegClass;
       else
         llvm_unreachable("RegVT not supported by FORMAL_ARGUMENTS Lowering");
 
+      MCPhysReg PhysReg = VA.getLocReg();
+      if (!MachineOperand::clobbersPhysReg(PreservedMask, PhysReg)) {
+        if (ParametersInCalleeSaveRegs.size() == 0)
+          ParametersInCalleeSaveRegs.resize(TRI.getNumRegs());
+        ParametersInCalleeSaveRegs.set(PhysReg);
+      }
+
       // Transform the arguments in physical registers into virtual ones.
-      unsigned Reg = MF.addLiveIn(VA.getLocReg(), RC);
+      unsigned Reg = MF.addLiveIn(PhysReg, RC);
       ArgValue = DAG.getCopyFromReg(Chain, DL, Reg, RegVT);
 
       // If this is an 8, 16 or 32-bit value, it is really passed promoted
       // to 64 bits.  Insert an assert[sz]ext to capture this, then
       // truncate to the right size.
       switch (VA.getLocInfo()) {
       default:
         llvm_unreachable("Unknown loc info!");
if (DoesCalleeRestoreStack(CallConv, TailCallOpt)) {
     // This realignment carries over to the available bytes below. Our own
     // callers will guarantee the space is free by giving an aligned value to
     // CALLSEQ_START.
   }
   // Even if we're not expected to free up the space, it's useful to know how
   // much is there while considering tail calls (because we can reuse it).
   FuncInfo->setBytesInStackArgArea(StackArgSize);
 
-  // Should we save callee save registers via vregs?
-  if (CallConv == CallingConv::CXX_FAST_TLS &&
+  const MCPhysReg *CalleeSavedRegs =
+    AArch64RegisterInfo::getBasicCalleeSavedRegs(&MF);
+  const MCPhysReg *SavedViaCopy = nullptr;
+
+  // Should we save some callee saved registers in vregs?
+  if ((CallConv == CallingConv::CXX_FAST_TLS ||
+       ParametersInCalleeSaveRegs.size() > 0) &&
       MF.getFunction()->hasFnAttribute(Attribute::NoUnwind) &&
       !MF.getTarget().Options.EnableFastISel && mayUseSplitCSR(MF)) {
-    const MCPhysReg *SavedViaCopy =
-      AArch64RegisterInfo::getFastTLSSavedViaCopy();
+    if (CallConv == CallingConv::CXX_FAST_TLS) {
+      SavedViaCopy = AArch64RegisterInfo::getFastTLSSavedViaCopy();
+      CalleeSavedRegs = AArch64RegisterInfo::getFastTLSSaved();
+    } else {
+      assert(ParametersInCalleeSaveRegs.size() > 0 && "must have csr params");
+      // Split callee saved registers into SavedViaCopy and the rest.
+      unsigned NCalleeSavedRegs = 0;
+      for (const MCPhysReg *CSR = CalleeSavedRegs; *CSR != 0; ++CSR)
+        ++NCalleeSavedRegs;
+      unsigned NParamsInCSR = ParametersInCalleeSaveRegs.count();
+
+      MCPhysReg *NewCalleeSavedRegs =
+        MF.getAllocator().Allocate<MCPhysReg>(NCalleeSavedRegs-NParamsInCSR+1);
+      MCPhysReg *NewSavedViaCopy =
+        MF.getAllocator().Allocate<MCPhysReg>(NParamsInCSR+1);
+      unsigned C0 = 0;
+      unsigned C1 = 0;
+      for (const MCPhysReg *CSR = CalleeSavedRegs; *CSR != 0; ++CSR) {
+        MCPhysReg Reg = *CSR;
+        if (ParametersInCalleeSaveRegs[Reg])
+          NewSavedViaCopy[C0++] = Reg;
+        else
+          NewCalleeSavedRegs[C1++] = Reg;
+      }
+      NewSavedViaCopy[C0++] = 0;
+      NewCalleeSavedRegs[C1++] = 0;
+      assert(C0 == NParamsInCSR+1 && "Correct register count");
+      assert(C1 == NCalleeSavedRegs-NParamsInCSR+1 && "Correct register count");
+      CalleeSavedRegs = NewCalleeSavedRegs;
+      SavedViaCopy = NewSavedViaCopy;
+    }
     FuncInfo->setCSRSavedViaCopy(SavedViaCopy);
 
-    // Add live-ins.
+    // Add live-ins for SavedViaCopy.
     for (const MCPhysReg *I = SavedViaCopy; *I; ++I) {
       const MCPhysReg Reg = *I;
       const TargetRegisterClass *RC;
       if (AArch64::GPR64RegClass.contains(Reg))
         RC = &AArch64::GPR64RegClass;
       else if (AArch64::FPR64RegClass.contains(Reg))
         RC = &AArch64::FPR64RegClass;
       else
         llvm_unreachable("Unexpected register class in CSRsViaCopy!");
       MF.addLiveIn(Reg, RC);
     }
   }
+  FuncInfo->setCalleeSavedRegs(CalleeSavedRegs);
 
   return Chain;
 }
 
 void AArch64TargetLowering::saveVarArgRegisters(CCState &CCInfo,
                                                 SelectionDAG &DAG, SDLoc DL,
                                                 SDValue &Chain) const {
   MachineFunction &MF = DAG.getMachineFunction();

lib/Target/AArch64/AArch64MachineFunctionInfo.h

class AArch64FunctionInfo : public MachineFunctionInfo {
 
   /// zero-terminated list of callee saved registers saved by copying to a vreg.
   const MCPhysReg *CSRSavedViaCopy = nullptr;
 
   /// True when the stack gets realigned dynamically because the size of stack
   /// frame is unknown at compile time. e.g., in case of VLAs.
   bool StackRealigned;
 
+  const MCPhysReg *CalleeSavedRegs = nullptr;
+
 public:
   AArch64FunctionInfo()
       : BytesInStackArgArea(0), ArgumentStackToRestore(0), HasStackFrame(false),
         NumLocalDynamicTLSAccesses(0), VarArgsStackIndex(0), VarArgsGPRIndex(0),
         VarArgsGPRSize(0), VarArgsFPRIndex(0), VarArgsFPRSize(0),
         StackRealigned(false) {}
 
   explicit AArch64FunctionInfo(MachineFunction &MF)
public:
   void setVarArgsGPRSize(unsigned Size) { VarArgsGPRSize = Size; }
 
   int getVarArgsFPRIndex() const { return VarArgsFPRIndex; }
   void setVarArgsFPRIndex(int Index) { VarArgsFPRIndex = Index; }
 
   unsigned getVarArgsFPRSize() const { return VarArgsFPRSize; }
   void setVarArgsFPRSize(unsigned Size) { VarArgsFPRSize = Size; }
 
+  const MCPhysReg *getCalleeSavedRegs() const { return CalleeSavedRegs; }
+  void setCalleeSavedRegs(const MCPhysReg *CSR) { CalleeSavedRegs = CSR; }
+
   typedef SmallPtrSet<const MachineInstr *, 16> SetOfInstructions;
 
   const SetOfInstructions &getLOHRelated() const { return LOHRelated; }
 
   // Shortcuts for LOH related types.
   class MILOHDirective {
     MCLOHType Kind;
 

lib/Target/AArch64/AArch64RegisterInfo.h

 private:
   const Triple &TT;
 
 public:
   AArch64RegisterInfo(const Triple &TT);
 
   bool isReservedReg(const MachineFunction &MF, unsigned Reg) const;
 
-  /// Code Generation virtual methods...
+  /// Return a zero teerminated list of callee saved registers for function \p
+  /// MF. This does not include callee saved registers handled by the SplitCSR
+  /// mechanism \see getCalleeSavedRegsViaCopy().
   const MCPhysReg *getCalleeSavedRegs(const MachineFunction *MF) const override;
   const uint32_t *getCallPreservedMask(const MachineFunction &MF,
                                        CallingConv::ID) const override;
 
+  /// Return a zero terminated list of callee saved registers for function
+  /// \p MF, also includes registers handled by SplitCSR.
+  static const MCPhysReg *getBasicCalleeSavedRegs(const MachineFunction *MF);
+
   unsigned getCSRFirstUseCost() const override {
     // The cost will be compared against BlockFrequency where entry has the
     // value of 1 << 14. A value of 5 will choose to spill or split really
     // cold path instead of using a callee-saved register.
     return 5;
   }
 
   // Calls involved in thread-local variable lookup save more registers than
public:
   // Debug information queries.
   unsigned getFrameRegister(const MachineFunction &MF) const override;
 
   unsigned getRegPressureLimit(const TargetRegisterClass *RC,
                                MachineFunction &MF) const override;
 
   /// Returns CSR_AArch64_CXX_TLS_Darwin_ViaCopy_SaveList.
   static const MCPhysReg *getFastTLSSavedViaCopy();
+  /// Returns CSR_AArch64_CXX_TLS_Darwin_PE_SaveList.
+  static const MCPhysReg *getFastTLSSaved();
 };
 
 } // end namespace llvm
 
 #endif

lib/Target/AArch64/AArch64RegisterInfo.cpp

 
 #define GET_REGINFO_TARGET_DESC
 #include "AArch64GenRegisterInfo.inc"
 
 AArch64RegisterInfo::AArch64RegisterInfo(const Triple &TT)
     : AArch64GenRegisterInfo(AArch64::LR), TT(TT) {}
 
 const MCPhysReg *
-AArch64RegisterInfo::getCalleeSavedRegs(const MachineFunction *MF) const {
+AArch64RegisterInfo::getBasicCalleeSavedRegs(const MachineFunction *MF) {
   assert(MF && "Invalid MachineFunction pointer.");
   if (MF->getFunction()->getCallingConv() == CallingConv::GHC)
     // GHC set of callee saved regs is empty as all those regs are
     // used for passing STG regs around
     return CSR_AArch64_NoRegs_SaveList;
   if (MF->getFunction()->getCallingConv() == CallingConv::AnyReg)
     return CSR_AArch64_AllRegs_SaveList;
   if (MF->getFunction()->getCallingConv() == CallingConv::CXX_FAST_TLS)
-    return MF->getInfo<AArch64FunctionInfo>()->getCSRSavedViaCopy() ?
-           CSR_AArch64_CXX_TLS_Darwin_PE_SaveList :
-           CSR_AArch64_CXX_TLS_Darwin_SaveList;
+    return CSR_AArch64_CXX_TLS_Darwin_SaveList;
   if (MF->getSubtarget<AArch64Subtarget>().getTargetLowering()
           ->supportSwiftError() &&
       MF->getFunction()->getAttributes().hasAttrSomewhere(
           Attribute::SwiftError))
     return CSR_AArch64_AAPCS_SwiftError_SaveList;
   if (MF->getFunction()->getCallingConv() == CallingConv::PreserveMost)
     return CSR_AArch64_RT_MostRegs_SaveList;
   else
     return CSR_AArch64_AAPCS_SaveList;
 }
 
+const MCPhysReg *
+AArch64RegisterInfo::getCalleeSavedRegs(const MachineFunction *MF) const {
+  const AArch64FunctionInfo &FuncInfo = *MF->getInfo<AArch64FunctionInfo>();
+  const MCPhysReg *CalleeSavedRegs = FuncInfo.getCalleeSavedRegs();
+  if (CalleeSavedRegs)
+    return CalleeSavedRegs;
+  return getBasicCalleeSavedRegs(MF);
+}
+
 const MCPhysReg *AArch64RegisterInfo::getFastTLSSavedViaCopy() {
   return CSR_AArch64_CXX_TLS_Darwin_ViaCopy_SaveList;
 }
 
+const MCPhysReg *AArch64RegisterInfo::getFastTLSSaved() {
+  return CSR_AArch64_CXX_TLS_Darwin_PE_SaveList;
+}
+
 const uint32_t *
 AArch64RegisterInfo::getCallPreservedMask(const MachineFunction &MF,
                                           CallingConv::ID CC) const {
   if (CC == CallingConv::GHC)
     // This is academic becase all GHC calls are (supposed to be) tail calls
     return CSR_AArch64_NoRegs_RegMask;
   if (CC == CallingConv::AnyReg)
     return CSR_AArch64_AllRegs_RegMask;

test/CodeGen/AArch64/swiftself.ll

 ; CHECK: mov x19, x0
 ; CHECK: bl {{_?}}swiftself_param
 ; CHECK: ret
 define i8 *@call_swiftself(i8* %arg) {
   %res = call i8 *@swiftself_param(i8* swiftself %arg)
   ret i8 *%res
 }
 
+declare void @func(i8 *)
+
+; CHECK-LABEL: swiftself_call_normal:
+; OPT-NOT: {{stp|str}} {{.*}}x19
+; OPT: mov x0, x19
+; OPT-NEXT: bl {{_?}}func
+; OPT-NOT: {{ldp|ldr}} {{.*}}x19
+; OPT: ret
+define void @swiftself_call_normal(i8* swiftself %addr0) nounwind {
+  call void @func(i8 *%addr0)
+  ret void
+}
+
 ; x19 should be saved by the callee even if used for swiftself
 ; CHECK-LABEL: swiftself_clobber:
 ; CHECK: {{stp|str}} {{.*}}x19{{.*}}sp
 ; ...
 ; CHECK: {{ldp|ldr}} {{.*}}x19{{.*}}sp
 ; CHECK: ret
 define i8 *@swiftself_clobber(i8* swiftself %addr0) {
   call void asm sideeffect "", "~{x19}"()
   ret i8 *%addr0
 }
 
 ; Demonstrate that we do not need any movs when calling multiple functions
 ; with swiftself argument.
 ; CHECK-LABEL: swiftself_passthrough:
+; OPT-NOT: stp {{.*}}x19
 ; OPT: bl {{_?}}swiftself_param
 ; OPT-NEXT: bl {{_?}}swiftself_param
 ; OPT: ret
 define void @swiftself_passthrough(i8* swiftself %addr0) {
   call i8 *@swiftself_param(i8* swiftself %addr0)
   call i8 *@swiftself_param(i8* swiftself %addr0)
   ret void
 }
 
 ; We can use a tail call if the callee swiftself is the same as the caller one.
 ; CHECK-LABEL: swiftself_tail:
 ; OPT: b {{_?}}swiftself_param
 ; OPT-NOT: ret
-define i8* @swiftself_tail(i8* swiftself %addr0) {
+define i8* @swiftself_tail(i8* swiftself %addr0) nounwind {
   call void asm sideeffect "", "~{x19}"()
   %res = tail call i8* @swiftself_param(i8* swiftself %addr0)
   ret i8* %res
 }
 
 ; We can not use a tail call if the callee swiftself is not the same as the
 ; caller one.
 ; CHECK-LABEL: swiftself_notail:
 ; CHECK: bl {{_?}}swiftself_param
 ; CHECK: ret
 define i8* @swiftself_notail(i8* swiftself %addr0, i8* %addr1) nounwind {
   %res = tail call i8* @swiftself_param(i8* swiftself %addr1)
   ret i8* %res
 }