[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

284

The Algorithm complexity is O(N*N)

[xiangzhangllvm]

llvm/lib/Target/X86/X86PreTileConfig.cpp

284

The Algorithm complexity is O(N*N)

284

The Algorithm complexity is O(N*N)

I think we can do this in 1 DFS.

[xiangzhangllvm]

llvm/lib/Target/X86/X86PreTileConfig.cpp

244

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

[pengfei]

llvm/lib/Target/X86/X86PreTileConfig.cpp

244

Good catch! Thanks.

284

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

232

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

[LuoYuanke]

llvm/lib/Target/X86/X86PreTileConfig.cpp

242

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

242

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

327

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

327

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

299

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

324

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

[pengfei]

Adress Yuanke's comments.

[LuoYuanke]

llvm/lib/Target/X86/X86PreTileConfig.cpp

337

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

337

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

263

I'm not sure there is only one exit.

333

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

340

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

[xiangzhangllvm]

llvm/lib/Target/X86/X86PreTileConfig.cpp

259

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

282

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

291

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

297

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

[pengfei]

llvm/lib/Target/X86/X86PreTileConfig.cpp

282

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

291

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.

297

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

228

LGTM for the logic.

234

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

246

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

229

Better to have a diagram comments to illustrate the state.

234

!WorkList.empty()?

234

Agree. It is easy to mix with HasBeforeCallAMX.

246

WorkList should be clear before this loop?

249

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

260

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

265

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

[pengfei]

llvm/lib/Target/X86/X86PreTileConfig.cpp

228

Thanks for the review.

234

Should IsCallBeforeAMX better?

[LuoYuanke]

llvm/lib/Target/X86/X86PreTileConfig.cpp

251

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

234

Maybe HasCallBeforeAMXInBB.

[pengfei]

llvm/lib/Target/X86/X86PreTileConfig.cpp

234

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

251

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.

260

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

260

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

260

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

246

Not necessary. But adding it seems more robust.

265

Add some comments on line 290.

[LuoYuanke]

llvm/lib/Target/X86/X86PreTileConfig.cpp

261

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

269

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

261

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

269

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

228

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

228

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

228

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

 //===------- X86ExpandPseudo.cpp - Expand pseudo instructions -------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
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

 //===-- X86FrameLowering.cpp - X86 Frame Information ----------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
if (Terminator == MBB.end() || !isTailCallOpcode(Terminator->getOpcode())) {
     int Offset = -1 * X86FI->getTCReturnAddrDelta();
     assert(Offset >= 0 && "TCDelta should never be positive");
     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))
-    BuildMI(MBB, Terminator, DL, TII.get(X86::TILERELEASE));
 }
 
 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

 //===- X86ISelDAGToDAG.cpp - A DAG pattern matching inst selector for X86 -===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
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

 //===-- X86InstrInfo.cpp - X86 Instruction Information --------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
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

 //===-- X86PreTileConfig.cpp - Tile Register Configure---------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
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) {
+typedef std::pair<MachineInstr *, MachineBasicBlock::iterator> CallPointPair;
 
-    // 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;
+
+static bool findAllSuccSlot(CallPointPair CPP,
+                            std::queue<CallPointPair> &WorkList,
+                            SmallSet<MachineInstr *, 8> &LastAMXSet) {
+  bool HasAMX = false;
+  MachineInstr *LastAMX = CPP.first;
+  MachineBasicBlock *MBB = CPP.second->getParent();
+  for (MachineBasicBlock::iterator I = CPP.second, E = MBB->end(); I != E;
+       ++I) {
+    if (isAMXInstruction(I)) {
+      HasAMX = true;
+      LastAMX = &*I;
+      LastAMXSet.erase(CPP.first);
+    } else if (I->isCall()) {
+      if (std::next(I) != E)
+        WorkList.push(CallPointPair(LastAMX, I));
+      return HasAMX;
+    }
+  }
+
+  if (HasAMX)
+    LastAMXSet.insert(LastAMX);
+
+  for (MachineBasicBlock::succ_iterator I = MBB->succ_begin(),
+                                        E = MBB->succ_end();
+       I != E; ++I)
+    if (findAllSuccSlot(CallPointPair(LastAMX, (*I)->begin()), WorkList,
+                        LastAMXSet))
+      HasAMX = true;
+
+  return HasAMX;
+}
+
+static void reloadTileConfig(MachineInstr *MI, int FI,
+                             const TargetInstrInfo *TII,
+                             const TargetRegisterInfo *TRI) {
+  SmallSet<MachineInstr *, 8> LastAMXSet;
+  SmallSet<MachineInstr *, 8> CallVisited;
+  SmallSet<MachineInstr *, 8> CfgNeedInsert;
+
+  std::queue<CallPointPair> WorkList;
+  WorkList.push(CallPointPair(nullptr, MI->getIterator()));
+
+  while (!WorkList.empty()) {
+    CallPointPair CPP = WorkList.front();
+    WorkList.pop();
+    MachineInstr *Call = &*CPP.second;
+    if (CallVisited.count(Call)) {
+      if (CfgNeedInsert.count(Call))
+        LastAMXSet.erase(CPP.first);
+      continue;
     }
+    CallVisited.insert(Call);
+    CPP.second++;
+    if (findAllSuccSlot(CPP, WorkList, LastAMXSet) && Call != MI)
+      CfgNeedInsert.insert(Call);
+  }
+
+  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);
+  }
+
+  for (auto I : LastAMXSet)
+    BuildMI(*I->getParent(), ++I->getIterator(), DebugLoc(),
+            TII->get(X86::TILERELEASE));
 }
 
 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

 //===-- X86TileConfig.cpp - Tile Register Configure----------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 /// \file Pass to config the shape of AMX physical registers
 /// AMX register need to be configured before use. In X86PreTileConfig pass
 /// the pldtilecfg instruction is inserted, however at that time we don't
 /// know the shape of each physical tile registers, because the register
 /// allocation is not done yet. This pass runs after egister allocation
 /// pass. It collects the shape information of each physical tile register
 /// and store the shape in the stack slot that is allocated for load config
 /// to tile config register.
 //
 //===----------------------------------------------------------------------===//
 
 #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:    tilerelease
 ; 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:    tilerelease
