[llvm] [ctx_prof] Profile flatterner (PR #104539)

Mircea Trofin via llvm-commits llvm-commits at lists.llvm.org
Fri Aug 16 13:40:34 PDT 2024


https://github.com/mtrofin updated https://github.com/llvm/llvm-project/pull/104539

>From d65b9e556a7e4e21f4e77884ab185085a331eda5 Mon Sep 17 00:00:00 2001
From: Mircea Trofin <mtrofin at google.com>
Date: Thu, 15 Aug 2024 10:28:04 -0700
Subject: [PATCH 1/2] [ctx_prof] Add analysis utility to fetch ID of a callsite

---
 llvm/include/llvm/Analysis/CtxProfAnalysis.h  |   4 +
 llvm/lib/Analysis/CtxProfAnalysis.cpp         |   7 +
 llvm/unittests/Analysis/CMakeLists.txt        |   1 +
 .../Analysis/CtxProfAnalysisTest.cpp          | 145 ++++++++++++++++++
 4 files changed, 157 insertions(+)
 create mode 100644 llvm/unittests/Analysis/CtxProfAnalysisTest.cpp

diff --git a/llvm/include/llvm/Analysis/CtxProfAnalysis.h b/llvm/include/llvm/Analysis/CtxProfAnalysis.h
index d0fb99fe1966a6..25a469bb428b99 100644
--- a/llvm/include/llvm/Analysis/CtxProfAnalysis.h
+++ b/llvm/include/llvm/Analysis/CtxProfAnalysis.h
@@ -11,6 +11,8 @@
 
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/IR/GlobalValue.h"
+#include "llvm/IR/InstrTypes.h"
+#include "llvm/IR/IntrinsicInst.h"
 #include "llvm/IR/PassManager.h"
 #include "llvm/ProfileData/PGOCtxProfReader.h"
 
@@ -84,6 +86,8 @@ class CtxProfAnalysis : public AnalysisInfoMixin<CtxProfAnalysis> {
   using Result = PGOContextualProfile;
 
   PGOContextualProfile run(Module &M, ModuleAnalysisManager &MAM);
+
+  static InstrProfCallsite *getCallsiteInstrumentation(CallBase &CB);
 };
 
 class CtxProfAnalysisPrinterPass
diff --git a/llvm/lib/Analysis/CtxProfAnalysis.cpp b/llvm/lib/Analysis/CtxProfAnalysis.cpp
index d0ccf4ba537f84..51663196b13070 100644
--- a/llvm/lib/Analysis/CtxProfAnalysis.cpp
+++ b/llvm/lib/Analysis/CtxProfAnalysis.cpp
@@ -186,3 +186,10 @@ PreservedAnalyses CtxProfAnalysisPrinterPass::run(Module &M,
   OS << "\n";
   return PreservedAnalyses::all();
 }
+
+InstrProfCallsite *CtxProfAnalysis::getCallsiteInstrumentation(CallBase &CB) {
+  while (auto *Prev = CB.getPrevNode())
+    if (auto *IPC = dyn_cast<InstrProfCallsite>(Prev))
+      return IPC;
+  return nullptr;
+}
diff --git a/llvm/unittests/Analysis/CMakeLists.txt b/llvm/unittests/Analysis/CMakeLists.txt
index d9eb81faac42ad..1dec41972b3578 100644
--- a/llvm/unittests/Analysis/CMakeLists.txt
+++ b/llvm/unittests/Analysis/CMakeLists.txt
@@ -22,6 +22,7 @@ set(ANALYSIS_TEST_SOURCES
   CFGTest.cpp
   CGSCCPassManagerTest.cpp
   ConstraintSystemTest.cpp
+  CtxProfAnalysisTest.cpp
   DDGTest.cpp
   DomTreeUpdaterTest.cpp
   DXILResourceTest.cpp
diff --git a/llvm/unittests/Analysis/CtxProfAnalysisTest.cpp b/llvm/unittests/Analysis/CtxProfAnalysisTest.cpp
new file mode 100644
index 00000000000000..40c00b33be1dcd
--- /dev/null
+++ b/llvm/unittests/Analysis/CtxProfAnalysisTest.cpp
@@ -0,0 +1,145 @@
+//===--- CtxProfAnalysisTest.cpp ------------------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/Analysis/CtxProfAnalysis.h"
+#include "llvm/Analysis/BlockFrequencyInfo.h"
+#include "llvm/Analysis/BranchProbabilityInfo.h"
+#include "llvm/Analysis/CGSCCPassManager.h"
+#include "llvm/Analysis/LoopAnalysisManager.h"
+#include "llvm/Analysis/TargetLibraryInfo.h"
+#include "llvm/AsmParser/Parser.h"
+#include "llvm/IR/Analysis.h"
+#include "llvm/IR/Module.h"
+#include "llvm/IR/PassInstrumentation.h"
+#include "llvm/IR/PassManager.h"
+#include "llvm/Passes/PassBuilder.h"
+#include "llvm/Support/SourceMgr.h"
+#include "llvm/Transforms/Instrumentation/PGOInstrumentation.h"
+#include "gtest/gtest.h"
+
+using namespace llvm;
+
+namespace {
+
+class CtxProfAnalysisTest : public testing::Test {
+  static constexpr auto *IR = R"IR(
+declare void @bar()
+
+define private void @foo(i32 %a, ptr %fct) #0 !guid !0 {
+  %t = icmp eq i32 %a, 0
+  br i1 %t, label %yes, label %no
+yes:
+  call void %fct(i32 %a)
+  br label %exit
+no:
+  call void @bar()
+  br label %exit
+exit:
+  ret void
+}
+
+define void @an_entrypoint(i32 %a) {
+  %t = icmp eq i32 %a, 0
+  br i1 %t, label %yes, label %no
+
+yes:
+  call void @foo(i32 1, ptr null)
+  ret void
+no:
+  ret void
+}
+
+define void @another_entrypoint_no_callees(i32 %a) {
+  %t = icmp eq i32 %a, 0
+  br i1 %t, label %yes, label %no
+
+yes:
+  ret void
+no:
+  ret void
+}
+
+attributes #0 = { noinline }
+!0 = !{ i64 11872291593386833696 }
+)IR";
+
+protected:
+  LLVMContext C;
+  PassBuilder PB;
+  ModuleAnalysisManager MAM;
+  FunctionAnalysisManager FAM;
+  CGSCCAnalysisManager CGAM;
+  LoopAnalysisManager LAM;
+  std::unique_ptr<Module> M;
+
+  void SetUp() override {
+    SMDiagnostic Err;
+    M = parseAssemblyString(IR, Err, C);
+    if (!M)
+      Err.print("CtxProfAnalysisTest", errs());
+  }
+
+public:
+  CtxProfAnalysisTest() {
+    PB.registerModuleAnalyses(MAM);
+    PB.registerCGSCCAnalyses(CGAM);
+    PB.registerFunctionAnalyses(FAM);
+    PB.registerLoopAnalyses(LAM);
+    PB.crossRegisterProxies(LAM, FAM, CGAM, MAM);
+  }
+};
+
+TEST_F(CtxProfAnalysisTest, GetCallsiteIDTest) {
+  ASSERT_TRUE(!!M);
+  ModulePassManager MPM;
+  MPM.addPass(PGOInstrumentationGen(PGOInstrumentationType::CTXPROF));
+  EXPECT_FALSE(MPM.run(*M, MAM).areAllPreserved());
+  auto *F = M->getFunction("foo");
+  ASSERT_NE(F, nullptr);
+  CallBase *IndCall = nullptr;
+  CallBase *DirCall = nullptr;
+  for (auto &BB : *F)
+    for (auto &I : BB)
+      if (auto *CB = dyn_cast<CallBase>(&I)) {
+        if (CB->isIndirectCall()) {
+          EXPECT_EQ(IndCall, nullptr);
+          IndCall = CB;
+        } else if (!CB->getCalledFunction()->isIntrinsic()) {
+          EXPECT_EQ(DirCall, nullptr);
+          DirCall = CB;
+        }
+      }
+  EXPECT_NE(IndCall, nullptr);
+  EXPECT_NE(DirCall, nullptr);
+  auto *IndIns = CtxProfAnalysis::getCallsiteInstrumentation(*IndCall);
+  ASSERT_NE(IndIns, nullptr);
+  EXPECT_EQ(IndIns->getIndex()->getZExtValue(), 0U);
+  auto *DirIns = CtxProfAnalysis::getCallsiteInstrumentation(*DirCall);
+  ASSERT_NE(DirIns, nullptr);
+  EXPECT_EQ(DirIns->getIndex()->getZExtValue(), 1U);
+}
+
+TEST_F(CtxProfAnalysisTest, GetCallsiteIDNegativeTest) {
+  ASSERT_TRUE(!!M);
+  auto *F = M->getFunction("foo");
+  ASSERT_NE(F, nullptr);
+  CallBase *FirstCall = nullptr;
+  for (auto &BB : *F)
+    for (auto &I : BB)
+      if (auto *CB = dyn_cast<CallBase>(&I)) {
+        if (CB->isIndirectCall() || !CB->getCalledFunction()->isIntrinsic()) {
+          FirstCall = CB;
+          break;
+        }
+      }
+  EXPECT_NE(FirstCall, nullptr);
+  auto *IndIns = CtxProfAnalysis::getCallsiteInstrumentation(*FirstCall);
+  ASSERT_EQ(IndIns, nullptr);
+}
+
+} // namespace

