[JTS][NFC] Optimize guid fetching

[mtrofin]

It's unnecessary to build the whole symtable, and on top of everything, un-optimal to do so for every function. All we really need is the instrumented PGO name - considering also LTO-ness - and then we can compute the function name.

[mtrofin]

• [JTS][NFC] Optimize guid fetching #161612 👈 (View in Graphite)
• main

This stack of pull requests is managed by Graphite. Learn more about stacking.

[llvmbot]

@llvm/pr-subscribers-llvm-transforms

Author: Mircea Trofin (mtrofin)

Changes

It's unnecessary to build the whole symtable, and on top of everything, un-optimal to do so for every function. All we really need is the instrumented PGO name - considering also LTO-ness - and then we can compute the function name.

---

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

3 Files Affected:

• (modified) llvm/include/llvm/Transforms/Scalar/JumpTableToSwitch.h (+6-1)
• (modified) llvm/lib/Passes/PassBuilderPipelines.cpp (+3-1)
• (modified) llvm/lib/Transforms/Scalar/JumpTableToSwitch.cpp (+8-14)

diff --git a/llvm/include/llvm/Transforms/Scalar/JumpTableToSwitch.h b/llvm/include/llvm/Transforms/Scalar/JumpTableToSwitch.h
index 61786227d7a33..71bf2c643d813 100644
--- a/llvm/include/llvm/Transforms/Scalar/JumpTableToSwitch.h
+++ b/llvm/include/llvm/Transforms/Scalar/JumpTableToSwitch.h
@@ -15,7 +15,12 @@ namespace llvm {
 
 class Function;
 
-struct JumpTableToSwitchPass : PassInfoMixin<JumpTableToSwitchPass> {
+class JumpTableToSwitchPass : public PassInfoMixin<JumpTableToSwitchPass> {
+  // Necessary until we switch to GUIDs as metadata, after which we can drop it.
+  const bool InLTO;
+
+public:
+  JumpTableToSwitchPass(bool InLTO = false) : InLTO(InLTO) {}
   /// Run the pass over the function.
   PreservedAnalyses run(Function &F, FunctionAnalysisManager &AM);
 };
diff --git a/llvm/lib/Passes/PassBuilderPipelines.cpp b/llvm/lib/Passes/PassBuilderPipelines.cpp
index 256cf9d4cd1ce..0b93dda4087e3 100644
--- a/llvm/lib/Passes/PassBuilderPipelines.cpp
+++ b/llvm/lib/Passes/PassBuilderPipelines.cpp
@@ -610,7 +610,9 @@ PassBuilder::buildFunctionSimplificationPipeline(OptimizationLevel Level,
 
   // Jump table to switch conversion.
   if (EnableJumpTableToSwitch)
-    FPM.addPass(JumpTableToSwitchPass());
+    FPM.addPass(JumpTableToSwitchPass(
+        /*InLTO=*/Phase == ThinOrFullLTOPhase::ThinLTOPostLink ||
+        Phase == ThinOrFullLTOPhase::FullLTOPostLink));
 
   FPM.addPass(
       SimplifyCFGPass(SimplifyCFGOptions().convertSwitchRangeToICmp(true)));
diff --git a/llvm/lib/Transforms/Scalar/JumpTableToSwitch.cpp b/llvm/lib/Transforms/Scalar/JumpTableToSwitch.cpp
index 2025fbbf05973..5afdef9e95fae 100644
--- a/llvm/lib/Transforms/Scalar/JumpTableToSwitch.cpp
+++ b/llvm/lib/Transforms/Scalar/JumpTableToSwitch.cpp
@@ -201,14 +201,12 @@ PreservedAnalyses JumpTableToSwitchPass::run(Function &F,
   PostDominatorTree *PDT = AM.getCachedResult<PostDominatorTreeAnalysis>(F);
   DomTreeUpdater DTU(DT, PDT, DomTreeUpdater::UpdateStrategy::Lazy);
   bool Changed = false;
-  InstrProfSymtab Symtab;
-  if (auto E = Symtab.create(*F.getParent()))
-    F.getContext().emitError(
-        "Could not create indirect call table, likely corrupted IR" +
-        toString(std::move(E)));
-  DenseMap<const Function *, GlobalValue::GUID> FToGuid;
-  for (const auto &[G, FPtr] : Symtab.getIDToNameMap())
-    FToGuid.insert({FPtr, G});
+  auto FuncToGuid = [&](const Function &Fct) {
+    if (Fct.getMetadata(AssignGUIDPass::GUIDMetadataName))
+      return AssignGUIDPass::getGUID(Fct);
+
+    return Function::getGUIDAssumingExternalLinkage(getIRPGOFuncName(F, InLTO));
+  };
 
   for (BasicBlock &BB : make_early_inc_range(F)) {
     BasicBlock *CurrentBB = &BB;
@@ -230,12 +228,8 @@ PreservedAnalyses JumpTableToSwitchPass::run(Function &F,
         std::optional<JumpTableTy> JumpTable = parseJumpTable(GEP, PtrTy);
         if (!JumpTable)
           continue;
-        SplittedOutTail = expandToSwitch(
-            Call, *JumpTable, DTU, ORE, [&](const Function &Fct) {
-              if (Fct.getMetadata(AssignGUIDPass::GUIDMetadataName))
-                return AssignGUIDPass::getGUID(Fct);
-              return FToGuid.lookup_or(&Fct, 0U);
-            });
+        SplittedOutTail =
+            expandToSwitch(Call, *JumpTable, DTU, ORE, FuncToGuid);
         Changed = true;
         break;
       }