[NFCI]Add SampleRecord::serialize and LineLocation::serialize to simplify FunctionSamples serialization

[mingmingl-llvm]

No description provided.

[llvmbot]

@llvm/pr-subscribers-pgo

Author: Mingming Liu (mingmingl-llvm)

Changes

---

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

3 Files Affected:

• (modified) llvm/include/llvm/ProfileData/SampleProf.h (+10-1)
• (modified) llvm/lib/ProfileData/SampleProf.cpp (+25)
• (modified) llvm/lib/ProfileData/SampleProfWriter.cpp (+3-13)

diff --git a/llvm/include/llvm/ProfileData/SampleProf.h b/llvm/include/llvm/ProfileData/SampleProf.h
index c7155b3f982c6..3704de6e49f40 100644
--- a/llvm/include/llvm/ProfileData/SampleProf.h
+++ b/llvm/include/llvm/ProfileData/SampleProf.h
@@ -15,17 +15,18 @@
 #define LLVM_PROFILEDATA_SAMPLEPROF_H
 
 #include "llvm/ADT/DenseSet.h"
+#include "llvm/ADT/MapVector.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/GlobalValue.h"
 #include "llvm/ProfileData/FunctionId.h"
+#include "llvm/ProfileData/HashKeyMap.h"
 #include "llvm/Support/Allocator.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/ErrorOr.h"
 #include "llvm/Support/MathExtras.h"
-#include "llvm/ProfileData/HashKeyMap.h"
 #include <algorithm>
 #include <cstdint>
 #include <list>
@@ -283,6 +284,9 @@ struct LineLocation {
   void print(raw_ostream &OS) const;
   void dump() const;
 
+  // Serialize the line location to the output stream using ULEB128 encoding.
+  void serialize(raw_ostream &OS);
+
   bool operator<(const LineLocation &O) const {
     return LineOffset < O.LineOffset ||
            (LineOffset == O.LineOffset && Discriminator < O.Discriminator);
@@ -427,6 +431,11 @@ class SampleRecord {
   sampleprof_error merge(const SampleRecord &Other, uint64_t Weight = 1);
   void print(raw_ostream &OS, unsigned Indent) const;
   void dump() const;
+  /// Serialize the sample record to the output stream using ULEB128 encoding.
+  /// The \p NameTable is used to map function names to their IDs.
+  std::error_code
+  serialize(raw_ostream &OS,
+            const MapVector<FunctionId, uint32_t> &NameTable) const;
 
   bool operator==(const SampleRecord &Other) const {
     return NumSamples == Other.NumSamples && CallTargets == Other.CallTargets;
diff --git a/llvm/lib/ProfileData/SampleProf.cpp b/llvm/lib/ProfileData/SampleProf.cpp
index 4d48de9bc7d63..0c47df3c55dec 100644
--- a/llvm/lib/ProfileData/SampleProf.cpp
+++ b/llvm/lib/ProfileData/SampleProf.cpp
@@ -20,6 +20,7 @@
 #include "llvm/Support/Compiler.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/ErrorHandling.h"
+#include "llvm/Support/LEB128.h"
 #include "llvm/Support/raw_ostream.h"
 #include <string>
 #include <system_error>
@@ -126,10 +127,34 @@ sampleprof_error SampleRecord::merge(const SampleRecord &Other,
   return Result;
 }
 
+std::error_code SampleRecord::serialize(
+    raw_ostream &OS, const MapVector<FunctionId, uint32_t> &NameTable) const {
+  encodeULEB128(getSamples(), OS);
+  encodeULEB128(getCallTargets().size(), OS);
+  for (const auto &J : getSortedCallTargets()) {
+    FunctionId Callee = J.first;
+    uint64_t CalleeSamples = J.second;
+    if (auto NameIndexIter = NameTable.find(Callee);
+        NameIndexIter != NameTable.end()) {
+      encodeULEB128(NameIndexIter->second, OS);
+    } else {
+      // If the callee is not in the name table, we cannot serialize it.
+      return sampleprof_error::truncated_name_table;
+    }
+    encodeULEB128(CalleeSamples, OS);
+  }
+  return sampleprof_error::success;
+}
+
 #if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP)
 LLVM_DUMP_METHOD void LineLocation::dump() const { print(dbgs()); }
 #endif
 
+void LineLocation::serialize(raw_ostream &OS) {
+  encodeULEB128(LineOffset, OS);
+  encodeULEB128(Discriminator, OS);
+}
+
 /// Print the sample record to the stream \p OS indented by \p Indent.
 void SampleRecord::print(raw_ostream &OS, unsigned Indent) const {
   OS << NumSamples;
diff --git a/llvm/lib/ProfileData/SampleProfWriter.cpp b/llvm/lib/ProfileData/SampleProfWriter.cpp
index 8b164f24e06e1..71d2f522fc295 100644
--- a/llvm/lib/ProfileData/SampleProfWriter.cpp
+++ b/llvm/lib/ProfileData/SampleProfWriter.cpp
@@ -827,17 +827,8 @@ std::error_code SampleProfileWriterBinary::writeBody(const FunctionSamples &S) {
   for (const auto &I : S.getBodySamples()) {
     LineLocation Loc = I.first;
     const SampleRecord &Sample = I.second;
-    encodeULEB128(Loc.LineOffset, OS);
-    encodeULEB128(Loc.Discriminator, OS);
-    encodeULEB128(Sample.getSamples(), OS);
-    encodeULEB128(Sample.getCallTargets().size(), OS);
-    for (const auto &J : Sample.getSortedCallTargets()) {
-      FunctionId Callee = J.first;
-      uint64_t CalleeSamples = J.second;
-      if (std::error_code EC = writeNameIdx(Callee))
-        return EC;
-      encodeULEB128(CalleeSamples, OS);
-    }
+    Loc.serialize(OS);
+    Sample.serialize(OS, getNameTable());
   }
 
   // Recursively emit all the callsite samples.
@@ -849,8 +840,7 @@ std::error_code SampleProfileWriterBinary::writeBody(const FunctionSamples &S) {
     for (const auto &FS : J.second) {
       LineLocation Loc = J.first;
       const FunctionSamples &CalleeSamples = FS.second;
-      encodeULEB128(Loc.LineOffset, OS);
-      encodeULEB128(Loc.Discriminator, OS);
+      Loc.serialize(OS);
       if (std::error_code EC = writeBody(CalleeSamples))
         return EC;
     }

[mingmingl-llvm]

Thanks for the reviews!

Test failure in LLVM.CodeGen/PowerPC/atomic-compare-exchange-weak.ll (https://buildkite.com/llvm-project/github-pull-requests/builds/183626#019713a8-f79f-4303-b34c-4e958ffc7758) are unrelated. Submitting now.