[NFC][SampleFDO] Move some common stuff from SampleProfileReaderExtBinary/SampleProfileWriterExtBinary to their parent classes.

[wmi]

Summary

SampleProfileReaderExtBinary/SampleProfileWriterExtBinary specify the typical section layout currently used by SampleFDO. Currently a lot of section readers/writers stay in the two classes. However, as we expect to have more types of SampleFDO profiles, we hope those new types of profiles can share the common sections while configuring their own sections easily with minimal change. That is why I move some common stuff from SampleProfileReaderExtBinary/SampleProfileWriterExtBinary to SampleProfileReaderExtBinaryBase/SampleProfileWriterExtBinaryBase so new profiles class inheriting from the base class can reuse them.

[hoy]

llvm/lib/ProfileData/SampleProfWriter.cpp

232

Why is this separated from writeOneSection?

247

Thanks for the refactoring which looks much cleaner.

254

Can this be moved into writeProfileSymbolListSection?

[wenlei]

Thanks for the refactoring. I'm wondering what is most flexible way for us to take advantage of extended binary format. I thought being able to mix and use different section is really flexible, but if we tie non-common section to specific new format/sub class, it may restrict the way we use this format. E.g. if ExtFormatA has SecProfileSymbolList, and ExtFormatB has something like SecFuncMetadata (which is what have for CSSPGO internally), then consider if there's a need for a profile to have both SecProfileSymbolList and SecFuncMetadata - implementing that through inheritance could be a bit cumbersome.

I'm thinking about two alternatives:

1. Follow similar model like ELF, adding or removing a section does not change the format "type". With that, we can pass in SectionHdrLayout as parameter to ctor of reader and writer, and use the layout to drive the reading and writing. In this case, we don't need to use inheritance and all new section and its reader/writer functions are always registered in a single class.
2. Still use inheritance to allow new combination of sections as a new sub format, but always register all new section and its reader/writer functions in the base class. The subclasses only define unique SectionHdrLayout. With this, the example above will not need to inherit from both ExtFormatA and ExtFormatB.

It's nice to keep the ability to use sections orthogonally, instead of in hierarchical manner. What do you think?

[wmi]

In D89524#2337297, @wenlei wrote:

Thanks for the refactoring. I'm wondering what is most flexible way for us to take advantage of extended binary format. I thought being able to mix and use different section is really flexible, but if we tie non-common section to specific new format/sub class, it may restrict the way we use this format. E.g. if ExtFormatA has SecProfileSymbolList, and ExtFormatB has something like SecFuncMetadata (which is what have for CSSPGO internally), then consider if there's a need for a profile to have both SecProfileSymbolList and SecFuncMetadata - implementing that through inheritance could be a bit cumbersome.

I'm thinking about two alternatives:

1. Follow similar model like ELF, adding or removing a section does not change the format "type". With that, we can pass in SectionHdrLayout as parameter to ctor of reader and writer, and use the layout to drive the reading and writing. In this case, we don't need to use inheritance and all new section and its reader/writer functions are always registered in a single class.
2. Still use inheritance to allow new combination of sections as a new sub format, but always register all new section and its reader/writer functions in the base class. The subclasses only define unique SectionHdrLayout. With this, the example above will not need to inherit from both ExtFormatA and ExtFormatB.

It's nice to keep the ability to use sections orthogonally, instead of in hierarchical manner. What do you think?

Good point. Thanks to raise it up. I considered those two alternatives and the current solution. They all have their own advantages and disadvantages.

For 1, it is flexible. Any Reader/Writer users can define their layout they want to have without having to define a new subclass. The disadvantage is the layout reflects the internal logic about how the sections will be implemented and used. Passing in SectionHdrLayout means to expose many details to the code using the reader/writer. Another disadvantage is that putting all the section reader/writer in a single class may be cumbersome.

For 2, it is also very flexible. The disadvantage is putting all the section read/writers in a class.

The current patch is an extension of 2. We still have inheritance and we register the section reader/writers at a proper level that all the subclasses can get/share the support they want (Each class can register their section reader/writer in readOtherSection/writeOtherSection). We may put a section in a subclass at the beginning but move the section up when more formats are added.

A common disadvantage of 2 and 3 are every time we need to define a new subclass, there is some routinary code we need to add.

So current patch is a compromise of the above consideration. I like to hear your opinions since CSSPGO will have some immediate requirement to define some new sections.

