[AArch64] Handle more zero reg cases in foldMemoryOperandImpl

[gberry]

Summary

Extend AArch64 foldMemoryOperandImpl() to handle folding spills of COPYs like:

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

by generating:

STRXui %XZR <fi#0>

These cases show up mostly due to zero extended values that get
coalesced and subsequently spilled.

[gberry]

Note that this change depends on two other in-review changes noted above

[MatzeB]

LGTM, with nitpicks. And obviously needs https://reviews.llvm.org/D26521 to be accepted first.

lib/Target/AArch64/AArch64InstrInfo.cpp

2636–2637

Maybe use reference variables to indicate they cannot contain nullptr.

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

13

A sentece describing what you are testing here would be nice.

15

Do the defs need to be early clobbers? I would have assumed it doesn't matter...

[mkuper]

Note that since D26521 only enables folding of stack slots for X86, it changed the MachineMemOperand size handling only in the FI version of TTI::foldMemoryOperand().
If you're enabling it for LoadMI here, you may need to do something similar for the LoadMI version (unless it happens to just work).

[gberry]

I'm abandoning this patch in favor of https://reviews.llvm.org/D27425, which is more general

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;
 
+  bool isSubregFoldable(MachineOperand &MO,
+                        MachineInstr *LoadMI) const override;
+
   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

void llvm::emitFrameOffset(MachineBasicBlock &MBB,
   }
   BuildMI(MBB, MBBI, DL, TII->get(Opc), DestReg)
       .addReg(SrcReg)
       .addImm(Offset)
       .addImm(AArch64_AM::getShifterImm(AArch64_AM::LSL, 0))
       .setMIFlag(Flag);
 }
 
+bool AArch64InstrInfo::isSubregFoldable(MachineOperand &MO,
+                                        MachineInstr *LoadMI) const {
+  return true;
+}
+
 MachineInstr *AArch64InstrInfo::foldMemoryOperandImpl(
     MachineFunction &MF, MachineInstr &MI, ArrayRef<unsigned> Ops,
     MachineBasicBlock::iterator InsertPt, int FrameIndex,
     LiveIntervals *LIS) const {
   // This is a bit of a hack. Consider this instruction:
   //
   //   %vreg0<def> = COPY %SP; GPR64all:%vreg0
   //
   // We explicitly chose GPR64all for the virtual register so such a copy might
   // 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 &&
         TargetRegisterInfo::isVirtualRegister(SrcReg)) {
       MF.getRegInfo().constrainRegClass(SrcReg, &AArch64::GPR64RegClass);
       return nullptr;
     }
   }
 
   // Handle the case where a WZR/XZR copy is being spilled but the destination
   // register class doesn't contain WZR/XZR.  For example:
   //
   //   %vreg0<def> = COPY %XZR; GPR64common:%vreg0
   //
   // In this case we can still safely fold away the COPY and generate the
   // following spill code:
   //
   //   STRXui %XZR, <fi#0>
   //
