[llvm] [X86] Skip AMX type lowering when AMX is not used (PR #92910)

[via llvm-commits]

https://github.com/aengelke created https://github.com/llvm/llvm-project/pull/92910

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.

>From bb03f53d5024bfcf07c03448847c0e03ee3f1d8c Mon Sep 17 00:00:00 2001
From: Alexis Engelke <engelke at in.tum.de>
Date: Tue, 21 May 2024 14:50:16 +0200
Subject: [PATCH] [X86] Skip AMX type lowering when AMX is not used

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.
---
 llvm/lib/Target/X86/X86LowerAMXType.cpp | 16 ++++++++++++++++
 1 file changed, 16 insertions(+)

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 =