[ctxprof] ProfileWriter abstraction #129590

mtrofin · 2025-03-03T21:01:11Z

Introduce a ProfileWriter abstraction to replace the callback passed to __llvm_ctx_profile_fetch. Subsequent changes will add support for flat profile collection (as in, collection of non-contextual profile for those functions not under a contextual root), which require also a change in the profile format. The abstraction makes it easy to add "write flat" - related capabilities without constantly complicating the signature of __llvm_ctx_profile_fetch.

mtrofin · 2025-03-03T21:01:31Z

[ctxprof] ProfileWriter abstraction #129590 : 2 dependent PRs (#129592 , #129623 ) 👈 (View in Graphite)
main

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

kazutakahirata

LGTM.

llvmbot · 2025-03-04T02:03:59Z

@llvm/pr-subscribers-pgo

Author: Mircea Trofin (mtrofin)

Changes

Introduce a ProfileWriter abstraction to replace the callback passed to __llvm_ctx_profile_fetch. Subsequent changes will add support for flat profile collection (as in, collection of non-contextual profile for those functions not under a contextual root), which require also a change in the profile format. The abstraction makes it easy to add "write flat" - related capabilities without constantly complicating the signature of __llvm_ctx_profile_fetch.

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

6 Files Affected:

(modified) compiler-rt/lib/ctx_profile/CtxInstrContextNode.h (+8)
(modified) compiler-rt/lib/ctx_profile/CtxInstrProfiling.cpp (+4-5)
(modified) compiler-rt/lib/ctx_profile/CtxInstrProfiling.h (+1-2)
(modified) compiler-rt/lib/ctx_profile/tests/CtxInstrProfilingTest.cpp (+6-9)
(modified) compiler-rt/test/ctx_profile/TestCases/generate-context.cpp (+28-24)
(modified) llvm/include/llvm/ProfileData/CtxInstrContextNode.h (+8)

diff --git a/compiler-rt/lib/ctx_profile/CtxInstrContextNode.h b/compiler-rt/lib/ctx_profile/CtxInstrContextNode.h
index 36a996632b71e..6b020733e1f37 100644
--- a/compiler-rt/lib/ctx_profile/CtxInstrContextNode.h
+++ b/compiler-rt/lib/ctx_profile/CtxInstrContextNode.h
@@ -112,6 +112,14 @@ class ContextNode final {
 
   uint64_t entrycount() const { return counters()[0]; }
 };
+
+class ProfileWriter {
+public:
+  virtual void startContextSection() = 0;
+  virtual void writeContextual(const ctx_profile::ContextNode &RootNode) = 0;
+  virtual void endContextSection() = 0;
+  virtual ~ProfileWriter() = default;
+};
 } // namespace ctx_profile
 } // namespace llvm
 #endif
diff --git a/compiler-rt/lib/ctx_profile/CtxInstrProfiling.cpp b/compiler-rt/lib/ctx_profile/CtxInstrProfiling.cpp
index df30986cdfc69..992aa94a6631d 100644
--- a/compiler-rt/lib/ctx_profile/CtxInstrProfiling.cpp
+++ b/compiler-rt/lib/ctx_profile/CtxInstrProfiling.cpp
@@ -294,12 +294,11 @@ void __llvm_ctx_profile_start_collection() {
   __sanitizer::Printf("[ctxprof] Initial NumMemUnits: %zu \n", NumMemUnits);
 }
 