>From f74a44bcf35c71023a077f655ca18cdf08e78ec8 Mon Sep 17 00:00:00 2001
From: Mircea Trofin <mtrofin at google.com>
Date: Thu, 15 Aug 2024 19:03:30 -0700
Subject: [PATCH 2/2] [ctx_prof] Profile flatterner

---
 llvm/include/llvm/Analysis/CtxProfAnalysis.h  |  6 ++
 llvm/lib/Analysis/CtxProfAnalysis.cpp         | 40 ++++++++++++
 .../Analysis/CtxProfAnalysis/full-cycle.ll    | 65 ++++++++++++++++++-
 llvm/test/Analysis/CtxProfAnalysis/load.ll    |  5 ++
 4 files changed, 113 insertions(+), 3 deletions(-)

diff --git a/llvm/include/llvm/Analysis/CtxProfAnalysis.h b/llvm/include/llvm/Analysis/CtxProfAnalysis.h
index 25a469bb428b99..060bbcb5e12f3f 100644
--- a/llvm/include/llvm/Analysis/CtxProfAnalysis.h
+++ b/llvm/include/llvm/Analysis/CtxProfAnalysis.h
@@ -10,6 +10,7 @@
 #define LLVM_ANALYSIS_CTXPROFANALYSIS_H
 
 #include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/ilist_node.h"
 #include "llvm/IR/GlobalValue.h"
 #include "llvm/IR/InstrTypes.h"
 #include "llvm/IR/IntrinsicInst.h"
