[AArch64] Add an option to reserve physical registers from RA

[Carrot]

Summary

This patch adds an option --reserve-regs-for-regalloc, so we can reserve a list of physical registers. These registers will not be used by register allocator, but can still be used as ABI requests such as passing arguments to function call.

Its main purpose is simulating high register pressure by reserving many physical registers. So it will be much easier to test and debug register allocation changes.

[Carrot]

ping

[efriedma]

Do we actually need getReservedRegsForRA() to be something different from getReservedRegs()? Having getReservedRegs() return bogus values won't actually hurt anything, I think. The issue with the other patch was just that we don't want to expose feature flags that we know don't actually work.

[Carrot]

It's not necessary strictly speaking. The only problem of getReservedRegs is it is used by isAnyArgRegReserved. Testing this new kind of reserved register in isAnyArgRegReserved doesn't make sense. A possible solution is moving current content of getReservedRegs to a new function getStrictlyReservedRegs, and call this function from getReservedRegs and isAnyArgRegReserved. I will try it.

[efriedma]

llvm/lib/Target/AArch64/AArch64Subtarget.cpp

65

Maybe note in the description that this is intended for testing the register allocator.

303

Please use StringMap (see https://llvm.org/docs/ProgrammersManual.html#llvm-adt-stringmap-h )

[Carrot]

llvm/lib/Target/AArch64/AArch64Subtarget.cpp

303

It looks StringSet is more appropriate.

[efriedma]

LGTM

llvm/lib/Target/AArch64/AArch64RegisterInfo.h

public:
   AArch64RegisterInfo(const Triple &TT);
 
   // FIXME: This should be tablegen'd like getDwarfRegNum is
   int getSEHRegNum(unsigned i) const {
     return getEncodingValue(i);
   }
 
   bool isReservedReg(const MachineFunction &MF, MCRegister Reg) const;
+  bool isStrictlyReservedReg(const MachineFunction &MF, MCRegister Reg) const;
   bool isAnyArgRegReserved(const MachineFunction &MF) const;
   void emitReservedArgRegCallError(const MachineFunction &MF) const;
 
   void UpdateCustomCalleeSavedRegs(MachineFunction &MF) const;
   void UpdateCustomCallPreservedMask(MachineFunction &MF,
                                      const uint32_t **Mask) const;
 
   /// Code Generation virtual methods...
public:
   /// Should return NULL in the case that the calling convention does not have
   /// this property
   const uint32_t *getThisReturnPreservedMask(const MachineFunction &MF,
                                              CallingConv::ID) const;
 
   /// Stack probing calls preserve different CSRs to the normal CC.
   const uint32_t *getWindowsStackProbePreservedMask() const;
 
+  BitVector getStrictlyReservedRegs(const MachineFunction &MF) const;
   BitVector getReservedRegs(const MachineFunction &MF) const override;
   bool isAsmClobberable(const MachineFunction &MF,
                        MCRegister PhysReg) const override;
   const TargetRegisterClass *
   getPointerRegClass(const MachineFunction &MF,
                      unsigned Kind = 0) const override;
   const TargetRegisterClass *
   getCrossCopyRegClass(const TargetRegisterClass *RC) const override;

llvm/lib/Target/AArch64/AArch64RegisterInfo.cpp

AArch64RegisterInfo::getThisReturnPreservedMask(const MachineFunction &MF,
   return CSR_AArch64_AAPCS_ThisReturn_RegMask;
 }
 
 const uint32_t *AArch64RegisterInfo::getWindowsStackProbePreservedMask() const {
   return CSR_AArch64_StackProbe_Windows_RegMask;
 }
 
 BitVector
-AArch64RegisterInfo::getReservedRegs(const MachineFunction &MF) const {
+AArch64RegisterInfo::getStrictlyReservedRegs(const MachineFunction &MF) const {
   const AArch64FrameLowering *TFI = getFrameLowering(MF);
 
   // FIXME: avoid re-calculating this every time.
   BitVector Reserved(getNumRegs());
   markSuperRegs(Reserved, AArch64::WSP);
   markSuperRegs(Reserved, AArch64::WZR);
 
   if (TFI->hasFP(MF) || TT.isOSDarwin())
for (MCSubRegIterator SubReg(AArch64::ZA, this, /*self=*/true);
          SubReg.isValid(); ++SubReg)
       Reserved.set(*SubReg);
   }
 
   assert(checkAllSuperRegsMarked(Reserved));
   return Reserved;
 }
 
+BitVector
+AArch64RegisterInfo::getReservedRegs(const MachineFunction &MF) const {
+  BitVector Reserved = getStrictlyReservedRegs(MF);
+
+  for (size_t i = 0; i < AArch64::GPR32commonRegClass.getNumRegs(); ++i) {
+    if (MF.getSubtarget<AArch64Subtarget>().isXRegisterReservedForRA(i))
+      markSuperRegs(Reserved, AArch64::GPR32commonRegClass.getRegister(i));
+  }
+
+  assert(checkAllSuperRegsMarked(Reserved));
+  return Reserved;
+}
+
 bool AArch64RegisterInfo::isReservedReg(const MachineFunction &MF,
                                         MCRegister Reg) const {
   return getReservedRegs(MF)[Reg];
 }
 
+bool AArch64RegisterInfo::isStrictlyReservedReg(const MachineFunction &MF,
+                                                MCRegister Reg) const {
+  return getStrictlyReservedRegs(MF)[Reg];
+}
+
 bool AArch64RegisterInfo::isAnyArgRegReserved(const MachineFunction &MF) const {
   return llvm::any_of(*AArch64::GPR64argRegClass.MC, [this, &MF](MCPhysReg r) {
-    return isReservedReg(MF, r);
+    return isStrictlyReservedReg(MF, r);
   });
 }
 
 void AArch64RegisterInfo::emitReservedArgRegCallError(
     const MachineFunction &MF) const {
   const Function &F = MF.getFunction();
   F.getContext().diagnose(DiagnosticInfoUnsupported{F, ("AArch64 doesn't support"
     " function calls if any of the argument registers is reserved.")});

llvm/lib/Target/AArch64/AArch64Subtarget.h

#include "AArch64GenSubtargetInfo.inc"
   unsigned PrefFunctionLogAlignment = 0;
   unsigned PrefLoopLogAlignment = 0;
   unsigned MaxBytesForLoopAlignment = 0;
   unsigned MaxJumpTableSize = 0;
 
   // ReserveXRegister[i] - X#i is not available as a general purpose register.
   BitVector ReserveXRegister;
 
+  // ReserveXRegisterForRA[i] - X#i is not available for register allocator.
+  BitVector ReserveXRegisterForRA;
+
   // CustomCallUsedXRegister[i] - X#i call saved.
   BitVector CustomCallSavedXRegs;
 
   bool IsLittle;
 
   unsigned MinSVEVectorSizeInBits;
   unsigned MaxSVEVectorSizeInBits;
   unsigned VScaleForTuning = 2;
#include "AArch64GenSubtargetInfo.inc"
 
   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 isXRegisterReservedForRA(size_t i) const { return ReserveXRegisterForRA[i]; }
+  unsigned getNumXRegisterReserved() const {
+    BitVector AllReservedX(AArch64::GPR64commonRegClass.getNumRegs());
+    AllReservedX |= ReserveXRegister;
+    AllReservedX |= ReserveXRegisterForRA;
+    return AllReservedX.count();
+  }
   bool isXRegCustomCalleeSaved(size_t i) const {
     return CustomCallSavedXRegs[i];
   }
   bool hasCustomCallingConv() const { return CustomCallSavedXRegs.any(); }
 
   /// Return true if the CPU supports any kind of instruction fusion.
   bool hasFusion() const {
     return hasArithmeticBccFusion() || hasArithmeticCbzFusion() ||

llvm/lib/Target/AArch64/AArch64Subtarget.cpp

 
 static cl::opt<bool> UseAA("aarch64-use-aa", cl::init(true),
                            cl::desc("Enable the use of AA during codegen."));
 
 static cl::opt<unsigned> OverrideVectorInsertExtractBaseCost(
     "aarch64-insert-extract-base-cost",
     cl::desc("Base cost of vector insert/extract element"), cl::Hidden);
 
+// Reserve a list of X# registers, so they are unavailable for register
+// allocator, but can still be used as ABI requests, such as passing arguments
+// to function call.
+static cl::list<std::string>
+ReservedRegsForRA("reserve-regs-for-regalloc", cl::desc("Reserve physical "
+                  "registers, so they can't be used by register allocator. "
+                  "Should only be used for testing register allocator."),
+                  cl::CommaSeparated, cl::Hidden);
+
 unsigned AArch64Subtarget::getVectorInsertExtractBaseCost() const {
   if (OverrideVectorInsertExtractBaseCost.getNumOccurrences() > 0)
     return OverrideVectorInsertExtractBaseCost;
   return VectorInsertExtractBaseCost;
 }
 
 AArch64Subtarget &AArch64Subtarget::initializeSubtargetDependencies(
     StringRef FS, StringRef CPUString, StringRef TuneCPUString) {

 AArch64Subtarget::AArch64Subtarget(const Triple &TT, const std::string &CPU,
                                    const std::string &TuneCPU,
                                    const std::string &FS,
                                    const TargetMachine &TM, bool LittleEndian,
                                    unsigned MinSVEVectorSizeInBitsOverride,
                                    unsigned MaxSVEVectorSizeInBitsOverride)
     : AArch64GenSubtargetInfo(TT, CPU, TuneCPU, FS),
       ReserveXRegister(AArch64::GPR64commonRegClass.getNumRegs()),
+      ReserveXRegisterForRA(AArch64::GPR64commonRegClass.getNumRegs()),
       CustomCallSavedXRegs(AArch64::GPR64commonRegClass.getNumRegs()),
       IsLittle(LittleEndian),
       MinSVEVectorSizeInBits(MinSVEVectorSizeInBitsOverride),
       MaxSVEVectorSizeInBits(MaxSVEVectorSizeInBitsOverride), TargetTriple(TT),
       InstrInfo(initializeSubtargetDependencies(FS, CPU, TuneCPU)),
       TLInfo(TM, *this) {
   if (AArch64::isX18ReservedByDefault(TT))
     ReserveXRegister.set(18);
 
   CallLoweringInfo.reset(new AArch64CallLowering(*getTargetLowering()));
   InlineAsmLoweringInfo.reset(new InlineAsmLowering(getTargetLowering()));
   Legalizer.reset(new AArch64LegalizerInfo(*this));
 
   auto *RBI = new AArch64RegisterBankInfo(*getRegisterInfo());
 
   // FIXME: At this point, we can't rely on Subtarget having RBI.
   // It's awkward to mix passing RBI and the Subtarget; should we pass
   // TII/TRI as well?
   InstSelector.reset(createAArch64InstructionSelector(
       *static_cast<const AArch64TargetMachine *>(&TM), *this, *RBI));
 
   RegBankInfo.reset(RBI);
+
+  auto TRI = getRegisterInfo();
+  StringSet<> ReservedRegNames;
+  ReservedRegNames.insert(ReservedRegsForRA.begin(), ReservedRegsForRA.end());
+  for (unsigned i = 0; i < 31; ++i) {
+    if (ReservedRegNames.count(TRI->getName(AArch64::X0 + i)))
+      ReserveXRegisterForRA.set(i);
+  }
 }
 
 const CallLowering *AArch64Subtarget::getCallLowering() const {
   return CallLoweringInfo.get();
 }
 
 const InlineAsmLowering *AArch64Subtarget::getInlineAsmLowering() const {
   return InlineAsmLoweringInfo.get();

llvm/test/CodeGen/AArch64/arm64-platform-reg.ll

 ; RUN: llc -mtriple=arm64-apple-ios -mattr=+reserve-x18 -o - %s | FileCheck %s --check-prefix=CHECK-RESERVE --check-prefix=CHECK-RESERVE-X18
 ; RUN: llc -mtriple=arm64-freebsd-gnu -mattr=+reserve-x18 -o - %s | FileCheck %s --check-prefix=CHECK-RESERVE --check-prefix=CHECK-RESERVE-X18
 ; RUN: llc -mtriple=arm64-linux-gnu -o - %s | FileCheck %s
 ; RUN: llc -mtriple=aarch64-linux-android -o - %s | FileCheck %s --check-prefix=CHECK-RESERVE --check-prefix=CHECK-RESERVE-X18
 ; RUN: llc -mtriple=aarch64-fuchsia -o - %s | FileCheck %s --check-prefix=CHECK-RESERVE --check-prefix=CHECK-RESERVE-X18
 ; RUN: llc -mtriple=aarch64-windows -o - %s | FileCheck %s --check-prefix=CHECK-RESERVE --check-prefix=CHECK-RESERVE-X18
 
 ; Test reserve-x# options individually.
 ; RUN: llc -mtriple=arm64-linux-gnu -mattr=+reserve-x1 -o - %s | FileCheck %s --check-prefixes=CHECK-RESERVE,CHECK-RESERVE-X1
 ; RUN: llc -mtriple=arm64-linux-gnu -mattr=+reserve-x2 -o - %s | FileCheck %s --check-prefixes=CHECK-RESERVE,CHECK-RESERVE-X2
 ; RUN: llc -mtriple=arm64-linux-gnu -mattr=+reserve-x3 -o - %s | FileCheck %s --check-prefixes=CHECK-RESERVE,CHECK-RESERVE-X3
 ; RUN: llc -mtriple=arm64-linux-gnu -mattr=+reserve-x4 -o - %s | FileCheck %s --check-prefixes=CHECK-RESERVE,CHECK-RESERVE-X4
 ; RUN: llc -mtriple=arm64-linux-gnu -mattr=+reserve-x5 -o - %s | FileCheck %s --check-prefixes=CHECK-RESERVE,CHECK-RESERVE-X5
 ; RUN: llc -mtriple=arm64-linux-gnu -mattr=+reserve-x6 -o - %s | FileCheck %s --check-prefixes=CHECK-RESERVE,CHECK-RESERVE-X6
 ; RUN: llc -mtriple=arm64-linux-gnu -mattr=+reserve-x7 -o - %s | FileCheck %s --check-prefixes=CHECK-RESERVE,CHECK-RESERVE-X7
+; RUN: llc -mtriple=arm64-linux-gnu -reserve-regs-for-regalloc=X8 -o - %s | FileCheck %s --check-prefixes=CHECK-RESERVE,CHECK-RESERVE-X8
 ; RUN: llc -mtriple=arm64-linux-gnu -mattr=+reserve-x9 -o - %s | FileCheck %s --check-prefixes=CHECK-RESERVE,CHECK-RESERVE-X9
 ; RUN: llc -mtriple=arm64-linux-gnu -mattr=+reserve-x10 -o - %s | FileCheck %s --check-prefixes=CHECK-RESERVE,CHECK-RESERVE-X10
 ; RUN: llc -mtriple=arm64-linux-gnu -mattr=+reserve-x11 -o - %s | FileCheck %s --check-prefixes=CHECK-RESERVE,CHECK-RESERVE-X11
 ; RUN: llc -mtriple=arm64-linux-gnu -mattr=+reserve-x12 -o - %s | FileCheck %s --check-prefixes=CHECK-RESERVE,CHECK-RESERVE-X12
 ; RUN: llc -mtriple=arm64-linux-gnu -mattr=+reserve-x13 -o - %s | FileCheck %s --check-prefixes=CHECK-RESERVE,CHECK-RESERVE-X13
 ; RUN: llc -mtriple=arm64-linux-gnu -mattr=+reserve-x14 -o - %s | FileCheck %s --check-prefixes=CHECK-RESERVE,CHECK-RESERVE-X14
 ; RUN: llc -mtriple=arm64-linux-gnu -mattr=+reserve-x15 -o - %s | FileCheck %s --check-prefixes=CHECK-RESERVE,CHECK-RESERVE-X15
+; RUN: llc -mtriple=arm64-linux-gnu -reserve-regs-for-regalloc=X16 -o - %s | FileCheck %s --check-prefixes=CHECK-RESERVE,CHECK-RESERVE-X16
+; RUN: llc -mtriple=arm64-linux-gnu -reserve-regs-for-regalloc=X17 -o - %s | FileCheck %s --check-prefixes=CHECK-RESERVE,CHECK-RESERVE-X17
+; RUN: llc -mtriple=arm64-linux-gnu -reserve-regs-for-regalloc=X19 -o - %s | FileCheck %s --check-prefixes=CHECK-RESERVE,CHECK-RESERVE-X19
 ; RUN: llc -mtriple=arm64-linux-gnu -mattr=+reserve-x20 -o - %s | FileCheck %s --check-prefixes=CHECK-RESERVE,CHECK-RESERVE-X20
 ; RUN: llc -mtriple=arm64-linux-gnu -mattr=+reserve-x21 -o - %s | FileCheck %s --check-prefixes=CHECK-RESERVE,CHECK-RESERVE-X21
 ; RUN: llc -mtriple=arm64-linux-gnu -mattr=+reserve-x22 -o - %s | FileCheck %s --check-prefixes=CHECK-RESERVE,CHECK-RESERVE-X22
 ; RUN: llc -mtriple=arm64-linux-gnu -mattr=+reserve-x23 -o - %s | FileCheck %s --check-prefixes=CHECK-RESERVE,CHECK-RESERVE-X23
 ; RUN: llc -mtriple=arm64-linux-gnu -mattr=+reserve-x24 -o - %s | FileCheck %s --check-prefixes=CHECK-RESERVE,CHECK-RESERVE-X24
 ; RUN: llc -mtriple=arm64-linux-gnu -mattr=+reserve-x25 -o - %s | FileCheck %s --check-prefixes=CHECK-RESERVE,CHECK-RESERVE-X25
 ; RUN: llc -mtriple=arm64-linux-gnu -mattr=+reserve-x26 -o - %s | FileCheck %s --check-prefixes=CHECK-RESERVE,CHECK-RESERVE-X26
 ; RUN: llc -mtriple=arm64-linux-gnu -mattr=+reserve-x27 -o - %s | FileCheck %s --check-prefixes=CHECK-RESERVE,CHECK-RESERVE-X27

 ; RUN: -mattr=+reserve-x22 \
 ; RUN: -mattr=+reserve-x23 \
 ; RUN: -mattr=+reserve-x24 \
 ; RUN: -mattr=+reserve-x25 \
 ; RUN: -mattr=+reserve-x26 \
 ; RUN: -mattr=+reserve-x27 \
 ; RUN: -mattr=+reserve-x28 \
 ; RUN: -mattr=+reserve-x30 \
+; RUN: -reserve-regs-for-regalloc=X8,X16,X17,X19 \
 ; RUN: -o - %s | FileCheck %s \
 ; RUN: --check-prefix=CHECK-RESERVE \
 ; RUN: --check-prefix=CHECK-RESERVE-X1 \
 ; RUN: --check-prefix=CHECK-RESERVE-X2 \
 ; RUN: --check-prefix=CHECK-RESERVE-X3 \
 ; RUN: --check-prefix=CHECK-RESERVE-X4 \
 ; RUN: --check-prefix=CHECK-RESERVE-X5 \
 ; RUN: --check-prefix=CHECK-RESERVE-X6 \
 ; RUN: --check-prefix=CHECK-RESERVE-X7 \
+; RUN: --check-prefix=CHECK-RESERVE-X8 \
 ; RUN: --check-prefix=CHECK-RESERVE-X9 \
 ; RUN: --check-prefix=CHECK-RESERVE-X10 \
 ; RUN: --check-prefix=CHECK-RESERVE-X11 \
 ; RUN: --check-prefix=CHECK-RESERVE-X12 \
 ; RUN: --check-prefix=CHECK-RESERVE-X13 \
 ; RUN: --check-prefix=CHECK-RESERVE-X14 \
 ; RUN: --check-prefix=CHECK-RESERVE-X15 \
+; RUN: --check-prefix=CHECK-RESERVE-X16 \
+; RUN: --check-prefix=CHECK-RESERVE-X17 \
 ; RUN: --check-prefix=CHECK-RESERVE-X18 \
+; RUN: --check-prefix=CHECK-RESERVE-X19 \
 ; RUN: --check-prefix=CHECK-RESERVE-X20 \
 ; RUN: --check-prefix=CHECK-RESERVE-X21 \
 ; RUN: --check-prefix=CHECK-RESERVE-X22 \
 ; RUN: --check-prefix=CHECK-RESERVE-X23 \
 ; RUN: --check-prefix=CHECK-RESERVE-X24 \
 ; RUN: --check-prefix=CHECK-RESERVE-X25 \
 ; RUN: --check-prefix=CHECK-RESERVE-X26 \
 ; RUN: --check-prefix=CHECK-RESERVE-X27 \

 define void @keep_live() {
   %val = load volatile [30 x i64], [30 x i64]* @var
   store volatile [30 x i64] %val, [30 x i64]* @var
 
 ; CHECK: ldr x18
 ; CHECK: str x18
 
 ; CHECK-RESERVE-NOT: ldr fp
+; CHECK-RESERVE-X8-NOT: adrp x8
+; CHECK-RESERVE-X8-NOT: ldr x8
 ; CHECK-RESERVE-X1-NOT: ldr x1,
 ; CHECK-RESERVE-X2-NOT: ldr x2,
 ; CHECK-RESERVE-X3-NOT: ldr x3,
 ; CHECK-RESERVE-X4-NOT: ldr x4,
 ; CHECK-RESERVE-X5-NOT: ldr x5,
 ; CHECK-RESERVE-X6-NOT: ldr x6,
 ; CHECK-RESERVE-X7-NOT: ldr x7,
 ; CHECK-RESERVE-X9-NOT: ldr x9,
 ; CHECK-RESERVE-X10-NOT: ldr x10,
 ; CHECK-RESERVE-X11-NOT: ldr x11,
 ; CHECK-RESERVE-X12-NOT: ldr x12,
 ; CHECK-RESERVE-X13-NOT: ldr x13,
 ; CHECK-RESERVE-X14-NOT: ldr x14,
 ; CHECK-RESERVE-X15-NOT: ldr x15,
+; CHECK-RESERVE-X16-NOT: ldr x16
+; CHECK-RESERVE-X17-NOT: ldr x17
 ; CHECK-RESERVE-X18-NOT: ldr x18
+; CHECK-RESERVE-X19-NOT: ldr x19
 ; CHECK-RESERVE-X20-NOT: ldr x20
 ; CHECK-RESERVE-X21-NOT: ldr x21
 ; CHECK-RESERVE-X22-NOT: ldr x22
 ; CHECK-RESERVE-X23-NOT: ldr x23
 ; CHECK-RESERVE-X24-NOT: ldr x24
 ; CHECK-RESERVE-X25-NOT: ldr x25
 ; CHECK-RESERVE-X26-NOT: ldr x26
 ; CHECK-RESERVE-X27-NOT: ldr x27

 ; CHECK-RESERVE-X7-NOT: ldr x7,
 ; CHECK-RESERVE-X9-NOT: ldr x9,
 ; CHECK-RESERVE-X10-NOT: ldr x10,
 ; CHECK-RESERVE-X11-NOT: ldr x11,
 ; CHECK-RESERVE-X12-NOT: ldr x12,
 ; CHECK-RESERVE-X13-NOT: ldr x13,
 ; CHECK-RESERVE-X14-NOT: ldr x14,
 ; CHECK-RESERVE-X15-NOT: ldr x15,
+; CHECK-RESERVE-X16-NOT: ldr x16
+; CHECK-RESERVE-X17-NOT: ldr x17
 ; CHECK-RESERVE-X18-NOT: ldr x18
+; CHECK-RESERVE-X19-NOT: ldr x19
 ; CHECK-RESERVE-X20-NOT: ldr x20
 ; CHECK-RESERVE-X21-NOT: ldr x21
 ; CHECK-RESERVE-X22-NOT: ldr x22
 ; CHECK-RESERVE-X23-NOT: ldr x23
 ; CHECK-RESERVE-X24-NOT: ldr x24
 ; CHECK-RESERVE-X25-NOT: ldr x25
 ; CHECK-RESERVE-X26-NOT: ldr x26
 ; CHECK-RESERVE-X27-NOT: ldr x27
 ; CHECK-RESERVE-X28-NOT: ldr x28
 ; CHECK-RESERVE-X30-NOT: ldr x30
 ; CHECK-RESERVE: ret
   ret void
 }