[AArch64, TargetRegisterInfo, TargetInstrInfo] Add hook to check if physreg spillable

[gberry]

Summary

Add TargetRegisterInfo hook canSpillPhysReg() and use it in
TargetTransformInfo::canFoldCopy to catch more cases where the COPY
being spilled is copying from a spill-able physical register, but the
register is not in the COPY destination register's regclass.

For example, when spilling:

%vreg0 = COPY %XZR ; %vreg0:GPR64common

without this change, canFoldCopy would return false since %XZR is not in
GPR64common. So the spill code generated would be:

%vreg0 = COPY %XZR
STR %vreg

instead of

STR %ZXR

Implement canSpillPhysReg() for AArch64, where all phys regs are
spillable other than SP/WSP.

[gberry]

I'm thinking it might be cleaner to just implement this in AArch64InstrInfo::foldMemoryOperandImpl to avoid needing a new hook. I'm going to pursue that for a bit.

[gberry]

Superseded by https://reviews.llvm.org/D26976

include/llvm/Target/TargetRegisterInfo.h

public:
   /// considered unavailable at all times, e.g. SP, RA. This is
   /// used by register scavenger to determine what registers are free.
   virtual BitVector getReservedRegs(const MachineFunction &MF) const = 0;
 
   /// Returns true if PhysReg is unallocatable and constant throughout the
   /// function.  Used by MachineRegisterInfo::isConstantPhysReg().
   virtual bool isConstantPhysReg(unsigned PhysReg) const { return false; }
 
+  /// Returns true if PhysReg can be spilled.  It is safe to define this
+  /// function conservatively, i.e. it does not modify the usual rule that all
+  /// allocatable registers can be spilled.
+  virtual bool canSpillPhysReg(unsigned PhysReg) const { return false; }
+
   /// Prior to adding the live-out mask to a stackmap or patchpoint
   /// instruction, provide the target the opportunity to adjust it (mainly to
   /// remove pseudo-registers that should be ignored).
   virtual void adjustStackMapLiveOutMask(uint32_t *Mask) const { }
 
   /// Return a super-register of the specified register
   /// Reg so its sub-register of index SubIdx is Reg.
   unsigned getMatchingSuperReg(unsigned Reg, unsigned SubIdx,

lib/CodeGen/TargetInstrInfo.cpp

MachineInstr *TargetInstrInfo::duplicate(MachineInstr &Orig,
                                          MachineFunction &MF) const {
   assert(!Orig.isNotDuplicable() && "Instruction cannot be duplicated");
   return MF.CloneMachineInstr(&Orig);
 }
 
 // If the COPY instruction in MI can be folded to a stack operation, return
 // the register class to use.
 static const TargetRegisterClass *canFoldCopy(const MachineInstr &MI,
-                                              unsigned FoldIdx) {
+                                              unsigned FoldIdx,
+                                              const TargetRegisterInfo &TRI) {
   assert(MI.isCopy() && "MI must be a COPY instruction");
   if (MI.getNumOperands() != 2)
     return nullptr;
   assert(FoldIdx<2 && "FoldIdx refers no nonexistent operand");
 
   const MachineOperand &FoldOp = MI.getOperand(FoldIdx);
   const MachineOperand &LiveOp = MI.getOperand(1 - FoldIdx);
 
   if (FoldOp.getSubReg() || LiveOp.getSubReg())
     return nullptr;
 
   unsigned FoldReg = FoldOp.getReg();
   unsigned LiveReg = LiveOp.getReg();
 
   assert(TargetRegisterInfo::isVirtualRegister(FoldReg) &&
          "Cannot fold physregs");
 
   const MachineRegisterInfo &MRI = MI.getParent()->getParent()->getRegInfo();
   const TargetRegisterClass *RC = MRI.getRegClass(FoldReg);
 
-  if (TargetRegisterInfo::isPhysicalRegister(LiveOp.getReg()))
-    return RC->contains(LiveOp.getReg()) ? RC : nullptr;
+  if (TargetRegisterInfo::isPhysicalRegister(LiveOp.getReg())) {
+    if (RC->contains(LiveOp.getReg()))
+      return RC;
+    const TargetRegisterClass *LivePhysRC =
+        TRI.getMinimalPhysRegClass(LiveOp.getReg());
+    if (LivePhysRC->getSize() == RC->getSize() && TRI.canSpillPhysReg(LiveOp.getReg()))
+      return LivePhysRC;
+    return nullptr;
+  }
 
   if (RC->hasSubClassEq(MRI.getRegClass(LiveReg)))
     return RC;
 
   // FIXME: Allow folding when register classes are memory compatible.
   return nullptr;
 }
 
if (NewMI) {
 
     return NewMI;
   }
 
   // Straight COPY may fold as load/store.
   if (!MI.isCopy() || Ops.size() != 1)
     return nullptr;
 
-  const TargetRegisterClass *RC = canFoldCopy(MI, Ops[0]);
+  const TargetRegisterInfo *TRI = MF.getSubtarget().getRegisterInfo();
+  const TargetRegisterClass *RC = canFoldCopy(MI, Ops[0], *TRI);
   if (!RC)
     return nullptr;
 
   const MachineOperand &MO = MI.getOperand(1 - Ops[0]);
   MachineBasicBlock::iterator Pos = MI;
-  const TargetRegisterInfo *TRI = MF.getSubtarget().getRegisterInfo();
 
   if (Flags == MachineMemOperand::MOStore)
     storeRegToStackSlot(*MBB, Pos, MO.getReg(), MO.isKill(), FI, RC, TRI);
   else
     loadRegFromStackSlot(*MBB, Pos, MO.getReg(), FI, RC, TRI);
   return &*--Pos;
 }
 

