[llvm] [ProfileData] Avoid magic layout indices in SampleProfileWriter (NFC) (PR #215981)

[via llvm-commits]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-pgo

Author: Kazu Hirata (kazutakahirata)

<details>
<summary>Changes</summary>

This patch removes hardcoded section layout indices in
SampleProfileWriterExtBinary in favor of finding the first unwritten
matching entry in SectionHdrLayout.

Previously, writeDefaultLayout and writeCtxSplitLayout passed magic
integer indices to writeOneSection to specify the position of each
section in SectionHdrLayout.

This patch introduces findUnwrittenEntry to dynamically look up the
matching unwritten section entry in SectionHdrLayout and determine its
ordinal among sections of the same type.  This allows
writeDefaultLayout and writeCtxSplitLayout to iterate over declarative
arrays of sections without magic indices, and allows removing the
index-based addSectionFlag overload.

Assisted-by: Antigravity


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


2 Files Affected:

- (modified) llvm/include/llvm/ProfileData/SampleProfWriter.h (+6-8) 
- (modified) llvm/lib/ProfileData/SampleProfWriter.cpp (+41-43) 


``````````diff
diff --git a/llvm/include/llvm/ProfileData/SampleProfWriter.h b/llvm/include/llvm/ProfileData/SampleProfWriter.h
index db59fb7579967..c0ebfdcc749a8 100644
--- a/llvm/include/llvm/ProfileData/SampleProfWriter.h
+++ b/llvm/include/llvm/ProfileData/SampleProfWriter.h
@@ -337,11 +337,6 @@ class LLVM_ABI SampleProfileWriterExtBinaryBase
         addSecFlag(Entry, Flag);
     }
   }
-  template <class SecFlagType>
-  void addSectionFlag(uint32_t SectionIdx, SecFlagType Flag) {
-    addSecFlag(SectionHdrLayout[SectionIdx], Flag);
-  }
-
   void addContext(const SampleContext &Context) override;
 
   // placeholder for subclasses to dispatch their own section writers.
@@ -352,9 +347,12 @@ class LLVM_ABI SampleProfileWriterExtBinaryBase
   // specify the order to write sections.
   virtual std::error_code writeSections(const SampleProfileMap &ProfileMap) = 0;
 
