[X86] Fix tile config register spill issue.

[pengfei]

Summary

This is an optimized approach for D94155.

Previous code build the model that tile config register is the user of
each AMX instruction. There is a problem for the tile config register
spill. When across function, the ldtilecfg instruction may be inserted
on each AMX instruction which use tile config register. This cause all
tile data register clobber.

To fix this issue, we remove the model of tile config register. We
analyze the AMX instructions between one call to another. We will insert
ldtilecfg after the first call if we find any AMX instructions.

We also optimized the inserting of tilerelease by moving it just after
the last AMX instruction of each branch. We will insert it to its
successors if the BB contains the last AMX instruction is in a loop.

[pengfei]

Fix Lint warnings.

[xiangzhangllvm]

llvm/lib/Target/X86/X86PreTileConfig.cpp

321

The Algorithm complexity is O(N*N)

[xiangzhangllvm]

llvm/lib/Target/X86/X86PreTileConfig.cpp

321

The Algorithm complexity is O(N*N)

321

The Algorithm complexity is O(N*N)

I think we can do this in 1 DFS.

[xiangzhangllvm]

llvm/lib/Target/X86/X86PreTileConfig.cpp

281

If the call at the bottom of the BB, we may need insert tile_ld_config too.

[pengfei]

llvm/lib/Target/X86/X86PreTileConfig.cpp

281

Good catch! Thanks.

321

The Algorithm complexity is O(N*N)

I think we can do this in 1 DFS.

The Algorithm complexity is O(N) ideally. E.g. Each BB has at least one call. If not, the BB (and its successor) might be iterated more than once. I can improve it future.

[LuoYuanke]

llvm/lib/Target/X86/X86PreTileConfig.cpp

269

WorkList and LastAMXSet may be a class member to avoid parameter passing.

[LuoYuanke]

llvm/lib/Target/X86/X86PreTileConfig.cpp

279

This is not clear to me. In below example, do we only insert tilerelease in amx1? Does call2 have chance to be iterated? We may add more test cases for different scenario.

       call0
    /        \
amx1       amx2
  call1       call2

[pengfei]

1. Address review comments from Xiang and Yuanke.
2. Optimize the algorithm to a single DFS.
3. Fix bugs when counting the place for tilerelease.

[pengfei]

llvm/lib/Target/X86/X86PreTileConfig.cpp

279

The previous code doesn't have these problem. Anyway, I changed the algorithm. It should be more clear now. Could you have a new review?

[pengfei]

Remove an unnecessary ++MII;

[pengfei]

A minor refactor.

[pengfei]

This patch reduces algorithm complexity of function findNeedInsertBB, while increases a bit code complexity compared with previous one.

[pengfei]

A minor change.

[LuoYuanke]

llvm/lib/Target/X86/X86PreTileConfig.cpp

354

It seems a lot of change in the patch. If you drop tilerelease, will the patch be simpler. I would like solve tile config issue first and move forward.

[pengfei]

llvm/lib/Target/X86/X86PreTileConfig.cpp

354

Yes, it will be. But the former BFS + DFS algorithm is still complex than the single DFS one, which means we still need major changes.
I would like to keep the tilerelease handling here. I didn't find bug of it for now and the calculating math is straightforward. I can show you the 4 status in below:

    ______________________________
S3 |begin ... AMX ... call ... end|
    ------------------------------
    ______________________________
S2 |begin ... call ... AMX ... end|
    ------------------------------
           _____________
S1        |begin ... end|
           -------------
          _/_       _\_
   max(S(|BB0|), S(|BB1|), ...) >= 1
          ---       ---

S0 is similar with S1 where all its successors' status are all 0. It doesn't matter if there's loop in the graph. Because when calculating the status, we assign 0 to current BB at first (if there's no AMX in it). Then we will iterate all its successors. It turns to 1 if one of other successors is not 0, and keeps 0 if all other successors are 0.
We insert tilerelease after the last AMX of current BB only if all its successors' status are 0. Otherwise, we will iterate all its first level successors and insert tilerelease at the begining of BB whoes status is 0. We also need to consider the BB whoes status is 1. We will recursively find their successors and insert tilerelease for BBs whoes status is 0.
All BB will be calculated for tilerelease inserting case once at most, since we record them in BBNeedInsertVisited. There's no big algorithm complexity increasement here.

[LuoYuanke]

llvm/lib/Target/X86/X86PreTileConfig.cpp

334

May add more test case for tilerelease. We may insert multi tilerelease in different block, better to have test case to cover that.

351

The 4 state may be a enum to be more readable.

[pengfei]

Adress Yuanke's comments.

[LuoYuanke]

llvm/lib/Target/X86/X86PreTileConfig.cpp

364

This confused me. This looks check call after amx. Should we instead check if there is amx instruction after call to determine if ldtilecfg is needed?

[pengfei]

Address Yuanke's comment. Moving the insertion to check AMX block.

[pengfei]

llvm/lib/Target/X86/X86PreTileConfig.cpp

364

I moved it to AMX checking block now.
The previous code try to reduce multi insertion by puting it in the call checking block.

[pengfei]

Fix a missing case in function updateBBStatus.

[pengfei]

Small refactor: Change recursive DFS to iterative DFS when insert tilerelease.

[LuoYuanke]

llvm/lib/Target/X86/X86PreTileConfig.cpp

264

I'm not sure there is only one exit.

360

Not sure about this condition. If there is no call in its direct successor, but there is amx before call in its successor's successor.

      call
    /      \
BB1      BB2
  |
 amx
 call

367

I didn't figure out this function yet. @xiangzhangllvm, would you help to review this function?

[xiangzhangllvm]

llvm/lib/Target/X86/X86PreTileConfig.cpp

260

We just need to collect insert point after call which appended by AMX instruction,
I suggest a clear O(N) logic in 1 DFS:

ReloadTileConfig (MBB ) {  // stark from tilecfg MBB
     For each MI in MBB {
        CallNum = 0;
        if MI is Call 
            push Call into CallStack
            ++CallNum;
        if MI is AMX
            Mark the Call in Top of CallStack is NeedInsert;
      }
      For MBBS in MBB->Succs
          ReloadTileConfig (MBB )
      While CallNum--
          if CallStack .top is marked NeedInsert
             Save CallStack .top into NeedInsertPosSet //NeedInsertPosSet will collect all the insert points
          CallStack.pop 
}

[pengfei]

Fix a bug when calculating indirect successors.

[pengfei]

Reduce BB status to 3 cases, simplify code based on it.

[pengfei]

A minor refactor.

[pengfei]

Change enum names to make them more accurate.

[xiangzhangllvm]

llvm/lib/Target/X86/X86PreTileConfig.cpp

283

I am not sure the order of iterate BBVisitedInfo is right.
In map it may use "<" to sort its elements.

292

What does HasBeforeCallAMX really mean, why set it once BB has AMX ?

298

Not sure directly continue is right, if the status required by the visited predecessors.

[pengfei]

llvm/lib/Target/X86/X86PreTileConfig.cpp

283

The order doesn't matter. Just make sure all BBs are iterated.

292

It will be updated soon in line 293. We can assume it alway has AMX before call first. Then together with no AMX case, we check if there's call before AMX. If there is, we update it to HasAfterCallAMX.

298

We don't need to care about the BB which is not found when we do searching from top to bottom.
Becaused these BBs are never to be a successor of any AMX instruction.
We might meet them sometimes since we are searching from bottom to top. Ignoring them is safe.

[pengfei]

Remove the code of inserting tilerelease for easy review.

[xiangzhangllvm]

llvm/lib/Target/X86/X86PreTileConfig.cpp

229

LGTM for the logic.

235

Can use a better name for HasCallBeforeAMX? it is very easy to mix with HasAfterCallAMX when I read here

283

The logic here is complex, better to add clear comment to show how we update the Status for BBs.

[pengfei]

Address Xiang's comments.

[LuoYuanke]

llvm/lib/Target/X86/X86PreTileConfig.cpp

230

Better to have a diagram comments to illustrate the state.

235

Agree. It is easy to mix with HasBeforeCallAMX.

271

!WorkList.empty()?

283

WorkList should be clear before this loop?

286

(!WorkList.empty()) seems more readable.

297

Not quite understand this. Isn't all BBs added at line 275?

302

Don't understand this. Can you add a diagram to illustrate it?

[pengfei]

llvm/lib/Target/X86/X86PreTileConfig.cpp

229

Thanks for the review.

235

Should IsCallBeforeAMX better?

[LuoYuanke]

llvm/lib/Target/X86/X86PreTileConfig.cpp

252

What about multi call and multi AMX in one BB? For below example, should we insert ldtilecfg after call1?

call1
AMX1
call2
AMX2

[LuoYuanke]

llvm/lib/Target/X86/X86PreTileConfig.cpp

235

Maybe HasCallBeforeAMXInBB.

[pengfei]

llvm/lib/Target/X86/X86PreTileConfig.cpp

235

HasCallBeforeAMXInBB is still easy mix up with HasAfterCallAMX, besides, we need to mark it true even there's no AMX in the BB.

252

Yes, we insert for call1 and call2 in line 247.
We only add the BB to BBUnsolved if there's no AMX after the last call.

297