@@ -20,6 +21,9 @@ namespace llvm {
 
 class CtxProfAnalysis;
 
+using CtxProfFlatProfile =
+    DenseMap<GlobalValue::GUID, SmallVector<uint64_t, 16>>;
+
 /// The instrumented contextual profile, produced by the CtxProfAnalysis.
 class PGOContextualProfile {
   friend class CtxProfAnalysis;
@@ -67,6 +71,8 @@ class PGOContextualProfile {
     return FuncInfo.find(getDefinedFunctionGUID(F))->second.NextCallsiteIndex++;
   }
 
+  const CtxProfFlatProfile flatten() const;
+
   bool invalidate(Module &, const PreservedAnalyses &PA,
                   ModuleAnalysisManager::Invalidator &) {
     // Check whether the analysis has been explicitly invalidated. Otherwise,
diff --git a/llvm/lib/Analysis/CtxProfAnalysis.cpp b/llvm/lib/Analysis/CtxProfAnalysis.cpp
index 51663196b13070..837ffc43a30235 100644
--- a/llvm/lib/Analysis/CtxProfAnalysis.cpp
+++ b/llvm/lib/Analysis/CtxProfAnalysis.cpp
@@ -184,6 +184,14 @@ PreservedAnalyses CtxProfAnalysisPrinterPass::run(Module &M,
   OS << "\nCurrent Profile:\n";
   OS << formatv("{0:2}", JSONed);
   OS << "\n";
+  OS << "\nFlat Profile:\n";
+  auto Flat = C.flatten();
+  for (const auto &[Guid, Counters] : Flat) {
+    OS << Guid << " : ";
+    for (auto V : Counters)
+      OS << V << " ";
+    OS << "\n";
+  }
   return PreservedAnalyses::all();
 }
 
@@ -193,3 +201,35 @@ InstrProfCallsite *CtxProfAnalysis::getCallsiteInstrumentation(CallBase &CB) {
       return IPC;
   return nullptr;
 }
+
+static void
+preorderVisit(const PGOCtxProfContext::CallTargetMapTy &Profiles,
+              function_ref<void(const PGOCtxProfContext &)> Visitor) {
+  std::function<void(const PGOCtxProfContext &)> Traverser =
+      [&](const auto &Ctx) {
+        Visitor(Ctx);
+        for (const auto &[_, SubCtxSet] : Ctx.callsites())
+          for (const auto &[__, Subctx] : SubCtxSet)
+            Traverser(Subctx);
+      };
+  for (const auto &[_, P] : Profiles)
+    Traverser(P);
+}
+
+const CtxProfFlatProfile PGOContextualProfile::flatten() const {
+  assert(Profiles.has_value());
+  CtxProfFlatProfile Flat;
+  preorderVisit(*Profiles, [&](const PGOCtxProfContext &Ctx) {
+    auto [It, Ins] = Flat.insert({Ctx.guid(), {}});
+    if (Ins) {
+      llvm::append_range(It->second, Ctx.counters());
+    } else {
+      assert(It->second.size() == Ctx.counters().size() &&
+             "All contexts corresponding to a function should have the exact "
+             "same nr of counters.");
+      for (size_t I = 0, E = It->second.size(); I < E; ++I)
+        It->second[I] += Ctx.counters()[I];
+    }
+  });
+  return Flat;
+}
diff --git a/llvm/test/Analysis/CtxProfAnalysis/full-cycle.ll b/llvm/test/Analysis/CtxProfAnalysis/full-cycle.ll
index 0cdf82bd96efcb..06ba8b3542f7d5 100644
--- a/llvm/test/Analysis/CtxProfAnalysis/full-cycle.ll
+++ b/llvm/test/Analysis/CtxProfAnalysis/full-cycle.ll
@@ -4,6 +4,9 @@
 ; RUN: split-file %s %t
 ;
 ; Test that the GUID metadata survives through thinlink.
+; Also test that the flattener works correctly. f2 is called in 2 places, with
+; different counter values, and we expect resulting flat profile to be the sum
+; (of values at the same index).
 ;
 ; RUN: llvm-ctxprof-util fromJSON --input=%t/profile.json --output=%t/profile.ctxprofdata
 ;
@@ -17,7 +20,9 @@
 ; RUN: llvm-lto2 run %t/m1.bc %t/m2.bc -o %t/ -thinlto-distributed-indexes \
 ; RUN:  -use-ctx-profile=%t/profile.ctxprofdata \
 ; RUN:  -r %t/m1.bc,f1,plx \
+; RUN:  -r %t/m1.bc,f3,plx \
 ; RUN:  -r %t/m2.bc,f1 \
+; RUN:  -r %t/m2.bc,f3 \
 ; RUN:  -r %t/m2.bc,entrypoint,plx
 ; RUN: opt --passes='function-import,require<ctx-prof-analysis>,print<ctx-prof-analysis>' \
 ; RUN:  -summary-file=%t/m2.bc.thinlto.bc -use-ctx-profile=%t/profile.ctxprofdata %t/m2.bc \
@@ -38,6 +43,11 @@ define void @f1() #0 {
   ret void
 }
 
+define void @f3() #0 {
+  call void @f2()
+  ret void
+}
+
 attributes #0 = { noinline }
 !0 = !{ i64 3087265239403591524 }
 
@@ -48,9 +58,11 @@ target triple = "x86_64-pc-linux-gnu"
 source_filename = "random_path/m2.cc"
 
 declare void @f1()
+declare void @f3()
 
 define void @entrypoint() {
   call void @f1()
+  call void @f3()
   ret void
 }
 ;--- profile.json
@@ -63,7 +75,8 @@ define void @entrypoint() {
             [
               {
                 "Counters": [
-                  10
+                  10,
+                  7
                 ],
                 "Guid": 3087265239403591524
               }
@@ -74,6 +87,25 @@ define void @entrypoint() {
           ],
           "Guid": 2072045998141807037
         }
+      ],
+      [
+        {
+          "Callsites": [
+            [
+              {
+                "Counters": [
+                  1,
+                  2
+                ],
+                "Guid": 3087265239403591524
+              }
+            ]
+          ],
+          "Counters": [
+            2
+          ],
+          "Guid": 4197650231481825559
+        }
       ]
     ],
     "Counters": [
@@ -84,8 +116,9 @@ define void @entrypoint() {
 ]
 ;--- expected.txt
 Function Info:
-10507721908651011566 : entrypoint. MaxCounterID: 1. MaxCallsiteID: 1
+10507721908651011566 : entrypoint. MaxCounterID: 1. MaxCallsiteID: 2
 3087265239403591524 : f2.llvm.0. MaxCounterID: 1. MaxCallsiteID: 0
+4197650231481825559 : f3. MaxCounterID: 1. MaxCallsiteID: 1
 2072045998141807037 : f1. MaxCounterID: 1. MaxCallsiteID: 1
 
 Current Profile:
@@ -98,7 +131,8 @@ Current Profile:
             [
               {
                 "Counters": [
-                  10
+                  10,
+                  7
                 ],
                 "Guid": 3087265239403591524
               }
@@ -109,6 +143,25 @@ Current Profile:
           ],
           "Guid": 2072045998141807037
         }
