[JTS] Propagate profile info

Author: mtrofin

llvm/lib/Transforms/Scalar/JumpTableToSwitch.cpp

@@ -7,14 +7,23 @@
 //===----------------------------------------------------------------------===//
 
 #include "llvm/Transforms/Scalar/JumpTableToSwitch.h"
+#include "llvm/ADT/DenseSet.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SmallSet.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/Analysis/ConstantFolding.h"
 #include "llvm/Analysis/DomTreeUpdater.h"
 #include "llvm/Analysis/OptimizationRemarkEmitter.h"
 #include "llvm/Analysis/PostDominators.h"
 #include "llvm/IR/IRBuilder.h"
+#include "llvm/IR/LLVMContext.h"
+#include "llvm/IR/ProfDataUtils.h"
+#include "llvm/ProfileData/InstrProf.h"
 #include "llvm/Support/CommandLine.h"
+#include "llvm/Support/Error.h"
+#include "llvm/Transforms/Instrumentation/PGOInstrumentation.h"
 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
+#include <limits>
 
 using namespace llvm;
 
@@ -92,7 +101,8 @@ static std::optional<JumpTableTy> parseJumpTable(GetElementPtrInst *GEP,
 
 static BasicBlock *expandToSwitch(CallBase *CB, const JumpTableTy &JT,
                                   DomTreeUpdater &DTU,
-                                  OptimizationRemarkEmitter &ORE) {
+                                  OptimizationRemarkEmitter &ORE,
+                                  const InstrProfSymtab &Symtab) {
   const bool IsVoid = CB->getType() == Type::getVoidTy(CB->getContext());
 
   SmallVector<DominatorTree::UpdateType, 8> DTUpdates;
@@ -115,7 +125,32 @@ static BasicBlock *expandToSwitch(CallBase *CB, const JumpTableTy &JT,
   IRBuilder<> BuilderTail(CB);
   PHINode *PHI =
       IsVoid ? nullptr : BuilderTail.CreatePHI(CB->getType(), JT.Funcs.size());
-
+  const auto *ProfMD = CB->getMetadata(LLVMContext::MD_prof);
+
+  if (isValueProfileMD(ProfMD)) {
+    SmallVector<uint64_t> BranchWeights;
+    BranchWeights.reserve(JT.Funcs.size());
+    uint64_t TotalCount = 0;
+    auto Targets = getValueProfDataFromInst(
+        *CB, InstrProfValueKind::IPVK_IndirectCallTarget,
+        std::numeric_limits<uint32_t>::max(), TotalCount);
+    DenseMap<const Function *, uint64_t> FToC;
+    for (const auto &[G, C] : Targets)
+      if (const auto *F = Symtab.getFunction(G)) {
+        auto It = FToC.insert({F, C});
+        if (!It.second)
+          It.second += C; // Unexpected, but likely the right way forward
+      }
+    // In the same order as the entries in JT.Funcs, which is the order in which
+    // we'll add cases to the switch statement, get the branch weight of the
+    // corresponding indirect call target, or 0 if we didn't have a profile for
+    // it.
+    for (const auto *F : JT.Funcs)
+      BranchWeights.push_back(FToC.lookup_or(F, 0U));
+    if (!BranchWeights.empty())
+      setProfMetadata(F.getParent(), Switch, BranchWeights,
+                      *llvm::max_element(BranchWeights));
+  }
   for (auto [Index, Func] : llvm::enumerate(JT.Funcs)) {
     BasicBlock *B = BasicBlock::Create(Func->getContext(),
                                        "call." + Twine(Index), &F, Tail);
@@ -150,6 +185,11 @@ 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)));
   for (BasicBlock &BB : make_early_inc_range(F)) {
     BasicBlock *CurrentBB = &BB;
     while (CurrentBB) {
@@ -170,7 +210,7 @@ PreservedAnalyses JumpTableToSwitchPass::run(Function &F,
         std::optional<JumpTableTy> JumpTable = parseJumpTable(GEP, PtrTy);
         if (!JumpTable)
           continue;
-        SplittedOutTail = expandToSwitch(Call, *JumpTable, DTU, ORE);
+        SplittedOutTail = expandToSwitch(Call, *JumpTable, DTU, ORE, Symtab);
         Changed = true;
         break;
       }