[X86][AMX][fastalloc] Allocate tile register based on its shape.

[LuoYuanke]

Summary

Allocating tile register in a separate pass have 3 benifits.

1. When spill tile register in fast register allocation, it would create

a virtual register for stride. The instruction that def the virtual
register is inserted after current instruction which is allocating
physical register, so fast register allocation won't access the new
instrution and the virtual register won't be allocated with physical
register.

2. When fill the shape information of tile configure, we need to know

the tile physical register. If the shape (row, column) is allocated with
physical register, the physical register may be re-defined before
writing it to stack. The re-def may happen during split or spill
registers. If the shape is in virtual register, there is no such
problem.
%al = 1
mov %al, %bl
%al = 2
store stack.cfg, %al ; config
ldtilecfg

3. Single configure may fail. In that case we can fallback to

multi-config and allocate tile register seperately, while leave other
registers be allocated by greedy RA.

[LuoYuanke]

Fix lit test failure.

[LuoYuanke]

Rebase.

[LuoYuanke]

Distinguish shape when allocating 2D tile register.

[LuoYuanke]

Add test case for reuse physical tile register.

[LuoYuanke]

Add test case across config.

[MatzeB]

There's a lot going on here.

• Could you extract the ShouldAllocClass fixes for fastregalloc into a separate diff so we can get discuss them separately and get feedback from AMDGPU folks who introduced this and are the major other user of this AFAIK.
• I am still wrapping my head around tile registers and tile register configs; specifically I wonder if the support for that really needs to be integrated into the generic register allocation code or whether there is a way to materialize the register configurations in a post-pass. For example do you know how the legacy x86 x87-FPU support works, where we use the register allocator to allocate pseudo FP register fp0-fp7 and then use a post-pass in X86FloatingPoint to insert the necessary stack management operations after the fact. I am not saying it's the same problem, but it is an example of an instance where we managed to have the regalloc allocate to some intermediate pseudo registers and adapt to the complications (in that case register stacks) in a target-specific pass so we wouldn't need to introduce the concept of register stacks to the generic code.

[LuoYuanke]

In D125602#3547595, @MatzeB wrote:

There's a lot going on here.

• Could you extract the ShouldAllocClass fixes for fastregalloc into a separate diff so we can get discuss them separately and get feedback from AMDGPU folks who introduced this and are the major other user of this AFAIK.

Thank you for review. Sure, I'll extract ShouldAllocClass fixes for fastregalloc into a separate patch.

• I am still wrapping my head around tile registers and tile register configs; specifically I wonder if the support for that really needs to be integrated into the generic register allocation code or whether there is a way to materialize the register configurations in a post-pass. For example do you know how the legacy x86 x87-FPU support works, where we use the register allocator to allocate pseudo FP register fp0-fp7 and then use a post-pass in X86FloatingPoint to insert the necessary stack management operations after the fact. I am not saying it's the same problem, but it is an example of an instance where we managed to have the regalloc allocate to some intermediate pseudo registers and adapt to the complications (in that case register stacks) in a target-specific pass so we wouldn't need to introduce the concept of register stacks to the generic code.

We have pre-pass to insert config instructions and post-pass to fill the shape information of each physical tile register and feed it to config instruction.
I think they are different problem for tile register allocation and x87-FPU register allocation. The problem of x87-FPU register is that register is arranged in stack order in HW and I'm not sure stackify pass is very efficient to convert pseudo instruction to x87 instruction without inserting extra instructions to adjust stack order. The problem for tile registers is introduced by configure. Before accessing any tile register, they should be configured to specify the shape (row, column) information. HW would base on the shape information to operate AMX intruction. The configure instruction would clear all the data of tile registers. If we add a post-fixup pass like stackify pass, we need reconfig registers that is allocated to the same physical tile register but have different shape, and the reconfig would clobber all tile registers. That may generate too many config instruction and spill/reload in post-fix pass. Nevertheless I am happy to work with community for any suggestion to improve the solution for AMX register allocation.

llvm/lib/CodeGen/RegAllocFast.cpp

getMBBBeginInsertionPoint(MachineBasicBlock &MBB,
                               SmallSet<Register, 2> &PrologLiveIns) const;
 
     void reloadAtBegin(MachineBasicBlock &MBB);
     void setPhysReg(MachineInstr &MI, MachineOperand &MO, MCPhysReg PhysReg);
 
     Register traceCopies(Register VirtReg) const;
     Register traceCopyChain(Register Reg) const;
 
+    bool ShouldAllocateRegister(const Register Reg) const;
     int getStackSpaceFor(Register VirtReg);
     void spill(MachineBasicBlock::iterator Before, Register VirtReg,
                MCPhysReg AssignedReg, bool Kill, bool LiveOut);
     void reload(MachineBasicBlock::iterator Before, Register VirtReg,
                 MCPhysReg PhysReg);
 
     bool mayLiveOut(Register VirtReg);
     bool mayLiveIn(Register VirtReg);
 
     void dumpState() const;
   };
 
 } // end anonymous namespace
 
 char RegAllocFast::ID = 0;
 
 INITIALIZE_PASS(RegAllocFast, "regallocfast", "Fast Register Allocator", false,
                 false)
 
