[SampleFDO] Add profile remapping support for profile on-demand loading used by ExtBinary format profile

[wmi]

Summary

profile on-demand loading was added for ExtBinary format profile in https://reviews.llvm.org/rL374233, but currently profile on-demand loading doesn't work well with profile remapping. The patch adds the support.

Suppose a function in the current module has outline instance in the profile. The function name in the module is different from the name of the outline instance, but remapper knows the two names are equal. When loading profile on-demand, the outline instance has to be loaded with remapper's help.

Before the patch, the steps to read the profile is as follows:

• create the profile reader
• profile reader read the profile.
• create the profile remapper
• remapper set the underlying reader.
• reset the profile reader to remapper.

With the patch, the steps to read the profile is changed to:

• create the profile reader
• create the profile remapper
• profile reader set the underlying remapper.
• profile reader read the profile.
• remapper set the underlying reader.
• reset the profile reader to remapper.

[wenlei]

llvm/include/llvm/ProfileData/SampleProfReader.h

723

while turning this Remappings object into unique_ptr decouples things a bit and enables moving error handling of reading remapping file earlier, there's now one more thing to check, Remappings could be null.

725

Check/assert clean state at the beginning? e.g. disallow setting/updating underlying reader repeatedly - was guaranteed when this was done in ctor.

llvm/lib/ProfileData/SampleProfReader.cpp

1212–1213

nit: assert(getFormat() != SPF_None && "UnderlyingReader must be setup before remapper reads")

llvm/lib/Transforms/IPO/SampleProfile.cpp

1700–1701

Wondering if we can move this up to line 1705, by forwarding SampleProfileReaderItaniumRemapper's read(), getProfileSymbolList(), etc. to the underlying reader.

If we can do that, we would be able to avoid the interleaved structure of setting up RemapReader, Reader, and let client of Reader to call read(), etc., then swap the two readers. Without that interleaved structure, we could hide this complexity inside SampleProfileReader::create by letting it take profile file name, as well as remapping file name, and return SampleProfileReaderItaniumRemapper (proxy) or a regular reader depending on whether remapping file is used.

[davidxl]

llvm/lib/Transforms/IPO/SampleProfile.cpp

1691

Better to move line 1705 together with block 1678 to 1693 and make block a helper function 'createRemapper' which does the following

1. create the remapper
2. set reader for the remapper
3. set remapper for the reader.

The call to createRemapper can be moved to line 1705.

[wmi]

Thanks Wenlei and David for the reviews.

It is a good suggestion to try to make remapper and reader less interleaving. That leads me to find a more fundamental problem. Current mechanism to make SampleProfileReaderItaniumRemapper a proxy of underlying reader is easy to have bugs. Currently, some internal data structures like Profiles and NameTables of underlyingReader have been transfered to SampleProfileReaderItaniumRemapper in order for SampleProfileReaderItaniumRemapper to function as a normal reader. However, for every virtual function in SampleProfileReader, we need to have an override function in SampleProfileReaderItaniumRemapper and delegate the function to the underlying reader, otherwise, it will call the wrong version in SampleProfileReader. It is confusing and it is easy to have bugs -- the underlying reader is ExtBinary format but SampleProfileReaderItaniumRemapper behaves like a SampleProfileReader.

So I plan to make SampleProfileReaderItaniumRemapper a helper inside SampleProfileReader instead of a proxy. From what I already saw that will simplify the code a lot. Will update the patch when it is done.

[wmi]

Address Wenlei and David's comments.
Change SampleProfileReaderItaniumRemapper from a proxy of SampleProfileReader to a helper object inside of SampleProfileReader, and it is used in two places: one in SampleProfileReader::getSamplesFor for looking up a function sample and one in SampleProfileReaderExtBinary::readFuncProfiles for loading profiles on demand.

[wenlei]

llvm/include/llvm/ProfileData/SampleProfReader.h

363–369

If client calls getSamplesFor before reading profile, RemappingApplied would be set without doing anything. Then even if client later calls read followed by getSamplesFor, remapping will never be applied. Probably not a big deal though as it's kind of corner case. But alternatively we could eliminate that possibility by calling applyRemapping right after reading profile, e.g. have read from base class call something like readImpl and applyRemapping, and different readers only override readImpl instead of read.

llvm/lib/ProfileData/SampleProfReader.cpp

1323