[wmi]

llvm/lib/ProfileData/SampleProfWriter.cpp

232

That is somewhat explained in the reply to Wenlei's mail. I want to put a section to an inheritance level when foreseeably new subclasses created in other inheritance branch will not use it. Surely whether SecProfileSymbolList will not be used by other new formats are questionable. I just use it as an example and it will be easy to move the section up.

254

It cannot. Whether the section will be compressed or not has to be decided before markSectionStart is called.

[hoy]

llvm/lib/ProfileData/SampleProfWriter.cpp

232

I see. Can I ask if you are going to introduce a new ExtBinary variant? Since we are also adding a new section (for pseudo probe) , I'm not sure if it should be treated as an optional section of ExtBinaryBase or a required section of a new variant.

[wmi]

llvm/lib/ProfileData/SampleProfWriter.cpp

232

Yes, I am going to add a new ExtBinary variant.

To introduce it a little bit, we have a synthesized profile which is merged from many sample profiles and we use it for targets without regular sample profile. The merged profile is quite large. Previously inline instances in the synthesized profile are all flattened, i.e., all inline instances in the profile are extracted and merged to their outline instance. We found selectively flattening instead of fully flattening the synthesized profile is helpful for performance, however it may increase compile time significantly in the postlink phase during profile loading. To solve that, we are trying to split the profile into two parts, one part containing inline instance and the other part not containing inline instance. When thinlto is used, the second profile loading pass in postlink phase will only have to load the part containing inline instance, which will minimize the compile time impact.

For the new variant we are going to propose, it doesn't need SecProfileSymbolList for the moment because the synthesized profile doesn't need SecProfileSymbolList. But we find in some initial experiment that we can selectively flatten regular sampleFDO profile without performance impact but with much smaller profile size, so the new variant could be possibly useful for regular sampleFDO too in the future. In that case, SecProfileSymbolList will be needed.

With that being said, maybe SecProfileSymbolList is not a good example to be put in a subclass, we can move SecProfileSymbolList to the Base class if you agree.

[hoy]

llvm/lib/ProfileData/SampleProfWriter.cpp

232

Thanks for providing more context which is helpful. I'm not particularly concerned with where to place SecProfileSymbolList. I'm still trying to understand class hierarchy. Are we separating the variants so that a section ends up being required but not optional in every variant? If a section can be optional, do we still need the inheritance?

[wmi]

llvm/lib/ProfileData/SampleProfWriter.cpp

232

Currently if a section is registered in SectionHdrLayout of a class, it means the variant will have a place for the section in section table in the file header, but if you don't need it, just don't write anything to the section and it will be an empty section. That make section support essentially optional.

I tend to keep inheritance is because if we just use one class, we end up putting the requirement of every variant in the same class and I am worried that may be too much for the single class as the formats evolve.

Surely we also don't want to have too many classes either. So maybe we can support multiple layouts for a certain subclass if the layouts are similar. Only if the new variant is significantly different from existing one, we create a new class. What do you think?

[wenlei]

Thanks for the context! So the intention is to use subclass to hide extra complexity of uncommon sections/variants and its reader/writer into subclass. Then this refactoring makes sense.

Just to confirm, this use case you described is the partial profile that was introduced a while ago, correct? And new variant will have new sections for flattened and hierarchical profile?

For SecProfileSymbolList specifically, I think it's common enough that justifies it to be in the base.

Surely we also don't want to have too many classes either. So maybe we can support multiple layouts for a certain subclass if the layouts are similar. Only if the new variant is significantly different from existing one, we create a new class. What do you think?

Sounds good to me.

[hoy]

llvm/lib/ProfileData/SampleProfWriter.cpp

232

I see your point now. Yes, it's reasonable to not mix custom section implementation with the common base class.

Should the base implementation SampleProfileWriterExtBinaryBase::writeOtherSection be also called here since a subclass inherits all sections from the base class?

Nit: consider naming it writeCustomSection if that sounds good to you?

[wmi]

In D89524#2343611, @wenlei wrote:

Thanks for the context! So the intention is to use subclass to hide extra complexity of uncommon sections/variants and its reader/writer into subclass. Then this refactoring makes sense.

Just to confirm, this use case you described is the partial profile that was introduced a while ago, correct? And new variant will have new sections for flattened and hierarchical profile?