+bool RegAllocFast::ShouldAllocateRegister(const Register Reg) const {
+  const TargetRegisterClass &RC = *MRI->getRegClass(Reg);
+  return ShouldAllocateClass(*TRI, RC);
+}
+
 void RegAllocFast::setPhysRegState(MCPhysReg PhysReg, unsigned NewState) {
   for (MCRegUnitIterator UI(PhysReg, TRI); UI.isValid(); ++UI)
     RegUnitStates[*UI] = NewState;
 }
 
 bool RegAllocFast::isPhysRegFree(MCPhysReg PhysReg) const {
   for (MCRegUnitIterator UI(PhysReg, TRI); UI.isValid(); ++UI) {
     if (RegUnitStates[*UI] != regFree)

 #endif
 
 /// Count number of defs consumed from each register class by \p Reg
 void RegAllocFast::addRegClassDefCounts(std::vector<unsigned> &RegClassDefCounts,
                                         Register Reg) const {
   assert(RegClassDefCounts.size() == TRI->getNumRegClasses());
 
   if (Reg.isVirtual()) {
+    if (ShouldAllocateRegister(Reg))
+      return;
     const TargetRegisterClass *OpRC = MRI->getRegClass(Reg);
     for (unsigned RCIdx = 0, RCIdxEnd = TRI->getNumRegClasses();
          RCIdx != RCIdxEnd; ++RCIdx) {
       const TargetRegisterClass *IdxRC = TRI->getRegClass(RCIdx);
       // FIXME: Consider aliasing sub/super registers.
       if (OpRC->hasSubClassEq(IdxRC))
         ++RegClassDefCounts[RCIdx];
     }
void RegAllocFast::allocateInstruction(MachineInstr &MI) {
   bool HasDef = false;
   bool HasEarlyClobber = false;
   bool NeedToAssignLiveThroughs = false;
   for (unsigned I = 0; I < MI.getNumOperands(); ++I) {
     MachineOperand &MO = MI.getOperand(I);
     if (MO.isReg()) {
       Register Reg = MO.getReg();
       if (Reg.isVirtual()) {
+        if (!ShouldAllocateRegister(Reg))
+          continue;
         if (MO.isDef()) {
           HasDef = true;
           HasVRegDef = true;
           if (MO.isEarlyClobber()) {
             HasEarlyClobber = true;
             NeedToAssignLiveThroughs = true;
           }
           if ((MO.isTied() && !TiedOpIsUndef(MO, I)) ||
if (HasVRegDef) {
             if (Reg.isPhysical()) {
               LLVM_DEBUG(dbgs() << "mark extra used: " << printReg(Reg, TRI)
                                 << '\n');
               markPhysRegUsedInInstr(Reg);
             }
           }
 
           if (MO.isDef()) {
-            if (Reg.isVirtual())
+            if (Reg.isVirtual() && ShouldAllocateRegister(Reg))
               DefOperandIndexes.push_back(I);
 
             addRegClassDefCounts(RegClassDefCounts, Reg);
           }
         }
 
         llvm::sort(DefOperandIndexes, [&](uint16_t I0, uint16_t I1) {
           const MachineOperand &MO0 = MI.getOperand(I0);
if (HasVRegDef) {
         }
       } else {
         // Assign virtual register defs.
         for (unsigned I = 0, E = MI.getNumOperands(); I < E; ++I) {
           MachineOperand &MO = MI.getOperand(I);
           if (!MO.isReg() || !MO.isDef())
             continue;
           Register Reg = MO.getReg();
-          if (Reg.isVirtual())
+          if (Reg.isVirtual() && ShouldAllocateRegister(Reg))
             defineVirtReg(MI, I, Reg);
         }
       }
     }
 
     // Free registers occupied by defs.
     // Iterate operands in reverse order, so we see the implicit super register
     // defs first (we added them earlier in case of <def,read-undef>).
for (signed I = MI.getNumOperands() - 1; I >= 0; --I) {
              "tied def assigned to clobbered register");
 
       // Do not free tied operands and early clobbers.
       if ((MO.isTied() && !TiedOpIsUndef(MO, I)) || MO.isEarlyClobber())
         continue;
       Register Reg = MO.getReg();
       if (!Reg)
         continue;
+      if (Reg.isVirtual() && !ShouldAllocateRegister(Reg))
+        continue;
       assert(Reg.isPhysical());
       if (MRI->isReserved(Reg))
         continue;
       freePhysReg(Reg);
       unmarkRegUsedInInstr(Reg);
     }
   }
 
void RegAllocFast::allocateInstruction(MachineInstr &MI) {
 
   // Allocate virtreg uses and insert reloads as necessary.
   bool HasUndefUse = false;
   for (unsigned I = 0; I < MI.getNumOperands(); ++I) {
     MachineOperand &MO = MI.getOperand(I);
     if (!MO.isReg() || !MO.isUse())
       continue;
     Register Reg = MO.getReg();
-    if (!Reg.isVirtual())
+    if (!Reg.isVirtual() || !ShouldAllocateRegister(Reg))
       continue;
 
     if (MO.isUndef()) {
       HasUndefUse = true;
       continue;
     }
 
 
void RegAllocFast::allocateInstruction(MachineInstr &MI) {
   // Allocate undef operands. This is a separate step because in a situation
   // like  ` = OP undef %X, %X`    both operands need the same register assign
   // so we should perform the normal assignment first.
   if (HasUndefUse) {
     for (MachineOperand &MO : MI.uses()) {
       if (!MO.isReg() || !MO.isUse())
         continue;
       Register Reg = MO.getReg();
-      if (!Reg.isVirtual())
+      if (!Reg.isVirtual() || !ShouldAllocateRegister(Reg))
         continue;
 
       assert(MO.isUndef() && "Should only have undef virtreg uses left");
       allocVirtRegUndef(MO);
     }
   }
 
   // Free early clobbers.
void RegAllocFast::handleBundle(MachineInstr &MI) {
   MachineBasicBlock::instr_iterator BundledMI = MI.getIterator();
   ++BundledMI;
   while (BundledMI->isBundledWithPred()) {
     for (MachineOperand &MO : BundledMI->operands()) {
       if (!MO.isReg())
         continue;
 
       Register Reg = MO.getReg();
-      if (!Reg.isVirtual())
+      if (!Reg.isVirtual() || !ShouldAllocateRegister(Reg))
         continue;
 
       DenseMap<Register, MCPhysReg>::iterator DI;
       DI = BundleVirtRegsMap.find(Reg);
       assert(DI != BundleVirtRegsMap.end() && "Unassigned virtual register");
 
       setPhysReg(MI, MO, DI->second);
     }

llvm/lib/Target/X86/X86LowerAMXType.cpp

bool runOnFunction(Function &F) override {
     C |= LAC.transformAllAMXCast();
 
     X86LowerAMXType LAT(F);
     C |= LAT.visit();
 
     // Prepare for fast register allocation at O0.
     // Todo: May better check the volatile model of AMX code, not just
     // by checking Attribute::OptimizeNone and CodeGenOpt::None.
-    if (TM->getOptLevel() == CodeGenOpt::None) {
+    if (0) { // (TM->getOptLevel() == CodeGenOpt::None) {
       // If Front End not use O0 but the Mid/Back end use O0, (e.g.
       // "Clang -O2 -S -emit-llvm t.c" + "llc t.ll") we should make
       // sure the amx data is volatile, that is nessary for AMX fast
       // register allocation.
       if (!F.hasFnAttribute(Attribute::OptimizeNone)) {
         X86VolatileTileData VTD(F);
         C = VTD.volatileTileData() || C;
       }

llvm/lib/Target/X86/X86RegisterInfo.h

public:
   /// register scavenger to determine what registers are free.
   BitVector getReservedRegs(const MachineFunction &MF) const override;
 
   /// isArgumentReg - Returns true if Reg can be used as an argument to a
   /// function.
   bool isArgumentRegister(const MachineFunction &MF,
                           MCRegister Reg) const override;
 
+  bool isTileRegisterClass(const TargetRegisterClass *RC) const;
+
   /// Returns true if PhysReg is a fixed register.
   bool isFixedRegister(const MachineFunction &MF,
                        MCRegister PhysReg) const override;
 
   void adjustStackMapLiveOutMask(uint32_t *Mask) const override;
 
   bool hasBasePointer(const MachineFunction &MF) const;
 

llvm/lib/Target/X86/X86RegisterInfo.cpp

bool X86RegisterInfo::isFixedRegister(const MachineFunction &MF,
   // Don't use the frame pointer if it's being used.
   const X86FrameLowering &TFI = *getFrameLowering(MF);
   if (TFI.hasFP(MF) && TRI.isSuperOrSubRegisterEq(X86::RBP, PhysReg))
     return true;
 
   return X86GenRegisterInfo::isFixedRegister(MF, PhysReg);
 }
 
+bool X86RegisterInfo::isTileRegisterClass(const TargetRegisterClass *RC) const {
+  return RC->getID() == X86::TILERegClassID;
+}
+
 void X86RegisterInfo::adjustStackMapLiveOutMask(uint32_t *Mask) const {
   // Check if the EFLAGS register is marked as live-out. This shouldn't happen,
   // because the calling convention defines the EFLAGS register as NOT
   // preserved.
   //
   // Unfortunatelly the EFLAGS show up as live-out after branch folding. Adding
   // an assert to track this and clear the register afterwards to avoid
   // unnecessary crashes during release builds.

llvm/lib/Target/X86/X86TargetMachine.cpp

public:
   void addMachineSSAOptimization() override;
   void addPreRegAlloc() override;
   bool addPostFastRegAllocRewrite() override;
   void addPostRegAlloc() override;
   void addPreEmitPass() override;
   void addPreEmitPass2() override;
   void addPreSched2() override;
   bool addPreRewrite() override;
+  bool addRegAssignAndRewriteFast() override;
 
   std::unique_ptr<CSEConfigBase> getCSEConfig() const override;
 };
 
 class X86ExecutionDomainFix : public ExecutionDomainFix {
 public:
   static char ID;
   X86ExecutionDomainFix() : ExecutionDomainFix(ID, X86::VR128XRegClass) {}
addPass(createUnpackMachineBundles([](const MachineFunction &MF) {
       // present in the module.
       const Function &F = MF.getFunction();
       const Module *M = F.getParent();
       return M->getFunction("objc_retainAutoreleasedReturnValue") ||
              M->getFunction("objc_unsafeClaimAutoreleasedReturnValue");
     }));
 }
 
+static bool onlyAllocateTileRegisters(const TargetRegisterInfo &TRI,
+                                      const TargetRegisterClass &RC) {
+  return static_cast<const X86RegisterInfo &>(TRI).isTileRegisterClass(&RC);
+}
+
+bool X86PassConfig::addRegAssignAndRewriteFast() {
+  // Allocate AMX registers separately.
+  addPass(createFastRegisterAllocator(onlyAllocateTileRegisters, false));
+  return TargetPassConfig::addRegAssignAndRewriteFast();
+}
+
 bool X86PassConfig::addPostFastRegAllocRewrite() {
   addPass(createX86FastTileConfigPass());
   return true;
 }
 
 bool X86PassConfig::addPreRewrite() {
   addPass(createX86TileConfigPass());
   return true;
 }
 
 std::unique_ptr<CSEConfigBase> X86PassConfig::getCSEConfig() const {
   return getStandardCSEConfigForOpt(TM->getOptLevel());
 }

llvm/test/CodeGen/X86/AMX/amx-configO2toO0.ll

 ; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
 ; RUN: llc < %s -O0 -mtriple=x86_64-unknown-unknown -mattr=+amx-int8 -mattr=+avx512f | FileCheck %s --check-prefix=AVX512
 
 @buf = dso_local global [1024 x i8] zeroinitializer, align 16
 @buf2 = dso_local global [1024 x i8] zeroinitializer, align 16
 
 define dso_local void @test_api(i32 %cond, i16 signext %row, i16 signext %col) nounwind {
 ; AVX512-LABEL: test_api:
 ; AVX512:       # %bb.0: # %entry
-; AVX512-NEXT:    pushq %rbp
-; AVX512-NEXT:    movq %rsp, %rbp
-; AVX512-NEXT:    andq $-1024, %rsp # imm = 0xFC00
-; AVX512-NEXT:    subq $8192, %rsp # imm = 0x2000
+; AVX512-NEXT:    pushq %r15
+; AVX512-NEXT:    pushq %r14
+; AVX512-NEXT:    pushq %rbx
+; AVX512-NEXT:    subq $9056, %rsp # imm = 0x2360
 ; AVX512-NEXT:    vxorps %xmm0, %xmm0, %xmm0
 ; AVX512-NEXT:    vmovups %zmm0, {{[0-9]+}}(%rsp)
 ; AVX512-NEXT:    movb $1, {{[0-9]+}}(%rsp)
 ; AVX512-NEXT:    movw %dx, %ax
 ; AVX512-NEXT:    movw %ax, {{[-0-9]+}}(%r{{[sb]}}p) # 2-byte Spill
 ; AVX512-NEXT:    movw %si, %ax
 ; AVX512-NEXT:    movw %ax, {{[-0-9]+}}(%r{{[sb]}}p) # 2-byte Spill
-; AVX512-NEXT:    leaq {{[0-9]+}}(%rsp), %rax
-; AVX512-NEXT:    movq %rax, {{[-0-9]+}}(%r{{[sb]}}p) # 8-byte Spill
-; AVX512-NEXT:    leaq {{[0-9]+}}(%rsp), %rax
-; AVX512-NEXT:    movq %rax, {{[-0-9]+}}(%r{{[sb]}}p) # 8-byte Spill
-; AVX512-NEXT:    leaq {{[0-9]+}}(%rsp), %rax
-; AVX512-NEXT:    movq %rax, {{[-0-9]+}}(%r{{[sb]}}p) # 8-byte Spill
-; AVX512-NEXT:    leaq {{[0-9]+}}(%rsp), %rax
-; AVX512-NEXT:    movq %rax, {{[-0-9]+}}(%r{{[sb]}}p) # 8-byte Spill
 ; AVX512-NEXT:    cmpl $0, %edi
 ; AVX512-NEXT:    je .LBB0_2
 ; AVX512-NEXT:  # %bb.1: # %if.then
-; AVX512-NEXT:    movw {{[-0-9]+}}(%r{{[sb]}}p), %ax # 2-byte Reload
-; AVX512-NEXT:    movw {{[-0-9]+}}(%r{{[sb]}}p), %cx # 2-byte Reload
-; AVX512-NEXT:    movq {{[-0-9]+}}(%r{{[sb]}}p), %rdx # 8-byte Reload
-; AVX512-NEXT:    movq {{[-0-9]+}}(%r{{[sb]}}p), %rdi # 8-byte Reload
-; AVX512-NEXT:    movq {{[-0-9]+}}(%r{{[sb]}}p), %r8 # 8-byte Reload
-; AVX512-NEXT:    movl $buf, %r9d
+; AVX512-NEXT:    movw {{[-0-9]+}}(%r{{[sb]}}p), %si # 2-byte Reload
+; AVX512-NEXT:    movw {{[-0-9]+}}(%r{{[sb]}}p), %r8w # 2-byte Reload
+; AVX512-NEXT:    movl $buf, %eax
+; AVX512-NEXT:    movl $32, %ecx
+; AVX512-NEXT:    movw $8, %r9w
+; AVX512-NEXT:    movb %r8b, {{[0-9]+}}(%rsp)
+; AVX512-NEXT:    movw %r9w, {{[0-9]+}}(%rsp)
+; AVX512-NEXT:    ldtilecfg {{[0-9]+}}(%rsp)
+; AVX512-NEXT:    leaq {{[0-9]+}}(%rsp), %rdi
+; AVX512-NEXT:    tileloadd (%rax,%rcx), %tmm0
+; AVX512-NEXT:    movabsq $64, %rax
+; AVX512-NEXT:    tilestored %tmm0, 3968(%rsp,%rax) # 1024-byte Folded Spill
+; AVX512-NEXT:    movl $buf, %eax
 ; AVX512-NEXT:    movl $32, %r10d
-; AVX512-NEXT:    movw $8, %si
-; AVX512-NEXT:    movb %al, {{[0-9]+}}(%rsp)
+; AVX512-NEXT:    movw $8, %dx
+; AVX512-NEXT:    movb %r8b, {{[0-9]+}}(%rsp)
+; AVX512-NEXT:    movw %si, {{[0-9]+}}(%rsp)
+; AVX512-NEXT:    movb %dl, {{[0-9]+}}(%rsp)
 ; AVX512-NEXT:    movw %si, {{[0-9]+}}(%rsp)
 ; AVX512-NEXT:    ldtilecfg {{[0-9]+}}(%rsp)
-; AVX512-NEXT:    tileloadd (%r9,%r10), %tmm0
-; AVX512-NEXT:    movl $64, %r9d
-; AVX512-NEXT:    movw $8, %si
-; AVX512-NEXT:    tilestored %tmm0, (%r8,%r9)
-; AVX512-NEXT:    movl $buf, %r8d
-; AVX512-NEXT:    movl $32, %r9d
-; AVX512-NEXT:    movw $8, %si
-; AVX512-NEXT:    movb %al, {{[0-9]+}}(%rsp)
-; AVX512-NEXT:    movw %cx, {{[0-9]+}}(%rsp)
-; AVX512-NEXT:    movb %sil, {{[0-9]+}}(%rsp)
-; AVX512-NEXT:    movw %cx, {{[0-9]+}}(%rsp)
-; AVX512-NEXT:    ldtilecfg {{[0-9]+}}(%rsp)
-; AVX512-NEXT:    tileloadd (%r8,%r9), %tmm0
-; AVX512-NEXT:    movl $64, %r8d
-; AVX512-NEXT:    movw $8, %si
-; AVX512-NEXT:    tilestored %tmm0, (%rdi,%r8)
-; AVX512-NEXT:    movl $buf, %esi
-; AVX512-NEXT:    movl $32, %edi
-; AVX512-NEXT:    tileloadd (%rsi,%rdi), %tmm0
-; AVX512-NEXT:    movl $64, %esi
-; AVX512-NEXT:    tilestored %tmm0, (%rdx,%rsi)
+; AVX512-NEXT:    leaq {{[0-9]+}}(%rsp), %rcx
+; AVX512-NEXT:    tileloadd (%rax,%r10), %tmm0
+; AVX512-NEXT:    movabsq $64, %rax
+; AVX512-NEXT:    tilestored %tmm0, 6016(%rsp,%rax) # 1024-byte Folded Spill
+; AVX512-NEXT:    movl $buf, %r10d
+; AVX512-NEXT:    movl $32, %r11d
+; AVX512-NEXT:    leaq {{[0-9]+}}(%rsp), %rax
+; AVX512-NEXT:    tileloadd (%r10,%r11), %tmm0
+; AVX512-NEXT:    movabsq $64, %r10
+; AVX512-NEXT:    tilestored %tmm0, 7040(%rsp,%r10) # 1024-byte Folded Spill
+; AVX512-NEXT:    movw %r9w, {{[-0-9]+}}(%r{{[sb]}}p) # 2-byte Spill
+; AVX512-NEXT:    movw %r8w, {{[-0-9]+}}(%r{{[sb]}}p) # 2-byte Spill
+; AVX512-NEXT:    movq %rdi, {{[-0-9]+}}(%r{{[sb]}}p) # 8-byte Spill
+; AVX512-NEXT:    movw %si, {{[-0-9]+}}(%r{{[sb]}}p) # 2-byte Spill
+; AVX512-NEXT:    movw %dx, {{[-0-9]+}}(%r{{[sb]}}p) # 2-byte Spill
+; AVX512-NEXT:    movq %rcx, {{[-0-9]+}}(%r{{[sb]}}p) # 8-byte Spill
+; AVX512-NEXT:    movq %rax, {{[-0-9]+}}(%r{{[sb]}}p) # 8-byte Spill
 ; AVX512-NEXT:    jmp .LBB0_3
 ; AVX512-NEXT:  .LBB0_2: # %if.else
-; AVX512-NEXT:    movw {{[-0-9]+}}(%r{{[sb]}}p), %ax # 2-byte Reload
-; AVX512-NEXT:    movw {{[-0-9]+}}(%r{{[sb]}}p), %cx # 2-byte Reload
-; AVX512-NEXT:    movq {{[-0-9]+}}(%r{{[sb]}}p), %rdx # 8-byte Reload
-; AVX512-NEXT:    movq {{[-0-9]+}}(%r{{[sb]}}p), %rdi # 8-byte Reload
-; AVX512-NEXT:    movq {{[-0-9]+}}(%r{{[sb]}}p), %r8 # 8-byte Reload
-; AVX512-NEXT:    movl $buf2, %r9d
+; AVX512-NEXT:    movw {{[-0-9]+}}(%r{{[sb]}}p), %si # 2-byte Reload
+; AVX512-NEXT:    movw {{[-0-9]+}}(%r{{[sb]}}p), %r8w # 2-byte Reload
+; AVX512-NEXT:    movl $buf2, %eax
+; AVX512-NEXT:    movl $32, %ecx
+; AVX512-NEXT:    movw $8, %r9w
+; AVX512-NEXT:    movb %r8b, {{[0-9]+}}(%rsp)
+; AVX512-NEXT:    movw %r9w, {{[0-9]+}}(%rsp)
+; AVX512-NEXT:    ldtilecfg {{[0-9]+}}(%rsp)
+; AVX512-NEXT:    leaq {{[0-9]+}}(%rsp), %rdi
+; AVX512-NEXT:    tileloadd (%rax,%rcx), %tmm0
+; AVX512-NEXT:    movabsq $64, %rax
+; AVX512-NEXT:    tilestored %tmm0, 896(%rsp,%rax) # 1024-byte Folded Spill
+; AVX512-NEXT:    movl $buf2, %eax
 ; AVX512-NEXT:    movl $32, %r10d
-; AVX512-NEXT:    movw $8, %si
-; AVX512-NEXT:    movb %al, {{[0-9]+}}(%rsp)
+; AVX512-NEXT:    movw $8, %dx
+; AVX512-NEXT:    movb %r8b, {{[0-9]+}}(%rsp)
+; AVX512-NEXT:    movw %si, {{[0-9]+}}(%rsp)
+; AVX512-NEXT:    movb %dl, {{[0-9]+}}(%rsp)
 ; AVX512-NEXT:    movw %si, {{[0-9]+}}(%rsp)
 ; AVX512-NEXT:    ldtilecfg {{[0-9]+}}(%rsp)
-; AVX512-NEXT:    tileloadd (%r9,%r10), %tmm0
-; AVX512-NEXT:    movl $64, %r9d
-; AVX512-NEXT:    movw $8, %si
-; AVX512-NEXT:    tilestored %tmm0, (%r8,%r9)
-; AVX512-NEXT:    movl $buf2, %r8d
-; AVX512-NEXT:    movl $32, %r9d
-; AVX512-NEXT:    movw $8, %si
-; AVX512-NEXT:    movb %al, {{[0-9]+}}(%rsp)
-; AVX512-NEXT:    movw %cx, {{[0-9]+}}(%rsp)
-; AVX512-NEXT:    movb %sil, {{[0-9]+}}(%rsp)
-; AVX512-NEXT:    movw %cx, {{[0-9]+}}(%rsp)
-; AVX512-NEXT:    ldtilecfg {{[0-9]+}}(%rsp)
-; AVX512-NEXT:    tileloadd (%r8,%r9), %tmm0
-; AVX512-NEXT:    movl $64, %r8d
-; AVX512-NEXT:    movw $8, %si
-; AVX512-NEXT:    tilestored %tmm0, (%rdi,%r8)
-; AVX512-NEXT:    movl $buf2, %esi
-; AVX512-NEXT:    movl $32, %edi
-; AVX512-NEXT:    tileloadd (%rsi,%rdi), %tmm0
-; AVX512-NEXT:    movl $64, %esi
-; AVX512-NEXT:    tilestored %tmm0, (%rdx,%rsi)
+; AVX512-NEXT:    leaq {{[0-9]+}}(%rsp), %rcx
+; AVX512-NEXT:    tileloadd (%rax,%r10), %tmm0
+; AVX512-NEXT:    movabsq $64, %rax
+; AVX512-NEXT:    tilestored %tmm0, 1920(%rsp,%rax) # 1024-byte Folded Spill
+; AVX512-NEXT:    movl $buf2, %r10d
+; AVX512-NEXT:    movl $32, %r11d
+; AVX512-NEXT:    leaq {{[0-9]+}}(%rsp), %rax
+; AVX512-NEXT:    tileloadd (%r10,%r11), %tmm0
+; AVX512-NEXT:    movabsq $64, %r10
+; AVX512-NEXT:    tilestored %tmm0, 2944(%rsp,%r10) # 1024-byte Folded Spill
+; AVX512-NEXT:    movw %r9w, {{[-0-9]+}}(%r{{[sb]}}p) # 2-byte Spill
+; AVX512-NEXT:    movw %r8w, {{[-0-9]+}}(%r{{[sb]}}p) # 2-byte Spill
+; AVX512-NEXT:    movq %rdi, {{[-0-9]+}}(%r{{[sb]}}p) # 8-byte Spill
+; AVX512-NEXT:    movw %si, {{[-0-9]+}}(%r{{[sb]}}p) # 2-byte Spill
+; AVX512-NEXT:    movw %dx, {{[-0-9]+}}(%r{{[sb]}}p) # 2-byte Spill
+; AVX512-NEXT:    movq %rcx, {{[-0-9]+}}(%r{{[sb]}}p) # 8-byte Spill
+; AVX512-NEXT:    movq %rax, {{[-0-9]+}}(%r{{[sb]}}p) # 8-byte Spill
 ; AVX512-NEXT:  .LBB0_3: # %if.end
 ; AVX512-NEXT:    movw {{[-0-9]+}}(%r{{[sb]}}p), %ax # 2-byte Reload
 ; AVX512-NEXT:    movw {{[-0-9]+}}(%r{{[sb]}}p), %cx # 2-byte Reload
-; AVX512-NEXT:    movq {{[-0-9]+}}(%r{{[sb]}}p), %rdx # 8-byte Reload
-; AVX512-NEXT:    movq {{[-0-9]+}}(%r{{[sb]}}p), %rsi # 8-byte Reload
-; AVX512-NEXT:    movq {{[-0-9]+}}(%r{{[sb]}}p), %r9 # 8-byte Reload
-; AVX512-NEXT:    movq {{[-0-9]+}}(%r{{[sb]}}p), %r8 # 8-byte Reload
-; AVX512-NEXT:    movl $64, %r10d
-; AVX512-NEXT:    movw $8, %di
-; AVX512-NEXT:    movb %al, {{[0-9]+}}(%rsp)
-; AVX512-NEXT:    movw %di, {{[0-9]+}}(%rsp)
-; AVX512-NEXT:    ldtilecfg {{[0-9]+}}(%rsp)
-; AVX512-NEXT:    tileloadd (%r8,%r10), %tmm0
-; AVX512-NEXT:    movabsq $64, %r8
-; AVX512-NEXT:    tilestored %tmm0, 1024(%rsp,%r8) # 1024-byte Folded Spill
-; AVX512-NEXT:    movl $64, %r10d
-; AVX512-NEXT:    movw $8, %r8w
-; AVX512-NEXT:    movb %al, {{[0-9]+}}(%rsp)
-; AVX512-NEXT:    movw %cx, {{[0-9]+}}(%rsp)
-; AVX512-NEXT:    movb %al, {{[0-9]+}}(%rsp)
-; AVX512-NEXT:    movw %cx, {{[0-9]+}}(%rsp)
-; AVX512-NEXT:    movb %al, {{[0-9]+}}(%rsp)
-; AVX512-NEXT:    movw %di, {{[0-9]+}}(%rsp)
+; AVX512-NEXT:    movw {{[-0-9]+}}(%r{{[sb]}}p), %bx # 2-byte Reload
+; AVX512-NEXT:    movw {{[-0-9]+}}(%r{{[sb]}}p), %r11w # 2-byte Reload
+; AVX512-NEXT:    movq {{[-0-9]+}}(%r{{[sb]}}p), %r14 # 8-byte Reload
+; AVX512-NEXT:    movw {{[-0-9]+}}(%r{{[sb]}}p), %r9w # 2-byte Reload
+; AVX512-NEXT:    movw {{[-0-9]+}}(%r{{[sb]}}p), %r8w # 2-byte Reload
+; AVX512-NEXT:    movq {{[-0-9]+}}(%r{{[sb]}}p), %r10 # 8-byte Reload
+; AVX512-NEXT:    movq {{[-0-9]+}}(%r{{[sb]}}p), %rdi # 8-byte Reload
+; AVX512-NEXT:    movw %r11w, %dx
+; AVX512-NEXT:    movw %r9w, %si
 ; AVX512-NEXT:    movb %al, {{[0-9]+}}(%rsp)
 ; AVX512-NEXT:    movw %cx, {{[0-9]+}}(%rsp)
+; AVX512-NEXT:    movb %dl, {{[0-9]+}}(%rsp)
+; AVX512-NEXT:    movw %si, {{[0-9]+}}(%rsp)
 ; AVX512-NEXT:    movb %r8b, {{[0-9]+}}(%rsp)
-; AVX512-NEXT:    movw %cx, {{[0-9]+}}(%rsp)
-; AVX512-NEXT:    ldtilecfg {{[0-9]+}}(%rsp)
-; AVX512-NEXT:    tileloadd (%r9,%r10), %tmm2
-; AVX512-NEXT:    movl $64, %r8d
-; AVX512-NEXT:    tileloadd (%rsi,%r8), %tmm0
-; AVX512-NEXT:    movw $8, %si
+; AVX512-NEXT:    movw %r9w, {{[0-9]+}}(%rsp)
+; AVX512-NEXT:    movb %r11b, {{[0-9]+}}(%rsp)
+; AVX512-NEXT:    movw %bx, {{[0-9]+}}(%rsp)
+; AVX512-NEXT:    ldtilecfg {{[0-9]+}}(%rsp)
+; AVX512-NEXT:    movabsq $64, %r15
+; AVX512-NEXT:    tileloadd (%r14,%r15), %tmm1
+; AVX512-NEXT:    movabsq $64, %r11
+; AVX512-NEXT:    tileloadd (%r10,%r11), %tmm2
 ; AVX512-NEXT:    movabsq $64, %r8
-; AVX512-NEXT:    tileloadd 1024(%rsp,%r8), %tmm1 # 1024-byte Folded Reload
+; AVX512-NEXT:    tileloadd (%rdi,%r8), %tmm0
+; AVX512-NEXT:    movw $8, %dx
 ; AVX512-NEXT:    tdpbssd %tmm2, %tmm1, %tmm0
-; AVX512-NEXT:    movl $64, %esi
-; AVX512-NEXT:    tilestored %tmm0, (%rdx,%rsi)
-; AVX512-NEXT:    movl $64, %esi
-; AVX512-NEXT:    tileloadd (%rdx,%rsi), %tmm0
 ; AVX512-NEXT:    movl $buf, %edx
 ; AVX512-NEXT:    movl $32, %esi
 ; AVX512-NEXT:    tilestored %tmm0, (%rdx,%rsi)
-; AVX512-NEXT:    movq %rbp, %rsp
-; AVX512-NEXT:    popq %rbp
+; AVX512-NEXT:    addq $9056, %rsp # imm = 0x2360
+; AVX512-NEXT:    popq %rbx
+; AVX512-NEXT:    popq %r14
+; AVX512-NEXT:    popq %r15
 ; AVX512-NEXT:    tilerelease
 ; AVX512-NEXT:    vzeroupper
 ; AVX512-NEXT:    retq
 entry:
   %tobool.not = icmp eq i32 %cond, 0
   br i1 %tobool.not, label %if.else, label %if.then
 
 if.then:                                          ; preds = %entry

llvm/test/CodeGen/X86/AMX/amx-zero-config.ll

 ; SSE2-NEXT:    tilezero %tmm0
 ; SSE2-NEXT:    movl $1024, %edx # imm = 0x400
 ; SSE2-NEXT:    tilestored %tmm0, (%rdi,%rdx)
 ; SSE2-NEXT:    tilerelease
 ; SSE2-NEXT:    retq
 ;
 ; AVX512-O0-LABEL: foo:
 ; AVX512-O0:       # %bb.0: # %entry
-; AVX512-O0-NEXT:    pushq %rbp
-; AVX512-O0-NEXT:    movq %rsp, %rbp
-; AVX512-O0-NEXT:    andq $-1024, %rsp # imm = 0xFC00
-; AVX512-O0-NEXT:    subq $3072, %rsp # imm = 0xC00
 ; AVX512-O0-NEXT:    vxorps %xmm0, %xmm0, %xmm0
-; AVX512-O0-NEXT:    vmovups %zmm0, {{[0-9]+}}(%rsp)
-; AVX512-O0-NEXT:    movb $1, {{[0-9]+}}(%rsp)
-; AVX512-O0-NEXT:    leaq {{[0-9]+}}(%rsp), %rdx
+; AVX512-O0-NEXT:    vmovups %zmm0, -{{[0-9]+}}(%rsp)
+; AVX512-O0-NEXT:    movb $1, -{{[0-9]+}}(%rsp)
 ; AVX512-O0-NEXT:    movw $32, %cx
 ; AVX512-O0-NEXT:    movw $8, %ax
-; AVX512-O0-NEXT:    movb %al, {{[0-9]+}}(%rsp)
-; AVX512-O0-NEXT:    movw %cx, {{[0-9]+}}(%rsp)
-; AVX512-O0-NEXT:    ldtilecfg {{[0-9]+}}(%rsp)
+; AVX512-O0-NEXT:    movb %al, -{{[0-9]+}}(%rsp)
+; AVX512-O0-NEXT:    movw %cx, -{{[0-9]+}}(%rsp)
+; AVX512-O0-NEXT:    ldtilecfg -{{[0-9]+}}(%rsp)
 ; AVX512-O0-NEXT:    tilezero %tmm0
-; AVX512-O0-NEXT:    movl $64, %esi
-; AVX512-O0-NEXT:    movw $32, %cx
-; AVX512-O0-NEXT:    movw $8, %ax
-; AVX512-O0-NEXT:    tilestored %tmm0, (%rdx,%rsi)
-; AVX512-O0-NEXT:    movl $64, %esi
-; AVX512-O0-NEXT:    movw $32, %cx
-; AVX512-O0-NEXT:    movw $8, %ax
-; AVX512-O0-NEXT:    movb %al, {{[0-9]+}}(%rsp)
-; AVX512-O0-NEXT:    movw %cx, {{[0-9]+}}(%rsp)
-; AVX512-O0-NEXT:    ldtilecfg {{[0-9]+}}(%rsp)
-; AVX512-O0-NEXT:    tileloadd (%rdx,%rsi), %tmm0
 ; AVX512-O0-NEXT:    movl $1024, %edx # imm = 0x400
 ; AVX512-O0-NEXT:    movw $32, %cx
 ; AVX512-O0-NEXT:    movw $8, %ax
 ; AVX512-O0-NEXT:    tilestored %tmm0, (%rdi,%rdx)
-; AVX512-O0-NEXT:    movq %rbp, %rsp
-; AVX512-O0-NEXT:    popq %rbp
 ; AVX512-O0-NEXT:    tilerelease
 ; AVX512-O0-NEXT:    vzeroupper
 ; AVX512-O0-NEXT:    retq
 ;
 ; AVX2-O0-LABEL: foo:
 ; AVX2-O0:       # %bb.0: # %entry
-; AVX2-O0-NEXT:    pushq %rbp
-; AVX2-O0-NEXT:    movq %rsp, %rbp
-; AVX2-O0-NEXT:    andq $-1024, %rsp # imm = 0xFC00
-; AVX2-O0-NEXT:    subq $3072, %rsp # imm = 0xC00
 ; AVX2-O0-NEXT:    vxorps %xmm0, %xmm0, %xmm0
-; AVX2-O0-NEXT:    vmovups %ymm0, {{[0-9]+}}(%rsp)
-; AVX2-O0-NEXT:    vmovups %ymm0, {{[0-9]+}}(%rsp)
-; AVX2-O0-NEXT:    movb $1, {{[0-9]+}}(%rsp)
-; AVX2-O0-NEXT:    leaq {{[0-9]+}}(%rsp), %rdx
+; AVX2-O0-NEXT:    vmovups %ymm0, -{{[0-9]+}}(%rsp)
+; AVX2-O0-NEXT:    vmovups %ymm0, -{{[0-9]+}}(%rsp)
+; AVX2-O0-NEXT:    movb $1, -{{[0-9]+}}(%rsp)
 ; AVX2-O0-NEXT:    movw $32, %cx
 ; AVX2-O0-NEXT:    movw $8, %ax
-; AVX2-O0-NEXT:    movb %al, {{[0-9]+}}(%rsp)
-; AVX2-O0-NEXT:    movw %cx, {{[0-9]+}}(%rsp)
-; AVX2-O0-NEXT:    ldtilecfg {{[0-9]+}}(%rsp)
+; AVX2-O0-NEXT:    movb %al, -{{[0-9]+}}(%rsp)
+; AVX2-O0-NEXT:    movw %cx, -{{[0-9]+}}(%rsp)
+; AVX2-O0-NEXT:    ldtilecfg -{{[0-9]+}}(%rsp)
 ; AVX2-O0-NEXT:    tilezero %tmm0
-; AVX2-O0-NEXT:    movl $64, %esi
-; AVX2-O0-NEXT:    movw $32, %cx
-; AVX2-O0-NEXT:    movw $8, %ax
-; AVX2-O0-NEXT:    tilestored %tmm0, (%rdx,%rsi)
-; AVX2-O0-NEXT:    movl $64, %esi
-; AVX2-O0-NEXT:    movw $32, %cx
-; AVX2-O0-NEXT:    movw $8, %ax
-; AVX2-O0-NEXT:    movb %al, {{[0-9]+}}(%rsp)
-; AVX2-O0-NEXT:    movw %cx, {{[0-9]+}}(%rsp)
-; AVX2-O0-NEXT:    ldtilecfg {{[0-9]+}}(%rsp)
-; AVX2-O0-NEXT:    tileloadd (%rdx,%rsi), %tmm0
 ; AVX2-O0-NEXT:    movl $1024, %edx # imm = 0x400
 ; AVX2-O0-NEXT:    movw $32, %cx
 ; AVX2-O0-NEXT:    movw $8, %ax
 ; AVX2-O0-NEXT:    tilestored %tmm0, (%rdi,%rdx)
-; AVX2-O0-NEXT:    movq %rbp, %rsp
-; AVX2-O0-NEXT:    popq %rbp
 ; AVX2-O0-NEXT:    tilerelease
 ; AVX2-O0-NEXT:    vzeroupper
 ; AVX2-O0-NEXT:    retq
 ;
 ; SSE2-O0-LABEL: foo:
 ; SSE2-O0:       # %bb.0: # %entry
-; SSE2-O0-NEXT:    pushq %rbp
-; SSE2-O0-NEXT:    movq %rsp, %rbp
-; SSE2-O0-NEXT:    andq $-1024, %rsp # imm = 0xFC00
-; SSE2-O0-NEXT:    subq $3072, %rsp # imm = 0xC00
 ; SSE2-O0-NEXT:    xorps %xmm0, %xmm0
-; SSE2-O0-NEXT:    movups %xmm0, {{[0-9]+}}(%rsp)
-; SSE2-O0-NEXT:    movups %xmm0, {{[0-9]+}}(%rsp)
-; SSE2-O0-NEXT:    movups %xmm0, {{[0-9]+}}(%rsp)
-; SSE2-O0-NEXT:    movups %xmm0, {{[0-9]+}}(%rsp)
-; SSE2-O0-NEXT:    movb $1, {{[0-9]+}}(%rsp)
-; SSE2-O0-NEXT:    leaq {{[0-9]+}}(%rsp), %rdx
+; SSE2-O0-NEXT:    movups %xmm0, -{{[0-9]+}}(%rsp)
+; SSE2-O0-NEXT:    movups %xmm0, -{{[0-9]+}}(%rsp)
+; SSE2-O0-NEXT:    movups %xmm0, -{{[0-9]+}}(%rsp)
+; SSE2-O0-NEXT:    movups %xmm0, -{{[0-9]+}}(%rsp)
+; SSE2-O0-NEXT:    movb $1, -{{[0-9]+}}(%rsp)
 ; SSE2-O0-NEXT:    movw $32, %cx
 ; SSE2-O0-NEXT:    movw $8, %ax
-; SSE2-O0-NEXT:    movb %al, {{[0-9]+}}(%rsp)
-; SSE2-O0-NEXT:    movw %cx, {{[0-9]+}}(%rsp)
-; SSE2-O0-NEXT:    ldtilecfg {{[0-9]+}}(%rsp)
+; SSE2-O0-NEXT:    movb %al, -{{[0-9]+}}(%rsp)
+; SSE2-O0-NEXT:    movw %cx, -{{[0-9]+}}(%rsp)
+; SSE2-O0-NEXT:    ldtilecfg -{{[0-9]+}}(%rsp)
 ; SSE2-O0-NEXT:    tilezero %tmm0
-; SSE2-O0-NEXT:    movl $64, %esi
-; SSE2-O0-NEXT:    movw $32, %cx
-; SSE2-O0-NEXT:    movw $8, %ax
-; SSE2-O0-NEXT:    tilestored %tmm0, (%rdx,%rsi)
-; SSE2-O0-NEXT:    movl $64, %esi
-; SSE2-O0-NEXT:    movw $32, %cx
-; SSE2-O0-NEXT:    movw $8, %ax
-; SSE2-O0-NEXT:    movb %al, {{[0-9]+}}(%rsp)
-; SSE2-O0-NEXT:    movw %cx, {{[0-9]+}}(%rsp)
-; SSE2-O0-NEXT:    ldtilecfg {{[0-9]+}}(%rsp)
-; SSE2-O0-NEXT:    tileloadd (%rdx,%rsi), %tmm0
 ; SSE2-O0-NEXT:    movl $1024, %edx # imm = 0x400
 ; SSE2-O0-NEXT:    movw $32, %cx
 ; SSE2-O0-NEXT:    movw $8, %ax
 ; SSE2-O0-NEXT:    tilestored %tmm0, (%rdi,%rdx)
-; SSE2-O0-NEXT:    movq %rbp, %rsp
-; SSE2-O0-NEXT:    popq %rbp
 ; SSE2-O0-NEXT:    tilerelease
 ; SSE2-O0-NEXT:    retq
 entry:
   %t = call x86_amx @llvm.x86.tilezero.internal(i16 8, i16 32)
   call void @llvm.x86.tilestored64.internal(i16 8, i16 32, i8* %buf, i64 1024, x86_amx %t)
   ret void
 }
 
 declare x86_amx @llvm.x86.tilezero.internal(i16, i16)
 declare void @llvm.x86.tilestored64.internal(i16, i16, i8*, i64, x86_amx)