comment needs update for the new parameter (there're a few other places too).

[wmi]

Address Wenlei's comment.

llvm/include/llvm/ProfileData/SampleProfReader.h

363–369

Good point. Done.

llvm/lib/ProfileData/SampleProfReader.cpp

1323

Done.

[wenlei]

Thanks for making the changes, LGTM!

llvm/include/llvm/ProfileData/SampleProfReader.h

 #include <vector>
 
 namespace llvm {
 
 class raw_ostream;
 
 namespace sampleprof {
 
+class SampleProfileReader;
+
+/// SampleProfileReaderItaniumRemapper remaps the profile data from a
+/// sample profile data reader, by applying a provided set of equivalences
+/// between components of the symbol names in the profile.
+class SampleProfileReaderItaniumRemapper {
+public:
+  SampleProfileReaderItaniumRemapper(std::unique_ptr<MemoryBuffer> B,
+                                     std::unique_ptr<SymbolRemappingReader> SRR,
+                                     SampleProfileReader &R)
+      : Buffer(std::move(B)), Remappings(std::move(SRR)), Reader(R) {
+    assert(Remappings && "Remappings cannot be nullptr");
+  }
+
+  /// Create a remapper from the given remapping file. The remapper will
+  /// be used for profile read in by Reader.
+  static ErrorOr<std::unique_ptr<SampleProfileReaderItaniumRemapper>>
+  create(const std::string Filename, SampleProfileReader &Reader,
+         LLVMContext &C);
+
+  /// Create a remapper from the given Buffer. The remapper will
+  /// be used for profile read in by Reader.
+  static ErrorOr<std::unique_ptr<SampleProfileReaderItaniumRemapper>>
+  create(std::unique_ptr<MemoryBuffer> &B, SampleProfileReader &Reader,
+         LLVMContext &C);
+
+  /// Apply remappings to the profile read by Reader.
+  void applyRemapping(LLVMContext &Ctx);
+
+  bool hasApplied() { return RemappingApplied; }
+
+  /// Insert function name into remapper.
+  void insert(StringRef FunctionName) { Remappings->insert(FunctionName); }
+
+  /// Query whether there is equivalent in the remapper which has been
+  /// inserted.
+  bool exist(StringRef FunctionName) {
+    return Remappings->lookup(FunctionName);
+  }
+
+  /// Return the samples collected for function \p F if remapper knows
+  /// it is present in SampleMap.
+  FunctionSamples *getSamplesFor(StringRef FunctionName);
+
+private:
+  // The buffer holding the content read from remapping file.
+  std::unique_ptr<MemoryBuffer> Buffer;
+  std::unique_ptr<SymbolRemappingReader> Remappings;
+  DenseMap<SymbolRemappingReader::Key, FunctionSamples *> SampleMap;
+  // The Reader the remapper is servicing.
+  SampleProfileReader &Reader;
+  // Indicate whether remapping has been applied to the profile read
+  // by Reader -- by calling applyRemapping.
+  bool RemappingApplied = false;
+};
+
 /// Sample-based profile reader.
 ///
 /// Each profile contains sample counts for all the functions
 /// executed. Inside each function, statements are annotated with the
 /// collected samples on all the instructions associated with that
 /// statement.
 ///
 /// For this to produce meaningful data, the program needs to be
SampleProfileReader(std::unique_ptr<MemoryBuffer> B, LLVMContext &C,
                       SampleProfileFormat Format = SPF_None)
       : Profiles(0), Ctx(C), Buffer(std::move(B)), Format(Format) {}
 
   virtual ~SampleProfileReader() = default;
 
   /// Read and validate the file header.
   virtual std::error_code readHeader() = 0;
 
-  /// Read sample profiles from the associated file.
-  virtual std::error_code read() = 0;
+  /// The interface to read sample profiles from the associated file.
+  std::error_code read() {
+    if (std::error_code EC = readImpl())
+      return EC;
+    if (Remapper)
+      Remapper->applyRemapping(Ctx);
+    return sampleprof_error::success;
+  }
+
+  /// The implementaion to read sample profiles from the associated file.
+  virtual std::error_code readImpl() = 0;
 
   /// Print the profile for \p FName on stream \p OS.
   void dumpFunctionProfile(StringRef FName, raw_ostream &OS = dbgs());
 
   virtual void collectFuncsFrom(const Module &M) {}
 
   /// Print all the profiles on stream \p OS.
   void dump(raw_ostream &OS = dbgs());
 
   /// Return the samples collected for function \p F.
   FunctionSamples *getSamplesFor(const Function &F) {
     // The function name may have been updated by adding suffix. Call
     // a helper to (optionally) strip off suffixes so that we can
     // match against the original function name in the profile.
     StringRef CanonName = FunctionSamples::getCanonicalFnName(F);
     return getSamplesFor(CanonName);
   }
 
   /// Return the samples collected for function \p F.
   virtual FunctionSamples *getSamplesFor(StringRef Fname) {
+    if (Remapper) {
+      if (auto FS = Remapper->getSamplesFor(Fname))
+        return FS;
+    }
     std::string FGUID;
     Fname = getRepInFormat(Fname, getFormat(), FGUID);
     auto It = Profiles.find(Fname);
     if (It != Profiles.end())
       return &It->second;
     return nullptr;
   }
 
   /// Return all the profiles.
   StringMap<FunctionSamples> &getProfiles() { return Profiles; }
 
   /// Report a parse error message.
   void reportError(int64_t LineNumber, Twine Msg) const {
     Ctx.diagnose(DiagnosticInfoSampleProfile(Buffer->getBufferIdentifier(),
                                              LineNumber, Msg));
   }
 
   /// Create a sample profile reader appropriate to the file format.
+  /// Create a remapper underlying if RemapFilename is not empty.
   static ErrorOr<std::unique_ptr<SampleProfileReader>>
-  create(const Twine &Filename, LLVMContext &C);
+  create(const std::string Filename, LLVMContext &C,
+         const std::string RemapFilename = "");
 
   /// Create a sample profile reader from the supplied memory buffer.
+  /// Create a remapper underlying if RemapFilename is not empty.
   static ErrorOr<std::unique_ptr<SampleProfileReader>>
-  create(std::unique_ptr<MemoryBuffer> &B, LLVMContext &C);
+  create(std::unique_ptr<MemoryBuffer> &B, LLVMContext &C,
+         const std::string RemapFilename = "");
 
   /// Return the profile summary.
-  ProfileSummary &getSummary() { return *(Summary.get()); }
+  ProfileSummary &getSummary() const { return *(Summary.get()); }
+
+  MemoryBuffer *getBuffer() const { return Buffer.get(); }
 
   /// \brief Return the profile format.
-  SampleProfileFormat getFormat() { return Format; }
+  SampleProfileFormat getFormat() const { return Format; }
 
   virtual std::unique_ptr<ProfileSymbolList> getProfileSymbolList() {
     return nullptr;
   };
 
   /// It includes all the names that have samples either in outline instance
   /// or inline instance.
   virtual std::vector<StringRef> *getNameTable() { return nullptr; }
protected:
   static std::unique_ptr<ProfileSummary>
   takeSummary(SampleProfileReader &Reader) {
     return std::move(Reader.Summary);
   }
 
   /// Compute summary for this profile.
   void computeSummary();
 
+  std::unique_ptr<SampleProfileReaderItaniumRemapper> Remapper;
+
   /// \brief The format of sample.
   SampleProfileFormat Format = SPF_None;
 };
 
 class SampleProfileReaderText : public SampleProfileReader {
 public:
   SampleProfileReaderText(std::unique_ptr<MemoryBuffer> B, LLVMContext &C)
       : SampleProfileReader(std::move(B), C, SPF_Text) {}
 
   /// Read and validate the file header.
   std::error_code readHeader() override { return sampleprof_error::success; }
 
   /// Read sample profiles from the associated file.
-  std::error_code read() override;
+  std::error_code readImpl() override;
 
   /// Return true if \p Buffer is in the format supported by this class.
   static bool hasFormat(const MemoryBuffer &Buffer);
 };
 
 class SampleProfileReaderBinary : public SampleProfileReader {
 public:
   SampleProfileReaderBinary(std::unique_ptr<MemoryBuffer> B, LLVMContext &C,
                             SampleProfileFormat Format = SPF_None)
       : SampleProfileReader(std::move(B), C, Format) {}
 
   /// Read and validate the file header.
   virtual std::error_code readHeader() override;
 
   /// Read sample profiles from the associated file.
-  std::error_code read() override;
+  std::error_code readImpl() override;
 
   /// It includes all the names that have samples either in outline instance
   /// or inline instance.
   virtual std::vector<StringRef> *getNameTable() override { return &NameTable; }
 
 protected:
   /// Read a numeric value of type T from the profile.
   ///
virtual std::error_code readOneSection(const uint8_t *Start, uint64_t Size,
                                          SecType Type) = 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 read() override;
+  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;
 };
 
public:
   SampleProfileReaderCompactBinary(std::unique_ptr<MemoryBuffer> B,
                                    LLVMContext &C)
       : SampleProfileReaderBinary(std::move(B), C, SPF_Compact_Binary) {}
 
   /// \brief Return true if \p Buffer is in the format supported by this class.
   static bool hasFormat(const MemoryBuffer &Buffer);
 
   /// Read samples only for functions to use.
-  std::error_code read() override;
+  std::error_code readImpl() override;
 
   /// Collect functions to be used when compiling Module \p M.
   void collectFuncsFrom(const Module &M) override;
 };
 
 using InlineCallStack = SmallVector<FunctionSamples *, 10>;
 
 // Supported histogram types in GCC.  Currently, we only need support for
public:
   SampleProfileReaderGCC(std::unique_ptr<MemoryBuffer> B, LLVMContext &C)
       : SampleProfileReader(std::move(B), C, SPF_GCC),
         GcovBuffer(Buffer.get()) {}
 
   /// Read and validate the file header.
   std::error_code readHeader() override;
 
   /// Read sample profiles from the associated file.
-  std::error_code read() override;
+  std::error_code readImpl() override;
 
   /// Return true if \p Buffer is in the format supported by this class.
   static bool hasFormat(const MemoryBuffer &Buffer);
 
 protected:
   std::error_code readNameTable();
   std::error_code readOneFunctionProfile(const InlineCallStack &InlineStack,
                                          bool Update, uint32_t Offset);
protected:
   /// Function names in this profile.
   std::vector<std::string> Names;
 
   /// GCOV tags used to separate sections in the profile file.
   static const uint32_t GCOVTagAFDOFileNames = 0xaa000000;
   static const uint32_t GCOVTagAFDOFunction = 0xac000000;
 };
 
