[X86][AMX] Hoist ldtilecfg

[pengfei]

Summary

The previous code calculated the first ldtilecfg by dominating all AMX registers' def. This may result in the ldtilecfg being inserted into a loop.

This patch try to calculate the nearest point where post dominats all shapes of AMX registers.

[pengfei]

Since D98845 is landed. I'd like to do the ldtilecfg hoist together with this patch. WIP.

[pengfei]

Implement ldtilecfg hoist.

[pengfei]

Minor fix.

[pengfei]

Fix a bug in calculating HasAMXBeforeCall.

[pengfei]

Minor fix.

[LuoYuanke]

llvm/lib/Target/X86/X86PreTileConfig.cpp

8–9

Since we change the algorithm, we need to update the pass description.

97

In line 68 Pos is size_t, but the input is unsigned. Need to align the type?

97–128

We can initialize at the beginning. Maybe at runOnMachineFunction().

112

Better to comment on each condition.

call
here we need to reload tile config.
amx instr

115

Does HasAMXBeforeCall mean amx config register live in MBB?

127

PostDominateAllShapes.

128

This looks to represent the tile register live through MBB in the code. Basically we need to reconfig after last call if the config register live out the MBB.

155

Does HasAMXBeforeCall mean amx tile config register live through the MBB?

164

Maybe rename it as "update(MachineBasicBlock *Pred, MachineBasicBlock *MBB)"

[pengfei]

Address Yuanke and Xiang's comments.

llvm/lib/Target/X86/X86PreTileConfig.cpp

128

Yes, But it's equal to requiring instead of describing tile config register live into MBB.

    +----+
BB0 |Cfg |
    +----+
    /    \
+----+  +----+ 
|Call|  |AMX |
|AMX |  |Call|
+----+  +----+
  BB1     BB2

For example, tilecfg lives out BB0 and lives in BB1 and BB2, but we require it lives in for BB2 only. So,

BB1.NeedTileCfgLiveIn = false;
BB2.NeedTileCfgLiveIn = true;

164

How about MBB and Succ, I'd like to use MBB to respect the current node that is operating.
It's more clear to me if we consider UpdateShapeDominators is more independent with DFSPredsUpdater.

[pengfei]

Minor fix.

[LuoYuanke]

Perhaps we need more comments and more test cases (maybe in a sperate file) to cover those scenario.

llvm/lib/Target/X86/X86PreTileConfig.cpp

149

Better to split the function into 2 function. One is isAMXInstruction(), the other is collectShapeInfo().

164

Better to add comments to say shape is defined in Succ, so the pred won't postdominate the shape.

166

Here need comments to explain why the except is needed.

174

Add comment to explain why the loop is needed, and how to find a position to insert.

176

Add comment to explain why it fails.

178

If possible, better add a test case for it.

190

Is it possible that LastShape is null?

[pengfei]

Address Yuanke's comments.

llvm/lib/Target/X86/X86PreTileConfig.cpp

178

Will do later.

190

Yes, compare with MIRef() is allowed. The Pos of it is 0. So the condition is false here.

[pengfei]

Minor fix.
Add a test to verify the sink code.

[pengfei]

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

139

This seems another bug, I'll investigate it.

[pengfei]

Fix a bug found by test.

[pengfei]

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

139

Fixed.

[pengfei]

Refactor.

[pengfei]

Fixed the problem when the sink need to be forked. I.e.

     +------+
     |Entry | BB0
     +------+
     /     \
+------+  +------+
|Shape1|  |Shape2| BB2
+------+  +------+
BB1  \       /
     +------+
     | AMX  | BB3
     +------+

If BB1 and BB2 don't have a call, we will try to insert ldtilecfg from BB0.
Since BB0 doesn't dominate all shapes, we will sink the insert point to its successors.
The previous code has limitation that only one of BB0's successors has NeedTileCfgLiveIn = true.
This update fixes the problem without major changes. So I amended it instead of putting into a new revision.

[pengfei]

The algorithm for updating shape postdominate BBs is buggy. For a given ShapeBB, clear all its predecessors flag is not enough since its unreachable BBs are also need to clear.
I have thought out a new method but need major refactor. Stay tuned~

[pengfei]

The algorithm for updating shape postdominate BBs is buggy. For a given ShapeBB, clear all its predecessors flag is not enough since its unreachable BBs are also need to clear.

I have thought out a new method but need major refactor. Stay tuned~

Worked it out.

[pengfei]

A litter format & comments refactor.

[pengfei]

Fix a bug in collectShapeInfo.

[LuoYuanke]

llvm/lib/Target/X86/X86PreTileConfig.cpp

16–18

Do you mean all tile registers are caller saved?

185

Do we need to recursively handle phi instruction? In below example, record "def shape" instruction?

def shape

\
phi
  \
 phi

253

So we would try to insert ldtilecfg at the entry if there is not shape dependency issue?

294

May update the comments. All shapes are reachable.

305

Can we sink the ldtilecfg in below case? We need insert ldtilecfg to B4.

     Entry
   /         \
 B1         B2
call        call
   \           /
       B3
         |
    B4 shape defs
         |
      AMX

329

Is this semicolon necessary?

[pengfei]

Address Yuanke's comments.

llvm/lib/Target/X86/X86PreTileConfig.cpp

16–18

No, tile registers are visible to RA, we don't need to restore them ourselves.
We are mainly doing the restore for the shape register in this pass.

185

I don't think so. Each phi will be iterated once. So we just need to peek its latest level and ignore the phi in it.

s1 s2
 \ /
  p1  s3  ...
   \  /   /
    p2   p3
     \  /
      p4

For exampe, if we meet p4 firstly, we find all its defs are phis, we don't do anything. Then, we meet p2 in another iteration and record s3. Then p1 etc. The order doesn't metter, you can see all the real shape defs will be recorded without recursion.

253

Yes. Here the insert points are candidates. We will sink them based on the shape defs. We will insert at the entry if shape dependency satisfies.

305

