[AArch64] Fold some refilled/spilled subreg COPYs

[gberry]

Summary

Extend AArch64 foldMemoryOperandImpl() to handle folding spills of
subreg COPYs with read-undef defs like:

%vreg0:sub_32<def,read-undef> = COPY %WZR; GPR64:%vreg0

by widening the spilled physical source reg and generating:

STRXui %XZR <fi#0>

as well as folding refills of similar COPYs like:

%vreg0:sub_32<def,read-undef> = COPY %vreg1; GPR64:%vreg0, GPR32:%vreg1

by generating:

%vreg0:sub_32<def,read-undef> = LDRWui <fi#0>

[gberry]

This is a more general version of D27002

[MatzeB]

Generally looks fine, but I have some remarks.

lib/Target/AArch64/AArch64InstrInfo.cpp

2629

Either use the term fill or reload (I personally prefer the latter) but not refill. Should probably change that in the comment as well.

I wonder if the code would become slightly cleaner if you put the spill/refill check into a single outer if and move the other ifs inside of that accordingly (you have to split up/duplicate the DstMO.getSubReg() == 0 && SrcMO.getSubReg() == 0 check but that should be fine).

2642–2643

Maybe you can delay the computation of the register classes until one of the if conditions is fulfilled (because getMinimalPhysRegClass() isn't that fast).

2646

I think DestRC.getSize() == SrcRC.getSize() must be true for a copy without subreg indexes, so it can be an assert.

2664–2665

TargetRegisterInfo::isPhysicalRegister(SrcReg) should imply SrcMO.getSubReg() == 0, so the latter can be an assert.

2666

This condition be combined with the preceding if.

2670–2675

The listing of GPR64commonRegClass, GPR64spRegClass, ... seems a bit arbitrary (if someone adds a new class he will probably miss updating this piece of code). Also I can see this failing if the copy is actually from SP and you change GPR64spRegClass to GPR64.

Maybe what we really want here is something like

if (GPR64RegClass.contains(SrcReg))
  SpillRC = GPR64RegClass;
else if (GPR32RegClass.contains(SrcReg))
  SpillRc = GPR32RegClass;
else
  SpillRc = &DstRC;

?

2686

Use fill or reload.

2697–2708

I wish we had a function that given a regclass + subregindex would return the class that results from applying the subregindex to all registers in regclass. Surprisingly I couldn't find such a function, may be worth investigating at some point what it takes to create that. (Of course this is independent of this commit, we can leave this for now I think).

2714

Use a reference?

[gberry]

lib/Target/AArch64/AArch64InstrInfo.cpp

2629

I don't think this refactoring makes it much easier to read, especially given the fixups that have to be done to the load in the fill, read-undef case, but let me know if you disagree.

2670–2675

I believe I have addressed this in my update.

2697–2708

Yeah, I also looked for just such a function, but this code is slightly more general in that it will handle cross-bank COPYs as well.

[gberry]

New version addressing Matthias' comments.

[MatzeB]

LGTM

lib/Target/AArch64/AArch64InstrInfo.h

void storeRegToStackSlot(MachineBasicBlock &MBB,
                            const TargetRegisterClass *RC,
                            const TargetRegisterInfo *TRI) const override;
 
   void loadRegFromStackSlot(MachineBasicBlock &MBB,
                             MachineBasicBlock::iterator MBBI, unsigned DestReg,
                             int FrameIndex, const TargetRegisterClass *RC,
                             const TargetRegisterInfo *TRI) const override;
 
+  // This tells target independent code that it is okay to pass instructions
+  // with subreg operands to foldMemoryOperandImpl.
+  bool isSubregFoldable() const override { return true; }
+
   using TargetInstrInfo::foldMemoryOperandImpl;
   MachineInstr *
   foldMemoryOperandImpl(MachineFunction &MF, MachineInstr &MI,
                         ArrayRef<unsigned> Ops,
                         MachineBasicBlock::iterator InsertPt, int FrameIndex,
                         LiveIntervals *LIS = nullptr) const override;
 
   /// \returns true if a branch from an instruction with opcode \p BranchOpc

lib/Target/AArch64/AArch64InstrInfo.cpp

MachineInstr *AArch64InstrInfo::foldMemoryOperandImpl(
   // be eliminated by RegisterCoalescer. However, that may not be possible, and
   // %vreg0 may even spill. We can't spill %SP, and since it is in the GPR64all
   // register class, TargetInstrInfo::foldMemoryOperand() is going to try.
   //
   // To prevent that, we are going to constrain the %vreg0 register class here.
   //
   // <rdar://problem/11522048>
   //
-  if (MI.isCopy()) {
+  if (MI.isFullCopy()) {
     unsigned DstReg = MI.getOperand(0).getReg();
     unsigned SrcReg = MI.getOperand(1).getReg();
     if (SrcReg == AArch64::SP &&
         TargetRegisterInfo::isVirtualRegister(DstReg)) {
       MF.getRegInfo().constrainRegClass(DstReg, &AArch64::GPR64RegClass);
       return nullptr;
     }
     if (DstReg == AArch64::SP &&
MachineInstr *AArch64InstrInfo::foldMemoryOperandImpl(
   //
   //   LDRDui %vreg0, fi<#0>
   //
   // instead of
   //
   //   LDRXui %vregTemp, fi<#0>
   //   %vreg0 = FMOV %vregTemp
   //
-  if (MI.isFullCopy() && Ops.size() == 1 &&
+  if (MI.isCopy() && Ops.size() == 1 &&
       // Make sure we're only folding the explicit COPY defs/uses.
       (Ops[0] == 0 || Ops[0] == 1)) {
+    bool IsSpill = Ops[0] == 0;
+    bool IsRefill = !IsSpill;
     const TargetRegisterInfo &TRI = *MF.getSubtarget().getRegisterInfo();
     const MachineRegisterInfo &MRI = MF.getRegInfo();
     MachineBasicBlock &MBB = *MI.getParent();
     const MachineOperand &DstMO = MI.getOperand(0);
     const MachineOperand &SrcMO = MI.getOperand(1);
     unsigned DstReg = DstMO.getReg();
     unsigned SrcReg = SrcMO.getReg();
     auto getRegClass = [&](unsigned Reg) {
       return TargetRegisterInfo::isVirtualRegister(Reg)
                  ? MRI.getRegClass(Reg)
                  : TRI.getMinimalPhysRegClass(Reg);
     };
     const TargetRegisterClass &DstRC = *getRegClass(DstReg);
     const TargetRegisterClass &SrcRC = *getRegClass(SrcReg);
-    if (DstRC.getSize() == SrcRC.getSize()) {
-      if (Ops[0] == 0)
+
+    if (DstMO.getSubReg() == 0 && SrcMO.getSubReg() == 0 &&
+        DstRC.getSize() == SrcRC.getSize()) {
+      if (IsSpill)
         storeRegToStackSlot(MBB, InsertPt, SrcReg, SrcMO.isKill(), FrameIndex,
                             &SrcRC, &TRI);
       else
         loadRegFromStackSlot(MBB, InsertPt, DstReg, FrameIndex, &DstRC, &TRI);
       return &*--InsertPt;
     }
+
+    // Handle cases like spilling def of:
+    //
+    //   %vreg0:sub_32<def,read-undef> = COPY %WZR; GPR64common:%vreg0
+    //
+    // where the physical register source can be widened and stored to the full
+    // virtual reg destination stack slot, in this case producing:
+    //
+    //   STRXui %XZR, <fi#0>
+    //
+    if (IsSpill && SrcMO.getSubReg() == 0 && DstMO.getSubReg() != 0 &&
+        DstMO.isUndef() && TargetRegisterInfo::isPhysicalRegister(SrcReg)) {
+      if (DstMO.getSubReg() == AArch64::sub_32 ||
+          DstMO.getSubReg() == AArch64::ssub ||
+          DstMO.getSubReg() == AArch64::dsub) {
+        const TargetRegisterClass *SpillRC = &DstRC;
+        if (SpillRC == &AArch64::GPR64commonRegClass ||
+            SpillRC == &AArch64::GPR64spRegClass)
+          SpillRC = &AArch64::GPR64RegClass;
+        else if (SpillRC == &AArch64::GPR32commonRegClass ||
+                 SpillRC == &AArch64::GPR32spRegClass)
+          SpillRC = &AArch64::GPR32RegClass;
+        unsigned WidenedSrcReg =
+            TRI.getMatchingSuperReg(SrcReg, DstMO.getSubReg(), SpillRC);
+        if (WidenedSrcReg) {
+          storeRegToStackSlot(MBB, InsertPt, WidenedSrcReg, SrcMO.isKill(),
+                              FrameIndex, SpillRC, &TRI);
+          return &*--InsertPt;
+        }
+      }
+    }
+
+    // Handle cases like refilling use of:
+    //
+    //   %vreg0:sub_32<def,read-undef> = COPY %vreg1; GPR64:%vreg0, GPR32:%vreg1
+    //
+    // where we can load the full virtual reg source stack slot, into the subreg
+    // destination, in this case producing:
+    //
+    //   LDRWui %vreg0:sub_32<def,read-undef>, <fi#0>
+    //
+    if (IsRefill && SrcMO.getSubReg() == 0 && DstMO.getSubReg() != 0 &&
+        DstMO.isUndef()) {
+      auto getRefillRC = [](unsigned SubReg) -> const TargetRegisterClass* {
+        switch (SubReg) {
+        default:
+          return nullptr;
+        case AArch64::sub_32:
+          return &AArch64::GPR32RegClass;
+        case AArch64::ssub:
+          return &AArch64::FPR32RegClass;
+        case AArch64::dsub:
+          return &AArch64::FPR64RegClass;
+        }
+      };
+      if (const TargetRegisterClass *RefillRC =
+              getRefillRC(DstMO.getSubReg())) {
+        assert(SrcRC.getSize() == RefillRC->getSize() &&
+               "Mismatched regclass size on folded subreg COPY");
+        loadRegFromStackSlot(MBB, InsertPt, DstReg, FrameIndex, RefillRC, &TRI);
+        MachineInstr *LoadMI = &*--InsertPt;
+        MachineOperand &LoadDst = LoadMI->getOperand(0);
+        assert(LoadDst.getSubReg() == 0 && "unexpected subreg on refill load");
+        LoadDst.setSubReg(DstMO.getSubReg());
+        LoadDst.setIsUndef();
+        return LoadMI;
+      }
+    }
   }
 
   // Cannot fold.
   return nullptr;
 }
 
 int llvm::isAArch64FrameOffsetLegal(const MachineInstr &MI, int &Offset,
                                     bool *OutUseUnscaledOp,

test/CodeGen/MIR/AArch64/spill-fold.mir

+# RUN: llc -mtriple=aarch64-none-linux-gnu -run-pass greedy -verify-machineinstrs  -o - %s | FileCheck %s
+--- |
+  define i64 @test_subreg_spill_fold() { ret i64 0 }
+  define i64 @test_subreg_spill_fold2() { ret i64 0 }
+  define i64 @test_subreg_refill_fold() { ret i64 0 }
+  define double @test_subreg_refill_fold2() { ret double 0.0 }
+...
+---
+# CHECK-LABEL: name: test_subreg_spill_fold
+# Ensure that the spilled subreg COPY is eliminated and folded into the spill store.
+name:            test_subreg_spill_fold
+registers:
+  - { id: 0, class: gpr64 }
+body:             |
+  bb.0:
+    ; CHECK: STRXui %xzr, %stack.0, 0 :: (store 8 into %stack.0)
+    undef %0.sub_32 = COPY %wzr
+    INLINEASM $nop, 1, 12, implicit-def dead %x0, 12, implicit-def dead %x1, 12, implicit-def dead %x2, 12, implicit-def dead %x3, 12, implicit-def dead %x4, 12, implicit-def dead %x5, 12, implicit-def dead %x6, 12, implicit-def dead %x7, 12, implicit-def dead %x8, 12, implicit-def dead %x9, 12, implicit-def dead %x10, 12, implicit-def dead %x11, 12, implicit-def dead %x12, 12, implicit-def dead %x13, 12, implicit-def dead %x14, 12, implicit-def dead %x15, 12, implicit-def dead %x16, 12, implicit-def dead %x17, 12, implicit-def dead %x18, 12, implicit-def dead %x19, 12, implicit-def dead %x20, 12, implicit-def dead %x21, 12, implicit-def dead %x22, 12, implicit-def dead %x23, 12, implicit-def dead %x24, 12, implicit-def dead %x25, 12, implicit-def dead %x26, 12, implicit-def dead %x27, 12, implicit-def dead %x28, 12, implicit-def dead %fp, 12, implicit-def dead %lr, 12, implicit-def %sp
+    %x0 = COPY %0
+    RET_ReallyLR implicit %x0
+...
+---
+# CHECK-LABEL: name: test_subreg_spill_fold2
+# Similar to test_subreg_spill_fold, but with a vreg0 register class not containing %WZR.
+name:            test_subreg_spill_fold2
+registers:
+  - { id: 0, class: gpr64sp }
+body:             |
+  bb.0:
+    ; CHECK: STRXui %xzr, %stack.0, 0 :: (store 8 into %stack.0)
+    undef %0.sub_32 = COPY %wzr
+    INLINEASM $nop, 1, 12, implicit-def dead %x0, 12, implicit-def dead %x1, 12, implicit-def dead %x2, 12, implicit-def dead %x3, 12, implicit-def dead %x4, 12, implicit-def dead %x5, 12, implicit-def dead %x6, 12, implicit-def dead %x7, 12, implicit-def dead %x8, 12, implicit-def dead %x9, 12, implicit-def dead %x10, 12, implicit-def dead %x11, 12, implicit-def dead %x12, 12, implicit-def dead %x13, 12, implicit-def dead %x14, 12, implicit-def dead %x15, 12, implicit-def dead %x16, 12, implicit-def dead %x17, 12, implicit-def dead %x18, 12, implicit-def dead %x19, 12, implicit-def dead %x20, 12, implicit-def dead %x21, 12, implicit-def dead %x22, 12, implicit-def dead %x23, 12, implicit-def dead %x24, 12, implicit-def dead %x25, 12, implicit-def dead %x26, 12, implicit-def dead %x27, 12, implicit-def dead %x28, 12, implicit-def dead %fp, 12, implicit-def dead %lr, 12, implicit-def %sp
+    %x0 = ADDXri %0, 1, 0
+    RET_ReallyLR implicit %x0
+...
+---
+# CHECK-LABEL: name: test_subreg_refill_fold
+# Ensure that the refilled COPY is eliminated and folded into the refill load.
+name:            test_subreg_refill_fold
+registers:
+  - { id: 0, class: gpr32 }
+  - { id: 1, class: gpr64 }
+body:             |
+  bb.0:
+    %0 = COPY %wzr
+    INLINEASM $nop, 1, 12, implicit-def dead %x0, 12, implicit-def dead %x1, 12, implicit-def dead %x2, 12, implicit-def dead %x3, 12, implicit-def dead %x4, 12, implicit-def dead %x5, 12, implicit-def dead %x6, 12, implicit-def dead %x7, 12, implicit-def dead %x8, 12, implicit-def dead %x9, 12, implicit-def dead %x10, 12, implicit-def dead %x11, 12, implicit-def dead %x12, 12, implicit-def dead %x13, 12, implicit-def dead %x14, 12, implicit-def dead %x15, 12, implicit-def dead %x16, 12, implicit-def dead %x17, 12, implicit-def dead %x18, 12, implicit-def dead %x19, 12, implicit-def dead %x20, 12, implicit-def dead %x21, 12, implicit-def dead %x22, 12, implicit-def dead %x23, 12, implicit-def dead %x24, 12, implicit-def dead %x25, 12, implicit-def dead %x26, 12, implicit-def dead %x27, 12, implicit-def dead %x28, 12, implicit-def dead %fp, 12, implicit-def dead %lr, 12, implicit-def %sp
+    ; CHECK: undef %1.sub_32 = LDRWui %stack.0, 0 :: (load 4 from %stack.0)
+    undef %1.sub_32 = COPY %0
+    %x0 = COPY %1
+    RET_ReallyLR implicit %x0
+...
+---
+# CHECK-LABEL: name: test_subreg_refill_fold2
+# Similar to test_subreg_refill_fold, but with a cross-class copy.
+name:            test_subreg_refill_fold2
+registers:
+  - { id: 0, class: gpr32 }
+  - { id: 1, class: fpr64 }
+body:             |
+  bb.0:
+    %0 = COPY %wzr
+    INLINEASM $nop, 1, 12, implicit-def dead %x0, 12, implicit-def dead %x1, 12, implicit-def dead %x2, 12, implicit-def dead %x3, 12, implicit-def dead %x4, 12, implicit-def dead %x5, 12, implicit-def dead %x6, 12, implicit-def dead %x7, 12, implicit-def dead %x8, 12, implicit-def dead %x9, 12, implicit-def dead %x10, 12, implicit-def dead %x11, 12, implicit-def dead %x12, 12, implicit-def dead %x13, 12, implicit-def dead %x14, 12, implicit-def dead %x15, 12, implicit-def dead %x16, 12, implicit-def dead %x17, 12, implicit-def dead %x18, 12, implicit-def dead %x19, 12, implicit-def dead %x20, 12, implicit-def dead %x21, 12, implicit-def dead %x22, 12, implicit-def dead %x23, 12, implicit-def dead %x24, 12, implicit-def dead %x25, 12, implicit-def dead %x26, 12, implicit-def dead %x27, 12, implicit-def dead %x28, 12, implicit-def dead %fp, 12, implicit-def dead %lr, 12, implicit-def %sp
+    ; CHECK: undef %1.ssub = LDRSui %stack.0, 0 :: (load 4 from %stack.0)
+    undef %1.ssub = COPY %0
+    %d0 = COPY %1
+    RET_ReallyLR implicit %d0
+...