1 files changed, 12 insertions, 71 deletions

diff --git a/llvm/lib/Target/X86/X86TileConfig.cpp b/llvm/lib/Target/X86/X86TileConfig.cpp
index 17a44dd..09ef8fb 100644
--- a/llvm/lib/Target/X86/X86TileConfig.cpp
+++ b/llvm/lib/Target/X86/X86TileConfig.cpp
@@ -74,63 +74,6 @@ INITIALIZE_PASS_DEPENDENCY(VirtRegMapWrapperLegacy)
 INITIALIZE_PASS_END(X86TileConfig, DEBUG_TYPE, "Tile Register Configure", false,
                     false)
 
-unsigned getAMXRegNum(MachineRegisterInfo *MRI, Register Reg) {
-  if (Reg.isVirtual()) {
-    unsigned RegClassID = MRI->getRegClass(Reg)->getID();
-    if (RegClassID == X86::TILERegClassID)
-      return 1;
-    if (RegClassID == X86::TILEPAIRRegClassID)
-      return 2;
-  } else {
-    if (Reg >= X86::TMM0 && Reg <= X86::TMM7)
-      return 1;
-    if (Reg >= X86::TMM0_TMM1 && Reg <= X86::TMM6_TMM7)
-      return 2;
-  }
-  return 0;
-}
-
-static void collectVirtRegShapes(MachineRegisterInfo *MRI, VirtRegMap &VRM,
-                                 Register VirtReg,
-                                 SmallVector<ShapeT, 8> &Phys2Shapes) {
-  unsigned Num = getAMXRegNum(MRI, VirtReg);
-  MCRegister PhysReg = VRM.getPhys(VirtReg);
-  if (!PhysReg)
-    return;
-
-  if (Num == 1) {
-    unsigned Index = PhysReg - X86::TMM0;
-    if (!Phys2Shapes[Index].isValid()) {
-      ShapeT Shape = VRM.getShape(VirtReg);
-      Phys2Shapes[Index] = std::move(Shape);
-      return;
-    }
-  }
-  // Split tile pair shape info to 2 single tile shape info. e.g:
-  // Put TMM0_TMM1's Shape to TMM0's shape + TMM1's Shape in Phys2Shapes.
-  if (Num == 2) {
-    unsigned Index0 = (PhysReg - X86::TMM0_TMM1) * 2;
-    unsigned Index1 = (PhysReg - X86::TMM0_TMM1) * 2 + 1;
-
-    ShapeT Shape = VRM.getShape(VirtReg);
-    assert(Shape.getShapeNum() == 2 && "Unexpected shape number!");
-
-    if (!Phys2Shapes[Index0].isValid()) {
-      ShapeT Shape0(Shape.getRow(0), Shape.getCol(0), MRI);
-      Phys2Shapes[Index0] = std::move(Shape0);
-    }
-
-    if (!Phys2Shapes[Index1].isValid()) {
-      ShapeT Shape1(Shape.getRow(1), Shape.getCol(1), MRI);
-      Phys2Shapes[Index1] = std::move(Shape1);
-    }
-  }
-}
-
-static bool isAMXRegClass(MachineRegisterInfo *MRI, Register Reg) {
-  return getAMXRegNum(MRI, Reg) > 0;
-}
-
 bool X86TileConfig::runOnMachineFunction(MachineFunction &MF) {
   X86MachineFunctionInfo *X86FI = MF.getInfo<X86MachineFunctionInfo>();
   // Early exit in the common case of non-AMX code.
@@ -138,7 +81,7 @@ bool X86TileConfig::runOnMachineFunction(MachineFunction &MF) {
     return false;
 
   const X86Subtarget &ST = MF.getSubtarget<X86Subtarget>();
-  const TargetRegisterInfo *TRI = ST.getRegisterInfo();
+  const X86RegisterInfo *TRI = ST.getRegisterInfo();
   const TargetInstrInfo *TII = ST.getInstrInfo();
   MachineRegisterInfo &MRI = MF.getRegInfo();
   LiveIntervals &LIS = getAnalysis<LiveIntervalsWrapperPass>().getLIS();
@@ -176,24 +119,29 @@ bool X86TileConfig::runOnMachineFunction(MachineFunction &MF) {
   assert(ConstMI && "Cannot find an insertion point");
 
   unsigned AMXRegNum = TRI->getRegClass(X86::TILERegClassID)->getNumRegs();
-  SmallVector<ShapeT, 8> Phys2Shapes(AMXRegNum, ShapeT());
+  SmallVector<Register, 8> Phys2Virt(AMXRegNum, 0);
   for (unsigned I = 0, E = MRI.getNumVirtRegs(); I != E; ++I) {
     Register VirtReg = Register::index2VirtReg(I);
     if (MRI.reg_nodbg_empty(VirtReg))
       continue;
-    if (!isAMXRegClass(&MRI, VirtReg))
+    if (!TRI->isTileRegisterClass(MRI.getRegClass(VirtReg)))
+      continue;
+    MCRegister PhysReg = VRM.getPhys(VirtReg);
+    if (!PhysReg)
       continue;
-    collectVirtRegShapes(&MRI, VRM, VirtReg, Phys2Shapes);
+    unsigned Index = PhysReg - X86::TMM0;
+    if (!Phys2Virt[Index])
+      Phys2Virt[Index] = VirtReg;
   }
 
   // Fill in the shape of each tile physical register.
   for (unsigned I = 0; I < AMXRegNum; ++I) {
-    ShapeT Shape = Phys2Shapes[I];
-    if (!Shape.isValid())
+    if (!Phys2Virt[I])
       continue;
     DebugLoc DL;
     bool IsRow = true;
     MachineInstr *NewMI = nullptr;
+    ShapeT Shape = VRM.getShape(Phys2Virt[I]);
     for (auto &R : {Shape.getRow()->getReg(), Shape.getCol()->getReg()}) {
       // Here is the data format for the tile config.
       // 0      palette
@@ -222,14 +170,7 @@ bool X86TileConfig::runOnMachineFunction(MachineFunction &MF) {
                    "Cannot initialize with different shapes");
             continue;
           }
-          if (DefMI.getOperand(1).isImm()) {
-            Imm = DefMI.getOperand(1).getImm();
-          } else {
-            assert(DefMI.getOpcode() == X86::MOV32r0 &&
-                   "The opcode is assumed to be MOV32r0 if the operand is not "
-                   "immediate.");
-            Imm = 0;
-          }
+          Imm = DefMI.getOperand(1).getImm();
 
           NewMI = addFrameReference(
                       BuildMI(MF.front(), ++ConstMI->getIterator(), DL,