[Spill2Reg] Use AVX opcodes when available

This patch updates the vector spill/reload instructions to use the AVX
opcodes by default if the targets supports it. This can be turned off
with the -spill2reg-no-avx flag.

Original review: https://reviews.llvm.org/D118951

Author: Vasileios Porpodas

llvm/include/llvm/CodeGen/TargetInstrInfo.h

@@ -2303,21 +2303,19 @@ class TargetInstrInfo : public MCInstrInfo {
   }
 
   /// Inserts \p SrcReg into the destination Spill2Reg register \p DstReg.
-  virtual MachineInstr *
-  spill2RegInsertToS2RReg(Register S2RReg, Register SrcReg, int OperationBits,
-                          MachineBasicBlock *MBB,
-                          MachineBasicBlock::iterator InsertBeforeIt,
-                          const TargetRegisterInfo *TRI) const {
+  virtual MachineInstr *spill2RegInsertToS2RReg(
+      Register S2RReg, Register SrcReg, int OperationBits,
+      MachineBasicBlock *MBB, MachineBasicBlock::iterator InsertBeforeIt,
+      const TargetRegisterInfo *TRI, const TargetSubtargetInfo *STI) const {
     llvm_unreachable(
         "Target didn't implement TargetInstrInfo::spill2RegInsertToS2RReg!");
   }
 
   /// Extracts from \p S2RReg into \p DstReg.
-  virtual MachineInstr *
-  spill2RegExtractFromS2RReg(Register DstReg, Register S2RReg,
-                             int OperationBits, MachineBasicBlock *InsertMBB,
-                             MachineBasicBlock::iterator InsertBeforeIt,
-                             const TargetRegisterInfo *TRI) const {
+  virtual MachineInstr *spill2RegExtractFromS2RReg(
+      Register DstReg, Register S2RReg, int OperationBits,
+      MachineBasicBlock *InsertMBB, MachineBasicBlock::iterator InsertBeforeIt,
+      const TargetRegisterInfo *TRI, const TargetSubtargetInfo *STI) const {
     llvm_unreachable("Target didn't implement "
                      "TargetInstrInfo::spill2RegExtractFromS2RReg!");
   }

llvm/include/llvm/CodeGen/TargetRegisterInfo.h

@@ -1249,6 +1249,7 @@ class TargetRegisterInfo : public MCRegisterInfo {
 
   virtual const TargetRegisterClass *
   getCandidateRegisterClassForSpill2Reg(const TargetRegisterInfo *TRI,
+                                        const TargetSubtargetInfo *STI,
                                         Register SpilledReg) const {
     llvm_unreachable("Target didn't implement "
                      "TargetInstrInfo::getCandidateRegisterClassForSpill2Reg!");

llvm/lib/CodeGen/Spill2Reg.cpp

@@ -99,9 +99,10 @@ class Spill2Reg : public MachineFunctionPass {
     /// \Returns the register class of the register being spilled.
     const TargetRegisterClass *
     getSpilledRegClass(const TargetInstrInfo *TII,
-                       const TargetRegisterInfo *TRI) const {
+                       const TargetRegisterInfo *TRI,
+                       const TargetSubtargetInfo *STI) const {
       auto Reg0 = Spills.front().MO->getReg();
-      return TRI->getCandidateRegisterClassForSpill2Reg(TRI, Reg0);
+      return TRI->getCandidateRegisterClassForSpill2Reg(TRI, STI, Reg0);
     }
 
 #ifndef NDEBUG
@@ -220,15 +221,6 @@ void Spill2Reg::collectSpillsAndReloads() {
         }
         unsigned SpillBits = TRI->getRegSizeInBits(MO->getReg(), *MRI);
         Entry.Spills.emplace_back(Spill, MO, SpillBits);
-
-        // If any of the reloads collected so far is in the same MBB then mark
-        // it as non live-in. This is used in `updateLiveIns()` where we update
-        // the liveins of MBBs to include the new vector register. Doing this
-        // now avoids an MBB walk in `updateLiveIns()` which should save
-        // compilation time.
-        for (auto &MID : Entry.Reloads)
-          if (MID.MI->getParent() == &MBB)
-            MID.IsLiveIn = false;
       } else if (const MachineOperand *MO =
                      TII->isLoadFromStackSlotMO(MI, StackSlot)) {
         MachineInstr *Reload = &MI;
@@ -352,7 +344,7 @@ void Spill2Reg::replaceStackWithReg(StackSlotDataEntry &Entry,
 
     TII->spill2RegInsertToS2RReg(
         VectorReg, OldReg, SpillData.SpillBits, StackSpill->getParent(),
-        /*InsertBeforeIt=*/StackSpill->getIterator(), TRI);
+        /*InsertBeforeIt=*/StackSpill->getIterator(), TRI, &MF->getSubtarget());
 
     // Mark VectorReg as live in the instr's BB.
     LRUs[StackSpill->getParent()].addReg(VectorReg);
@@ -369,7 +361,8 @@ void Spill2Reg::replaceStackWithReg(StackSlotDataEntry &Entry,
 
     TII->spill2RegExtractFromS2RReg(
         OldReg, VectorReg, ReloadData.SpillBits, StackReload->getParent(),
-        /*InsertBeforeIt=*/StackReload->getIterator(), TRI);
+        /*InsertBeforeIt=*/StackReload->getIterator(), TRI,
+        &MF->getSubtarget());
 
     // Mark VectorReg as live in the instr's BB.
     LRUs[StackReload->getParent()].addReg(VectorReg);
@@ -480,8 +473,8 @@ void Spill2Reg::generateCode() {
     calculateLiveRegs(Entry, LRU);
 
     // Look for a physical register that is not in LRU.
-    std::optional<MCRegister> PhysVectorRegOpt =
-        tryGetFreePhysicalReg(Entry.getSpilledRegClass(TII, TRI), LRU);
+    std::optional<MCRegister> PhysVectorRegOpt = tryGetFreePhysicalReg(
+        Entry.getSpilledRegClass(TII, TRI, &MF->getSubtarget()), LRU);
     if (!PhysVectorRegOpt)
       continue;
 

llvm/lib/Target/X86/X86InstrInfo.cpp

@@ -11018,14 +11018,24 @@ bool X86InstrInfo::isSpill2RegProfitable(const MachineInstr *MI,
   return MemHeuristic && VecHeuristic;
 }
 
-static unsigned getInsertOrExtractOpcode(unsigned Bits, bool Insert) {
+extern bool useAVX(const TargetSubtargetInfo *STI);
+
+static unsigned getInsertOrExtractOpcode(unsigned Bits, bool Insert,
+                                         const TargetSubtargetInfo *STI) {
+  bool UseAVX = useAVX(STI);
   switch (Bits) {
   case 8:
   case 16:
   case 32:
-    return Insert ? X86::MOVDI2PDIrr : X86::MOVPDI2DIrr;
+    if (UseAVX)
+      return Insert ? X86::VMOVDI2PDIZrr : X86::VMOVPDI2DIZrr;
+    else
+      return Insert ? X86::MOVDI2PDIrr : X86::MOVPDI2DIrr;
   case 64:
-    return Insert ? X86::MOV64toPQIrr : X86::MOVPQIto64rr;
+    if (UseAVX)
+      return Insert ? X86::VMOV64toPQIZrr : X86::VMOVPQIto64Zrr;
+    else
+      return Insert ? X86::MOV64toPQIrr : X86::MOVPQIto64rr;
   default:
     llvm_unreachable("Unsupported bits");
   }
@@ -11063,11 +11073,11 @@ X86InstrInfo::getMovdCompatibleReg(MCRegister OldReg, uint32_t OldRegBits,
 
 MachineInstr *X86InstrInfo::spill2RegInsertToS2RReg(
     Register S2RReg, Register SrcReg, int OperationBits, MachineBasicBlock *MBB,
-    MachineBasicBlock::iterator InsertBeforeIt,
-    const TargetRegisterInfo *TRI) const {
+    MachineBasicBlock::iterator InsertBeforeIt, const TargetRegisterInfo *TRI,
+    const TargetSubtargetInfo *STI) const {
   DebugLoc DL;
   unsigned InsertOpcode =
-      getInsertOrExtractOpcode(OperationBits, true /*insert*/);
+      getInsertOrExtractOpcode(OperationBits, true /*insert*/, STI);
   const MCInstrDesc &InsertMCID = get(InsertOpcode);
   // `movd` does not support 8/16 bit operands. Instead, we use a 32-bit
   // register. For example:
@@ -11083,10 +11093,10 @@ MachineInstr *X86InstrInfo::spill2RegInsertToS2RReg(
 MachineInstr *X86InstrInfo::spill2RegExtractFromS2RReg(
     Register DstReg, Register S2RReg, int OperationBits,
     MachineBasicBlock *InsertMBB, MachineBasicBlock::iterator InsertBeforeIt,
-    const TargetRegisterInfo *TRI) const {
+    const TargetRegisterInfo *TRI, const TargetSubtargetInfo *STI) const {
   DebugLoc DL;
   unsigned ExtractOpcode =
-      getInsertOrExtractOpcode(OperationBits, false /*extract*/);
+      getInsertOrExtractOpcode(OperationBits, false /*extract*/, STI);
   const MCInstrDesc &ExtractMCID = get(ExtractOpcode);
   // `movd` does not support 8/16 bit operands. Instead, we use a 32-bit
   // register. For example:

llvm/lib/Target/X86/X86InstrInfo.h

@@ -752,13 +752,15 @@ class X86InstrInfo final : public X86GenInstrInfo {
   spill2RegInsertToS2RReg(Register S2RReg, Register SrcReg, int OperationBits,
                           MachineBasicBlock *MBB,
                           MachineBasicBlock::iterator InsertBeforeIt,
-                          const TargetRegisterInfo *TRI) const override;
+                          const TargetRegisterInfo *TRI,
+                          const TargetSubtargetInfo *STI) const override;
 
   MachineInstr *
   spill2RegExtractFromS2RReg(Register DstReg, Register S2RReg,
                              int OperationBits, MachineBasicBlock *InsertMBB,
                              MachineBasicBlock::iterator InsertBeforeIt,
-                             const TargetRegisterInfo *TRI) const override;
+                             const TargetRegisterInfo *TRI,
+                             const TargetSubtargetInfo *STI) const override;
 };
 } // namespace llvm
 

llvm/lib/Target/X86/X86RegisterInfo.cpp

@@ -50,6 +50,10 @@ static cl::opt<bool>
                             cl::desc("Disable two address hints for register "
                                      "allocation"));
 
+static cl::opt<bool> Spill2RegNoAVX(
+    "spill2reg-no-avx", cl::Hidden, cl::init(false),
+    cl::desc("Don't use AVX instructions even if the targets supports them."));
+
 X86RegisterInfo::X86RegisterInfo(const Triple &TT)
     : X86GenRegisterInfo((TT.isArch64Bit() ? X86::RIP : X86::EIP),
                          X86_MC::getDwarfRegFlavour(TT, false),
@@ -1279,10 +1283,17 @@ bool X86RegisterInfo::targetSupportsSpill2Reg(
   return X86STI->hasSSE41();
 }
 
+bool useAVX(const TargetSubtargetInfo *STI) {
+  const X86Subtarget *X86STI = static_cast<const X86Subtarget *>(STI);
+  bool UseAVX = X86STI->hasAVX() && !Spill2RegNoAVX;
+  return UseAVX;
+}
+
 const TargetRegisterClass *
 X86RegisterInfo::getCandidateRegisterClassForSpill2Reg(
-    const TargetRegisterInfo *TRI, Register SpilledReg) const {
-  const TargetRegisterClass *VecRegClass =
-      TRI->getRegClass(X86::VR128RegClassID);
+    const TargetRegisterInfo *TRI, const TargetSubtargetInfo *STI,
+    Register SpilledReg) const {
+  const TargetRegisterClass *VecRegClass = TRI->getRegClass(
+      useAVX(STI) ? X86::VR128XRegClassID : X86::VR128RegClassID);
   return VecRegClass;
 }

llvm/lib/Target/X86/X86RegisterInfo.h

@@ -184,6 +184,7 @@ class X86RegisterInfo final : public X86GenRegisterInfo {
 
   const TargetRegisterClass *
   getCandidateRegisterClassForSpill2Reg(const TargetRegisterInfo *TRI,
+                                        const TargetSubtargetInfo *STI,
                                         Register SpilledReg) const override;
 };
 

llvm/test/CodeGen/X86/spill2reg_end_to_end_16bit.ll

@@ -1,5 +1,6 @@
 ; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
 ; RUN: llc %s -o - -mtriple=x86_64-unknown-linux -enable-spill2reg -mattr=+sse4.1 | FileCheck %s
+; RUN: llc %s -o - -mtriple=x86_64-unknown-linux -enable-spill2reg -mattr=+avx | FileCheck --check-prefix=AVX %s
 
 ; End-to-end check that Spill2Reg works with 16-bit registers.
 
@@ -130,6 +131,90 @@ define dso_local void @_Z5spillv() local_unnamed_addr #0 {
 ; CHECK-NEXT:    popq %rbp
 ; CHECK-NEXT:    .cfi_def_cfa_offset 8
 ; CHECK-NEXT:    retq
+;
+; AVX-LABEL: _Z5spillv:
+; AVX:       # %bb.0: # %entry
+; AVX-NEXT:    pushq %rbp
+; AVX-NEXT:    .cfi_def_cfa_offset 16
+; AVX-NEXT:    pushq %r15
+; AVX-NEXT:    .cfi_def_cfa_offset 24
+; AVX-NEXT:    pushq %r14
+; AVX-NEXT:    .cfi_def_cfa_offset 32
+; AVX-NEXT:    pushq %r13
+; AVX-NEXT:    .cfi_def_cfa_offset 40
+; AVX-NEXT:    pushq %r12
+; AVX-NEXT:    .cfi_def_cfa_offset 48
+; AVX-NEXT:    pushq %rbx
+; AVX-NEXT:    .cfi_def_cfa_offset 56
+; AVX-NEXT:    .cfi_offset %rbx, -56
+; AVX-NEXT:    .cfi_offset %r12, -48
+; AVX-NEXT:    .cfi_offset %r13, -40
+; AVX-NEXT:    .cfi_offset %r14, -32
+; AVX-NEXT:    .cfi_offset %r15, -24
+; AVX-NEXT:    .cfi_offset %rbp, -16
+; AVX-NEXT:    movw D0(%rip), %ax
+; AVX-NEXT:    movw %ax, {{[-0-9]+}}(%r{{[sb]}}p) # 2-byte Spill
+; AVX-NEXT:    movzwl D1(%rip), %ecx
+; AVX-NEXT:    movzwl D2(%rip), %edx
+; AVX-NEXT:    movzwl D3(%rip), %esi
+; AVX-NEXT:    movzwl D4(%rip), %edi
+; AVX-NEXT:    movzwl D5(%rip), %r8d
+; AVX-NEXT:    movzwl D6(%rip), %r9d
+; AVX-NEXT:    movzwl D7(%rip), %r10d
+; AVX-NEXT:    movzwl D8(%rip), %r11d
+; AVX-NEXT:    movzwl D9(%rip), %ebx
+; AVX-NEXT:    movzwl D10(%rip), %ebp
+; AVX-NEXT:    movzwl D11(%rip), %r14d
+; AVX-NEXT:    movzwl D12(%rip), %r15d
+; AVX-NEXT:    movzwl D13(%rip), %r12d
+; AVX-NEXT:    movzwl D14(%rip), %r13d
+; AVX-NEXT:    movw D15(%rip), %ax
+; AVX-NEXT:    vmovd %eax, %xmm0
+; AVX-NEXT:    movw D16(%rip), %ax
+; AVX-NEXT:    movw %ax, {{[-0-9]+}}(%r{{[sb]}}p) # 2-byte Spill
+; AVX-NEXT:    movw D17(%rip), %ax
+; AVX-NEXT:    vmovd %eax, %xmm1
+; AVX-NEXT:    movzwl D18(%rip), %eax
+; AVX-NEXT:    movw %ax, {{[-0-9]+}}(%r{{[sb]}}p) # 2-byte Spill
+; AVX-NEXT:    #APP
+; AVX-NEXT:    #NO_APP
+; AVX-NEXT:    movzwl {{[-0-9]+}}(%r{{[sb]}}p), %eax # 2-byte Folded Reload
+; AVX-NEXT:    movw %ax, U0(%rip)
+; AVX-NEXT:    movw %cx, U1(%rip)
+; AVX-NEXT:    movw %dx, U2(%rip)
+; AVX-NEXT:    movw %si, U3(%rip)
+; AVX-NEXT:    movw %di, U4(%rip)
+; AVX-NEXT:    movw %r8w, U5(%rip)
+; AVX-NEXT:    movw %r9w, U6(%rip)
+; AVX-NEXT:    movw %r10w, U7(%rip)
+; AVX-NEXT:    movw %r11w, U8(%rip)
+; AVX-NEXT:    movw %bx, U9(%rip)
+; AVX-NEXT:    movw %bp, U10(%rip)
+; AVX-NEXT:    movw %r14w, U11(%rip)
+; AVX-NEXT:    movw %r15w, U12(%rip)
+; AVX-NEXT:    movw %r12w, U13(%rip)
+; AVX-NEXT:    movw %r13w, U14(%rip)
+; AVX-NEXT:    vmovd %xmm0, %eax
+; AVX-NEXT:    movw %ax, U15(%rip)
+; AVX-NEXT:    movzwl {{[-0-9]+}}(%r{{[sb]}}p), %eax # 2-byte Folded Reload
+; AVX-NEXT:    movw %ax, U16(%rip)
+; AVX-NEXT:    vmovd %xmm1, %eax
+; AVX-NEXT:    movw %ax, U17(%rip)
+; AVX-NEXT:    movzwl {{[-0-9]+}}(%r{{[sb]}}p), %eax # 2-byte Folded Reload
+; AVX-NEXT:    movw %ax, U18(%rip)
+; AVX-NEXT:    popq %rbx
+; AVX-NEXT:    .cfi_def_cfa_offset 48
+; AVX-NEXT:    popq %r12
+; AVX-NEXT:    .cfi_def_cfa_offset 40
+; AVX-NEXT:    popq %r13
+; AVX-NEXT:    .cfi_def_cfa_offset 32
+; AVX-NEXT:    popq %r14
+; AVX-NEXT:    .cfi_def_cfa_offset 24
+; AVX-NEXT:    popq %r15
+; AVX-NEXT:    .cfi_def_cfa_offset 16
+; AVX-NEXT:    popq %rbp
+; AVX-NEXT:    .cfi_def_cfa_offset 8
+; AVX-NEXT:    retq
 entry:
   %0 = load i16, i16* @D0
   %1 = load i16, i16* @D1