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

[xur]

Summary

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.

[wmi]

llvm/include/llvm/Support/Discriminator.h

56

Can we use unsigned int instead so no implicit conversion when FSPassHighBit is used? And how about using std::array so we can use FSPassHighBit.at(i) in FSPassBitBegin and FSPassBitEnd which will have array bounds check.

90–92

Nit: For N==31, (1U << (N + 1)) - 1 is 0xFFFFFFFF so no if (N == 31) is needed.

llvm/lib/ProfileData/SampleProfReader.cpp

312–314

Nit: Discrriminator &= getDiscriminatorMask();

526–527

Nit: DiscrriminatorVal &= getDiscriminatorMask();

565–566

Nit: DiscrriminatorVal &= getDiscriminatorMask();

[xur]

llvm/include/llvm/Support/Discriminator.h

56

OK. This sounds good to me. Will change.

90–92

Yes. using 1U is a good idea. I will change.

[xur]

Thank Wei for the suggestions. Here is the new patch integrated his review.

[hoy]

llvm/include/llvm/Support/Discriminator.h

64

Does it make sense to include 0 in the FSPassHighBit array so that the check can be omitted?

How about add a assert of I <= getNumFSPasses() ?

llvm/lib/ProfileData/SampleProfReader.cpp

312

Assuming the sample loader runs in multiple places of the pipeline, does it make sense to assert the length of Discriminator matches the length of the mask? IIUC, if the number of valid bits of Discriminator is more than that of the mask, it means we are loading the profile of an earlier stage and we should use a longer mask?

[wenlei]

llvm/include/llvm/ProfileData/SampleProfReader.h

355

This is an API that client of profile reader will need to call, so I'm wondering if we can make it higher level - i.e, instead of requiring the client to be aware and pass in the bit-wise mask, just let client set it for base discriminator (i == 0) or i-th FS discriminator?

llvm/include/llvm/Support/Discriminator.h

56

Not a big deal but do we actually need a map here? Could just use getBaseDiscriminatorBits + i * FSDiscriminatorWidth in a helper?

72

There're mentions of FS passes and FD passes in multiple places in the comment, are they referring to the same thing?

llvm/lib/CodeGen/TargetPassConfig.cpp

1181

A comment explaining the purpose of final discriminator pass would be helpful here.

llvm/lib/ProfileData/SampleProfReader.cpp

564

We could move the ProfileIsFS check into getDiscriminatorMask and make the mask all 1s for non-FS case?

[xur]

Thanks to Hongtao and Weilei for the review!

llvm/include/llvm/Support/Discriminator.h

56

Yes. I think we can do that. The reason we have a map for the bits is because I used variable length of bit for different FS discriminator passes in my earlier implementation /experiments.
I think we will end up with using fixed length here. I'll change this.

64

This part is changed to go with Wenlei's suggestion.
I will assert that I >0 and I<=0 getNumFSPasses(). This way we don't need the check. We don't have a caller with 0 as the argument anyway.

72

Sorry. FD is a typo. Fixed.

llvm/lib/CodeGen/TargetPassConfig.cpp

1181

Will do.

llvm/lib/ProfileData/SampleProfReader.cpp

312

I'm not sure if I understand the question.
Hers the "discriminator" is the one from profile which can have all the discriminator bits (i.e. including the later passes). Here we use the mask to remove the later bits and merge them together.

564

Done.

[xur]

Integrated review comments from Hongtao and Weilei.

@Wei: I had to add back the special check for N==31 in getN1Bits().
My stock g++, gcc version 10.2.1 20210110 (Debian 10.2.1-6+build2)
folds "(1U << (N + 1)) - 1" to 0.

I think this is a gcc bug.

With this code check, the test case will break.

[hoy]

llvm/include/llvm/Support/Discriminator.h

17

Nit: this is not needed anymore.

llvm/lib/ProfileData/SampleProfReader.cpp

312

You are right. I was confused.

[xur]

Remove unneeded head files (thanks to Hongtao).