-  // Dispatch section writer for each section. \p LayoutIdx is the sequence
-  // number indicating where the section is located in SectionHdrLayout.
-  virtual std::error_code writeOneSection(SecType Type, uint32_t LayoutIdx,
+  // Find the first unwritten entry in SectionHdrLayout matching Type, returning
+  // its layout index.
+  unsigned findUnwrittenEntry(SecType Type);
+
+  // Dispatch section writer for each section.
+  virtual std::error_code writeOneSection(SecType Type,
                                           const SampleProfileMap &ProfileMap);
 
   // Helper function to write name table.
diff --git a/llvm/lib/ProfileData/SampleProfWriter.cpp b/llvm/lib/ProfileData/SampleProfWriter.cpp
index 5c6ca653b0cd9..c624bf5ccc4b7 100644
--- a/llvm/lib/ProfileData/SampleProfWriter.cpp
+++ b/llvm/lib/ProfileData/SampleProfWriter.cpp
@@ -631,8 +631,27 @@ SampleProfileWriterExtBinaryBase::writeMD5ProfileSymbolListSection() {
   return sampleprof_error::success;
 }
 
+unsigned SampleProfileWriterExtBinaryBase::findUnwrittenEntry(SecType Type) {
+  auto WrittenIndices =
+      llvm::map_range(SecHdrTable, &SecHdrTableEntry::LayoutIndex);
+  for (auto [I, Entry] : llvm::enumerate(SectionHdrLayout))
+    if (Entry.Type == Type && !llvm::is_contained(WrittenIndices, I))
+      return I;
+  llvm_unreachable("Matching section not found in SectionHdrLayout");
+}
+
 std::error_code SampleProfileWriterExtBinaryBase::writeOneSection(
-    SecType Type, uint32_t LayoutIdx, const SampleProfileMap &ProfileMap) {
+    SecType Type, const SampleProfileMap &ProfileMap) {
+  unsigned LayoutIdx = findUnwrittenEntry(Type);
+  SecHdrTableEntry &Entry = SectionHdrLayout[LayoutIdx];
+
+  // In the context-split layout, the second instance of a repeated section type
+  // (such as LBR profile and function offset table) contains flat profiles.
+  if (SecLayout == CtxSplitLayout &&
+      llvm::count_if(SecHdrTable,
+                     [&](const auto &E) { return E.Type == Type; }) == 1)
+    addSecFlag(Entry, SecCommonFlags::SecFlagFlat);
+
   // The setting of SecFlagCompress should happen before markSectionStart.
   if (Type == SecProfileSymbolList && ProfSymList && ProfSymList->toCompress())
     setToCompressSection(SecProfileSymbolList);
@@ -676,8 +695,7 @@ std::error_code SampleProfileWriterExtBinaryBase::writeOneSection(
       return EC;
     break;
   case SecFuncOffsetTable: {
-    bool IsFlat =
-        hasSecFlag(SectionHdrLayout[LayoutIdx], SecCommonFlags::SecFlagFlat);
+    bool IsFlat = hasSecFlag(Entry, SecCommonFlags::SecFlagFlat);
     // An unflagged function offset table inherently indexes the primary
     // Nested symbol span.
     bool IsNested = !IsFlat;
@@ -711,24 +729,13 @@ SampleProfileWriterExtBinary::SampleProfileWriterExtBinary(
 
 std::error_code SampleProfileWriterExtBinary::writeDefaultLayout(
     const SampleProfileMap &ProfileMap) {
-  // The const indices passed to writeOneSection below are specifying the
-  // positions of the sections in SectionHdrLayout. Look at
-  // initSectionHdrLayout to find out where each section is located in
-  // SectionHdrLayout.
-  if (auto EC = writeOneSection(SecProfSummary, 0, ProfileMap))
-    return EC;
-  if (auto EC = writeOneSection(SecNameTable, 1, ProfileMap))
-    return EC;
-  if (auto EC = writeOneSection(SecCSNameTable, 2, ProfileMap))
-    return EC;
-  if (auto EC = writeOneSection(SecLBRProfile, 4, ProfileMap))
-    return EC;
-  if (auto EC = writeOneSection(SecProfileSymbolList, 5, ProfileMap))
-    return EC;
-  if (auto EC = writeOneSection(SecFuncOffsetTable, 3, ProfileMap))
-    return EC;
-  if (auto EC = writeOneSection(SecFuncMetadata, 6, ProfileMap))
-    return EC;
+  static constexpr SecType Sections[] = {
+      SecProfSummary,       SecNameTable,       SecCSNameTable,  SecLBRProfile,
+      SecProfileSymbolList, SecFuncOffsetTable, SecFuncMetadata,
+  };
+  for (SecType Type : Sections)
+    if (std::error_code EC = writeOneSection(Type, ProfileMap))
+      return EC;
   return sampleprof_error::success;
 }
 
@@ -748,28 +755,19 @@ std::error_code SampleProfileWriterExtBinary::writeCtxSplitLayout(
   SampleProfileMap NestedProfileMap, FlatProfileMap;
   splitProfileMapToTwo(ProfileMap, NestedProfileMap, FlatProfileMap);
 
-  if (auto EC = writeOneSection(SecProfSummary, 0, ProfileMap))
-    return EC;
-  if (auto EC = writeOneSection(SecNameTable, 1, ProfileMap))
-    return EC;
-  if (auto EC = writeOneSection(SecLBRProfile, 3, NestedProfileMap))
-    return EC;
-  if (auto EC = writeOneSection(SecFuncOffsetTable, 2, NestedProfileMap))
-    return EC;
-  // Mark the section as flat (without callsite samples). Note section flag
-  // needs to be set before writing the section.
-  addSectionFlag(5, SecCommonFlags::SecFlagFlat);
-  if (auto EC = writeOneSection(SecLBRProfile, 5, FlatProfileMap))
-    return EC;
-  // Mark the section as flat (without callsite samples). Note section flag
-  // needs to be set before writing the section.
-  addSectionFlag(4, SecCommonFlags::SecFlagFlat);
-  if (auto EC = writeOneSection(SecFuncOffsetTable, 4, FlatProfileMap))
-    return EC;
-  if (auto EC = writeOneSection(SecProfileSymbolList, 6, ProfileMap))
-    return EC;
-  if (auto EC = writeOneSection(SecFuncMetadata, 7, ProfileMap))
-    return EC;
+  const std::pair<SecType, const SampleProfileMap &> Sections[] = {
+      {SecProfSummary, ProfileMap},
+      {SecNameTable, ProfileMap},
+      {SecLBRProfile, NestedProfileMap},
+      {SecFuncOffsetTable, NestedProfileMap},
+      {SecLBRProfile, FlatProfileMap},
+      {SecFuncOffsetTable, FlatProfileMap},
+      {SecProfileSymbolList, ProfileMap},
+      {SecFuncMetadata, ProfileMap},
+  };
+  for (const auto &[Type, Map] : Sections)
+    if (std::error_code EC = writeOneSection(Type, Map))
+      return EC;
 
   return sampleprof_error::success;
 }

``````````

</details>


https://github.com/llvm/llvm-project/pull/215981