+      ],
+      [
+        {
+          "Callsites": [
+            [
+              {
+                "Counters": [
+                  1,
+                  2
+                ],
+                "Guid": 3087265239403591524
+              }
+            ]
+          ],
+          "Counters": [
+            2
+          ],
+          "Guid": 4197650231481825559
+        }
       ]
     ],
     "Counters": [
@@ -117,3 +170,9 @@ Current Profile:
     "Guid": 10507721908651011566
   }
 ]
+
+Flat Profile:
+10507721908651011566 : 1 
+3087265239403591524 : 11 9 
+4197650231481825559 : 2 
+2072045998141807037 : 7 
diff --git a/llvm/test/Analysis/CtxProfAnalysis/load.ll b/llvm/test/Analysis/CtxProfAnalysis/load.ll
index 69806e334aaec9..fa09474f433151 100644
--- a/llvm/test/Analysis/CtxProfAnalysis/load.ll
+++ b/llvm/test/Analysis/CtxProfAnalysis/load.ll
@@ -86,6 +86,11 @@ Current Profile:
     "Guid": 12074870348631550642
   }
 ]
+
+Flat Profile:
+728453322856651412 : 6 7 
+12074870348631550642 : 5 
+11872291593386833696 : 1 
 ;--- example.ll
 declare void @bar()
 



More information about the llvm-commits mailing list