[SampleFDO] Add FunctionOffsetTable in compact binary format profile.

[wmi]

Summary

The patch saves a function offset table which maps function name index to the offset of its function profile to the start of the binary profile. By using the function offset table, for those function profiles which will not be used when compiling a module, the profile reader does't have to read them. For profile size around 10~20M, it saves ~10% compile time.

[wmi]

Add some cleanup, please take another look:
Add virtual function collectFuncsToUse to remove the awareness of format change in sample profile loader pass.

llvm/trunk/include/llvm/ProfileData/SampleProf.h

enum class sampleprof_error {
   unsupported_version,
   too_large,
   truncated,
   malformed,
   unrecognized_format,
   unsupported_writing_format,
   truncated_name_table,
   not_implemented,
-  counter_overflow
+  counter_overflow,
+  ostream_seek_unsupported
 };
 
 inline std::error_code make_error_code(sampleprof_error E) {
   return std::error_code(static_cast<int>(E), sampleprof_category());
 }
 
 inline sampleprof_error MergeResult(sampleprof_error &Accumulator,
                                     sampleprof_error Result) {

llvm/trunk/include/llvm/ProfileData/SampleProfReader.h

public:
   virtual std::error_code readHeader() = 0;
 
   /// Read sample profiles from the associated file.
   virtual std::error_code read() = 0;
 
   /// Print the profile for \p FName on stream \p OS.
   void dumpFunctionProfile(StringRef FName, raw_ostream &OS = dbgs());
 
+  virtual void collectFuncsToUse(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. In sample
     // profile, the function names are all stripped, so we need to strip
     // the function name suffix before matching with profile.

 
 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.
-  std::error_code readHeader() override;
+  virtual std::error_code readHeader() override;
 
   /// Read sample profiles from the associated file.
   std::error_code read() override;
 
 protected:
   /// Read a numeric value of type T from the profile.
   ///
   /// If an error occurs during decoding, a diagnostic message is emitted and
   /// EC is set.
   ///
   /// \returns the read value.
   template <typename T> ErrorOr<T> readNumber();
 
+  /// Read a numeric value of type T from the profile. The value is saved
+  /// without encoded.
+  template <typename T> ErrorOr<T> readUnencodedNumber();
+
   /// Read a string from the profile.
   ///
   /// If an error occurs during decoding, a diagnostic message is emitted and
   /// EC is set.
   ///
   /// \returns the read value.
   ErrorOr<StringRef> readString();
 
   /// Read the string index and check whether it overflows the table.
   template <typename T> inline ErrorOr<uint32_t> readStringIndex(T &Table);
 
   /// Return true if we've reached the end of file.
   bool at_eof() const { return Data >= End; }
 
+  /// Read the next function profile instance.
+  std::error_code readFuncProfile();
+
   /// Read the contents of the given profile instance.
   std::error_code readProfile(FunctionSamples &FProfile);
 
   /// Points to the current location in the buffer.
   const uint8_t *Data = nullptr;
 
   /// Points to the end of the buffer.
   const uint8_t *End = nullptr;
public:
   /// \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.
+  DenseMap<StringRef, uint64_t> FuncOffsetTable;
+  /// The set containing the functions to use when compiling a module.
+  DenseSet<StringRef> FuncsToUse;
   virtual std::error_code verifySPMagic(uint64_t Magic) override;
   virtual std::error_code readNameTable() override;
   /// Read a string indirectly via the name table.
   virtual ErrorOr<StringRef> readStringFromTable() override;
+  virtual std::error_code readHeader() override;
+  std::error_code readFuncOffsetTable();
 
 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;
+
+  /// Collect functions to be used when compiling Module \p M.
+  void collectFuncsToUse(const Module &M) override;
 };
 
 using InlineCallStack = SmallVector<FunctionSamples *, 10>;
 
 // Supported histogram types in GCC.  Currently, we only need support for
 // call target histograms.
 enum HistType {
   HIST_TYPE_INTERVAL,

llvm/trunk/include/llvm/ProfileData/SampleProfWriter.h

public:
   /// Write sample profiles in \p S.
   ///
   /// \returns status code of the file update operation.
   virtual std::error_code write(const FunctionSamples &S) = 0;
 
   /// Write all the sample profiles in the given map of samples.
   ///
   /// \returns status code of the file update operation.
-  std::error_code write(const StringMap<FunctionSamples> &ProfileMap);
+  virtual std::error_code write(const StringMap<FunctionSamples> &ProfileMap);
 
   raw_ostream &getOutputStream() { return *OutputStream; }
 
   /// Profile writer factory.
   ///
   /// Create a new file writer based on the value of \p Format.
   static ErrorOr<std::unique_ptr<SampleProfileWriter>>
   create(StringRef Filename, SampleProfileFormat Format);
private:
   friend ErrorOr<std::unique_ptr<SampleProfileWriter>>
   SampleProfileWriter::create(std::unique_ptr<raw_ostream> &OS,
                               SampleProfileFormat Format);
 };
 
 /// Sample-based profile writer (binary format).
 class SampleProfileWriterBinary : public SampleProfileWriter {
 public:
-  std::error_code write(const FunctionSamples &S) override;
+  virtual std::error_code write(const FunctionSamples &S) override;
   SampleProfileWriterBinary(std::unique_ptr<raw_ostream> &OS)
       : SampleProfileWriter(OS) {}
 
 protected:
   virtual std::error_code writeNameTable() = 0;
   virtual std::error_code writeMagicIdent() = 0;
-  std::error_code writeHeader(const StringMap<FunctionSamples> &ProfileMap) override;
+  virtual std::error_code
+  writeHeader(const StringMap<FunctionSamples> &ProfileMap) override;
   std::error_code writeSummary();
   std::error_code writeNameIdx(StringRef FName);
   std::error_code writeBody(const FunctionSamples &S);
   inline void stablizeNameTable(std::set<StringRef> &V);
 
   MapVector<StringRef, uint32_t> NameTable;
 
 private:
   void addName(StringRef FName);
   void addNames(const FunctionSamples &S);
 
   friend ErrorOr<std::unique_ptr<SampleProfileWriter>>
   SampleProfileWriter::create(std::unique_ptr<raw_ostream> &OS,
                               SampleProfileFormat Format);
 };
 
 class SampleProfileWriterRawBinary : public SampleProfileWriterBinary {
   using SampleProfileWriterBinary::SampleProfileWriterBinary;
 
 protected:
   virtual std::error_code writeNameTable() override;
   virtual std::error_code writeMagicIdent() override;
 };
 
+// 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,
+// so by using the function offset table, for those function profiles which
+// will not be needed when compiling a module, the profile reader does't
+// have to read them and it saves compile time if the profile size is huge.
+// The layout of the compact format is shown as follows:
+//
+//    Part1: Profile header, the same as binary format, containing magic
+//           number, version, summary, name table...
+//    Part2: Function Offset Table Offset, which saves the position of
+//           Part4.
+//    Part3: Function profile collection
+//             function1 profile start
+//                 ....
+//             function2 profile start
+//                 ....
+//             function3 profile start
+//                 ....
+//                ......
+//    Part4: Function Offset Table
+//             function1 name index --> function1 profile start
+//             function2 name index --> function2 profile start
+//             function3 name index --> function3 profile start
+//
+// We need Part2 because profile reader can use it to find out and read
+// function offset table without reading Part3 first.
 class SampleProfileWriterCompactBinary : public SampleProfileWriterBinary {
   using SampleProfileWriterBinary::SampleProfileWriterBinary;
 
+public:
+  virtual std::error_code write(const FunctionSamples &S) override;
+  virtual std::error_code
+  write(const StringMap<FunctionSamples> &ProfileMap) override;
+
 protected:
+  /// The table mapping from function name to the offset of its FunctionSample
+  /// towards profile start.
+  MapVector<StringRef, uint64_t> FuncOffsetTable;
+  /// The offset of the slot to be filled with the offset of FuncOffsetTable
+  /// towards profile start.
+  uint64_t TableOffset;
   virtual std::error_code writeNameTable() override;
   virtual std::error_code writeMagicIdent() override;
+  virtual std::error_code
+  writeHeader(const StringMap<FunctionSamples> &ProfileMap) override;
+  std::error_code writeFuncOffsetTable();
 };
 
 } // end namespace sampleprof
 } // end namespace llvm
 
 #endif // LLVM_PROFILEDATA_SAMPLEPROFWRITER_H

