[SampleFDO] New hierarchical discriminator for FS SampleFDO (ProfileData part)

This patch was split from https://reviews.llvm.org/D102246
[SampleFDO] New hierarchical discriminator for Flow Sensitive SampleFDO
This is mainly for ProfileData part of change. It will load
FS Profile when such profile is detected. For an extbinary format profile,
create_llvm_prof tool will add a flag to profile summary section.
For other format profiles, the users need to use an internal option
(-profile-isfs) to tell the compiler that the profile uses FS discriminators.

This patch also simplified the bit API used by FS discriminators.

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

include/llvm/CodeGen/MIRFSDiscriminator.h

@@ -41,6 +41,7 @@
 
 namespace llvm {
 
+using namespace sampleprof;
 class MIRAddFSDiscriminators : public MachineFunctionPass {
   MachineFunction *MF;
   unsigned LowBit;
@@ -48,10 +49,11 @@ class MIRAddFSDiscriminators : public MachineFunctionPass {
 
 public:
   static char ID;
-  /// FS bits that will be used in this pass (numbers are 0 based and
-  /// inclusive).
-  MIRAddFSDiscriminators(unsigned LowBit = 0, unsigned HighBit = 0)
-      : MachineFunctionPass(ID), LowBit(LowBit), HighBit(HighBit) {
+  /// PassNum is the sequence number this pass is called, start from 1.
+  MIRAddFSDiscriminators(FSDiscriminatorPass P = FSDiscriminatorPass::Pass1)
+      : MachineFunctionPass(ID) {
+    LowBit = getFSPassBitBegin(P);
+    HighBit = getFSPassBitEnd(P);
     assert(LowBit < HighBit && "HighBit needs to be greater than Lowbit");
   }
 

include/llvm/CodeGen/Passes.h

@@ -15,6 +15,7 @@
 #define LLVM_CODEGEN_PASSES_H
 
 #include "llvm/Support/CodeGen.h"
+#include "llvm/Support/Discriminator.h"
 #include <functional>
 #include <string>
 
@@ -490,9 +491,10 @@ namespace llvm {
   /// Create IR Type Promotion pass. \see TypePromotion.cpp
   FunctionPass *createTypePromotionPass();
 
-  /// Add Flow Sensitive Discriminators.
-  FunctionPass *createMIRAddFSDiscriminatorsPass(unsigned LowBit,
-                                                 unsigned HighBit);
+  /// Add Flow Sensitive Discriminators. PassNum specifies the
+  /// sequence number of this pass (starting from 1).
+  FunctionPass *
+  createMIRAddFSDiscriminatorsPass(sampleprof::FSDiscriminatorPass P);
 
   /// Creates MIR Debugify pass. \see MachineDebugify.cpp
   ModulePass *createDebugifyMachineModulePass();

include/llvm/IR/DebugInfoMetadata.h

@@ -1741,7 +1741,7 @@ public:
   }
 
   /// Return the bits used for base discriminators.
-  static unsigned getBaseDiscriminatorBits() { return BASE_DIS_BIT_END; }
+  static unsigned getBaseDiscriminatorBits() { return getBaseFSBitEnd(); }
 
   /// Returns the base discriminator for a given encoded discriminator \p D.
   static unsigned getBaseDiscriminatorFromDiscriminator(unsigned D) {

include/llvm/ProfileData/SampleProf.h

@@ -190,7 +190,10 @@ enum class SecProfSummaryFlags : uint32_t {
   SecFlagPartial = (1 << 0),
   /// SecFlagContext means this is context-sensitive profile for
   /// CSSPGO
-  SecFlagFullContext = (1 << 1)
+  SecFlagFullContext = (1 << 1),
+  /// SecFlagFSDiscriminator means this profile uses flow-sensitive
+  /// discriminators.
+  SecFlagFSDiscriminator = (1 << 2)
 };
 
 enum class SecFuncMetadataFlags : uint32_t {
@@ -891,6 +894,9 @@ public:
   /// Whether the profile contains any ".__uniq." suffix in a name.
   static bool HasUniqSuffix;
 
+  /// If this profile uses flow sensitive discriminators.
+  static bool ProfileIsFS;
+
   /// GUIDToFuncNameMap saves the mapping from GUID to the symbol name, for
   /// all the function symbols defined or declared in current module.
   DenseMap<uint64_t, StringRef> *GUIDToFuncNameMap = nullptr;

include/llvm/ProfileData/SampleProfReader.h

@@ -236,6 +236,7 @@
 #include "llvm/ProfileData/GCOV.h"
 #include "llvm/ProfileData/SampleProf.h"
 #include "llvm/Support/Debug.h"
+#include "llvm/Support/Discriminator.h"
 #include "llvm/Support/ErrorOr.h"
 #include "llvm/Support/MemoryBuffer.h"
 #include "llvm/Support/SymbolRemappingReader.h"
@@ -350,6 +351,26 @@ public:
   /// Read and validate the file header.
   virtual std::error_code readHeader() = 0;
 
+  /// Set the bits for FS discriminators. Parameter Pass specify the sequence
+  /// number, Pass == i is for the i-th round of adding FS discriminators.
+  /// Pass == 0 is for using base discriminators.
+  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.
+  uint32_t getDiscriminatorMask() const {
+    if (!ProfileIsFS)
+      return 0xFFFFFFFF;
+    assert((MaskedBitFrom != 0) && "MaskedBitFrom is not set properly");
+    return getN1Bits(MaskedBitFrom);
+  }
+
   /// The interface to read sample profiles from the associated file.
   std::error_code read() {
     if (std::error_code EC = readImpl())
@@ -505,6 +526,9 @@ protected:
   /// Number of context-sensitive profiles.
   uint32_t CSProfileCount = 0;
 
+  /// Whether the function profiles use FS discriminators.
+  bool ProfileIsFS = false;
+
   /// \brief The format of sample.
   SampleProfileFormat Format = SPF_None;
 
@@ -512,6 +536,10 @@ protected:
   /// is used by compiler. If SampleProfileReader is used by other
   /// tools which are not compiler, M is usually nullptr.
   const Module *M = nullptr;
+
+  /// Zero out the discriminator bits higher than bit MaskedBitFrom (0 based).
+  /// The default is to keep all the bits.
+  uint32_t MaskedBitFrom = 31;
 };
 
 class SampleProfileReaderText : public SampleProfileReader {

include/llvm/Support/Discriminator.h

@@ -13,6 +13,9 @@
 #ifndef LLVM_SUPPORT_DISCRIMINATOR_H
 #define LLVM_SUPPORT_DISCRIMINATOR_H
 
+#include "llvm/Support/Error.h"
+#include <assert.h>
+
 // Utility functions for encoding / decoding discriminators.
 /// With a given unsigned int \p U, use up to 13 bits to represent it.
 /// old_bit 1~5  --> new_bit 1~5
@@ -48,26 +51,83 @@ static inline unsigned encodingBits(unsigned C) {
 }
 
 // Some constants used in FS Discriminators.
-#define BASE_DIS_BIT_BEG 0
-#define BASE_DIS_BIT_END 7
+//
+namespace llvm {
+namespace sampleprof {
+enum class FSDiscriminatorPass : unsigned {
+  Base = 0,
+  Pass0 = 0,
+  Pass1 = 1,
+  Pass2 = 2,
+  Pass3 = 3,
+  Pass4 = 4,
+  PassLast = 4,
+};
+} // namespace sampleprof
 
-#define PASS_1_DIS_BIT_BEG 8
-#define PASS_1_DIS_BIT_END 13
+using namespace sampleprof;
 
-#define PASS_2_DIS_BIT_BEG 14
-#define PASS_2_DIS_BIT_END 19
+// The number of bits reserved for the base discrimininator. The base
+// discriminaitor starts from bit 0.
+static const unsigned BaseDiscriminatorBitWidth = 8;
 
-#define PASS_3_DIS_BIT_BEG 20
-#define PASS_3_DIS_BIT_END 25
+// The number of bits reserved for each FS discriminator pass.
+static const unsigned FSDiscriminatorBitWidth = 6;
 
-#define PASS_LAST_DIS_BIT_BEG 26
-#define PASS_LAST_DIS_BIT_END 31
-
-// Set bits range [0 .. n] to 1. Used in FS Discriminators.
-static inline unsigned getN1Bits(int N) {
-  if (N >= 31)
-    return 0xFFFFFFFF;
-  return (1 << (N + 1)) - 1;
+// Return the number of FS passes, excluding the pass adding the base
+// discriminators.
+// The number of passes for FS discriminators. Note that the total
+// number of discriminaitor bits, i.e.
+// BaseDiscriminatorBitWidth
+//  + FSDiscriminatorBitWidth * getNumFSPasses()
+// needs to fit in an unsigned int type.
+static inline unsigned getNumFSPasses() {
+  return static_cast<unsigned>(FSDiscriminatorPass::PassLast);
 }
 
+// Return the ending bit for FSPass P.
+static inline unsigned getFSPassBitEnd(FSDiscriminatorPass P) {
+  unsigned I = static_cast<unsigned>(P);
+  assert(I <= getNumFSPasses() && "Invalid FS discriminator pass number.");
+  return BaseDiscriminatorBitWidth + I * FSDiscriminatorBitWidth - 1;
+}
+
+// Return the begining bit for FSPass P.
+static inline unsigned getFSPassBitBegin(FSDiscriminatorPass P) {
+  if (P == FSDiscriminatorPass::Base)
+    return 0;
+  unsigned I = static_cast<unsigned>(P);
+  assert(I <= getNumFSPasses() && "Invalid FS discriminator pass number.");
+  return getFSPassBitEnd(static_cast<FSDiscriminatorPass>(I - 1)) + 1;
+}
+
+// Return the beginning bit for the last FSPass.
+static inline int getLastFSPassBitBegin() {
+  return getFSPassBitBegin(static_cast<FSDiscriminatorPass>(getNumFSPasses()));
+}
+
+// Return the ending bit for the last FSPass.
+static inline unsigned getLastFSPassBitEnd() {
+  return getFSPassBitEnd(static_cast<FSDiscriminatorPass>(getNumFSPasses()));
+}
+
+// Return the beginning bit for the base (first) FSPass.
+static inline unsigned getBaseFSBitBegin() { return 0; }
+
+// Return the ending bit for the base (first) FSPass.
+static inline unsigned getBaseFSBitEnd() {
+  return BaseDiscriminatorBitWidth - 1;
+}
+
+// Set bits in range of [0 .. n] to 1. Used in FS Discriminators.
+static inline unsigned getN1Bits(int N) {
+  // Work around the g++ bug that folding "(1U << (N + 1)) - 1" to 0.
+  if (N == 31)
+    return 0xFFFFFFFF;
+  assert((N < 32) && "N is invalid");
+  return (1U << (N + 1)) - 1;
+}
+
+} // namespace llvm
+
 #endif /* LLVM_SUPPORT_DISCRIMINATOR_H */

lib/CodeGen/MIRFSDiscriminator.cpp

@@ -22,6 +22,7 @@
 #include <unordered_map>
 
 using namespace llvm;
+using namespace sampleprof;
 
 #define DEBUG_TYPE "mirfs-discriminators"
 
@@ -33,9 +34,8 @@ INITIALIZE_PASS(MIRAddFSDiscriminators, DEBUG_TYPE,
 
 char &llvm::MIRAddFSDiscriminatorsID = MIRAddFSDiscriminators::ID;
 
-FunctionPass *llvm::createMIRAddFSDiscriminatorsPass(unsigned LowBit,
-                                                     unsigned HighBit) {
-  return new MIRAddFSDiscriminators(LowBit, HighBit);
+FunctionPass *llvm::createMIRAddFSDiscriminatorsPass(FSDiscriminatorPass P) {
+  return new MIRAddFSDiscriminators(P);
 }
 
 // Compute a hash value using debug line number, and the line numbers from the

lib/CodeGen/TargetPassConfig.cpp

@@ -1177,8 +1177,12 @@ void TargetPassConfig::addMachinePasses() {
   addPass(&PatchableFunctionID);
 
   if (EnableFSDiscriminator && !FSNoFinalDiscrim)
-    addPass(createMIRAddFSDiscriminatorsPass(PASS_LAST_DIS_BIT_BEG,
-                                             PASS_LAST_DIS_BIT_END));
+    // Add FS discriminators here so that all the instruction duplicates
+    // in different BBs get their own discriminators. With this, we can "sum"
+    // the SampleFDO counters instead of using MAX. This will improve the
+    // SampleFDO profile quality.
+    addPass(createMIRAddFSDiscriminatorsPass(
+        sampleprof::FSDiscriminatorPass::PassLast));
 
   addPreEmitPass();
 

lib/ProfileData/SampleProf.cpp

@@ -42,6 +42,7 @@ bool FunctionSamples::ProfileIsProbeBased = false;
 bool FunctionSamples::ProfileIsCS = false;
 bool FunctionSamples::UseMD5 = false;
 bool FunctionSamples::HasUniqSuffix = true;
+bool FunctionSamples::ProfileIsFS = false;
 } // namespace sampleprof
 } // namespace llvm
 
@@ -232,9 +233,15 @@ const FunctionSamples *FunctionSamples::findFunctionSamples(
 
   const DILocation *PrevDIL = DIL;
   for (DIL = DIL->getInlinedAt(); DIL; DIL = DIL->getInlinedAt()) {
-    S.push_back(std::make_pair(
-        LineLocation(getOffset(DIL), DIL->getBaseDiscriminator()),
-        PrevDIL->getScope()->getSubprogram()->getLinkageName()));
+    unsigned Discriminator;
+    if (ProfileIsFS)
+      Discriminator = DIL->getDiscriminator();
+    else
+      Discriminator = DIL->getBaseDiscriminator();
+
+    S.push_back(
+        std::make_pair(LineLocation(getOffset(DIL), Discriminator),
+                       PrevDIL->getScope()->getSubprogram()->getLinkageName()));
     PrevDIL = DIL;
   }
   if (S.size() == 0)

lib/ProfileData/SampleProfReader.cpp

@@ -26,6 +26,7 @@
 #include "llvm/IR/ProfileSummary.h"
 #include "llvm/ProfileData/ProfileCommon.h"
 #include "llvm/ProfileData/SampleProf.h"
+#include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Compression.h"
 #include "llvm/Support/ErrorOr.h"
 #include "llvm/Support/LEB128.h"
@@ -45,6 +46,15 @@
 using namespace llvm;
 using namespace sampleprof;
 
+#define DEBUG_TYPE "samplepgo-reader"
+
+// This internal option specifies if the profile uses FS discriminators.
+// It only applies to text, binary and compact binary format profiles.
+// For ext-binary format profiles, the flag is set in the summary.
+static cl::opt<bool> ProfileIsFSDisciminator(
+    "profile-isfs", cl::Hidden, cl::init(false),
+    cl::desc("Profile uses flow senstive discriminators"));
+
 /// Dump the function profile for \p FName.
 ///
 /// \param FName Name of the function to print.
@@ -238,6 +248,7 @@ std::error_code SampleProfileReaderText::readImpl() {
   // top-level function profile.
   bool SeenMetadata = false;
 
+  ProfileIsFS = ProfileIsFSDisciminator;
   for (; !LineIt.is_at_eof(); ++LineIt) {
     if ((*LineIt)[(*LineIt).find_first_not_of(' ')] == '#')
       continue;
@@ -295,6 +306,10 @@ std::error_code SampleProfileReaderText::readImpl() {
                     "Found non-metadata after metadata: " + *LineIt);
         return sampleprof_error::malformed;
       }
+
+      // Here we handle FS discriminators.
+      Discriminator &= getDiscriminatorMask();
+
       while (InlineStack.size() > Depth) {
         InlineStack.pop_back();
       }
@@ -504,6 +519,9 @@ SampleProfileReaderBinary::readProfile(FunctionSamples &FProfile) {
     if (std::error_code EC = NumCalls.getError())
       return EC;
 
+    // Here we handle FS discriminators:
+    uint32_t DiscriminatorVal = (*Discriminator) & getDiscriminatorMask();
+
     for (uint32_t J = 0; J < *NumCalls; ++J) {
       auto CalledFunction(readStringFromTable());
       if (std::error_code EC = CalledFunction.getError())
@@ -513,11 +531,11 @@ SampleProfileReaderBinary::readProfile(FunctionSamples &FProfile) {
       if (std::error_code EC = CalledFunctionSamples.getError())
         return EC;
 
-      FProfile.addCalledTargetSamples(*LineOffset, *Discriminator,
+      FProfile.addCalledTargetSamples(*LineOffset, DiscriminatorVal,
                                       *CalledFunction, *CalledFunctionSamples);
     }
 
-    FProfile.addBodySamples(*LineOffset, *Discriminator, *NumSamples);
+    FProfile.addBodySamples(*LineOffset, DiscriminatorVal, *NumSamples);
   }
 
   // Read all the samples for inlined function calls.
@@ -538,8 +556,11 @@ SampleProfileReaderBinary::readProfile(FunctionSamples &FProfile) {
     if (std::error_code EC = FName.getError())
       return EC;
 
+    // Here we handle FS discriminators:
+    uint32_t DiscriminatorVal = (*Discriminator) & getDiscriminatorMask();
+
     FunctionSamples &CalleeProfile = FProfile.functionSamplesAt(
-        LineLocation(*LineOffset, *Discriminator))[std::string(*FName)];
+        LineLocation(*LineOffset, DiscriminatorVal))[std::string(*FName)];
     CalleeProfile.setName(*FName);
     if (std::error_code EC = readProfile(CalleeProfile))
       return EC;
@@ -575,6 +596,7 @@ SampleProfileReaderBinary::readFuncProfile(const uint8_t *Start) {
 }
 
 std::error_code SampleProfileReaderBinary::readImpl() {
+  ProfileIsFS = ProfileIsFSDisciminator;
   while (!at_eof()) {
     if (std::error_code EC = readFuncProfile(Data))
       return EC;
@@ -595,6 +617,8 @@ std::error_code SampleProfileReaderExtBinaryBase::readOneSection(
       Summary->setPartialProfile(true);
     if (hasSecFlag(Entry, SecProfSummaryFlags::SecFlagFullContext))
       FunctionSamples::ProfileIsCS = ProfileIsCS = true;
+    if (hasSecFlag(Entry, SecProfSummaryFlags::SecFlagFSDiscriminator))
+      FunctionSamples::ProfileIsFS = ProfileIsFS = true;
     break;
   case SecNameTable: {
     FixedLengthMD5 =
@@ -860,7 +884,7 @@ std::error_code SampleProfileReaderCompactBinary::readImpl() {
   // Collect functions used by current module if the Reader has been
   // given a module.
   bool LoadFuncsToBeUsed = collectFuncsFromModule();
-
+  ProfileIsFS = ProfileIsFSDisciminator;
   std::vector<uint64_t> OffsetsToUse;
   if (!LoadFuncsToBeUsed) {
     // load all the function profiles.
@@ -1105,6 +1129,8 @@ static std::string getSecFlagsStr(const SecHdrTableEntry &Entry) {
       Flags.append("partial,");
     if (hasSecFlag(Entry, SecProfSummaryFlags::SecFlagFullContext))
       Flags.append("context,");
+    if (hasSecFlag(Entry, SecProfSummaryFlags::SecFlagFSDiscriminator))
+      Flags.append("fs-discriminator,");
     break;
   default:
     break;
@@ -1521,6 +1547,7 @@ std::error_code SampleProfileReaderGCC::readOneFunctionProfile(
 /// This format is generated by the Linux Perf conversion tool at
 /// https://github.com/google/autofdo.
 std::error_code SampleProfileReaderGCC::readImpl() {
+  assert(!ProfileIsFSDisciminator && "Gcc profiles not support FSDisciminator");
   // Read the string table.
   if (std::error_code EC = readNameTable())
     return EC;

lib/ProfileData/SampleProfWriter.cpp

@@ -252,6 +252,8 @@ std::error_code SampleProfileWriterExtBinaryBase::writeOneSection(
     addSectionFlag(SecProfSummary, SecProfSummaryFlags::SecFlagFullContext);
   if (Type == SecFuncMetadata && FunctionSamples::ProfileIsCS)
     addSectionFlag(SecFuncMetadata, SecFuncMetadataFlags::SecFlagHasAttribute);
+  if (Type == SecProfSummary && FunctionSamples::ProfileIsFS)
+    addSectionFlag(SecProfSummary, SecProfSummaryFlags::SecFlagFSDiscriminator);
 
   uint64_t SectionStart = markSectionStart(Type, LayoutIdx);
   switch (Type) {

lib/Target/X86/X86InsertPrefetch.cpp

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

lib/Transforms/IPO/SampleProfile.cpp

@@ -1769,6 +1769,7 @@ 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/Transforms/SampleProfile/Inputs/fsafdo.prof

@@ -0,0 +1,35 @@
+work:33383580:1068858
+ 1: 981870
+ 5: 981870
+foo:22388581:3449
+ 0: 3449
+ 2.1: 204820
+ 4: 213086 bar:205247
+ 4.2013265920: 222893 bar:218378
+ 4.2281701376: 214552 bar:217479
+ 4.2550136832: 210692 bar:220056
+ 5: 213086
+ 5.1207959552: 210692
+ 5.1610612736: 202301
+ 5.2952790016: 222893
+ 6: 4780
+ 6.268435456: 202301 work:198259
+ 6.1073741824: 222893 work:231680
+ 6.2147483648: 4780
+ 7: 219065
+ 7.134217728: 217053
+ 7.2013265920: 183304
+ 7.3758096384: 222101
+ 8: 4780
+ 8.2818572288: 222101 work:238765
+ 8.3489660928: 183304 work:181615
+ 8.4160749568: 217053 work:218539
+ 10: 3281
+bar:7622325:861160
+ 2: 846925
+ 3: 846925
+main:16419:0
+ 0: 0
+ 2.1: 3280
+ 3: 3299 foo:3449
+ 5: 0

test/Transforms/SampleProfile/fsafdo_test.ll

@@ -0,0 +1,230 @@
+; RUN: opt < %s -sample-profile -sample-profile-file=%S/Inputs/fsafdo.extbinary.afdo | opt -analyze -branch-prob -enable-new-pm=0 | FileCheck %s
+; RUN: opt < %s -sample-profile -profile-isfs -sample-profile-file=%S/Inputs/fsafdo.prof | opt -analyze -branch-prob -enable-new-pm=0 | FileCheck %s
+
+target triple = "x86_64-unknown-linux-gnu"
+
+@sum = dso_local local_unnamed_addr global i32 0, align 4
+
+declare i32 @bar(i32 %i) #0
+declare void @work(i32 %i) #2
+
+define dso_local void @foo() #0 !dbg !29 {
+; CHECK: Printing analysis {{.*}} for function 'foo':
+
+entry:
+  br label %for.cond1.preheader, !dbg !30
+; CHECK: edge entry -> for.cond1.preheader probability is 0x80000000 / 0x80000000 = 100.00% [HOT edge]
+
+for.cond1.preheader:
+  %j.012 = phi i32 [ 0, %entry ], [ %inc11, %if.end9.3 ]
+  %mul = mul nuw nsw i32 %j.012, 48
+  %call = tail call i32 @bar(i32 %mul), !dbg !32
+  %0 = and i32 %call, 1, !dbg !33
+  %tobool.not = icmp eq i32 %0, 0, !dbg !33
+  br i1 %tobool.not, label %if.end, label %if.then, !dbg !35
+; CHECK: edge for.cond1.preheader -> if.end probability is 0x3f6262b8 / 0x80000000 = 49.52%
+; CHECK: edge for.cond1.preheader -> if.then probability is 0x409d9d48 / 0x80000000 = 50.48%
+
+
+if.then:
+  %mul4 = shl nsw i32 %call, 1, !dbg !36
+  tail call void @work(i32 %mul4), !dbg !37
+  br label %if.end, !dbg !38
+; CHECK: edge if.then -> if.end probability is 0x80000000 / 0x80000000 = 100.00% [HOT edge]
+
+if.end:
+  %1 = and i32 %call, 3, !dbg !39
+  %tobool6.not = icmp eq i32 %1, 0, !dbg !39
+  br i1 %tobool6.not, label %if.end9, label %if.then7, !dbg !40
+; CHECK: edge if.end -> if.end9 probability is 0x22c6bac3 / 0x80000000 = 27.17%
+; CHECK: edge if.end -> if.then7 probability is 0x5d39453d / 0x80000000 = 72.83%
+
+
+if.then7:
+  %mul8 = mul nsw i32 %call, 3, !dbg !41
+  tail call void @work(i32 %mul8), !dbg !42
+  br label %if.end9, !dbg !43
+; CHECK: edge if.then7 -> if.end9 probability is 0x80000000 / 0x80000000 = 100.00% [HOT edge]
+
+if.end9:
+  %add.1 = or i32 %mul, 1, !dbg !44
+  %call.1 = tail call i32 @bar(i32 %add.1), !dbg !32
+  %2 = and i32 %call.1, 1, !dbg !33
+  %tobool.not.1 = icmp eq i32 %2, 0, !dbg !33
+  br i1 %tobool.not.1, label %if.end.1, label %if.then.1, !dbg !35
+; CHECK: edge if.end9 -> if.end.1 probability is 0x3f6262b8 / 0x80000000 = 49.52%
+; CHECK: edge if.end9 -> if.then.1 probability is 0x409d9d48 / 0x80000000 = 50.48%
+
+for.end12:
+  ret void, !dbg !45
+
+if.then.1:
+  %mul4.1 = shl nsw i32 %call.1, 1, !dbg !36
+  tail call void @work(i32 %mul4.1), !dbg !37
+  br label %if.end.1, !dbg !38
+; CHECK: edge if.then.1 -> if.end.1 probability is 0x80000000 / 0x80000000 = 100.00% [HOT edge]
+
+if.end.1:
+  %3 = and i32 %call.1, 3, !dbg !39
+  %tobool6.not.1 = icmp eq i32 %3, 0, !dbg !39
+  br i1 %tobool6.not.1, label %if.end9.1, label %if.then7.1, !dbg !40
+; CHECK: edge if.end.1 -> if.end9.1 probability is 0x22c6bac3 / 0x80000000 = 27.17%
+; CHECK: edge if.end.1 -> if.then7.1 probability is 0x5d39453d / 0x80000000 = 72.83%
+
+if.then7.1:
+  %mul8.1 = mul nsw i32 %call.1, 3, !dbg !41
+  tail call void @work(i32 %mul8.1), !dbg !42
+  br label %if.end9.1, !dbg !43
+; CHECK: edge if.then7.1 -> if.end9.1 probability is 0x80000000 / 0x80000000 = 100.00% [HOT edge]
+
+if.end9.1:
+  %add.2 = or i32 %mul, 2, !dbg !44
+  %call.2 = tail call i32 @bar(i32 %add.2), !dbg !32
+  %4 = and i32 %call.2, 1, !dbg !33
+  %tobool.not.2 = icmp eq i32 %4, 0, !dbg !33
+  br i1 %tobool.not.2, label %if.end.2, label %if.then.2, !dbg !35
+; CHECK: edge if.end9.1 -> if.end.2 probability is 0x3f6262b8 / 0x80000000 = 49.52%
+; CHECK: edge if.end9.1 -> if.then.2 probability is 0x409d9d48 / 0x80000000 = 50.48%
+
+if.then.2:
+  %mul4.2 = shl nsw i32 %call.2, 1, !dbg !36
+  tail call void @work(i32 %mul4.2), !dbg !37
+  br label %if.end.2, !dbg !38
+; CHECK: edge if.then.2 -> if.end.2 probability is 0x80000000 / 0x80000000 = 100.00% [HOT edge]
+
+if.end.2:
+  %5 = and i32 %call.2, 3, !dbg !39
+  %tobool6.not.2 = icmp eq i32 %5, 0, !dbg !39
+  br i1 %tobool6.not.2, label %if.end9.2, label %if.then7.2, !dbg !40
+; CHECK: edge if.end.2 -> if.end9.2 probability is 0x22c6bac3 / 0x80000000 = 27.17%
+; CHECK: edge if.end.2 -> if.then7.2 probability is 0x5d39453d / 0x80000000 = 72.83%
+
+if.then7.2:
+  %mul8.2 = mul nsw i32 %call.2, 3, !dbg !41
+  tail call void @work(i32 %mul8.2), !dbg !42
+  br label %if.end9.2, !dbg !43
+; CHECK: edge if.then7.2 -> if.end9.2 probability is 0x80000000 / 0x80000000 = 100.00% [HOT edge]
+
+if.end9.2:
+  %add.3 = or i32 %mul, 3, !dbg !44
+  %call.3 = tail call i32 @bar(i32 %add.3), !dbg !32
+  %6 = and i32 %call.3, 1, !dbg !33
+  %tobool.not.3 = icmp eq i32 %6, 0, !dbg !33
+  br i1 %tobool.not.3, label %if.end.3, label %if.then.3, !dbg !35
+; CHECK: edge if.end9.2 -> if.end.3 probability is 0x3f6262b8 / 0x80000000 = 49.52%
+; CHECK: edge if.end9.2 -> if.then.3 probability is 0x409d9d48 / 0x80000000 = 50.48%
+
+if.then.3:
+  %mul4.3 = shl nsw i32 %call.3, 1, !dbg !36
+  tail call void @work(i32 %mul4.3), !dbg !37
+  br label %if.end.3, !dbg !38
+; CHECK: edge if.then.3 -> if.end.3 probability is 0x80000000 / 0x80000000 = 100.00% [HOT edge]
+
+if.end.3:
+  %7 = and i32 %call.3, 3, !dbg !39
+  %tobool6.not.3 = icmp eq i32 %7, 0, !dbg !39
+  br i1 %tobool6.not.3, label %if.end9.3, label %if.then7.3, !dbg !40
+; CHECK: edge if.end.3 -> if.end9.3 probability is 0x22c6bac3 / 0x80000000 = 27.17%
+; CHECK: edge if.end.3 -> if.then7.3 probability is 0x5d39453d / 0x80000000 = 72.83%
+
+if.then7.3:
+  %mul8.3 = mul nsw i32 %call.3, 3, !dbg !41
+  tail call void @work(i32 %mul8.3), !dbg !42
+  br label %if.end9.3, !dbg !43
+; CHECK: edge if.then7.3 -> if.end9.3 probability is 0x80000000 / 0x80000000 = 100.00% [HOT edge]
+
+if.end9.3:
+  %inc11 = add nuw nsw i32 %j.012, 1, !dbg !46
+  %exitcond.not = icmp eq i32 %inc11, 48, !dbg !48
+  br i1 %exitcond.not, label %for.end12, label %for.cond1.preheader, !dbg !30, !llvm.loop !49
+; CHECK: edge if.end9.3 -> for.end12 probability is 0x00834dd9 / 0x80000000 = 0.40%
+; CHECK: edge if.end9.3 -> for.cond1.preheader probability is 0x7f7cb227 / 0x80000000 = 99.60% [HOT edge]
+}
+
+define dso_local i32 @main() #3 !dbg !52 {
+entry:
+  br label %for.body, !dbg !53
+
+for.body:
+  %i.03 = phi i32 [ 0, %entry ], [ %inc, %for.body ]
+  tail call void @foo(), !dbg !55
+  %inc = add nuw nsw i32 %i.03, 1, !dbg !56
+  %exitcond.not = icmp eq i32 %inc, 10000000, !dbg !58
+  br i1 %exitcond.not, label %for.end, label %for.body, !dbg !53, !llvm.loop !60
+
+for.end:
+  ret i32 0, !dbg !63
+}
+
+
+attributes #0 = { noinline nounwind uwtable "disable-tail-calls"="false" "frame-pointer"="none" "less-precise-fpmad"="false" "min-legal-vector-width"="0" "no-infs-fp-math"="false" "no-jump-tables"="false" "no-nans-fp-math"="false" "no-signed-zeros-fp-math"="false" "no-trapping-math"="true" "stack-protector-buffer-size"="8" "target-cpu"="x86-64" "target-features"="+cx8,+fxsr,+mmx,+sse,+sse2,+x87" "tune-cpu"="generic" "unsafe-fp-math"="false" "use-soft-float"="false" "use-sample-profile"}
+attributes #1 = { argmemonly nounwind willreturn }
+attributes #2 = { nofree noinline norecurse nounwind uwtable "disable-tail-calls"="false" "frame-pointer"="none" "less-precise-fpmad"="false" "min-legal-vector-width"="0" "no-infs-fp-math"="false" "no-jump-tables"="false" "no-nans-fp-math"="false" "no-signed-zeros-fp-math"="false" "no-trapping-math"="true" "stack-protector-buffer-size"="8" "target-cpu"="x86-64" "target-features"="+cx8,+fxsr,+mmx,+sse,+sse2,+x87" "tune-cpu"="generic" "unsafe-fp-math"="false" "use-soft-float"="false" }
+attributes #3 = { nounwind uwtable "disable-tail-calls"="false" "frame-pointer"="none" "less-precise-fpmad"="false" "min-legal-vector-width"="0" "no-infs-fp-math"="false" "no-jump-tables"="false" "no-nans-fp-math"="false" "no-signed-zeros-fp-math"="false" "no-trapping-math"="true" "stack-protector-buffer-size"="8" "target-cpu"="x86-64" "target-features"="+cx8,+fxsr,+mmx,+sse,+sse2,+x87" "tune-cpu"="generic" "unsafe-fp-math"="false" "use-soft-float"="false" }
+
+!llvm.dbg.cu = !{!0}
+!llvm.module.flags = !{!3, !4, !5}
+
+!0 = distinct !DICompileUnit(language: DW_LANG_C99, file: !1, isOptimized: true, runtimeVersion: 0, emissionKind: LineTablesOnly, enums: !2, splitDebugInlining: false, debugInfoForProfiling: true, nameTableKind: None)
+!1 = !DIFile(filename: "unroll.c", directory: "a/")
+!2 = !{}
+!3 = !{i32 7, !"Dwarf Version", i32 4}
+!4 = !{i32 2, !"Debug Info Version", i32 3}
+!5 = !{i32 1, !"wchar_size", i32 4}
+!7 = distinct !DISubprogram(name: "bar", scope: !1, file: !1, line: 3, type: !8, scopeLine: 3, flags: DIFlagPrototyped | DIFlagAllCallsDescribed, spFlags: DISPFlagDefinition | DISPFlagOptimized, unit: !0, retainedNodes: !2)
+!8 = !DISubroutineType(types: !2)
+!9 = !DILocation(line: 4, column: 3, scope: !7)
+!10 = !DILocation(line: 5, column: 5, scope: !7)
+!11 = !{!12, !12, i64 0}
+!12 = !{!"int", !13, i64 0}
+!13 = !{!"omnipotent char", !14, i64 0}
+!14 = !{!"Simple C/C++ TBAA"}
+!15 = !DILocation(line: 6, column: 10, scope: !7)
+!16 = !DILocation(line: 7, column: 1, scope: !7)
+!17 = !DILocation(line: 6, column: 3, scope: !18)
+!18 = !DILexicalBlockFile(scope: !7, file: !1, discriminator: 1)
+!19 = distinct !DISubprogram(name: "work", scope: !1, file: !1, line: 10, type: !8, scopeLine: 10, flags: DIFlagPrototyped | DIFlagAllCallsDescribed, spFlags: DISPFlagDefinition | DISPFlagOptimized, unit: !0, retainedNodes: !2)
+!20 = !DILocation(line: 11, column: 7, scope: !19)
+!21 = !DILocation(line: 11, column: 11, scope: !22)
+!22 = !DILexicalBlockFile(scope: !19, file: !1, discriminator: 1)
+!23 = !DILocation(line: 11, column: 11, scope: !24)
+!24 = !DILexicalBlockFile(scope: !19, file: !1, discriminator: 2)
+!25 = !DILocation(line: 11, column: 7, scope: !26)
+!26 = !DILexicalBlockFile(scope: !19, file: !1, discriminator: 3)
+!27 = !DILocation(line: 0, scope: !22)
+!28 = !DILocation(line: 15, column: 1, scope: !19)
+!29 = distinct !DISubprogram(name: "foo", scope: !1, file: !1, line: 17, type: !8, scopeLine: 17, flags: DIFlagAllCallsDescribed, spFlags: DISPFlagDefinition | DISPFlagOptimized, unit: !0, retainedNodes: !2)
+!30 = !DILocation(line: 19, column: 3, scope: !31)
+!31 = !DILexicalBlockFile(scope: !29, file: !1, discriminator: 2)
+!32 = !DILocation(line: 21, column: 16, scope: !31)
+!33 = !DILocation(line: 22, column: 14, scope: !34)
+!34 = !DILexicalBlockFile(scope: !29, file: !1, discriminator: 1)
+!35 = !DILocation(line: 22, column: 11, scope: !31)
+!36 = !DILocation(line: 23, column: 16, scope: !29)
+!37 = !DILocation(line: 23, column: 9, scope: !34)
+!38 = !DILocation(line: 23, column: 9, scope: !31)
+!39 = !DILocation(line: 24, column: 14, scope: !34)
+!40 = !DILocation(line: 24, column: 11, scope: !31)
+!41 = !DILocation(line: 25, column: 16, scope: !29)
+!42 = !DILocation(line: 25, column: 9, scope: !34)
+!43 = !DILocation(line: 25, column: 9, scope: !31)
+!44 = !DILocation(line: 21, column: 21, scope: !34)
+!45 = !DILocation(line: 27, column: 1, scope: !29)
+!46 = !DILocation(line: 19, column: 24, scope: !47)
+!47 = !DILexicalBlockFile(scope: !29, file: !1, discriminator: 3)
+!48 = !DILocation(line: 19, column: 17, scope: !34)
+!49 = distinct !{!49, !50, !51}
+!50 = !DILocation(line: 19, column: 3, scope: !29)
+!51 = !DILocation(line: 26, column: 3, scope: !29)
+!52 = distinct !DISubprogram(name: "main", scope: !1, file: !1, line: 29, type: !8, scopeLine: 29, flags: DIFlagAllCallsDescribed, spFlags: DISPFlagDefinition | DISPFlagOptimized, unit: !0, retainedNodes: !2)
+!53 = !DILocation(line: 31, column: 3, scope: !54)
+!54 = !DILexicalBlockFile(scope: !52, file: !1, discriminator: 2)
+!55 = !DILocation(line: 32, column: 5, scope: !52)
+!56 = !DILocation(line: 31, column: 30, scope: !57)
+!57 = !DILexicalBlockFile(scope: !52, file: !1, discriminator: 3)
+!58 = !DILocation(line: 31, column: 17, scope: !59)
+!59 = !DILexicalBlockFile(scope: !52, file: !1, discriminator: 1)
+!60 = distinct !{!60, !61, !62}
+!61 = !DILocation(line: 31, column: 3, scope: !52)
+!62 = !DILocation(line: 33, column: 3, scope: !52)
+!63 = !DILocation(line: 34, column: 1, scope: !52)

unittests/ProfileData/SampleProfTest.cpp

@@ -9,6 +9,7 @@
 #include "llvm/ProfileData/SampleProf.h"
 #include "llvm/ADT/StringMap.h"
 #include "llvm/ADT/StringRef.h"
+#include "llvm/IR/DebugInfoMetadata.h"
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/Metadata.h"
 #include "llvm/IR/Module.h"
@@ -61,6 +62,7 @@ struct SampleProfTest : ::testing::Test {
     ASSERT_TRUE(NoError(ReaderOrErr.getError()));
     Reader = std::move(ReaderOrErr.get());
     Reader->setModule(&M);
+    Reader->setBaseDiscriminatorMask();
   }
 
   TempFile createRemapFile() {