[CSSPGO] Consume pseudo-probe-based AutoFDO profile

This change enables pseudo-probe-based sample counts to be consumed by the sample profile loader under the regular `-fprofile-sample-use` switch with minimal adjustments to the existing sample file formats. After the counts are imported, a probe helper, aka, a `PseudoProbeManager` object, is automatically launched to verify the CFG checksum of every function in the current compilation against the corresponding checksum from the profile. Mismatched checksums will cause a function profile to be slipped. A `SampleProfileProber` pass is scheduled before any of the `SampleProfileLoader` instances so that the CFG checksums as well as probe mappings are available during the profile loading time. The `PseudoProbeManager` object is set up right after the profile reading is done. In the future a CFG-based fuzzy matching could be done in `PseudoProbeManager`. Samples will be applied only to pseudo probe instructions as well as probed callsites once the checksum verification goes through. Those instructions are processed in the same way that regular instructions would be processed in the line-number-based scenario. In other words, a function is processed in a regular way as if it was reduced to just containing pseudo probes (block probes and callsites). **Adjustment to profile format ** A CFG checksum field is being added to the existing AutoFDO profile formats. So far only the text format and the extended binary format are supported. For the text format, a new line like ``` !CFGChecksum: 12345 ``` is added to the end of the body sample lines. For the extended binary profile format, we introduce a metadata section to store the checksum map from function names to their CFG checksums. Differential Revision: https://reviews.llvm.org/D92347
YuriPlyakhin · Dec 16, 2020 · ac068e0 · ac068e0
1 parent 687e80b
commit ac068e0
Show file tree

Hide file tree

Showing 19 changed files with 654 additions and 38 deletions.
diff --git a/llvm/include/llvm/IR/PseudoProbe.h b/llvm/include/llvm/IR/PseudoProbe.h
@@ -13,11 +13,14 @@
 #ifndef LLVM_IR_PSEUDOPROBE_H
 #define LLVM_IR_PSEUDOPROBE_H
 
+#include "llvm/ADT/Optional.h"
 #include <cassert>
 #include <cstdint>
 
 namespace llvm {
 
+class Instruction;
+
 constexpr const char *PseudoProbeDescMetadataName = "llvm.pseudo_probe_desc";
 
 enum class PseudoProbeType { Block = 0, IndirectCall, DirectCall };
@@ -49,6 +52,15 @@ struct PseudoProbeDwarfDiscriminator {
     return (Value >> 29) & 0x7;
   }
 };
+
+struct PseudoProbe {
+  uint32_t Id;
+  uint32_t Type;
+  uint32_t Attr;
+};
+
+Optional<PseudoProbe> extractProbe(const Instruction &Inst);
+
 } // end namespace llvm
 
 #endif // LLVM_IR_PSEUDOPROBE_H
