llvm/lib/ProfileData/SampleProfWriter.cpp

   1 //===- SampleProfWriter.cpp - Write LLVM sample profile data --------------===//
   2 //
   3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
   4 // See https://llvm.org/LICENSE.txt for license information.
   5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
   6 //
   7 //===----------------------------------------------------------------------===//
   8 //
   9 // This file implements the class that writes LLVM sample profiles. It
  10 // supports two file formats: text and binary. The textual representation
  11 // is useful for debugging and testing purposes. The binary representation
  12 // is more compact, resulting in smaller file sizes. However, they can
  13 // both be used interchangeably.
  14 //
  15 // See lib/ProfileData/SampleProfReader.cpp for documentation on each of the
  16 // supported formats.
  17 //
  18 //===----------------------------------------------------------------------===//
  19
  20 #include "llvm/ProfileData/SampleProfWriter.h"
  21 #include "llvm/ADT/StringRef.h"
  22 #include "llvm/ADT/StringSet.h"
  23 #include "llvm/ProfileData/ProfileCommon.h"
  24 #include "llvm/ProfileData/SampleProf.h"
  25 #include "llvm/Support/Compression.h"
  26 #include "llvm/Support/Endian.h"
  27 #include "llvm/Support/EndianStream.h"
  28 #include "llvm/Support/ErrorOr.h"
  29 #include "llvm/Support/FileSystem.h"
  30 #include "llvm/Support/LEB128.h"
  31 #include "llvm/Support/MD5.h"
  32 #include "llvm/Support/raw_ostream.h"
  33 #include <algorithm>
  34 #include <cstdint>
  35 #include <memory>
  36 #include <set>
  37 #include <system_error>
  38 #include <utility>
  39 #include <vector>
  40
  41 using namespace llvm;
  42 using namespace sampleprof;
  43
  44 std::error_code
  45 SampleProfileWriter::writeFuncProfiles(const SampleProfileMap &ProfileMap) {
  46   std::vector<NameFunctionSamples> V;
  47   sortFuncProfiles(ProfileMap, V);
  48   for (const auto &I : V) {
  49     if (std::error_code EC = writeSample(*I.second))
  50       return EC;
  51   }
  52   return sampleprof_error::success;
  53 }
  54
  55 std::error_code SampleProfileWriter::write(const SampleProfileMap &ProfileMap) {
  56   if (std::error_code EC = writeHeader(ProfileMap))
  57     return EC;
  58
  59   if (std::error_code EC = writeFuncProfiles(ProfileMap))
  60     return EC;
  61
  62   return sampleprof_error::success;
  63 }
  64
  65 /// Return the current position and prepare to use it as the start
  66 /// position of a section given the section type \p Type and its position
  67 /// \p LayoutIdx in SectionHdrLayout.
  68 uint64_t
  69 SampleProfileWriterExtBinaryBase::markSectionStart(SecType Type,
  70                                                    uint32_t LayoutIdx) {
  71   uint64_t SectionStart = OutputStream->tell();
  72   assert(LayoutIdx < SectionHdrLayout.size() && "LayoutIdx out of range");
  73   const auto &Entry = SectionHdrLayout[LayoutIdx];
  74   assert(Entry.Type == Type && "Unexpected section type");
  75   // Use LocalBuf as a temporary output for writting data.
  76   if (hasSecFlag(Entry, SecCommonFlags::SecFlagCompress))
  77     LocalBufStream.swap(OutputStream);
  78   return SectionStart;
  79 }
  80
  81 std::error_code SampleProfileWriterExtBinaryBase::compressAndOutput() {
  82   if (!llvm::zlib::isAvailable())
  83     return sampleprof_error::zlib_unavailable;
  84   std::string &UncompressedStrings =
  85       static_cast<raw_string_ostream *>(LocalBufStream.get())->str();
  86   if (UncompressedStrings.size() == 0)
  87     return sampleprof_error::success;
  88   auto &OS = *OutputStream;
  89   SmallString<128> CompressedStrings;
  90   llvm::Error E = zlib::compress(UncompressedStrings, CompressedStrings,
  91                                  zlib::BestSizeCompression);
  92   if (E)
  93     return sampleprof_error::compress_failed;
  94   encodeULEB128(UncompressedStrings.size(), OS);
  95   encodeULEB128(CompressedStrings.size(), OS);
  96   OS << CompressedStrings.str();
  97   UncompressedStrings.clear();
  98   return sampleprof_error::success;
  99 }
 100
 101 /// Add a new section into section header table given the section type
 102 /// \p Type, its position \p LayoutIdx in SectionHdrLayout and the
 103 /// location \p SectionStart where the section should be written to.
 104 std::error_code SampleProfileWriterExtBinaryBase::addNewSection(
 105     SecType Type, uint32_t LayoutIdx, uint64_t SectionStart) {
 106   assert(LayoutIdx < SectionHdrLayout.size() && "LayoutIdx out of range");
 107   const auto &Entry = SectionHdrLayout[LayoutIdx];
 108   assert(Entry.Type == Type && "Unexpected section type");
 109   if (hasSecFlag(Entry, SecCommonFlags::SecFlagCompress)) {
 110     LocalBufStream.swap(OutputStream);
 111     if (std::error_code EC = compressAndOutput())
 112       return EC;
 113   }
 114   SecHdrTable.push_back({Type, Entry.Flags, SectionStart - FileStart,
 115                          OutputStream->tell() - SectionStart, LayoutIdx});
 116   return sampleprof_error::success;
 117 }
 118
 119 std::error_code
 120 SampleProfileWriterExtBinaryBase::write(const SampleProfileMap &ProfileMap) {
 121   if (std::error_code EC = writeHeader(ProfileMap))
 122     return EC;
 123
 124   std::string LocalBuf;
 125   LocalBufStream = std::make_unique<raw_string_ostream>(LocalBuf);
 126   if (std::error_code EC = writeSections(ProfileMap))
 127     return EC;
 128
 129   if (std::error_code EC = writeSecHdrTable())
 130     return EC;
 131
 132   return sampleprof_error::success;
 133 }
 134
 135 std::error_code SampleProfileWriterExtBinaryBase::writeContextIdx(
 136     const SampleContext &Context) {
 137   if (Context.hasContext())
 138     return writeCSNameIdx(Context);
 139   else
 140     return SampleProfileWriterBinary::writeNameIdx(Context.getName());
 141 }
 142
 143 std::error_code
 144 SampleProfileWriterExtBinaryBase::writeCSNameIdx(const SampleContext &Context) {
 145   const auto &Ret = CSNameTable.find(Context);
 146   if (Ret == CSNameTable.end())
 147     return sampleprof_error::truncated_name_table;
 148   encodeULEB128(Ret->second, *OutputStream);
 149   return sampleprof_error::success;
 150 }
 151
 152 std::error_code
 153 SampleProfileWriterExtBinaryBase::writeSample(const FunctionSamples &S) {
 154   uint64_t Offset = OutputStream->tell();
 155   auto &Context = S.getContext();
 156   FuncOffsetTable[Context] = Offset - SecLBRProfileStart;
 157   encodeULEB128(S.getHeadSamples(), *OutputStream);
 158   return writeBody(S);
 159 }
 160
 161 std::error_code SampleProfileWriterExtBinaryBase::writeFuncOffsetTable() {
 162   auto &OS = *OutputStream;
 163
 164   // Write out the table size.
 165   encodeULEB128(FuncOffsetTable.size(), OS);
 166
 167   // Write out FuncOffsetTable.
 168   auto WriteItem = [&](const SampleContext &Context, uint64_t Offset) {
 169     if (std::error_code EC = writeContextIdx(Context))
 170       return EC;
 171     encodeULEB128(Offset, OS);
 172     return (std::error_code)sampleprof_error::success;
 173   };
 174
 175   if (FunctionSamples::ProfileIsCSFlat) {
 176     // Sort the contexts before writing them out. This is to help fast load all
 177     // context profiles for a function as well as their callee contexts which
 178     // can help profile-guided importing for ThinLTO.
 179     std::map<SampleContext, uint64_t> OrderedFuncOffsetTable(
 180         FuncOffsetTable.begin(), FuncOffsetTable.end());
 181     for (const auto &Entry : OrderedFuncOffsetTable) {
 182       if (std::error_code EC = WriteItem(Entry.first, Entry.second))
 183         return EC;
 184     }
 185     addSectionFlag(SecFuncOffsetTable, SecFuncOffsetFlags::SecFlagOrdered);
 186   } else {
 187     for (const auto &Entry : FuncOffsetTable) {
 188       if (std::error_code EC = WriteItem(Entry.first, Entry.second))
 189         return EC;
 190     }
 191   }
 192
 193   FuncOffsetTable.clear();
 194   return sampleprof_error::success;
 195 }
 196
 197 std::error_code SampleProfileWriterExtBinaryBase::writeFuncMetadata(
 198     const FunctionSamples &FunctionProfile) {
 199   auto &OS = *OutputStream;
 200   if (std::error_code EC = writeContextIdx(FunctionProfile.getContext()))
 201     return EC;
 202
 203   if (FunctionSamples::ProfileIsProbeBased)
 204     encodeULEB128(FunctionProfile.getFunctionHash(), OS);
 205   if (FunctionSamples::ProfileIsCSFlat || FunctionSamples::ProfileIsCSNested) {
 206     encodeULEB128(FunctionProfile.getContext().getAllAttributes(), OS);
 207   }
 208
 209   if (!FunctionSamples::ProfileIsCSFlat) {
 210     // Recursively emit attributes for all callee samples.
 211     uint64_t NumCallsites = 0;
 212     for (const auto &J : FunctionProfile.getCallsiteSamples())
 213       NumCallsites += J.second.size();
 214     encodeULEB128(NumCallsites, OS);
 215     for (const auto &J : FunctionProfile.getCallsiteSamples()) {
 216       for (const auto &FS : J.second) {
 217         LineLocation Loc = J.first;
 218         encodeULEB128(Loc.LineOffset, OS);
 219         encodeULEB128(Loc.Discriminator, OS);
 220         if (std::error_code EC = writeFuncMetadata(FS.second))
 221           return EC;
 222       }
 223     }
 224   }
 225
 226   return sampleprof_error::success;
 227 }
 228
 229 std::error_code SampleProfileWriterExtBinaryBase::writeFuncMetadata(
 230     const SampleProfileMap &Profiles) {
 231   if (!FunctionSamples::ProfileIsProbeBased &&
 232       !FunctionSamples::ProfileIsCSFlat && !FunctionSamples::ProfileIsCSNested)
 233     return sampleprof_error::success;
 234   for (const auto &Entry : Profiles) {
 235     if (std::error_code EC = writeFuncMetadata(Entry.second))
 236       return EC;
 237   }
 238   return sampleprof_error::success;
 239 }
 240
 241 std::error_code SampleProfileWriterExtBinaryBase::writeNameTable() {
 242   if (!UseMD5)
 243     return SampleProfileWriterBinary::writeNameTable();
 244
 245   auto &OS = *OutputStream;
 246   std::set<StringRef> V;
 247   stablizeNameTable(NameTable, V);
 248
 249   // Write out the MD5 name table. We wrote unencoded MD5 so reader can
 250   // retrieve the name using the name index without having to read the
 251   // whole name table.
 252   encodeULEB128(NameTable.size(), OS);
 253   support::endian::Writer Writer(OS, support::little);
 254   for (auto N : V)
 255     Writer.write(MD5Hash(N));
 256   return sampleprof_error::success;
 257 }
 258
 259 std::error_code SampleProfileWriterExtBinaryBase::writeNameTableSection(
 260     const SampleProfileMap &ProfileMap) {
 261   for (const auto &I : ProfileMap) {
 262     assert(I.first == I.second.getContext() && "Inconsistent profile map");
 263     addContext(I.second.getContext());
 264     addNames(I.second);
 265   }
 266
 267   // If NameTable contains ".__uniq." suffix, set SecFlagUniqSuffix flag
 268   // so compiler won't strip the suffix during profile matching after
 269   // seeing the flag in the profile.
 270   for (const auto &I : NameTable) {
 271     if (I.first.contains(FunctionSamples::UniqSuffix)) {
 272       addSectionFlag(SecNameTable, SecNameTableFlags::SecFlagUniqSuffix);
 273       break;
 274     }
 275   }
 276
 277   if (auto EC = writeNameTable())
 278     return EC;
 279   return sampleprof_error::success;
 280 }
 281
 282 std::error_code SampleProfileWriterExtBinaryBase::writeCSNameTableSection() {
 283   // Sort the names to make CSNameTable deterministic.
 284   std::set<SampleContext> OrderedContexts;
 285   for (const auto &I : CSNameTable)
 286     OrderedContexts.insert(I.first);
 287   assert(OrderedContexts.size() == CSNameTable.size() &&
 288          "Unmatched ordered and unordered contexts");
 289   uint64_t I = 0;
 290   for (auto &Context : OrderedContexts)
 291     CSNameTable[Context] = I++;
 292
 293   auto &OS = *OutputStream;
 294   encodeULEB128(OrderedContexts.size(), OS);
 295   support::endian::Writer Writer(OS, support::little);
 296   for (auto Context : OrderedContexts) {
 297     auto Frames = Context.getContextFrames();
 298     encodeULEB128(Frames.size(), OS);
 299     for (auto &Callsite : Frames) {
 300       if (std::error_code EC = writeNameIdx(Callsite.FuncName))
 301         return EC;
 302       encodeULEB128(Callsite.Location.LineOffset, OS);
 303       encodeULEB128(Callsite.Location.Discriminator, OS);
 304     }
 305   }
 306
 307   return sampleprof_error::success;
 308 }
 309
 310 std::error_code
 311 SampleProfileWriterExtBinaryBase::writeProfileSymbolListSection() {
 312   if (ProfSymList && ProfSymList->size() > 0)
 313     if (std::error_code EC = ProfSymList->write(*OutputStream))
 314       return EC;
 315
 316   return sampleprof_error::success;
 317 }
 318
 319 std::error_code SampleProfileWriterExtBinaryBase::writeOneSection(
 320     SecType Type, uint32_t LayoutIdx, const SampleProfileMap &ProfileMap) {
 321   // The setting of SecFlagCompress should happen before markSectionStart.
 322   if (Type == SecProfileSymbolList && ProfSymList && ProfSymList->toCompress())
 323     setToCompressSection(SecProfileSymbolList);
 324   if (Type == SecFuncMetadata && FunctionSamples::ProfileIsProbeBased)
 325     addSectionFlag(SecFuncMetadata, SecFuncMetadataFlags::SecFlagIsProbeBased);
 326   if (Type == SecFuncMetadata && FunctionSamples::ProfileIsCSNested)
 327     addSectionFlag(SecFuncMetadata, SecFuncMetadataFlags::SecFlagIsCSNested);
 328   if (Type == SecFuncMetadata &&
 329       (FunctionSamples::ProfileIsCSFlat || FunctionSamples::ProfileIsCSNested))
 330     addSectionFlag(SecFuncMetadata, SecFuncMetadataFlags::SecFlagHasAttribute);
 331   if (Type == SecProfSummary && FunctionSamples::ProfileIsCSFlat)
 332     addSectionFlag(SecProfSummary, SecProfSummaryFlags::SecFlagFullContext);
 333   if (Type == SecProfSummary && FunctionSamples::ProfileIsFS)
 334     addSectionFlag(SecProfSummary, SecProfSummaryFlags::SecFlagFSDiscriminator);
 335
 336   uint64_t SectionStart = markSectionStart(Type, LayoutIdx);
 337   switch (Type) {
 338   case SecProfSummary:
 339     computeSummary(ProfileMap);
 340     if (auto EC = writeSummary())
 341       return EC;
 342     break;
 343   case SecNameTable:
 344     if (auto EC = writeNameTableSection(ProfileMap))
 345       return EC;
 346     break;
 347   case SecCSNameTable:
 348     if (auto EC = writeCSNameTableSection())
 349       return EC;
 350     break;
 351   case SecLBRProfile:
 352     SecLBRProfileStart = OutputStream->tell();
 353     if (std::error_code EC = writeFuncProfiles(ProfileMap))
 354       return EC;
 355     break;
 356   case SecFuncOffsetTable:
 357     if (auto EC = writeFuncOffsetTable())
 358       return EC;
 359     break;
 360   case SecFuncMetadata:
 361     if (std::error_code EC = writeFuncMetadata(ProfileMap))
 362       return EC;
 363     break;
 364   case SecProfileSymbolList:
 365     if (auto EC = writeProfileSymbolListSection())
 366       return EC;
 367     break;
 368   default:
 369     if (auto EC = writeCustomSection(Type))
 370       return EC;
 371     break;
 372   }
 373   if (std::error_code EC = addNewSection(Type, LayoutIdx, SectionStart))
 374     return EC;
 375   return sampleprof_error::success;
 376 }
 377
 378 std::error_code SampleProfileWriterExtBinary::writeDefaultLayout(
 379     const SampleProfileMap &ProfileMap) {
 380   // The const indices passed to writeOneSection below are specifying the
 381   // positions of the sections in SectionHdrLayout. Look at
 382   // initSectionHdrLayout to find out where each section is located in
 383   // SectionHdrLayout.
 384   if (auto EC = writeOneSection(SecProfSummary, 0, ProfileMap))
 385     return EC;
 386   if (auto EC = writeOneSection(SecNameTable, 1, ProfileMap))
 387     return EC;
 388   if (auto EC = writeOneSection(SecCSNameTable, 2, ProfileMap))
 389     return EC;
 390   if (auto EC = writeOneSection(SecLBRProfile, 4, ProfileMap))
 391     return EC;
 392   if (auto EC = writeOneSection(SecProfileSymbolList, 5, ProfileMap))
 393     return EC;
 394   if (auto EC = writeOneSection(SecFuncOffsetTable, 3, ProfileMap))
 395     return EC;
 396   if (auto EC = writeOneSection(SecFuncMetadata, 6, ProfileMap))
 397     return EC;
 398   return sampleprof_error::success;
 399 }
 400
 401 static void splitProfileMapToTwo(const SampleProfileMap &ProfileMap,
 402                                  SampleProfileMap &ContextProfileMap,
 403                                  SampleProfileMap &NoContextProfileMap) {
 404   for (const auto &I : ProfileMap) {
 405     if (I.second.getCallsiteSamples().size())
 406       ContextProfileMap.insert({I.first, I.second});
 407     else
 408       NoContextProfileMap.insert({I.first, I.second});
 409   }
 410 }
 411
 412 std::error_code SampleProfileWriterExtBinary::writeCtxSplitLayout(
 413     const SampleProfileMap &ProfileMap) {
 414   SampleProfileMap ContextProfileMap, NoContextProfileMap;
 415   splitProfileMapToTwo(ProfileMap, ContextProfileMap, NoContextProfileMap);
 416
 417   if (auto EC = writeOneSection(SecProfSummary, 0, ProfileMap))
 418     return EC;
 419   if (auto EC = writeOneSection(SecNameTable, 1, ProfileMap))
 420     return EC;
 421   if (auto EC = writeOneSection(SecLBRProfile, 3, ContextProfileMap))
 422     return EC;
 423   if (auto EC = writeOneSection(SecFuncOffsetTable, 2, ContextProfileMap))
 424     return EC;
 425   // Mark the section to have no context. Note section flag needs to be set
 426   // before writing the section.
 427   addSectionFlag(5, SecCommonFlags::SecFlagFlat);
 428   if (auto EC = writeOneSection(SecLBRProfile, 5, NoContextProfileMap))
 429     return EC;
 430   // Mark the section to have no context. Note section flag needs to be set
 431   // before writing the section.
 432   addSectionFlag(4, SecCommonFlags::SecFlagFlat);
 433   if (auto EC = writeOneSection(SecFuncOffsetTable, 4, NoContextProfileMap))
 434     return EC;
 435   if (auto EC = writeOneSection(SecProfileSymbolList, 6, ProfileMap))
 436     return EC;
 437   if (auto EC = writeOneSection(SecFuncMetadata, 7, ProfileMap))
 438     return EC;
 439
 440   return sampleprof_error::success;
 441 }
 442
 443 std::error_code SampleProfileWriterExtBinary::writeSections(
 444     const SampleProfileMap &ProfileMap) {
 445   std::error_code EC;
 446   if (SecLayout == DefaultLayout)
 447     EC = writeDefaultLayout(ProfileMap);
 448   else if (SecLayout == CtxSplitLayout)
 449     EC = writeCtxSplitLayout(ProfileMap);
 450   else
 451     llvm_unreachable("Unsupported layout");
 452   return EC;
 453 }
 454
 455 std::error_code
 456 SampleProfileWriterCompactBinary::write(const SampleProfileMap &ProfileMap) {
 457   if (std::error_code EC = SampleProfileWriter::write(ProfileMap))
 458     return EC;
 459   if (std::error_code EC = writeFuncOffsetTable())
 460     return EC;
 461   return sampleprof_error::success;
 462 }
 463
 464 /// Write samples to a text file.
 465 ///
 466 /// Note: it may be tempting to implement this in terms of
 467 /// FunctionSamples::print().  Please don't.  The dump functionality is intended
 468 /// for debugging and has no specified form.
 469 ///
 470 /// The format used here is more structured and deliberate because
 471 /// it needs to be parsed by the SampleProfileReaderText class.
 472 std::error_code SampleProfileWriterText::writeSample(const FunctionSamples &S) {
 473   auto &OS = *OutputStream;
 474   if (FunctionSamples::ProfileIsCSFlat)
 475     OS << "[" << S.getContext().toString() << "]:" << S.getTotalSamples();
 476   else
 477     OS << S.getName() << ":" << S.getTotalSamples();
 478
 479   if (Indent == 0)
 480     OS << ":" << S.getHeadSamples();
 481   OS << "\n";
 482
 483   SampleSorter<LineLocation, SampleRecord> SortedSamples(S.getBodySamples());
 484   for (const auto &I : SortedSamples.get()) {
 485     LineLocation Loc = I->first;
 486     const SampleRecord &Sample = I->second;
 487     OS.indent(Indent + 1);
 488     if (Loc.Discriminator == 0)
 489       OS << Loc.LineOffset << ": ";
 490     else
 491       OS << Loc.LineOffset << "." << Loc.Discriminator << ": ";
 492
 493     OS << Sample.getSamples();
 494
 495     for (const auto &J : Sample.getSortedCallTargets())
 496       OS << " " << J.first << ":" << J.second;
 497     OS << "\n";
 498   }
 499
 500   SampleSorter<LineLocation, FunctionSamplesMap> SortedCallsiteSamples(
 501       S.getCallsiteSamples());
 502   Indent += 1;
 503   for (const auto &I : SortedCallsiteSamples.get())
 504     for (const auto &FS : I->second) {
 505       LineLocation Loc = I->first;
 506       const FunctionSamples &CalleeSamples = FS.second;
 507       OS.indent(Indent);
 508       if (Loc.Discriminator == 0)
 509         OS << Loc.LineOffset << ": ";
 510       else
 511         OS << Loc.LineOffset << "." << Loc.Discriminator << ": ";
 512       if (std::error_code EC = writeSample(CalleeSamples))
 513         return EC;
 514     }
 515   Indent -= 1;
 516
 517   if (FunctionSamples::ProfileIsProbeBased) {
 518     OS.indent(Indent + 1);
 519     OS << "!CFGChecksum: " << S.getFunctionHash() << "\n";
 520   }
 521
 522   if (S.getContext().getAllAttributes()) {
 523     OS.indent(Indent + 1);
 524     OS << "!Attributes: " << S.getContext().getAllAttributes() << "\n";
 525   }
 526
 527   return sampleprof_error::success;
 528 }
 529
 530 std::error_code
 531 SampleProfileWriterBinary::writeContextIdx(const SampleContext &Context) {
 532   assert(!Context.hasContext() && "cs profile is not supported");
 533   return writeNameIdx(Context.getName());
 534 }
 535
 536 std::error_code SampleProfileWriterBinary::writeNameIdx(StringRef FName) {
 537   auto &NTable = getNameTable();
 538   const auto &Ret = NTable.find(FName);
 539   if (Ret == NTable.end())
 540     return sampleprof_error::truncated_name_table;
 541   encodeULEB128(Ret->second, *OutputStream);
 542   return sampleprof_error::success;
 543 }
 544
 545 void SampleProfileWriterBinary::addName(StringRef FName) {
 546   auto &NTable = getNameTable();
 547   NTable.insert(std::make_pair(FName, 0));
 548 }
 549
 550 void SampleProfileWriterBinary::addContext(const SampleContext &Context) {
 551   addName(Context.getName());
 552 }
 553
 554 void SampleProfileWriterBinary::addNames(const FunctionSamples &S) {
 555   // Add all the names in indirect call targets.
 556   for (const auto &I : S.getBodySamples()) {
 557     const SampleRecord &Sample = I.second;
 558     for (const auto &J : Sample.getCallTargets())
 559       addName(J.first());
 560   }
 561
 562   // Recursively add all the names for inlined callsites.
 563   for (const auto &J : S.getCallsiteSamples())
 564     for (const auto &FS : J.second) {
 565       const FunctionSamples &CalleeSamples = FS.second;
 566       addName(CalleeSamples.getName());
 567       addNames(CalleeSamples);
 568     }
 569 }
 570
 571 void SampleProfileWriterExtBinaryBase::addContext(
 572     const SampleContext &Context) {
 573   if (Context.hasContext()) {
 574     for (auto &Callsite : Context.getContextFrames())
 575       SampleProfileWriterBinary::addName(Callsite.FuncName);
 576     CSNameTable.insert(std::make_pair(Context, 0));
 577   } else {
 578     SampleProfileWriterBinary::addName(Context.getName());
 579   }
 580 }
 581
 582 void SampleProfileWriterBinary::stablizeNameTable(
 583     MapVector<StringRef, uint32_t> &NameTable, std::set<StringRef> &V) {
 584   // Sort the names to make NameTable deterministic.
 585   for (const auto &I : NameTable)
 586     V.insert(I.first);
 587   int i = 0;
 588   for (const StringRef &N : V)
 589     NameTable[N] = i++;
 590 }
 591
 592 std::error_code SampleProfileWriterBinary::writeNameTable() {
 593   auto &OS = *OutputStream;
 594   std::set<StringRef> V;
 595   stablizeNameTable(NameTable, V);
 596
 597   // Write out the name table.
 598   encodeULEB128(NameTable.size(), OS);
 599   for (auto N : V) {
 600     OS << N;
 601     encodeULEB128(0, OS);
 602   }
 603   return sampleprof_error::success;
 604 }
 605
 606 std::error_code SampleProfileWriterCompactBinary::writeFuncOffsetTable() {
 607   auto &OS = *OutputStream;
 608
 609   // Fill the slot remembered by TableOffset with the offset of FuncOffsetTable.
 610   auto &OFS = static_cast<raw_fd_ostream &>(OS);
 611   uint64_t FuncOffsetTableStart = OS.tell();
 612   if (OFS.seek(TableOffset) == (uint64_t)-1)
 613     return sampleprof_error::ostream_seek_unsupported;
 614   support::endian::Writer Writer(*OutputStream, support::little);
 615   Writer.write(FuncOffsetTableStart);
 616   if (OFS.seek(FuncOffsetTableStart) == (uint64_t)-1)
 617     return sampleprof_error::ostream_seek_unsupported;
 618
 619   // Write out the table size.
 620   encodeULEB128(FuncOffsetTable.size(), OS);
 621
 622   // Write out FuncOffsetTable.
 623   for (auto Entry : FuncOffsetTable) {
 624     if (std::error_code EC = writeNameIdx(Entry.first))
 625       return EC;
 626     encodeULEB128(Entry.second, OS);
 627   }
 628   return sampleprof_error::success;
 629 }
 630
 631 std::error_code SampleProfileWriterCompactBinary::writeNameTable() {
 632   auto &OS = *OutputStream;
 633   std::set<StringRef> V;
 634   stablizeNameTable(NameTable, V);
 635
 636   // Write out the name table.
 637   encodeULEB128(NameTable.size(), OS);
 638   for (auto N : V) {
 639     encodeULEB128(MD5Hash(N), OS);
 640   }
 641   return sampleprof_error::success;
 642 }
 643
 644 std::error_code
 645 SampleProfileWriterBinary::writeMagicIdent(SampleProfileFormat Format) {
 646   auto &OS = *OutputStream;
 647   // Write file magic identifier.
 648   encodeULEB128(SPMagic(Format), OS);
 649   encodeULEB128(SPVersion(), OS);
 650   return sampleprof_error::success;
 651 }
 652
 653 std::error_code
 654 SampleProfileWriterBinary::writeHeader(const SampleProfileMap &ProfileMap) {
 655   writeMagicIdent(Format);
 656
 657   computeSummary(ProfileMap);
 658   if (auto EC = writeSummary())
 659     return EC;
 660
 661   // Generate the name table for all the functions referenced in the profile.
 662   for (const auto &I : ProfileMap) {
 663     assert(I.first == I.second.getContext() && "Inconsistent profile map");
 664     addContext(I.first);
 665     addNames(I.second);
 666   }
 667
 668   writeNameTable();
 669   return sampleprof_error::success;
 670 }
 671
 672 void SampleProfileWriterExtBinaryBase::setToCompressAllSections() {
 673   for (auto &Entry : SectionHdrLayout)
 674     addSecFlag(Entry, SecCommonFlags::SecFlagCompress);
 675 }
 676
 677 void SampleProfileWriterExtBinaryBase::setToCompressSection(SecType Type) {
 678   addSectionFlag(Type, SecCommonFlags::SecFlagCompress);
 679 }
 680
 681 void SampleProfileWriterExtBinaryBase::allocSecHdrTable() {
 682   support::endian::Writer Writer(*OutputStream, support::little);
 683
 684   Writer.write(static_cast<uint64_t>(SectionHdrLayout.size()));
 685   SecHdrTableOffset = OutputStream->tell();
 686   for (uint32_t i = 0; i < SectionHdrLayout.size(); i++) {
 687     Writer.write(static_cast<uint64_t>(-1));
 688     Writer.write(static_cast<uint64_t>(-1));
 689     Writer.write(static_cast<uint64_t>(-1));
 690     Writer.write(static_cast<uint64_t>(-1));
 691   }
 692 }
 693
 694 std::error_code SampleProfileWriterExtBinaryBase::writeSecHdrTable() {
 695   auto &OFS = static_cast<raw_fd_ostream &>(*OutputStream);
 696   uint64_t Saved = OutputStream->tell();
 697
 698   // Set OutputStream to the location saved in SecHdrTableOffset.
 699   if (OFS.seek(SecHdrTableOffset) == (uint64_t)-1)
 700     return sampleprof_error::ostream_seek_unsupported;
 701   support::endian::Writer Writer(*OutputStream, support::little);
 702
 703   assert(SecHdrTable.size() == SectionHdrLayout.size() &&
 704          "SecHdrTable entries doesn't match SectionHdrLayout");
 705   SmallVector<uint32_t, 16> IndexMap(SecHdrTable.size(), -1);
 706   for (uint32_t TableIdx = 0; TableIdx < SecHdrTable.size(); TableIdx++) {
 707     IndexMap[SecHdrTable[TableIdx].LayoutIndex] = TableIdx;
 708   }
 709
 710   // Write the section header table in the order specified in
 711   // SectionHdrLayout. SectionHdrLayout specifies the sections
 712   // order in which profile reader expect to read, so the section
 713   // header table should be written in the order in SectionHdrLayout.
 714   // Note that the section order in SecHdrTable may be different
 715   // from the order in SectionHdrLayout, for example, SecFuncOffsetTable
 716   // needs to be computed after SecLBRProfile (the order in SecHdrTable),
 717   // but it needs to be read before SecLBRProfile (the order in
 718   // SectionHdrLayout). So we use IndexMap above to switch the order.
 719   for (uint32_t LayoutIdx = 0; LayoutIdx < SectionHdrLayout.size();
 720        LayoutIdx++) {
 721     assert(IndexMap[LayoutIdx] < SecHdrTable.size() &&
 722            "Incorrect LayoutIdx in SecHdrTable");
 723     auto Entry = SecHdrTable[IndexMap[LayoutIdx]];
 724     Writer.write(static_cast<uint64_t>(Entry.Type));
 725     Writer.write(static_cast<uint64_t>(Entry.Flags));
 726     Writer.write(static_cast<uint64_t>(Entry.Offset));
 727     Writer.write(static_cast<uint64_t>(Entry.Size));
 728   }
 729
 730   // Reset OutputStream.
 731   if (OFS.seek(Saved) == (uint64_t)-1)
 732     return sampleprof_error::ostream_seek_unsupported;
 733
 734   return sampleprof_error::success;
 735 }
 736
 737 std::error_code SampleProfileWriterExtBinaryBase::writeHeader(
 738     const SampleProfileMap &ProfileMap) {
 739   auto &OS = *OutputStream;
 740   FileStart = OS.tell();
 741   writeMagicIdent(Format);
 742
 743   allocSecHdrTable();
 744   return sampleprof_error::success;
 745 }
 746
 747 std::error_code SampleProfileWriterCompactBinary::writeHeader(
 748     const SampleProfileMap &ProfileMap) {
 749   support::endian::Writer Writer(*OutputStream, support::little);
 750   if (auto EC = SampleProfileWriterBinary::writeHeader(ProfileMap))
 751     return EC;
 752
 753   // Reserve a slot for the offset of function offset table. The slot will
 754   // be populated with the offset of FuncOffsetTable later.
 755   TableOffset = OutputStream->tell();
 756   Writer.write(static_cast<uint64_t>(-2));
 757   return sampleprof_error::success;
 758 }
 759
 760 std::error_code SampleProfileWriterBinary::writeSummary() {
 761   auto &OS = *OutputStream;
 762   encodeULEB128(Summary->getTotalCount(), OS);
 763   encodeULEB128(Summary->getMaxCount(), OS);
 764   encodeULEB128(Summary->getMaxFunctionCount(), OS);
 765   encodeULEB128(Summary->getNumCounts(), OS);
 766   encodeULEB128(Summary->getNumFunctions(), OS);
 767   const std::vector<ProfileSummaryEntry> &Entries =
 768       Summary->getDetailedSummary();
 769   encodeULEB128(Entries.size(), OS);
 770   for (auto Entry : Entries) {
 771     encodeULEB128(Entry.Cutoff, OS);
 772     encodeULEB128(Entry.MinCount, OS);
 773     encodeULEB128(Entry.NumCounts, OS);
 774   }
 775   return sampleprof_error::success;
 776 }
 777 std::error_code SampleProfileWriterBinary::writeBody(const FunctionSamples &S) {
 778   auto &OS = *OutputStream;
 779   if (std::error_code EC = writeContextIdx(S.getContext()))
 780     return EC;
 781
 782   encodeULEB128(S.getTotalSamples(), OS);
 783
 784   // Emit all the body samples.
 785   encodeULEB128(S.getBodySamples().size(), OS);
 786   for (const auto &I : S.getBodySamples()) {
 787     LineLocation Loc = I.first;
 788     const SampleRecord &Sample = I.second;
 789     encodeULEB128(Loc.LineOffset, OS);
 790     encodeULEB128(Loc.Discriminator, OS);
 791     encodeULEB128(Sample.getSamples(), OS);
 792     encodeULEB128(Sample.getCallTargets().size(), OS);
 793     for (const auto &J : Sample.getSortedCallTargets()) {
 794       StringRef Callee = J.first;
 795       uint64_t CalleeSamples = J.second;
 796       if (std::error_code EC = writeNameIdx(Callee))
 797         return EC;
 798       encodeULEB128(CalleeSamples, OS);
 799     }
 800   }
 801
 802   // Recursively emit all the callsite samples.
 803   uint64_t NumCallsites = 0;
 804   for (const auto &J : S.getCallsiteSamples())
 805     NumCallsites += J.second.size();
 806   encodeULEB128(NumCallsites, OS);
 807   for (const auto &J : S.getCallsiteSamples())
 808     for (const auto &FS : J.second) {
 809       LineLocation Loc = J.first;
 810       const FunctionSamples &CalleeSamples = FS.second;
 811       encodeULEB128(Loc.LineOffset, OS);
 812       encodeULEB128(Loc.Discriminator, OS);
 813       if (std::error_code EC = writeBody(CalleeSamples))
 814         return EC;
 815     }
 816
 817   return sampleprof_error::success;
 818 }
 819
 820 /// Write samples of a top-level function to a binary file.
 821 ///
 822 /// \returns true if the samples were written successfully, false otherwise.
 823 std::error_code
 824 SampleProfileWriterBinary::writeSample(const FunctionSamples &S) {
 825   encodeULEB128(S.getHeadSamples(), *OutputStream);
 826   return writeBody(S);
 827 }
 828
 829 std::error_code
 830 SampleProfileWriterCompactBinary::writeSample(const FunctionSamples &S) {
 831   uint64_t Offset = OutputStream->tell();
 832   StringRef Name = S.getName();
 833   FuncOffsetTable[Name] = Offset;
 834   encodeULEB128(S.getHeadSamples(), *OutputStream);
 835   return writeBody(S);
 836 }
 837
 838 /// Create a sample profile file writer based on the specified format.
 839 ///
 840 /// \param Filename The file to create.
 841 ///
 842 /// \param Format Encoding format for the profile file.
 843 ///
 844 /// \returns an error code indicating the status of the created writer.
 845 ErrorOr<std::unique_ptr<SampleProfileWriter>>
 846 SampleProfileWriter::create(StringRef Filename, SampleProfileFormat Format) {
 847   std::error_code EC;
 848   std::unique_ptr<raw_ostream> OS;
 849   if (Format == SPF_Binary || Format == SPF_Ext_Binary ||
 850       Format == SPF_Compact_Binary)
 851     OS.reset(new raw_fd_ostream(Filename, EC, sys::fs::OF_None));
 852   else
 853     OS.reset(new raw_fd_ostream(Filename, EC, sys::fs::OF_TextWithCRLF));
 854   if (EC)
 855     return EC;
 856
 857   return create(OS, Format);
 858 }
 859
 860 /// Create a sample profile stream writer based on the specified format.
 861 ///
 862 /// \param OS The output stream to store the profile data to.
 863 ///
 864 /// \param Format Encoding format for the profile file.
 865 ///
 866 /// \returns an error code indicating the status of the created writer.
 867 ErrorOr<std::unique_ptr<SampleProfileWriter>>
 868 SampleProfileWriter::create(std::unique_ptr<raw_ostream> &OS,
 869                             SampleProfileFormat Format) {
 870   std::error_code EC;
 871   std::unique_ptr<SampleProfileWriter> Writer;
 872
 873   // Currently only Text and Extended Binary format are supported for CSSPGO.
 874   if ((FunctionSamples::ProfileIsCSFlat ||
 875        FunctionSamples::ProfileIsProbeBased) &&
 876       (Format == SPF_Binary || Format == SPF_Compact_Binary))
 877     return sampleprof_error::unsupported_writing_format;
 878
 879   if (Format == SPF_Binary)
 880     Writer.reset(new SampleProfileWriterRawBinary(OS));
 881   else if (Format == SPF_Ext_Binary)
 882     Writer.reset(new SampleProfileWriterExtBinary(OS));
 883   else if (Format == SPF_Compact_Binary)
 884     Writer.reset(new SampleProfileWriterCompactBinary(OS));
 885   else if (Format == SPF_Text)
 886     Writer.reset(new SampleProfileWriterText(OS));
 887   else if (Format == SPF_GCC)
 888     EC = sampleprof_error::unsupported_writing_format;
 889   else
 890     EC = sampleprof_error::unrecognized_format;
 891
 892   if (EC)
 893     return EC;
 894
 895   Writer->Format = Format;
 896   return std::move(Writer);
 897 }
 898
 899 void SampleProfileWriter::computeSummary(const SampleProfileMap &ProfileMap) {
 900   SampleProfileSummaryBuilder Builder(ProfileSummaryBuilder::DefaultCutoffs);
 901   Summary = Builder.computeSummaryForProfiles(ProfileMap);
 902 }