[wmi]

LGTM.

[wenlei]

lgtm, thanks.

llvm/include/llvm/ProfileData/SampleProfReader.h

357

maybe worth making a enum for FS-Passes. So it's obviously how we organize FS passes, and we don't need to pass int literals around.

[xur]

Take Wenlei's suggestion to use enum (instead of using integer literal).
The change is bigger than I initially thought. So probably needs another look from the reviewers.

llvm/include/llvm/IR/DebugInfoMetadata.h

public:
   /// Return the masked discriminator value for an input discrimnator value D
   /// (i.e. zero out the (B+1)-th and above bits for D (B is 0-base).
   // Example: an input of (0x1FF, 7) returns 0xFF.
   static unsigned getMaskedDiscriminator(unsigned D, unsigned B) {
     return (D & getN1Bits(B));
   }
 
   /// 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) {
     if (EnableFSDiscriminator)
       return getMaskedDiscriminator(D, getBaseDiscriminatorBits());
     return getUnsignedFromPrefixEncoding(D);
   }
 

llvm/include/llvm/ProfileData/SampleProf.h

 enum class SecProfSummaryFlags : uint32_t {
   SecFlagInValid = 0,
   /// SecFlagPartial means the profile is for common/shared code.
   /// The common profile is usually merged from profiles collected
   /// from running other targets.
   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 {
   SecFlagInvalid = 0,
   SecFlagIsProbeBased = (1 << 0),
   SecFlagHasAttribute = (1 << 1)
 };
 
public:
   static SampleProfileFormat Format;
 
   /// Whether the profile uses MD5 to represent string.
   static bool UseMD5;
 
   /// 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;
 
   // Assume the input \p Name is a name coming from FunctionSamples itself.
   // If UseMD5 is true, the name is already a GUID and we
   // don't want to return the GUID of GUID.
   static uint64_t getGUID(StringRef Name) {

llvm/include/llvm/ProfileData/SampleProfReader.h

 #include "llvm/ADT/StringRef.h"
 #include "llvm/IR/DiagnosticInfo.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/ProfileSummary.h"
 #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"
 #include <algorithm>
 #include <cstdint>
 #include <memory>
 #include <string>
 #include <system_error>
SampleProfileReader(std::unique_ptr<MemoryBuffer> B, LLVMContext &C,
                       SampleProfileFormat Format = SPF_None)
       : Profiles(0), Ctx(C), Buffer(std::move(B)), Format(Format) {}
 
   virtual ~SampleProfileReader() = default;
 
   /// Read and validate the file header.
   virtual std::error_code readHeader() = 0;
 
+  /// Set the bits for FS discriminators.
+  void setDiscriminatorMaskedBitFrom(uint32_t B) { MaskedBitFrom = B; }
+
+  /// Get the bitmask for FS discriminators.
+  uint32_t getDiscriminatorMask() const {
+    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())
       return EC;
     if (Remapper)
       Remapper->applyRemapping(Ctx);
     FunctionSamples::UseMD5 = useMD5();
     return sampleprof_error::success;
protected:
   bool ProfileIsProbeBased = false;
 
   /// Whether function profiles are context-sensitive.
   bool ProfileIsCS = false;
 
   /// 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;
 
   /// \brief The current module being compiled if SampleProfileReader
   /// 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 {
 public:
   SampleProfileReaderText(std::unique_ptr<MemoryBuffer> B, LLVMContext &C)
       : SampleProfileReader(std::move(B), C, SPF_Text) {}
 
   /// Read and validate the file header.

llvm/include/llvm/Support/Discriminator.h

 //===---- llvm/Support/Discriminator.h -- Discriminator Utils ---*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file defines the constants and utility functions for discriminators.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_SUPPORT_DISCRIMINATOR_H
 #define LLVM_SUPPORT_DISCRIMINATOR_H
 
+#include "llvm/Support/Error.h"
+#include <array>
+#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
 /// old_bit 6~12 --> new_bit 7~13
 /// new_bit_6 is 0 if higher bits (7~13) are all 0
 static inline unsigned getPrefixEncodingFromUnsigned(unsigned U) {
   U &= 0xfff;
   return U > 0x1f ? (((U & 0xfe0) << 1) | (U & 0x1f) | 0x20) : U;
static inline unsigned encodeComponent(unsigned C) {
   return (C == 0) ? 1U : (getPrefixEncodingFromUnsigned(C) << 1);
 }
 
 static inline unsigned encodingBits(unsigned C) {
   return (C == 0) ? 1 : (C > 0x1f ? 14 : 7);
 }
 
 // Some constants used in FS Discriminators.
-#define BASE_DIS_BIT_BEG 0
-#define BASE_DIS_BIT_END 7
+// Define the High bit for each FS pass.
+const std::array<unsigned, 5> FSPassHighBit = {7, 13, 19, 25, 31};
+
+// Return the number of FS passes, excluding the pass adding the base
+// discriminators.
+static inline unsigned getNumFSPasses() { return FSPassHighBit.size() - 1; }
+
+// Return the begining bit for the i-th FS pass.
+static inline unsigned getFSPassBitBegin(unsigned I) {
+  if (I == 0)
+    return 0;
+  return FSPassHighBit[I - 1] + 1;
+}
 
-#define PASS_1_DIS_BIT_BEG 8
-#define PASS_1_DIS_BIT_END 13
+// Return the ending bit for the i-th FS pass.
+static inline unsigned getFSPassBitEnd(unsigned I) { return FSPassHighBit[I]; }
 
-#define PASS_2_DIS_BIT_BEG 14
-#define PASS_2_DIS_BIT_END 19
+// Return the beginning bit for the last FD pass.
+static inline int getLastFSPassBitBegin() {
+  return getFSPassBitBegin(getNumFSPasses());
+}
+
+// Return the ending bit for the last FD pass.
+static inline unsigned getLastFSPassBitEnd() {
+  return getFSPassBitEnd(getNumFSPasses());
+}
 
-#define PASS_3_DIS_BIT_BEG 20
-#define PASS_3_DIS_BIT_END 25
+// Return the beginning bit for the base (first) FD pass.
+static inline unsigned getBaseFSBitBegin() { return 0; }
 
-#define PASS_LAST_DIS_BIT_BEG 26
-#define PASS_LAST_DIS_BIT_END 31
+// Return the ending bit for the base (first) FD pass.
+static inline unsigned getBaseFSBitEnd() { return FSPassHighBit[0]; }
 
-// Set bits range [0 .. n] to 1. Used in FS Discriminators.
+// Set bits in range of [0 .. n] to 1. Used in FS Discriminators.
 static inline unsigned getN1Bits(int N) {
-  if (N >= 31)
-    return 0xFFFFFFFF;
-  return (1 << (N + 1)) - 1;
+  assert((N < 32) && "N is invalid");
+  return (1U << (N + 1)) - 1;
 }
 
 #endif /* LLVM_SUPPORT_DISCRIMINATOR_H */

llvm/lib/CodeGen/TargetPassConfig.cpp

void TargetPassConfig::addMachinePasses() {
 
   // Insert before XRay Instrumentation.
   addPass(&FEntryInserterID);
 
   addPass(&XRayInstrumentationID);
   addPass(&PatchableFunctionID);
 
   if (EnableFSDiscriminator && !FSNoFinalDiscrim)
-    addPass(createMIRAddFSDiscriminatorsPass(PASS_LAST_DIS_BIT_BEG,
-                                             PASS_LAST_DIS_BIT_END));
+    addPass(createMIRAddFSDiscriminatorsPass(getLastFSPassBitBegin(),
+                                             getLastFSPassBitEnd()));
 
   addPreEmitPass();
 
   if (TM->Options.EnableIPRA)
     // Collect register usage information and produce a register mask of
     // clobbered registers, to be used to optimize call sites.
     addPass(createRegUsageInfoCollector());
 

llvm/lib/ProfileData/SampleProf.cpp

 
 namespace llvm {
 namespace sampleprof {
 SampleProfileFormat FunctionSamples::Format;
 bool FunctionSamples::ProfileIsProbeBased = false;
 bool FunctionSamples::ProfileIsCS = false;
 bool FunctionSamples::UseMD5 = false;
 bool FunctionSamples::HasUniqSuffix = true;
+bool FunctionSamples::ProfileIsFS = false;
 } // namespace sampleprof
 } // namespace llvm
 
 namespace {
 
 // FIXME: This class is only here to support the transition to llvm::Error. It
 // will be removed once this transition is complete. Clients should prefer to
 // deal with the Error value directly, rather than converting to error_code.

 
 const FunctionSamples *FunctionSamples::findFunctionSamples(
     const DILocation *DIL, SampleProfileReaderItaniumRemapper *Remapper) const {
   assert(DIL);
   SmallVector<std::pair<LineLocation, StringRef>, 10> S;
 
   const DILocation *PrevDIL = DIL;
   for (DIL = DIL->getInlinedAt(); DIL; DIL = DIL->getInlinedAt()) {
-    S.push_back(std::make_pair(
-        LineLocation(getOffset(DIL), DIL->getBaseDiscriminator()),
+    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)
     return this;
   const FunctionSamples *FS = this;
   for (int i = S.size() - 1; i >= 0 && FS != nullptr; i--) {
     FS = FS->findFunctionSamplesAt(S[i].first, S[i].second, Remapper);
   }

llvm/lib/ProfileData/SampleProfReader.cpp

 
 #include "llvm/ProfileData/SampleProfReader.h"
 #include "llvm/ADT/DenseMap.h"
 #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/CommandLine.h"
 #include "llvm/Support/Compression.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 <set>
 #include <system_error>
 #include <vector>
 
 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.
 /// \param OS Stream to emit the output to.
 void SampleProfileReader::dumpFunctionProfile(StringRef FName,
                                               raw_ostream &OS) {
   OS << "Function: " << FName << ": " << Profiles[FName];
 }
std::error_code SampleProfileReaderText::readImpl() {
 
   InlineCallStack InlineStack;
   uint32_t ProbeProfileCount = 0;
 
   // SeenMetadata tracks whether we have processed metadata for the current
   // top-level function profile.
   bool SeenMetadata = false;
 
+  ProfileIsFS = ProfileIsFSDisciminator;
   for (; !LineIt.is_at_eof(); ++LineIt) {
     if ((*LineIt)[(*LineIt).find_first_not_of(' ')] == '#')
       continue;
     // Read the header of each function.
     //
     // Note that for function identifiers we are actually expecting
     // mangled names, but we may not always get them. This happens when
     // the compiler decides not to emit the function (e.g., it was inlined
if ((*LineIt)[0] != ' ') {
         return sampleprof_error::malformed;
       }
       if (SeenMetadata && LineTy != LineType::Metadata) {
         // Metadata must be put at the end of a function profile.
         reportError(LineIt.line_number(),
                     "Found non-metadata after metadata: " + *LineIt);
         return sampleprof_error::malformed;
       }
+
+      // Here we handle FS discriminators.
+      if (ProfileIsFS)
+        Discriminator &= getDiscriminatorMask();
+
       while (InlineStack.size() > Depth) {
         InlineStack.pop_back();
       }
       switch (LineTy) {
       case LineType::CallSiteProfile: {
         FunctionSamples &FSamples = InlineStack.back()->functionSamplesAt(
             LineLocation(LineOffset, Discriminator))[std::string(FName)];
         FSamples.setName(FName);
         MergeResult(Result, FSamples.addTotalSamples(NumSamples));
for (uint32_t I = 0; I < *NumRecords; ++I) {
     auto NumSamples = readNumber<uint64_t>();
     if (std::error_code EC = NumSamples.getError())
       return EC;
 
     auto NumCalls = readNumber<uint32_t>();
     if (std::error_code EC = NumCalls.getError())
       return EC;
 
+    // Here we handle FS discriminators:
+    uint32_t DiscriminatorVal = *Discriminator;
+    if (ProfileIsFS)
+      DiscriminatorVal &= getDiscriminatorMask();
+
     for (uint32_t J = 0; J < *NumCalls; ++J) {
       auto CalledFunction(readStringFromTable());
       if (std::error_code EC = CalledFunction.getError())
         return EC;
 
       auto CalledFunctionSamples = readNumber<uint64_t>();
       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.
   auto NumCallsites = readNumber<uint32_t>();
   if (std::error_code EC = NumCallsites.getError())
     return EC;
 
   for (uint32_t J = 0; J < *NumCallsites; ++J) {
     auto LineOffset = readNumber<uint64_t>();
     if (std::error_code EC = LineOffset.getError())
       return EC;
 
     auto Discriminator = readNumber<uint64_t>();
     if (std::error_code EC = Discriminator.getError())
       return EC;
 
     auto FName(readStringFromTable());
     if (std::error_code EC = FName.getError())
       return EC;
 
+    // Here we handle FS discriminators:
+    uint32_t DiscriminatorVal = *Discriminator;
+    if (ProfileIsFS)
+      DiscriminatorVal &= 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;
   }
 
   return sampleprof_error::success;
 }
 
if (FContext.hasContext())
     CSProfileCount++;
 
   if (std::error_code EC = readProfile(FProfile))
     return EC;
   return sampleprof_error::success;
 }
 
 std::error_code SampleProfileReaderBinary::readImpl() {
+  ProfileIsFS = ProfileIsFSDisciminator;
   while (!at_eof()) {
     if (std::error_code EC = readFuncProfile(Data))
       return EC;
   }
 
   return sampleprof_error::success;
 }
 
 std::error_code SampleProfileReaderExtBinaryBase::readOneSection(
     const uint8_t *Start, uint64_t Size, const SecHdrTableEntry &Entry) {
   Data = Start;
   End = Start + Size;
   switch (Entry.Type) {
   case SecProfSummary:
     if (std::error_code EC = readSummary())
       return EC;
     if (hasSecFlag(Entry, SecProfSummaryFlags::SecFlagPartial))
       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 =
         hasSecFlag(Entry, SecNameTableFlags::SecFlagFixedLengthMD5);
     bool UseMD5 = hasSecFlag(Entry, SecNameTableFlags::SecFlagMD5Name);
     assert((!FixedLengthMD5 || UseMD5) &&
            "If FixedLengthMD5 is true, UseMD5 has to be true");
     FunctionSamples::HasUniqSuffix =
std::error_code SampleProfileReaderExtBinaryBase::readImpl() {
 
   return sampleprof_error::success;
 }
 
 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.
     for (auto FuncEntry : FuncOffsetTable) {
       OffsetsToUse.push_back(FuncEntry.second);
     }
   } else {
     // load function profiles on demand.
case SecNameTable:
     if (hasSecFlag(Entry, SecNameTableFlags::SecFlagUniqSuffix))
       Flags.append("uniq,");
     break;
   case SecProfSummary:
     if (hasSecFlag(Entry, SecProfSummaryFlags::SecFlagPartial))
       Flags.append("partial,");
     if (hasSecFlag(Entry, SecProfSummaryFlags::SecFlagFullContext))
       Flags.append("context,");
+    if (hasSecFlag(Entry, SecProfSummaryFlags::SecFlagFSDiscriminator))
+      Flags.append("fs-discriminator,");
     break;
   default:
     break;
   }
   char &last = Flags.back();
   if (last == ',')
     last = '}';
   else
std::error_code SampleProfileReaderGCC::readOneFunctionProfile(
   return sampleprof_error::success;
 }
 
 /// Read a GCC AutoFDO profile.
 ///
 /// This format is generated by the Linux Perf conversion tool at
 /// https://github.com/google/autofdo.
 std::error_code SampleProfileReaderGCC::readImpl() {
+  assert(!ProfileIsFSDisciminator && "Gcc profiles not support FSDisciminator");
   // Read the string table.
   if (std::error_code EC = readNameTable())
     return EC;
 
   // Read the source profile.
   if (std::error_code EC = readFunctionProfiles())
     return EC;
 

llvm/lib/ProfileData/SampleProfWriter.cpp

std::error_code SampleProfileWriterExtBinaryBase::writeOneSection(
   if (Type == SecProfileSymbolList && ProfSymList && ProfSymList->toCompress())
     setToCompressSection(SecProfileSymbolList);
   if (Type == SecFuncMetadata && FunctionSamples::ProfileIsProbeBased)
     addSectionFlag(SecFuncMetadata, SecFuncMetadataFlags::SecFlagIsProbeBased);
   if (Type == SecProfSummary && FunctionSamples::ProfileIsCS)
     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) {
   case SecProfSummary:
     computeSummary(ProfileMap);
     if (auto EC = writeSummary())
       return EC;
     break;

llvm/lib/Target/X86/X86InsertPrefetch.cpp

ErrorOr<std::unique_ptr<SampleProfileReader>> 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,
                                              DiagnosticSeverity::DS_Warning));
     return false;
   }
   Reader = std::move(ReaderOrErr.get());
+  Reader->setDiscriminatorMaskedBitFrom(DILocation::getBaseDiscriminatorBits());
   Reader->read();
   return true;
 }
 
 void X86InsertPrefetch::getAnalysisUsage(AnalysisUsage &AU) const {
   AU.setPreservesAll();
   MachineFunctionPass::getAnalysisUsage(AU);
 }

llvm/lib/Transforms/IPO/SampleProfile.cpp

if (std::error_code EC = ReaderOrErr.getError()) {
     Ctx.diagnose(DiagnosticInfoSampleProfile(Filename, Msg));
     return false;
   }
   Reader = std::move(ReaderOrErr.get());
   Reader->setSkipFlatProf(LTOPhase == ThinOrFullLTOPhase::ThinLTOPostLink);
   // 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->setDiscriminatorMaskedBitFrom(DILocation::getBaseDiscriminatorBits());
   if (std::error_code EC = Reader->read()) {
     std::string Msg = "profile reading failed: " + EC.message();
     Ctx.diagnose(DiagnosticInfoSampleProfile(Filename, Msg));
     return false;
   }
 
   PSL = Reader->getProfileSymbolList();
 

llvm/test/Transforms/SampleProfile/Inputs/fsafdo.prof

+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

llvm/test/Transforms/SampleProfile/fsafdo_test.ll

+; 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)

llvm/unittests/ProfileData/SampleProfTest.cpp

 //===- unittest/ProfileData/SampleProfTest.cpp ------------------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #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"
 #include "llvm/ProfileData/SampleProfReader.h"
 #include "llvm/ProfileData/SampleProfWriter.h"
 #include "llvm/Support/Casting.h"
 #include "llvm/Support/ErrorOr.h"
 #include "llvm/Support/FileSystem.h"
struct SampleProfTest : ::testing::Test {
 
   void readProfile(const Module &M, StringRef Profile,
                    StringRef RemapFile = "") {
     auto ReaderOrErr = SampleProfileReader::create(
         std::string(Profile), Context, std::string(RemapFile));
     ASSERT_TRUE(NoError(ReaderOrErr.getError()));
     Reader = std::move(ReaderOrErr.get());
     Reader->setModule(&M);
+    Reader->setDiscriminatorMaskedBitFrom(
+        DILocation::getBaseDiscriminatorBits());
   }
 
   TempFile createRemapFile() {
     return TempFile("remapfile", "", R"(
       # Types 'int' and 'long' are equivalent
       type i l
       # Function names 'foo' and 'faux' are equivalent
       name 3foo 4faux