diff --git a/llvm/include/llvm/ProfileData/SampleProf.h b/llvm/include/llvm/ProfileData/SampleProf.h
@@ -54,7 +54,8 @@ enum class sampleprof_error {
   ostream_seek_unsupported,
   compress_failed,
   uncompress_failed,
-  zlib_unavailable
+  zlib_unavailable,
+  hash_mismatch
 };
 
 inline std::error_code make_error_code(sampleprof_error E) {
@@ -120,6 +121,7 @@ enum SecType {
   SecNameTable = 2,
   SecProfileSymbolList = 3,
   SecFuncOffsetTable = 4,
+  SecFuncMetadata = 5,
   // marker for the first type of profile.
   SecFuncProfileFirst = 32,
   SecLBRProfile = SecFuncProfileFirst
@@ -137,6 +139,8 @@ static inline std::string getSecName(SecType Type) {
     return "ProfileSymbolListSection";
   case SecFuncOffsetTable:
     return "FuncOffsetTableSection";
+  case SecFuncMetadata:
+    return "FunctionMetadata";
   case SecLBRProfile:
     return "LBRProfileSection";
   }
@@ -178,6 +182,11 @@ enum class SecProfSummaryFlags : uint32_t {
   SecFlagPartial = (1 << 0)
 };
 
+enum class SecFuncMetadataFlags : uint32_t {
+  SecFlagInvalid = 0,
+  SecFlagIsProbeBased = (1 << 0),
+};
+
 // Verify section specific flag is used for the correct section.
 template <class SecFlagType>
 static inline void verifySecFlag(SecType Type, SecFlagType Flag) {
@@ -194,6 +203,9 @@ static inline void verifySecFlag(SecType Type, SecFlagType Flag) {
   case SecProfSummary:
     IsFlagLegal = std::is_same<SecProfSummaryFlags, SecFlagType>();
     break;
+  case SecFuncMetadata:
+    IsFlagLegal = std::is_same<SecFuncMetadataFlags, SecFlagType>();
+    break;
   default:
     break;
   }
@@ -502,6 +514,8 @@ class FunctionSamples {
                       : sampleprof_error::success;
   }
 
+  void setTotalSamples(uint64_t Num) { TotalSamples = Num; }
+
   sampleprof_error addHeadSamples(uint64_t Num, uint64_t Weight = 1) {
     bool Overflowed;
     TotalHeadSamples =
@@ -537,6 +551,12 @@ class FunctionSamples {
       if (ProfileIsCS)
         return 0;
       return std::error_code();
+      // A missing counter for a probe likely means the probe was not executed.
+      // Treat it as a zero count instead of an unknown count to help edge
+      // weight inference.
+      if (FunctionSamples::ProfileIsProbeBased)
+        return 0;
+      return std::error_code();
     } else {
       return ret->second.getSamples();
     }
@@ -553,6 +573,16 @@ class FunctionSamples {
     return ret->second.getCallTargets();
   }
 
+  /// Returns the call target map collected at a given location specified by \p
+  /// CallSite. If the location is not found in profile, return error.
+  ErrorOr<SampleRecord::CallTargetMap>
+  findCallTargetMapAt(const LineLocation &CallSite) const {
+    const auto &Ret = BodySamples.find(CallSite);
+    if (Ret == BodySamples.end())
+      return std::error_code();
+    return Ret->second.getCallTargets();
+  }
+
   /// Return the function samples at the given callsite location.
   FunctionSamplesMap &functionSamplesAt(const LineLocation &Loc) {
     return CallsiteSamples[Loc];
@@ -641,6 +671,21 @@ class FunctionSamples {
     Name = Other.getName();
     if (!GUIDToFuncNameMap)
       GUIDToFuncNameMap = Other.GUIDToFuncNameMap;
+
+    if (FunctionHash == 0) {
+      // Set the function hash code for the target profile.
+      FunctionHash = Other.getFunctionHash();
+    } else if (FunctionHash != Other.getFunctionHash()) {
+      // The two profiles coming with different valid hash codes indicates
+      // either:
+      // 1. They are same-named static functions from different compilation
+      // units (without using -unique-internal-linkage-names), or
+      // 2. They are really the same function but from different compilations.
+      // Let's bail out in either case for now, which means one profile is
+      // dropped.
+      return sampleprof_error::hash_mismatch;
+    }
+
     MergeResult(Result, addTotalSamples(Other.getTotalSamples(), Weight));
     MergeResult(Result, addHeadSamples(Other.getHeadSamples(), Weight));
     for (const auto &I : Other.getBodySamples()) {
@@ -700,6 +745,10 @@ class FunctionSamples {
   /// Return the original function name.
   StringRef getFuncName() const { return getFuncName(Name); }
 
+  void setFunctionHash(uint64_t Hash) { FunctionHash = Hash; }
+
+  uint64_t getFunctionHash() const { return FunctionHash; }
+
   /// Return the canonical name for a function, taking into account
   /// suffix elision policy attributes.
   static StringRef getCanonicalFnName(const Function &F) {
@@ -754,6 +803,12 @@ class FunctionSamples {
   /// We assume that a single function will not exceed 65535 LOC.
   static unsigned getOffset(const DILocation *DIL);
 
+  /// Returns a unique call site identifier for a given debug location of a call
+  /// instruction. This is wrapper of two scenarios, the probe-based profile and
+  /// regular profile, to hide implementation details from the sample loader and
+  /// the context tracker.
+  static LineLocation getCallSiteIdentifier(const DILocation *DIL);
+
   /// Get the FunctionSamples of the inline instance where DIL originates
   /// from.
   ///
@@ -769,6 +824,8 @@ class FunctionSamples {
       const DILocation *DIL,
       SampleProfileReaderItaniumRemapper *Remapper = nullptr) const;
 
+  static bool ProfileIsProbeBased;
+
   static bool ProfileIsCS;
 
   SampleContext &getContext() const { return Context; }
@@ -799,6 +856,9 @@ class FunctionSamples {
   /// Mangled name of the function.
   StringRef Name;
 
+  /// CFG hash value for the function.
+  uint64_t FunctionHash = 0;
+
   /// Calling context for function profile
   mutable SampleContext Context;
 

diff --git a/llvm/include/llvm/ProfileData/SampleProfReader.h b/llvm/include/llvm/ProfileData/SampleProfReader.h
@@ -27,8 +27,9 @@
 //      offsetA[.discriminator]: fnA:num_of_total_samples
 //       offsetA1[.discriminator]: number_of_samples [fn7:num fn8:num ... ]
 //       ...
+//      !CFGChecksum: num
 //
-// This is a nested tree in which the identation represents the nesting level
+// This is a nested tree in which the indentation represents the nesting level
 // of the inline stack. There are no blank lines in the file. And the spacing
 // within a single line is fixed. Additional spaces will result in an error
 // while reading the file.
@@ -47,10 +48,11 @@
 // in the prologue of the function (second number). This head sample
 // count provides an indicator of how frequently the function is invoked.
 //
-// There are two types of lines in the function body.
+// There are three types of lines in the function body.
 //
 // * Sampled line represents the profile information of a source location.
 // * Callsite line represents the profile information of a callsite.
+// * Metadata line represents extra metadata of the function.
 //
 // Each sampled line may contain several items. Some are optional (marked
 // below):
@@ -114,6 +116,18 @@
 //    total number of samples collected for the inlined instance at this
 //    callsite
 //
+// Metadata line can occur in lines with one indent only, containing extra
+// information for the top-level function. Furthermore, metadata can only
+// occur after all the body samples and callsite samples.
+// Each metadata line may contain a particular type of metadata, marked by
+// the starting characters annotated with !. We process each metadata line
+// independently, hence each metadata line has to form an independent piece
+// of information that does not require cross-line reference.
+// We support the following types of metadata:
+//
+// a. CFG Checksum (a.k.a. function hash):
+//   !CFGChecksum: 12345
+//
 //
 // Binary format
 // -------------
@@ -419,7 +433,10 @@ class SampleProfileReader {
   /// \brief Return the profile format.
   SampleProfileFormat getFormat() const { return Format; }
 
-  /// Whether input profile is fully context-sensitie
+  /// Whether input profile is based on pseudo probes.
+  bool profileIsProbeBased() const { return ProfileIsProbeBased; }
+
+  /// Whether input profile is fully context-sensitive
   bool profileIsCS() const { return ProfileIsCS; }
 
   virtual std::unique_ptr<ProfileSymbolList> getProfileSymbolList() {
@@ -464,6 +481,9 @@ class SampleProfileReader {
 
   std::unique_ptr<SampleProfileReaderItaniumRemapper> Remapper;
 
+  /// \brief Whether samples are collected based on pseudo probes.
+  bool ProfileIsProbeBased = false;
+
   bool ProfileIsCS = false;
 
   /// \brief The format of sample.
@@ -606,6 +626,7 @@ class SampleProfileReaderExtBinaryBase : public SampleProfileReaderBinary {
   std::error_code readSecHdrTableEntry();
   std::error_code readSecHdrTable();
 
+  std::error_code readFuncMetadata();
   std::error_code readFuncOffsetTable();
   std::error_code readFuncProfiles();
   std::error_code readMD5NameTable();

diff --git a/llvm/include/llvm/ProfileData/SampleProfWriter.h b/llvm/include/llvm/ProfileData/SampleProfWriter.h
@@ -200,6 +200,8 @@ class SampleProfileWriterExtBinaryBase : public SampleProfileWriterBinary {
   // Helper function to write name table.
   virtual std::error_code writeNameTable() override;
 
+  std::error_code writeFuncMetadata(const StringMap<FunctionSamples> &Profiles);
+
   // Functions to write various kinds of sections.
   std::error_code
   writeNameTableSection(const StringMap<FunctionSamples> &ProfileMap);
@@ -270,11 +272,10 @@ class SampleProfileWriterExtBinary : public SampleProfileWriterExtBinaryBase {
     // SecFuncOffsetTable section is written after SecLBRProfile in the
     // profile because FuncOffsetTable needs to be populated while section
     // SecLBRProfile is written.
-    SectionHdrLayout = {{SecProfSummary, 0, 0, 0},
-                        {SecNameTable, 0, 0, 0},
-                        {SecFuncOffsetTable, 0, 0, 0},
-                        {SecLBRProfile, 0, 0, 0},
-                        {SecProfileSymbolList, 0, 0, 0}};
+    SectionHdrLayout = {
+        {SecProfSummary, 0, 0, 0},       {SecNameTable, 0, 0, 0},
+        {SecFuncOffsetTable, 0, 0, 0},   {SecLBRProfile, 0, 0, 0},
+        {SecProfileSymbolList, 0, 0, 0}, {SecFuncMetadata, 0, 0, 0}};
   };
   virtual std::error_code
   writeSections(const StringMap<FunctionSamples> &ProfileMap) override;

diff --git a/llvm/include/llvm/Transforms/IPO/SampleProfileProbe.h b/llvm/include/llvm/Transforms/IPO/SampleProfileProbe.h
@@ -18,18 +18,44 @@
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/IR/PassManager.h"
 #include "llvm/IR/PseudoProbe.h"
+#include "llvm/ProfileData/SampleProf.h"
 #include "llvm/Target/TargetMachine.h"
 #include <unordered_map>
 
 namespace llvm {
 
 class Module;
 
+using namespace sampleprof;
 using BlockIdMap = std::unordered_map<BasicBlock *, uint32_t>;
 using InstructionIdMap = std::unordered_map<Instruction *, uint32_t>;
 
 enum class PseudoProbeReservedId { Invalid = 0, Last = Invalid };
 
+class PseudoProbeDescriptor {
+  uint64_t FunctionGUID;
+  uint64_t FunctionHash;
+
+public:
+  PseudoProbeDescriptor(uint64_t GUID, uint64_t Hash)
+      : FunctionGUID(GUID), FunctionHash(Hash) {}
+  uint64_t getFunctionGUID() const { return FunctionGUID; }
+  uint64_t getFunctionHash() const { return FunctionHash; }
+};
+
+// This class serves sample counts correlation for SampleProfileLoader by
+// analyzing pseudo probes and their function descriptors injected by
+// SampleProfileProber.
+class PseudoProbeManager {
+  DenseMap<uint64_t, PseudoProbeDescriptor> GUIDToProbeDescMap;
+
+  const PseudoProbeDescriptor *getDesc(const Function &F) const;
+
+public:
+  PseudoProbeManager(const Module &M);
+  bool moduleIsProbed(const Module &M) const;
+  bool profileIsValid(const Function &F, const FunctionSamples &Samples) const;
+};
 
 /// Sample profile pseudo prober.
 ///

diff --git a/llvm/lib/IR/CMakeLists.txt b/llvm/lib/IR/CMakeLists.txt
@@ -48,6 +48,7 @@ add_llvm_component_library(LLVMCore
   PrintPasses.cpp
   SafepointIRVerifier.cpp
   ProfileSummary.cpp
+  PseudoProbe.cpp
   Statepoint.cpp
   StructuralHash.cpp
   Type.cpp

diff --git a/llvm/lib/IR/PseudoProbe.cpp b/llvm/lib/IR/PseudoProbe.cpp
@@ -0,0 +1,58 @@
+//===- PseudoProbe.cpp - Pseudo Probe Helpers -----------------------------===//
+//
+// 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 implements the helpers to manipulate pseudo probe IR intrinsic
+// calls.
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/IR/PseudoProbe.h"
+#include "llvm/IR/DebugInfoMetadata.h"
+#include "llvm/IR/IRBuilder.h"
+#include "llvm/IR/Instruction.h"
+
+using namespace llvm;
+
+namespace llvm {
+
+Optional<PseudoProbe> extractProbeFromDiscriminator(const Instruction &Inst) {
+  assert(isa<CallBase>(&Inst) && !isa<IntrinsicInst>(&Inst) &&
+         "Only call instructions should have pseudo probe encodes as their "
+         "Dwarf discriminators");
+  if (const DebugLoc &DLoc = Inst.getDebugLoc()) {
+    const DILocation *DIL = DLoc;
+    auto Discriminator = DIL->getDiscriminator();
+    if (DILocation::isPseudoProbeDiscriminator(Discriminator)) {
+      PseudoProbe Probe;
+      Probe.Id =
+          PseudoProbeDwarfDiscriminator::extractProbeIndex(Discriminator);
+      Probe.Type =
+          PseudoProbeDwarfDiscriminator::extractProbeType(Discriminator);
+      Probe.Attr =
+          PseudoProbeDwarfDiscriminator::extractProbeAttributes(Discriminator);
+      return Probe;
+    }
+  }
+  return None;
+}
+
+Optional<PseudoProbe> extractProbe(const Instruction &Inst) {
+  if (const auto *II = dyn_cast<PseudoProbeInst>(&Inst)) {
+    PseudoProbe Probe;
+    Probe.Id = II->getIndex()->getZExtValue();
+    Probe.Type = (uint32_t)PseudoProbeType::Block;
+    Probe.Attr = II->getAttributes()->getZExtValue();
+    return Probe;
+  }
+
+  if (isa<CallBase>(&Inst) && !isa<IntrinsicInst>(&Inst))
+    return extractProbeFromDiscriminator(Inst);
+
+  return None;
+}
+} // namespace llvm