Weigh program profile matches by the profile's counts (#2528)

native/llvm-codegen/src/emit.cpp

@@ -252,15 +252,20 @@ struct UnnameStructTypes : PassInfoMixin<UnnameStructTypes> {
 // over partitions: functions that received its counts, functions whose
 // control flow changed since it was collected (hash mismatch), and functions
 // it has no record of (new since it was collected, or never run in training
-// when the profile was merged with --sparse). A partition reused from the cache without its counts is counted
-// in Unknown. Recovered counts the mismatched functions that --profile-blocks
-// annotated, so it never exceeds Mismatched.
+// when the profile was merged with --sparse). A partition reused from the cache
+// without its counts is counted in Unknown. Recovered counts the mismatched
+// functions that --profile-blocks annotated, so it never exceeds Mismatched.
+// The two weights sum the profile's counters for the annotated and for the
+// mismatched functions, so a changed hot function weighs more than a changed
+// cold one.
 struct ProfileMatch {
   std::atomic<uint64_t> Annotated{0};
   std::atomic<uint64_t> Mismatched{0};
   std::atomic<uint64_t> Missing{0};
   std::atomic<uint64_t> Recovered{0};
   std::atomic<uint64_t> Unknown{0};
+  std::atomic<uint64_t> AnnotatedCounts{0};
+  std::atomic<uint64_t> MismatchedCounts{0};
   // Set once the annotation before inlining has run. The context-sensitive
   // annotation after inlining, which a single-module pipeline also runs,
   // reports its mismatched and missing functions again; they are not
@@ -272,27 +277,32 @@ struct ProfileMatch {
     Mismatched += Other.Mismatched;
     Missing += Other.Missing;
     Recovered += Other.Recovered;
+    AnnotatedCounts += Other.AnnotatedCounts;
+    MismatchedCounts += Other.MismatchedCounts;
   }
 
   std::string serialize() const {
-    return formatv("{0} {1} {2} {3}", Annotated.load(), Mismatched.load(),
-                   Missing.load(), Recovered.load())
+    return formatv("{0} {1} {2} {3} {4} {5}", Annotated.load(),
+                   Mismatched.load(), Missing.load(), Recovered.load(),
+                   AnnotatedCounts.load(), MismatchedCounts.load())
         .str();
   }
 
   bool deserialize(StringRef Text) {
-    SmallVector<StringRef, 4> Parts;
+    SmallVector<StringRef, 6> Parts;
     Text.split(Parts, ' ');
-    uint64_t Values[4];
-    if (Parts.size() != 4)
+    uint64_t Values[6];
+    if (Parts.size() != 6)
       return false;
-    for (size_t I = 0; I < 4; ++I)
+    for (size_t I = 0; I < 6; ++I)
       if (Parts[I].getAsInteger(10, Values[I]))
         return false;
     Annotated += Values[0];
     Mismatched += Values[1];
     Missing += Values[2];
     Recovered += Values[3];
+    AnnotatedCounts += Values[4];
+    MismatchedCounts += Values[5];
     return true;
   }
 };
@@ -430,11 +440,77 @@ void reportProfileMatch(const ProfileMatch &Match) {
                        {"mismatched", Match.Mismatched.load()},
                        {"missing", Match.Missing.load()},
                        {"recovered", Match.Recovered.load()},
-                       {"unknown_partitions", Match.Unknown.load()}};
+                       {"unknown_partitions", Match.Unknown.load()},
+                       {"annotated_counts", Match.AnnotatedCounts.load()},
+                       {"mismatched_counts", Match.MismatchedCounts.load()}};
   json::Object Response{{"ok", true}, {"profile", std::move(Profile)}};
   outs() << formatv("{0}\n", json::Value(std::move(Response)));
 }
 