-/// A profile data reader proxy that remaps the profile data from another
-/// sample profile data reader, by applying a provided set of equivalences
-/// between components of the symbol names in the profile.
-class SampleProfileReaderItaniumRemapper : public SampleProfileReader {
-public:
-  SampleProfileReaderItaniumRemapper(
-      std::unique_ptr<MemoryBuffer> B, LLVMContext &C,
-      std::unique_ptr<SampleProfileReader> Underlying)
-      : SampleProfileReader(std::move(B), C, Underlying->getFormat()) {
-    Profiles = std::move(Underlying->getProfiles());
-    Summary = takeSummary(*Underlying);
-    // Keep the underlying reader alive; the profile data may contain
-    // StringRefs referencing names in its name table.
-    UnderlyingReader = std::move(Underlying);
-  }
-
-  /// Create a remapped sample profile from the given remapping file and
-  /// underlying samples.
-  static ErrorOr<std::unique_ptr<SampleProfileReader>>
-  create(const Twine &Filename, LLVMContext &C,
-         std::unique_ptr<SampleProfileReader> Underlying);
-
-  /// Read and validate the file header.
-  std::error_code readHeader() override { return sampleprof_error::success; }
-
-  /// Read remapping file and apply it to the sample profile.
-  std::error_code read() override;
-
-  /// Return the samples collected for function \p F.
-  FunctionSamples *getSamplesFor(StringRef FunctionName) override;
-  using SampleProfileReader::getSamplesFor;
-
-private:
-  SymbolRemappingReader Remappings;
-  DenseMap<SymbolRemappingReader::Key, FunctionSamples*> SampleMap;
-  std::unique_ptr<SampleProfileReader> UnderlyingReader;
-};
-
 } // end namespace sampleprof
 
 } // end namespace llvm
 
 #endif // LLVM_PROFILEDATA_SAMPLEPROFREADER_H

