[X86] Early exit MIR AMX passes when AMX is unused

[aengelke]

Follow-up of #94358. Do the checks even before calling getRegisterInfo
etc., because some of these are virtual function calls.

Depends on #94988.

[llvmbot]

@llvm/pr-subscribers-backend-x86

Author: None (aengelke)

Changes

Follow-up of #94358. Do the checks even before calling getRegisterInfo
etc., because some of these are virtual function calls.

Depends on #94988.

---

Full diff: https://github.com/llvm/llvm-project/pull/94989.diff

4 Files Affected:

• (modified) llvm/lib/Target/X86/X86FastPreTileConfig.cpp (+5-13)
• (modified) llvm/lib/Target/X86/X86FastTileConfig.cpp (+5-1)
• (modified) llvm/lib/Target/X86/X86PreTileConfig.cpp (+5-1)
• (modified) llvm/lib/Target/X86/X86TileConfig.cpp (+5)

diff --git a/llvm/lib/Target/X86/X86FastPreTileConfig.cpp b/llvm/lib/Target/X86/X86FastPreTileConfig.cpp
index 7708089074c0a..d50a4d3b23ae2 100644
--- a/llvm/lib/Target/X86/X86FastPreTileConfig.cpp
+++ b/llvm/lib/Target/X86/X86FastPreTileConfig.cpp
@@ -653,28 +653,20 @@ bool X86FastPreTileConfig::configBasicBlock(MachineBasicBlock &MBB) {
 }
 
 bool X86FastPreTileConfig::runOnMachineFunction(MachineFunction &MFunc) {
+  X86FI = MFunc.getInfo<X86MachineFunctionInfo>();
+  // Early exit in the common case of non-AMX code.
+  if (X86FI->getAMXProgModel() != AMXProgModelEnum::ManagedRA)
+    return false;
+
   MF = &MFunc;
   MRI = &MFunc.getRegInfo();
   ST = &MFunc.getSubtarget<X86Subtarget>();
   TII = ST->getInstrInfo();
-  X86FI = MFunc.getInfo<X86MachineFunctionInfo>();
   MFI = &MFunc.getFrameInfo();
   TRI = ST->getRegisterInfo();
   CfgSS = -1;
 
   unsigned NumVirtRegs = MRI->getNumVirtRegs();
-  // Abandon early if there is no tile register to config.
-  bool HasVirtTileReg = false;
-  for (unsigned I = 0, E = NumVirtRegs; I != E; ++I) {
-    Register VirtReg = Register::index2VirtReg(I);
-    if (!MRI->reg_nodbg_empty(VirtReg) &&
-        MRI->getRegClass(VirtReg)->getID() == X86::TILERegClassID) {
-      HasVirtTileReg = true;
-      break;
-    }
-  }
-  if (!HasVirtTileReg)
-    return false;
 
   StackSlotForVirtReg.resize(NumVirtRegs);
   MayLiveAcrossBlocks.clear();
diff --git a/llvm/lib/Target/X86/X86FastTileConfig.cpp b/llvm/lib/Target/X86/X86FastTileConfig.cpp
index 2a20cd13791de..cc59b97f17ab2 100644
--- a/llvm/lib/Target/X86/X86FastTileConfig.cpp
+++ b/llvm/lib/Target/X86/X86FastTileConfig.cpp
@@ -168,12 +168,16 @@ bool X86FastTileConfig::configBasicBlock(MachineBasicBlock &MBB) {
 }
 
 bool X86FastTileConfig::runOnMachineFunction(MachineFunction &MFunc) {
+  X86FI = MFunc.getInfo<X86MachineFunctionInfo>();
+  // Early exit in the common case of non-AMX code.
+  if (X86FI->getAMXProgModel() != AMXProgModelEnum::ManagedRA)
+    return false;
+
   MF = &MFunc;
   MRI = &MFunc.getRegInfo();
   const TargetSubtargetInfo *ST = &MFunc.getSubtarget<X86Subtarget>();
   TRI = ST->getRegisterInfo();
   TII = MFunc.getSubtarget().getInstrInfo();
-  X86FI = MFunc.getInfo<X86MachineFunctionInfo>();
   bool Change = false;
 
   // Loop over all of the basic blocks, eliminating virtual register references
diff --git a/llvm/lib/Target/X86/X86PreTileConfig.cpp b/llvm/lib/Target/X86/X86PreTileConfig.cpp
index 75ad58e5cdcb7..1e3b72e7aed0d 100644
--- a/llvm/lib/Target/X86/X86PreTileConfig.cpp
+++ b/llvm/lib/Target/X86/X86PreTileConfig.cpp
@@ -237,11 +237,15 @@ void X86PreTileConfig::collectShapeInfo(MachineInstr &MI) {
 }
 
 bool X86PreTileConfig::runOnMachineFunction(MachineFunction &MF) {
+  X86MachineFunctionInfo *X86FI = MF.getInfo<X86MachineFunctionInfo>();
+  // Early exit in the common case of non-AMX code.
+  if (X86FI->getAMXProgModel() != AMXProgModelEnum::ManagedRA)
+    return false;
+
   const X86Subtarget &ST = MF.getSubtarget<X86Subtarget>();
   const TargetInstrInfo *TII = ST.getInstrInfo();
   const TargetRegisterInfo *TRI = ST.getRegisterInfo();
   const TargetRegisterClass *RC = TRI->getRegClass(X86::TILERegClassID);
-  X86MachineFunctionInfo *X86FI = MF.getInfo<X86MachineFunctionInfo>();
 
   BitVector AMXRegs(TRI->getNumRegs());
   for (unsigned I = 0; I < RC->getNumRegs(); I++)
diff --git a/llvm/lib/Target/X86/X86TileConfig.cpp b/llvm/lib/Target/X86/X86TileConfig.cpp
index 5cada924e0064..ebe48910225f9 100644
--- a/llvm/lib/Target/X86/X86TileConfig.cpp
+++ b/llvm/lib/Target/X86/X86TileConfig.cpp
@@ -77,6 +77,11 @@ INITIALIZE_PASS_END(X86TileConfig, DEBUG_TYPE, "Tile Register Configure", false,
                     false)
 
 bool X86TileConfig::runOnMachineFunction(MachineFunction &MF) {
+  X86MachineFunctionInfo *X86FI = MF.getInfo<X86MachineFunctionInfo>();
+  // Early exit in the common case of non-AMX code.
+  if (X86FI->getAMXProgModel() != AMXProgModelEnum::ManagedRA)
+    return false;
+
   const X86Subtarget &ST = MF.getSubtarget<X86Subtarget>();
   const TargetRegisterInfo *TRI = ST.getRegisterInfo();
   const TargetInstrInfo *TII = ST.getInstrInfo();

[KanRobert]

LGTM except you need to update some MIR tests @aengelke

[aengelke]

Updating MIR tests is #94988 -- this was a bit more work than I expected (needs YAML infrastructure for X86FunctionInfo), so I split it out in a separate PR.

[phoebewang]

Just found we have set/hasVirtualTileReg interfaces. We can remove it and replace with X86FI->getAMXProgModel() == AMXProgModelEnum::ManagedRA too.

[aengelke]

Good point, opened #95105.

[phoebewang]

LGTM, thanks!