[llvm-profdata] Enabled functionality to write split-layout profile (#101795)

Using the flag `-split_layout` in llvm-profdata merge, the output
profile can write profiles with and without inlined function into two
different extbinary sections (and their FuncOffsetTable too). The
section without inlined functions are marked with `SecFlagFlat` and is
skipped by ThinLTO because it provides no useful info.

The split layout feature was already implemented in SampleProfWriter but
previously there is no way to use it from llvm-profdata.

Author: huangjd

llvm/docs/CommandGuide/llvm-profdata.rst

@@ -162,6 +162,12 @@ OPTIONS
  coverage for the optimized target. This option can only be used with
  sample-based profile in extbinary format.
 
+.. option:: --split-layout=[true|false]
+
+ Split the profile data section to two with one containing sample profiles with
+ inlined functions and the other not. This option can only be used with
+ sample-based profile in extbinary format.
+
 .. option:: --convert-sample-profile-layout=[nest|flat]
 
  Convert the merged profile into a profile with a new layout. Supported

llvm/include/llvm/ProfileData/SampleProfReader.h

@@ -495,9 +495,9 @@ class SampleProfileReader {
   /// are present.
   virtual void setProfileUseMD5() { ProfileIsMD5 = true; }
 
-  /// Don't read profile without context if the flag is set. This is only meaningful
-  /// for ExtBinary format.
-  virtual void setSkipFlatProf(bool Skip) {}
+  /// Don't read profile without context if the flag is set.
+  void setSkipFlatProf(bool Skip) { SkipFlatProf = Skip; }
+
   /// Return whether any name in the profile contains ".__uniq." suffix.
   virtual bool hasUniqSuffix() { return false; }
 
@@ -581,6 +581,10 @@ class SampleProfileReader {
   /// Whether the profile uses MD5 for Sample Contexts and function names. This
   /// can be one-way overriden by the user to force use MD5.
   bool ProfileIsMD5 = false;
+
+  /// If SkipFlatProf is true, skip functions marked with !Flat in text mode or
+  /// sections with SecFlagFlat flag in ExtBinary mode.
+  bool SkipFlatProf = false;
 };
 
 class SampleProfileReaderText : public SampleProfileReader {
@@ -789,10 +793,6 @@ class SampleProfileReaderExtBinaryBase : public SampleProfileReaderBinary {
   /// The set containing the functions to use when compiling a module.
   DenseSet<StringRef> FuncsToUse;
 
-  /// If SkipFlatProf is true, skip the sections with
-  /// SecFlagFlat flag.
-  bool SkipFlatProf = false;
-
 public:
   SampleProfileReaderExtBinaryBase(std::unique_ptr<MemoryBuffer> B,
                                    LLVMContext &C, SampleProfileFormat Format)
@@ -815,8 +815,6 @@ class SampleProfileReaderExtBinaryBase : public SampleProfileReaderBinary {
     return std::move(ProfSymList);
   };
 
-  void setSkipFlatProf(bool Skip) override { SkipFlatProf = Skip; }
-
 private:
   /// Read the profiles on-demand for the given functions. This is used after
   /// stale call graph matching finds new functions whose profiles aren't loaded

llvm/include/llvm/ProfileData/SampleProfWriter.h

@@ -28,9 +28,9 @@ namespace sampleprof {
 
 enum SectionLayout {
   DefaultLayout,
-  // The layout splits profile with context information from profile without
-  // context information. When Thinlto is enabled, ThinLTO postlink phase only
-  // has to load profile with context information and can skip the other part.
+  // The layout splits profile with inlined functions from profile without
+  // inlined functions. When Thinlto is enabled, ThinLTO postlink phase only
+  // has to load profile with inlined functions and can skip the other part.
   CtxSplitLayout,
   NumOfLayout,
 };
@@ -128,7 +128,7 @@ class SampleProfileWriter {
   virtual void setToCompressAllSections() {}
   virtual void setUseMD5() {}
   virtual void setPartialProfile() {}
-  virtual void resetSecLayout(SectionLayout SL) {}
+  virtual void setUseCtxSplitLayout() {}
 
 protected:
   SampleProfileWriter(std::unique_ptr<raw_ostream> &OS)
@@ -176,12 +176,22 @@ class SampleProfileWriterText : public SampleProfileWriter {
     return sampleprof_error::success;
   }
 
+  void setUseCtxSplitLayout() override {
+    MarkFlatProfiles = true;
+  }
+
 private:
   /// Indent level to use when writing.
   ///
   /// This is used when printing inlined callees.
   unsigned Indent = 0;
 
+  /// If set, writes metadata "!Flat" to functions without inlined functions.
+  /// This flag is for manual inspection only, it has no effect for the profile
+  /// reader because a text sample profile is read sequentially and functions
+  /// cannot be skipped.
+  bool MarkFlatProfiles = false;
+
   friend ErrorOr<std::unique_ptr<SampleProfileWriter>>
   SampleProfileWriter::create(std::unique_ptr<raw_ostream> &OS,
                               SampleProfileFormat Format);
@@ -242,11 +252,11 @@ const std::array<SmallVector<SecHdrTableEntry, 8>, NumOfLayout>
         // CtxSplitLayout
         SmallVector<SecHdrTableEntry, 8>({{SecProfSummary, 0, 0, 0, 0},
                                           {SecNameTable, 0, 0, 0, 0},
-                                          // profile with context
+                                          // profile with inlined functions
                                           // for next two sections
                                           {SecFuncOffsetTable, 0, 0, 0, 0},
                                           {SecLBRProfile, 0, 0, 0, 0},
-                                          // profile without context
+                                          // profile without inlined functions
                                           // for next two sections
                                           {SecFuncOffsetTable, 0, 0, 0, 0},
                                           {SecLBRProfile, 0, 0, 0, 0},
@@ -283,7 +293,11 @@ class SampleProfileWriterExtBinaryBase : public SampleProfileWriterBinary {
     ProfSymList = PSL;
   };
 
-  void resetSecLayout(SectionLayout SL) override {
+  void setUseCtxSplitLayout() override {
+    resetSecLayout(SectionLayout::CtxSplitLayout);
+  }
+
+  void resetSecLayout(SectionLayout SL) {
     verifySecLayout(SL);
 #ifndef NDEBUG
     // Make sure resetSecLayout is called before any flag setting.

llvm/lib/ProfileData/SampleProfReader.cpp

@@ -215,14 +215,22 @@ static bool ParseLine(const StringRef &Input, LineType &LineTy, uint32_t &Depth,
                       uint64_t &NumSamples, uint32_t &LineOffset,
                       uint32_t &Discriminator, StringRef &CalleeName,
                       DenseMap<StringRef, uint64_t> &TargetCountMap,
-                      uint64_t &FunctionHash, uint32_t &Attributes) {
+                      uint64_t &FunctionHash, uint32_t &Attributes,
+                      bool &IsFlat) {
   for (Depth = 0; Input[Depth] == ' '; Depth++)
     ;
   if (Depth == 0)
     return false;
 
   if (Input[Depth] == '!') {
     LineTy = LineType::Metadata;
+    // This metadata is only for manual inspection only. We already created a
+    // FunctionSamples and put it in the profile map, so there is no point
+    // to skip profiles even they have no use for ThinLTO.
+    if (Input == StringRef(" !Flat")) {
+      IsFlat = true;
+      return true;
+    }
     return parseMetadata(Input.substr(Depth), FunctionHash, Attributes);
   }
 
@@ -325,6 +333,8 @@ std::error_code SampleProfileReaderText::readImpl() {
   // top-level or nested function profile.
   uint32_t DepthMetadata = 0;
 
+  std::vector<SampleContext *> FlatSamples;
+
   ProfileIsFS = ProfileIsFSDisciminator;
   FunctionSamples::ProfileIsFS = ProfileIsFS;
   for (; !LineIt.is_at_eof(); ++LineIt) {
@@ -368,9 +378,10 @@ std::error_code SampleProfileReaderText::readImpl() {
       LineType LineTy;
       uint64_t FunctionHash = 0;
       uint32_t Attributes = 0;
+      bool IsFlat = false;
       if (!ParseLine(*LineIt, LineTy, Depth, NumSamples, LineOffset,
                      Discriminator, FName, TargetCountMap, FunctionHash,
-                     Attributes)) {
+                     Attributes, IsFlat)) {
         reportError(LineIt.line_number(),
                     "Expected 'NUM[.NUM]: NUM[ mangled_name:NUM]*', found " +
                         *LineIt);
@@ -426,12 +437,26 @@ std::error_code SampleProfileReaderText::readImpl() {
         if (Attributes & (uint32_t)ContextShouldBeInlined)
           ProfileIsPreInlined = true;
         DepthMetadata = Depth;
+        if (IsFlat) {
+          if (Depth == 1)
+            FlatSamples.push_back(&FProfile.getContext());
+          else
+            Ctx.diagnose(DiagnosticInfoSampleProfile(
+                Buffer->getBufferIdentifier(), LineIt.line_number(),
+                "!Flat may only be used at top level function.", DS_Warning));
+        }
         break;
       }
       }
     }
   }
 
+  // Honor the option to skip flat functions. Since they are already added to
+  // the profile map, remove them all here.
+  if (SkipFlatProf)
+    for (SampleContext *FlatSample : FlatSamples)
+      Profiles.erase(*FlatSample);
+
   assert((CSProfileCount == 0 || CSProfileCount == Profiles.size()) &&
          "Cannot have both context-sensitive and regular profile");
   ProfileIsCS = (CSProfileCount > 0);
@@ -1026,7 +1051,7 @@ std::error_code SampleProfileReaderExtBinaryBase::readImpl() {
     if (!Entry.Size)
       continue;
 
-    // Skip sections without context when SkipFlatProf is true.
+    // Skip sections without inlined functions when SkipFlatProf is true.
     if (SkipFlatProf && hasSecFlag(Entry, SecCommonFlags::SecFlagFlat))
       continue;
 

llvm/lib/ProfileData/SampleProfWriter.cpp

@@ -624,6 +624,9 @@ std::error_code SampleProfileWriterText::writeSample(const FunctionSamples &S) {
     LineCount++;
   }
 
+  if (Indent == 0 && MarkFlatProfiles && S.getCallsiteSamples().size() == 0)
+    OS << " !Flat\n";
+
   return sampleprof_error::success;
 }
 

llvm/test/tools/llvm-profdata/Inputs/split-layout.profdata

@@ -0,0 +1,22 @@
+RUN: llvm-profdata merge --sample --extbinary --split-layout %p/Inputs/sample-profile.proftext -o %t-output
+RUN: diff %t-output %p/Inputs/split-layout.profdata
+
+RUN: llvm-profdata merge --sample --text --split-layout %t-output | FileCheck %s
+CHECK: main:184019:0
+CHECK-NEXT:  4: 534
+CHECK-NEXT:  4.2: 534
+CHECK-NEXT:  5: 1075
+CHECK-NEXT:  5.1: 1075
+CHECK-NEXT:  6: 2080
+CHECK-NEXT:  7: 534
+CHECK-NEXT:  9: 2064 _Z3bari:1471 _Z3fooi:631
+CHECK-NEXT:  10: inline1:1000
+CHECK-NEXT:   1: 1000
+CHECK-NEXT:  10: inline2:2000
+CHECK-NEXT:   1: 2000
+CHECK-NEXT: _Z3bari:20301:1437
+CHECK-NEXT:  1: 1437
+CHECK-NEXT:  !Flat
+CHECK-NEXT: _Z3fooi:7711:610
+CHECK-NEXT:  1: 610
+CHECK-NEXT:  !Flat

llvm/test/tools/llvm-profdata/sample-split-layout.test

@@ -207,6 +207,12 @@ cl::opt<bool> GenPartialProfile(
     "gen-partial-profile", cl::init(false), cl::Hidden,
     cl::sub(MergeSubcommand),
     cl::desc("Generate a partial profile (only meaningful for -extbinary)"));
+cl::opt<bool> SplitLayout(
+    "split-layout", cl::init(false), cl::Hidden,
+    cl::sub(MergeSubcommand),
+    cl::desc("Split the profile to two sections with one containing sample "
+             "profiles with inlined functions and the other without (only "
+             "meaningful for -extbinary)"));
 cl::opt<std::string> SupplInstrWithSample(
     "supplement-instr-with-sample", cl::init(""), cl::Hidden,
     cl::sub(MergeSubcommand),
@@ -1467,6 +1473,13 @@ static void handleExtBinaryWriter(sampleprof::SampleProfileWriter &Writer,
                                   sampleprof::ProfileSymbolList &WriterList,
                                   bool CompressAllSections, bool UseMD5,
                                   bool GenPartialProfile) {
+  if (SplitLayout) {
+    if (OutputFormat == PF_Binary)
+      warn("-split-layout is ignored. Specify -extbinary to enable it");
+    else
+      Writer.setUseCtxSplitLayout();
+  }
+
   populateProfileSymbolList(Buffer, WriterList);
   if (WriterList.size() > 0 && OutputFormat != PF_Ext_Binary)
     warn("Profile Symbol list is not empty but the output format is not "