+// The sum of each function's counters in the indexed profile at `Path`, by
+// profile name, read once per process. An unreadable profile has no counts
+// (the annotation reports it).
+const StringMap<uint64_t> &profileCounts(const std::string &Path) {
+  static std::once_flag Once;
+  static StringMap<uint64_t> Counts;
+  std::call_once(Once, [&Path] {
+    Expected<std::unique_ptr<IndexedInstrProfReader>> Reader =
+        IndexedInstrProfReader::create(Path, *vfs::getRealFileSystem());
+    if (!Reader) {
+      consumeError(Reader.takeError());
+      return;
+    }
+    for (const NamedInstrProfRecord &Record : **Reader) {
+      uint64_t Sum = 0;
+      for (uint64_t Count : Record.Counts)
+        Sum += Count;
+      Counts[Record.Name] += Sum;
+    }
+    consumeError((*Reader)->getError());
+  });
+  return Counts;
+}
+
+// Weighs one module's profile match by the profile's counters: collects the
+// annotated and the mismatched functions while the module is optimized,
+// then adds their counts to the match.
+class MatchWeigher {
+public:
+  explicit MatchWeigher(const Module &M) {
+    for (const Function &F : M)
+      if (!F.isDeclaration())
+        ProfileNames[F.getName()] = getIRPGOFuncName(F);
+  }
+
+  void weigh(const std::string &ProfileUse, ProfileMatch &Match) const {
+    const StringMap<uint64_t> &Counts = profileCounts(ProfileUse);
+    auto CountsOf = [&Counts](StringRef Name) -> uint64_t {
+      auto It = Counts.find(Name);
+      return It == Counts.end() ? 0 : It->second;
+    };
+    // A context-sensitive profile annotates the module twice; each function
+    // counts once.
+    StringSet<> Seen;
+    for (const std::string &Name : Annotated)
+      if (Seen.insert(Name).second)
+        Match.AnnotatedCounts += CountsOf(Name);
+    // The mismatch warnings name functions by symbol.
+    for (const std::string &Name : Mismatched) {
+      auto It = ProfileNames.find(Name);
+      StringRef ProfileName =
+          It == ProfileNames.end() ? StringRef(Name) : StringRef(It->second);
+      if (Seen.insert(ProfileName).second)
+        Match.MismatchedCounts += CountsOf(ProfileName);
+    }
+  }
+
+  std::vector<std::string> Annotated;
+  std::vector<std::string> Mismatched;
+
+private:
+  StringMap<std::string> ProfileNames;
+};
+
 // Reads the profile's header once, before any partition needs it: it must
 // hold IR-level counts. Context-sensitive counts lay out the code by them;
 // a profile that has them cannot take a second set.
@@ -533,13 +609,15 @@ struct Pipeline {
         });
   }
 
