[SampleFDO] New hierarchical discriminator for FS SampleFDO (llvm-profdata part)

This patch was split from https://reviews.llvm.org/D102246
[SampleFDO] New hierarchical discriminator for Flow Sensitive SampleFDO
This is for llvm-profdata part of change. It sets the bit masks for the
profile reader in llvm-profdata. Also add an internal option
"-fs-discriminator-pass" for show and merge command to process the profile
offline.

This patch also moved setDiscriminatorMaskedBitFrom() to
SampleProfileReader::create() to simplify the interface.

Differential Revision: https://reviews.llvm.org/D103550

include/llvm/ProfileData/SampleProfReader.h

@@ -357,10 +357,6 @@ public:
   void setDiscriminatorMaskedBitFrom(FSDiscriminatorPass P) {
     MaskedBitFrom = getFSPassBitEnd(P);
   }
-  /// Set the bits for using base discriminators.
-  void setBaseDiscriminatorMask() {
-    setDiscriminatorMaskedBitFrom(FSDiscriminatorPass::Base);
-  }
 
   /// Get the bitmask the discriminators: For FS profiles, return the bit
   /// mask for this pass. For non FS profiles, return (unsigned) -1.
@@ -443,14 +439,18 @@ public:
 
   /// Create a sample profile reader appropriate to the file format.
   /// Create a remapper underlying if RemapFilename is not empty.
+  /// Parameter P specifies the FSDiscriminatorPass.
   static ErrorOr<std::unique_ptr<SampleProfileReader>>
   create(const std::string Filename, LLVMContext &C,
+         FSDiscriminatorPass P = FSDiscriminatorPass::Base,
          const std::string RemapFilename = "");
 
   /// Create a sample profile reader from the supplied memory buffer.
   /// Create a remapper underlying if RemapFilename is not empty.
+  /// Parameter P specifies the FSDiscriminatorPass.
   static ErrorOr<std::unique_ptr<SampleProfileReader>>
   create(std::unique_ptr<MemoryBuffer> &B, LLVMContext &C,
+         FSDiscriminatorPass P = FSDiscriminatorPass::Base,
          const std::string RemapFilename = "");
 
   /// Return the profile summary.

include/llvm/Support/Discriminator.h