No, we only added successors since ldtilecfg BB. BB0, BB1 and BB2 in below graph won't be added.

  BB0   ldtilecfg
   |      / \
  BB1   BB4 BB5
  / \   /   /
BB2  BB3  BB6
   \  |  /
     BB7

[LuoYuanke]

llvm/lib/Target/X86/X86PreTileConfig.cpp

297

There is no AMX instruction in BB0, BB1, BB2, so we don't care about BB0, BB1, BB2?

[pengfei]

llvm/lib/Target/X86/X86PreTileConfig.cpp

297

Yes. BB ldtilecfg dominates all BB that have AMX. So, there's no AMX in BB0, BB1 and BB2. In fact, BB3 and BB7 don't have either.
We calculate the status here which are used by the BB in BBUnsolved, which will never use the status of BB0~2. Nevertheless, we know their stutus are always NoAMXAtAll.

[pengfei]

Address Yuanke's comments.

[pengfei]

llvm/lib/Target/X86/X86PreTileConfig.cpp

283

Not necessary. But adding it seems more robust.

302

Add some comments on line 290.

[LuoYuanke]

llvm/lib/Target/X86/X86PreTileConfig.cpp

298

State HasAfterCallAMX seems useless. I didn't find we check this state anywhere.

306

Why if BBVisitedInfo[*I].HasAMX is true, we don't update it? I think when BBVisitedInfo[*I].MaxSucc changes, we need update its predecessor.

[pengfei]

llvm/lib/Target/X86/X86PreTileConfig.cpp

298

Yes, it was used for inserting tilerelease. Now we just need 2 status, I simplified code base it.

306

It is just an optimization.
If *I has AMX, the status it contributes to its predecessors won't be changed with its MaxSucc, e.g.
When we have a BB which has AMX and its MaxSucc is S0, it updates its predecessors with S2 in the first time. So next time, when its MaxSucc is changed to S1 by its successor, it still updates its predecessors with S2. This is not necessary, so it is skipped.
Anyway, since we are using 2 status now, this code is also been simplified.

[pengfei]

Reduce one unused status. Simplify code based on it.

[LuoYuanke]

LGTM. Thank you!

[xiangzhangllvm]

+1

[xiangzhangllvm]

llvm/lib/Target/X86/X86PreTileConfig.cpp

229

Not try to replace this patch, I update the O(N) algorithm in https://reviews.llvm.org/D95508 just for interest.
In Non-loop case, it cost O(N), in loop case, it cost O(2N)

[pengfei]

llvm/lib/Target/X86/X86PreTileConfig.cpp

229

So it is related on the number of loops? But in the worst case, a graph may have O(N!) loops? The worst complexity is terrible.

[xiangzhangllvm]

llvm/lib/Target/X86/X86PreTileConfig.cpp

229

So it is not related on the number of loops, in loop case the max cost is O(2N), (we see O(2N) == O(N) in program)

llvm/lib/Target/X86/X86ExpandPseudo.cpp

case X86::MWAITX_SAVE_RBX: {
     TII->copyPhysReg(MBB, MBBI, DL, X86::RBX, SaveRbx, /*SrcIsKill*/ true);
     // Delete the pseudo.
     MBBI->eraseFromParent();
     return true;
   }
   case TargetOpcode::ICALL_BRANCH_FUNNEL:
     ExpandICallBranchFunnel(&MBB, MBBI);
     return true;
-  case X86::PLDTILECFG: {
-    MI.RemoveOperand(0);
-    MI.setDesc(TII->get(X86::LDTILECFG));
-    return true;
-  }
-  case X86::PSTTILECFG: {
-    MI.RemoveOperand(MI.getNumOperands() - 1); // Remove $tmmcfg
-    MI.setDesc(TII->get(X86::STTILECFG));
-    return true;
-  }
   case X86::PTILELOADDV: {
-    MI.RemoveOperand(8); // Remove $tmmcfg
     for (unsigned i = 2; i > 0; --i)
       MI.RemoveOperand(i);
     MI.setDesc(TII->get(X86::TILELOADD));
     return true;
   }
   case X86::PTDPBSSDV: {
-    MI.RemoveOperand(7); // Remove $tmmcfg
     MI.untieRegOperand(4);
     for (unsigned i = 3; i > 0; --i)
       MI.RemoveOperand(i);
     MI.setDesc(TII->get(X86::TDPBSSD));
     MI.tieOperands(0, 1);
     return true;
   }
   case X86::PTILESTOREDV: {
-    MI.RemoveOperand(8); // Remove $tmmcfg
     for (int i = 1; i >= 0; --i)
       MI.RemoveOperand(i);
     MI.setDesc(TII->get(X86::TILESTORED));
     return true;
   }
   case X86::PTILEZEROV: {
-    for (int i = 3; i > 0; --i) // Remove row, col, $tmmcfg
+    for (int i = 2; i > 0; --i) // Remove row, col
       MI.RemoveOperand(i);
     MI.setDesc(TII->get(X86::TILEZERO));
     return true;
   }
   }
   llvm_unreachable("Previous switch has a fallthrough?");
 }
 

llvm/lib/Target/X86/X86FrameLowering.cpp

if (Offset) {
       // Check for possible merge with preceding ADD instruction.
       Offset += mergeSPUpdates(MBB, Terminator, true);
       emitSPUpdate(MBB, Terminator, DL, Offset, /*InEpilogue=*/true);
     }
   }
 
   // Emit tilerelease for AMX kernel.
   const MachineRegisterInfo &MRI = MF.getRegInfo();
-  if (!MRI.reg_nodbg_empty(X86::TMMCFG))
+  const TargetRegisterClass *RC = TRI->getRegClass(X86::TILERegClassID);
+  for (unsigned I = 0; I < RC->getNumRegs(); I++)
+    if (!MRI.reg_nodbg_empty(X86::TMM0 + I)) {
       BuildMI(MBB, Terminator, DL, TII.get(X86::TILERELEASE));
+      break;
+    }
 }
 
 StackOffset X86FrameLowering::getFrameIndexReference(const MachineFunction &MF,
                                                      int FI,
                                                      Register &FrameReg) const {
   const MachineFrameInfo &MFI = MF.getFrameInfo();
 
   bool IsFixed = MFI.isFixedObjectIndex(FI);

llvm/lib/Target/X86/X86ISelDAGToDAG.cpp

case Intrinsic::x86_tileloadd64_internal: {
         break;
       unsigned Opc = X86::PTILELOADDV;
       // _tile_loadd_internal(row, col, buf, STRIDE)
       SDValue Base = Node->getOperand(4);
       SDValue Scale = getI8Imm(1, dl);
       SDValue Index = Node->getOperand(5);
       SDValue Disp = CurDAG->getTargetConstant(0, dl, MVT::i32);
       SDValue Segment = CurDAG->getRegister(0, MVT::i16);
-      SDValue CFG = CurDAG->getRegister(0, MVT::Untyped);
       SDValue Chain = Node->getOperand(0);
       MachineSDNode *CNode;
       SDValue Ops[] = {Node->getOperand(2),
                        Node->getOperand(3),
                        Base,
                        Scale,
                        Index,
                        Disp,
                        Segment,
-                       CFG,
                        Chain};
       CNode = CurDAG->getMachineNode(Opc, dl, {MVT::x86amx, MVT::Other}, Ops);
       ReplaceNode(Node, CNode);
       return;
     }
     case Intrinsic::x86_tdpbssd_internal: {
       if (!Subtarget->hasAMXTILE())
         break;
       SDValue Chain = Node->getOperand(0);
       unsigned Opc = X86::PTDPBSSDV;
-      SDValue CFG = CurDAG->getRegister(0, MVT::Untyped);
       SDValue Ops[] = {Node->getOperand(2),
                        Node->getOperand(3),
                        Node->getOperand(4),
                        Node->getOperand(5),
                        Node->getOperand(6),
                        Node->getOperand(7),
-                       CFG,
                        Chain};
       MachineSDNode *CNode =
           CurDAG->getMachineNode(Opc, dl, {MVT::x86amx, MVT::Other}, Ops);
       ReplaceNode(Node, CNode);
       return;
     }
     case Intrinsic::x86_tilezero_internal: {
       if (!Subtarget->hasAMXTILE())
         break;
       unsigned Opc = X86::PTILEZEROV;
       SDValue Chain = Node->getOperand(0);
-      SDValue CFG = CurDAG->getRegister(0, MVT::Untyped);
-      SDValue Ops[] = {Node->getOperand(2), Node->getOperand(3), CFG, Chain};
+      SDValue Ops[] = {Node->getOperand(2), Node->getOperand(3), Chain};
       MachineSDNode *CNode =
           CurDAG->getMachineNode(Opc, dl, {MVT::x86amx, MVT::Other}, Ops);
       ReplaceNode(Node, CNode);
       return;
     }
     }
     break;
   }