-  // Adds to Names the profile names of the functions the annotation gave
-  // counts to, as soon as that pass has run.
+  // Adds to Names the profile names of the functions the annotation before
+  // inlining gave counts to, as soon as that pass has run.
   void collectAnnotated(std::vector<std::string> &Names) {
     Callbacks.registerAfterPassCallback(
-        [&Names](StringRef PassID, Any IR, const PreservedAnalyses &) {
-          if (PassID != "PGOInstrumentationUse")
+        [&Names, Done = false](StringRef PassID, Any IR,
+                               const PreservedAnalyses &) mutable {
+          if (Done || PassID != "PGOInstrumentationUse")
             return;
+          Done = true;
           if (const auto *M = llvm::any_cast<const Module *>(&IR))
             for (const Function &F : **M) {
               std::optional<Function::ProfileCount> Entry = F.getEntryCount();
@@ -732,13 +810,21 @@ emitModule(Module &M, TargetMachine &Machine, OptimizationLevel Level,
            StringRef OutputPath, std::vector<SmallString<256>> &Temporaries,
            ProfileMatch &Match) {
   bool Annotating = Profile && Profile->Action == PGOOptions::IRUse;
+  std::string ProfileUse = Annotating ? Profile->ProfileFile : "";
+  MatchWeigher Weigher(M);
   if (Annotating)
     M.getContext().setDiagnosticHandler(
-        std::make_unique<ProfileMatchHandler>(Match));
+        std::make_unique<ProfileMatchHandler>(Match, &Weigher.Mismatched));
   Pipeline P(Machine, Level, std::move(Profile));
-  if (Annotating)
+  if (Annotating) {
+    // Registered first, so it lists the functions the profile itself
+    // annotated, before block counts recover any.
+    P.collectAnnotated(Weigher.Annotated);
     P.afterProfileAnnotation(Match);
+  }
   P.PB.buildPerModuleDefaultPipeline(Level).run(M, P.MAM);
+  if (Annotating)
+    Weigher.weigh(ProfileUse, Match);
   if (std::optional<EmitFailure> Failure =
           verify(M, "post_optimization_verification_failed"))
     return Failure;
@@ -929,14 +1015,16 @@ compileOptimizedPartitions(const EmitOptions &Options,
       return Failures.set({"target_machine_failed", LookupError});
     LLVMContext Context;
     ProfileMatch PartitionMatch;
+    std::vector<std::string> MismatchedNames;
     if (Annotating)
-      Context.setDiagnosticHandler(
-          std::make_unique<ProfileMatchHandler>(PartitionMatch));
+      Context.setDiagnosticHandler(std::make_unique<ProfileMatchHandler>(
+          PartitionMatch, &MismatchedNames));
     Expected<std::unique_ptr<Module>> Loaded =
         loadPartition(Options, Plan.Modules[I], I, Context, *Machine);
     if (!Loaded)
       return Failures.set({"partition_failed", toString(Loaded.takeError())});
     Module &Mod = **Loaded;
+    MatchWeigher Weigher(Mod);
     if (!Options.ImportBitcode.empty()) {
       // See emit(): the attribute is inert without a sanitizer pass and
       // would block inlining of the imported runtime bodies.
@@ -955,8 +1043,12 @@ compileOptimizedPartitions(const EmitOptions &Options,
       // Profile instrumentation or annotation happens here, before the thin
       // link; the post-link optimization reads the annotated weights.
       Pipeline P(*Machine, Level, profileOptions(Options));
-      if (Annotating)
+      if (Annotating) {
+        // Registered first, so it lists the functions the profile itself
+        // annotated, before block counts recover any.
+        P.collectAnnotated(Weigher.Annotated);
         P.afterProfileAnnotation(PartitionMatch);
+      }
       ModulePassManager MPM = P.PB.buildThinLTOPreLinkDefaultPipeline(Level);
       MPM.addPass(UnnameStructTypes());
       // The module hash is set at the thin link instead (see Names).
@@ -966,6 +1058,8 @@ compileOptimizedPartitions(const EmitOptions &Options,
       MPM.run(Mod, P.MAM);
     }
     if (Annotating) {
+      Weigher.Mismatched = std::move(MismatchedNames);
+      Weigher.weigh(Options.ProfileUse, PartitionMatch);
       Match.add(PartitionMatch);
       // Stored before the summary, so a reused summary finds its counts.
       Cache.store(Key, "match", PartitionMatch.serialize());
@@ -1463,15 +1557,22 @@ std::optional<EmitFailure>
 emitSplitCodegen(const EmitOptions &Options, Module &M, TargetMachine &Machine,
                  OptimizationLevel Level, ProfileMatch &Match) {
   bool Annotating = !Options.ProfileUse.empty();
+  MatchWeigher Weigher(M);
   if (Annotating)
     M.getContext().setDiagnosticHandler(
-        std::make_unique<ProfileMatchHandler>(Match));
+        std::make_unique<ProfileMatchHandler>(Match, &Weigher.Mismatched));
   {
     Pipeline P(Machine, Level, profileOptions(Options));
-    if (Annotating)
+    if (Annotating) {
+      // Registered first, so it lists the functions the profile itself
+      // annotated, before block counts recover any.
+      P.collectAnnotated(Weigher.Annotated);
       P.afterProfileAnnotation(Match);
+    }
     P.PB.buildPerModuleDefaultPipeline(Level).run(M, P.MAM);
   }
+  if (Annotating)
+    Weigher.weigh(Options.ProfileUse, Match);
   if (std::optional<EmitFailure> Failure =
           verify(M, "post_optimization_verification_failed"))
     return Failure;

packages/compiler/src/backend/native-codegen.ts

@@ -563,6 +563,9 @@ function helperProfileMatch(match: NativeProfileMatch): Record<string, number> {
     missing: match.missing,
     recovered: match.recovered,
     unknown_partitions: match.unknownPartitions,
+    ...(match.annotatedCounts === undefined || match.mismatchedCounts === undefined
+      ? {}
+      : { annotated_counts: match.annotatedCounts, mismatched_counts: match.mismatchedCounts }),
   };
 }
 

packages/compiler/src/backend/native-profile.test.ts

@@ -145,6 +145,20 @@ test("the helper's profile match is read from its output", () => {
       '{"ok":true,"profile":{"annotated":1,"mismatched":0,"missing":0,"recovered":-1,"unknown_partitions":0}}\n',
     ),
   ).toBeNull();
+  // With the profile's counters on the annotated and the mismatched functions.
+  expect(
+    parseNativeProfileMatch(
+      '{"ok":true,"profile":{"annotated":120,"mismatched":7,"missing":40,"recovered":0,"unknown_partitions":0,"annotated_counts":9000,"mismatched_counts":1000}}\n',
+    ),
+  ).toEqual({
+    annotated: 120,
+    mismatched: 7,
+    missing: 40,
+    recovered: 0,
+    unknownPartitions: 0,
+    annotatedCounts: 9000,
+    mismatchedCounts: 1000,
+  });
   // A helper that predates the report prints nothing.
   expect(parseNativeProfileMatch("")).toBeNull();
   expect(parseNativeProfileMatch('{"ok":true}\n')).toBeNull();

packages/compiler/src/backend/native-profile.ts

@@ -111,13 +111,16 @@ export function nativeProfileHelperArgs(profile: NativeProfileOptions | undefine
  * profile has no record of (new, or never run while it was collected),
  * mismatched functions that the block-count file annotated (never more
  * than the mismatched ones), and partitions reused from the cache without
- * their counts. */
+ * their counts. Helpers that weigh the match also report the profile's
+ * counters on the annotated and on the mismatched functions. */
 export interface NativeProfileMatch {
   annotated: number;
   mismatched: number;
   missing: number;
   recovered: number;
   unknownPartitions: number;
+  annotatedCounts?: number;
+  mismatchedCounts?: number;
 }
 
 /** The profile match in a successful helper emit's output, or null when the
@@ -150,7 +153,18 @@ export function parseNativeProfileMatch(stdout: string): NativeProfileMatch | nu
       unknownPartitions === null
     )
       continue;
-    return { annotated, mismatched, missing, recovered, unknownPartitions };
+    const annotatedCounts = count("annotated_counts");
+    const mismatchedCounts = count("mismatched_counts");
+    return {
+      annotated,
+      mismatched,
+      missing,
+      recovered,
+      unknownPartitions,
+      ...(annotatedCounts === null || mismatchedCounts === null
+        ? {}
+        : { annotatedCounts, mismatchedCounts }),
+    };
   }
   return null;
 }

packages/compiler/test/native-profile.test.ts

@@ -509,6 +509,27 @@ describe.runIf(helper !== null)("LLVM profile-guided optimization", () => {
       expect(await run(await link(dir, objects, "edited"), 2)).toBe(expectedExit(2, 7));
     });
 
+    test("the match report weighs the functions that changed by their counts", async () => {
+      const dir = await mkdtemp(join(tmpdir(), "scriptc-profile-weights-"));
+      dirs.push(dir);
+      const { profile } = await collect(dir, program(), 1, "train");
+      // f4 gained a branch, so the profile no longer describes it.
+      const edited = program(0, true);
+      const { outputs, stdout } = await emitReporting(dir, edited, {
+        outputs: 1,
+        name: "edited",
+        use: profile,
+      });
+      expect(await run(await link(dir, outputs, "edited"), 2)).toBe(expectedExit(2));
+      const { profile: match } = JSON.parse(stdout.trim().split("\n").at(-1)!) as {
+        profile: Record<string, number>;
+      };
+      expect(match["mismatched"]).toBe(1);
+      // main and the other f<index> keep their counts; f4's are reported apart.
+      expect(match["mismatched_counts"]).toBeGreaterThan(0);
+      expect(match["annotated_counts"]).toBeGreaterThan(match["mismatched_counts"]!);
+    });
+
     test("the match report counts annotated functions that inlining removes", async () => {
       const dir = await mkdtemp(join(tmpdir(), "scriptc-profile-match-"));
       dirs.push(dir);
@@ -521,7 +542,11 @@ describe.runIf(helper !== null)("LLVM profile-guided optimization", () => {
       expect(await run(await link(dir, outputs, "guided"), 2)).toBe(inlinedExit(2));
       // Every g<index> and main received counts, though optimization then
       // inlined the helpers into main and deleted them.
-      expect(JSON.parse(stdout.trim().split("\n").at(-1)!)).toEqual({
+      const report = JSON.parse(stdout.trim().split("\n").at(-1)!) as {
+        ok: boolean;
+        profile: Record<string, number>;
+      };
+      expect(report).toMatchObject({
         ok: true,
         profile: {
           annotated: INLINED + 1,
@@ -529,8 +554,10 @@ describe.runIf(helper !== null)("LLVM profile-guided optimization", () => {
           missing: 0,
           recovered: 0,
           unknown_partitions: 0,
+          mismatched_counts: 0,
         },
       });
+      expect(report.profile["annotated_counts"]).toBeGreaterThan(0);
       const symbols = await execFileAsync("nm", [outputs[0]!]);
       expect(symbols.stdout).not.toMatch(/\bg\d\b/);
     });

scripts/build-native-cli.mts

@@ -190,8 +190,17 @@ const match = useProfile && "artifact" in result ? result.artifact.profileMatch
 if (match !== undefined) {
   const profiled = match.annotated + match.mismatched;
   const share = profiled === 0 ? 0 : match.annotated / profiled;
+  // The same by the profile's counters: a changed hot function loses more
+  // of the profile's gain than a changed cold one.
+  const counts =
+    match.annotatedCounts === undefined || match.mismatchedCounts === undefined
+      ? 0
+      : match.annotatedCounts + match.mismatchedCounts;
   console.log(
     `native compiler profile: ${match.annotated} of ${profiled} profiled functions match (${(100 * share).toFixed(1)}%); ${match.missing} functions have no profile` +
+      (counts > 0
+        ? `; changed functions hold ${((100 * match.mismatchedCounts!) / counts).toFixed(1)}% of its counts`
+        : "") +
       (match.recovered > 0
         ? `; ${match.recovered} functions annotated from the profile's block counts`
         : "") +