[AArch64] Support adding X[8-15,18] registers as CSRs.

[trong]

Summary

Specifying X[8-15,18] registers as callee-saved is used to support
CONFIG_ARM64_LSE_ATOMICS in Linux kernel. As part of this patch we:

• use custom CSR list/mask when user specifies custom CSRs
• update Machine Register Info's list of CSRs with additional custom CSRs in

LowerCall and LowerFormalArguments.

[nickdesaulniers]

lib/Target/AArch64/AArch64ISelLowering.cpp

4054–4058

This pattern is repeated, a lot.

Maybe you could make a single function like:

void UpdateMaskCustomCall (const MachineFunction& MF, uint32_t** Mask) {
  ...
}
...
{
   uint32_t* Mask = ...;
   MachineFunction &MF = ...;
   UpdateMaskCustomCall(MF, &Mask);
}

?

[nickdesaulniers]

lib/Target/AArch64/AArch64.td

109

Since x18 can now be call-saved, any thoughts to how llvm should handle x18 being both call-saved and reserved (above line 102)? Maybe a test can be added for this case?

[trong]

• Added UpdateMaskCustomCall() util function to AArch64ISelLowering.
• Added test for case when x18 is specified as both reserved and call-saved.

[trong]

lib/Target/AArch64/AArch64.td

109

I think -fcall-saved-x18 should have no effect in presence of -ffixed-x18. If x18 is reserved, callee shouldn't allocate x18. So there shouldn't be a reason to save it. Added a test case.

lib/Target/AArch64/AArch64ISelLowering.cpp

4054–4058

MachineFunction will have to be passed non-const to make MF.allocateRegMask() call. So the signature would look something like:

void UpdateMaskCustomCall (MachineFunction &MF, uint32_t** Mask)

AArch64RegisterInfo is probably not a good place to change the state of a MachineFunction, so added a function to AArch64ISelLowering to do this. I left AArch64CallLowering as is since it only gets the custom mask once.

[nickdesaulniers]

lib/Target/AArch64/AArch64CallLowering.cpp

343

This looks very very similar to UpdateMaskCustomCall. Is there a way to convert a uint32_t[] to a SmallVector<MCPhysReg, 32>? That might help cut down on code duplication further.

390–394

UpdateMaskCustomCall?

lib/Target/AArch64/AArch64ISelLowering.cpp

3110

Don't we already have a reference to the Subtarget on the previous line? auto TRI = Subtarget.getRegisterInfo();?

7783–7784

It seems like all invocations of UpdateMaskCustomCall have this pattern:

if (x) UpdateMaskCustomCall(...);

where x is the same condition. Can the condition be moved in to UpdateMaskCustomCall and then UpdateMaskCustomCall called unconditionally, only actually doing anything if (x)?

lib/Target/AArch64/AArch64ISelLowering.h

550 ↗