Yes, that is the partial profile. Previously the partial profile have all flattened profile in SecLBRProfile section. The new variant will try to split the flattened and hierarchical profile into two sections, and so as the SecNameTable and SecFuncOffsetTable sections.

llvm/lib/ProfileData/SampleProfWriter.cpp

232

SampleProfileWriterExtBinaryBase is a pure virtual class so I expect it will not have customized sections. I think I should change writeOtherSection to a pure virtual function instead of leaving it as an empty function.

writeCustomSection is a better name. Thanks!

[hoy]

llvm/lib/ProfileData/SampleProfWriter.cpp

232

Yeah, a pure virtual function sounds good.

[wmi]

Address Hongtao and Wenlei's comments.

[wenlei]

LGTM. Thanks!

llvm/include/llvm/ProfileData/SampleProfReader.h

std::error_code decompressSection(const uint8_t *SecStart,
                                     const uint64_t SecSize,
                                     const uint8_t *&DecompressBuf,
                                     uint64_t &DecompressBufSize);
 
   BumpPtrAllocator Allocator;
 
 protected:
   std::vector<SecHdrTableEntry> SecHdrTable;
-  std::unique_ptr<ProfileSymbolList> ProfSymList;
   std::error_code readSecHdrTableEntry();
   std::error_code readSecHdrTable();
-  virtual std::error_code readHeader() override;
-  virtual std::error_code verifySPMagic(uint64_t Magic) override = 0;
-  virtual std::error_code readOneSection(const uint8_t *Start, uint64_t Size,
-                                         const SecHdrTableEntry &Entry) = 0;
-
-public:
-  SampleProfileReaderExtBinaryBase(std::unique_ptr<MemoryBuffer> B,
-                                   LLVMContext &C, SampleProfileFormat Format)
-      : SampleProfileReaderBinary(std::move(B), C, Format) {}
 