llvm/trunk/lib/ProfileData/SampleProf.cpp

std::string message(int IE) const override {
     case sampleprof_error::unsupported_writing_format:
       return "Profile encoding format unsupported for writing operations";
     case sampleprof_error::truncated_name_table:
       return "Truncated function name table";
     case sampleprof_error::not_implemented:
       return "Unimplemented feature";
     case sampleprof_error::counter_overflow:
       return "Counter overflow";
+    case sampleprof_error::ostream_seek_unsupported:
+      return "Ostream does not support seek";
     }
     llvm_unreachable("A value of sampleprof_error has no message.");
   }
 };
 
 } // end anonymous namespace
 
 static ManagedStatic<SampleProfErrorCategoryType> ErrorCategory;

llvm/trunk/lib/ProfileData/SampleProfReader.cpp

 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/IR/ProfileSummary.h"
 #include "llvm/ProfileData/ProfileCommon.h"
 #include "llvm/ProfileData/SampleProf.h"
 #include "llvm/Support/ErrorOr.h"
 #include "llvm/Support/LEB128.h"
 #include "llvm/Support/LineIterator.h"
+#include "llvm/Support/MD5.h"
 #include "llvm/Support/MemoryBuffer.h"
 #include "llvm/Support/raw_ostream.h"
 #include <algorithm>
 #include <cstddef>
 #include <cstdint>
 #include <limits>
 #include <memory>
 #include <system_error>
