Separate CSRs into a subset handled via explicit copies and a subset handled by prologue/epilogue (AArch64)

[manmanren]

Summary

This is the target-specific part of http://reviews.llvm.org/D15340

[hfinkel]

lib/Target/AArch64/AArch64ISelLowering.cpp

10055 ↗

(On Diff #42199)

Do you need to emit any kind of CFI pseudo-instructions for these so that unwinding works?

[manmanren]

lib/Target/AArch64/AArch64ISelLowering.cpp

10055 ↗

(On Diff #42199)

This is a good point. I think we should but I am not so sure.

Maybe Nick is a better person to answer the question on Darwin platform.

Specific to C++-style TLS access functions, they should be nounwind. Throwing an exception from the initializer of a thread-local variable is defined to immediately result in a call to std::terminate(); there’s no legal way for unwinding to actually leave such a function. (quote from John McCall) :]

[hfinkel]

lib/Target/AArch64/AArch64ISelLowering.cpp

10055 ↗

(On Diff #42199)

In that case, please add a comment about this. If someone goes to generalize this later (or port it to another target), it should be obvious that this is something they'll need to think about.

[manmanren]

Addressing Hal's comment. Patch is also updated to work with the target independent portion (http://reviews.llvm.org/D15340).

[qcolombet]

Hi Manman,

Thanks for your patience.

Here are a few comments.

Cheers,
-Quentin

lib/Target/AArch64/AArch64ISelLowering.cpp

3286 ↗

(On Diff #42539)

Why is this using a pointer type instead of an i64?

3287 ↗

(On Diff #42539)

No need for {}

3290 ↗

(On Diff #42539)

Ditto.

3293 ↗

(On Diff #42539)

Was this clang-formated?
The layout looks funny to me.

10069 ↗

(On Diff #42539)

Could you assert that the MachineFunction have the nounwind attribute?

lib/Target/AArch64/AArch64ISelLowering.h

392 ↗

(On Diff #42539)

I think we should change the hook to take a MachineFunction.
That way we could check that the function does have the nounwind attribute and if it does not, we can say this is not supported.

lib/Target/AArch64/AArch64MachineFunctionInfo.h

75 ↗

(On Diff #42539)

Add a comment for that field.

lib/Target/AArch64/AArch64RegisterInfo.cpp

61 ↗

(On Diff #42539)

Isn’t this redundant with the isSplitCSR check?
I mean, if we want to support more calling convention, we should minimize the places we have to patch to do so.

test/CodeGen/AArch64/cxx-tlscc.ll

50 ↗

(On Diff #42539)

Nice! :)

[manmanren]

Thanks Quentin for reviewing the patch!

Addressing his comments with the updated patch.

[manmanren]

Addressing review comments,

Cheers,
Manman

lib/Target/AArch64/AArch64RegisterInfo.cpp

61 ↗

(On Diff #42893)

If we are going to support other calling conventions with split CSR, it will be listed here with "else if", so we know which corresponding ViaCopy_SaveList to use. Similarly in getCalleeSavedRegs above.

[qcolombet]

Hi Manman,

LGTM.

Thanks,
-Quentin

llvm/trunk/lib/Target/AArch64/AArch64CallingConvention.td

 // should belong to the same class as well. Since the access function on the
 // fast path calls a function that follows CSR_AArch64_TLS_Darwin,
 // CSR_AArch64_CXX_TLS_Darwin should be a subset of CSR_AArch64_TLS_Darwin.
 def CSR_AArch64_CXX_TLS_Darwin
     : CalleeSavedRegs<(add CSR_AArch64_AAPCS,
                            (sub (sequence "X%u", 1, 28), X15, X16, X17, X18),
                            (sequence "D%u", 0, 31))>;
 
+// CSRs that are handled by prologue, epilogue.
+def CSR_AArch64_CXX_TLS_Darwin_PE
+    : CalleeSavedRegs<(add LR, FP)>;
+
+// CSRs that are handled explicitly via copies.
+def CSR_AArch64_CXX_TLS_Darwin_ViaCopy
+    : CalleeSavedRegs<(sub CSR_AArch64_CXX_TLS_Darwin, LR, FP)>;
+
 // The ELF stub used for TLS-descriptor access saves every feasible
 // register. Only X0 and LR are clobbered.
 def CSR_AArch64_TLS_ELF
     : CalleeSavedRegs<(add (sequence "X%u", 1, 28), FP,
                            (sequence "Q%u", 0, 31))>;
 
 def CSR_AArch64_AllRegs
     : CalleeSavedRegs<(add (sequence "W%u", 0, 30), WSP,
                            (sequence "X%u", 0, 28), FP, LR, SP,
                            (sequence "B%u", 0, 31), (sequence "H%u", 0, 31),
                            (sequence "S%u", 0, 31), (sequence "D%u", 0, 31),
                            (sequence "Q%u", 0, 31))>;
 
 def CSR_AArch64_NoRegs : CalleeSavedRegs<(add)>;

llvm/trunk/lib/Target/AArch64/AArch64FastISel.cpp

bool AArch64FastISel::selectRet(const Instruction *I) {
   const Function &F = *I->getParent()->getParent();
 
   if (!FuncInfo.CanLowerReturn)
     return false;
 
   if (F.isVarArg())
     return false;
 
+  if (TLI.supportSplitCSR(FuncInfo.MF))
+    return false;
+
   // Build a list of return value registers.
   SmallVector<unsigned, 4> RetRegs;
 
   if (Ret->getNumOperands() > 0) {
     CallingConv::ID CC = F.getCallingConv();
     SmallVector<ISD::OutputArg, 4> Outs;
     GetReturnInfo(F.getReturnType(), F.getAttributes(), Outs, TLI, DL);
 

llvm/trunk/lib/Target/AArch64/AArch64ISelLowering.h

public:
 
   bool isCheapToSpeculateCttz() const override {
     return true;
   }
 
   bool isCheapToSpeculateCtlz() const override {
     return true;
   }
+  bool supportSplitCSR(MachineFunction *MF) const override {
+    return MF->getFunction()->getCallingConv() == CallingConv::CXX_FAST_TLS &&
+           MF->getFunction()->hasFnAttribute(Attribute::NoUnwind);
+  }
+  void initializeSplitCSR(MachineBasicBlock *Entry) const override;
+  void insertCopiesSplitCSR(
+      MachineBasicBlock *Entry,
+      const SmallVectorImpl<MachineBasicBlock *> &Exits) const override;
 
 private:
   bool isExtFreeImpl(const Instruction *Ext) const override;
 
   /// Keep a pointer to the AArch64Subtarget around so that we can
   /// make the right decision when generating code for different targets.
   const AArch64Subtarget *Subtarget;
 

llvm/trunk/lib/Target/AArch64/AArch64ISelLowering.cpp

case CCValAssign::BCvt:
       Arg = DAG.getNode(ISD::BITCAST, DL, VA.getLocVT(), Arg);
       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());
+  if (I) {
+    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!");
+    }
+  }
 
   RetOps[0] = Chain; // Update chain.
 
   // Add the flag if we have it.
   if (Flag.getNode())
     RetOps.push_back(Flag);
 
   return DAG.getNode(AArch64ISD::RET_FLAG, DL, MVT::Other, RetOps);
Value *AArch64TargetLowering::getSafeStackPointerLocation(IRBuilder<> &IRB) const {
   const unsigned TlsOffset = 0x48;
   Module *M = IRB.GetInsertBlock()->getParent()->getParent();
   Function *ThreadPointerFunc =
       Intrinsic::getDeclaration(M, Intrinsic::aarch64_thread_pointer);
   return IRB.CreatePointerCast(
       IRB.CreateConstGEP1_32(IRB.CreateCall(ThreadPointerFunc), TlsOffset),
       Type::getInt8PtrTy(IRB.getContext())->getPointerTo(0));
 }
+
+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 {
+  const AArch64RegisterInfo *TRI = Subtarget->getRegisterInfo();
+  const MCPhysReg *IStart = TRI->getCalleeSavedRegsViaCopy(Entry->getParent());
+  if (!IStart)
+    return;
+
+  const TargetInstrInfo *TII = Subtarget->getInstrInfo();
+  MachineRegisterInfo *MRI = &Entry->getParent()->getRegInfo();
+  for (const MCPhysReg *I = IStart; *I; ++I) {
+    const TargetRegisterClass *RC = nullptr;
+    if (AArch64::GPR64RegClass.contains(*I))
+      RC = &AArch64::GPR64RegClass;
+    else if (AArch64::FPR64RegClass.contains(*I))
+      RC = &AArch64::FPR64RegClass;
+    else
+      llvm_unreachable("Unexpected register class in CSRsViaCopy!");
+
+    unsigned NewVR = MRI->createVirtualRegister(RC);
+    // Create copy from CSR to a virtual register.
+    // FIXME: this currently does not emit CFI pseudo-instructions, it works
+    // fine for CXX_FAST_TLS since the C++-style TLS access functions should be
+    // nounwind. If we want to generalize this later, we may need to emit
+    // CFI pseudo-instructions.
+    assert(Entry->getParent()->getFunction()->hasFnAttribute(
+               Attribute::NoUnwind) &&
+           "Function should be nounwind in insertCopiesSplitCSR!");
+    Entry->addLiveIn(*I);
+    BuildMI(*Entry, Entry->begin(), DebugLoc(), TII->get(TargetOpcode::COPY),
+            NewVR)
+        .addReg(*I);
+
+    for (auto *Exit : Exits)
+      BuildMI(*Exit, Exit->begin(), DebugLoc(), TII->get(TargetOpcode::COPY),
+              *I)
+          .addReg(NewVR);
+  }
+}

llvm/trunk/lib/Target/AArch64/AArch64MachineFunctionInfo.h

class AArch64FunctionInfo : public MachineFunctionInfo {
   /// \brief FrameIndex for start of varargs area for arguments passed in
   /// floating-point registers.
   int VarArgsFPRIndex;
 
   /// \brief Size of the varargs area for arguments passed in floating-point
   /// registers.
   unsigned VarArgsFPRSize;
 
+  /// True if this function has a subset of CSRs that is handled explicitly via
+  /// copies.
+  bool IsSplitCSR;
+
 public:
   AArch64FunctionInfo()
       : BytesInStackArgArea(0), ArgumentStackToRestore(0), HasStackFrame(false),
         NumLocalDynamicTLSAccesses(0), VarArgsStackIndex(0), VarArgsGPRIndex(0),
-        VarArgsGPRSize(0), VarArgsFPRIndex(0), VarArgsFPRSize(0) {}
+        VarArgsGPRSize(0), VarArgsFPRIndex(0), VarArgsFPRSize(0),
+        IsSplitCSR(false) {}
 
   explicit AArch64FunctionInfo(MachineFunction &MF)
       : BytesInStackArgArea(0), ArgumentStackToRestore(0), HasStackFrame(false),
         NumLocalDynamicTLSAccesses(0), VarArgsStackIndex(0), VarArgsGPRIndex(0),
-        VarArgsGPRSize(0), VarArgsFPRIndex(0), VarArgsFPRSize(0) {
+        VarArgsGPRSize(0), VarArgsFPRIndex(0), VarArgsFPRSize(0),
+        IsSplitCSR(false) {
     (void)MF;
   }
 
   unsigned getBytesInStackArgArea() const { return BytesInStackArgArea; }
   void setBytesInStackArgArea(unsigned bytes) { BytesInStackArgArea = bytes; }
 
   unsigned getArgumentStackToRestore() const { return ArgumentStackToRestore; }
   void setArgumentStackToRestore(unsigned bytes) {
     ArgumentStackToRestore = bytes;
   }
 
   bool hasStackFrame() const { return HasStackFrame; }
   void setHasStackFrame(bool s) { HasStackFrame = s; }
 
+  bool isSplitCSR() const { return IsSplitCSR; }
+  void setIsSplitCSR(bool s) { IsSplitCSR = s; }
+
   void setLocalStackSize(unsigned Size) { LocalStackSize = Size; }
   unsigned getLocalStackSize() const { return LocalStackSize; }
 
   void incNumLocalDynamicTLSAccesses() { ++NumLocalDynamicTLSAccesses; }
   unsigned getNumLocalDynamicTLSAccesses() const {
     return NumLocalDynamicTLSAccesses;
   }
 

llvm/trunk/lib/Target/AArch64/AArch64RegisterInfo.h

 
 public:
   AArch64RegisterInfo(const Triple &TT);
 
   bool isReservedReg(const MachineFunction &MF, unsigned Reg) const;
 
   /// Code Generation virtual methods...
   const MCPhysReg *getCalleeSavedRegs(const MachineFunction *MF) const override;
+  const MCPhysReg *
+  getCalleeSavedRegsViaCopy(const MachineFunction *MF) const override;
   const uint32_t *getCallPreservedMask(const MachineFunction &MF,
                                        CallingConv::ID) const override;
 
   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;

llvm/trunk/lib/Target/AArch64/AArch64RegisterInfo.cpp

 // This file contains the AArch64 implementation of the TargetRegisterInfo
 // class.
 //
 //===----------------------------------------------------------------------===//
 
 #include "AArch64RegisterInfo.h"
 #include "AArch64FrameLowering.h"
 #include "AArch64InstrInfo.h"
+#include "AArch64MachineFunctionInfo.h"
 #include "AArch64Subtarget.h"
 #include "MCTargetDesc/AArch64AddressingModes.h"
 #include "llvm/ADT/BitVector.h"
 #include "llvm/ADT/Triple.h"
 #include "llvm/CodeGen/MachineFrameInfo.h"
 #include "llvm/CodeGen/MachineInstrBuilder.h"
 #include "llvm/CodeGen/MachineRegisterInfo.h"
 #include "llvm/CodeGen/RegisterScavenging.h"
AArch64RegisterInfo::getCalleeSavedRegs(const MachineFunction *MF) const {
   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 CSR_AArch64_CXX_TLS_Darwin_SaveList;
+    return MF->getInfo<AArch64FunctionInfo>()->isSplitCSR() ?
+           CSR_AArch64_CXX_TLS_Darwin_PE_SaveList :
+           CSR_AArch64_CXX_TLS_Darwin_SaveList;
   else
     return CSR_AArch64_AAPCS_SaveList;
 }
 
+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;
+}
+
 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;

llvm/trunk/test/CodeGen/AArch64/cxx-tlscc.ll

 ; RUN: llc < %s -mtriple=aarch64-apple-ios | FileCheck %s
 ; RUN: llc < %s -mtriple=aarch64-apple-ios -enable-shrink-wrap=true | FileCheck --check-prefix=CHECK %s
 ; Shrink wrapping currently does not kick in because we have a TLS CALL
 ; in the entry block and it will clobber the link register.
 
 %struct.S = type { i8 }
 
 @sg = internal thread_local global %struct.S zeroinitializer, align 1
 @__dso_handle = external global i8
 @__tls_guard = internal thread_local unnamed_addr global i1 false
 
 declare %struct.S* @_ZN1SC1Ev(%struct.S* returned)
 declare %struct.S* @_ZN1SD1Ev(%struct.S* returned)
 declare i32 @_tlv_atexit(void (i8*)*, i8*, i8*)
 
-define cxx_fast_tlscc nonnull %struct.S* @_ZTW2sg() {
+define cxx_fast_tlscc nonnull %struct.S* @_ZTW2sg() nounwind {
   %.b.i = load i1, i1* @__tls_guard, align 1
   br i1 %.b.i, label %__tls_init.exit, label %init.i
 
 init.i:
   store i1 true, i1* @__tls_guard, align 1
   %call.i.i = tail call %struct.S* @_ZN1SC1Ev(%struct.S* nonnull @sg)
   %1 = tail call i32 @_tlv_atexit(void (i8*)* nonnull bitcast (%struct.S* (%struct.S*)* @_ZN1SD1Ev to void (i8*)*), i8* nonnull getelementptr inbounds (%struct.S, %struct.S* @sg, i64 0, i32 0), i8* nonnull @__dso_handle)
   br label %__tls_init.exit
 
 __tls_init.exit:
   ret %struct.S* @sg
 }
 
 ; CHECK-LABEL: _ZTW2sg
-; CHECK-DAG: stp d31, d30
-; CHECK-DAG: stp d29, d28
-; CHECK-DAG: stp d27, d26
-; CHECK-DAG: stp d25, d24
-; CHECK-DAG: stp d23, d22
-; CHECK-DAG: stp d21, d20
-; CHECK-DAG: stp d19, d18
-; CHECK-DAG: stp d17, d16
-; CHECK-DAG: stp d7, d6
-; CHECK-DAG: stp d5, d4
-; CHECK-DAG: stp d3, d2
-; CHECK-DAG: stp d1, d0
-; CHECK-DAG: stp x20, x19
-; CHECK-DAG: stp x14, x13
-; CHECK-DAG: stp x12, x11
-; CHECK-DAG: stp x10, x9
-; CHECK-DAG: stp x8, x7
-; CHECK-DAG: stp x6, x5
-; CHECK-DAG: stp x4, x3
-; CHECK-DAG: stp x2, x1
-; CHECK-DAG: stp x29, x30
+; CHECK-NOT: stp d31, d30
+; CHECK-NOT: stp d29, d28
+; CHECK-NOT: stp d27, d26
+; CHECK-NOT: stp d25, d24
+; CHECK-NOT: stp d23, d22
+; CHECK-NOT: stp d21, d20
+; CHECK-NOT: stp d19, d18
+; CHECK-NOT: stp d17, d16
+; CHECK-NOT: stp d7, d6
+; CHECK-NOT: stp d5, d4
+; CHECK-NOT: stp d3, d2
+; CHECK-NOT: stp d1, d0
+; CHECK-NOT: stp x20, x19
+; CHECK-NOT: stp x14, x13
+; CHECK-NOT: stp x12, x11
+; CHECK-NOT: stp x10, x9
+; CHECK-NOT: stp x8, x7
+; CHECK-NOT: stp x6, x5
+; CHECK-NOT: stp x4, x3
+; CHECK-NOT: stp x2, x1
 ; CHECK: blr
 ; CHECK: tbnz w{{.*}}, #0, [[BB_end:.?LBB0_[0-9]+]]
 ; CHECK: blr
 ; CHECK: tlv_atexit
 ; CHECK: [[BB_end]]:
 ; CHECK: blr
-; CHECK-DAG: ldp x2, x1
-; CHECK-DAG: ldp x4, x3
-; CHECK-DAG: ldp x6, x5
-; CHECK-DAG: ldp x8, x7
-; CHECK-DAG: ldp x10, x9
-; CHECK-DAG: ldp x12, x11
-; CHECK-DAG: ldp x14, x13
-; CHECK-DAG: ldp x20, x19
-; CHECK-DAG: ldp d1, d0
-; CHECK-DAG: ldp d3, d2
-; CHECK-DAG: ldp d5, d4
-; CHECK-DAG: ldp d7, d6
-; CHECK-DAG: ldp d17, d16
-; CHECK-DAG: ldp d19, d18
-; CHECK-DAG: ldp d21, d20
-; CHECK-DAG: ldp d23, d22
-; CHECK-DAG: ldp d25, d24
-; CHECK-DAG: ldp d27, d26
-; CHECK-DAG: ldp d29, d28
-; CHECK-DAG: ldp d31, d30
+; CHECK-NOT: ldp x2, x1
+; CHECK-NOT: ldp x4, x3
+; CHECK-NOT: ldp x6, x5
+; CHECK-NOT: ldp x8, x7
+; CHECK-NOT: ldp x10, x9
+; CHECK-NOT: ldp x12, x11
+; CHECK-NOT: ldp x14, x13
+; CHECK-NOT: ldp x20, x19
+; CHECK-NOT: ldp d1, d0
+; CHECK-NOT: ldp d3, d2
+; CHECK-NOT: ldp d5, d4
+; CHECK-NOT: ldp d7, d6
+; CHECK-NOT: ldp d17, d16
+; CHECK-NOT: ldp d19, d18
+; CHECK-NOT: ldp d21, d20
+; CHECK-NOT: ldp d23, d22
+; CHECK-NOT: ldp d25, d24
+; CHECK-NOT: ldp d27, d26
+; CHECK-NOT: ldp d29, d28
+; CHECK-NOT: ldp d31, d30