lib/Target/AArch64/AArch64RegisterInfo.h

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;
 
   BitVector getReservedRegs(const MachineFunction &MF) const override;
   bool isConstantPhysReg(unsigned PhysReg) const override;
+  bool canSpillPhysReg(unsigned PhysReg) const override;
   const TargetRegisterClass *
   getPointerRegClass(const MachineFunction &MF,
                      unsigned Kind = 0) const override;
   const TargetRegisterClass *
   getCrossCopyRegClass(const TargetRegisterClass *RC) const override;
 
   bool requiresRegisterScavenging(const MachineFunction &MF) const override;
   bool useFPForScavengingIndex(const MachineFunction &MF) const override;

lib/Target/AArch64/AArch64RegisterInfo.cpp

bool AArch64RegisterInfo::isReservedReg(const MachineFunction &MF,
 
   return false;
 }
 
 bool AArch64RegisterInfo::isConstantPhysReg(unsigned PhysReg) const {
   return PhysReg == AArch64::WZR || PhysReg == AArch64::XZR;
 }
 
+// This is needed to catch cases of spill copy folding where the COPY
+// destination register class is incompatible with PhysReg, most commonly when
+// spilling a COPY of WZR/XZR.
+bool AArch64RegisterInfo::canSpillPhysReg(unsigned PhysReg) const {
+  switch (PhysReg) {
+  default:
+    return true;
+  case AArch64::WSP:
+  case AArch64::SP:
+    return false;
+  }
+}
+
 const TargetRegisterClass *
 AArch64RegisterInfo::getPointerRegClass(const MachineFunction &MF,
                                       unsigned Kind) const {
   return &AArch64::GPR64RegClass;
 }
 
 const TargetRegisterClass *
 AArch64RegisterInfo::getCrossCopyRegClass(const TargetRegisterClass *RC) const {

test/CodeGen/AArch64/zero-reg.ll

 ; instruction (0b11111 in the Rn field would mean "sp").
   %addr = getelementptr i32, i32* null, i64 0
   store i32 %val, i32* %addr
 ; CHECK: str {{w[0-9]+}}, [{{x[0-9]+|sp}}]
 
   ret void
 ; CHECK: ret
 }
+
+declare i32 @bar()
+declare i32 @baz()
+
+; Check that the spill of the zero value gets stored directly instead
+; of being copied from wzr and then stored.
+define i32 @test_zr_spill_copyprop1(i1 %c) {
+; CHECK-LABEL: test_zr_spill_copyprop1:
+entry:
+  br i1 %c, label %if.else, label %if.then
+
+if.else:
+; CHECK: bl bar
+; CHECK-NEXT: str w0, [sp, #[[SLOT:[0-9]+]]]
+  %call1 = tail call i32 @bar()
+  br label %if.end
+
+if.then:
+; CHECK: bl baz
+; CHECK-NEXT: str wzr, [sp, #[[SLOT]]]
+  %call2 = tail call i32 @baz()
+  br label %if.end
+
+if.end:
+  %x.0 = phi i32 [ 0, %if.then ], [ %call1, %if.else ]
+  call void asm sideeffect "", "~{x0},~{x1},~{x2},~{x3},~{x4},~{x5},~{x6},~{x7},~{x8},~{x9},~{x10},~{x11},~{x12},~{x13},~{x14},~{x15},~{x16},~{x17},~{x18},~{x19},~{x20},~{x21},~{x22},~{x23},~{x24},~{x25},~{x26},~{x27},~{x28},~{fp},~{lr},~{sp}"() nounwind
+  ret i32 %x.0
+}
+
+; Similar to test_zr_spill_copyprop1, but with mis-matched register
+; class between %x.0 and the 0 from %if.then.
+define i32 @test_zr_spill_copyprop2(i1 %c) {
+; CHECK-LABEL: test_zr_spill_copyprop2:
+entry:
+  br i1 %c, label %if.else, label %if.then
+
+if.else:
+; CHECK: bl bar
+; CHECK-NEXT: str w0, [sp, #[[SLOT:[0-9]+]]]
+  %call1 = tail call i32 @bar()
+  br label %if.end
+
+if.then:
+; CHECK: bl baz
+; CHECK-NEXT: str wzr, [sp, #[[SLOT]]]
+  %call2 = tail call i32 @baz()
+  br label %if.end
+
+if.end:
+  %x.0 = phi i32 [ 0, %if.then ], [ %call1, %if.else ]
+  call void asm sideeffect "", "~{x0},~{x1},~{x2},~{x3},~{x4},~{x5},~{x6},~{x7},~{x8},~{x9},~{x10},~{x11},~{x12},~{x13},~{x14},~{x15},~{x16},~{x17},~{x18},~{x19},~{x20},~{x21},~{x22},~{x23},~{x24},~{x25},~{x26},~{x27},~{x28},~{fp},~{lr},~{sp}"() nounwind
+  %x.1 = add i32 %x.0, 1
+  ret i32 %x.1
+}