if (Data + Str.size() + 1 > End) {
     return EC;
   }
 
   Data += Str.size() + 1;
   return Str;
 }
 
 template <typename T>
+ErrorOr<T> SampleProfileReaderBinary::readUnencodedNumber() {
+  std::error_code EC;
+
+  if (Data + sizeof(T) > End) {
+    EC = sampleprof_error::truncated;
+    reportError(0, EC.message());
+    return EC;
+  }
+
+  using namespace support;
+  T Val = endian::readNext<T, little, unaligned>(Data);
+  return Val;
+}
+
+template <typename T>
 inline ErrorOr<uint32_t> SampleProfileReaderBinary::readStringIndex(T &Table) {
   std::error_code EC;
   auto Idx = readNumber<uint32_t>();
   if (std::error_code EC = Idx.getError())
     return EC;
   if (*Idx >= Table.size())
     return sampleprof_error::truncated_name_table;
   return *Idx;
for (uint32_t J = 0; J < *NumCallsites; ++J) {
     CalleeProfile.setName(*FName);
     if (std::error_code EC = readProfile(CalleeProfile))
       return EC;
   }
 
   return sampleprof_error::success;
 }
 
-std::error_code SampleProfileReaderBinary::read() {
-  while (!at_eof()) {
-    auto NumHeadSamples = readNumber<uint64_t>();
-    if (std::error_code EC = NumHeadSamples.getError())
-      return EC;
+std::error_code SampleProfileReaderBinary::readFuncProfile() {
 
-    auto FName(readStringFromTable());
-    if (std::error_code EC = FName.getError())
-      return EC;
 
-    Profiles[*FName] = FunctionSamples();
-    FunctionSamples &FProfile = Profiles[*FName];
-    FProfile.setName(*FName);
 
-    FProfile.addHeadSamples(*NumHeadSamples);
 
-    if (std::error_code EC = readProfile(FProfile))
-      return EC;
-  }
+  return sampleprof_error::success;
 
+std::error_code SampleProfileReaderBinary::read() {
+  while (!at_eof()) {
+    if (std::error_code EC = readFuncProfile())
+      return EC;
+  }
+
+  return sampleprof_error::success;
+}
+
+std::error_code SampleProfileReaderCompactBinary::read() {
+  for (auto Name : FuncsToUse) {
+    auto GUID = std::to_string(MD5Hash(Name));
+    auto iter = FuncOffsetTable.find(StringRef(GUID));
+    if (iter == FuncOffsetTable.end())
+      continue;
+    const uint8_t *SavedData = Data;
+    Data = reinterpret_cast<const uint8_t *>(Buffer->getBufferStart()) +
+           iter->second;
+    if (std::error_code EC = readFuncProfile())
+      return EC;
+    Data = SavedData;
+  }
   return sampleprof_error::success;
 }
 
 std::error_code SampleProfileReaderRawBinary::verifySPMagic(uint64_t Magic) {
   if (Magic == SPMagic())
     return sampleprof_error::success;
   return sampleprof_error::bad_magic;
 }
std::error_code SampleProfileReaderBinary::readHeader() {
   if (std::error_code EC = readSummary())
     return EC;
 
   if (std::error_code EC = readNameTable())
     return EC;
   return sampleprof_error::success;
 }
 
+std::error_code SampleProfileReaderCompactBinary::readHeader() {
+  SampleProfileReaderBinary::readHeader();
+  if (std::error_code EC = readFuncOffsetTable())
+    return EC;
+  return sampleprof_error::success;
+}
+
+std::error_code SampleProfileReaderCompactBinary::readFuncOffsetTable() {
+  auto TableOffset = readUnencodedNumber<uint64_t>();
+  if (std::error_code EC = TableOffset.getError())
+    return EC;
+
+  const uint8_t *SavedData = Data;
+  const uint8_t *TableStart =
+      reinterpret_cast<const uint8_t *>(Buffer->getBufferStart()) +
+      *TableOffset;
+  Data = TableStart;
+
+  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;
+  }
+  End = TableStart;
+  Data = SavedData;
+  return sampleprof_error::success;
+}
+
+void SampleProfileReaderCompactBinary::collectFuncsToUse(const Module &M) {
+  FuncsToUse.clear();
+  for (auto &F : M) {
+    StringRef Fname = F.getName().split('.').first;
+    FuncsToUse.insert(Fname);
+  }
+}
+
 std::error_code SampleProfileReaderBinary::readSummaryEntry(
     std::vector<ProfileSummaryEntry> &Entries) {
   auto Cutoff = readNumber<uint64_t>();
   if (std::error_code EC = Cutoff.getError())
     return EC;
 
   auto MinBlockCount = readNumber<uint64_t>();
   if (std::error_code EC = MinBlockCount.getError())

llvm/trunk/lib/ProfileData/SampleProfWriter.cpp

 // supported formats.
 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/ProfileData/SampleProfWriter.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/ProfileData/ProfileCommon.h"
 #include "llvm/ProfileData/SampleProf.h"
+#include "llvm/Support/Endian.h"
+#include "llvm/Support/EndianStream.h"
 #include "llvm/Support/ErrorOr.h"
 #include "llvm/Support/FileSystem.h"
 #include "llvm/Support/LEB128.h"
 #include "llvm/Support/MD5.h"
 #include "llvm/Support/raw_ostream.h"
 #include <algorithm>
 #include <cstdint>
 #include <memory>
SampleProfileWriter::write(const StringMap<FunctionSamples> &ProfileMap) {
 
   for (const auto &I : V) {
     if (std::error_code EC = write(*I.second))
       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())
+    return EC;
+  return sampleprof_error::success;
+}
+
 /// Write samples to a text file.
 ///
 /// Note: it may be tempting to implement this in terms of
 /// FunctionSamples::print().  Please don't.  The dump functionality is intended
 /// for debugging and has no specified form.
 ///
 /// The format used here is more structured and deliberate because
 /// it needs to be parsed by the SampleProfileReaderText class.
std::error_code SampleProfileWriterRawBinary::writeNameTable() {
   encodeULEB128(NameTable.size(), OS);
   for (auto N : V) {
     OS << N;
     encodeULEB128(0, OS);
   }
   return sampleprof_error::success;
 }
 
+std::error_code SampleProfileWriterCompactBinary::writeFuncOffsetTable() {
+  auto &OS = *OutputStream;
+
+  // Fill the slot remembered by TableOffset with the offset of FuncOffsetTable.
+  auto &OFS = static_cast<raw_fd_ostream &>(OS);
+  uint64_t FuncOffsetTableStart = OS.tell();
+  if (OFS.seek(TableOffset) == (uint64_t)-1)
+    return sampleprof_error::ostream_seek_unsupported;
+  support::endian::Writer Writer(*OutputStream, support::little);
+  Writer.write(FuncOffsetTableStart);
+  if (OFS.seek(FuncOffsetTableStart) == (uint64_t)-1)
+    return sampleprof_error::ostream_seek_unsupported;
+
+  // 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 SampleProfileWriterCompactBinary::writeNameTable() {
   auto &OS = *OutputStream;
   std::set<StringRef> V;
   stablizeNameTable(V);
 
   // Write out the name table.
   encodeULEB128(NameTable.size(), OS);
   for (auto N : V) {
for (const auto &I : ProfileMap) {
     addName(I.first());
     addNames(I.second);
   }
 
   writeNameTable();
   return sampleprof_error::success;
 }
 
+std::error_code SampleProfileWriterCompactBinary::writeHeader(
+    const StringMap<FunctionSamples> &ProfileMap) {
+  support::endian::Writer Writer(*OutputStream, support::little);
+  if (auto EC = SampleProfileWriterBinary::writeHeader(ProfileMap))
+    return EC;
+
+  // Reserve a slot for the offset of function offset table. The slot will
+  // be populated with the offset of FuncOffsetTable later.
+  TableOffset = OutputStream->tell();
+  Writer.write(static_cast<uint64_t>(-2));
+  return sampleprof_error::success;
+}
+
 std::error_code SampleProfileWriterBinary::writeSummary() {
   auto &OS = *OutputStream;
   encodeULEB128(Summary->getTotalCount(), OS);
   encodeULEB128(Summary->getMaxCount(), OS);
   encodeULEB128(Summary->getMaxFunctionCount(), OS);
   encodeULEB128(Summary->getNumCounts(), OS);
   encodeULEB128(Summary->getNumFunctions(), OS);
   std::vector<ProfileSummaryEntry> &Entries = Summary->getDetailedSummary();

 /// Write samples of a top-level function to a binary file.
 ///
 /// \returns true if the samples were written successfully, false otherwise.
 std::error_code SampleProfileWriterBinary::write(const FunctionSamples &S) {
   encodeULEB128(S.getHeadSamples(), *OutputStream);
   return writeBody(S);
 }
 
+std::error_code
+SampleProfileWriterCompactBinary::write(const FunctionSamples &S) {
+  uint64_t Offset = OutputStream->tell();
+  StringRef Name = S.getName();
+  FuncOffsetTable[Name] = Offset;
+  encodeULEB128(S.getHeadSamples(), *OutputStream);
+  return writeBody(S);
+}
+
 /// Create a sample profile file writer based on the specified format.
 ///
 /// \param Filename The file to create.
 ///
 /// \param Format Encoding format for the profile file.
 ///
 /// \returns an error code indicating the status of the created writer.
 ErrorOr<std::unique_ptr<SampleProfileWriter>>

llvm/trunk/lib/Transforms/IPO/SampleProfile.cpp

bool SampleProfileLoader::doInitialization(Module &M) {
   auto &Ctx = M.getContext();
   auto ReaderOrErr = SampleProfileReader::create(Filename, Ctx);
   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->collectFuncsToUse(M);
   ProfileIsValid = (Reader->read() == sampleprof_error::success);
   return true;
 }
 
 ModulePass *llvm::createSampleProfileLoaderPass() {
   return new SampleProfileLoaderLegacyPass(SampleProfileFile);
 }
 

llvm/trunk/unittests/ProfileData/SampleProfTest.cpp

if (!EC)
     return ::testing::AssertionSuccess();
   return ::testing::AssertionFailure() << "error " << EC.value() << ": "
                                        << EC.message();
 }
 
 namespace {
 
 struct SampleProfTest : ::testing::Test {
-  std::string Data;
   LLVMContext Context;
+  std::string Profile;
   std::unique_ptr<raw_ostream> OS;
   std::unique_ptr<SampleProfileWriter> Writer;
   std::unique_ptr<SampleProfileReader> Reader;
+  std::error_code EC;
 
   SampleProfTest()
-      : Data(), OS(new raw_string_ostream(Data)), Writer(), Reader() {}
+      : Profile("profile"),
+        OS(new raw_fd_ostream(Profile, EC, sys::fs::F_None)), Writer(),
+        Reader() {}
 
   void createWriter(SampleProfileFormat Format) {
     auto WriterOrErr = SampleProfileWriter::create(OS, Format);
     ASSERT_TRUE(NoError(WriterOrErr.getError()));
     Writer = std::move(WriterOrErr.get());
   }
 
-  void readProfile(std::unique_ptr<MemoryBuffer> &Profile) {
+  void readProfile(const Module &M) {
     auto ReaderOrErr = SampleProfileReader::create(Profile, Context);
     ASSERT_TRUE(NoError(ReaderOrErr.getError()));
     Reader = std::move(ReaderOrErr.get());
+    Reader->collectFuncsToUse(M);
   }
 
   void testRoundTrip(SampleProfileFormat Format) {
     createWriter(Format);
 
     StringRef FooName("_Z3fooi");
     FunctionSamples FooSamples;
     FooSamples.setName(FooName);
void testRoundTrip(SampleProfileFormat Format) {
     BarSamples.addHeadSamples(1437);
     BarSamples.addBodySamples(1, 0, 1437);
     // Test how reader/writer handles unmangled names.
     StringRef MconstructName("_M_construct<char *>");
     StringRef StringviewName("string_view<std::allocator<char> >");
     BarSamples.addCalledTargetSamples(1, 0, MconstructName, 1000);
     BarSamples.addCalledTargetSamples(1, 0, StringviewName, 437);
 
+    Module M("my_module", Context);
+    FunctionType *fn_type =
+        FunctionType::get(Type::getVoidTy(Context), {}, false);
+    M.getOrInsertFunction(FooName, fn_type);
+    M.getOrInsertFunction(BarName, fn_type);
+
     StringMap<FunctionSamples> Profiles;
     Profiles[FooName] = std::move(FooSamples);
     Profiles[BarName] = std::move(BarSamples);
 
     std::error_code EC;
     EC = Writer->write(Profiles);
     ASSERT_TRUE(NoError(EC));
 
     Writer->getOutputStream().flush();
 
-    auto Profile = MemoryBuffer::getMemBufferCopy(Data);
-    readProfile(Profile);
+    readProfile(M);
 
     EC = Reader->read();
     ASSERT_TRUE(NoError(EC));
 
     StringMap<FunctionSamples> &ReadProfiles = Reader->getProfiles();
     ASSERT_EQ(2u, ReadProfiles.size());
 
     std::string FooGUID;
void testRoundTrip(SampleProfileFormat Format) {
     Metadata *MD = Summary.getMD(Context);
     ASSERT_TRUE(MD);
     ProfileSummary *PS = ProfileSummary::getFromMD(MD);
     ASSERT_TRUE(PS);
     VerifySummary(*PS);
     delete PS;
 
     // Test that summary can be attached to and read back from module.
-    Module M("my_module", Context);
     M.setProfileSummary(MD);
     MD = M.getProfileSummary();
     ASSERT_TRUE(MD);
     PS = ProfileSummary::getFromMD(MD);
     ASSERT_TRUE(PS);
     VerifySummary(*PS);
     delete PS;
   }