[X86] Skip AMX type lowering when AMX is not used

[aengelke]

The pass iterates over the IR multiple times, but most code doesn't use AMX. Therefore, do a single iteration in advance to check whether a function uses AMX at all, and exit early if it doesn't. This makes the function-has-AMX path slightly more expensive, but AMX users probably care a lot less about compile time than JIT users (which tend to not use AMX).

For us, it reduces the time spent in this pass from 0.62% to 0.12%.

Ideally, we wouldn't even need to iterate over the function to determine that it doesn't use AMX.

[llvmbot]

@llvm/pr-subscribers-backend-x86

Author: None (aengelke)

Changes

The pass iterates over the IR multiple times, but most code doesn't use AMX. Therefore, do a single iteration in advance to check whether a function uses AMX at all, and exit early if it doesn't. This makes the function-has-AMX path slightly more expensive, but AMX users probably care a lot less about compile time than JIT users (which tend to not use AMX).

For us, it reduces the time spent in this pass from 0.62% to 0.12%.

Ideally, we wouldn't even need to iterate over the function to determine that it doesn't use AMX.

---

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

1 Files Affected:

• (modified) llvm/lib/Target/X86/X86LowerAMXType.cpp (+16)

diff --git a/llvm/lib/Target/X86/X86LowerAMXType.cpp b/llvm/lib/Target/X86/X86LowerAMXType.cpp
index b69058787a4e2..079ac983a8a01 100644
--- a/llvm/lib/Target/X86/X86LowerAMXType.cpp
+++ b/llvm/lib/Target/X86/X86LowerAMXType.cpp
@@ -92,6 +92,14 @@ static bool isAMXIntrinsic(Value *I) {
   return false;
 }
 
+static bool containsAMXCode(Function &F) {
+  for (BasicBlock &BB : F)
+    for (Instruction &I : BB)
+      if (I.getType()->isX86_AMXTy())
+        return true;
+  return false;
+}
+
 static AllocaInst *createAllocaInstAtEntry(IRBuilder<> &Builder, BasicBlock *BB,
                                            Type *Ty) {
   Function &F = *BB->getParent();
@@ -1230,6 +1238,14 @@ class X86LowerAMXTypeLegacyPass : public FunctionPass {
   }
 
   bool runOnFunction(Function &F) override {
+    // Performance optimization: most code doesn't use AMX, so return early if
+    // there are no instructions that produce AMX values. This is sufficient, as
+    // AMX arguments and constants are not allowed -- so any producer of an AMX
+    // value must be an instruction.
+    // TODO: find a cheaper way for this, without looking at all instructions.
+    if (!containsAMXCode(F))
+      return false;
+
     bool C = false;
     TargetMachine *TM = &getAnalysis<TargetPassConfig>().getTM<TargetMachine>();
     TargetLibraryInfo *TLI =

[phoebewang]

We have a couple passes handling AMX. Can this be done in an analysis pass and reuse the result else where?

[aengelke]

From what I see on AMX-related passes:

• IR LowerAMXTypes, always iterates over IR a few times (thus this PR)
• IR LowerAMXIntrinsics, off by default (technically it's in the pipeline, but it returns early if -enable-x86-scalar-amx is not set) => not relevant.
• MIR FastPreTileConfig (O0), iterates over all virtual regs, exits early if no tile reg exists (NB: iterating over registers is much faster than iterating over instructions)
• MIR PreTileConfig (non-O0), iterates over all machine instrs, returns early if no tile reg def or use exists (not benchmarkd)
• MIR FastTileConfig (O0), iterates over all machine instrs, returns early if no tile reg def or use exists (quite expensive for doing nothing)
• MIR TileConfig (non-O0), iterates over all machine instrs, returns early if VirtRegMap shape map is empty (not benchmarked, but probably ok)
• MIR LowerTileCopy, iterates over all MBB live ins, returns early if no live in is a tile reg

So there is one IR pass and three(five) MIR passes that would benefit from such information. In MIR, X86MachineFunctionInfo seems like a good place and we could trivially set a "uses AMX" flag during ISel when selecting intrinsics. (X86MFI already has AMX-related HasVirtualTileReg, so adding something like UsesAMX would seem good. I can (and probably will) implement such a change in a separate PR.) MFI is (AIUI) not available before ISel and therefore not in the IR pass.

Adding a new analysis pass just for the IR pass feels like overkill for 1 bit of information per function and from the existing immutable passes and passes preserved by MachineFunctionPass, nothing seems like a good fit.

[phoebewang]

Good summary! Feel free to refactor other MIR passes if you are interested :)

[phoebewang]

This LGTM now given we have #94358