llvm/lib/ProfileData/SampleProfReader.cpp

 }
 
 /// Load samples from a text file.
 ///
 /// See the documentation at the top of the file for an explanation of
 /// the expected format.
 ///
 /// \returns true if the file was loaded successfully, false otherwise.
-std::error_code SampleProfileReaderText::read() {
+std::error_code SampleProfileReaderText::readImpl() {
   line_iterator LineIt(*Buffer, /*SkipBlanks=*/true, '#');
   sampleprof_error Result = sampleprof_error::success;
 
   InlineCallStack InlineStack;
 
   for (; !LineIt.is_at_eof(); ++LineIt) {
     if ((*LineIt)[(*LineIt).find_first_not_of(' ')] == '#')
       continue;
SampleProfileReaderBinary::readFuncProfile(const uint8_t *Start) {
 
   FProfile.addHeadSamples(*NumHeadSamples);
 
   if (std::error_code EC = readProfile(FProfile))
     return EC;
   return sampleprof_error::success;
 }
 
-std::error_code SampleProfileReaderBinary::read() {
+std::error_code SampleProfileReaderBinary::readImpl() {
   while (!at_eof()) {
     if (std::error_code EC = readFuncProfile(Data))
       return EC;
   }
 
   return sampleprof_error::success;
 }
 
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;
   }
 
+  if (Remapper) {
     for (auto Name : FuncsToUse) {
-    auto iter = FuncOffsetTable.find(Name);
-    if (iter == FuncOffsetTable.end())
+      Remapper->insert(Name);
+    }
+  }
+
+  for (auto NameOffset : FuncOffsetTable) {
+    auto FuncName = NameOffset.first;
+    if (!FuncsToUse.count(FuncName) &&
+        (!Remapper || !Remapper->exist(FuncName)))
       continue;
-    const uint8_t *FuncProfileAddr = Start + iter->second;
+    const uint8_t *FuncProfileAddr = Start + NameOffset.second;
     assert(FuncProfileAddr < End && "out of LBRProfile section");
     if (std::error_code EC = readFuncProfile(FuncProfileAddr))
       return EC;
   }
+
   Data = End;
   return sampleprof_error::success;
 }
 
 std::error_code SampleProfileReaderExtBinary::readProfileSymbolList() {
   if (!ProfSymList)
     ProfSymList = std::make_unique<ProfileSymbolList>();
 
std::error_code SampleProfileReaderExtBinaryBase::decompressSection(
   llvm::Error E =
       zlib::uncompress(CompressedStrings, Buffer, UCSize);
   if (E)
     return sampleprof_error::uncompress_failed;
   DecompressBuf = reinterpret_cast<const uint8_t *>(Buffer);
   return sampleprof_error::success;
 }
 
-std::error_code SampleProfileReaderExtBinaryBase::read() {
+std::error_code SampleProfileReaderExtBinaryBase::readImpl() {
   const uint8_t *BufStart =
       reinterpret_cast<const uint8_t *>(Buffer->getBufferStart());
 
   for (auto &Entry : SecHdrTable) {
     // Skip empty section.
     if (!Entry.Size)
       continue;
 
if (isCompressed) {
       Data = BufStart + Entry.Offset;
       End = BufStart + Buffer->getBufferSize();
     }
   }
 
   return sampleprof_error::success;
 }
 
-std::error_code SampleProfileReaderCompactBinary::read() {
+std::error_code SampleProfileReaderCompactBinary::readImpl() {
   std::vector<uint64_t> OffsetsToUse;
   if (UseAllFuncs) {
     for (auto FuncEntry : FuncOffsetTable) {
       OffsetsToUse.push_back(FuncEntry.second);
     }
   }
   else {
     for (auto Name : FuncsToUse) {
std::error_code SampleProfileReaderGCC::readOneFunctionProfile(
 
   return sampleprof_error::success;
 }
 
 /// Read a GCC AutoFDO profile.
 ///
 /// This format is generated by the Linux Perf conversion tool at
 /// https://github.com/google/autofdo.
-std::error_code SampleProfileReaderGCC::read() {
+std::error_code SampleProfileReaderGCC::readImpl() {
   // Read the string table.
   if (std::error_code EC = readNameTable())
     return EC;
 
   // Read the source profile.
   if (std::error_code EC = readFunctionProfiles())
     return EC;
 
   return sampleprof_error::success;
 }
 
 bool SampleProfileReaderGCC::hasFormat(const MemoryBuffer &Buffer) {
   StringRef Magic(reinterpret_cast<const char *>(Buffer.getBufferStart()));
   return Magic == "adcg*704";
 }
 
-std::error_code SampleProfileReaderItaniumRemapper::read() {
-  // If the underlying data is in compact format, we can't remap it because
+void SampleProfileReaderItaniumRemapper::applyRemapping(LLVMContext &Ctx) {
+  // If the reader is in compact format, we can't remap it because
   // we don't know what the original function names were.
-  if (getFormat() == SPF_Compact_Binary) {
+  if (Reader.getFormat() == SPF_Compact_Binary) {
     Ctx.diagnose(DiagnosticInfoSampleProfile(
-        Buffer->getBufferIdentifier(),
+        Reader.getBuffer()->getBufferIdentifier(),
         "Profile data remapping cannot be applied to profile data "
         "in compact format (original mangled names are not available).",
         DS_Warning));
-    return sampleprof_error::success;
-  }
-
-  if (Error E = Remappings.read(*Buffer)) {
-    handleAllErrors(
-        std::move(E), [&](const SymbolRemappingParseError &ParseError) {
-          reportError(ParseError.getLineNum(), ParseError.getMessage());
-        });
-    return sampleprof_error::malformed;
+    return;
   }
 
-  for (auto &Sample : getProfiles())
-    if (auto Key = Remappings.insert(Sample.first()))
+  assert(Remappings && "should be initialized while creating remapper");
+  for (auto &Sample : Reader.getProfiles())
+    if (auto Key = Remappings->insert(Sample.first()))
       SampleMap.insert({Key, &Sample.second});
 
-  return sampleprof_error::success;
+  RemappingApplied = true;
 }
 
 FunctionSamples *
 SampleProfileReaderItaniumRemapper::getSamplesFor(StringRef Fname) {
-  if (auto Key = Remappings.lookup(Fname))
+  if (auto Key = Remappings->lookup(Fname))
     return SampleMap.lookup(Key);
-  return SampleProfileReader::getSamplesFor(Fname);
+  return nullptr;
 }
 
 /// Prepare a memory buffer for the contents of \p Filename.
 ///
 /// \returns an error code indicating the status of the buffer.
 static ErrorOr<std::unique_ptr<MemoryBuffer>>
 setupMemoryBuffer(const Twine &Filename) {
   auto BufferOrErr = MemoryBuffer::getFileOrSTDIN(Filename);

 }
 
 /// Create a sample profile reader based on the format of the input file.
 ///
 /// \param Filename The file to open.
 ///
 /// \param C The LLVM context to use to emit diagnostics.
 ///
+/// \param RemapFilename The file used for profile remapping.
+///
 /// \returns an error code indicating the status of the created reader.
 ErrorOr<std::unique_ptr<SampleProfileReader>>
-SampleProfileReader::create(const Twine &Filename, LLVMContext &C) {
+SampleProfileReader::create(const std::string Filename, LLVMContext &C,
+                            const std::string RemapFilename) {
   auto BufferOrError = setupMemoryBuffer(Filename);
   if (std::error_code EC = BufferOrError.getError())
     return EC;
-  return create(BufferOrError.get(), C);
+  return create(BufferOrError.get(), C, RemapFilename);
 }
 
 /// Create a sample profile remapper from the given input, to remap the
 /// function names in the given profile data.
 ///
 /// \param Filename The file to open.
 ///
-/// \param C The LLVM context to use to emit diagnostics.
+/// \param Reader The profile reader the remapper is going to be applied to.
 ///
-/// \param Underlying The underlying profile data reader to remap.
+/// \param C The LLVM context to use to emit diagnostics.
 ///
 /// \returns an error code indicating the status of the created reader.
-ErrorOr<std::unique_ptr<SampleProfileReader>>
-SampleProfileReaderItaniumRemapper::create(
-    const Twine &Filename, LLVMContext &C,
-    std::unique_ptr<SampleProfileReader> Underlying) {
+ErrorOr<std::unique_ptr<SampleProfileReaderItaniumRemapper>>
+SampleProfileReaderItaniumRemapper::create(const std::string Filename,
+                                           SampleProfileReader &Reader,
+                                           LLVMContext &C) {
   auto BufferOrError = setupMemoryBuffer(Filename);
   if (std::error_code EC = BufferOrError.getError())
     return EC;
+  return create(BufferOrError.get(), Reader, C);
+}
+
+/// Create a sample profile remapper from the given input, to remap the
+/// function names in the given profile data.
+///
+/// \param B The memory buffer to create the reader from (assumes ownership).
+///
+/// \param C The LLVM context to use to emit diagnostics.
+///
+/// \param Reader The profile reader the remapper is going to be applied to.
+///
+/// \returns an error code indicating the status of the created reader.
+ErrorOr<std::unique_ptr<SampleProfileReaderItaniumRemapper>>
+SampleProfileReaderItaniumRemapper::create(std::unique_ptr<MemoryBuffer> &B,
+                                           SampleProfileReader &Reader,
+                                           LLVMContext &C) {
+  auto Remappings = std::make_unique<SymbolRemappingReader>();
+  if (Error E = Remappings->read(*B.get())) {
+    handleAllErrors(
+        std::move(E), [&](const SymbolRemappingParseError &ParseError) {
+          C.diagnose(DiagnosticInfoSampleProfile(B->getBufferIdentifier(),
+                                                 ParseError.getLineNum(),
+                                                 ParseError.getMessage()));
+        });
+    return sampleprof_error::malformed;
+  }
+
   return std::make_unique<SampleProfileReaderItaniumRemapper>(
-      std::move(BufferOrError.get()), C, std::move(Underlying));
+      std::move(B), std::move(Remappings), Reader);
 }
 
 /// Create a sample profile reader based on the format of the input data.
 ///
 /// \param B The memory buffer to create the reader from (assumes ownership).
 ///
 /// \param C The LLVM context to use to emit diagnostics.
 ///
+/// \param RemapFilename The file used for profile remapping.
+///
 /// \returns an error code indicating the status of the created reader.
 ErrorOr<std::unique_ptr<SampleProfileReader>>
-SampleProfileReader::create(std::unique_ptr<MemoryBuffer> &B, LLVMContext &C) {
+SampleProfileReader::create(std::unique_ptr<MemoryBuffer> &B, LLVMContext &C,
+                            const std::string RemapFilename) {
   std::unique_ptr<SampleProfileReader> Reader;
   if (SampleProfileReaderRawBinary::hasFormat(*B))
     Reader.reset(new SampleProfileReaderRawBinary(std::move(B), C));
   else if (SampleProfileReaderExtBinary::hasFormat(*B))
     Reader.reset(new SampleProfileReaderExtBinary(std::move(B), C));
   else if (SampleProfileReaderCompactBinary::hasFormat(*B))
     Reader.reset(new SampleProfileReaderCompactBinary(std::move(B), C));
   else if (SampleProfileReaderGCC::hasFormat(*B))
     Reader.reset(new SampleProfileReaderGCC(std::move(B), C));
   else if (SampleProfileReaderText::hasFormat(*B))
     Reader.reset(new SampleProfileReaderText(std::move(B), C));
   else
     return sampleprof_error::unrecognized_format;
 
+  if (!RemapFilename.empty()) {
+    auto ReaderOrErr =
+        SampleProfileReaderItaniumRemapper::create(RemapFilename, *Reader, C);
+    if (std::error_code EC = ReaderOrErr.getError()) {
+      std::string Msg = "Could not create remapper: " + EC.message();
+      C.diagnose(DiagnosticInfoSampleProfile(RemapFilename, Msg));
+      return EC;
+    }
+    Reader->Remapper = std::move(ReaderOrErr.get());
+  }
+
   FunctionSamples::Format = Reader->getFormat();
   if (std::error_code EC = Reader->readHeader()) {
     return EC;
   }
 
   return std::move(Reader);
 }
 

llvm/lib/Transforms/IPO/SampleProfile.cpp

 INITIALIZE_PASS_DEPENDENCY(AssumptionCacheTracker)
 INITIALIZE_PASS_DEPENDENCY(TargetTransformInfoWrapperPass)
 INITIALIZE_PASS_DEPENDENCY(ProfileSummaryInfoWrapperPass)
 INITIALIZE_PASS_END(SampleProfileLoaderLegacyPass, "sample-profile",
                     "Sample Profile loader", false, false)
 
 bool SampleProfileLoader::doInitialization(Module &M) {
   auto &Ctx = M.getContext();
-  auto ReaderOrErr = SampleProfileReader::create(Filename, Ctx);
+
+  std::unique_ptr<SampleProfileReaderItaniumRemapper> RemapReader;
+  auto ReaderOrErr =
+      SampleProfileReader::create(Filename, Ctx, RemappingFilename);
   if (std::error_code EC = ReaderOrErr.getError()) {
     std::string Msg = "Could not open profile: " + EC.message();
     Ctx.diagnose(DiagnosticInfoSampleProfile(Filename, Msg));
     return false;
   }
   Reader = std::move(ReaderOrErr.get());
   Reader->collectFuncsFrom(M);
   ProfileIsValid = (Reader->read() == sampleprof_error::success);
   PSL = Reader->getProfileSymbolList();
 
   // While profile-sample-accurate is on, ignore symbol list.
   ProfAccForSymsInList =
       ProfileAccurateForSymsInList && PSL && !ProfileSampleAccurate;
   if (ProfAccForSymsInList) {
     NamesInProfile.clear();
     if (auto NameTable = Reader->getNameTable())
       NamesInProfile.insert(NameTable->begin(), NameTable->end());
   }
 
-  if (!RemappingFilename.empty()) {
-    // Apply profile remappings to the loaded profile data if requested.
-    // For now, we only support remapping symbols encoded using the Itanium
-    // C++ ABI's name mangling scheme.
-    ReaderOrErr = SampleProfileReaderItaniumRemapper::create(
-        RemappingFilename, Ctx, std::move(Reader));
-    if (std::error_code EC = ReaderOrErr.getError()) {
-      std::string Msg = "Could not open profile remapping file: " + EC.message();
-      Ctx.diagnose(DiagnosticInfoSampleProfile(Filename, Msg));
-      return false;
-    }
-    Reader = std::move(ReaderOrErr.get());
-    ProfileIsValid = (Reader->read() == sampleprof_error::success);
-  }
   return true;
 }
 
 ModulePass *llvm::createSampleProfileLoaderPass() {
   return new SampleProfileLoaderLegacyPass();
 }
 
 ModulePass *llvm::createSampleProfileLoaderPass(StringRef Name) {
   return new SampleProfileLoaderLegacyPass(Name);

llvm/test/Transforms/SampleProfile/remap.ll

 ; RUN: opt %s -passes=sample-profile -sample-profile-file=%S/Inputs/remap.prof -sample-profile-remapping-file=%S/Inputs/remap.map | opt -analyze -branch-prob | FileCheck %s
-
+;
+; Check whether profile remapping work with loading profile on demand used by extbinary format profile.
+; RUN: llvm-profdata merge -sample -extbinary %S/Inputs/remap.prof -o %t.extbinary.afdo
+; RUN: opt %s -passes=sample-profile -sample-profile-file=%S/Inputs/remap.prof -sample-profile-remapping-file=%S/Inputs/remap.map | opt -analyze -branch-prob | FileCheck %s
+;
 ; Reduced from branch.ll
 
 declare i1 @foo()
 
 define void @_ZN3foo3barERKN1M1XINS_6detail3quxEEE() !dbg !2 {
 ; CHECK: Printing analysis 'Branch Probability Analysis' for function '_ZN3foo3barERKN1M1XINS_6detail3quxEEE':
 
 entry:

llvm/unittests/ProfileData/SampleProfTest.cpp

void createWriter(SampleProfileFormat Format, StringRef Profile) {
     std::error_code EC;
     std::unique_ptr<raw_ostream> OS(
         new raw_fd_ostream(Profile, EC, sys::fs::OF_None));
     auto WriterOrErr = SampleProfileWriter::create(OS, Format);
     ASSERT_TRUE(NoError(WriterOrErr.getError()));
     Writer = std::move(WriterOrErr.get());
   }
 
-  void readProfile(const Module &M, StringRef Profile) {
-    auto ReaderOrErr = SampleProfileReader::create(Profile, Context);
+  void readProfile(const Module &M, StringRef Profile,
+                   StringRef RemapFile = "") {
+    auto ReaderOrErr = SampleProfileReader::create(Profile, Context, RemapFile);
     ASSERT_TRUE(NoError(ReaderOrErr.getError()));
     Reader = std::move(ReaderOrErr.get());
     Reader->collectFuncsFrom(M);
   }
 
+  void createRemapFile(SmallVectorImpl<char> &RemapPath, StringRef &RemapFile) {
+    std::error_code EC =
+        llvm::sys::fs::createTemporaryFile("remapfile", "", RemapPath);
+    ASSERT_TRUE(NoError(EC));
+    RemapFile = StringRef(RemapPath.data(), RemapPath.size());
+
+    std::unique_ptr<raw_fd_ostream> OS(
+        new raw_fd_ostream(RemapFile, EC, sys::fs::OF_None));
+    *OS << R"(
+      # Types 'int' and 'long' are equivalent
+      type i l
+      # Function names 'foo' and 'faux' are equivalent
+      name 3foo 4faux
+    )";
+    OS->close();
+  }
+
   void testRoundTrip(SampleProfileFormat Format, bool Remap) {
     SmallVector<char, 128> ProfilePath;
     ASSERT_TRUE(NoError(llvm::sys::fs::createTemporaryFile("profile", "", ProfilePath)));
     StringRef Profile(ProfilePath.data(), ProfilePath.size());
     createWriter(Format, Profile);
 
     StringRef FooName("_Z3fooi");
     FunctionSamples FooSamples;
void testRoundTrip(SampleProfileFormat Format, bool Remap) {
 
     StringRef BazName("_Z3bazi");
     FunctionSamples BazSamples;
     BazSamples.setName(BazName);
     BazSamples.addTotalSamples(12557);
     BazSamples.addHeadSamples(1257);
     BazSamples.addBodySamples(1, 0, 12557);
 
-    Module M("my_module", Context);
-    FunctionType *fn_type =
-        FunctionType::get(Type::getVoidTy(Context), {}, false);
-    M.getOrInsertFunction(FooName, fn_type);
-    M.getOrInsertFunction(BarName, fn_type);
+    StringRef BooName("_Z3booi");
+    FunctionSamples BooSamples;
+    BooSamples.setName(BooName);
+    BooSamples.addTotalSamples(1232);
+    BooSamples.addHeadSamples(1);
+    BooSamples.addBodySamples(1, 0, 1232);
 
     StringMap<FunctionSamples> Profiles;
     Profiles[FooName] = std::move(FooSamples);
     Profiles[BarName] = std::move(BarSamples);
     Profiles[BazName] = std::move(BazSamples);
+    Profiles[BooName] = std::move(BooSamples);
+
+    Module M("my_module", Context);
+    FunctionType *fn_type =
+        FunctionType::get(Type::getVoidTy(Context), {}, false);
+
+    SmallVector<char, 128> RemapPath;
+    StringRef RemapFile;
+    if (Remap) {
+      createRemapFile(RemapPath, RemapFile);
+      FooName = "_Z4fauxi";
+      BarName = "_Z3barl";
+    }
+
+    M.getOrInsertFunction(FooName, fn_type);
+    M.getOrInsertFunction(BarName, fn_type);
+    M.getOrInsertFunction(BooName, fn_type);
 
     ProfileSymbolList List;
     if (Format == SampleProfileFormat::SPF_Ext_Binary) {
       List.add("zoo", true);
       List.add("moo", true);
     }
     Writer->setProfileSymbolList(&List);
 
     std::error_code EC;
     EC = Writer->write(Profiles);
     ASSERT_TRUE(NoError(EC));
 
     Writer->getOutputStream().flush();
 
-    readProfile(M, Profile);
-
+    readProfile(M, Profile, RemapFile);
     EC = Reader->read();
     ASSERT_TRUE(NoError(EC));
 
     if (Format == SampleProfileFormat::SPF_Ext_Binary) {
       std::unique_ptr<ProfileSymbolList> ReaderList =
           Reader->getProfileSymbolList();
       ReaderList->contains("zoo");
       ReaderList->contains("moo");
     }
 
-    if (Remap) {
-      auto MemBuffer = llvm::MemoryBuffer::getMemBuffer(R"(
-        # Types 'int' and 'long' are equivalent
-        type i l
-        # Function names 'foo' and 'faux' are equivalent
-        name 3foo 4faux
-      )");
-      Reader.reset(new SampleProfileReaderItaniumRemapper(
-          std::move(MemBuffer), Context, std::move(Reader)));
-      FooName = "_Z4fauxi";
-      BarName = "_Z3barl";
-
-      EC = Reader->read();
-      ASSERT_TRUE(NoError(EC));
-    }
-
     FunctionSamples *ReadFooSamples = Reader->getSamplesFor(FooName);
     ASSERT_TRUE(ReadFooSamples != nullptr);
     if (Format != SampleProfileFormat::SPF_Compact_Binary) {
       ASSERT_EQ("_Z3fooi", ReadFooSamples->getName());
     }
     ASSERT_EQ(7711u, ReadFooSamples->getTotalSamples());
     ASSERT_EQ(610u, ReadFooSamples->getHeadSamples());
 
void testRoundTrip(SampleProfileFormat Format, bool Remap) {
 
     // Because _Z3bazi is not defined in module M, expect _Z3bazi's profile
     // is not loaded when the profile is ExtBinary or Compact format because
     // these formats support loading function profiles on demand.
     FunctionSamples *ReadBazSamples = Reader->getSamplesFor(BazName);
     if (Format == SampleProfileFormat::SPF_Ext_Binary ||
         Format == SampleProfileFormat::SPF_Compact_Binary) {
       ASSERT_TRUE(ReadBazSamples == nullptr);
-      ASSERT_EQ(2u, Reader->getProfiles().size());
+      ASSERT_EQ(3u, Reader->getProfiles().size());
     } else {
       ASSERT_TRUE(ReadBazSamples != nullptr);
       ASSERT_EQ(12557u, ReadBazSamples->getTotalSamples());
-      ASSERT_EQ(3u, Reader->getProfiles().size());
+      ASSERT_EQ(4u, Reader->getProfiles().size());
     }
 
+    FunctionSamples *ReadBooSamples = Reader->getSamplesFor(BooName);
+    ASSERT_TRUE(ReadBooSamples != nullptr);
+    ASSERT_EQ(1232u, ReadBooSamples->getTotalSamples());
+
     std::string MconstructGUID;
     StringRef MconstructRep =
         getRepInFormat(MconstructName, Format, MconstructGUID);
     std::string StringviewGUID;
     StringRef StringviewRep =
         getRepInFormat(StringviewName, Format, StringviewGUID);
     ASSERT_EQ(1000u, CTMap.get()[MconstructRep]);
     ASSERT_EQ(437u, CTMap.get()[StringviewRep]);
 
     auto VerifySummary = [](ProfileSummary &Summary) mutable {
       ASSERT_EQ(ProfileSummary::PSK_Sample, Summary.getKind());
-      ASSERT_EQ(136160u, Summary.getTotalCount());
-      ASSERT_EQ(7u, Summary.getNumCounts());
-      ASSERT_EQ(3u, Summary.getNumFunctions());
+      ASSERT_EQ(137392u, Summary.getTotalCount());
+      ASSERT_EQ(8u, Summary.getNumCounts());
+      ASSERT_EQ(4u, Summary.getNumFunctions());
       ASSERT_EQ(1437u, Summary.getMaxFunctionCount());
       ASSERT_EQ(60351u, Summary.getMaxCount());
 
       uint32_t Cutoff = 800000;
       auto Predicate = [&Cutoff](const ProfileSummaryEntry &PE) {
         return PE.Cutoff == Cutoff;
       };
       std::vector<ProfileSummaryEntry> &Details = Summary.getDetailedSummary();