+; IPRA-NEXT:    addq $72, %rsp
+; 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(...)
-
 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-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) local_unnamed_addr {
+; 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:    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) local_unnamed_addr {
+; CHECK-LABEL: test2:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    pushq %rbp
+; CHECK-NEXT:    .cfi_def_cfa_offset 16
+; CHECK-NEXT:    pushq %rbx
+; CHECK-NEXT:    .cfi_def_cfa_offset 24
+; CHECK-NEXT:    subq $72, %rsp
+; CHECK-NEXT:    .cfi_def_cfa_offset 96
+; CHECK-NEXT:    .cfi_offset %rbx, -24
+; CHECK-NEXT:    .cfi_offset %rbp, -16
+; 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:    tilerelease
+; CHECK-NEXT:    addq $72, %rsp
+; CHECK-NEXT:    .cfi_def_cfa_offset 24
+; CHECK-NEXT:    popq %rbx
+; CHECK-NEXT:    .cfi_def_cfa_offset 16
+; CHECK-NEXT:    popq %rbp
+; CHECK-NEXT:    .cfi_def_cfa_offset 8
+; 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() local_unnamed_addr
+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:    movabsq $64, %rax
 ; CHECK-NEXT:    tileloadd 2048(%rsp,%rax), %tmm5 # 1024-byte Folded Reload
 ; CHECK-NEXT:    tdpbssd %tmm4, %tmm6, %tmm5
 ; CHECK-NEXT:    movl $buf+2048, %eax
 ; CHECK-NEXT:    tilestored %tmm5, (%rax,%rcx)
+; CHECK-NEXT:    tilerelease
 ; 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
   %c = 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)
   br i1 undef, label %if.true, label %if.false
 if.true:
   %a1 = 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)
   %b1 = 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)
   %d1 = tail call x86_amx @llvm.x86.tdpbssd.internal(i16 %0, i16 %1, i16 8, x86_amx %c, x86_amx %a1, x86_amx %b1)
   tail call void (...) @foo()

 ; 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
 ; CHECK-NEXT:    cmpw $100, %bp
 ; CHECK-NEXT:    jl .LBB1_2
 ; CHECK-NEXT:  .LBB1_3: # %exit
 ; CHECK-NEXT:    addq $3032, %rsp # imm = 0xBD8
 ; CHECK-NEXT:    popq %rbx
 ; CHECK-NEXT:    popq %r14
 ; CHECK-NEXT:    popq %r15
 ; CHECK-NEXT:    popq %rbp
-; CHECK-NEXT:    tilerelease
 ; CHECK-NEXT:    vzeroupper
 ; CHECK-NEXT:    retq
 entry:
   %t5 = tail call x86_amx @llvm.x86.tilezero.internal(i16 8, i16 8)
   br i1 undef, label %loop.header, label %exit
 
 loop.header:
   %ivphi = phi i16 [0, %entry], [%iv, %loop.latch]

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

 ; CHECK-NEXT:    movabsq $64, %rax
 ; CHECK-NEXT:    tileloadd 896(%rsp,%rax), %tmm7 # 1024-byte Folded Reload
 ; CHECK-NEXT:    tdpbssd %tmm7, %tmm1, %tmm3
 ; CHECK-NEXT:    tdpbssd %tmm4, %tmm3, %tmm2
 ; CHECK-NEXT:    tdpbssd %tmm5, %tmm2, %tmm0
 ; CHECK-NEXT:    movl $buf, %eax
 ; CHECK-NEXT:    movl $32, %ecx
 ; CHECK-NEXT:    tilestored %tmm0, (%rax,%rcx)
-; CHECK-NEXT:    addq $2936, %rsp # imm = 0xB78
 ; CHECK-NEXT:    tilerelease
+; CHECK-NEXT:    addq $2936, %rsp # imm = 0xB78
 ; CHECK-NEXT:    vzeroupper
 ; CHECK-NEXT:    retq
   %4 = tail call x86_amx @llvm.x86.tileloadd64.internal(i16 %1, i16 %2, i8* getelementptr inbounds ([1024 x i8], [1024 x i8]* @buf, i64 0, i64 0), i64 32)
   %5 = tail call x86_amx @llvm.x86.tileloadd64.internal(i16 %1, i16 %2, i8* getelementptr inbounds ([1024 x i8], [1024 x i8]* @buf, i64 0, i64 0), i64 32)
   %6 = tail call x86_amx @llvm.x86.tileloadd64.internal(i16 %1, i16 %2, i8* getelementptr inbounds ([1024 x i8], [1024 x i8]* @buf, i64 0, i64 0), i64 32)
   %7 = tail call x86_amx @llvm.x86.tileloadd64.internal(i16 %2, i16 %2, i8* getelementptr inbounds ([1024 x i8], [1024 x i8]* @buf, i64 0, i64 0), i64 32)
   %8 = tail call x86_amx @llvm.x86.tileloadd64.internal(i16 %2, i16 %2, i8* getelementptr inbounds ([1024 x i8], [1024 x i8]* @buf, i64 0, i64 0), i64 32)
   %9 = tail call x86_amx @llvm.x86.tileloadd64.internal(i16 %2, i16 %2, i8* getelementptr inbounds ([1024 x i8], [1024 x i8]* @buf, i64 0, i64 0), i64 32)