@@ -54,7 +54,7 @@ static inline unsigned encodingBits(unsigned C) {
 //
 namespace llvm {
 namespace sampleprof {
-enum class FSDiscriminatorPass : unsigned {
+enum FSDiscriminatorPass {
   Base = 0,
   Pass0 = 0,
   Pass1 = 1,

lib/ProfileData/SampleProfReader.cpp

@@ -1620,16 +1620,19 @@ setupMemoryBuffer(const Twine &Filename) {
 ///
 /// \param C The LLVM context to use to emit diagnostics.
 ///
+/// \param P The FSDiscriminatorPass.
+///
 /// \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 std::string Filename, LLVMContext &C,
+                            FSDiscriminatorPass P,
                             const std::string RemapFilename) {
   auto BufferOrError = setupMemoryBuffer(Filename);
   if (std::error_code EC = BufferOrError.getError())
     return EC;
-  return create(BufferOrError.get(), C, RemapFilename);
+  return create(BufferOrError.get(), C, P, RemapFilename);
 }
 
 /// Create a sample profile remapper from the given input, to remap the
@@ -1687,11 +1690,14 @@ SampleProfileReaderItaniumRemapper::create(std::unique_ptr<MemoryBuffer> &B,
 ///
 /// \param C The LLVM context to use to emit diagnostics.
 ///
+/// \param P The FSDiscriminatorPass.
+///
 /// \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,
+                            FSDiscriminatorPass P,
                             const std::string RemapFilename) {
   std::unique_ptr<SampleProfileReader> Reader;
   if (SampleProfileReaderRawBinary::hasFormat(*B))
@@ -1723,6 +1729,8 @@ SampleProfileReader::create(std::unique_ptr<MemoryBuffer> &B, LLVMContext &C,
     return EC;
   }
 
+  Reader->setDiscriminatorMaskedBitFrom(P);
+
   return std::move(Reader);
 }
 

lib/Target/X86/X86InsertPrefetch.cpp

@@ -167,7 +167,6 @@ bool X86InsertPrefetch::doInitialization(Module &M) {
     return false;
   }
   Reader = std::move(ReaderOrErr.get());
-  Reader->setBaseDiscriminatorMask();
   Reader->read();
   return true;
 }

lib/Transforms/IPO/SampleProfile.cpp

@@ -1757,8 +1757,8 @@ bool SampleProfileLoader::doInitialization(Module &M,
                                            FunctionAnalysisManager *FAM) {
   auto &Ctx = M.getContext();
 
-  auto ReaderOrErr =
+  auto ReaderOrErr = SampleProfileReader::create(
-      SampleProfileReader::create(Filename, Ctx, RemappingFilename);
+      Filename, Ctx, FSDiscriminatorPass::Base, RemappingFilename);
   if (std::error_code EC = ReaderOrErr.getError()) {
     std::string Msg = "Could not open profile: " + EC.message();
     Ctx.diagnose(DiagnosticInfoSampleProfile(Filename, Msg));
@@ -1769,7 +1769,6 @@ bool SampleProfileLoader::doInitialization(Module &M,
   // set module before reading the profile so reader may be able to only
   // read the function profiles which are used by the current module.
   Reader->setModule(&M);
-  Reader->setBaseDiscriminatorMask();
   if (std::error_code EC = Reader->read()) {
     std::string Msg = "profile reading failed: " + EC.message();
     Ctx.diagnose(DiagnosticInfoSampleProfile(Filename, Msg));

test/tools/llvm-profdata/Inputs/sample-fs.proftext

@@ -0,0 +1,7 @@
+main:6436:0
+ 4: 534
+ 4.2: 534
+ 4.738209026: 1068
+ 5: 1075
+ 5.1: 1075
+ 5.738209025: 2150

test/tools/llvm-profdata/sample-fs.test

@@ -0,0 +1,54 @@
+Basic tests for sample profiles using fs discriminators.
+
+1- Show command and keep all the discrimiantor bits
+RUN: llvm-profdata show --sample %p/Inputs/sample-fs.proftext -profile-isfs | FileCheck %s --check-prefix=SHOW1
+RUN: llvm-profdata show --sample %p/Inputs/sample-fs.proftext -profile-isfs -fs-discriminator-pass=PassLast | FileCheck %s --check-prefix=SHOW1
+SHOW1: Function: main: 6436, 0, 6 sampled lines
+SHOW1: Samples collected in the function's body {
+SHOW1:   4: 534
+SHOW1:   4.2: 534
+SHOW1:   4.738209026: 1068
+SHOW1:   5: 1075
+SHOW1:   5.1: 1075
+SHOW1:   5.738209025: 2150
+SHOW1: }
+
+2- Show command and keep only the base discriminator bits
+RUN: llvm-profdata show --sample %p/Inputs/sample-fs.proftext -profile-isfs -fs-discriminator-pass=Base | FileCheck %s --check-prefix=SHOW2
+SHOW2: Function: main: 6436, 0, 4 sampled lines
+SHOW2: Samples collected in the function's body {
+SHOW2:   4: 534
+SHOW2:   4.2: 1602
+SHOW2:   5: 1075
+SHOW2:   5.1: 3225
+SHOW2: }
+
+3- Show command and keep only the base discriminator bits and first pass of FS discriminator
+RUN: llvm-profdata show --sample %p/Inputs/sample-fs.proftext -profile-isfs -fs-discriminator-pass=Pass1 | FileCheck %s --check-prefix=SHOW3
+Function: main: 6436, 0, 6 sampled lines
+SHOW3: Samples collected in the function's body {
+SHOW3:   4: 534
+SHOW3:   4.2: 534
+SHOW3:   4.11522: 1068
+SHOW3:   5: 1075
+SHOW3:   5.1: 1075
+SHOW3:   5.11521: 2150
+SHOW3: }
+
+4- Merge command and keep all the discrimiantor bits
+RUN: llvm-profdata merge --sample %p/Inputs/sample-fs.proftext -profile-isfs -fs-discriminator-pass=PassLast --binary -o - | llvm-profdata show --sample - -o %t1-binary_1
+RUN: llvm-profdata merge --sample %p/Inputs/sample-fs.proftext -profile-isfs --binary -o - | llvm-profdata show --sample - -o %t1-binary_2
+RUN: llvm-profdata show --sample %p/Inputs/sample-fs.proftext -profile-isfs -o %t1-text
+RUN: diff %t1-binary_1 %t1-text
+RUN: diff %t1-binary_2 %t1-text
+
+2- Merge command and keep only the base discriminator bits
+RUN: llvm-profdata merge --sample %p/Inputs/sample-fs.proftext -profile-isfs -fs-discriminator-pass=Base --binary -o - | llvm-profdata show --sample - -o %t2-binary
+RUN: llvm-profdata show --sample %p/Inputs/sample-fs.proftext -profile-isfs -fs-discriminator-pass=Base -o %t2-text
+RUN: diff %t2-binary %t2-text
+
+3- Merge command and keep only the base discriminator bits and first pass of FS discriminator
+RUN: llvm-profdata merge --sample %p/Inputs/sample-fs.proftext -profile-isfs -fs-discriminator-pass=Pass1 --binary -o - | llvm-profdata show --sample - -o %t3-binary
+RUN: llvm-profdata show --sample %p/Inputs/sample-fs.proftext -profile-isfs -fs-discriminator-pass=Pass1 -o %t3-text
+RUN: diff %t3-binary %t3-text
+

tools/llvm-profdata/llvm-profdata.cpp

@@ -20,6 +20,7 @@
 #include "llvm/ProfileData/SampleProfReader.h"
 #include "llvm/ProfileData/SampleProfWriter.h"
 #include "llvm/Support/CommandLine.h"
+#include "llvm/Support/Discriminator.h"
 #include "llvm/Support/Errc.h"
 #include "llvm/Support/FileSystem.h"
 #include "llvm/Support/Format.h"
@@ -451,6 +452,25 @@ static void updateInstrProfileEntry(InstrProfileEntry &IFE,
 const uint64_t ColdPercentileIdx = 15;
 const uint64_t HotPercentileIdx = 11;
 
+using sampleprof::FSDiscriminatorPass;
+
+// Internal options to set FSDiscriminatorPass. Used in merge and show
+// commands.
+static cl::opt<FSDiscriminatorPass> FSDiscriminatorPassOption(
+    "fs-discriminator-pass", cl::init(PassLast), cl::Hidden,
+    cl::desc("Zero out the discriminator bits for the FS discrimiantor "
+             "pass beyond this value. The enum values are defined in "
+             "Support/Discriminator.h"),
+    cl::values(clEnumVal(Base, "Use base discriminators only"),
+               clEnumVal(Pass1, "Use base and pass 1 discriminators"),
+               clEnumVal(Pass2, "Use base and pass 1-2 discriminators"),
+               clEnumVal(Pass3, "Use base and pass 1-3 discriminators"),
+               clEnumVal(PassLast, "Use all discriminator bits (default)")));
+
+static unsigned getDiscriminatorMask() {
+  return getN1Bits(getFSPassBitEnd(FSDiscriminatorPassOption.getValue()));
+}
+
 /// Adjust the instr profile in \p WC based on the sample profile in
 /// \p Reader.
 static void
@@ -542,8 +562,8 @@ static void supplementInstrProfile(
 
   // Read sample profile.
   LLVMContext Context;
-  auto ReaderOrErr =
+  auto ReaderOrErr = sampleprof::SampleProfileReader::create(
-      sampleprof::SampleProfileReader::create(SampleFilename.str(), Context);
+      SampleFilename.str(), Context, FSDiscriminatorPassOption);
   if (std::error_code EC = ReaderOrErr.getError())
     exitWithErrorCode(EC, SampleFilename);
   auto Reader = std::move(ReaderOrErr.get());
@@ -574,12 +594,13 @@ remapSamples(const sampleprof::FunctionSamples &Samples,
   Result.addTotalSamples(Samples.getTotalSamples());
   Result.addHeadSamples(Samples.getHeadSamples());
   for (const auto &BodySample : Samples.getBodySamples()) {
-    Result.addBodySamples(BodySample.first.LineOffset,
+    uint32_t MaskedDiscriminator =
-                          BodySample.first.Discriminator,
+        BodySample.first.Discriminator & getDiscriminatorMask();
+    Result.addBodySamples(BodySample.first.LineOffset, MaskedDiscriminator,
                           BodySample.second.getSamples());
     for (const auto &Target : BodySample.second.getCallTargets()) {
       Result.addCalledTargetSamples(BodySample.first.LineOffset,
-                                    BodySample.first.Discriminator,
+                                    MaskedDiscriminator,
                                     Remapper(Target.first()), Target.second);
     }
   }
@@ -677,7 +698,8 @@ mergeSampleProfile(const WeightedFileVector &Inputs, SymbolRemapper *Remapper,
   Optional<bool> ProfileIsProbeBased;
   Optional<bool> ProfileIsCS;
   for (const auto &Input : Inputs) {
-    auto ReaderOrErr = SampleProfileReader::create(Input.Filename, Context);
+    auto ReaderOrErr = SampleProfileReader::create(Input.Filename, Context,
+                                                   FSDiscriminatorPassOption);
     if (std::error_code EC = ReaderOrErr.getError()) {
       warnOrExitGivenError(FailMode, EC, Input.Filename);
       continue;
@@ -907,16 +929,16 @@ static int merge_main(int argc, const char *argv[]) {
                "sample profile, if the ratio of the number of zero counters "
                "divided by the the total number of counters is above the "
                "threshold, the profile of the function will be regarded as "
-               "being harmful for performance and will be dropped. "));
+               "being harmful for performance and will be dropped."));
   cl::opt<unsigned> SupplMinSizeThreshold(
       "suppl-min-size-threshold", cl::init(10), cl::Hidden,
       cl::desc("If the size of a function is smaller than the threshold, "
                "assume it can be inlined by PGO early inliner and it won't "
-               "be adjusted based on sample profile. "));
+               "be adjusted based on sample profile."));
   cl::opt<unsigned> InstrProfColdThreshold(
       "instr-prof-cold-threshold", cl::init(0), cl::Hidden,
       cl::desc("User specified cold threshold for instr profile which will "
-               "override the cold threshold got from profile summary. "));
+               "override the cold threshold got from profile summary."));
 
   cl::ParseCommandLineOptions(argc, argv, "LLVM profile data merger\n");
 
@@ -1859,11 +1881,13 @@ std::error_code SampleOverlapAggregator::loadProfiles() {
   using namespace sampleprof;
 
   LLVMContext Context;
-  auto BaseReaderOrErr = SampleProfileReader::create(BaseFilename, Context);
+  auto BaseReaderOrErr = SampleProfileReader::create(BaseFilename, Context,
+                                                     FSDiscriminatorPassOption);
   if (std::error_code EC = BaseReaderOrErr.getError())
     exitWithErrorCode(EC, BaseFilename);
 
-  auto TestReaderOrErr = SampleProfileReader::create(TestFilename, Context);
+  auto TestReaderOrErr = SampleProfileReader::create(TestFilename, Context,
+                                                     FSDiscriminatorPassOption);
   if (std::error_code EC = TestReaderOrErr.getError())
     exitWithErrorCode(EC, TestFilename);
 
@@ -2375,12 +2399,12 @@ static int showSampleProfile(const std::string &Filename, bool ShowCounts,
                              raw_fd_ostream &OS) {
   using namespace sampleprof;
   LLVMContext Context;
-  auto ReaderOrErr = SampleProfileReader::create(Filename, Context);
+  auto ReaderOrErr =
+      SampleProfileReader::create(Filename, Context, FSDiscriminatorPassOption);
   if (std::error_code EC = ReaderOrErr.getError())
     exitWithErrorCode(EC, Filename);
 
   auto Reader = std::move(ReaderOrErr.get());
-
   if (ShowSectionInfoOnly) {
     showSectionInfo(Reader.get(), OS);
     return 0;

unittests/ProfileData/SampleProfTest.cpp

@@ -58,11 +58,11 @@ struct SampleProfTest : ::testing::Test {
   void readProfile(const Module &M, StringRef Profile,
                    StringRef RemapFile = "") {
     auto ReaderOrErr = SampleProfileReader::create(
-        std::string(Profile), Context, std::string(RemapFile));
+        std::string(Profile), Context, FSDiscriminatorPass::Base,
+        std::string(RemapFile));
     ASSERT_TRUE(NoError(ReaderOrErr.getError()));
     Reader = std::move(ReaderOrErr.get());
     Reader->setModule(&M);
-    Reader->setBaseDiscriminatorMask();
   }
 
   TempFile createRemapFile() {