case ISD::INTRINSIC_VOID: {
     case Intrinsic::x86_tilestored64_internal: {
       unsigned Opc = X86::PTILESTOREDV;
       // _tile_stored_internal(row, col, buf, STRIDE, c)
       SDValue Base = Node->getOperand(4);
       SDValue Scale = getI8Imm(1, dl);
       SDValue Index = Node->getOperand(5);
       SDValue Disp = CurDAG->getTargetConstant(0, dl, MVT::i32);
       SDValue Segment = CurDAG->getRegister(0, MVT::i16);
-      SDValue CFG = CurDAG->getRegister(0, MVT::Untyped);
       SDValue Chain = Node->getOperand(0);
       MachineSDNode *CNode;
       SDValue Ops[] = {Node->getOperand(2),
                        Node->getOperand(3),
                        Base,
                        Scale,
                        Index,
                        Disp,
                        Segment,
                        Node->getOperand(6),
-                       CFG,
                        Chain};
       CNode = CurDAG->getMachineNode(Opc, dl, MVT::Other, Ops);
       ReplaceNode(Node, CNode);
       return;
     }
     case Intrinsic::x86_tileloadd64:
     case Intrinsic::x86_tileloaddt164:
     case Intrinsic::x86_tilestored64: {

llvm/lib/Target/X86/X86InstrAMX.td

def TILESTORED : I<0x4b, MRMDestMemFSIB, (outs),
                        (ins sibmem:$dst, TILE:$src),
                        "tilestored\t{$src, $dst|$dst, $src}", []>,
                        VEX, T8XS;
     def TILEZERO : I<0x49, MRMr0, (outs TILE:$dst), (ins),
                      "tilezero\t$dst", []>,
                      VEX, T8XD;
 
     // Pseduo instruction for RA.
-    let hasSideEffects = 1, mayLoad = 1,
-        Defs = [TMM0,TMM1,TMM2,TMM3,TMM4,TMM5,TMM6,TMM7] in
-    def PLDTILECFG : PseudoI <(outs TILECFG:$cfg), (ins opaquemem:$src), []>;
-
-    let hasSideEffects = 1, mayStore = 1 in
-    def PSTTILECFG : PseudoI<(outs), (ins opaquemem:$dst, TILECFG:$cfg), []>;
-
     def PTILELOADDV : PseudoI<(outs TILE: $dst), (ins GR16:$src1,
                                                       GR16:$src2,
-                                                      opaquemem:$src3,
-                                                      TILECFG:$cfg), []>;
+                                                      opaquemem:$src3), []>;
     def PTILESTOREDV : PseudoI<(outs), (ins GR16:$src1,
                                             GR16:$src2, opaquemem:$src3,
-                                            TILE:$src4, TILECFG:$cfg), []>;
+                                            TILE:$src4), []>;
     def PTILEZEROV : PseudoI<(outs TILE: $dst), (ins GR16:$src1,
-                                                     GR16:$src2,
-                                                     TILECFG:$cfg), []>;
+                                                     GR16:$src2), []>;
 
     let usesCustomInserter = 1 in {
       // Pseudo instructions, using immediates instead of tile registers.
       // To be translated to the actual instructions in X86ISelLowering.cpp
       def PTILELOADD : PseudoI<(outs), (ins u8imm:$src1, sibmem:$src2), []>;
       def PTILELOADDT1 : PseudoI<(outs), (ins u8imm:$src1,
                                           sibmem:$src2), []>;
       def PTILESTORED : PseudoI<(outs), (ins i8mem:$dst, u8imm:$src), []>;
let Constraints = "$src1 = $dst" in {
                       "tdpbuud\t{$src3, $src2, $dst|$dst, $src2, $src3}", []>,
                       VEX_4V, T8PS;
     }
 
     // Pseduo instruction for RA.
     let Constraints = "$src4 = $dst" in
     def PTDPBSSDV : PseudoI<(outs TILE: $dst), (ins GR16:$src1,
                             GR16:$src2, GR16:$src3, TILE:$src4,
-                            TILE:$src5, TILE:$src6, TILECFG:$cfg), []>;
+                            TILE:$src5, TILE:$src6), []>;
 
     let usesCustomInserter = 1 in {
       // Pseudo instructions, using immediates instead of tile registers.
       // To be translated to the actual instructions in X86ISelLowering.cpp
       def PTDPBSSD : PseudoI<(outs), (ins u8imm:$src1,
                              u8imm:$src2, u8imm:$src3),
                              [(int_x86_tdpbssd timm:$src1,
                                timm:$src2, timm:$src3)]>;

llvm/lib/Target/X86/X86InstrInfo.cpp

if (RC->getID() == X86::TILERegClassID) {
     Register VirtReg = RegInfo.createVirtualRegister(&X86::GR64_NOSPRegClass);
     BuildMI(MBB, MI, DebugLoc(), get(X86::MOV64ri), VirtReg).addImm(64);
     MachineInstr *NewMI =
         addFrameReference(BuildMI(MBB, MI, DebugLoc(), get(Opc)), FrameIdx)
             .addReg(SrcReg, getKillRegState(isKill));
     MachineOperand &MO = NewMI->getOperand(2);
     MO.setReg(VirtReg);
     MO.setIsKill(true);
-  } else if (RC->getID() == X86::TILECFGRegClassID) {
-    unsigned Opc = X86::PSTTILECFG;
-    addFrameReference(BuildMI(MBB, MI, DebugLoc(), get(Opc)), FrameIdx)
-        .addReg(SrcReg, getKillRegState(isKill));
   } else {
     unsigned Alignment = std::max<uint32_t>(TRI->getSpillSize(*RC), 16);
     bool isAligned =
         (Subtarget.getFrameLowering()->getStackAlign() >= Alignment) ||
         RI.canRealignStack(MF);
     unsigned Opc = getStoreRegOpcode(SrcReg, RC, isAligned, Subtarget);
     addFrameReference(BuildMI(MBB, MI, DebugLoc(), get(Opc)), FrameIdx)
         .addReg(SrcReg, getKillRegState(isKill));
if (RC->getID() == X86::TILERegClassID) {
     Register VirtReg = RegInfo.createVirtualRegister(&X86::GR64_NOSPRegClass);
     MachineInstr *NewMI =
         BuildMI(MBB, MI, DebugLoc(), get(X86::MOV64ri), VirtReg).addImm(64);
     NewMI = addFrameReference(BuildMI(MBB, MI, DebugLoc(), get(Opc), DestReg),
                               FrameIdx);
     MachineOperand &MO = NewMI->getOperand(3);
     MO.setReg(VirtReg);
     MO.setIsKill(true);
-  } else if (RC->getID() == X86::TILECFGRegClassID) {
-    unsigned Opc = X86::PLDTILECFG;
-    addFrameReference(BuildMI(MBB, MI, DebugLoc(), get(Opc), DestReg),
-                      FrameIdx);
   } else {
     const MachineFunction &MF = *MBB.getParent();
     unsigned Alignment = std::max<uint32_t>(TRI->getSpillSize(*RC), 16);
     bool isAligned =
         (Subtarget.getFrameLowering()->getStackAlign() >= Alignment) ||
         RI.canRealignStack(MF);
     unsigned Opc = getLoadRegOpcode(DestReg, RC, isAligned, Subtarget);
     addFrameReference(BuildMI(MBB, MI, DebugLoc(), get(Opc), DestReg),

llvm/lib/Target/X86/X86PreTileConfig.cpp

INITIALIZE_PASS_END(X86PreTileConfig, "tilepreconfig",
                     "Tile Register Configure", false, false)
 
 void X86PreTileConfig::getAnalysisUsage(AnalysisUsage &AU) const {
   AU.setPreservesAll();
   AU.addRequired<MachineDominatorTree>();
   MachineFunctionPass::getAnalysisUsage(AU);
 }
 
-static Register buildConfigMI(MachineBasicBlock::iterator MI, int FrameIdx,
-                              const TargetInstrInfo *TII,
-                              MachineRegisterInfo *MRI,
-                              const X86Subtarget *ST) {
+static void buildConfigMI(MachineBasicBlock::iterator MI, int FrameIdx,
+                          const TargetInstrInfo *TII, MachineRegisterInfo *MRI,
   auto *MBB = MI->getParent();
 
   // FIXME: AMX should assume AVX512 enabled.
   if (ST->hasAVX512()) {
     // Zero stack slot.
     Register Zmm = MRI->createVirtualRegister(&X86::VR512RegClass);
     BuildMI(*MBB, MI, DebugLoc(), TII->get(X86::VPXORDZrr), Zmm)
         .addReg(Zmm, RegState::Undef)
         .addReg(Zmm, RegState::Undef);
     addFrameReference(BuildMI(*MBB, MI, DebugLoc(), TII->get(X86::VMOVUPSZmr)),
                       FrameIdx)
         .addReg(Zmm);
   }
 
   // build psuedo ldtilecfg
-  Register VReg = MRI->createVirtualRegister(&X86::TILECFGRegClass);
-
-  addFrameReference(
-      BuildMI(*MBB, MI, DebugLoc(), TII->get(X86::PLDTILECFG), VReg), FrameIdx);
-
-  return VReg;
+  addFrameReference(BuildMI(*MBB, MI, DebugLoc(), TII->get(X86::LDTILECFG)),
+                    FrameIdx);
 }
 
 static ShapeT getShape(const MachineInstr &MI, MachineRegisterInfo *MRI) {
   unsigned Opcode = MI.getOpcode();
   switch (Opcode) {
   default:
     llvm_unreachable("Unexpected machine instruction on tile");
   case X86::PTILELOADDV:
if (I == E)
     MII = MBB->getFirstNonPHI();
   else {
     MII = MachineBasicBlock::iterator(&*I);
     MII++;
   }
   return &*MII;
 }
 
-static void addTileCFGUse(MachineFunction &MF, Register CFG) {
-  for (MachineBasicBlock &MBB : MF) {
-
-    // Traverse the basic block.
-    for (MachineInstr &MI : MBB) {
-      unsigned Opcode = MI.getOpcode();
-      switch (Opcode) {
-      default:
-        break;
-      case X86::PTILELOADDV:
-      case X86::PTILESTOREDV:
-      case X86::PTDPBSSDV:
-      case X86::PTILEZEROV:
-        unsigned NumOperands = MI.getNumOperands();
-        MI.RemoveOperand(NumOperands - 1);
-        MI.addOperand(MF, MachineOperand::CreateReg(CFG, false));
-        break;
-      }
-    }
+static bool isAMXInstruction(MachineBasicBlock::iterator MII) {
+  switch (MII->getOpcode()) {
+    return false;
+    return true;
+
+struct ReloadTileConfig {
+  enum SuccStatus { NoAMXAtAll = 0, HasAfterCallAMX, HasBeforeCallAMX };
+
+  struct BBInfo {
+    bool HasAMX = false;
+    bool NeedUpdatePred = true;
+    bool IsCallBeforeAMX = false;
+    SuccStatus MaxSucc = NoAMXAtAll;
+    MachineInstr *LastCallMayHasAMX = nullptr;
+
+    BBInfo() = default;
+    BBInfo(ReloadTileConfig *RTC, MachineBasicBlock *MBB,
+           MachineInstr *MI = nullptr) {
+      MachineBasicBlock::iterator MII = MI ? MI->getIterator() : MBB->begin();
+      for (auto E = MBB->end(); MII != E; ++MII) {
+        if (isAMXInstruction(MII)) {
+          HasAMX = true;
+          if (LastCallMayHasAMX)
+            RTC->CfgNeedInsert.insert(LastCallMayHasAMX);
+          LastCallMayHasAMX = nullptr;
+        } else if (MII->isCall()) {
+          if (!HasAMX)
+            IsCallBeforeAMX = true;
+          LastCallMayHasAMX = &*MII;
         }
       }
+      if (LastCallMayHasAMX)
+        RTC->BBUnsolved.insert(MBB);
+    }
+  };
+
+  SmallSet<MachineInstr *, 8> CfgNeedInsert;
+  SmallSet<MachineBasicBlock *, 8> BBUnsolved;
+  DenseMap<MachineBasicBlock *, BBInfo> BBVisitedInfo;
+
+  ReloadTileConfig(MachineInstr *MI, int FI, const TargetInstrInfo *TII,
+                   const TargetRegisterInfo *TRI) {
+    MachineBasicBlock *MBB = MI->getParent();
+    BBVisitedInfo[MBB] = BBInfo(this, MBB, MI);
+    SmallVector<MachineBasicBlock *, 8> WorkList;
+
+    WorkList.push_back(MBB);
+    while (WorkList.size()) {
+      MBB = WorkList.pop_back_val();
+      for (auto I = MBB->succ_begin(), E = MBB->succ_end(); I != E; ++I) {
+        if (!BBVisitedInfo.count(*I)) {
+          BBVisitedInfo[*I] = BBInfo(this, *I);
+          WorkList.push_back(*I);
+        }
+      }
+    }
+
+    // FIXME: Maybe using the last BB in the MF is enough if it dominates all
+    // other BBs. But iterating all BBs isn't too expensive here.
+    // We calculate the MaxSucc of each BB here based on a bottom to up search.
+    // If a BB has AMX instruction, it only has to update its predecessors once.
+    // Otherwise, it updates its predecessors based on MaxSucc itself, which may
+    // be changed. So we update its predecessors again once its MaxSucc changed.
+    for (auto I : BBVisitedInfo) {
+      if (I.second.NeedUpdatePred)
+        WorkList.push_back(I.first);
+      while (WorkList.size()) {
+        MBB = WorkList.pop_back_val();
+        BBVisitedInfo[MBB].NeedUpdatePred = false;
+
+        SuccStatus S = BBVisitedInfo[MBB].HasAMX ? HasBeforeCallAMX
+                                                 : BBVisitedInfo[MBB].MaxSucc;
+        if (BBVisitedInfo[MBB].IsCallBeforeAMX)
+          S = std::min(S, HasAfterCallAMX);
+
+        for (auto I = MBB->pred_begin(), E = MBB->pred_end(); I != E; ++I) {
+          if (!BBVisitedInfo.count(*I))
+            continue;
+          if (S > BBVisitedInfo[*I].MaxSucc) {
+            BBVisitedInfo[*I].MaxSucc = S;
+            // If BB has AMX, its predecessors don't need to be updated again.
+            if (!BBVisitedInfo[*I].HasAMX)
+              BBVisitedInfo[*I].NeedUpdatePred = true;
+          }
+          if (BBVisitedInfo[*I].NeedUpdatePred)
+            WorkList.push_back(*I);
+        }
+      }
+    }
+
+    for (auto *MBB : BBUnsolved)
+      if (BBVisitedInfo[MBB].MaxSucc == HasBeforeCallAMX)
+        CfgNeedInsert.insert(BBVisitedInfo[MBB].LastCallMayHasAMX);
+
+    for (auto *I : CfgNeedInsert) {
+      BitVector UsableRegs(TRI->getNumRegs());
+      const TargetRegisterClass *RC = TRI->getRegClass(X86::TILERegClassID);
+      for (unsigned J = 0; J < RC->getNumRegs(); J++)
+        UsableRegs.set(X86::TMM0 + J);
+      for (MachineOperand &CallMO : I->operands()) {
+        if (CallMO.isRegMask())
+          UsableRegs.clearBitsInMask(CallMO.getRegMask());
+      }
+      if (!UsableRegs.none())
+        addFrameReference(BuildMI(*I->getParent(), ++I->getIterator(),
+                                  DebugLoc(), TII->get(X86::LDTILECFG)),
+                          FI);
+    }
+  }
+};
 
 bool X86PreTileConfig::runOnMachineFunction(MachineFunction &mf) {
   MF = &mf;
   MRI = &mf.getRegInfo();
   ST = &mf.getSubtarget<X86Subtarget>();
   TRI = ST->getRegisterInfo();
   TII = mf.getSubtarget().getInstrInfo();
   DomTree = &getAnalysis<MachineDominatorTree>();
 
   MachineInstr *MI = getTileConfigPoint();
   if (!MI)
     return false;
   unsigned Size = ST->getTileConfigSize();
   Align Alignment = ST->getTileConfigAlignment();
   int SS = mf.getFrameInfo().CreateStackObject(Size, Alignment, false);
-  Register CFG = buildConfigMI(MI, SS, TII, MRI, ST);
-  addTileCFGUse(mf, CFG);
+  buildConfigMI(MI, SS, TII, MRI, ST);
+  ReloadTileConfig(MI, SS, TII, TRI);
   return true;
 }
 
 FunctionPass *llvm::createX86PreTileConfigPass() {
   return new X86PreTileConfig();
 }

llvm/lib/Target/X86/X86RegisterInfo.td

 
 // Bound registers
 def BNDR : RegisterClass<"X86", [v2i64], 128, (sequence "BND%u", 0, 3)>;
 
 // Tiles
 let CopyCost = -1 in // Don't allow copying of tile registers
 def TILE : RegisterClass<"X86", [x86amx], 8192,
                          (sequence "TMM%u", 0, 7)> {let Size = 8192;}
-def TILECFG : RegisterClass<"X86", [untyped], 512, (add TMMCFG)> {
-  let CopyCost = -1;  // Don't allow copying of tile config registers.
-  let isAllocatable = 1;
-  let Size = 512;
-}

llvm/lib/Target/X86/X86TileConfig.cpp

 //
 //===----------------------------------------------------------------------===//
 
 #include "X86.h"
 #include "X86InstrBuilder.h"
 #include "X86MachineFunctionInfo.h"
 #include "X86RegisterInfo.h"
 #include "X86Subtarget.h"
+#include "llvm/ADT/PostOrderIterator.h"
 #include "llvm/CodeGen/LiveIntervals.h"
 #include "llvm/CodeGen/MachineDominators.h"
 #include "llvm/CodeGen/MachineFrameInfo.h"
 #include "llvm/CodeGen/MachineFunctionPass.h"
 #include "llvm/CodeGen/MachineInstr.h"
 #include "llvm/CodeGen/MachineRegisterInfo.h"
 #include "llvm/CodeGen/Passes.h"
 #include "llvm/CodeGen/TargetInstrInfo.h"
static MachineInstr *storeImmToStackSlot(MachineBasicBlock &MBB,
                                          const TargetInstrInfo *TII) {
   unsigned Opc = (BitSize == 8) ? X86::MOV8mi : X86::MOV16mi;
   return addFrameReference(BuildMI(MBB, MI, DebugLoc(), TII->get(Opc)),
                            FrameIdx, Offset)
       .addImm(Imm);
 }
 
 MachineInstr *X86TileConfig::getTileConfigPoint() {
-  for (MachineBasicBlock &MBB : *MF) {
-
-    // Traverse the basic block.
-    for (MachineInstr &MI : MBB)
+  MachineBasicBlock *Entry = &*MF->begin();
+  ReversePostOrderTraversal<MachineBasicBlock *> RPOT(Entry);
+  for (MachineBasicBlock *MBB : RPOT) {
+    for (MachineInstr &MI : *MBB)
       // Refer X86PreTileConfig.cpp.
-      // We only support one tile config for now.
-      if (MI.getOpcode() == X86::PLDTILECFG)
+      // We only support one tile config for now. The other ldtilecfg
+      // is for spill purpose and is dominated by the first ldtilecfg.
+      if (MI.getOpcode() == X86::LDTILECFG)
         return &MI;
   }
 
   return nullptr;
 }
 
 void X86TileConfig::tileConfig() {
   MachineInstr *MI = getTileConfigPoint();
   if (!MI)
     return;
   MachineBasicBlock *MBB = MI->getParent();
-  int SS = MI->getOperand(1).getIndex();
+  int SS = MI->getOperand(0).getIndex();
   BitVector PhysRegs(TRI->getNumRegs());
 
   // Fill in the palette first.
   auto *NewMI = storeImmToStackSlot(*MBB, *MI, 1, 8, SS, 0, TII);
   LIS->InsertMachineInstrInMaps(*NewMI);
   // Fill in the shape of each tile physical register.
   for (unsigned i = 0, e = MRI->getNumVirtRegs(); i != e; ++i) {
     Register VirtReg = Register::index2VirtReg(i);

llvm/test/CodeGen/X86/AMX/amx-across-func.ll

 ; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
 ; RUN: llc < %s -mtriple=x86_64-unknown-unknown -mattr=+amx-int8 -mattr=+avx512f -verify-machineinstrs | FileCheck %s
-
-%struct.__tile_str = type <{ i16, i16, [60 x i8], <256 x i32> }>
+; RUN: llc < %s -mtriple=x86_64-unknown-unknown -mattr=+amx-int8 -mattr=+avx512f -verify-machineinstrs -enable-ipra | FileCheck -check-prefix=IPRA %s
 
 @buf = dso_local global [3072 x i8] zeroinitializer, align 64
 
+define internal void @foo() {
+; CHECK-LABEL: foo:
+; CHECK:       # %bb.0: # %entry
+; CHECK-NEXT:    retq
+;
+; IPRA-LABEL: foo:
+; IPRA:       # %bb.0: # %entry
+; IPRA-NEXT:    retq
+entry:
+  ret void
+}
+
 define dso_local void @test_api(i16 signext %0, i16 signext %1) nounwind {
 ; CHECK-LABEL: test_api:
 ; CHECK:       # %bb.0:
 ; CHECK-NEXT:    pushq %rbp
 ; CHECK-NEXT:    pushq %r15
 ; CHECK-NEXT:    pushq %r14
 ; CHECK-NEXT:    pushq %rbx
 ; CHECK-NEXT:    subq $4056, %rsp # imm = 0xFD8
 ; CHECK-NEXT:    movl %esi, %ebx
 ; CHECK-NEXT:    movl %edi, %ebp
 ; CHECK-NEXT:    vpxord %zmm0, %zmm0, %zmm0
 ; CHECK-NEXT:    vmovdqu64 %zmm0, {{[0-9]+}}(%rsp)
 ; CHECK-NEXT:    movb $1, {{[0-9]+}}(%rsp)
 ; CHECK-NEXT:    movb %bpl, {{[0-9]+}}(%rsp)
 ; CHECK-NEXT:    movw %bx, {{[0-9]+}}(%rsp)
 ; CHECK-NEXT:    movb %bpl, {{[0-9]+}}(%rsp)
 ; CHECK-NEXT:    movw $8, {{[0-9]+}}(%rsp)
 ; CHECK-NEXT:    movb $8, {{[0-9]+}}(%rsp)
 ; CHECK-NEXT:    movw %bx, {{[0-9]+}}(%rsp)
 ; CHECK-NEXT:    ldtilecfg {{[0-9]+}}(%rsp)
-; CHECK-NEXT:    sttilecfg {{[-0-9]+}}(%r{{[sb]}}p) # 64-byte Folded Spill
 ; CHECK-NEXT:    movl $buf, %eax
 ; CHECK-NEXT:    movl $32, %r14d
 ; CHECK-NEXT:    movw $8, %r15w
 ; CHECK-NEXT:    tileloadd (%rax,%r14), %tmm1
 ; CHECK-NEXT:    movabsq $64, %rax
 ; CHECK-NEXT:    tilestored %tmm1, 2048(%rsp,%rax) # 1024-byte Folded Spill
 ; CHECK-NEXT:    movl $buf+1024, %eax
 ; CHECK-NEXT:    tileloadd (%rax,%r14), %tmm2
 ; CHECK-NEXT:    movabsq $64, %rax
 ; CHECK-NEXT:    tilestored %tmm2, 1024(%rsp,%rax) # 1024-byte Folded Spill
-; CHECK-NEXT:    xorl %eax, %eax
 ; CHECK-NEXT:    vzeroupper
 ; CHECK-NEXT:    callq foo
+; CHECK-NEXT:    ldtilecfg {{[0-9]+}}(%rsp)
 ; CHECK-NEXT:    movl $buf+2048, %eax
-; CHECK-NEXT:    ldtilecfg {{[-0-9]+}}(%r{{[sb]}}p) # 64-byte Folded Reload
 ; CHECK-NEXT:    tileloadd (%rax,%r14), %tmm0
 ; CHECK-NEXT:    movabsq $64, %rcx
 ; CHECK-NEXT:    tileloadd 2048(%rsp,%rcx), %tmm1 # 1024-byte Folded Reload
 ; CHECK-NEXT:    movabsq $64, %rcx
 ; CHECK-NEXT:    tileloadd 1024(%rsp,%rcx), %tmm2 # 1024-byte Folded Reload
 ; CHECK-NEXT:    tdpbssd %tmm2, %tmm1, %tmm0
 ; CHECK-NEXT:    tilestored %tmm0, (%rax,%r14)
 ; CHECK-NEXT:    addq $4056, %rsp # imm = 0xFD8
 ; CHECK-NEXT:    popq %rbx
 ; CHECK-NEXT:    popq %r14
 ; CHECK-NEXT:    popq %r15
 ; CHECK-NEXT:    popq %rbp
 ; CHECK-NEXT:    tilerelease
 ; CHECK-NEXT:    retq
+;
+; IPRA-LABEL: test_api:
+; IPRA:       # %bb.0:
+; IPRA-NEXT:    subq $72, %rsp
+; IPRA-NEXT:    vpxord %zmm0, %zmm0, %zmm0
+; IPRA-NEXT:    vmovdqu64 %zmm0, {{[0-9]+}}(%rsp)
+; IPRA-NEXT:    movb $1, {{[0-9]+}}(%rsp)
+; IPRA-NEXT:    movb %dil, {{[0-9]+}}(%rsp)
+; IPRA-NEXT:    movw $8, {{[0-9]+}}(%rsp)
+; IPRA-NEXT:    movb $8, {{[0-9]+}}(%rsp)
+; IPRA-NEXT:    movw %si, {{[0-9]+}}(%rsp)
+; IPRA-NEXT:    movb %dil, {{[0-9]+}}(%rsp)
+; IPRA-NEXT:    movw %si, {{[0-9]+}}(%rsp)
+; IPRA-NEXT:    ldtilecfg {{[0-9]+}}(%rsp)
+; IPRA-NEXT:    movl $buf, %eax
+; IPRA-NEXT:    movl $32, %ecx
+; IPRA-NEXT:    movw $8, %dx
+; IPRA-NEXT:    tileloadd (%rax,%rcx), %tmm0
+; IPRA-NEXT:    movl $buf+1024, %eax
+; IPRA-NEXT:    tileloadd (%rax,%rcx), %tmm1
+; IPRA-NEXT:    callq foo
+; IPRA-NEXT:    movl $buf+2048, %eax
+; IPRA-NEXT:    tileloadd (%rax,%rcx), %tmm2
+; IPRA-NEXT:    tdpbssd %tmm1, %tmm0, %tmm2
+; IPRA-NEXT:    tilestored %tmm2, (%rax,%rcx)
+; IPRA-NEXT:    addq $72, %rsp
+; IPRA-NEXT:    tilerelease
+; IPRA-NEXT:    vzeroupper
+; IPRA-NEXT:    retq
   %3 = tail call x86_amx @llvm.x86.tileloadd64.internal(i16 %0, i16 8, i8* getelementptr inbounds ([3072 x i8], [3072 x i8]* @buf, i64 0, i64 0), i64 32)
   %4 = tail call x86_amx @llvm.x86.tileloadd64.internal(i16 8, i16 %1, i8* getelementptr inbounds ([3072 x i8], [3072 x i8]* @buf, i64 0, i64 1024), i64 32)
-  tail call void (...) @foo()
+  call void @foo()
   %5 = tail call x86_amx @llvm.x86.tileloadd64.internal(i16 %0, i16 %1, i8* getelementptr inbounds ([3072 x i8], [3072 x i8]* @buf, i64 0, i64 2048), i64 32)
   %6 = tail call x86_amx @llvm.x86.tdpbssd.internal(i16 %0, i16 %1, i16 8, x86_amx %5, x86_amx %3, x86_amx %4)
   tail call void @llvm.x86.tilestored64.internal(i16 %0, i16 %1, i8* getelementptr inbounds ([3072 x i8], [3072 x i8]* @buf, i64 0, i64 2048), i64 32, x86_amx %6)
   ret void
 }
 
-declare dso_local void @foo(...)
+define dso_local i32 @test_loop(i32 %0) nounwind {
+; CHECK-LABEL: test_loop:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    pushq %rbp
+; CHECK-NEXT:    pushq %r15
+; CHECK-NEXT:    pushq %r14
+; CHECK-NEXT:    pushq %r13
+; CHECK-NEXT:    pushq %r12
+; CHECK-NEXT:    pushq %rbx
+; CHECK-NEXT:    subq $3016, %rsp # imm = 0xBC8
+; CHECK-NEXT:    movl %edi, %r14d
+; CHECK-NEXT:    callq foo
+; CHECK-NEXT:    vpxord %zmm0, %zmm0, %zmm0
+; CHECK-NEXT:    vmovdqu64 %zmm0, {{[0-9]+}}(%rsp)
+; CHECK-NEXT:    movb $1, {{[0-9]+}}(%rsp)
+; CHECK-NEXT:    movb $8, {{[0-9]+}}(%rsp)
+; CHECK-NEXT:    movw $8, {{[0-9]+}}(%rsp)
+; CHECK-NEXT:    ldtilecfg {{[0-9]+}}(%rsp)
+; CHECK-NEXT:    testl %r14d, %r14d
+; CHECK-NEXT:    jg .LBB2_4
+; CHECK-NEXT:  # %bb.1: # %.preheader
+; CHECK-NEXT:    movl $7, %ebp
+; CHECK-NEXT:    movl $buf, %r15d
+; CHECK-NEXT:    movl $32, %r12d
+; CHECK-NEXT:    movw $8, %bx
+; CHECK-NEXT:    movl $buf+2048, %r13d
+; CHECK-NEXT:    .p2align 4, 0x90
+; CHECK-NEXT:  .LBB2_2: # =>This Inner Loop Header: Depth=1
+; CHECK-NEXT:    tileloadd (%r15,%r12), %tmm0
+; CHECK-NEXT:    movabsq $64, %rax
+; CHECK-NEXT:    tilestored %tmm0, 1024(%rsp,%rax) # 1024-byte Folded Spill
+; CHECK-NEXT:    vzeroupper
+; CHECK-NEXT:    callq foo
+; CHECK-NEXT:    ldtilecfg {{[0-9]+}}(%rsp)
+; CHECK-NEXT:    movabsq $64, %rax
+; CHECK-NEXT:    tileloadd 1024(%rsp,%rax), %tmm0 # 1024-byte Folded Reload
+; CHECK-NEXT:    tilestored %tmm0, (%r13,%r12)
+; CHECK-NEXT:    callq foo
+; CHECK-NEXT:    ldtilecfg {{[0-9]+}}(%rsp)
+; CHECK-NEXT:    decl %ebp
+; CHECK-NEXT:    cmpl $7, %ebp
+; CHECK-NEXT:    jne .LBB2_2
+; CHECK-NEXT:  # %bb.3:
+; CHECK-NEXT:    cmpl $3, %r14d
+; CHECK-NEXT:    jne .LBB2_4
+; CHECK-NEXT:  # %bb.6:
+; CHECK-NEXT:    testl %ebp, %ebp
+; CHECK-NEXT:    jne .LBB2_5
+; CHECK-NEXT:  # %bb.7:
+; CHECK-NEXT:    incl %r14d
+; CHECK-NEXT:    jmp .LBB2_8
+; CHECK-NEXT:  .LBB2_4:
+; CHECK-NEXT:    vzeroupper
+; CHECK-NEXT:    callq foo
+; CHECK-NEXT:    ldtilecfg {{[0-9]+}}(%rsp)
+; CHECK-NEXT:    movl $32, %eax
+; CHECK-NEXT:    movl $buf+1024, %ecx
+; CHECK-NEXT:    movw $8, %dx
+; CHECK-NEXT:    tileloadd (%rcx,%rax), %tmm0
+; CHECK-NEXT:    tilestored %tmm0, (%rcx,%rax)
+; CHECK-NEXT:  .LBB2_5:
+; CHECK-NEXT:    decl %r14d
+; CHECK-NEXT:  .LBB2_8:
+; CHECK-NEXT:    movl %r14d, %eax
+; CHECK-NEXT:    addq $3016, %rsp # imm = 0xBC8
+; CHECK-NEXT:    popq %rbx
+; CHECK-NEXT:    popq %r12
+; CHECK-NEXT:    popq %r13
+; CHECK-NEXT:    popq %r14
+; CHECK-NEXT:    popq %r15
+; CHECK-NEXT:    popq %rbp
+; CHECK-NEXT:    tilerelease
+; CHECK-NEXT:    retq
+;
+; IPRA-LABEL: test_loop:
+; IPRA:       # %bb.0:
+; IPRA-NEXT:    subq $72, %rsp
+; IPRA-NEXT:    movl %edi, %eax
+; IPRA-NEXT:    callq foo
+; IPRA-NEXT:    vpxord %zmm0, %zmm0, %zmm0
+; IPRA-NEXT:    vmovdqu64 %zmm0, {{[0-9]+}}(%rsp)
+; IPRA-NEXT:    movb $1, {{[0-9]+}}(%rsp)
+; IPRA-NEXT:    movb $8, {{[0-9]+}}(%rsp)
+; IPRA-NEXT:    movw $8, {{[0-9]+}}(%rsp)
+; IPRA-NEXT:    ldtilecfg {{[0-9]+}}(%rsp)
+; IPRA-NEXT:    testl %edi, %edi
+; IPRA-NEXT:    jg .LBB2_4
+; IPRA-NEXT:  # %bb.1: # %.preheader
+; IPRA-NEXT:    movl $7, %ecx
+; IPRA-NEXT:    movl $buf, %r8d
+; IPRA-NEXT:    movl $32, %esi
+; IPRA-NEXT:    movw $8, %di
+; IPRA-NEXT:    movl $buf+2048, %edx
+; IPRA-NEXT:    .p2align 4, 0x90
+; IPRA-NEXT:  .LBB2_2: # =>This Inner Loop Header: Depth=1
+; IPRA-NEXT:    tileloadd (%r8,%rsi), %tmm0
+; IPRA-NEXT:    callq foo
+; IPRA-NEXT:    tilestored %tmm0, (%rdx,%rsi)
+; IPRA-NEXT:    callq foo
+; IPRA-NEXT:    decl %ecx
+; IPRA-NEXT:    cmpl $7, %ecx
+; IPRA-NEXT:    jne .LBB2_2
+; IPRA-NEXT:  # %bb.3:
+; IPRA-NEXT:    cmpl $3, %eax
+; IPRA-NEXT:    jne .LBB2_4
+; IPRA-NEXT:  # %bb.6:
+; IPRA-NEXT:    testl %ecx, %ecx
+; IPRA-NEXT:    jne .LBB2_5
+; IPRA-NEXT:  # %bb.7:
+; IPRA-NEXT:    incl %eax
+; IPRA-NEXT:    jmp .LBB2_8
+; IPRA-NEXT:  .LBB2_4:
+; IPRA-NEXT:    callq foo
+; IPRA-NEXT:    movl $32, %ecx
+; IPRA-NEXT:    movl $buf+1024, %edx
+; IPRA-NEXT:    movw $8, %si
+; IPRA-NEXT:    tileloadd (%rdx,%rcx), %tmm0
+; IPRA-NEXT:    tilestored %tmm0, (%rdx,%rcx)
+; IPRA-NEXT:  .LBB2_5:
+; IPRA-NEXT:    decl %eax
+; IPRA-NEXT:  .LBB2_8:
+; IPRA-NEXT:    addq $72, %rsp
+; IPRA-NEXT:    tilerelease
+; IPRA-NEXT:    vzeroupper
+; IPRA-NEXT:    retq
+  call void @foo()
+  br label %2
+2:
+  %3 = icmp sgt i32 %0, 0
+  br i1 %3, label %11, label %6
+4:
+  %5 = icmp eq i32 %0, 3
+  br i1 %5, label %13, label %11
+6:
+  %7 = phi i32 [ %9, %6 ], [ 0, %2 ]
+  %8 = tail call x86_amx @llvm.x86.tileloadd64.internal(i16 8, i16 8, i8* getelementptr inbounds ([3072 x i8], [3072 x i8]* @buf, i64 0, i64 0), i64 32)
+  call void @foo()
+  tail call void @llvm.x86.tilestored64.internal(i16 8, i16 8, i8* getelementptr inbounds ([3072 x i8], [3072 x i8]* @buf, i64 0, i64 2048), i64 32, x86_amx %8)
+  call void @foo()
+  %9 = add i32 %7, 1
+  %10 = icmp eq i32 %9, 0
+  br i1 %10, label %4, label %6
+11:
+  call void @foo()
+  %12 = tail call x86_amx @llvm.x86.tileloadd64.internal(i16 8, i16 8, i8* getelementptr inbounds ([3072 x i8], [3072 x i8]* @buf, i64 0, i64 1024), i64 32)
+  tail call void @llvm.x86.tilestored64.internal(i16 8, i16 8, i8* getelementptr inbounds ([3072 x i8], [3072 x i8]* @buf, i64 0, i64 1024), i64 32, x86_amx %12)
+  br label %17
+13:
+  %14 = icmp eq i32 %9, 7
+  br i1 %14, label %15, label %17
+15:
+  %16 = add i32 %0, 1
+  br label %19
+17:
+  %18 = sub i32 %0, 1
+  br label %19
+19:
+  %20 = phi i32 [ %16, %15 ], [ %18, %17 ]
+  ret i32 %20
+}
 
 declare x86_amx @llvm.x86.tileloadd64.internal(i16, i16, i8*, i64)
 declare x86_amx @llvm.x86.tdpbssd.internal(i16, i16, i16, x86_amx, x86_amx, x86_amx)
 declare void @llvm.x86.tilestored64.internal(i16, i16, i8*, i64, x86_amx)

llvm/test/CodeGen/X86/AMX/amx-bf16-intrinsics.ll

 ; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
 ; RUN: llc < %s -mtriple=x86_64-unknown-unknown -mattr=+amx-tile -mattr=+amx-bf16 -verify-machineinstrs | FileCheck %s
 
 define void @test_amx() {
 ; CHECK-LABEL: test_amx:
 ; CHECK:       # %bb.0:
 ; CHECK-NEXT:    tdpbf16ps %tmm7, %tmm4, %tmm3
+; CHECK-NEXT:    tilerelease
 ; CHECK-NEXT:    retq
   call void @llvm.x86.tdpbf16ps(i8 3, i8 4, i8 7)
   ret void
 }
 
 declare void @llvm.x86.tdpbf16ps(i8 %tile0, i8 %tile1, i8 %tile2)

llvm/test/CodeGen/X86/AMX/amx-ldtilecfg-insert.ll

+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
+; RUN: llc < %s -mtriple=x86_64-unknown-unknown -mattr=+amx-int8 -mattr=+avx512f -verify-machineinstrs | FileCheck %s
+@buf = dso_local global [3072 x i8] zeroinitializer, align 16
+
+define dso_local void @test1(i16 signext %0, i16 signext %1) nounwind {
+; CHECK-LABEL: test1:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    vpxord %zmm0, %zmm0, %zmm0
+; CHECK-NEXT:    vmovdqu64 %zmm0, -{{[0-9]+}}(%rsp)
+; CHECK-NEXT:    movb $1, -{{[0-9]+}}(%rsp)
+; CHECK-NEXT:    movb %dil, -{{[0-9]+}}(%rsp)
+; CHECK-NEXT:    movw $8, -{{[0-9]+}}(%rsp)
+; CHECK-NEXT:    movb $8, -{{[0-9]+}}(%rsp)
+; CHECK-NEXT:    movw %si, -{{[0-9]+}}(%rsp)
+; CHECK-NEXT:    movb %dil, -{{[0-9]+}}(%rsp)
+; CHECK-NEXT:    movw %si, -{{[0-9]+}}(%rsp)
+; CHECK-NEXT:    ldtilecfg -{{[0-9]+}}(%rsp)
+; CHECK-NEXT:    movl $buf, %eax
+; CHECK-NEXT:    movl $32, %ecx
+; CHECK-NEXT:    movw $8, %dx
+; CHECK-NEXT:    tileloadd (%rax,%rcx), %tmm0
+; CHECK-NEXT:    movl $buf+1024, %eax
+; CHECK-NEXT:    tileloadd (%rax,%rcx), %tmm1
+; CHECK-NEXT:    movl $buf+2048, %eax
+; CHECK-NEXT:    tileloadd (%rax,%rcx), %tmm2
+; CHECK-NEXT:    tdpbssd %tmm1, %tmm0, %tmm2
+; CHECK-NEXT:    tilestored %tmm2, (%rax,%rcx)
+; CHECK-NEXT:    tilerelease
+; CHECK-NEXT:    vzeroupper
+; CHECK-NEXT:    jmp foo # TAILCALL
+  %3 = tail call x86_amx @llvm.x86.tileloadd64.internal(i16 %0, i16 8, i8* getelementptr inbounds ([3072 x i8], [3072 x i8]* @buf, i64 0, i64 0), i64 32)
+  %4 = tail call x86_amx @llvm.x86.tileloadd64.internal(i16 8, i16 %1, i8* getelementptr inbounds ([3072 x i8], [3072 x i8]* @buf, i64 0, i64 1024), i64 32)
+  %5 = tail call x86_amx @llvm.x86.tileloadd64.internal(i16 %0, i16 %1, i8* getelementptr inbounds ([3072 x i8], [3072 x i8]* @buf, i64 0, i64 2048), i64 32)
+  %6 = tail call x86_amx @llvm.x86.tdpbssd.internal(i16 %0, i16 %1, i16 8, x86_amx %5, x86_amx %3, x86_amx %4)
+  tail call void @llvm.x86.tilestored64.internal(i16 %0, i16 %1, i8* getelementptr inbounds ([3072 x i8], [3072 x i8]* @buf, i64 0, i64 2048), i64 32, x86_amx %6)
+  tail call void @foo()
+  ret void
+}
+
+define dso_local void @test2(i16 signext %0, i16 signext %1) nounwind {
+; CHECK-LABEL: test2:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    pushq %rbp
+; CHECK-NEXT:    pushq %rbx
+; CHECK-NEXT:    subq $72, %rsp
+; CHECK-NEXT:    movl %esi, %ebx
+; CHECK-NEXT:    movl %edi, %ebp
+; CHECK-NEXT:    vpxord %zmm0, %zmm0, %zmm0
+; CHECK-NEXT:    vmovdqu64 %zmm0, {{[0-9]+}}(%rsp)
+; CHECK-NEXT:    movb $1, {{[0-9]+}}(%rsp)
+; CHECK-NEXT:    movb $8, {{[0-9]+}}(%rsp)
+; CHECK-NEXT:    movw %bx, {{[0-9]+}}(%rsp)
+; CHECK-NEXT:    movb %bpl, {{[0-9]+}}(%rsp)
+; CHECK-NEXT:    movw %bx, {{[0-9]+}}(%rsp)
+; CHECK-NEXT:    movb %bpl, {{[0-9]+}}(%rsp)
+; CHECK-NEXT:    movw $8, {{[0-9]+}}(%rsp)
+; CHECK-NEXT:    movb $8, {{[0-9]+}}(%rsp)
+; CHECK-NEXT:    movw %bx, {{[0-9]+}}(%rsp)
+; CHECK-NEXT:    movb %bpl, {{[0-9]+}}(%rsp)
+; CHECK-NEXT:    movw $8, {{[0-9]+}}(%rsp)
+; CHECK-NEXT:    ldtilecfg {{[0-9]+}}(%rsp)
+; CHECK-NEXT:    vzeroupper
+; CHECK-NEXT:    callq foo
+; CHECK-NEXT:    ldtilecfg {{[0-9]+}}(%rsp)
+; CHECK-NEXT:    xorl %eax, %eax
+; CHECK-NEXT:    testb %al, %al
+; CHECK-NEXT:    jne .LBB1_3
+; CHECK-NEXT:  # %bb.1: # %if.true
+; CHECK-NEXT:    movw $8, %ax
+; CHECK-NEXT:    tilezero %tmm0
+; CHECK-NEXT:    movl $32, %ecx
+; CHECK-NEXT:    movl $buf+1024, %edx
+; CHECK-NEXT:    tileloadd (%rdx,%rcx), %tmm1
+; CHECK-NEXT:    movl $buf+2048, %edx
+; CHECK-NEXT:    tileloadd (%rdx,%rcx), %tmm2
+; CHECK-NEXT:    tdpbssd %tmm2, %tmm1, %tmm0
+; CHECK-NEXT:    tilestored %tmm0, (%rdx,%rcx)
+; CHECK-NEXT:    jmp .LBB1_2
+; CHECK-NEXT:  .LBB1_3: # %if.false
+; CHECK-NEXT:    movl $buf, %eax
+; CHECK-NEXT:    movl $32, %ecx
+; CHECK-NEXT:    movw $8, %dx
+; CHECK-NEXT:    tileloadd (%rax,%rcx), %tmm3
+; CHECK-NEXT:    movl $buf+1024, %eax
+; CHECK-NEXT:    tileloadd (%rax,%rcx), %tmm4
+; CHECK-NEXT:    movl $buf+2048, %eax
+; CHECK-NEXT:    tileloadd (%rax,%rcx), %tmm2
+; CHECK-NEXT:    tdpbssd %tmm2, %tmm4, %tmm3
+; CHECK-NEXT:    tilestored %tmm3, (%rax,%rcx)
+; CHECK-NEXT:  .LBB1_2: # %if.true
+; CHECK-NEXT:    addq $72, %rsp
+; CHECK-NEXT:    popq %rbx
+; CHECK-NEXT:    popq %rbp
+; CHECK-NEXT:    tilerelease
+; CHECK-NEXT:    retq
+  call void @foo()
+  br i1 undef, label %if.true, label %if.false
+
+if.true:
+  %t1 = tail call x86_amx @llvm.x86.tilezero.internal(i16 %0, i16 8)
+  %t2 = tail call x86_amx @llvm.x86.tileloadd64.internal(i16 8, i16 %1, i8* getelementptr inbounds ([3072 x i8], [3072 x i8]* @buf, i64 0, i64 1024), i64 32)
+  %t3 = tail call x86_amx @llvm.x86.tileloadd64.internal(i16 %0, i16 %1, i8* getelementptr inbounds ([3072 x i8], [3072 x i8]* @buf, i64 0, i64 2048), i64 32)
+  %t4 = tail call x86_amx @llvm.x86.tdpbssd.internal(i16 %0, i16 %1, i16 8, x86_amx %t1, x86_amx %t2, x86_amx %t3)
+  tail call void @llvm.x86.tilestored64.internal(i16 %0, i16 %1, i8* getelementptr inbounds ([3072 x i8], [3072 x i8]* @buf, i64 0, i64 2048), i64 32, x86_amx %t4)
+  br label %exit
+
+if.false:
+  %t5 = tail call x86_amx @llvm.x86.tileloadd64.internal(i16 %0, i16 8, i8* getelementptr inbounds ([3072 x i8], [3072 x i8]* @buf, i64 0, i64 0), i64 32)
+  %t6 = tail call x86_amx @llvm.x86.tileloadd64.internal(i16 8, i16 %1, i8* getelementptr inbounds ([3072 x i8], [3072 x i8]* @buf, i64 0, i64 1024), i64 32)
+  %t7 = tail call x86_amx @llvm.x86.tileloadd64.internal(i16 %0, i16 %1, i8* getelementptr inbounds ([3072 x i8], [3072 x i8]* @buf, i64 0, i64 2048), i64 32)
+  %t8 = tail call x86_amx @llvm.x86.tdpbssd.internal(i16 %0, i16 %1, i16 8, x86_amx %t5, x86_amx %t6, x86_amx %t7)
+  tail call void @llvm.x86.tilestored64.internal(i16 %0, i16 %1, i8* getelementptr inbounds ([3072 x i8], [3072 x i8]* @buf, i64 0, i64 2048), i64 32, x86_amx %t8)
+  br label %exit
+
+exit:
+  ret void
+}
+
+declare dso_local void @foo() nounwind
+declare x86_amx @llvm.x86.tilezero.internal(i16, i16)
+declare x86_amx @llvm.x86.tileloadd64.internal(i16, i16, i8*, i64)
+declare x86_amx @llvm.x86.tdpbssd.internal(i16, i16, i16, x86_amx, x86_amx, x86_amx)
+declare void @llvm.x86.tilestored64.internal(i16, i16, i8*, i64, x86_amx)

llvm/test/CodeGen/X86/AMX/amx-spill-merge.ll

 ; CHECK-NEXT:    movb %bpl, {{[0-9]+}}(%rsp)
 ; CHECK-NEXT:    movw %bx, {{[0-9]+}}(%rsp)
 ; CHECK-NEXT:    ldtilecfg {{[0-9]+}}(%rsp)
 ; CHECK-NEXT:    movl $32, %r14d
 ; CHECK-NEXT:    movl $buf+2048, %r15d
 ; CHECK-NEXT:    tileloadd (%r15,%r14), %tmm5
 ; CHECK-NEXT:    xorl %eax, %eax
 ; CHECK-NEXT:    testb %al, %al
-; CHECK-NEXT:    sttilecfg {{[-0-9]+}}(%r{{[sb]}}p) # 64-byte Folded Spill
-; CHECK-NEXT:    movl $buf, %eax
-; CHECK-NEXT:    movw $8, %cx
 ; CHECK-NEXT:    jne .LBB0_2
 ; CHECK-NEXT:  # %bb.1: # %if.true
+; CHECK-NEXT:    movl $buf, %eax
+; CHECK-NEXT:    movw $8, %cx
 ; CHECK-NEXT:    tileloadd (%rax,%r14), %tmm0
 ; CHECK-NEXT:    movl $buf+1024, %eax
 ; CHECK-NEXT:    tileloadd (%rax,%r14), %tmm1
 ; CHECK-NEXT:    movabsq $64, %rax
 ; CHECK-NEXT:    tilestored %tmm5, 2048(%rsp,%rax) # 1024-byte Folded Spill
 ; CHECK-NEXT:    tdpbssd %tmm1, %tmm0, %tmm5
 ; CHECK-NEXT:    movabsq $64, %rax
 ; CHECK-NEXT:    tilestored %tmm5, 1024(%rsp,%rax) # 1024-byte Folded Spill
 ; CHECK-NEXT:    xorl %eax, %eax
 ; CHECK-NEXT:    vzeroupper
 ; CHECK-NEXT:    callq foo
-; CHECK-NEXT:    ldtilecfg {{[-0-9]+}}(%r{{[sb]}}p) # 64-byte Folded Reload
+; CHECK-NEXT:    ldtilecfg {{[0-9]+}}(%rsp)
 ; CHECK-NEXT:    movabsq $64, %rax
 ; CHECK-NEXT:    tileloadd 1024(%rsp,%rax), %tmm6 # 1024-byte Folded Reload
 ; CHECK-NEXT:    jmp .LBB0_3
 ; CHECK-NEXT:  .LBB0_2: # %if.false
+; CHECK-NEXT:    movl $buf, %eax
+; CHECK-NEXT:    movw $8, %cx
 ; CHECK-NEXT:    tileloadd (%rax,%r14), %tmm2
 ; CHECK-NEXT:    movl $buf+1024, %eax
 ; CHECK-NEXT:    tileloadd (%rax,%r14), %tmm3
 ; CHECK-NEXT:    movabsq $64, %rax
 ; CHECK-NEXT:    tilestored %tmm5, 2048(%rsp,%rax) # 1024-byte Folded Spill
 ; CHECK-NEXT:    tdpbssd %tmm3, %tmm2, %tmm5
 ; CHECK-NEXT:    movabsq $64, %rax
 ; CHECK-NEXT:    tilestored %tmm5, 1024(%rsp,%rax) # 1024-byte Folded Spill
 ; CHECK-NEXT:    xorl %eax, %eax
 ; CHECK-NEXT:    vzeroupper
 ; CHECK-NEXT:    callq foo
-; CHECK-NEXT:    ldtilecfg {{[-0-9]+}}(%r{{[sb]}}p) # 64-byte Folded Reload
+; CHECK-NEXT:    ldtilecfg {{[0-9]+}}(%rsp)
 ; CHECK-NEXT:    movabsq $64, %rax
 ; CHECK-NEXT:    tileloadd 1024(%rsp,%rax), %tmm6 # 1024-byte Folded Reload
 ; CHECK-NEXT:    tilestored %tmm6, (%r15,%r14)
 ; CHECK-NEXT:  .LBB0_3: # %exit
 ; CHECK-NEXT:    movl $buf, %eax
 ; CHECK-NEXT:    movl $32, %ecx
 ; CHECK-NEXT:    movw $8, %dx
 ; CHECK-NEXT:    tileloadd (%rax,%rcx), %tmm4

 ; CHECK-NEXT:    tilezero %tmm0
 ; CHECK-NEXT:    xorl %eax, %eax
 ; CHECK-NEXT:    testb %al, %al
 ; CHECK-NEXT:    jne .LBB1_3
 ; CHECK-NEXT:  # %bb.1: # %loop.header.preheader
 ; CHECK-NEXT:    movq %rdi, %rbx
 ; CHECK-NEXT:    movl $32, %r14d
 ; CHECK-NEXT:    xorl %ebp, %ebp
-; CHECK-NEXT:    sttilecfg {{[-0-9]+}}(%r{{[sb]}}p) # 64-byte Folded Spill
 ; CHECK-NEXT:    .p2align 4, 0x90
 ; CHECK-NEXT:  .LBB1_2: # %loop.header
 ; CHECK-NEXT:    # =>This Inner Loop Header: Depth=1
 ; CHECK-NEXT:    movabsq $64, %rax
 ; CHECK-NEXT:    tilestored %tmm0, 1024(%rsp,%rax) # 1024-byte Folded Spill
 ; CHECK-NEXT:    tilestored %tmm0, (%rbx,%r14)
 ; CHECK-NEXT:    xorl %eax, %eax
 ; CHECK-NEXT:    vzeroupper
 ; CHECK-NEXT:    callq foo
-; CHECK-NEXT:    ldtilecfg {{[-0-9]+}}(%r{{[sb]}}p) # 64-byte Folded Reload
+; CHECK-NEXT:    ldtilecfg {{[0-9]+}}(%rsp)
 ; CHECK-NEXT:    tilezero %tmm0
 ; CHECK-NEXT:    tileloadd (%rbx,%r14), %tmm1
 ; CHECK-NEXT:    tileloadd (%rbx,%r14), %tmm2
 ; CHECK-NEXT:    tdpbssd %tmm2, %tmm1, %tmm0
 ; CHECK-NEXT:    tilestored %tmm0, (%rbx,%r14)
 ; CHECK-NEXT:    movabsq $64, %rax
 ; CHECK-NEXT:    tileloadd 1024(%rsp,%rax), %tmm0 # 1024-byte Folded Reload
 ; CHECK-NEXT:    incl %ebp