-  if (MI.isFullCopy() && Ops.size() == 1 && Ops[0] == 0) {
-    MachineBasicBlock *MBB = MI.getParent();
+  // Also handle cases like:
+  //
+  //   %vreg0:sub_32<def,read-undef> = COPY %WZR; GPR64common:%vreg0
+  //
+  // by generating:
+  //
+  //   STRXui %XZR <fi#0>
+  //
+  if (MI.isCopy() && Ops.size() == 1 && Ops[0] == 0) {
+    const MachineOperand &DstMO = MI.getOperand(0);
+    assert(TargetRegisterInfo::isVirtualRegister(DstMO.getReg()) &&
+           "expected virtual register as destination of COPY");
+    // Only handle full writes to the destination register, or subreg writes
+    // that are undef for the non-written parts.
+    if (DstMO.getSubReg() == 0 || DstMO.isUndef()) {
       const MachineOperand &SrcMO = MI.getOperand(1);
       unsigned SrcReg = SrcMO.getReg();
       if (SrcReg == AArch64::WZR || SrcReg == AArch64::XZR) {
+        MachineBasicBlock &MBB = *MI.getParent();
+        MachineRegisterInfo &MRI = MBB.getParent()->getRegInfo();
         const TargetRegisterInfo *TRI = &getRegisterInfo();
-      const TargetRegisterClass *RC = SrcReg == AArch64::WZR
-                                          ? &AArch64::GPR32RegClass
-                                          : &AArch64::GPR64RegClass;
-      storeRegToStackSlot(*MBB, InsertPt, SrcReg, SrcMO.isKill(), FrameIndex,
-                          RC, TRI);
+        const TargetRegisterClass *DstRC = MRI.getRegClass(DstMO.getReg());
+        const TargetRegisterClass *SpillRC;
+        unsigned SpillReg;
+        if (DstRC->getSize() == 4) {
+          SpillRC = &AArch64::GPR32RegClass;
+          SpillReg = AArch64::WZR;
+        } else {
+          assert(DstRC->getSize() == 8);
+          SpillRC = &AArch64::GPR64RegClass;
+          SpillReg = AArch64::XZR;
+        }
+        storeRegToStackSlot(MBB, InsertPt, SpillReg, SrcMO.isKill(), FrameIndex,
+                            SpillRC, TRI);
         return &*--InsertPt;
       }
     }
+  }
 
   // Cannot fold.
   return nullptr;
 }
 
 int llvm::isAArch64FrameOffsetLegal(const MachineInstr &MI, int &Offset,
                                     bool *OutUseUnscaledOp,
                                     unsigned *OutUnscaledOp,

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

+# RUN: llc -mtriple=aarch64-none-linux-gnu -run-pass greedy -verify-machineinstrs  -o - %s | FileCheck %s
+
+--- |
+  define i64 @test_zr_spill_fold(i1 %c) { ret i64 0 }
+...
+---
+# CHECK-LABEL: name: test_zr_spill_fold
+name:            test_zr_spill_fold
+registers:       
+  - { id: 0, class: gpr64 }
+body:             |
+  bb.0:
+    ; CHECK: STRXui %xzr
+    undef %0.sub_32 = COPY %wzr
+    INLINEASM $nop, 1, 12, implicit-def dead early-clobber %x0, 12, implicit-def dead early-clobber %x1, 12, implicit-def dead early-clobber %x2, 12, implicit-def dead early-clobber %x3, 12, implicit-def dead early-clobber %x4, 12, implicit-def dead early-clobber %x5, 12, implicit-def dead early-clobber %x6, 12, implicit-def dead early-clobber %x7, 12, implicit-def dead early-clobber %x8, 12, implicit-def dead early-clobber %x9, 12, implicit-def dead early-clobber %x10, 12, implicit-def dead early-clobber %x11, 12, implicit-def dead early-clobber %x12, 12, implicit-def dead early-clobber %x13, 12, implicit-def dead early-clobber %x14, 12, implicit-def dead early-clobber %x15, 12, implicit-def dead early-clobber %x16, 12, implicit-def dead early-clobber %x17, 12, implicit-def dead early-clobber %x18, 12, implicit-def dead early-clobber %x19, 12, implicit-def dead early-clobber %x20, 12, implicit-def dead early-clobber %x21, 12, implicit-def dead early-clobber %x22, 12, implicit-def dead early-clobber %x23, 12, implicit-def dead early-clobber %x24, 12, implicit-def dead early-clobber %x25, 12, implicit-def dead early-clobber %x26, 12, implicit-def dead early-clobber %x27, 12, implicit-def dead early-clobber %x28, 12, implicit-def dead early-clobber %fp, 12, implicit-def dead early-clobber %lr, 12, implicit-def early-clobber %sp
+    %x0 = COPY %0
+    RET_ReallyLR implicit %x0
+...