[ctx_prof] Extend `WorkloadImportsManager` to use the contextual profile #98682

mtrofin · 2024-07-12T19:14:56Z

Keeping the json-based input as it's useful for diagnostics or for driving the import by other means than contextual composition.

The support for the contextual profile is just another modality for constructing the import list (WorkloadImportsManager::Workloads). Everything else - i.e. the actual importing logic - is already independent from how that list was obtained.

Issue #89287

Keeping the json-based input as it's useful for diagnostics or for driving the import by other means than contextual composition. The support for the contextual profile is just another modality for constructing the import list (`WorkloadImportsManager::Workloads`). Everything else - i.e. the actual importing logic - is already independent from how that list was obtained.

llvmbot · 2024-07-12T19:15:28Z

@llvm/pr-subscribers-lto
@llvm/pr-subscribers-pgo

@llvm/pr-subscribers-llvm-transforms

Author: Mircea Trofin (mtrofin)

Changes

Keeping the json-based input as it's useful for diagnostics or for driving the import by other means than contextual composition.

The support for the contextual profile is just another modality for constructing the import list (WorkloadImportsManager::Workloads). Everything else - i.e. the actual importing logic - is already independent from how that list was obtained.

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

2 Files Affected:

(modified) llvm/lib/Transforms/IPO/FunctionImport.cpp (+68-8)
(modified) llvm/unittests/ProfileData/PGOCtxProfReaderWriterTest.cpp (+9)

diff --git a/llvm/lib/Transforms/IPO/FunctionImport.cpp b/llvm/lib/Transforms/IPO/FunctionImport.cpp
index 038785114a0cf..985e419c2bb5e 100644
--- a/llvm/lib/Transforms/IPO/FunctionImport.cpp
+++ b/llvm/lib/Transforms/IPO/FunctionImport.cpp
@@ -18,6 +18,7 @@
 #include "llvm/ADT/Statistic.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Bitcode/BitcodeReader.h"
+#include "llvm/Bitstream/BitstreamWriter.h"
 #include "llvm/IR/AutoUpgrade.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/Function.h"
@@ -30,6 +31,7 @@
 #include "llvm/IR/ModuleSummaryIndex.h"
 #include "llvm/IRReader/IRReader.h"
 #include "llvm/Linker/IRMover.h"
+#include "llvm/ProfileData/PGOCtxProfReader.h"
 #include "llvm/Support/Casting.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
@@ -174,6 +176,10 @@ static cl::opt<std::string> WorkloadDefinitions(
              "}"),
     cl::Hidden);
 
+static cl::opt<std::string>
+    ContextualProfile("thinlto-pgo-ctx-prof",
+                      cl::desc("Path to a contextual profile."), cl::Hidden);
+
 namespace llvm {
 extern cl::opt<bool> EnableMemProfContextDisambiguation;
 }
@@ -586,13 +592,7 @@ class WorkloadImportsManager : public ModuleImportsManager {
     LLVM_DEBUG(dbgs() << "[Workload] Done\n");
   }
 
-public:
-  WorkloadImportsManager(
-      function_ref<bool(GlobalValue::GUID, const GlobalValueSummary *)>
-          IsPrevailing,
-      const ModuleSummaryIndex &Index,
-      DenseMap<StringRef, FunctionImporter::ExportSetTy> *ExportLists)
-      : ModuleImportsManager(IsPrevailing, Index, ExportLists) {
+  void loadFromJson() {
     // Since the workload def uses names, we need a quick lookup
     // name->ValueInfo.
     StringMap<ValueInfo> NameToValueInfo;
@@ -672,6 +672,66 @@ class WorkloadImportsManager : public ModuleImportsManager {
       });
     }
   }