Yes. We will sink two candidates from B1 anf B2. (If there's no call in B2, the candidates are B1 and Entry) We avoid BB being multi visited. So the second sink will stop at B3. This first sink will cross B3 and reach B4.

329

Good catch. It's residue from previous lambda expression.

[xiangzhangllvm]

llvm/lib/Target/X86/X86PreTileConfig.cpp

55

use C++ comments // between line 54,55

67

use getMI() or directly p->MI is more readable.

108

I suggest use tile operand to judge if a MI is AMX Instruction or not. and spilt the collection of shape to another special function.

180

here may be a constant imm in PHI's incomings.

185

1. MI2->isMoveImmediate() may need to be handled, because the imm may always keep in a register after PHI.

In other wards, if AMX MI use the imm shape, they must use the register after PHI.

2. the recursive PHI case is very rare, we can let it easy, not handle it, there is not correction problem.

[xiangzhangllvm]

llvm/lib/Target/X86/X86PreTileConfig.cpp

275

Not sure the non-bottom BBs in loop can both with "fall through + continue".

[pengfei]

Address Xiang's comments.

llvm/lib/Target/X86/X86PreTileConfig.cpp

108

But we still need to judge if it is PHI. I don't know if there're other exceptions.
We have split it to collectShapeInfo.

180

Seems IR parser will normalize this case, so we don't need to handle it.

185

1. The PostRA pass can handle this case.
2. We have covered this case here.

[pengfei]

llvm/lib/Target/X86/X86PreTileConfig.cpp

275

non-bottom BBs can propagate until the bottom BB.

[xiangzhangllvm]

llvm/lib/Target/X86/X86PreTileConfig.cpp

108

Yes, PHI is exception, the exceptions is very few, more better than list all the AMX opcodes.

185

I don't think "Each phi will be iterated once."

MRI->def_operands(ShapeReg) only contains the shape reg used in AMX MI

s1 s2
 \ /
  p1  s3  ...
   \  /   /
    p2   p3
     \  /
      p4
 then AMX MI use p4 as this shape,

only p4 in def_operands(ShapeReg).

185

But here missed the def shape BB, it is not meet the following logic of use "shapeBBNumber == ShapeReachedCount" to judge the BB is all shape def reachable.

   reg=imm        ...
   /     \         / 
amx          phi
              |
             BB

[pengfei]

llvm/lib/Target/X86/X86PreTileConfig.cpp

185

You are right, I excluded phi at first. I'll refactor for it. Thanks!

[xiangzhangllvm]

llvm/lib/Target/X86/X86PreTileConfig.cpp

320

I am a little puzzle here, The I.MBB is all shape def reachable, but how can we make sure the BBVisitedInfo[I.MBB].LastShape is shape def reachable ? They may be not same BB.

[pengfei]

llvm/lib/Target/X86/X86PreTileConfig.cpp

320

LastShape is the last shape def in a BB. It is only used when its BB is all shape def reachable. I.e. insert ldtilecfg after it. But it can be null if the BB has no shape def, which means we can insert at the begining of the BB.

[pengfei]

Address Xiang's comment.
Iterate all PHIs recursively to find all shapes.

[LuoYuanke]

llvm/lib/Target/X86/X86PreTileConfig.cpp

182

Better to add test case to cover the logic.

273

Better to add test case for it.

[pengfei]

Address Yuanke's comments.

Add test5 for checking shape peek's loop break.
Add test6 for checking shape reachability's loop break.

[pengfei]

Fix a silly bug that uses != in case == and a bug when records shape for phi.

This is found by work with PostRA implementation. This should be found by case amx-ldtilecfg-insert.ll:, but happened to get expected output with these 2 bugs.

[pengfei]

Fix crush when build with -g.

[pengfei]

Add back a condition deleted last time by mistaken.

[LuoYuanke]

llvm/lib/Target/X86/X86PreTileConfig.cpp

109

!MI.isPseudo()?

[xiangzhangllvm]

llvm/lib/Target/X86/X86PreTileConfig.cpp

174

This should be calculated in ShapeNum, for AMX always use reg as its shape after PHI.

[pengfei]

Refactor for excluding DBG_VALUE case.

llvm/lib/Target/X86/X86PreTileConfig.cpp

109

Yeah, I meant to. Surprisingly, we didn't set this flag when we define PseudoI. So the tests happened to pass. :)

174

The immediate value is specially handled in PostRA pass. Since the immediate value is always defined adjoining its use, we cannot calculate it as shape def. Otherwise, they will intersect with other AMX instructions.

[LuoYuanke]

LGTM. Thanks!

llvm/lib/Target/X86/X86PreTileConfig.cpp

 //===-- X86PreTileConfig.cpp - Tile Register Pre-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 pre-config the shape of AMX register
-/// AMX register need to be configured before use. The shape of AMX register
-/// is encoded in the 1st and 2nd machine operand of AMX pseudo instructions.
-/// The pldtilecfg is to config tile registers. It should dominator all AMX
-/// instructions. The pldtilecfg produce a virtual cfg register and the cfg
-/// register is used by all AMX instructions.
-/// This pass is to find the common dominator of all AMX instructions and
-/// insert the pldtilecfg instruction. Besides the cfg register that pldtilecfg
-/// produces is inserted as the last operand of each AMX instruction. We use
-/// this scheme to model the def-use relationship between AMX config instruction
-/// and other AMX instructions. Below is an example.
+/// \file Pass to pre-config the shapes of AMX registers
+/// AMX register needs to be configured before use. The shapes of AMX register
+/// are encoded in the 1st and 2nd machine operand of AMX pseudo instructions.
 ///
-///                        ----B1----
-///                       /           \
-///                      /             \
-///                    B2               B3
-///    %1:tile = PTILELOADDV        %2:tile = PTILELOADDV
+/// The instruction ldtilecfg is used to config the shapes. It must dominate all
+/// AMX instructions and postdominate all variable shapes.
 ///
-///  is transformed to
+/// The shape register is caller saved according to ABI. We need to insert
+/// ldtilecfg again after the call instruction if callee clobbers any AMX
+/// registers.
 ///
-///                            B1
-///                 %25:tilecfg = PLDTILECFG
-///                       /           \
-///                      /             \
-///  %1:tile = PTILELOADDV %25    %2:tile = PTILELOADDV %25
+/// This pass calculates all points that ldtilecfg need to be inserted to and
+/// insert them. It reports error if the domination conditions aren't met.
 //
 //===----------------------------------------------------------------------===//
 
 #include "X86.h"
 #include "X86InstrBuilder.h"
 #include "X86RegisterInfo.h"
 #include "X86Subtarget.h"
-#include "llvm/CodeGen/MachineDominators.h"
 #include "llvm/CodeGen/MachineFunctionPass.h"
 #include "llvm/CodeGen/MachineInstr.h"
+#include "llvm/CodeGen/MachineLoopInfo.h"
 #include "llvm/CodeGen/MachineRegisterInfo.h"
 #include "llvm/CodeGen/Passes.h"
 #include "llvm/CodeGen/TargetInstrInfo.h"
 #include "llvm/CodeGen/TargetRegisterInfo.h"
 #include "llvm/CodeGen/TileShapeInfo.h"
 #include "llvm/InitializePasses.h"
 
 using namespace llvm;
 
 #define DEBUG_TYPE "tile-pre-config"
+#define ASSERT_VALID_COMPARE                                                   \
+  assert((!MBB || !RHS.MBB || MBB == RHS.MBB) &&                               \
+         "Cannot compare between different BBs");
+#define REPORT_CONFIG_FAIL                                                     \
+  report_fatal_error(                                                          \
+      MF.getName() +                                                           \
+      ": Failed to config tile register, please define the shape earlier");
 
 namespace {
 
+struct MIRef {
+  MachineInstr *MI = nullptr;
+  MachineBasicBlock *MBB = nullptr;
+  size_t Pos = 0; /* A virtual position for instr will be inserted after MI */
+  MIRef() = default;
+  MIRef(MachineBasicBlock *MBB) : MBB(MBB) {
+    for (auto I = MBB->begin(), E = MBB->end(); I != E && I->isPHI();
+         ++I, ++Pos)
+      MI = &*I;
+  }
+  MIRef(MachineInstr *MI, MachineBasicBlock *MBB)
+      : MI(MI), MBB(MBB),
+        Pos(std::distance(MBB->instr_begin(), ++MI->getIterator())) {}
+  MIRef(MachineInstr *MI, MachineBasicBlock *MBB, size_t Pos)
+      : MI(MI), MBB(MBB), Pos(Pos) {}
+  MachineInstr *operator->() { return MI; }
+  operator bool() const { return MBB != nullptr; }
+  bool operator==(const MIRef &RHS) const {
+    return MI == RHS.MI && MBB == RHS.MBB;
+  }
+  bool operator<(const MIRef &RHS) const {
+    ASSERT_VALID_COMPARE;
+    return Pos < RHS.Pos;
+  }
+  bool operator>(const MIRef &RHS) const {
+    ASSERT_VALID_COMPARE;
+    return Pos > RHS.Pos;
+  }
+};
+
+struct BBInfo {
+  MIRef FirstAMX;
+  MIRef LastCall;
+  MIRef LastShape;
+  bool NeedTileCfgLiveIn = false;
+  unsigned ShapeReachedCount = 0;
+};
+
 class X86PreTileConfig : public MachineFunctionPass {
-  // context
-  MachineFunction *MF = nullptr;
-  const X86Subtarget *ST = nullptr;
-  const TargetRegisterInfo *TRI;
-  const TargetInstrInfo *TII;
-  MachineDominatorTree *DomTree = nullptr;
-  MachineRegisterInfo *MRI = nullptr;
+  MachineRegisterInfo *MRI;
+  SmallSet<MIRef, 8> CfgNeedInsert;
+  SmallSet<MachineBasicBlock *, 8> ShapeBBs;
+  SmallVector<MachineBasicBlock *, 8> CfgLiveInBBs;
+  DenseMap<MachineBasicBlock *, BBInfo> BBVisitedInfo;
 
-  MachineInstr *getTileConfigPoint();
+  /// Check if the callee will clobber AMX registers.
+  bool isDestructiveCall(MachineInstr &MI, BitVector UsableRegs) {
+    auto Iter = llvm::find_if(
+        MI.operands(), [](MachineOperand &MO) { return MO.isRegMask(); });
+    if (Iter == MI.operands_end())
+      return false;
+    UsableRegs.clearBitsInMask(Iter->getRegMask());
+    return !UsableRegs.none();
+  }
+
+  /// Check if MI is AMX instruction.
+  bool isAMXInstruction(MachineInstr &MI) {
+    switch (MI.getOpcode()) {
+    default:
+      return false;
+    case X86::PTILESTOREDV:
+      collectShapeInfo(ShapeT(&MI.getOperand(0), &MI.getOperand(1), MRI));
+      return true;
+    case X86::PTILELOADDV:
+    case X86::PTDPBSSDV:
+    case X86::PTDPBSUDV:
+    case X86::PTDPBUSDV:
+    case X86::PTDPBUUDV:
+    case X86::PTILEZEROV:
+    case X86::PTDPBF16PSV:
+      collectShapeInfo(ShapeT(&MI.getOperand(1), &MI.getOperand(2), MRI));
+      return true;
+    }
+  }
+
+  /// Collect the shape def information for later use.
+  void collectShapeInfo(ShapeT Shape);
 
 public:
   X86PreTileConfig() : MachineFunctionPass(ID) {}
 
   /// Return the pass name.
   StringRef getPassName() const override {
     return "Tile Register Pre-configure";
   }
 
   /// X86PreTileConfig analysis usage.
-  void getAnalysisUsage(AnalysisUsage &AU) const override;
+  void getAnalysisUsage(AnalysisUsage &AU) const override {
+    AU.setPreservesAll();
+    AU.addRequired<MachineLoopInfo>();
+    MachineFunctionPass::getAnalysisUsage(AU);
+  }
 
-  /// Perform register allocation.
-  bool runOnMachineFunction(MachineFunction &mf) override;
+  /// Clear MF related structures.
+  void releaseMemory() override {
+    ShapeBBs.clear();
+    CfgLiveInBBs.clear();
+    CfgNeedInsert.clear();
+    BBVisitedInfo.clear();
+  }
+
+  /// Perform ldtilecfg instructions inserting.
+  bool runOnMachineFunction(MachineFunction &MF) override;
 
   static char ID;
 };
 
 } // end anonymous namespace
 
 char X86PreTileConfig::ID = 0;
 
 INITIALIZE_PASS_BEGIN(X86PreTileConfig, "tilepreconfig",
                       "Tile Register Pre-configure", false, false)
-INITIALIZE_PASS_DEPENDENCY(MachineDominatorTree)
+INITIALIZE_PASS_DEPENDENCY(MachineLoopInfo)
 INITIALIZE_PASS_END(X86PreTileConfig, "tilepreconfig",
                     "Tile Register Pre-configure", false, false)
 
-void X86PreTileConfig::getAnalysisUsage(AnalysisUsage &AU) const {
-  AU.setPreservesAll();
-  AU.addRequired<MachineDominatorTree>();
-  MachineFunctionPass::getAnalysisUsage(AU);
-}
-
-static void buildConfigMI(MachineBasicBlock::iterator MI, int FrameIdx,
-                          const TargetInstrInfo *TII, MachineRegisterInfo *MRI,
-                          const X86Subtarget *ST) {
-  auto *MBB = MI->getParent();
-
-  // Zero stack slot.
-  if (ST->hasAVX512()) {
-    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);
-  } else if (ST->hasAVX2()) {
-    Register Ymm = MRI->createVirtualRegister(&X86::VR256RegClass);
-    BuildMI(*MBB, MI, DebugLoc(), TII->get(X86::VPXORYrr), Ymm)
-        .addReg(Ymm, RegState::Undef)
-        .addReg(Ymm, RegState::Undef);
-    addFrameReference(BuildMI(*MBB, MI, DebugLoc(), TII->get(X86::VMOVUPSYmr)),
-                      FrameIdx)
-        .addReg(Ymm);
-    addFrameReference(BuildMI(*MBB, MI, DebugLoc(), TII->get(X86::VMOVUPSYmr)),
-                      FrameIdx, 32)
-        .addReg(Ymm);
-  } else {
-    assert(ST->hasSSE2() && "AMX should assume SSE2 enabled");
-    Register Xmm = MRI->createVirtualRegister(&X86::VR128RegClass);
-    BuildMI(*MBB, MI, DebugLoc(), TII->get(X86::PXORrr), Xmm)
-        .addReg(Xmm, RegState::Undef)
-        .addReg(Xmm, RegState::Undef);
-    addFrameReference(BuildMI(*MBB, MI, DebugLoc(), TII->get(X86::MOVUPSmr)),
-                      FrameIdx)
-        .addReg(Xmm);
-    addFrameReference(BuildMI(*MBB, MI, DebugLoc(), TII->get(X86::MOVUPSmr)),
-                      FrameIdx, 16)
-        .addReg(Xmm);
-    addFrameReference(BuildMI(*MBB, MI, DebugLoc(), TII->get(X86::MOVUPSmr)),
-                      FrameIdx, 32)
-        .addReg(Xmm);
-    addFrameReference(BuildMI(*MBB, MI, DebugLoc(), TII->get(X86::MOVUPSmr)),
-                      FrameIdx, 48)
-        .addReg(Xmm);
-  }
-
-  // build psuedo ldtilecfg
-  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:
-  case X86::PTDPBSSDV:
-  case X86::PTDPBSUDV:
-  case X86::PTDPBUSDV:
-  case X86::PTDPBUUDV:
-  case X86::PTILEZEROV:
-  case X86::PTDPBF16PSV:
-    MachineOperand &MO1 = const_cast<MachineOperand &>(MI.getOperand(1));
-    MachineOperand &MO2 = const_cast<MachineOperand &>(MI.getOperand(2));
-    ShapeT Shape(&MO1, &MO2, MRI);
-    return Shape;
-  }
-}
-
-MachineInstr *X86PreTileConfig::getTileConfigPoint() {
-  DenseMap<Register, ShapeT> PhysShapeInfo;
-  MachineBasicBlock *MBB = nullptr;
-  DenseSet<const MachineInstr *> MIs;
-  for (unsigned i = 0, e = MRI->getNumVirtRegs(); i != e; ++i) {
-    Register VirtReg = Register::index2VirtReg(i);
-    if (MRI->reg_nodbg_empty(VirtReg))
+void X86PreTileConfig::collectShapeInfo(ShapeT Shape) {
+  for (auto *ShapeMO : {Shape.getRow(), Shape.getCol()}) {
+    Register ShapeReg = ShapeMO->getReg();
+    for (MachineOperand &MO : MRI->def_operands(ShapeReg)) {
+      if (MO.isUndef())
         continue;
-    const TargetRegisterClass &RC = *MRI->getRegClass(VirtReg);
-    if (RC.getID() != X86::TILERegClassID)
+      MachineInstr *MI = MO.getParent();
+      if (MI->isMoveImmediate())
         continue;
-
-    // Find the common dominator for all MI that define tile register.
-    for (const MachineOperand &MO : MRI->def_operands(VirtReg)) {
-      if (MO.isUndef())
+      if (MI->isPHI()) {
+        for (unsigned Index = 0; Index < MI->getNumOperands() / 2; ++Index) {
+          Register DefReg = MI->getOperand(Index * 2 + 1).getReg();
+          for (MachineOperand &MO2 : MRI->def_operands(DefReg)) {
+            if (MO2.isUndef())
               continue;
-      const auto *MI = MO.getParent();
-      // PHI or IMPLICIT_DEF instructiion.
-      // There must be a input tile before PHI instruction.
-      if (MI->isTransient())
+            MachineInstr *DefMI = MO2.getParent();
+            if (DefMI->isMoveImmediate() || DefMI->isPHI())
               continue;
-      if (!MBB)
-        MBB = const_cast<MachineBasicBlock *>(MI->getParent());
-      MBB = DomTree->findNearestCommonDominator(
-          MBB, const_cast<MachineBasicBlock *>(MI->getParent()));
-
-      // Collect the instructions that define shape.
-      ShapeT Shape = getShape(*MI, MRI);
-      std::array<MachineOperand *, 2> ShapeMOs = {Shape.getRow(),
-                                                  Shape.getCol()};
-      for (auto *ShapeMO : ShapeMOs) {
-        Register ShapeReg = ShapeMO->getReg();
-        for (const MachineOperand &MO : MRI->def_operands(ShapeReg)) {
-          const auto *ShapeMI = MO.getParent();
-          MIs.insert(ShapeMI);
-        }
-      }
+            MachineBasicBlock *MBB = MI->getParent();
+            MIRef MIR(MI, MBB);
+            if (BBVisitedInfo[MBB].LastShape < MIR)
+              BBVisitedInfo[MBB].LastShape = MIR;
+            ShapeBBs.insert(DefMI->getParent());
           }
         }
-  if (!MBB)
-    return nullptr;
-  // This pass is before the pass of eliminating PHI node, so it
-  // is in SSA form.
-  assert(MRI->isSSA() && "Not SSA form in pre-tile config");
-  // Shape def should dominate tile config MBB.
-  //    def s           s1    s2
-  //     / \             \   /
-  //    /   \             \ /
-  //  conf               s3=phi(s1,s2)
-  //                       |
-  //                       c
-  //
-  for (const auto *MI : MIs) {
-    const MachineBasicBlock *ShapeMBB = MI->getParent();
-    if (DomTree->dominates(ShapeMBB, MBB))
-      continue;
-    if (MI->isMoveImmediate())
         continue;
-    report_fatal_error(MF->getName() + ": Failed to config tile register, "
-                                       "please define the shape earlier");
       }
-
-  // ldtilecfg should be inserted after the MI that define the shape.
-  MachineBasicBlock::reverse_instr_iterator I, E;
-  for (I = MBB->instr_rbegin(), E = MBB->instr_rend(); I != E; ++I) {
-    auto *MI = &*I;
-    if (MIs.count(MI) && (!MI->isMoveImmediate()))
-      break;
+      MachineBasicBlock *MBB = MI->getParent();
+      MIRef MIR(MI, MBB);
+      if (BBVisitedInfo[MBB].LastShape < MIR)
+        BBVisitedInfo[MBB].LastShape = MIR;
+      ShapeBBs.insert(MBB);
     }
-  MachineBasicBlock::iterator MII;
-  if (I == E)
-    MII = MBB->getFirstNonPHI();
-  else {
-    MII = MachineBasicBlock::iterator(&*I);
-    MII++;
   }
-  return &*MII;
 }
 
-static bool isAMXInstruction(MachineBasicBlock::iterator MII) {
-  switch (MII->getOpcode()) {
-  default:
-    return false;
-  case X86::PTILELOADDV:
-  case X86::PTILESTOREDV:
-  case X86::PTDPBSSDV:
-  case X86::PTDPBSUDV:
-  case X86::PTDPBUSDV:
-  case X86::PTDPBUUDV:
-  case X86::PTILEZEROV:
-  case X86::PTDPBF16PSV:
-    return true;
+bool X86PreTileConfig::runOnMachineFunction(MachineFunction &MF) {
+  MRI = &MF.getRegInfo();
+  const X86Subtarget &ST = MF.getSubtarget<X86Subtarget>();
+  const TargetInstrInfo *TII = ST.getInstrInfo();
+  const TargetRegisterInfo *TRI = ST.getRegisterInfo();
+  const MachineLoopInfo *MLI = &getAnalysis<MachineLoopInfo>();
+  const TargetRegisterClass *RC = TRI->getRegClass(X86::TILERegClassID);
+
+  BitVector AMXRegs(TRI->getNumRegs());
+  for (unsigned I = 0; I < RC->getNumRegs(); I++)
+    AMXRegs.set(X86::TMM0 + I);
+
+  // Iterate MF to collect information.
+  for (auto &MBB : MF) {
+    size_t Pos = 0;
+    for (auto &MI : MBB) {
+      ++Pos;
+      if (isAMXInstruction(MI)) {
+        // If there's call before the AMX, we need to reload tile config.
+        if (BBVisitedInfo[&MBB].LastCall)
+          CfgNeedInsert.insert(BBVisitedInfo[&MBB].LastCall);
+        else /* Otherwise, we need tile config to live in this BB. */
+          BBVisitedInfo[&MBB].NeedTileCfgLiveIn = true;
+        // Always record the first AMX in case there's shape def after it.
+        if (!BBVisitedInfo[&MBB].FirstAMX)
+          BBVisitedInfo[&MBB].FirstAMX = MIRef(&MI, &MBB, Pos);
+      } else if (MI.isCall() && isDestructiveCall(MI, AMXRegs)) {
+        // Record the call only if the callee clobbers all AMX registers.
+        BBVisitedInfo[&MBB].LastCall = MIRef(&MI, &MBB, Pos);
+      }
+    }
+    if (BBVisitedInfo[&MBB].NeedTileCfgLiveIn) {
+      if (&MBB == &MF.front())
+        CfgNeedInsert.insert(MIRef(&MBB));
+      else
+        CfgLiveInBBs.push_back(&MBB);
     }
   }
 
-struct BBInfo {
-  bool HasAMX = false;
-  bool HasCallBeforeAMX = false;
-  bool HasAMXBeforeCallInSuccs = false;
-  MachineInstr *LastCall = nullptr;
-
-  BBInfo() = default;
-  BBInfo(SmallSet<MachineInstr *, 8> &CfgNeedInsert, 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 (LastCall)
-          CfgNeedInsert.insert(LastCall);
-      } else if (MII->isCall()) {
-        LastCall = &*MII;
-        if (!HasAMX)
-          HasCallBeforeAMX = true;
+  // Update NeedTileCfgLiveIn for predecessors.
+  while (!CfgLiveInBBs.empty()) {
+    MachineBasicBlock *MBB = CfgLiveInBBs.pop_back_val();
+    for (auto *Pred : MBB->predecessors()) {
+      if (BBVisitedInfo[Pred].LastCall) {
+        CfgNeedInsert.insert(BBVisitedInfo[Pred].LastCall);
+      } else if (!BBVisitedInfo[Pred].NeedTileCfgLiveIn) {
+        BBVisitedInfo[Pred].NeedTileCfgLiveIn = true;
+        if (Pred == &MF.front())
+          CfgNeedInsert.insert(MIRef(Pred));
+        else
+          CfgLiveInBBs.push_back(Pred);
       }
     }
   }
-};
 
-static void reloadTileConfig(MachineInstr *MI, int FI,
-                             const TargetInstrInfo *TII,
-                             const TargetRegisterInfo *TRI) {
-  SmallSet<MachineInstr *, 8> CfgNeedInsert;
-  SmallVector<MachineBasicBlock *, 8> WorkList;
-  DenseMap<MachineBasicBlock *, BBInfo> BBVisitedInfo;
-
-  MachineBasicBlock *MBB = MI->getParent();
-  BBVisitedInfo[MBB] = BBInfo(CfgNeedInsert, MBB, MI);
+  // There's no AMX instruction if we didn't find a tile config live in point.
+  if (CfgNeedInsert.empty())
+    return false;
 
-  WorkList.push_back(MBB);
+  // Calculate how many times the ShapeBB can reach to this BB.
+  unsigned ShapeBBNum = 0;
+  for (auto *MBB : ShapeBBs) {
+    SmallSet<MachineBasicBlock *, 8> VistedBB;
+    SmallVector<MachineBasicBlock *, 8> WorkList({MBB});
     while (!WorkList.empty()) {
-    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(CfgNeedInsert, *I);
-        WorkList.push_back(*I);
-      }
-    }
+      MachineBasicBlock *MBB = WorkList.pop_back_val();
+      ++BBVisitedInfo[MBB].ShapeReachedCount;
+      for (auto *Succ : MBB->successors())
+        if (VistedBB.insert(Succ).second &&
+            (!MLI->isLoopHeader(Succ) ||
+             MLI->getLoopFor(Succ)->getBottomBlock() != MBB))
+          WorkList.push_back(Succ);
     }
-
-  WorkList.clear();
-  for (auto I : BBVisitedInfo) {
-    WorkList.push_back(I.first);
+    ++ShapeBBNum;
+  }
+
+  DebugLoc DL;
+  SmallSet<MIRef, 8> VisitedOrInserted;
+  int SS = MF.getFrameInfo().CreateStackObject(
+      ST.getTileConfigSize(), ST.getTileConfigAlignment(), false);
+
+  // Try to insert for the tile config live in points.
+  for (auto I : CfgNeedInsert) {
+    SmallSet<MIRef, 8> InsertPoints;
+    SmallVector<MIRef, 8> WorkList({I});
     while (!WorkList.empty()) {
-      MBB = WorkList.pop_back_val();
-      if (BBVisitedInfo[MBB].HasCallBeforeAMX ||
-          (!BBVisitedInfo[MBB].HasAMX &&
-           !BBVisitedInfo[MBB].HasAMXBeforeCallInSuccs))
-        continue;
-      for (auto I = MBB->pred_begin(), E = MBB->pred_end(); I != E; ++I) {
-        if (!BBVisitedInfo.count(*I) ||
-            BBVisitedInfo[*I].HasAMXBeforeCallInSuccs)
-          continue;
-        if (BBVisitedInfo[*I].LastCall)
-          CfgNeedInsert.insert(BBVisitedInfo[*I].LastCall);
-        BBVisitedInfo[*I].HasAMXBeforeCallInSuccs = true;
-        WorkList.push_back(*I);
+      MIRef I = WorkList.pop_back_val();
+      if (!VisitedOrInserted.count(I)) {
+        if (BBVisitedInfo[I.MBB].ShapeReachedCount == ShapeBBNum) {
+          // If find a BB postdominates all shapes, stop sink and try to insert.
+          InsertPoints.insert(I);
+        } else {
+          // Avoid the BB to be multi visited.
+          VisitedOrInserted.insert(I);
+          // We cannot sink it across any AMX instruction.
+          if (BBVisitedInfo[I.MBB].FirstAMX)
+            REPORT_CONFIG_FAIL;
+          // Sink the inserting point along the chain with NeedTileCfgLiveIn =
+          // true when MBB doesn't postdominate all shapes.
+          for (auto *Succ : I.MBB->successors())
+            if (BBVisitedInfo[Succ].NeedTileCfgLiveIn)
+              WorkList.push_back(MIRef(Succ));
         }
       }
     }
 
-  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);
+    // A given point might be forked due to shape conditions are not met.
+    for (MIRef I : InsertPoints) {
+      // Even MBB postdominates all shapes, we still need to check if there's
+      // AMX that intersects with shapes in the same MBB.
+      if (BBVisitedInfo[I.MBB].FirstAMX &&
+          BBVisitedInfo[I.MBB].FirstAMX < BBVisitedInfo[I.MBB].LastShape)
+        REPORT_CONFIG_FAIL;
+      // Make sure we insert ldtilecfg after the last shape def in MBB.
+      if (I < BBVisitedInfo[I.MBB].LastShape)
+        I = BBVisitedInfo[I.MBB].LastShape;
+      // There're chances the MBB is sunk more than once. Record it to avoid
+      // multi insert.
+      if (VisitedOrInserted.insert(I).second) {
+        auto II = I.MI ? I->getIterator() : I.MBB->instr_begin();
+        addFrameReference(BuildMI(*I.MBB, ++II, DL, TII->get(X86::LDTILECFG)),
+                          SS);
       }
     }
+  };
 
-bool X86PreTileConfig::runOnMachineFunction(MachineFunction &mf) {
-  MF = &mf;
-  MRI = &mf.getRegInfo();
-  ST = &mf.getSubtarget<X86Subtarget>();
-  TRI = ST->getRegisterInfo();
-  TII = mf.getSubtarget().getInstrInfo();
-  DomTree = &getAnalysis<MachineDominatorTree>();
+  // Zero stack slot.
+  MachineBasicBlock &MBB = MF.front();
+  MachineInstr *MI = &*MBB.begin();
+  if (ST.hasAVX512()) {
+    Register Zmm = MRI->createVirtualRegister(&X86::VR512RegClass);
+    BuildMI(MBB, MI, DL, TII->get(X86::VPXORDZrr), Zmm)
+        .addReg(Zmm, RegState::Undef)
+        .addReg(Zmm, RegState::Undef);
+    addFrameReference(BuildMI(MBB, MI, DL, TII->get(X86::VMOVUPSZmr)), SS)
+        .addReg(Zmm);
+  } else if (ST.hasAVX2()) {
+    Register Ymm = MRI->createVirtualRegister(&X86::VR256RegClass);
+    BuildMI(MBB, MI, DL, TII->get(X86::VPXORYrr), Ymm)
+        .addReg(Ymm, RegState::Undef)
+        .addReg(Ymm, RegState::Undef);
+    addFrameReference(BuildMI(MBB, MI, DL, TII->get(X86::VMOVUPSYmr)), SS)
+        .addReg(Ymm);
+    addFrameReference(BuildMI(MBB, MI, DL, TII->get(X86::VMOVUPSYmr)), SS, 32)
+        .addReg(Ymm);
+  } else {
+    assert(ST.hasSSE2() && "AMX should assume SSE2 enabled");
+    Register Xmm = MRI->createVirtualRegister(&X86::VR128RegClass);
+    BuildMI(MBB, MI, DL, TII->get(X86::PXORrr), Xmm)
+        .addReg(Xmm, RegState::Undef)
+        .addReg(Xmm, RegState::Undef);
+    addFrameReference(BuildMI(MBB, MI, DL, TII->get(X86::MOVUPSmr)), SS)
+        .addReg(Xmm);
+    addFrameReference(BuildMI(MBB, MI, DL, TII->get(X86::MOVUPSmr)), SS, 16)
+        .addReg(Xmm);
+    addFrameReference(BuildMI(MBB, MI, DL, TII->get(X86::MOVUPSmr)), SS, 32)
+        .addReg(Xmm);
+    addFrameReference(BuildMI(MBB, MI, DL, TII->get(X86::MOVUPSmr)), SS, 48)
+        .addReg(Xmm);
+  }
 
-  MachineInstr *MI = getTileConfigPoint();
-  if (!MI)
-    return false;
-  unsigned Size = ST->getTileConfigSize();
-  Align Alignment = ST->getTileConfigAlignment();
-  int SS = mf.getFrameInfo().CreateStackObject(Size, Alignment, false);
-  buildConfigMI(MI, SS, TII, MRI, ST);
-  reloadTileConfig(MI, SS, TII, TRI);
   return true;
 }
 
 FunctionPass *llvm::createX86PreTileConfigPass() {
   return new X86PreTileConfig();
 }

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

 ; 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:    vzeroupper
+; CHECK-NEXT:    callq foo
 ; 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:    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:    callq foo
 ; 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

 ; CHECK:       # %bb.0:
 ; CHECK-NEXT:    pushq %rbp
 ; CHECK-NEXT:    pushq %r15
 ; CHECK-NEXT:    pushq %r14
 ; CHECK-NEXT:    pushq %r12
 ; CHECK-NEXT:    pushq %rbx
 ; CHECK-NEXT:    subq $3024, %rsp # imm = 0xBD0
 ; CHECK-NEXT:    movl %edi, %ebx
+; CHECK-NEXT:    vpxord %zmm0, %zmm0, %zmm0
+; CHECK-NEXT:    vmovdqu64 %zmm0, {{[0-9]+}}(%rsp)
 ; CHECK-NEXT:    movl $buf, %r14d
 ; CHECK-NEXT:    movl $32, %r15d
 ; CHECK-NEXT:    movw $8, %bp
 ; CHECK-NEXT:    movl $buf+2048, %r12d
 ; CHECK-NEXT:    .p2align 4, 0x90
 ; CHECK-NEXT:  .LBB3_1: # =>This Inner Loop Header: Depth=1
+; CHECK-NEXT:    vzeroupper
 ; CHECK-NEXT:    callq foo
 ; 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 %ebx, %ebx
 ; CHECK-NEXT:    jle .LBB3_3
 ; CHECK-NEXT:  # %bb.2: # in Loop: Header=BB3_1 Depth=1
-; CHECK-NEXT:    vpxord %zmm0, %zmm0, %zmm0
-; CHECK-NEXT:    vmovdqu64 %zmm0, {{[0-9]+}}(%rsp)
-; CHECK-NEXT:    ldtilecfg {{[0-9]+}}(%rsp)
 ; CHECK-NEXT:    tileloadd (%r14,%r15), %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, (%r12,%r15)
 ; CHECK-NEXT:    callq foo
 ; CHECK-NEXT:    jmp .LBB3_1
 ; CHECK-NEXT:  .LBB3_3:
 ; CHECK-NEXT:    addq $3024, %rsp # imm = 0xBD0
 ; CHECK-NEXT:    popq %rbx
 ; CHECK-NEXT:    popq %r12
 ; CHECK-NEXT:    popq %r14
 ; CHECK-NEXT:    popq %r15
 ; CHECK-NEXT:    popq %rbp
 ; CHECK-NEXT:    tilerelease
 ; CHECK-NEXT:    retq
 ;
 ; IPRA-LABEL: test_loop2:
 ; 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 $8, {{[0-9]+}}(%rsp)
+; IPRA-NEXT:    movw $8, {{[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:    movl $buf+2048, %esi
 ; IPRA-NEXT:    .p2align 4, 0x90
 ; IPRA-NEXT:  .LBB3_1: # =>This Inner Loop Header: Depth=1
 ; IPRA-NEXT:    callq foo
 ; IPRA-NEXT:    testl %edi, %edi
 ; IPRA-NEXT:    jle .LBB3_3
 ; IPRA-NEXT:  # %bb.2: # in Loop: Header=BB3_1 Depth=1
-; 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:    tileloadd (%rax,%rcx), %tmm0
 ; IPRA-NEXT:    callq foo
 ; IPRA-NEXT:    tilestored %tmm0, (%rsi,%rcx)
 ; IPRA-NEXT:    callq foo
 ; IPRA-NEXT:    jmp .LBB3_1
 ; IPRA-NEXT:  .LBB3_3:
 ; IPRA-NEXT:    addq $72, %rsp
 ; IPRA-NEXT:    tilerelease

llvm/test/CodeGen/X86/AMX/amx-config.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 --check-prefix=AVX512
 ; RUN: llc < %s -mtriple=x86_64-unknown-unknown -mattr=+amx-int8 -mattr=+avx2 -verify-machineinstrs | FileCheck %s --check-prefix=AVX2
 ; RUN: llc < %s -mtriple=x86_64-unknown-unknown -mattr=+amx-int8 -verify-machineinstrs | FileCheck %s --check-prefix=SSE2
 
 @buf = dso_local global [1024 x i8] zeroinitializer, align 64
 @buf2 = dso_local global [1024 x i8] zeroinitializer, align 64
 
 ; Function Attrs: nounwind uwtable
 define dso_local void @test_api(i32 %0, i16 signext %1, i16 signext %2) {
 ; AVX512-LABEL: test_api:
 ; AVX512:       # %bb.0:
-; AVX512-NEXT:    testl %edi, %edi
-; AVX512-NEXT:    movsbl %sil, %eax
 ; AVX512-NEXT:    vpxord %zmm0, %zmm0, %zmm0
 ; AVX512-NEXT:    vmovdqu64 %zmm0, -{{[0-9]+}}(%rsp)
+; AVX512-NEXT:    testl %edi, %edi
+; AVX512-NEXT:    movsbl %sil, %eax
 ; AVX512-NEXT:    movb $1, -{{[0-9]+}}(%rsp)
 ; AVX512-NEXT:    movb %al, -{{[0-9]+}}(%rsp)
 ; AVX512-NEXT:    movw %si, -{{[0-9]+}}(%rsp)
 ; AVX512-NEXT:    movb %al, -{{[0-9]+}}(%rsp)
 ; AVX512-NEXT:    movw %dx, -{{[0-9]+}}(%rsp)
 ; AVX512-NEXT:    movb %al, -{{[0-9]+}}(%rsp)
 ; AVX512-NEXT:    movw %dx, -{{[0-9]+}}(%rsp)
 ; AVX512-NEXT:    ldtilecfg -{{[0-9]+}}(%rsp)

 ; AVX512-NEXT:    movl $32, %esi
 ; AVX512-NEXT:    tilestored %tmm1, (%rcx,%rsi)
 ; AVX512-NEXT:    tilerelease
 ; AVX512-NEXT:    vzeroupper
 ; AVX512-NEXT:    retq
 ;
 ; AVX2-LABEL: test_api:
 ; AVX2:       # %bb.0:
-; AVX2-NEXT:    testl %edi, %edi
-; AVX2-NEXT:    movsbl %sil, %eax
 ; AVX2-NEXT:    vxorps %ymm0, %ymm0, %ymm0
 ; AVX2-NEXT:    vmovups %ymm0, -{{[0-9]+}}(%rsp)
 ; AVX2-NEXT:    vmovups %ymm0, -{{[0-9]+}}(%rsp)
+; AVX2-NEXT:    testl %edi, %edi
+; AVX2-NEXT:    movsbl %sil, %eax
 ; AVX2-NEXT:    movb $1, -{{[0-9]+}}(%rsp)
 ; AVX2-NEXT:    movb %al, -{{[0-9]+}}(%rsp)
 ; AVX2-NEXT:    movw %si, -{{[0-9]+}}(%rsp)
 ; AVX2-NEXT:    movb %al, -{{[0-9]+}}(%rsp)
 ; AVX2-NEXT:    movw %dx, -{{[0-9]+}}(%rsp)
 ; AVX2-NEXT:    movb %al, -{{[0-9]+}}(%rsp)
 ; AVX2-NEXT:    movw %dx, -{{[0-9]+}}(%rsp)
 ; AVX2-NEXT:    ldtilecfg -{{[0-9]+}}(%rsp)

 ; AVX2-NEXT:    movl $32, %esi
 ; AVX2-NEXT:    tilestored %tmm1, (%rcx,%rsi)
 ; AVX2-NEXT:    tilerelease
 ; AVX2-NEXT:    vzeroupper
 ; AVX2-NEXT:    retq
 ;
 ; SSE2-LABEL: test_api:
 ; SSE2:       # %bb.0:
-; SSE2-NEXT:    testl %edi, %edi
-; SSE2-NEXT:    movsbl %sil, %eax
 ; SSE2-NEXT:    xorps %xmm0, %xmm0
 ; SSE2-NEXT:    movups %xmm0, -{{[0-9]+}}(%rsp)
 ; SSE2-NEXT:    movups %xmm0, -{{[0-9]+}}(%rsp)
 ; SSE2-NEXT:    movups %xmm0, -{{[0-9]+}}(%rsp)
 ; SSE2-NEXT:    movups %xmm0, -{{[0-9]+}}(%rsp)
+; SSE2-NEXT:    testl %edi, %edi
+; SSE2-NEXT:    movsbl %sil, %eax
 ; SSE2-NEXT:    movb $1, -{{[0-9]+}}(%rsp)
 ; SSE2-NEXT:    movb %al, -{{[0-9]+}}(%rsp)
 ; SSE2-NEXT:    movw %si, -{{[0-9]+}}(%rsp)
 ; SSE2-NEXT:    movb %al, -{{[0-9]+}}(%rsp)
 ; SSE2-NEXT:    movw %dx, -{{[0-9]+}}(%rsp)
 ; SSE2-NEXT:    movb %al, -{{[0-9]+}}(%rsp)
 ; SSE2-NEXT:    movw %dx, -{{[0-9]+}}(%rsp)
 ; SSE2-NEXT:    ldtilecfg -{{[0-9]+}}(%rsp)

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

 ; 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:    vzeroupper
+; CHECK-NEXT:    callq foo
 ; 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
if.false:
   %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
 }
 
+define dso_local void @test3(i16 signext %0, i16 signext %1) nounwind {
+; CHECK-LABEL: test3:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    vpxord %zmm0, %zmm0, %zmm0
+; CHECK-NEXT:    vmovdqu64 %zmm0, -{{[0-9]+}}(%rsp)
+; CHECK-NEXT:    xorl %eax, %eax
+; CHECK-NEXT:    testb %al, %al
+; CHECK-NEXT:    jne .LBB2_2
+; CHECK-NEXT:  # %bb.1: # %if.true
+; CHECK-NEXT:    incl %edi
+; CHECK-NEXT:    jmp .LBB2_3
+; CHECK-NEXT:  .LBB2_2: # %if.false
+; CHECK-NEXT:    decl %edi
+; CHECK-NEXT:  .LBB2_3: # %exit
+; CHECK-NEXT:    movb $1, -{{[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:    tilezero %tmm0
+; CHECK-NEXT:    movl $buf, %eax
+; CHECK-NEXT:    movl $32, %ecx
+; CHECK-NEXT:    tilestored %tmm0, (%rax,%rcx)
+; CHECK-NEXT:    tilerelease
+; CHECK-NEXT:    vzeroupper
+; CHECK-NEXT:    retq
+  br i1 undef, label %if.true, label %if.false
+
+if.true:
+  %3 = add i16 %0, 1
+  br label %exit
+
+if.false:
+  %4 = sub i16 %0, 1
+  br label %exit
+
+exit:
+  %5 = phi i16 [ %3, %if.true ], [ %4, %if.false ]
+  %6 = tail call x86_amx @llvm.x86.tilezero.internal(i16 %5, i16 %1)
+  tail call void @llvm.x86.tilestored64.internal(i16 %5, i16 %1, i8* getelementptr inbounds ([3072 x i8], [3072 x i8]* @buf, i64 0, i64 0), i64 32, x86_amx %6)
+  ret void
+}
+
+; TODO: There's PRA Tile Register Configure bug needs to fix later.
+define dso_local void @test4(i16 signext %0, i16 signext %1) nounwind {
+; CHECK-LABEL: test4:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    vpxord %zmm0, %zmm0, %zmm0
+; CHECK-NEXT:    vmovdqu64 %zmm0, -{{[0-9]+}}(%rsp)
+; CHECK-NEXT:    xorl %eax, %eax
+; CHECK-NEXT:    testb %al, %al
+; CHECK-NEXT:    jne .LBB3_3
+; CHECK-NEXT:  # %bb.1: # %if.true
+; CHECK-NEXT:    incl %edi
+; CHECK-NEXT:    xorl %eax, %eax
+; CHECK-NEXT:    testb %al, %al
+; CHECK-NEXT:    jne .LBB3_4
+; CHECK-NEXT:  .LBB3_2: # %amx2
+; CHECK-NEXT:    movb $1, -{{[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 $32, %eax
+; CHECK-NEXT:    movl $buf+1024, %ecx
+; CHECK-NEXT:    tileloadd (%rcx,%rax), %tmm0
+; CHECK-NEXT:    movl $buf, %ecx
+; CHECK-NEXT:    tilestored %tmm0, (%rcx,%rax)
+; CHECK-NEXT:    tilerelease
+; CHECK-NEXT:    vzeroupper
+; CHECK-NEXT:    retq
+; CHECK-NEXT:  .LBB3_3: # %if.false
+; CHECK-NEXT:    decl %edi
+; CHECK-NEXT:    xorl %eax, %eax
+; CHECK-NEXT:    testb %al, %al
+; CHECK-NEXT:    jne .LBB3_2
+; CHECK-NEXT:  .LBB3_4: # %amx1
+; CHECK-NEXT:    ldtilecfg -{{[0-9]+}}(%rsp)
+; CHECK-NEXT:    tilezero %tmm0
+; CHECK-NEXT:    movl $buf, %eax
+; CHECK-NEXT:    movl $32, %ecx
+; CHECK-NEXT:    tilestored %tmm0, (%rax,%rcx)
+; CHECK-NEXT:    tilerelease
+; CHECK-NEXT:    vzeroupper
+; CHECK-NEXT:    retq
+  br i1 undef, label %if.true, label %if.false
+
+if.true:
+  %3 = add i16 %0, 1
+  br i1 undef, label %amx1, label %amx2
+
+if.false:
+  %4 = sub i16 %0, 1
+  br i1 undef, label %amx2, label %amx1
+
+amx1:
+  %5 = phi i16 [ %3, %if.true ], [ %4, %if.false ]
+  %6 = tail call x86_amx @llvm.x86.tilezero.internal(i16 %5, i16 %1)
+  tail call void @llvm.x86.tilestored64.internal(i16 %5, i16 %1, i8* getelementptr inbounds ([3072 x i8], [3072 x i8]* @buf, i64 0, i64 0), i64 32, x86_amx %6)
+  br label %exit
+
+amx2:
+  %7 = phi i16 [ %3, %if.true ], [ %4, %if.false ]
+  %8 = tail call x86_amx @llvm.x86.tileloadd64.internal(i16 %7, i16 %1, i8* getelementptr inbounds ([3072 x i8], [3072 x i8]* @buf, i64 0, i64 1024), i64 32)
+  tail call void @llvm.x86.tilestored64.internal(i16 %7, i16 %1, i8* getelementptr inbounds ([3072 x i8], [3072 x i8]* @buf, i64 0, i64 0), i64 32, x86_amx %8)
+  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/opt-pipeline.ll

 ; CHECK-NEXT:       X86 LEA Optimize
 ; CHECK-NEXT:       X86 Optimize Call Frame
 ; CHECK-NEXT:       X86 Avoid Store Forwarding Block
 ; CHECK-NEXT:       X86 speculative load hardening
 ; CHECK-NEXT:       MachineDominator Tree Construction
 ; CHECK-NEXT:       X86 EFLAGS copy lowering
 ; CHECK-NEXT:       X86 WinAlloca Expander
 ; CHECK-NEXT:       MachineDominator Tree Construction
+; CHECK-NEXT:       Machine Natural Loop Construction
 ; CHECK-NEXT:       Tile Register Pre-configure
 ; CHECK-NEXT:       Detect Dead Lanes
 ; CHECK-NEXT:       Process Implicit Definitions
 ; CHECK-NEXT:       Remove unreachable machine basic blocks
 ; CHECK-NEXT:       Live Variable Analysis
-; CHECK-NEXT:       Machine Natural Loop Construction
 ; CHECK-NEXT:       Eliminate PHI nodes for register allocation
 ; CHECK-NEXT:       Two-Address instruction pass
 ; CHECK-NEXT:       Slot index numbering
 ; CHECK-NEXT:       Live Interval Analysis
 ; CHECK-NEXT:       Simple Register Coalescing
 ; CHECK-NEXT:       Rename Disconnected Subregister Components
 ; CHECK-NEXT:       Machine Instruction Scheduler
 ; CHECK-NEXT:       Machine Block Frequency Analysis