[ctx_prof] Profile flatterner

llvm/include/llvm/Analysis/CtxProfAnalysis.h

@@ -9,6 +9,8 @@
 #ifndef LLVM_ANALYSIS_CTXPROFANALYSIS_H
 #define LLVM_ANALYSIS_CTXPROFANALYSIS_H
 
+#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"
@@ -18,6 +20,12 @@ namespace llvm {
 
 class CtxProfAnalysis;
 
+// Setting initial capacity to 1 because all contexts must have at least 1
+// counter, and then, because all contexts belonging to a function have the same
+// size, there'll be at most one other heap allocation.
+using CtxProfFlatProfile =
+    DenseMap<GlobalValue::GUID, SmallVector<uint64_t, 1>>;
+
 /// The instrumented contextual profile, produced by the CtxProfAnalysis.
 class PGOContextualProfile {
   friend class CtxProfAnalysis;
@@ -65,6 +73,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,

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());
+      return;
+    }
+    assert(It->second.size() == Ctx.counters().size() &&
+           "All contexts corresponding to a function should have the exact "
+           "same number of counters.");
+    for (size_t I = 0, E = It->second.size(); I < E; ++I)
+      It->second[I] += Ctx.counters()[I];
+  });
+  return Flat;
+}

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 

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()