-bool __llvm_ctx_profile_fetch(void *Data,
-                              bool (*Writer)(void *W, const ContextNode &)) {
-  assert(Writer);
+bool __llvm_ctx_profile_fetch(ProfileWriter &Writer) {
   __sanitizer::GenericScopedLock<__sanitizer::SpinMutex> Lock(
       &AllContextsMutex);
 
+  Writer.startContextSection();
   for (int I = 0, E = AllContextRoots.Size(); I < E; ++I) {
     auto *Root = AllContextRoots[I];
     __sanitizer::GenericScopedLock<__sanitizer::StaticSpinMutex> TakenLock(
@@ -308,9 +307,9 @@ bool __llvm_ctx_profile_fetch(void *Data,
       __sanitizer::Printf("[ctxprof] Contextual Profile is %s\n", "invalid");
       return false;
     }
-    if (!Writer(Data, *Root->FirstNode))
-      return false;
+    Writer.writeContextual(*Root->FirstNode);
   }
+  Writer.endContextSection();
   return true;
 }
 
diff --git a/compiler-rt/lib/ctx_profile/CtxInstrProfiling.h b/compiler-rt/lib/ctx_profile/CtxInstrProfiling.h
index 74d346d6e0a07..8a6949d4ec288 100644
--- a/compiler-rt/lib/ctx_profile/CtxInstrProfiling.h
+++ b/compiler-rt/lib/ctx_profile/CtxInstrProfiling.h
@@ -169,7 +169,6 @@ void __llvm_ctx_profile_free();
 /// The Writer's first parameter plays the role of closure for Writer, and is
 /// what the caller of __llvm_ctx_profile_fetch passes as the Data parameter.
 /// The second parameter is the root of a context tree.
-bool __llvm_ctx_profile_fetch(void *Data,
-                              bool (*Writer)(void *, const ContextNode &));
+bool __llvm_ctx_profile_fetch(ProfileWriter &);
 }
 #endif // CTX_PROFILE_CTXINSTRPROFILING_H_
diff --git a/compiler-rt/lib/ctx_profile/tests/CtxInstrProfilingTest.cpp b/compiler-rt/lib/ctx_profile/tests/CtxInstrProfilingTest.cpp
index d9f08b1e7efe8..d9183a8a0ff76 100644
--- a/compiler-rt/lib/ctx_profile/tests/CtxInstrProfilingTest.cpp
+++ b/compiler-rt/lib/ctx_profile/tests/CtxInstrProfilingTest.cpp
@@ -179,13 +179,15 @@ TEST_F(ContextTest, Dump) {
   (void)Subctx;
   __llvm_ctx_profile_release_context(&Root);
 
-  struct Writer {
+  struct Writer : public ProfileWriter {
     ContextRoot *const Root;
     const size_t Entries;
     bool State = false;
     Writer(ContextRoot *Root, size_t Entries) : Root(Root), Entries(Entries) {}
 
-    bool write(const ContextNode &Node) {
+    void startContextSection() override {}
+    void endContextSection() override {}
+    void writeContextual(const ContextNode &Node) override {
       EXPECT_FALSE(Root->Taken.TryLock());
       EXPECT_EQ(Node.guid(), 1U);
       EXPECT_EQ(Node.counters()[0], Entries);
@@ -202,22 +204,17 @@ TEST_F(ContextTest, Dump) {
       EXPECT_EQ(SN.callsites_size(), 1U);
       EXPECT_EQ(SN.subContexts()[0], nullptr);
       State = true;
-      return true;
     }
   };
   Writer W(&Root, 1);
   EXPECT_FALSE(W.State);
-  __llvm_ctx_profile_fetch(&W, [](void *W, const ContextNode &Node) -> bool {
-    return reinterpret_cast<Writer *>(W)->write(Node);
-  });
+  __llvm_ctx_profile_fetch(W);
   EXPECT_TRUE(W.State);
 
   // this resets all counters but not the internal structure.
   __llvm_ctx_profile_start_collection();
   Writer W2(&Root, 0);
   EXPECT_FALSE(W2.State);
-  __llvm_ctx_profile_fetch(&W2, [](void *W, const ContextNode &Node) -> bool {
-    return reinterpret_cast<Writer *>(W)->write(Node);
-  });
+  __llvm_ctx_profile_fetch(W2);
   EXPECT_TRUE(W2.State);
 }
diff --git a/compiler-rt/test/ctx_profile/TestCases/generate-context.cpp b/compiler-rt/test/ctx_profile/TestCases/generate-context.cpp
index 797b871860655..c26f4edfe783e 100644
--- a/compiler-rt/test/ctx_profile/TestCases/generate-context.cpp
+++ b/compiler-rt/test/ctx_profile/TestCases/generate-context.cpp
@@ -15,9 +15,7 @@
 #include <iostream>
 
 using namespace llvm::ctx_profile;
-extern "C" bool __llvm_ctx_profile_fetch(void *Data,
-                                         bool (*Writer)(void *,
-                                                        const ContextNode &));
+extern "C" bool __llvm_ctx_profile_fetch(ProfileWriter &);
 
 // avoid name mangling
 extern "C" {
@@ -46,22 +44,31 @@ __attribute__((noinline)) void theRoot() {
 // CHECK-NEXT: check even
 // CHECK-NEXT: check odd
 
-void printProfile(const ContextNode &Node, const std::string &Indent,
-                  const std::string &Increment) {
-  std::cout << Indent << "Guid: " << Node.guid() << std::endl;
-  std::cout << Indent << "Entries: " << Node.entrycount() << std::endl;
-  std::cout << Indent << Node.counters_size() << " counters and "
-            << Node.callsites_size() << " callsites" << std::endl;
-  std::cout << Indent << "Counter values: ";
-  for (uint32_t I = 0U; I < Node.counters_size(); ++I)
-    std::cout << Node.counters()[I] << " ";
-  std::cout << std::endl;
-  for (uint32_t I = 0U; I < Node.callsites_size(); ++I)
-    for (const auto *N = Node.subContexts()[I]; N; N = N->next()) {
-      std::cout << Indent << "At Index " << I << ":" << std::endl;
-      printProfile(*N, Indent + Increment, Increment);
-    }
-}
+class TestProfileWriter : public ProfileWriter {
+  void printProfile(const ContextNode &Node, const std::string &Indent,
+                    const std::string &Increment) {
+    std::cout << Indent << "Guid: " << Node.guid() << std::endl;
+    std::cout << Indent << "Entries: " << Node.entrycount() << std::endl;
+    std::cout << Indent << Node.counters_size() << " counters and "
+              << Node.callsites_size() << " callsites" << std::endl;
+    std::cout << Indent << "Counter values: ";
+    for (uint32_t I = 0U; I < Node.counters_size(); ++I)
+      std::cout << Node.counters()[I] << " ";
+    std::cout << std::endl;
+    for (uint32_t I = 0U; I < Node.callsites_size(); ++I)
+      for (const auto *N = Node.subContexts()[I]; N; N = N->next()) {
+        std::cout << Indent << "At Index " << I << ":" << std::endl;
+        printProfile(*N, Indent + Increment, Increment);
+      }
+  }
+
+public:
+  void startContextSection() override {}
+  void writeContextual(const ContextNode &RootNode) override {
+    printProfile(RootNode, "", "");
+  }
+  void endContextSection() override {}
+};
 
 // 8657661246551306189 is theRoot. We expect 2 callsites and 2 counters - one
 // for the entry basic block and one for the loop.
@@ -88,11 +95,8 @@ void printProfile(const ContextNode &Node, const std::string &Indent,
 // CHECK-NEXT:   Counter values: 2 1
 
 bool profileWriter() {
-  return __llvm_ctx_profile_fetch(
-      nullptr, +[](void *, const ContextNode &Node) {
-        printProfile(Node, "", "  ");
-        return true;
-      });
+  TestProfileWriter W;
+  return __llvm_ctx_profile_fetch(W);
 }
 
 int main(int argc, char **argv) {
diff --git a/llvm/include/llvm/ProfileData/CtxInstrContextNode.h b/llvm/include/llvm/ProfileData/CtxInstrContextNode.h
index 36a996632b71e..6b020733e1f37 100644
--- a/llvm/include/llvm/ProfileData/CtxInstrContextNode.h
+++ b/llvm/include/llvm/ProfileData/CtxInstrContextNode.h
@@ -112,6 +112,14 @@ class ContextNode final {
 
   uint64_t entrycount() const { return counters()[0]; }
 };
+
+class ProfileWriter {
+public:
+  virtual void startContextSection() = 0;
+  virtual void writeContextual(const ctx_profile::ContextNode &RootNode) = 0;
+  virtual void endContextSection() = 0;
+  virtual ~ProfileWriter() = default;
+};
 } // namespace ctx_profile
 } // namespace llvm
 #endif

snehasish · 2025-03-04T16:27:09Z

compiler-rt/test/ctx_profile/TestCases/generate-context.cpp

-    }
-}
+class TestProfileWriter : public ProfileWriter {
+  void printProfile(const ContextNode &Node, const std::string &Indent,


I think it might be usefult to pass in a llvm::raw_ostream& and then writing to it. I think this will give you more flexibility based when used in tests (e.g using a string buffer).

This is a .cpp file that gets compiled (and run) as part of the test, so we basically have access to std:: things.

To make things clear, here's an example of the usage I'm thinking of: https://github.com/llvm/llvm-project/blob/main/llvm/unittests/IR/MetadataTest.cpp#L164-L172

Just to make sure we're talking about the same thing: this here is a lit test. Note the RUN lines at the top of the file.

Discussed offline and since this in compiler-rt its not useful to use ostream.

snehasish · 2025-03-04T16:28:50Z

compiler-rt/lib/ctx_profile/tests/CtxInstrProfilingTest.cpp

@@ -179,13 +179,15 @@ TEST_F(ContextTest, Dump) {
  (void)Subctx;
  __llvm_ctx_profile_release_context(&Root);

-  struct Writer {
+  struct Writer : public ProfileWriter {


Its a bit strange to have Writer inherit from ProfileWriter. To me it makes more sense to call the interface Writer and the derived class / struct profile writer here. Wdyt?

Also a struct derived from a class seems strange to me.

Writer is too generic. I only use it here because this is a test. I can rename it to TestProfileWriter (and make it a class, that is weird, indeed)

snehasish · 2025-03-04T16:30:20Z

compiler-rt/lib/ctx_profile/CtxInstrContextNode.h

@@ -112,6 +112,12 @@ class ContextNode final {

  uint64_t entrycount() const { return counters()[0]; }
 };
+
+class ProfileWriter {


Add a comment to describe the interface and expected usage?

done - briefly. Updated further up the change stack, where more members appear.

snehasish · 2025-03-04T16:33:00Z

compiler-rt/test/ctx_profile/TestCases/generate-context.cpp

+  }
+
+public:
+  void startContextSection() override {}


startContextSection and endContextSection are empty funcs in both derived classes. Maybe introduce them when they are used? Also they aren't in the base class so override seems unneccessary?

Removed them.

github-actions · 2025-03-04T16:51:16Z

✅ With the latest revision this PR passed the C/C++ code formatter.

snehasish

lgtm

mtrofin · 2025-03-04T20:39:32Z

Merge activity

Mar 4, 3:39 PM EST: A user started a stack merge that includes this pull request via Graphite.
Mar 4, 3:41 PM EST: A user merged this pull request with Graphite.

Introduce a `ProfileWriter` abstraction to replace the callback passed to `__llvm_ctx_profile_fetch`. Subsequent changes will add support for flat profile collection (as in, collection of non-contextual profile for those functions not under a contextual root), which require also a change in the profile format. The abstraction makes it easy to add "write flat" - related capabilities without constantly complicating the signature of `__llvm_ctx_profile_fetch`.

mtrofin force-pushed the users/mtrofin/03-03-_ctxprof_profilewriter_abstraction branch 2 times, most recently from 91ab7e6 to 21d73ce Compare March 3, 2025 21:04

mtrofin requested review from kazutakahirata, snehasish and teresajohnson March 3, 2025 21:07

kazutakahirata approved these changes Mar 3, 2025

View reviewed changes

This was referenced Mar 3, 2025

[ctxprof] Profile format support for flat profiles #129592

Closed

[ctxprof][nfc] Prepare CtxProfAnalysis for flat profiles #129623

Merged

mtrofin force-pushed the users/mtrofin/03-03-_ctxprof_profilewriter_abstraction branch from 21d73ce to 5b223e7 Compare March 4, 2025 01:44

mtrofin mentioned this pull request Mar 4, 2025

[ctxprof] Prepare profile format for flat profiles #129626

Merged

mtrofin marked this pull request as ready for review March 4, 2025 02:03

llvmbot added compiler-rt PGO Profile Guided Optimizations labels Mar 4, 2025

mtrofin force-pushed the users/mtrofin/03-03-_ctxprof_profilewriter_abstraction branch from 5b223e7 to ce10a34 Compare March 4, 2025 03:05

snehasish reviewed Mar 4, 2025

View reviewed changes

mtrofin force-pushed the users/mtrofin/03-03-_ctxprof_profilewriter_abstraction branch from ce10a34 to 8fed7ea Compare March 4, 2025 16:47

mtrofin force-pushed the users/mtrofin/03-03-_ctxprof_profilewriter_abstraction branch 2 times, most recently from 8773741 to 7c64ba2 Compare March 4, 2025 18:40

[ctxprof] ProfileWriter abstraction

cb3945b

mtrofin force-pushed the users/mtrofin/03-03-_ctxprof_profilewriter_abstraction branch from 7c64ba2 to cb3945b Compare March 4, 2025 18:53

snehasish approved these changes Mar 4, 2025

View reviewed changes

mtrofin merged commit 1b46db7 into main Mar 4, 2025
9 of 10 checks passed

mtrofin deleted the users/mtrofin/03-03-_ctxprof_profilewriter_abstraction branch March 4, 2025 20:41

[ctxprof] ProfileWriter abstraction #129590

[ctxprof] ProfileWriter abstraction #129590

Uh oh!

Conversation

mtrofin commented Mar 3, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

mtrofin commented Mar 3, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

kazutakahirata left a comment

Choose a reason for hiding this comment

Uh oh!

llvmbot commented Mar 4, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

github-actions bot commented Mar 4, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

snehasish left a comment

Choose a reason for hiding this comment

Uh oh!

mtrofin commented Mar 4, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Merge activity

Uh oh!

Uh oh!

Uh oh!

mtrofin commented Mar 3, 2025 •

edited

Loading

mtrofin commented Mar 3, 2025 •

edited

Loading

github-actions bot commented Mar 4, 2025 •

edited

Loading

mtrofin commented Mar 4, 2025 •

edited

Loading