-  /// Read sample profiles in extensible format from the associated file.
-  std::error_code readImpl() override;
-
-  /// Get the total size of all \p Type sections.
-  uint64_t getSectionSize(SecType Type);
-  /// Get the total size of header and all sections.
-  uint64_t getFileSize();
-  virtual bool dumpSectionInfo(raw_ostream &OS = dbgs()) override;
-};
-
-class SampleProfileReaderExtBinary : public SampleProfileReaderExtBinaryBase {
-private:
-  virtual std::error_code verifySPMagic(uint64_t Magic) override;
-  virtual std::error_code
-  readOneSection(const uint8_t *Start, uint64_t Size,
-                 const SecHdrTableEntry &Entry) override;
-  std::error_code readProfileSymbolList();
   std::error_code readFuncOffsetTable();
   std::error_code readFuncProfiles();
   std::error_code readMD5NameTable();
   std::error_code readNameTableSec(bool IsMD5);
+  std::error_code readProfileSymbolList();
+
+  virtual std::error_code readHeader() override;
+  virtual std::error_code verifySPMagic(uint64_t Magic) override = 0;
+  virtual std::error_code readOneSection(const uint8_t *Start, uint64_t Size,
+                                         const SecHdrTableEntry &Entry);
+  // placeholder for subclasses to dispatch their own section readers.
+  virtual std::error_code readCustomSection(const SecHdrTableEntry &Entry) = 0;
+
+  std::unique_ptr<ProfileSymbolList> ProfSymList;
 
   /// The table mapping from function name to the offset of its FunctionSample
   /// towards file start.
   DenseMap<StringRef, uint64_t> FuncOffsetTable;
   /// The set containing the functions to use when compiling a module.
   DenseSet<StringRef> FuncsToUse;
   /// Use all functions from the input profile.
   bool UseAllFuncs = true;
 
   /// If MD5 is used in NameTable section, the section saves uint64_t data.
   /// The uint64_t data has to be converted to a string and then the string
   /// will be used to initialize StringRef in NameTable.
   /// Note NameTable contains StringRef so it needs another buffer to own
   /// the string data. MD5StringBuf serves as the string buffer that is
   /// referenced by NameTable (vector of StringRef). We make sure
   /// the lifetime of MD5StringBuf is not shorter than that of NameTable.
   std::unique_ptr<std::vector<std::string>> MD5StringBuf;
 
 public:
-  SampleProfileReaderExtBinary(std::unique_ptr<MemoryBuffer> B, LLVMContext &C,
-                               SampleProfileFormat Format = SPF_Ext_Binary)
-      : SampleProfileReaderExtBinaryBase(std::move(B), C, Format) {}
+  SampleProfileReaderExtBinaryBase(std::unique_ptr<MemoryBuffer> B,
+                                   LLVMContext &C, SampleProfileFormat Format)
+      : SampleProfileReaderBinary(std::move(B), C, Format) {}
 
-  /// \brief Return true if \p Buffer is in the format supported by this class.
-  static bool hasFormat(const MemoryBuffer &Buffer);
+  /// Read sample profiles in extensible format from the associated file.
+  std::error_code readImpl() override;
 
-  virtual std::unique_ptr<ProfileSymbolList> getProfileSymbolList() override {
-    return std::move(ProfSymList);
-  };
+  /// Get the total size of all \p Type sections.
+  uint64_t getSectionSize(SecType Type);
+  /// Get the total size of header and all sections.
+  uint64_t getFileSize();
+  virtual bool dumpSectionInfo(raw_ostream &OS = dbgs()) override;
 
   /// Collect functions with definitions in Module \p M.
   void collectFuncsFrom(const Module &M) override;
 
   /// Return whether names in the profile are all MD5 numbers.
   virtual bool useMD5() override {
     assert(!NameTable.empty() && "NameTable should have been initialized");
     return MD5StringBuf && !MD5StringBuf->empty();
   }
+
+  virtual std::unique_ptr<ProfileSymbolList> getProfileSymbolList() override {
+    return std::move(ProfSymList);
+  };
+};
+
+class SampleProfileReaderExtBinary : public SampleProfileReaderExtBinaryBase {
+private:
+  virtual std::error_code verifySPMagic(uint64_t Magic) override;
+  virtual std::error_code
+  readCustomSection(const SecHdrTableEntry &Entry) override {
+    return sampleprof_error::success;
+  };
+
+public:
+  SampleProfileReaderExtBinary(std::unique_ptr<MemoryBuffer> B, LLVMContext &C,
+                               SampleProfileFormat Format = SPF_Ext_Binary)
+      : SampleProfileReaderExtBinaryBase(std::move(B), C, Format) {}
+
+  /// \brief Return true if \p Buffer is in the format supported by this class.
+  static bool hasFormat(const MemoryBuffer &Buffer);
 };
 
 class SampleProfileReaderCompactBinary : public SampleProfileReaderBinary {
 private:
   /// Function name table.
   std::vector<std::string> NameTable;
   /// The table mapping from function name to the offset of its FunctionSample
   /// towards file start.

llvm/include/llvm/ProfileData/SampleProfWriter.h

 class SampleProfileWriterExtBinaryBase : public SampleProfileWriterBinary {
   using SampleProfileWriterBinary::SampleProfileWriterBinary;
 public:
   virtual std::error_code
   write(const StringMap<FunctionSamples> &ProfileMap) override;
 
   virtual void setToCompressAllSections() override;
   void setToCompressSection(SecType Type);
+  virtual std::error_code writeSample(const FunctionSamples &S) override;
+
+  // Set to use MD5 to represent string in NameTable.
+  virtual void setUseMD5() override {
+    UseMD5 = true;
+    addSectionFlag(SecNameTable, SecNameTableFlags::SecFlagMD5Name);
+  }
+
+  // Set the profile to be partial. It means the profile is for
+  // common/shared code. The common profile is usually merged from
+  // profiles collected from running other targets.
+  virtual void setPartialProfile() override {
+    addSectionFlag(SecProfSummary, SecProfSummaryFlags::SecFlagPartial);
+  }
+
+  virtual void setProfileSymbolList(ProfileSymbolList *PSL) override {
+    ProfSymList = PSL;
+  };
 
 protected:
   uint64_t markSectionStart(SecType Type);
   std::error_code addNewSection(SecType Sec, uint64_t SectionStart);
   template <class SecFlagType>
   void addSectionFlag(SecType Type, SecFlagType Flag) {
     for (auto &Entry : SectionHdrLayout) {
       if (Entry.Type == Type)
         addSecFlag(Entry, Flag);
     }
   }
 
+  // placeholder for subclasses to dispatch their own section writers.
+  virtual std::error_code writeCustomSection(SecType Type) = 0;
+
   virtual void initSectionHdrLayout() = 0;
+  // specify the order to write sections.
   virtual std::error_code
   writeSections(const StringMap<FunctionSamples> &ProfileMap) = 0;
 
+  // Dispatch section writer for each section.
+  virtual std::error_code
+  writeOneSection(SecType Type, const StringMap<FunctionSamples> &ProfileMap);
+
+  // Helper function to write name table.
+  virtual std::error_code writeNameTable() override;
+
+  // Functions to write various kinds of sections.
+  std::error_code
+  writeNameTableSection(const StringMap<FunctionSamples> &ProfileMap);
+  std::error_code writeFuncOffsetTable();
+  std::error_code writeProfileSymbolListSection();
+
   // Specifiy the order of sections in section header table. Note
   // the order of sections in the profile may be different that the
   // order in SectionHdrLayout. sample Reader will follow the order
   // in SectionHdrLayout to read each section.
   SmallVector<SecHdrTableEntry, 8> SectionHdrLayout;
 
+  // Save the start of SecLBRProfile so we can compute the offset to the
+  // start of SecLBRProfile for each Function's Profile and will keep it
+  // in FuncOffsetTable.
+  uint64_t SecLBRProfileStart = 0;
+
 private:
   void allocSecHdrTable();
   std::error_code writeSecHdrTable();
   virtual std::error_code
   writeHeader(const StringMap<FunctionSamples> &ProfileMap) override;
   SecHdrTableEntry &getEntryInLayout(SecType Type);
   std::error_code compressAndOutput();
 
   // We will swap the raw_ostream held by LocalBufStream and that
   // held by OutputStream if we try to add a section which needs
   // compression. After the swap, all the data written to output
   // will be temporarily buffered into the underlying raw_string_ostream
   // originally held by LocalBufStream. After the data writing for the
   // section is completed, compress the data in the local buffer,
   // swap the raw_ostream back and write the compressed data to the
   // real output.
   std::unique_ptr<raw_ostream> LocalBufStream;
   // The location where the output stream starts.
   uint64_t FileStart;
   // The location in the output stream where the SecHdrTable should be
   // written to.
   uint64_t SecHdrTableOffset;
   // Initial Section Flags setting.
   std::vector<SecHdrTableEntry> SecHdrTable;
+
+  // FuncOffsetTable maps function name to its profile offset in SecLBRProfile
+  // section. It is used to load function profile on demand.
+  MapVector<StringRef, uint64_t> FuncOffsetTable;
+  // Whether to use MD5 to represent string.
+  bool UseMD5 = false;
+
+  ProfileSymbolList *ProfSymList = nullptr;
 };
 
 class SampleProfileWriterExtBinary : public SampleProfileWriterExtBinaryBase {
 public:
   SampleProfileWriterExtBinary(std::unique_ptr<raw_ostream> &OS)
       : SampleProfileWriterExtBinaryBase(OS) {
     initSectionHdrLayout();
   }
 
-  virtual std::error_code writeSample(const FunctionSamples &S) override;
-  virtual void setProfileSymbolList(ProfileSymbolList *PSL) override {
-    ProfSymList = PSL;
-  };
-
-  // Set to use MD5 to represent string in NameTable.
-  virtual void setUseMD5() override {
-    UseMD5 = true;
-    addSectionFlag(SecNameTable, SecNameTableFlags::SecFlagMD5Name);
-  }
-
-  // Set the profile to be partial. It means the profile is for
-  // common/shared code. The common profile is usually merged from
-  // profiles collected from running other targets.
-  virtual void setPartialProfile() override {
-    addSectionFlag(SecProfSummary, SecProfSummaryFlags::SecFlagPartial);
-  }
-
 private:
   virtual void initSectionHdrLayout() override {
     // Note that SecFuncOffsetTable section is written after SecLBRProfile
     // in the profile, but is put before SecLBRProfile in SectionHdrLayout.
     //
     // This is because sample reader follows the order of SectionHdrLayout to
     // read each section, to read function profiles on demand sample reader
     // need to get the offset of each function profile first.
     //
     // SecFuncOffsetTable section is written after SecLBRProfile in the
     // profile because FuncOffsetTable needs to be populated while section
     // SecLBRProfile is written.
     SectionHdrLayout = {{SecProfSummary, 0, 0, 0},
                         {SecNameTable, 0, 0, 0},
                         {SecFuncOffsetTable, 0, 0, 0},
                         {SecLBRProfile, 0, 0, 0},
                         {SecProfileSymbolList, 0, 0, 0}};
   };
   virtual std::error_code
   writeSections(const StringMap<FunctionSamples> &ProfileMap) override;
 
-  std::error_code writeFuncOffsetTable();
-  virtual std::error_code writeNameTable() override;
-
-  ProfileSymbolList *ProfSymList = nullptr;
-
-  // Save the start of SecLBRProfile so we can compute the offset to the
-  // start of SecLBRProfile for each Function's Profile and will keep it
-  // in FuncOffsetTable.
-  uint64_t SecLBRProfileStart = 0;
-  // FuncOffsetTable maps function name to its profile offset in SecLBRProfile
-  // section. It is used to load function profile on demand.
-  MapVector<StringRef, uint64_t> FuncOffsetTable;
-  // Whether to use MD5 to represent string.
-  bool UseMD5 = false;
+  virtual std::error_code writeCustomSection(SecType Type) override {
+    return sampleprof_error::success;
+  };
 };
 
 // CompactBinary is a compact format of binary profile which both reduces
 // the profile size and the load time needed when compiling. It has two
 // major difference with Binary format.
 // 1. It represents all the strings in name table using md5 hash.
 // 2. It saves a function offset table which maps function name index to
 // the offset of its function profile to the start of the binary profile,

llvm/lib/ProfileData/SampleProfReader.cpp

std::error_code SampleProfileReaderBinary::readImpl() {
   while (!at_eof()) {
     if (std::error_code EC = readFuncProfile(Data))
       return EC;
   }
 
   return sampleprof_error::success;
 }
 
-std::error_code SampleProfileReaderExtBinary::readOneSection(
+std::error_code SampleProfileReaderExtBinaryBase::readOneSection(
     const uint8_t *Start, uint64_t Size, const SecHdrTableEntry &Entry) {
   Data = Start;
   End = Start + Size;
   switch (Entry.Type) {
   case SecProfSummary:
     if (std::error_code EC = readSummary())
       return EC;
     if (hasSecFlag(Entry, SecProfSummaryFlags::SecFlagPartial))
       Summary->setPartialProfile(true);
     break;
   case SecNameTable:
     if (std::error_code EC = readNameTableSec(
             hasSecFlag(Entry, SecNameTableFlags::SecFlagMD5Name)))
       return EC;
     break;
   case SecLBRProfile:
     if (std::error_code EC = readFuncProfiles())
       return EC;
     break;
-  case SecProfileSymbolList:
-    if (std::error_code EC = readProfileSymbolList())
-      return EC;
-    break;
   case SecFuncOffsetTable:
     if (std::error_code EC = readFuncOffsetTable())
       return EC;
     break;
+  case SecProfileSymbolList:
+    if (std::error_code EC = readProfileSymbolList())
+      return EC;
+    break;
   default:
+    if (std::error_code EC = readCustomSection(Entry))
+      return EC;
     break;
   }
   return sampleprof_error::success;
 }
 
-void SampleProfileReaderExtBinary::collectFuncsFrom(const Module &M) {
+void SampleProfileReaderExtBinaryBase::collectFuncsFrom(const Module &M) {
   UseAllFuncs = false;
   FuncsToUse.clear();
   for (auto &F : M)
     FuncsToUse.insert(FunctionSamples::getCanonicalFnName(F));
 }
 
-std::error_code SampleProfileReaderExtBinary::readFuncOffsetTable() {
+std::error_code SampleProfileReaderExtBinaryBase::readFuncOffsetTable() {
   auto Size = readNumber<uint64_t>();
   if (std::error_code EC = Size.getError())
     return EC;
 
   FuncOffsetTable.reserve(*Size);
   for (uint32_t I = 0; I < *Size; ++I) {
     auto FName(readStringFromTable());
     if (std::error_code EC = FName.getError())
       return EC;
 
     auto Offset = readNumber<uint64_t>();
     if (std::error_code EC = Offset.getError())
       return EC;
 
     FuncOffsetTable[*FName] = *Offset;
   }
   return sampleprof_error::success;
 }
 
-std::error_code SampleProfileReaderExtBinary::readFuncProfiles() {
+std::error_code SampleProfileReaderExtBinaryBase::readFuncProfiles() {
   const uint8_t *Start = Data;
   if (UseAllFuncs) {
     while (Data < End) {
       if (std::error_code EC = readFuncProfile(Data))
         return EC;
     }
     assert(Data == End && "More data is read than expected");
     return sampleprof_error::success;
for (auto NameOffset : FuncOffsetTable) {
         return EC;
     }
   }
 
   Data = End;
   return sampleprof_error::success;
 }
 
-std::error_code SampleProfileReaderExtBinary::readProfileSymbolList() {
+std::error_code SampleProfileReaderExtBinaryBase::readProfileSymbolList() {
   if (!ProfSymList)
     ProfSymList = std::make_unique<ProfileSymbolList>();
 
   if (std::error_code EC = ProfSymList->read(Data, End - Data))
     return EC;
 
   Data = End;
   return sampleprof_error::success;
for (uint32_t I = 0; I < *Size; ++I) {
     if (std::error_code EC = Name.getError())
       return EC;
     NameTable.push_back(*Name);
   }
 
   return sampleprof_error::success;
 }
 
-std::error_code SampleProfileReaderExtBinary::readMD5NameTable() {
+std::error_code SampleProfileReaderExtBinaryBase::readMD5NameTable() {
   auto Size = readNumber<uint64_t>();
   if (std::error_code EC = Size.getError())
     return EC;
   NameTable.reserve(*Size);
   MD5StringBuf = std::make_unique<std::vector<std::string>>();
   MD5StringBuf->reserve(*Size);
   for (uint32_t I = 0; I < *Size; ++I) {
     auto FID = readNumber<uint64_t>();
     if (std::error_code EC = FID.getError())
       return EC;
     MD5StringBuf->push_back(std::to_string(*FID));
     // NameTable is a vector of StringRef. Here it is pushing back a
     // StringRef initialized with the last string in MD5stringBuf.
     NameTable.push_back(MD5StringBuf->back());
   }
   return sampleprof_error::success;
 }
 
-std::error_code SampleProfileReaderExtBinary::readNameTableSec(bool IsMD5) {
+std::error_code SampleProfileReaderExtBinaryBase::readNameTableSec(bool IsMD5) {
   if (IsMD5)
     return readMD5NameTable();
   return SampleProfileReaderBinary::readNameTable();
 }
 
 std::error_code SampleProfileReaderCompactBinary::readNameTable() {
   auto Size = readNumber<uint64_t>();
   if (std::error_code EC = Size.getError())

llvm/lib/ProfileData/SampleProfWriter.cpp

std::error_code SampleProfileWriterExtBinaryBase::write(
 
   if (std::error_code EC = writeSecHdrTable())
     return EC;
 
   return sampleprof_error::success;
 }
 
 std::error_code
-SampleProfileWriterExtBinary::writeSample(const FunctionSamples &S) {
+SampleProfileWriterExtBinaryBase::writeSample(const FunctionSamples &S) {
   uint64_t Offset = OutputStream->tell();
   StringRef Name = S.getName();
   FuncOffsetTable[Name] = Offset - SecLBRProfileStart;
   encodeULEB128(S.getHeadSamples(), *OutputStream);
   return writeBody(S);
 }
 
-std::error_code SampleProfileWriterExtBinary::writeFuncOffsetTable() {
+std::error_code SampleProfileWriterExtBinaryBase::writeFuncOffsetTable() {
   auto &OS = *OutputStream;
 
   // Write out the table size.
   encodeULEB128(FuncOffsetTable.size(), OS);
 
   // Write out FuncOffsetTable.
   for (auto entry : FuncOffsetTable) {
     writeNameIdx(entry.first);
     encodeULEB128(entry.second, OS);
   }
   return sampleprof_error::success;
 }
 
-std::error_code SampleProfileWriterExtBinary::writeNameTable() {
+std::error_code SampleProfileWriterExtBinaryBase::writeNameTable() {
   if (!UseMD5)
     return SampleProfileWriterBinary::writeNameTable();
 
   auto &OS = *OutputStream;
   std::set<StringRef> V;
   stablizeNameTable(V);
 
   // Write out the name table.
   encodeULEB128(NameTable.size(), OS);
   for (auto N : V) {
     encodeULEB128(MD5Hash(N), OS);
   }
   return sampleprof_error::success;
 }
 
-std::error_code SampleProfileWriterExtBinary::writeSections(
+std::error_code SampleProfileWriterExtBinaryBase::writeNameTableSection(
     const StringMap<FunctionSamples> &ProfileMap) {
-  uint64_t SectionStart = markSectionStart(SecProfSummary);
-  computeSummary(ProfileMap);
-  if (auto EC = writeSummary())
-    return EC;
-  if (std::error_code EC = addNewSection(SecProfSummary, SectionStart))
-    return EC;
-
-  // Generate the name table for all the functions referenced in the profile.
-  SectionStart = markSectionStart(SecNameTable);
   for (const auto &I : ProfileMap) {
     addName(I.first());
     addNames(I.second);
   }
-  writeNameTable();
-  if (std::error_code EC = addNewSection(SecNameTable, SectionStart))
+  if (auto EC = writeNameTable())
     return EC;
+  return sampleprof_error::success;
+}
 
-  SectionStart = markSectionStart(SecLBRProfile);
-  SecLBRProfileStart = OutputStream->tell();
-  if (std::error_code EC = writeFuncProfiles(ProfileMap))
-    return EC;
-  if (std::error_code EC = addNewSection(SecLBRProfile, SectionStart))
+std::error_code
+SampleProfileWriterExtBinaryBase::writeProfileSymbolListSection() {
+  if (ProfSymList && ProfSymList->size() > 0)
+    if (std::error_code EC = ProfSymList->write(*OutputStream))
       return EC;
 
-  if (ProfSymList && ProfSymList->toCompress())
+  return sampleprof_error::success;
+}
+
+std::error_code SampleProfileWriterExtBinaryBase::writeOneSection(
+    SecType Type, const StringMap<FunctionSamples> &ProfileMap) {
+  // The setting of SecFlagCompress should happen before markSectionStart.
+  if (Type == SecProfileSymbolList && ProfSymList && ProfSymList->toCompress())
     setToCompressSection(SecProfileSymbolList);
 
-  SectionStart = markSectionStart(SecProfileSymbolList);
-  if (ProfSymList && ProfSymList->size() > 0)
-    if (std::error_code EC = ProfSymList->write(*OutputStream))
+  uint64_t SectionStart = markSectionStart(Type);
+  switch (Type) {
+  case SecProfSummary:
+    computeSummary(ProfileMap);
+    if (auto EC = writeSummary())
       return EC;
-  if (std::error_code EC = addNewSection(SecProfileSymbolList, SectionStart))
+    break;
+  case SecNameTable:
+    if (auto EC = writeNameTableSection(ProfileMap))
       return EC;
-
-  SectionStart = markSectionStart(SecFuncOffsetTable);
-  if (std::error_code EC = writeFuncOffsetTable())
+    break;
+  case SecLBRProfile:
+    SecLBRProfileStart = OutputStream->tell();
+    if (std::error_code EC = writeFuncProfiles(ProfileMap))
+      return EC;
+    break;
+  case SecFuncOffsetTable:
+    if (auto EC = writeFuncOffsetTable())
+      return EC;
+    break;
+  case SecProfileSymbolList:
+    if (auto EC = writeProfileSymbolListSection())
+      return EC;
+    break;
+  default:
+    if (auto EC = writeCustomSection(Type))
       return EC;
-  if (std::error_code EC = addNewSection(SecFuncOffsetTable, SectionStart))
+    break;
+  }
+  if (std::error_code EC = addNewSection(Type, SectionStart))
     return EC;
+  return sampleprof_error::success;
+}
 
+std::error_code SampleProfileWriterExtBinary::writeSections(
+    const StringMap<FunctionSamples> &ProfileMap) {
+  if (auto EC = writeOneSection(SecProfSummary, ProfileMap))
+    return EC;
+  if (auto EC = writeOneSection(SecNameTable, ProfileMap))
+    return EC;
+  if (auto EC = writeOneSection(SecLBRProfile, ProfileMap))
+    return EC;
+  if (auto EC = writeOneSection(SecProfileSymbolList, ProfileMap))
+    return EC;
+  if (auto EC = writeOneSection(SecFuncOffsetTable, ProfileMap))
+    return EC;
   return sampleprof_error::success;
 }
 
 std::error_code SampleProfileWriterCompactBinary::write(
     const StringMap<FunctionSamples> &ProfileMap) {
   if (std::error_code EC = SampleProfileWriter::write(ProfileMap))
     return EC;
   if (std::error_code EC = writeFuncOffsetTable())