+
+  void loadFromCtxProf() {
+    std::error_code EC;
+    auto BufferOrErr = MemoryBuffer::getFileOrSTDIN(ContextualProfile);
+    if (std::error_code EC = BufferOrErr.getError()) {
+      report_fatal_error("Failed to open contextual profile file");
+      return;
+    }
+    auto Buffer = std::move(BufferOrErr.get());
+
+    BitstreamCursor Cursor(*Buffer);
+    PGOCtxProfileReader Reader(Cursor);
+    auto Ctx = Reader.loadContexts();
+    if (!Ctx) {
+      report_fatal_error("Failed to parse contextual profiles");
+      return;
+    }
+    const auto &CtxMap = *Ctx;
+    for (const auto &[RootGuid, Root] : CtxMap) {
+      auto RootVI = Index.getValueInfo(RootGuid);
+      if (!RootVI) {
+        LLVM_DEBUG(dbgs() << "[Workload] Root " << RootGuid
+                          << " not found in this linkage unit.\n");
+        continue;
+      }
+      if (RootVI.getSummaryList().size() != 1) {
+        LLVM_DEBUG(dbgs() << "[Workload] Root " << RootGuid
+                          << " should have exactly one summary, but has "
+                          << RootVI.getSummaryList().size() << ". Skipping.\n");
+        continue;
+      }
+      StringRef RootDefiningModule =
+          RootVI.getSummaryList().front()->modulePath();
+      LLVM_DEBUG(dbgs() << "[Workload] Root defining module for " << RootGuid
+                        << " is : " << RootDefiningModule << "\n");
+      DenseSet<GlobalValue::GUID> ContainedGUIDs;
+      auto &Set = Workloads[RootDefiningModule];
+      Root.getContainedGuids(ContainedGUIDs);
+      for (auto Guid : ContainedGUIDs)
+        if (auto VI = Index.getValueInfo(Guid))
+          Set.insert(VI);
+    }
+  }
+
+public:
+  WorkloadImportsManager(
+      function_ref<bool(GlobalValue::GUID, const GlobalValueSummary *)>
+          IsPrevailing,
+      const ModuleSummaryIndex &Index,
+      DenseMap<StringRef, FunctionImporter::ExportSetTy> *ExportLists)
+      : ModuleImportsManager(IsPrevailing, Index, ExportLists) {
+    if (ContextualProfile.empty() == WorkloadDefinitions.empty()) {
+      report_fatal_error(
+          "Pass only one of: -thinlto-pgo-ctx-prof or -thinlto-workload-def");
+      return;
+    }
+    if (!ContextualProfile.empty())
+      loadFromCtxProf();
+    loadFromJson();
+  }
 };
 
 std::unique_ptr<ModuleImportsManager> ModuleImportsManager::create(
@@ -679,7 +739,7 @@ std::unique_ptr<ModuleImportsManager> ModuleImportsManager::create(
         IsPrevailing,
     const ModuleSummaryIndex &Index,
     DenseMap<StringRef, FunctionImporter::ExportSetTy> *ExportLists) {
-  if (WorkloadDefinitions.empty()) {
+  if (WorkloadDefinitions.empty() && ContextualProfile.empty()) {
     LLVM_DEBUG(dbgs() << "[Workload] Using the regular imports manager.\n");
     return std::unique_ptr<ModuleImportsManager>(
         new ModuleImportsManager(IsPrevailing, Index, ExportLists));
diff --git a/llvm/unittests/ProfileData/PGOCtxProfReaderWriterTest.cpp b/llvm/unittests/ProfileData/PGOCtxProfReaderWriterTest.cpp
index d2cdbb28e2fce..6c6798ded00b5 100644
--- a/llvm/unittests/ProfileData/PGOCtxProfReaderWriterTest.cpp
+++ b/llvm/unittests/ProfileData/PGOCtxProfReaderWriterTest.cpp
@@ -115,6 +115,15 @@ TEST_F(PGOCtxProfRWTest, RoundTrip) {
     EXPECT_EQ(Ctxes.size(), 2U);
     for (auto &[G, R] : roots())
       checkSame(*R, Ctxes.find(G)->second);
+
+    DenseSet<GlobalValue::GUID> Guids;
+    Ctxes.at(1U).getContainedGuids(Guids);
+    EXPECT_THAT(Guids,
+                testing::WhenSorted(testing::ElementsAre(1U, 2U, 4U, 5U)));
+
+    Guids.clear();
+    Ctxes.at(3U).getContainedGuids(Guids);
+    EXPECT_THAT(Guids, testing::ElementsAre(3U));
   }
 }

…ally) unused one to Constants.

teresajohnson · 2024-07-17T19:47:04Z

llvm/lib/Transforms/IPO/FunctionImport.cpp

@@ -174,6 +174,10 @@ static cl::opt<std::string> WorkloadDefinitions(
             "}"),
    cl::Hidden);

+static cl::opt<std::string>
+    ContextualProfile("thinlto-pgo-ctx-prof",


Can you add a test to exercise this handling?

We don't currently have a textual representation of the contextual profile, and the import manager is internal to this file, meaning we can't write a unittest. I'd do either, though, if the new functionality was more involved, but it's really just taking data from the ctx profile reader - which is tested - and filling a set with it.

I meant a lit test. I think it would be good to add (reading the binary format like I am assuming the ctx profile reader tests do?), otherwise this is untested functionality.

Right - need to figure out where to place it, because it would need to be a test that runs a program to generate a profile. Meaning it needs compiler-rt.

teresajohnson · 2024-07-17T19:53:03Z

llvm/lib/Transforms/IPO/FunctionImport.cpp

+      Root.getContainedGuids(ContainedGUIDs);
+      for (auto Guid : ContainedGUIDs)
+        if (auto VI = Index.getValueInfo(Guid))
+          Set.insert(VI);


The json case has some debug printing of the eventual Set contents - needed here too?

Moved in a common place.

teresajohnson · 2024-07-17T19:56:14Z

llvm/unittests/ProfileData/PGOCtxProfReaderWriterTest.cpp

@@ -115,6 +115,15 @@ TEST_F(PGOCtxProfRWTest, RoundTrip) {
    EXPECT_EQ(Ctxes.size(), 2U);
    for (auto &[G, R] : roots())
      checkSame(*R, Ctxes.find(G)->second);
+
+    DenseSet<GlobalValue::GUID> Guids;
+    Ctxes.at(1U).getContainedGuids(Guids);


It is unclear to me how this test change is related to the rest of the PR.

It's not, but this is the API that the rest of this patch hinges on. I could introduce the test separately.

I would just commit that separately (no need for another review, that part lgtm) since that isn't a new API here.

done (a742693)

teresajohnson

lgtm with a couple small fixes

teresajohnson · 2024-07-29T19:26:56Z

llvm/test/ThinLTO/X86/ctxprof.ll

+; RUN: mkdir -p %t_baseline
+; RUN: mkdir -p %t_exp
+;
+; Normal run. m1 shouldn't get m2_f1 because it's not referenced from there.


and ditto for m2 and m1_f1?

teresajohnson · 2024-07-29T19:27:15Z

llvm/lib/Transforms/IPO/FunctionImport.cpp

@@ -174,6 +174,17 @@ static cl::opt<std::string> WorkloadDefinitions(
             "}"),
    cl::Hidden);

+static cl::opt<bool> ImportAssumeUniqueLocal(


Unrelated change should be removed

this is from #100448?

Oh sorry, I was looking at diffs from my last review, and wasn't expecting things from other patches to show up (did it get rebased?)

Yes, I needed to so I could get the llvm-ctxprof-util tool

…ile (llvm#98682) Keeping the json-based input as it's useful for diagnostics or for driving the import by other means than contextual composition. The support for the contextual profile is just another modality for constructing the import list (`WorkloadImportsManager::Workloads`). Everything else - i.e. the actual importing logic - is already independent from how that list was obtained.

Reverts: depends upon earlier ctx_prof reverts c99bd3c [ctx_prof] Extend `WorkloadImportsManager` to use the contextual profile (llvm#98682) Change-Id: I602defa5d5aa32c272ba831165cb39cf8f83b5ba

…ile (llvm#98682) Keeping the json-based input as it's useful for diagnostics or for driving the import by other means than contextual composition. The support for the contextual profile is just another modality for constructing the import list (`WorkloadImportsManager::Workloads`). Everything else - i.e. the actual importing logic - is already independent from how that list was obtained. Change-Id: I267f13222a7007aeffb690feac00af90e779301a

mtrofin requested a review from teresajohnson July 12, 2024 19:14

llvmbot added PGO Profile Guided Optimizations LTO Link time optimization (regular/full LTO or ThinLTO) llvm:transforms labels Jul 12, 2024

mtrofin added 2 commits July 12, 2024 12:15

Removed spurious include of BitstreamWriter, as well as (oportunistic…

8e10550

…ally) unused one to Constants.

reuse the ContainedGUIDs container.

b6479f0

teresajohnson reviewed Jul 17, 2024

View reviewed changes

mtrofin added 4 commits July 17, 2024 13:31

llvm_dbgs

157d38c

Merge branch 'main' into thinlto

bfc7245

Merge branch 'main' into thinlto

ef3da26

added test

7cab592

teresajohnson approved these changes Jul 29, 2024

View reviewed changes

mtrofin added 2 commits July 29, 2024 12:31

fixes

e5f9ec4

Merge branch 'main' into thinlto

bb58e0f

mtrofin merged commit c99bd3c into llvm:main Jul 29, 2024
5 of 7 checks passed

mtrofin deleted the thinlto branch July 29, 2024 22:06

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[ctx_prof] Extend `WorkloadImportsManager` to use the contextual profile #98682

[ctx_prof] Extend `WorkloadImportsManager` to use the contextual profile #98682

mtrofin commented Jul 12, 2024 •

edited

Loading

llvmbot commented Jul 12, 2024 •

edited

Loading

teresajohnson Jul 17, 2024

mtrofin Jul 17, 2024 •

edited

Loading

teresajohnson Jul 17, 2024

mtrofin Jul 17, 2024

mtrofin Jul 29, 2024

teresajohnson Jul 17, 2024

mtrofin Jul 17, 2024

teresajohnson Jul 17, 2024

mtrofin Jul 17, 2024

teresajohnson Jul 17, 2024

mtrofin Jul 17, 2024

teresajohnson left a comment

teresajohnson Jul 29, 2024

mtrofin Jul 29, 2024

teresajohnson Jul 29, 2024

mtrofin Jul 29, 2024

teresajohnson Jul 29, 2024

mtrofin Jul 29, 2024

[ctx_prof] Extend WorkloadImportsManager to use the contextual profile #98682

[ctx_prof] Extend WorkloadImportsManager to use the contextual profile #98682

Conversation

mtrofin commented Jul 12, 2024 • edited Loading

llvmbot commented Jul 12, 2024 • edited Loading

Choose a reason for hiding this comment

mtrofin Jul 17, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

teresajohnson left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

[ctx_prof] Extend `WorkloadImportsManager` to use the contextual profile #98682

[ctx_prof] Extend `WorkloadImportsManager` to use the contextual profile #98682

mtrofin commented Jul 12, 2024 •

edited

Loading

llvmbot commented Jul 12, 2024 •

edited

Loading

mtrofin Jul 17, 2024 •

edited

Loading