(On Diff #166048)

Is there somewhere this can go so that it can continue to be used from lib/Target/AArch64/AArch64ISelLowering.cpp but additionally lib/Target/AArch64/AArch64CallLowering.cpp? See above comment.

It seems like Subtarget is accessible through the provided MF via MF.getSubtarget<AArch64Subtarget>().

lib/Target/AArch64/AArch64RegisterInfo.cpp

82–83

Can std::copy be used here?

85–89

Can std::copy_if be used here?

92

no need for empty return.

149–151

Just trying to understand what SubReg is in this case? I would think we only need to iterate over the number of registers once, so I don't understand the inner for loop. I've read the comment above the declaration of MCSubRegIterator, but I guess I don't understand what a subreg is in context of aarch64.

154

grep did not turn up matches for this method. Typo, or was a function you added then renamed?

160

Not an expert on the LLVM coding style, but no need for empty returns in void returning functions.

[trong]

• Added helper functions to AArch64RegisterInfo to update CSR list/mask with custom calling convention during instruction selection.

[trong]

lib/Target/AArch64/AArch64CallLowering.cpp

343

Since CSR masks and CSR lists have different format, we have to deal with them differently.

lib/Target/AArch64/AArch64ISelLowering.cpp

3110

right

7783–7784

Even though hasCustomCallingConv() checks are repetitive, I'd like to keep the check outside of UpdateMaskCustomCall for readability of instruction selection code.

lib/Target/AArch64/AArch64ISelLowering.h

550 ↗

(On Diff #166048)

Subtarget only exposes target information, so I don't think Subtarget is good place to make changes to MF. Added helper functions to AArch64RegisterInfo (not ideal either) so that they could be used by both isel implementations.

lib/Target/AArch64/AArch64RegisterInfo.cpp

82–83

We rely on last value of the list being 0. We'd have to look for the end of the list first to use std::copy, which is not as nice as looping through just once.

85–89

We don't have direct access to the copied container.

92

right

149–151

E.g. w0 (32-bit) is a subregister of x0 (64-bit)
CSR masks account for all subregisters:
https://github.com/llvm-mirror/llvm/blob/master/include/llvm/CodeGen/TargetRegisterInfo.h#L471
So while looping over X registers we also need to set the bit for the corresponding W register.

154

Yes, it's a typo, should be getCallPreservedMask.

[nickdesaulniers]

Thanks Tri, this looks good to me!

lib/Target/AArch64/AArch64.td

def FeatureStrictAlign : SubtargetFeature<"strict-align",
                                           "Disallow all unaligned memory "
                                           "access">;
 
 foreach i = {1-7,18,20} in
     def FeatureReserveX#i : SubtargetFeature<"reserve-x"#i, "ReserveXRegister["#i#"]", "true",
                                              "Reserve X"#i#", making it unavailable "
                                              "as a GPR">;
 
+foreach i = {8-15,18} in
+    def FeatureCallSavedX#i : SubtargetFeature<"call-saved-x"#i,
+         "CustomCallSavedXRegs["#i#"]", "true", "Make X"#i#" callee saved.">;
+
 def FeatureUseAA : SubtargetFeature<"use-aa", "UseAA", "true",
                                     "Use alias analysis during codegen">;
 
 def FeatureBalanceFPOps : SubtargetFeature<"balance-fp-ops", "BalanceFPOps",
     "true",
     "balance mix of odd and even D-registers for fp multiply(-accumulate) ops">;
 
 def FeaturePredictableSelectIsExpensive : SubtargetFeature<

lib/Target/AArch64/AArch64CallLowering.cpp

if (F.isVarArg()) {
     // We currently pass all varargs at 8-byte alignment.
     uint64_t StackOffset = alignTo(Handler.StackUsed, 8);
 
     auto &MFI = MIRBuilder.getMF().getFrameInfo();
     AArch64FunctionInfo *FuncInfo = MF.getInfo<AArch64FunctionInfo>();
     FuncInfo->setVarArgsStackIndex(MFI.CreateFixedObject(4, StackOffset, true));
   }
 
+  if (MF.getSubtarget<AArch64Subtarget>().hasCustomCallingConv()) {
+    auto TRI = MF.getSubtarget<AArch64Subtarget>().getRegisterInfo();
+    const MCPhysReg *CSRs = TRI->getCalleeSavedRegs(&MF);
+    SmallVector<MCPhysReg, 32> UpdatedCSRs;
+    TRI->getCustomCalleeSavedRegs(MF, CSRs, UpdatedCSRs);
+    MRI.setCalleeSavedRegs(UpdatedCSRs);
+  }
+
   // Move back to the end of the basic block.
   MIRBuilder.setMBB(MBB);
 
   return true;
 }
 
 bool AArch64CallLowering::lowerCall(MachineIRBuilder &MIRBuilder,
                                     CallingConv::ID CallConv,
bool AArch64CallLowering::lowerCall(MachineIRBuilder &MIRBuilder,
   // Create a temporarily-floating call instruction so we can add the implicit
   // uses of arg registers.
   auto MIB = MIRBuilder.buildInstrNoInsert(Callee.isReg() ? AArch64::BLR
                                                           : AArch64::BL);
   MIB.add(Callee);
 
   // Tell the call which registers are clobbered.
   auto TRI = MF.getSubtarget<AArch64Subtarget>().getRegisterInfo();
-  MIB.addRegMask(TRI->getCallPreservedMask(MF, F.getCallingConv()));
+  const uint32_t *Mask = TRI->getCallPreservedMask(MF, F.getCallingConv());
+  if (MF.getSubtarget<AArch64Subtarget>().hasCustomCallingConv()) {
+    uint32_t *UpdatedMask = MF.allocateRegMask();
+    TRI->getCustomCallPreservedMask(MF, Mask, UpdatedMask);
+    Mask = UpdatedMask;
+  }
+  MIB.addRegMask(Mask);
 
   if (TRI->isAnyArgRegReserved(MF))
     TRI->emitReservedArgRegCallError(MF);
 
   // Do the actual argument marshalling.
   SmallVector<unsigned, 8> PhysRegs;
   OutgoingArgHandler Handler(MIRBuilder, MRI, MIB, AssignFnFixed,
                              AssignFnVarArg);

lib/Target/AArch64/AArch64FastISel.cpp

bool AArch64FastISel::fastLowerArguments() {
   const Function *F = FuncInfo.Fn;
   if (F->isVarArg())
     return false;
 
   CallingConv::ID CC = F->getCallingConv();
   if (CC != CallingConv::C && CC != CallingConv::Swift)
     return false;
 
+  if (Subtarget->hasCustomCallingConv())
+    return false;
+
   // Only handle simple cases of up to 8 GPR and FPR each.
   unsigned GPRCnt = 0;
   unsigned FPRCnt = 0;
   for (auto const &Arg : F->args()) {
     if (Arg.hasAttribute(Attribute::ByVal) ||
         Arg.hasAttribute(Attribute::InReg) ||
         Arg.hasAttribute(Attribute::StructRet) ||
         Arg.hasAttribute(Attribute::SwiftSelf) ||

lib/Target/AArch64/AArch64FrameLowering.cpp

if (&MF->front() == MBB)
     return AArch64::X9;
 
   const AArch64Subtarget &Subtarget = MF->getSubtarget<AArch64Subtarget>();
   const AArch64RegisterInfo &TRI = *Subtarget.getRegisterInfo();
   LivePhysRegs LiveRegs(TRI);
   LiveRegs.addLiveIns(*MBB);
 
   // Mark callee saved registers as used so we will not choose them.
-  const MCPhysReg *CSRegs = TRI.getCalleeSavedRegs(MF);
+  const MCPhysReg *CSRegs = MF->getRegInfo().getCalleeSavedRegs();
   for (unsigned i = 0; CSRegs[i]; ++i)
     LiveRegs.addReg(CSRegs[i]);
 
   // Prefer X9 since it was historically used for the prologue scratch reg.
   const MachineRegisterInfo &MRI = MF->getRegInfo();
   if (LiveRegs.available(MRI, AArch64::X9))
     return AArch64::X9;
 
void AArch64FrameLowering::determineCalleeSaves(MachineFunction &MF,
   TargetFrameLowering::determineCalleeSaves(MF, SavedRegs, RS);
   const AArch64RegisterInfo *RegInfo = static_cast<const AArch64RegisterInfo *>(
       MF.getSubtarget().getRegisterInfo());
   AArch64FunctionInfo *AFI = MF.getInfo<AArch64FunctionInfo>();
   unsigned UnspilledCSGPR = AArch64::NoRegister;
   unsigned UnspilledCSGPRPaired = AArch64::NoRegister;
 
   MachineFrameInfo &MFI = MF.getFrameInfo();
-  const MCPhysReg *CSRegs = RegInfo->getCalleeSavedRegs(&MF);
+  const MCPhysReg *CSRegs = MF.getRegInfo().getCalleeSavedRegs();
 
   unsigned BasePointerReg = RegInfo->hasBasePointer(MF)
                                 ? RegInfo->getBaseRegister()
                                 : (unsigned)AArch64::NoRegister;
 
   unsigned ExtraCSSpill = 0;
   // Figure out which callee-saved registers to save/restore.
   for (unsigned i = 0; CSRegs[i]; ++i) {

lib/Target/AArch64/AArch64ISelLowering.cpp

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);
 
+  if (Subtarget->hasCustomCallingConv()) {
+    auto TRI = MF.getSubtarget<AArch64Subtarget>().getRegisterInfo();
+    const MCPhysReg *CSRs = TRI->getCalleeSavedRegs(&MF);
+    SmallVector<MCPhysReg, 32> UpdatedCSRs;
+    TRI->getCustomCalleeSavedRegs(MF, CSRs, UpdatedCSRs);
+    MF.getRegInfo().setCalleeSavedRegs(UpdatedCSRs);
+  }
+
   return Chain;
 }
 
 void AArch64TargetLowering::saveVarArgRegisters(CCState &CCInfo,
                                                 SelectionDAG &DAG,
                                                 const SDLoc &DL,
                                                 SDValue &Chain) const {
   MachineFunction &MF = DAG.getMachineFunction();
if (!CCState::resultsCompatible(CalleeCC, CallerCC, MF, C, Ins,
                                   CCAssignFnForCall(CalleeCC, isVarArg),
                                   CCAssignFnForCall(CallerCC, isVarArg)))
     return false;
   // The callee has to preserve all registers the caller needs to preserve.
   const AArch64RegisterInfo *TRI = Subtarget->getRegisterInfo();
   const uint32_t *CallerPreserved = TRI->getCallPreservedMask(MF, CallerCC);
   if (!CCMatch) {
     const uint32_t *CalleePreserved = TRI->getCallPreservedMask(MF, CalleeCC);
+    if (Subtarget->hasCustomCallingConv()) {
+      uint32_t *UpdatedCallerMask = MF.allocateRegMask();
+      uint32_t *UpdatedCalleeMask = MF.allocateRegMask();
+      TRI->getCustomCallPreservedMask(MF, CallerPreserved, UpdatedCallerMask);
+      TRI->getCustomCallPreservedMask(MF, CalleePreserved, UpdatedCalleeMask);
+      CallerPreserved = UpdatedCallerMask;
+      CalleePreserved = UpdatedCalleeMask;
+    }
     if (!TRI->regmaskSubsetEqual(CallerPreserved, CalleePreserved))
       return false;
   }
 
   // Nothing more to check if the callee is taking no arguments
   if (Outs.empty())
     return true;
 
if (IsThisReturn) {
     Mask = TRI->getThisReturnPreservedMask(MF, CallConv);
     if (!Mask) {
       IsThisReturn = false;
       Mask = TRI->getCallPreservedMask(MF, CallConv);
     }
   } else
     Mask = TRI->getCallPreservedMask(MF, CallConv);
 
+  if (MF.getSubtarget<AArch64Subtarget>().hasCustomCallingConv()) {
+    uint32_t *UpdatedMask = MF.allocateRegMask();
+    TRI->getCustomCallPreservedMask(MF, Mask, UpdatedMask);
+    Mask = UpdatedMask;
+  }
+
   if (TRI->isAnyArgRegReserved(MF))
     TRI->emitReservedArgRegCallError(MF);
 
   assert(Mask && "Missing call preserved mask for calling convention");
   Ops.push_back(DAG.getRegisterMask(Mask));
 
   if (InFlag.getNode())
     Ops.push_back(InFlag);

 }
 
 SDValue
 AArch64TargetLowering::LowerReturn(SDValue Chain, CallingConv::ID CallConv,
                                    bool isVarArg,
                                    const SmallVectorImpl<ISD::OutputArg> &Outs,
                                    const SmallVectorImpl<SDValue> &OutVals,
                                    const SDLoc &DL, SelectionDAG &DAG) const {
+  MachineFunction &MF = DAG.getMachineFunction();
   CCAssignFn *RetCC = CallConv == CallingConv::WebKit_JS
                           ? RetCC_AArch64_WebKit_JS
                           : RetCC_AArch64_AAPCS;
   SmallVector<CCValAssign, 16> RVLocs;
-  CCState CCInfo(CallConv, isVarArg, DAG.getMachineFunction(), RVLocs,
-                 *DAG.getContext());
+  CCState CCInfo(CallConv, isVarArg, MF, RVLocs, *DAG.getContext());
   CCInfo.AnalyzeReturn(Outs, RetCC);
 
   // Copy the result values into the output registers.
   SDValue Flag;
   SmallVector<SDValue, 4> RetOps(1, Chain);
   for (unsigned i = 0, realRVLocIdx = 0; i != RVLocs.size();
        ++i, ++realRVLocIdx) {
     CCValAssign &VA = RVLocs[i];
case CCValAssign::BCvt:
       break;
     }
 
     Chain = DAG.getCopyToReg(Chain, DL, VA.getLocReg(), Arg, Flag);
     Flag = Chain.getValue(1);
     RetOps.push_back(DAG.getRegister(VA.getLocReg(), VA.getLocVT()));
   }
   const AArch64RegisterInfo *TRI = Subtarget->getRegisterInfo();
-  const MCPhysReg *I =
-      TRI->getCalleeSavedRegsViaCopy(&DAG.getMachineFunction());
+  const MCPhysReg *I = TRI->getCalleeSavedRegsViaCopy(&MF);
   if (I) {
+    SmallVector<MCPhysReg, 32> UpdatedCSRs;
+    if (Subtarget->hasCustomCallingConv()) {
+      TRI->getCustomCalleeSavedRegs(MF, I, UpdatedCSRs);
+      I = UpdatedCSRs.begin();
+    }
     for (; *I; ++I) {
       if (AArch64::GPR64RegClass.contains(*I))
         RetOps.push_back(DAG.getRegister(*I, MVT::i64));
       else if (AArch64::FPR64RegClass.contains(*I))
         RetOps.push_back(DAG.getRegister(*I, MVT::getFloatingPointVT(64)));
       else
         llvm_unreachable("Unexpected register class in CSRsViaCopy!");
     }
AArch64TargetLowering::LowerDarwinGlobalTLSAddress(SDValue Op,
   MachineFrameInfo &MFI = DAG.getMachineFunction().getFrameInfo();
   MFI.setAdjustsStack(true);
 
   // TLS calls preserve all registers except those that absolutely must be
   // trashed: X0 (it takes an argument), LR (it's a call) and NZCV (let's not be
   // silly).
   const uint32_t *Mask =
       Subtarget->getRegisterInfo()->getTLSCallPreservedMask();
+  MachineFunction &MF = DAG.getMachineFunction();
+  const AArch64RegisterInfo *TRI = Subtarget->getRegisterInfo();
+  if (MF.getSubtarget<AArch64Subtarget>().hasCustomCallingConv()) {
+    uint32_t *UpdatedMask = MF.allocateRegMask();
+    TRI->getCustomCallPreservedMask(MF, Mask, UpdatedMask);
+    Mask = UpdatedMask;
+  }
 
   // Finally, we can make the call. This is just a degenerate version of a
   // normal AArch64 call node: x0 takes the address of the descriptor, and
   // returns the address of the variable in this thread.
   Chain = DAG.getCopyToReg(Chain, DL, AArch64::X0, DescAddr, SDValue());
   Chain =
       DAG.getNode(AArch64ISD::CALL, DL, DAG.getVTList(MVT::Other, MVT::Glue),
                   Chain, FuncTLVGet, DAG.getRegister(AArch64::X0, MVT::i64),

 SDValue AArch64TargetLowering::LowerWindowsDYNAMIC_STACKALLOC(
     SDValue Op, SDValue Chain, SDValue &Size, SelectionDAG &DAG) const {
   SDLoc dl(Op);
   EVT PtrVT = getPointerTy(DAG.getDataLayout());
   SDValue Callee = DAG.getTargetExternalSymbol("__chkstk", PtrVT, 0);
 
   const uint32_t *Mask =
       Subtarget->getRegisterInfo()->getWindowsStackProbePreservedMask();
+  MachineFunction &MF = DAG.getMachineFunction();
+  const AArch64RegisterInfo *TRI = Subtarget->getRegisterInfo();
+  if (MF.getSubtarget<AArch64Subtarget>().hasCustomCallingConv()) {
+    uint32_t *UpdatedMask = MF.allocateRegMask();
+    TRI->getCustomCallPreservedMask(MF, Mask, UpdatedMask);
+    Mask = UpdatedMask;
+  }
 
   Size = DAG.getNode(ISD::SRL, dl, MVT::i64, Size,
                      DAG.getConstant(4, dl, MVT::i64));
   Chain = DAG.getCopyToReg(Chain, dl, AArch64::X15, Size, SDValue());
   Chain =
       DAG.getNode(AArch64ISD::CALL, dl, DAG.getVTList(MVT::Other, MVT::Glue),
                   Chain, Callee, DAG.getRegister(AArch64::X15, MVT::i64),
                   DAG.getRegisterMask(Mask), Chain.getValue(1));
void AArch64TargetLowering::initializeSplitCSR(MachineBasicBlock *Entry) const {
   // Update IsSplitCSR in AArch64unctionInfo.
   AArch64FunctionInfo *AFI = Entry->getParent()->getInfo<AArch64FunctionInfo>();
   AFI->setIsSplitCSR(true);
 }
 
 void AArch64TargetLowering::insertCopiesSplitCSR(
     MachineBasicBlock *Entry,
     const SmallVectorImpl<MachineBasicBlock *> &Exits) const {
+  MachineFunction &MF = *Entry->getParent();
   const AArch64RegisterInfo *TRI = Subtarget->getRegisterInfo();
   const MCPhysReg *IStart = TRI->getCalleeSavedRegsViaCopy(Entry->getParent());
   if (!IStart)
     return;
+  SmallVector<MCPhysReg, 32> UpdatedCSRs;
+  if (Subtarget->hasCustomCallingConv()) {
+    TRI->getCustomCalleeSavedRegs(MF, IStart, UpdatedCSRs);
+    IStart = UpdatedCSRs.begin();
+  }
 
   const TargetInstrInfo *TII = Subtarget->getInstrInfo();
   MachineRegisterInfo *MRI = &Entry->getParent()->getRegInfo();
   MachineBasicBlock::iterator MBBI = Entry->begin();
   for (const MCPhysReg *I = IStart; *I; ++I) {
     const TargetRegisterClass *RC = nullptr;
     if (AArch64::GPR64RegClass.contains(*I))
       RC = &AArch64::GPR64RegClass;

lib/Target/AArch64/AArch64RegisterInfo.h

 
 public:
   AArch64RegisterInfo(const Triple &TT);
 
   bool isReservedReg(const MachineFunction &MF, unsigned Reg) const;
   bool isAnyArgRegReserved(const MachineFunction &MF) const;
   void emitReservedArgRegCallError(const MachineFunction &MF) const;
 
+  void getCustomCalleeSavedRegs(const MachineFunction &MF,
+                                const MCPhysReg *CSRs,
+                                SmallVectorImpl<MCPhysReg> &UpdatedCSRs) const;
+  void getCustomCallPreservedMask(const MachineFunction &MF,
+                                  const uint32_t *Mask,
+                                  uint32_t *UpdatedMask) const;
+
   /// Code Generation virtual methods...
   const MCPhysReg *getCalleeSavedRegs(const MachineFunction *MF) const override;
   const MCPhysReg *
   getCalleeSavedRegsViaCopy(const MachineFunction *MF) const;
   const uint32_t *getCallPreservedMask(const MachineFunction &MF,
                                        CallingConv::ID) const override;
 
   unsigned getCSRFirstUseCost() const override {

lib/Target/AArch64/AArch64RegisterInfo.cpp

const MCPhysReg *AArch64RegisterInfo::getCalleeSavedRegsViaCopy(
     const MachineFunction *MF) const {
   assert(MF && "Invalid MachineFunction pointer.");
   if (MF->getFunction().getCallingConv() == CallingConv::CXX_FAST_TLS &&
       MF->getInfo<AArch64FunctionInfo>()->isSplitCSR())
     return CSR_AArch64_CXX_TLS_Darwin_ViaCopy_SaveList;
   return nullptr;
 }
 
+void AArch64RegisterInfo::getCustomCalleeSavedRegs(
+      const MachineFunction &MF,
+      const MCPhysReg *CSRs,
+      SmallVectorImpl<MCPhysReg> &UpdatedCSRs) const {
+  for (const MCPhysReg *I = CSRs; *I; ++I)
+    UpdatedCSRs.push_back(*I);
+
+  for (size_t i = 0; i < AArch64::GPR64commonRegClass.getNumRegs(); ++i) {
+    if (MF.getSubtarget<AArch64Subtarget>().isXRegCustomCalleeSaved(i)) {
+      UpdatedCSRs.push_back(AArch64::GPR64commonRegClass.getRegister(i));
+    }
+  }
+  // Register lists are zero-terminated.
+  UpdatedCSRs.push_back(0);
+  return;
+}
+
 const TargetRegisterClass *
 AArch64RegisterInfo::getSubClassWithSubReg(const TargetRegisterClass *RC,
                                        unsigned Idx) const {
   // edge case for GPR/FPR register classes
   if (RC == &AArch64::GPR32allRegClass && Idx == AArch64::hsub)
     return &AArch64::FPR32RegClass;
   else if (RC == &AArch64::GPR64allRegClass && Idx == AArch64::hsub)
     return &AArch64::FPR64RegClass;

 const uint32_t *AArch64RegisterInfo::getTLSCallPreservedMask() const {
   if (TT.isOSDarwin())
     return CSR_AArch64_TLS_Darwin_RegMask;
 
   assert(TT.isOSBinFormatELF() && "Invalid target");
   return CSR_AArch64_TLS_ELF_RegMask;
 }
 
+void AArch64RegisterInfo::getCustomCallPreservedMask(
+    const MachineFunction &MF,
+    const uint32_t *Mask,
+    uint32_t *UpdatedMask) const {
+  unsigned RegMaskSize = MachineOperand::getRegMaskSize(getNumRegs());
+  memcpy(UpdatedMask, Mask, sizeof(Mask[0]) * RegMaskSize);
+
+  for (size_t i = 0; i < AArch64::GPR64commonRegClass.getNumRegs(); ++i) {
+    if (MF.getSubtarget<AArch64Subtarget>().isXRegCustomCalleeSaved(i)) {
+      for (MCSubRegIterator SubReg(AArch64::GPR64commonRegClass.getRegister(i),
+                                   this, true);
+           SubReg.isValid(); ++SubReg) {
+        // See TargetRegisterInfo::getCallPreserverMask for how to interpret the
+        // register mask.
+        UpdatedMask[*SubReg / 32] |= 1u << (*SubReg % 32);
+      }
+    }
+  }
+  return;
+}
+
 const uint32_t *
 AArch64RegisterInfo::getThisReturnPreservedMask(const MachineFunction &MF,
                                                 CallingConv::ID CC) const {
   // This should return a register mask that is the same as that returned by
   // getCallPreservedMask but that additionally preserves the register used for
   // the first i64 argument (which must also be the register used to return a
   // single i64 return value)
   //

lib/Target/AArch64/AArch64Subtarget.h

protected:
   unsigned PrefFunctionAlignment = 0;
   unsigned PrefLoopAlignment = 0;
   unsigned MaxJumpTableSize = 0;
   unsigned WideningBaseCost = 0;
 
   // ReserveXRegister[i] - X#i is not available as a general purpose register.
   BitVector ReserveXRegister;
 
+  // CustomCallUsedXRegister[i] - X#i call saved.
+  BitVector CustomCallSavedXRegs;
+
   bool IsLittle;
 
   /// TargetTriple - What processor and OS we're targeting.
   Triple TargetTriple;
 
   AArch64FrameLowering FrameLowering;
   AArch64InstrInfo InstrInfo;
   AArch64SelectionDAGInfo TSInfo;
public:
   bool isXRaySupported() const override { return true; }
 
   unsigned getMinVectorRegisterBitWidth() const {
     return MinVectorRegisterBitWidth;
   }
 
   bool isXRegisterReserved(size_t i) const { return ReserveXRegister[i]; }
   unsigned getNumXRegisterReserved() const { return ReserveXRegister.count(); }
+  bool isXRegCustomCalleeSaved(size_t i) const {
+    return CustomCallSavedXRegs[i];
+  }
+  bool hasCustomCallingConv() const { return CustomCallSavedXRegs.any(); }
   bool hasFPARMv8() const { return HasFPARMv8; }
   bool hasNEON() const { return HasNEON; }
   bool hasCrypto() const { return HasCrypto; }
   bool hasDotProd() const { return HasDotProd; }
   bool hasCRC() const { return HasCRC; }
   bool hasLSE() const { return HasLSE; }
   bool hasRAS() const { return HasRAS; }
   bool hasRDM() const { return HasRDM; }

lib/Target/AArch64/AArch64Subtarget.cpp

void AArch64Subtarget::initializeProperties() {
   }
 }
 
 AArch64Subtarget::AArch64Subtarget(const Triple &TT, const std::string &CPU,
                                    const std::string &FS,
                                    const TargetMachine &TM, bool LittleEndian)
     : AArch64GenSubtargetInfo(TT, CPU, FS),
       ReserveXRegister(AArch64::GPR64commonRegClass.getNumRegs()),
+      CustomCallSavedXRegs(AArch64::GPR64commonRegClass.getNumRegs()),
       IsLittle(LittleEndian),
       TargetTriple(TT), FrameLowering(),
       InstrInfo(initializeSubtargetDependencies(FS, CPU)), TSInfo(),
       TLInfo(TM, *this) {
   if (AArch64::isX18ReservedByDefault(TT))
     ReserveXRegister.set(18);
 
   CallLoweringInfo.reset(new AArch64CallLowering(*getTargetLowering()));

test/CodeGen/AArch64/arm64-custom-call-saved-reg.ll

+; RUN: llc -mtriple=arm64-linux-gnu -mattr=+call-saved-x8 -o - %s \
+; RUN: | FileCheck %s --check-prefix=CHECK-SAVED-X8
+
+; RUN: llc -mtriple=arm64-linux-gnu -mattr=+call-saved-x9 -o - %s \
+; RUN: | FileCheck %s --check-prefix=CHECK-SAVED-X9
+
+; RUN: llc -mtriple=arm64-linux-gnu -mattr=+call-saved-x10 -o - %s \
+; RUN: | FileCheck %s --check-prefix=CHECK-SAVED-X10
+
+; RUN: llc -mtriple=arm64-linux-gnu -mattr=+call-saved-x11 -o - %s \
+; RUN: | FileCheck %s --check-prefix=CHECK-SAVED-X11
+
+; RUN: llc -mtriple=arm64-linux-gnu -mattr=+call-saved-x12 -o - %s \
+; RUN: | FileCheck %s --check-prefix=CHECK-SAVED-X12
+
+; RUN: llc -mtriple=arm64-linux-gnu -mattr=+call-saved-x13 -o - %s \
+; RUN: | FileCheck %s --check-prefix=CHECK-SAVED-X13
+
+; RUN: llc -mtriple=arm64-linux-gnu -mattr=+call-saved-x14 -o - %s \
+; RUN: | FileCheck %s --check-prefix=CHECK-SAVED-X14
+
+; RUN: llc -mtriple=arm64-linux-gnu -mattr=+call-saved-x15 -o - %s \
+; RUN: | FileCheck %s --check-prefix=CHECK-SAVED-X15
+
+; RUN: llc -mtriple=arm64-linux-gnu -mattr=+call-saved-x18 -o - %s \
+; RUN: | FileCheck %s --check-prefix=CHECK-SAVED-X18
+
+; Test all call-saved-x# options together.
+; RUN: llc -mtriple=arm64-linux-gnu \
+; RUN: -mattr=+call-saved-x8 \
+; RUN: -mattr=+call-saved-x9 \
+; RUN: -mattr=+call-saved-x10 \
+; RUN: -mattr=+call-saved-x11 \
+; RUN: -mattr=+call-saved-x12 \
+; RUN: -mattr=+call-saved-x13 \
+; RUN: -mattr=+call-saved-x14 \
+; RUN: -mattr=+call-saved-x15 \
+; RUN: -mattr=+call-saved-x18 \
+; RUN: -o - %s | FileCheck %s \
+; RUN: --check-prefix=CHECK-SAVED-ALL
+
+; Test GlobalISel.
+; RUN: llc -mtriple=arm64-linux-gnu \
+; RUN: -mattr=+call-saved-x8 \
+; RUN: -mattr=+call-saved-x9 \
+; RUN: -mattr=+call-saved-x10 \
+; RUN: -mattr=+call-saved-x11 \
+; RUN: -mattr=+call-saved-x12 \
+; RUN: -mattr=+call-saved-x13 \
+; RUN: -mattr=+call-saved-x14 \
+; RUN: -mattr=+call-saved-x15 \
+; RUN: -mattr=+call-saved-x18 \
+; RUN: -global-isel \
+; RUN: -o - %s | FileCheck %s \
+; RUN: --check-prefix=CHECK-SAVED-ALL
+
+; Used to exhaust the supply of GPRs.
+@var = global [30 x i64] zeroinitializer
+
+; Check that callee preserves additional CSRs.
+define void @callee() {
+; CHECK-LABEL: callee
+
+; CHECK-SAVED-X8: str x8, [sp
+; CHECK-SAVED-X9: str x9, [sp
+; CHECK-SAVED-X10: str x10, [sp
+; CHECK-SAVED-X11: str x11, [sp
+; CHECK-SAVED-X12: str x12, [sp
+; CHECK-SAVED-X13: str x13, [sp
+; CHECK-SAVED-X14: str x14, [sp
+; CHECK-SAVED-X15: str x15, [sp
+; CHECK-SAVED-X18: str x18, [sp
+
+; CHECK-SAVED-ALL: str x18, [sp
+; CHECK-SAVED-ALL-NEXT: stp x15, x14, [sp
+; CHECK-SAVED-ALL-NEXT: stp x13, x12, [sp
+; CHECK-SAVED-ALL-NEXT: stp x11, x10, [sp
+; CHECK-SAVED-ALL-NEXT: stp x9, x8, [sp
+
+  %val = load volatile [30 x i64], [30 x i64]* @var
+  store volatile [30 x i64] %val, [30 x i64]* @var
+
+; CHECK-SAVED-ALL: ldp x9, x8, [sp
+; CHECK-SAVED-ALL-NEXT: ldp x11, x10, [sp
+; CHECK-SAVED-ALL-NEXT: ldp x13, x12, [sp
+; CHECK-SAVED-ALL-NEXT: ldp x15, x14, [sp
+; CHECK-SAVED-ALL-NEXT: ldr x18, [sp
+
+; CHECK-SAVED-X8: ldr x8, [sp
+; CHECK-SAVED-X9: ldr x9, [sp
+; CHECK-SAVED-X10: ldr x10, [sp
+; CHECK-SAVED-X11: ldr x11, [sp
+; CHECK-SAVED-X12: ldr x12, [sp
+; CHECK-SAVED-X13: ldr x13, [sp
+; CHECK-SAVED-X14: ldr x14, [sp
+; CHECK-SAVED-X15: ldr x15, [sp
+; CHECK-SAVED-X18: ldr x18, [sp
+
+  ret void
+}
+
+; Check that caller doesn't shy away from allocating additional CSRs.
+define void @caller() {
+; CHECK-LABEL: caller
+
+  %val = load volatile [30 x i64], [30 x i64]* @var
+; CHECK-SAVED-X8: adrp x8, var
+; CHECK-SAVED-X9: adrp x9, var
+; CHECK-SAVED-X10: adrp x10, var
+; CHECK-SAVED-X11: adrp x11, var
+; CHECK-SAVED-X12: adrp x12, var
+; CHECK-SAVED-X13: adrp x13, var
+; CHECK-SAVED-X14: adrp x14, var
+; CHECK-SAVED-X15: adrp x15, var
+; CHECK-SAVED-X18: adrp x18, var
+
+; CHECK-SAVED-ALL: adrp x8, var
+; CHECK-SAVED-ALL-DAG: ldr x9
+; CHECK-SAVED-ALL-DAG: ldr x10
+; CHECK-SAVED-ALL-DAG: ldr x11
+; CHECK-SAVED-ALL-DAG: ldr x12
+; CHECK-SAVED-ALL-DAG: ldr x13
+; CHECK-SAVED-ALL-DAG: ldr x14
+; CHECK-SAVED-ALL-DAG: ldr x15
+; CHECK-SAVED-ALL-DAG: ldr x18
+
+  call void @callee()
+; CHECK: bl callee
+
+  store volatile [30 x i64] %val, [30 x i64]* @var
+; CHECK-SAVED-ALL-DAG: str x9
+; CHECK-SAVED-ALL-DAG: str x10
+; CHECK-SAVED-ALL-DAG: str x11
+; CHECK-SAVED-ALL-DAG: str x12
+; CHECK-SAVED-ALL-DAG: str x13
+; CHECK-SAVED-ALL-DAG: str x14
+; CHECK-SAVED-ALL-DAG: str x15
+; CHECK-SAVED-ALL-DAG: str x18
+
+  ret void
+}