[CodeGen][InlineSpiller] Enable subreg spill-restore during RA

The register tuple spilling during RA is suboptimal when we restore
the entire tuple before the instructions that has only the subreg use
of the original tuple. This will lead to unwanted register pressure and
cases of unsuccessful allocation when deal with tests of large value types.

This patch allows targets to enable subreg reloads by constructing a
super class matching the subreg(s) involved in the MI in question and
eventually would improve the allocation by avoiding the unwanted subreg
reload components in the first place.

Author: cdevadas

llvm/include/llvm/CodeGen/LiveRangeEdit.h

@@ -154,8 +154,10 @@ class LiveRangeEdit : private MachineRegisterInfo::Delegate {
 
   ArrayRef<Register> regs() const { return ArrayRef(NewRegs).slice(FirstNew); }
 
-  /// createFrom - Create a new virtual register based on OldReg.
-  Register createFrom(Register OldReg);
+  /// createFrom - Create a new virtual register based on OldReg. If \p RC is
+  /// non-null, create the new virtual register from it instead. Subreg spills
+  /// will feed-in a subregclass derived from the regclass of OldReg.
+  Register createFrom(Register OldReg, const TargetRegisterClass *RC = nullptr);
 
   /// create - Create a new register with the same class and original slot as
   /// parent.

llvm/include/llvm/CodeGen/TargetInstrInfo.h

@@ -1215,6 +1215,7 @@ class LLVM_ABI TargetInstrInfo : public MCInstrInfo {
   virtual void loadRegFromStackSlot(
       MachineBasicBlock &MBB, MachineBasicBlock::iterator MI, Register DestReg,
       int FrameIndex, const TargetRegisterClass *RC, Register VReg,
+      unsigned SubReg = 0,
       MachineInstr::MIFlag Flags = MachineInstr::NoFlags) const {
     llvm_unreachable("Target didn't implement "
                      "TargetInstrInfo::loadRegFromStackSlot!");

llvm/include/llvm/CodeGen/TargetRegisterInfo.h

@@ -438,6 +438,8 @@ class LLVM_ABI TargetRegisterInfo : public MCRegisterInfo {
                                 LaneBitmask LaneMask,
                                 SmallVectorImpl<unsigned> &Indexes) const;
 
+  unsigned getSubRegIdxFromLaneMask(LaneBitmask LaneMask) const;
+
   /// The lane masks returned by getSubRegIndexLaneMask() above can only be
   /// used to determine if sub-registers overlap - they can't be used to
   /// determine if a set of sub-registers completely cover another
@@ -1225,6 +1227,11 @@ class LLVM_ABI TargetRegisterInfo : public MCRegisterInfo {
     return true;
   }
 
+  /// To enable the spill-restore of sub registers during RA. This would
+  /// eventually improve the register allocation for the functions that involve
+  /// subreg uses of register tuples.
+  virtual bool shouldEnableSubRegSpillRestore() const { return false; }
+
   /// When prioritizing live ranges in register allocation, if this hook returns
   /// true then the AllocationPriority of the register class will be treated as
   /// more important than whether the range is local to a basic block or global.

llvm/lib/CodeGen/InlineSpiller.cpp

@@ -217,7 +217,8 @@ class InlineSpiller : public Spiller {
   bool coalesceStackAccess(MachineInstr *MI, Register Reg);
   bool foldMemoryOperand(ArrayRef<std::pair<MachineInstr *, unsigned>>,
                          MachineInstr *LoadMI = nullptr);
-  void insertReload(Register VReg, SlotIndex, MachineBasicBlock::iterator MI);
+  void insertReload(Register VReg, unsigned SubReg, SlotIndex,
+                    MachineBasicBlock::iterator MI);
   void insertSpill(Register VReg, bool isKill, MachineBasicBlock::iterator MI);
 
   void spillAroundUses(Register Reg);
@@ -1112,14 +1113,14 @@ foldMemoryOperand(ArrayRef<std::pair<MachineInstr *, unsigned>> Ops,
   return true;
 }
 
-void InlineSpiller::insertReload(Register NewVReg,
+void InlineSpiller::insertReload(Register NewVReg, unsigned SubReg,
                                  SlotIndex Idx,
                                  MachineBasicBlock::iterator MI) {
   MachineBasicBlock &MBB = *MI->getParent();
 
   MachineInstrSpan MIS(MI, &MBB);
   TII.loadRegFromStackSlot(MBB, MI, NewVReg, StackSlot,
-                           MRI.getRegClass(NewVReg), Register());
+                           MRI.getRegClass(NewVReg), Register(), SubReg);
 
   LIS.InsertMachineInstrRangeInMaps(MIS.begin(), MI);
 
@@ -1248,18 +1249,50 @@ void InlineSpiller::spillAroundUses(Register Reg) {
 
     // Create a new virtual register for spill/fill.
     // FIXME: Infer regclass from instruction alone.
-    Register NewVReg = Edit->createFrom(Reg);
+
+    unsigned SubReg = 0;
+    LaneBitmask CoveringLanes = LaneBitmask::getNone();
+    // Identify the subreg use(s). Skip if the instruction defines the register.
+    // For copy bundles, get the covering lane masks.
+    if (TRI.shouldEnableSubRegSpillRestore() && !RI.Writes) {
+      for (auto [MI, OpIdx] : Ops) {
+        const MachineOperand &MO = MI->getOperand(OpIdx);
+        assert(MO.isReg() && MO.getReg() == Reg);
+        if (MO.isUse()) {
+          SubReg = MO.getSubReg();
+          CoveringLanes |= TRI.getSubRegIndexLaneMask(SubReg);
+        }
+      }
+    }
+
+    const TargetRegisterClass *OrigRC = MRI.getRegClass(Reg);
+    if (MI.isBundled() && CoveringLanes.any()) {
+      CoveringLanes = LaneBitmask(bit_ceil(CoveringLanes.getAsInteger()) - 1);
+      // Get the covering subreg index including the missing indices in the
+      // identified small range. Even if this is suboptimal, it is advantageous
+      // when the higher subreg components are not really involved in the bundle
+      // copy as we emit the subreg reload rather than the one for the entire
+      // tuple.
+      SubReg = TRI.getSubRegIdxFromLaneMask(CoveringLanes);
+    }
+
+    const TargetRegisterClass *NewRC =
+        SubReg ? TRI.getSubRegisterClass(OrigRC, SubReg) : nullptr;
+    Register NewVReg = Edit->createFrom(Reg, NewRC);
 
     if (RI.Reads)
-      insertReload(NewVReg, Idx, &MI);
+      insertReload(NewVReg, SubReg, Idx, &MI);
 
     // Rewrite instruction operands.
     bool hasLiveDef = false;
     for (const auto &OpPair : Ops) {
       MachineOperand &MO = OpPair.first->getOperand(OpPair.second);
       MO.setReg(NewVReg);
       if (MO.isUse()) {
-        if (!OpPair.first->isRegTiedToDefOperand(OpPair.second))
+        if (SubReg && !MI.isBundled())
+          MO.setSubReg(0);
+        if (!OpPair.first->isRegTiedToDefOperand(OpPair.second) ||
+            (SubReg && !MI.isBundled()))
           MO.setIsKill();
       } else {
         if (!MO.isDead())

llvm/lib/CodeGen/LiveRangeEdit.cpp

@@ -52,8 +52,10 @@ LiveInterval &LiveRangeEdit::createEmptyIntervalFrom(Register OldReg,
   return LI;
 }
 
-Register LiveRangeEdit::createFrom(Register OldReg) {
-  Register VReg = MRI.cloneVirtualRegister(OldReg);
+Register LiveRangeEdit::createFrom(Register OldReg,
+                                   const TargetRegisterClass *RC) {
+  Register VReg =
+      RC ? MRI.createVirtualRegister(RC) : MRI.cloneVirtualRegister(OldReg);
   if (VRM) {
     VRM->setIsSplitFromReg(VReg, VRM->getOriginal(OldReg));
   }

llvm/lib/CodeGen/TargetRegisterInfo.cpp

@@ -546,6 +546,16 @@ TargetRegisterInfo::getRegSizeInBits(Register Reg,
   return getRegSizeInBits(*RC);
 }
 
+unsigned
+TargetRegisterInfo::getSubRegIdxFromLaneMask(LaneBitmask LaneMask) const {
+  for (unsigned Idx = 1, E = getNumSubRegIndices(); Idx < E; ++Idx) {
+    if (getSubRegIndexLaneMask(Idx) == LaneMask)
+      return Idx;
+  }
+
+  return 0 /*NoSubRegister*/;
+}
+
 bool TargetRegisterInfo::getCoveringSubRegIndexes(
     const TargetRegisterClass *RC, LaneBitmask LaneMask,
     SmallVectorImpl<unsigned> &NeededIndexes) const {

llvm/lib/Target/AArch64/AArch64InstrInfo.cpp

@@ -6035,7 +6035,7 @@ void AArch64InstrInfo::loadRegFromStackSlot(MachineBasicBlock &MBB,
                                             MachineBasicBlock::iterator MBBI,
                                             Register DestReg, int FI,
                                             const TargetRegisterClass *RC,
-                                            Register VReg,
+                                            Register VReg, unsigned SubReg,
                                             MachineInstr::MIFlag Flags) const {
   MachineFunction &MF = *MBB.getParent();
   MachineFrameInfo &MFI = MF.getFrameInfo();

llvm/lib/Target/AArch64/AArch64InstrInfo.h

@@ -368,7 +368,7 @@ class AArch64InstrInfo final : public AArch64GenInstrInfo {
   void loadRegFromStackSlot(
       MachineBasicBlock &MBB, MachineBasicBlock::iterator MBBI,
       Register DestReg, int FrameIndex, const TargetRegisterClass *RC,
-      Register VReg,
+      Register VReg, unsigned SubReg = 0,
       MachineInstr::MIFlag Flags = MachineInstr::NoFlags) const override;
 
   // This tells target independent code that it is okay to pass instructions

llvm/lib/Target/AMDGPU/SIInstrInfo.cpp

@@ -1907,14 +1907,30 @@ void SIInstrInfo::loadRegFromStackSlot(MachineBasicBlock &MBB,
                                        MachineBasicBlock::iterator MI,
                                        Register DestReg, int FrameIndex,
                                        const TargetRegisterClass *RC,
-                                       Register VReg,
+                                       Register VReg, unsigned SubReg,
                                        MachineInstr::MIFlag Flags) const {
   MachineFunction *MF = MBB.getParent();
   SIMachineFunctionInfo *MFI = MF->getInfo<SIMachineFunctionInfo>();
   MachineFrameInfo &FrameInfo = MF->getFrameInfo();
   const DebugLoc &DL = MBB.findDebugLoc(MI);
   unsigned SpillSize = RI.getSpillSize(*RC);
 
+  assert(SubReg != AMDGPU::lo16 && SubReg != AMDGPU::hi16 &&
+         "unhandled 16-bit subregister spilling");
+
+  // For subreg reload, identify the start offset.
+  unsigned SubRegIdx =
+      SubReg
+          ? llvm::countr_zero(RI.getSubRegIndexLaneMask(SubReg).getAsInteger())
+          : 0;
+  // Each subreg consists of two bits in the RegMask. The SubRegIdx should be
+  // either zero or an even number. This assert is to ensure we will not have
+  // any 16-bit subreg access at this point.
+  assert(SubRegIdx % 2 == 0 && "expected an even number for the subreg index");
+
+  // Now get the actual subreg index.
+  SubRegIdx /= 2;
+
   MachinePointerInfo PtrInfo
     = MachinePointerInfo::getFixedStack(*MF, FrameIndex);
 
@@ -1939,19 +1955,23 @@ void SIInstrInfo::loadRegFromStackSlot(MachineBasicBlock &MBB,
     if (RI.spillSGPRToVGPR())
       FrameInfo.setStackID(FrameIndex, TargetStackID::SGPRSpill);
     BuildMI(MBB, MI, DL, OpDesc, DestReg)
-      .addFrameIndex(FrameIndex) // addr
-      .addMemOperand(MMO)
-      .addReg(MFI->getStackPtrOffsetReg(), RegState::Implicit);
+        .addFrameIndex(FrameIndex) // addr
+        .addImm(SubRegIdx)         // offset
+        .addMemOperand(MMO)
+        .addReg(MFI->getStackPtrOffsetReg(), RegState::Implicit);
 
     return;
   }
 
+  // Convert the subreg index to stack offset.
+  SubRegIdx *= 4;
+
   unsigned Opcode = getVectorRegSpillRestoreOpcode(VReg ? VReg : DestReg, RC,
                                                    SpillSize, *MFI);
   BuildMI(MBB, MI, DL, get(Opcode), DestReg)
       .addFrameIndex(FrameIndex)           // vaddr
       .addReg(MFI->getStackPtrOffsetReg()) // scratch_offset
-      .addImm(0)                           // offset
+      .addImm(SubRegIdx)                   // offset
       .addMemOperand(MMO);
 }
 

llvm/lib/Target/AMDGPU/SIInstrInfo.h

@@ -331,6 +331,7 @@ class SIInstrInfo final : public AMDGPUGenInstrInfo {
   void loadRegFromStackSlot(
       MachineBasicBlock &MBB, MachineBasicBlock::iterator MI, Register DestReg,
       int FrameIndex, const TargetRegisterClass *RC, Register VReg,
+      unsigned SubReg = 0,
       MachineInstr::MIFlag Flags = MachineInstr::NoFlags) const override;
 
   bool expandPostRAPseudo(MachineInstr &MI) const override;