[CSSPGO] Pseudo probes for function calls.

[hoy]

Summary

An indirect call site needs to be probed for its potential call targets. With CSSPGO a direct call also needs a probe so that a calling context can be represented by a stack of callsite probes. Unlike pseudo probes for basic blocks that are in form of standalone intrinsic call instructions, pseudo probes for callsites have to be attached to the call instruction, thus a separate instruction would not work.

One possible way of attaching a probe to a call instruction is to use a special metadata that carries information about the probe. The special metadata will have to make its way through the optimization pipeline down to object emission. This requires additional efforts to maintain the metadata in various places. Given that the !dbg metadata is a first-class metadata and has all essential support in place , leveraging the !dbg metadata as a channel to encode pseudo probe information is probably the easiest solution.

With the requirement of not inflating !dbg metadata that is allocated for almost every instruction, we found that the 32-bit DWARF discriminator field which mainly serves AutoFDO can be reused for pseudo probes. DWARF discriminators distinguish identical source locations between instructions and with pseudo probes such support is not required. In this change we are using the discriminator field to encode the ID and type of a callsite probe and the encoded value will be unpacked and consumed right before object emission. When a callsite is inlined, the callsite discriminator field will go with the inlined instructions. The !dbg metadata of an inlined instruction is in form of a scope stack. The top of the stack is the instruction's original !dbg metadata and the bottom of the stack is for the original callsite of the top-level inliner. Except for the top of the stack, all other elements of the stack actually refer to the nested inlined callsites whose discriminator field (which actually represents a calliste probe) can be used together to represent the inline context of an inlined PseudoProbeInst or CallInst.

To avoid collision with the baseline AutoFDO in various places that handles dwarf discriminators where a check against the -pseudo-probe-for-profiling switch is not available, a special encoding scheme is used to tell apart a pseudo probe discriminator from a regular discriminator. For the regular discriminator, if all lowest 3 bits are non-zero, it means the discriminator is basically empty and all higher 29 bits can be reversed for pseudo probe use.

Callsite pseudo probes are inserted in SampleProfileProbePass and a target-independent MIR pass PseudoProbeInserter is added to unpack the probe ID/type from !dbg.

Note that with this work the switch -debug-info-for-profiling will not work with -pseudo-probe-for-profiling anymore. They cannot be used at the same time.

[hoy]

Updating D91756: [CSSPGO] Pseudo probes for function calls.

[hoy]

Updating D91756: [CSSPGO] Pseudo probes for function calls.

Corresponding changes to support IR/MIR intrinsic attributes.

[wmi]

llvm/lib/CodeGen/PseudoProbeInserter.cpp

51

What is the usage of FirstInstr?

55

Will tailcall or other optimizations convert call into something else before PseudoProbeInserter pass?

llvm/lib/Transforms/IPO/SampleProfileProbe.cpp

137–139

for (auto &[Call, Index] : CallProbeIds) {

[hoy]

llvm/lib/CodeGen/PseudoProbeInserter.cpp

51

Good patch. It's not used in this change. It's used in an upcoming change where a probe probing an empty block will be marked specially for better counts inference.

55

Good point. A tailcall instruction (a special MIR jump instruction like TAILJMPd) is considered as a call and can be identified as MI.isCall() && MI.isTerminator() && MI.isReturn(). Tail calls are bad to the frame-pointer-based virtual unwinding since they may cause missing frames. An upcoming change will mark them specially and a tail call tracker in llvm-profgen will try to track the missing frames based on some static analyses.

llvm/lib/Transforms/IPO/SampleProfileProbe.cpp

137–139

Thanks for the suggestion. This is a C++17 usage and may cause a warning with the current build setup which defaults to C++14 (due to -Wc++17)

[hoy]

Updating D91756: [CSSPGO] Pseudo probes for function calls.

[wmi]

llvm/include/llvm/IR/PseudoProbe.h

34–35

Add assertion messages here and the other places.

llvm/lib/Transforms/IPO/SampleProfileProbe.cpp

137–139

Ah, right. I mistakenly remember I saw that somewhere in llvm codebase. llvm coding standard says we should use c++14.

[hoy]

Updating D91756: [CSSPGO] Pseudo probes for function calls.

Adding messages for asserts.

[wmi]

Another question. Sorry for not bringing it up earlier. When a call with probe metadata attached is inlined, the probe will be gone or it will be kept somehow? I think you want to keep the probe especially for inline instance to reconstruct the context but I didn't figure it out how you did that from the description.

[hoy]

In D91756#2427759, @wmi wrote:

Another question. Sorry for not bringing it up earlier. When a call with probe metadata attached is inlined, the probe will be gone or it will be kept somehow? I think you want to keep the probe especially for inline instance to reconstruct the context but I didn't figure it out how you did that from the description.

No problem. Sorry for not clarifying it in the description. When a callee is inlined, the probe metadata will go with the inlined instructions. The !dbg metadata of an inlined instruction is in form of a scope stack. The top of the stack is the instruction's original !dbg metadata and the bottom of the stack is the for the original callsite of the top-level inliner. Except for the top of the stack, all other elements of the stack actually refer to the nested inlined callsites whose discriminator fields (which actually represents a calliste probe) can be used to represent the inline context of an inlined PseudoProbeInst or a CallInst. I'll update the description.

[wmi]

In D91756#2427795, @hoy wrote:

In D91756#2427759, @wmi wrote:

Another question. Sorry for not bringing it up earlier. When a call with probe metadata attached is inlined, the probe will be gone or it will be kept somehow? I think you want to keep the probe especially for inline instance to reconstruct the context but I didn't figure it out how you did that from the description.

No problem. Sorry for not clarifying it in the description. When a callee is inlined, the probe metadata will go with the inlined instructions. The !dbg metadata of an inlined instruction is in form of a scope stack. The top of the stack is the instruction's original !dbg metadata and the bottom of the stack is the for the original callsite of the top-level inliner. Except for the top of the stack, all other elements of the stack actually refer to the nested inlined callsites whose discriminator fields (which actually represents a calliste probe) can be used to represent the inline context of an inlined PseudoProbeInst or a CallInst. I'll update the description.

I see. Thanks for the explanation.

clang/lib/CodeGen/BackendUtil.cpp

static bool initTargetOptions(DiagnosticsEngine &Diags,
   Options.EmulatedTLS = CodeGenOpts.EmulatedTLS;
   Options.ExplicitEmulatedTLS = CodeGenOpts.ExplicitEmulatedTLS;
   Options.DebuggerTuning = CodeGenOpts.getDebuggerTuning();
   Options.EmitStackSizeSection = CodeGenOpts.StackSizeSection;
   Options.EmitAddrsig = CodeGenOpts.Addrsig;
   Options.ForceDwarfFrameSection = CodeGenOpts.ForceDwarfFrameSection;
   Options.EmitCallSiteInfo = CodeGenOpts.EmitCallSiteInfo;
   Options.EnableAIXExtendedAltivecABI = CodeGenOpts.EnableAIXExtendedAltivecABI;
+  Options.PseudoProbeForProfiling = CodeGenOpts.PseudoProbeForProfiling;
   Options.ValueTrackingVariableLocations =
       CodeGenOpts.ValueTrackingVariableLocations;
   Options.XRayOmitFunctionIndex = CodeGenOpts.XRayOmitFunctionIndex;
 
   Options.MCOptions.SplitDwarfFile = CodeGenOpts.SplitDwarfFile;
   Options.MCOptions.MCRelaxAll = CodeGenOpts.RelaxAll;
   Options.MCOptions.MCSaveTempLabels = CodeGenOpts.SaveTempLabels;
   Options.MCOptions.MCUseDwarfDirectory = !CodeGenOpts.NoDwarfDirectoryAsm;

llvm/include/llvm/CodeGen/CommandFlags.h

 bool getEnableAddrsig();
 
 bool getEmitCallSiteInfo();
 
 bool getEnableMachineFunctionSplitter();
 
 bool getEnableDebugEntryValues();
 
+bool getPseudoProbeForProfiling();
+
 bool getValueTrackingVariableLocations();
 
 bool getForceDwarfFrameSection();
 
 bool getXRayOmitFunctionIndex();
 
 /// Create this object with static storage to register codegen-related command
 /// line options.

llvm/include/llvm/CodeGen/Passes.h

/// MachineDominanaceFrontier - This pass is a machine dominators analysis pass.
 
   /// Creates CFGuard longjmp target identification pass.
   /// \see CFGuardLongjmp.cpp
   FunctionPass *createCFGuardLongjmpPass();
 
   /// Create Hardware Loop pass. \see HardwareLoops.cpp
   FunctionPass *createHardwareLoopsPass();
 
+  /// This pass inserts pseudo probe annotation for callsite profiling.
+  FunctionPass *createPseudoProbeInserter();
+
   /// Create IR Type Promotion pass. \see TypePromotion.cpp
   FunctionPass *createTypePromotionPass();
 
   /// Creates MIR Debugify pass. \see MachineDebugify.cpp
   ModulePass *createDebugifyMachineModulePass();
 
   /// Creates MIR Strip Debug pass. \see MachineStripDebug.cpp
   /// If OnlyDebugified is true then it will only strip debug info if it was

llvm/include/llvm/IR/DebugInfoMetadata.h

public:
   /// component is also capped at 0x1ff, even in the case when both first
   /// components are 0, and we'd technically have 29 bits available.
   ///
   /// For precise control over the data being encoded in the discriminator,
   /// use encodeDiscriminator/decodeDiscriminator.
 
   inline unsigned getDiscriminator() const;
 
+  // For the regular discriminator, it stands for all empty components if all
+  // the lowest 3 bits are non-zero and all higher 29 bits are unused(zero by
+  // default). Here we fully leverage the higher 29 bits for pseudo probe use.
+  // This is the format:
+  // [2:0] - 0x7
+  // [31:3] - pseudo probe fields guaranteed to be non-zero as a whole
+  // So if the lower 3 bits is non-zero and the others has at least one
+  // non-zero bit, it guarantees to be a pseudo probe discriminator
+  inline static bool isPseudoProbeDiscriminator(unsigned Discriminator) {
+    return ((Discriminator & 0x7) == 0x7) && (Discriminator & 0xFFFFFFF8);
+  }
+
   /// Returns a new DILocation with updated \p Discriminator.
   inline const DILocation *cloneWithDiscriminator(unsigned Discriminator) const;
 
   /// Returns a new DILocation with updated base discriminator \p BD. Only the
   /// base discriminator is set in the new DILocation, the other encoded values
   /// are elided.
   /// If the discriminator cannot be encoded, the function returns None.
   inline Optional<const DILocation *> cloneWithBaseDiscriminator(unsigned BD) const;

llvm/include/llvm/IR/PseudoProbe.h

+//===- PseudoProbe.h - Pseudo Probe IR Helpers ------------------*- 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
+//
+//===----------------------------------------------------------------------===//
+//
+// Pseudo probe IR intrinsic and dwarf discriminator manipulation routines.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_IR_PSEUDOPROBE_H
+#define LLVM_IR_PSEUDOPROBE_H
+
+#include <cassert>
+#include <cstdint>
+
+namespace llvm {
+
+enum class PseudoProbeType { Block = 0, IndirectCall, DirectCall };
+
+struct PseudoProbeDwarfDiscriminator {
+  // The following APIs encodes/decodes per-probe information to/from a
+  // 32-bit integer which is organized as:
+  //  [2:0] - 0x7, this is reserved for regular discriminator,
+  //          see DWARF discriminator encoding rule
+  //  [18:3] - probe id
+  //  [25:19] - reserved
+  //  [28:26] - probe type, see PseudoProbeType
+  //  [31:29] - reserved for probe attributes
+  static uint32_t packProbeData(uint32_t Index, uint32_t Type) {
+    assert(Index <= 0xFFFF && "Probe index too big to encode, exceeding 2^16");
+    assert(Type <= 0x7 && "Probe type too big to encode, exceeding 7");
+    return (Index << 3) | (Type << 26) | 0x7;
+  }
+
+  static uint32_t extractProbeIndex(uint32_t Value) {
+    return (Value >> 3) & 0xFFFF;
+  }
+
+  static uint32_t extractProbeType(uint32_t Value) {
+    return (Value >> 26) & 0x7;
+  }
+
+  static uint32_t extractProbeAttributes(uint32_t Value) {
+    return (Value >> 29) & 0x7;
+  }
+};
+} // end namespace llvm
+
+#endif // LLVM_IR_PSEUDOPROBE_H

llvm/include/llvm/InitializePasses.h

 void initializePredicateInfoPrinterLegacyPassPass(PassRegistry&);
 void initializePrintFunctionPassWrapperPass(PassRegistry&);
 void initializePrintModulePassWrapperPass(PassRegistry&);
 void initializeProcessImplicitDefsPass(PassRegistry&);
 void initializeProfileSummaryInfoWrapperPassPass(PassRegistry&);
 void initializePromoteLegacyPassPass(PassRegistry&);
 void initializePruneEHPass(PassRegistry&);
 void initializeRABasicPass(PassRegistry&);
+void initializePseudoProbeInserterPass(PassRegistry &);
 void initializeRAGreedyPass(PassRegistry&);
 void initializeReachingDefAnalysisPass(PassRegistry&);
 void initializeReassociateLegacyPassPass(PassRegistry&);
 void initializeRedundantDbgInstEliminationPass(PassRegistry&);
 void initializeRegAllocFastPass(PassRegistry&);
 void initializeRegBankSelectPass(PassRegistry&);
 void initializeRegToMemLegacyPass(PassRegistry&);
 void initializeRegUsageInfoCollectorPass(PassRegistry&);

llvm/include/llvm/Passes/PassBuilder.h

PGOOptions(std::string ProfileFile = "", std::string CSProfileGenFile = "",
     // If CSAction is CSIRUse, PGOAction needs to be IRUse as they share
     // a profile.
     assert(this->CSAction != CSIRUse || this->Action == IRUse);
 
     // If neither Action nor CSAction, DebugInfoForProfiling or
     // PseudoProbeForProfiling needs to be true.
     assert(this->Action != NoAction || this->CSAction != NoCSAction ||
            this->DebugInfoForProfiling || this->PseudoProbeForProfiling);
+
+    // Pseudo probe emission does work with -fdebug-info-for-profiling since
+    // they both use the discriminator field of debug lines but for different
+    // purposes.
+    if (this->DebugInfoForProfiling && this->PseudoProbeForProfiling) {
+      report_fatal_error(
+          "Pseudo probes cannot be used with -debug-info-for-profiling", false);
+    }
   }
   std::string ProfileFile;
   std::string CSProfileGenFile;
   std::string ProfileRemappingFile;
   PGOAction Action;
   CSPGOAction CSAction;
   bool DebugInfoForProfiling;
   bool PseudoProbeForProfiling;

llvm/include/llvm/Target/TargetOptions.h

TargetOptions()
           IgnoreXCOFFVisibility(false), UniqueSectionNames(true),
           UniqueBasicBlockSectionNames(false), TrapUnreachable(false),
           NoTrapAfterNoreturn(false), TLSSize(0), EmulatedTLS(false),
           ExplicitEmulatedTLS(false), EnableIPRA(false),
           EmitStackSizeSection(false), EnableMachineOutliner(false),
           EnableMachineFunctionSplitter(false), SupportsDefaultOutlining(false),
           EmitAddrsig(false), EmitCallSiteInfo(false),
           SupportsDebugEntryValues(false), EnableDebugEntryValues(false),
-          ValueTrackingVariableLocations(false), ForceDwarfFrameSection(false),
-          XRayOmitFunctionIndex(false),
+          PseudoProbeForProfiling(false), ValueTrackingVariableLocations(false),
+          ForceDwarfFrameSection(false), XRayOmitFunctionIndex(false),
           FPDenormalMode(DenormalMode::IEEE, DenormalMode::IEEE) {}
 
     /// DisableFramePointerElim - This returns true if frame pointer elimination
     /// optimization should be disabled for the given machine function.
     bool DisableFramePointerElim(const MachineFunction &MF) const;
 
     /// UnsafeFPMath - This flag is enabled when the
     /// -enable-unsafe-fp-math flag is specified on the command line.  When
public:
     /// of debug entry values even if the target does not officially support
     /// it. Useful for testing purposes only. This flag should never be checked
     /// directly, always use \ref ShouldEmitDebugEntryValues instead.
      unsigned EnableDebugEntryValues : 1;
     /// NOTE: There are targets that still do not support the debug entry values
     /// production.
     bool ShouldEmitDebugEntryValues() const;
 
+    /// Emit pseudo probes into the binary for sample profiling
+    unsigned PseudoProbeForProfiling : 1;
+
     // When set to true, use experimental new debug variable location tracking,
     // which seeks to follow the values of variables rather than their location,
     // post isel.
     unsigned ValueTrackingVariableLocations : 1;
 
     /// Emit DWARF debug frame section.
     unsigned ForceDwarfFrameSection : 1;
 

llvm/include/llvm/Transforms/IPO/SampleProfileProbe.h

 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_TRANSFORMS_IPO_SAMPLEPROFILEPROBE_H
 #define LLVM_TRANSFORMS_IPO_SAMPLEPROFILEPROBE_H
 
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/IR/PassManager.h"
+#include "llvm/IR/PseudoProbe.h"
 #include "llvm/Target/TargetMachine.h"
 #include <unordered_map>
 
 namespace llvm {
 
 class Module;
 
 using BlockIdMap = std::unordered_map<BasicBlock *, uint32_t>;
+using InstructionIdMap = std::unordered_map<Instruction *, uint32_t>;
 
 enum class PseudoProbeReservedId { Invalid = 0, Last = Invalid };
 
-enum class PseudoProbeType { Block = 0 };
 
 /// Sample profile pseudo prober.
 ///
 /// Insert pseudo probes for block sampling and value sampling.
 class SampleProfileProber {
 public:
   // Give an empty module id when the prober is not used for instrumentation.
   SampleProfileProber(Function &F);
   void instrumentOneFunc(Function &F, TargetMachine *TM);
 
 private:
   Function *getFunction() const { return F; }
   uint32_t getBlockId(const BasicBlock *BB) const;
+  uint32_t getCallsiteId(const Instruction *Call) const;
   void computeProbeIdForBlocks();
+  void computeProbeIdForCallsites();
 
   Function *F;
 
   /// Map basic blocks to the their pseudo probe ids.
   BlockIdMap BlockProbeIds;
 
+  /// Map indirect calls to the their pseudo probe ids.
+  InstructionIdMap CallProbeIds;
+
   /// The ID of the last probe, Can be used to number a new probe.
   uint32_t LastProbeId;
 };
 
 class SampleProfileProbePass : public PassInfoMixin<SampleProfileProbePass> {
   TargetMachine *TM;
 
 public:
   SampleProfileProbePass(TargetMachine *TM) : TM(TM) {}
   PreservedAnalyses run(Module &M, ModuleAnalysisManager &AM);
 };
 
 } // end namespace llvm
 #endif // LLVM_TRANSFORMS_IPO_SAMPLEPROFILEPROBE_H

llvm/lib/CodeGen/CMakeLists.txt

add_llvm_component_library(LLVMCodeGen
   PeepholeOptimizer.cpp
   PHIElimination.cpp
   PHIEliminationUtils.cpp
   PostRAHazardRecognizer.cpp
   PostRASchedulerList.cpp
   PreISelIntrinsicLowering.cpp
   ProcessImplicitDefs.cpp
   PrologEpilogInserter.cpp
+  PseudoProbeInserter.cpp
   PseudoSourceValue.cpp
   RDFGraph.cpp
   RDFLiveness.cpp
   RDFRegisters.cpp
   ReachingDefAnalysis.cpp
   RegAllocBase.cpp
   RegAllocBasic.cpp
   RegAllocFast.cpp

llvm/lib/CodeGen/CommandFlags.cpp

 CGOPT(bool, UniqueBasicBlockSectionNames)
 CGOPT(EABI, EABIVersion)
 CGOPT(DebuggerKind, DebuggerTuningOpt)
 CGOPT(bool, EnableStackSizeSection)
 CGOPT(bool, EnableAddrsig)
 CGOPT(bool, EmitCallSiteInfo)
 CGOPT(bool, EnableMachineFunctionSplitter)
 CGOPT(bool, EnableDebugEntryValues)
+CGOPT(bool, PseudoProbeForProfiling)
 CGOPT(bool, ValueTrackingVariableLocations)
 CGOPT(bool, ForceDwarfFrameSection)
 CGOPT(bool, XRayOmitFunctionIndex)
 
 codegen::RegisterCodeGenFlags::RegisterCodeGenFlags() {
 #define CGBINDOPT(NAME)                                                        \
   do {                                                                         \
     NAME##View = std::addressof(NAME);                                         \
#define CGBINDOPT(NAME)                                                        \
   CGBINDOPT(EmitCallSiteInfo);
 
   static cl::opt<bool> EnableDebugEntryValues(
       "debug-entry-values",
       cl::desc("Enable debug info for the debug entry values."),
       cl::init(false));
   CGBINDOPT(EnableDebugEntryValues);
 
+  static cl::opt<bool> PseudoProbeForProfiling(
+      "pseudo-probe-for-profiling", cl::desc("Emit pseudo probes for AutoFDO"),
+      cl::init(false));
+  CGBINDOPT(PseudoProbeForProfiling);
+
   static cl::opt<bool> ValueTrackingVariableLocations(
       "experimental-debug-variable-locations",
       cl::desc("Use experimental new value-tracking variable locations"),
       cl::init(false));
   CGBINDOPT(ValueTrackingVariableLocations);
 
   static cl::opt<bool> EnableMachineFunctionSplitter(
       "split-machine-functions",
codegen::InitTargetOptionsFromCodeGenFlags(const Triple &TheTriple) {
   Options.EmulatedTLS = getEmulatedTLS();
   Options.ExplicitEmulatedTLS = EmulatedTLSView->getNumOccurrences() > 0;
   Options.ExceptionModel = getExceptionModel();
   Options.EmitStackSizeSection = getEnableStackSizeSection();
   Options.EnableMachineFunctionSplitter = getEnableMachineFunctionSplitter();
   Options.EmitAddrsig = getEnableAddrsig();
   Options.EmitCallSiteInfo = getEmitCallSiteInfo();
   Options.EnableDebugEntryValues = getEnableDebugEntryValues();
+  Options.PseudoProbeForProfiling = getPseudoProbeForProfiling();
   Options.ValueTrackingVariableLocations = getValueTrackingVariableLocations();
   Options.ForceDwarfFrameSection = getForceDwarfFrameSection();
   Options.XRayOmitFunctionIndex = getXRayOmitFunctionIndex();
 
   Options.MCOptions = mc::InitMCTargetOptionsFromFlags();
 
   Options.ThreadModel = getThreadModel();
   Options.EABIVersion = getEABIVersion();

llvm/lib/CodeGen/PseudoProbeInserter.cpp

+//===- PseudoProbeInserter.cpp - Insert annotation for callsite profiling -===//
+//
+// 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 PseudoProbeInserter pass, which inserts pseudo probe
+// annotations for call instructions with a pseudo-probe-specific dwarf
+// discriminator. such discriminator indicates that the call instruction comes
+// with a pseudo probe, and the discriminator value holds information to
+// identify the corresponding counter.
+//===----------------------------------------------------------------------===//
+
+#include "llvm/CodeGen/MachineBasicBlock.h"
+#include "llvm/CodeGen/MachineFunctionPass.h"
+#include "llvm/CodeGen/MachineInstr.h"
+#include "llvm/CodeGen/TargetInstrInfo.h"
+#include "llvm/IR/DebugInfoMetadata.h"
+#include "llvm/IR/PseudoProbe.h"
+#include "llvm/InitializePasses.h"
+#include "llvm/Target/TargetMachine.h"
+#include <unordered_map>
+
+#define DEBUG_TYPE "pseudo-probe-inserter"
+
+using namespace llvm;
+
+namespace {
+class PseudoProbeInserter : public MachineFunctionPass {
+public:
+  static char ID;
+
+  PseudoProbeInserter() : MachineFunctionPass(ID) {
+    initializePseudoProbeInserterPass(*PassRegistry::getPassRegistry());
+  }
+
+  StringRef getPassName() const override { return "Pseudo Probe Inserter"; }
+
+  void getAnalysisUsage(AnalysisUsage &AU) const override {
+    AU.setPreservesAll();
+    MachineFunctionPass::getAnalysisUsage(AU);
+  }
+
+  bool runOnMachineFunction(MachineFunction &MF) override {
+    const TargetInstrInfo *TII = MF.getSubtarget().getInstrInfo();
+    bool Changed = false;
+    for (MachineBasicBlock &MBB : MF) {
+      for (MachineInstr &MI : MBB) {
+        if (MI.isCall()) {
+          if (DILocation *DL = MI.getDebugLoc()) {
+            auto Value = DL->getDiscriminator();
+            if (DILocation::isPseudoProbeDiscriminator(Value)) {
+              BuildMI(MBB, MI, DL, TII->get(TargetOpcode::PSEUDO_PROBE))
+                  .addImm(getFuncGUID(MF.getFunction().getParent(), DL))
+                  .addImm(
+                      PseudoProbeDwarfDiscriminator::extractProbeIndex(Value))
+                  .addImm(
+                      PseudoProbeDwarfDiscriminator::extractProbeType(Value))
+                  .addImm(PseudoProbeDwarfDiscriminator::extractProbeAttributes(
+                      Value));
+              Changed = true;
+            }
+          }
+        }
+      }
+    }
+
+    return Changed;
+  }
+
+private:
+  uint64_t getFuncGUID(Module *M, DILocation *DL) {
+    auto *SP = DL->getScope()->getSubprogram();
+    auto Name = SP->getLinkageName();
+    if (Name.empty())
+      Name = SP->getName();
+    return Function::getGUID(Name);
+  }
+};
+} // namespace
+
+char PseudoProbeInserter::ID = 0;
+INITIALIZE_PASS_BEGIN(PseudoProbeInserter, DEBUG_TYPE,
+                      "Insert pseudo probe annotations for value profiling",
+                      false, false)
+INITIALIZE_PASS_DEPENDENCY(TargetPassConfig)
+INITIALIZE_PASS_END(PseudoProbeInserter, DEBUG_TYPE,
+                    "Insert pseudo probe annotations for value profiling",
+                    false, false)
+
+FunctionPass *llvm::createPseudoProbeInserter() {
+  return new PseudoProbeInserter();
+}

llvm/lib/CodeGen/SelectionDAG/InstrEmitter.cpp

 #include "llvm/CodeGen/MachineRegisterInfo.h"
 #include "llvm/CodeGen/SelectionDAG.h"
 #include "llvm/CodeGen/StackMaps.h"
 #include "llvm/CodeGen/TargetInstrInfo.h"
 #include "llvm/CodeGen/TargetLowering.h"
 #include "llvm/CodeGen/TargetSubtargetInfo.h"
 #include "llvm/IR/DataLayout.h"
 #include "llvm/IR/DebugInfo.h"
+#include "llvm/IR/PseudoProbe.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/MathExtras.h"
 #include "llvm/Target/TargetMachine.h"
 using namespace llvm;
 
 #define DEBUG_TYPE "instr-emitter"
 
case ISD::PSEUDO_PROBE: {
     unsigned TarOp = TargetOpcode::PSEUDO_PROBE;
     auto Guid = cast<PseudoProbeSDNode>(Node)->getGuid();
     auto Index = cast<PseudoProbeSDNode>(Node)->getIndex();
     auto Attr = cast<PseudoProbeSDNode>(Node)->getAttributes();
 
     BuildMI(*MBB, InsertPos, Node->getDebugLoc(), TII->get(TarOp))
         .addImm(Guid)
         .addImm(Index)
-        .addImm(0) // 0 for block probes
+        .addImm((uint8_t)PseudoProbeType::Block)
         .addImm(Attr);
     break;
   }
 
   case ISD::INLINEASM:
   case ISD::INLINEASM_BR: {
     unsigned NumOps = Node->getNumOperands();
     if (Node->getOperand(NumOps-1).getValueType() == MVT::Glue)

llvm/lib/CodeGen/TargetPassConfig.cpp

if (TM->Options.EnableMachineFunctionSplitter ||
     addPass(createMachineFunctionSplitterPass());
   } else if (TM->getBBSectionsType() != llvm::BasicBlockSection::None) {
     addPass(llvm::createBasicBlockSectionsPass(TM->getBBSectionsFuncListBuf()));
   }
 
   // Add passes that directly emit MI after all other MI passes.
   addPreEmitPass2();
 
+  // Insert pseudo probe annotation for callsite profiling
+  if (TM->Options.PseudoProbeForProfiling)
+    addPass(createPseudoProbeInserter());
+
   AddingMachinePasses = false;
 }
 
 /// Add passes that optimize machine instructions in SSA form.
 void TargetPassConfig::addMachineSSAOptimization() {
   // Pre-ra tail duplication.
   addPass(&EarlyTailDuplicateID);
 

llvm/lib/Target/X86/X86TargetMachine.cpp

extern "C" LLVM_EXTERNAL_VISIBILITY void LLVMInitializeX86Target() {
   initializeX86AvoidTrailingCallPassPass(PR);
   initializeX86SpeculativeLoadHardeningPassPass(PR);
   initializeX86SpeculativeExecutionSideEffectSuppressionPass(PR);
   initializeX86FlagsCopyLoweringPassPass(PR);
   initializeX86LoadValueInjectionLoadHardeningPassPass(PR);
   initializeX86LoadValueInjectionRetHardeningPassPass(PR);
   initializeX86OptimizeLEAPassPass(PR);
   initializeX86PartialReductionPass(PR);
+  initializePseudoProbeInserterPass(PR);
 }
 
 static std::unique_ptr<TargetLoweringObjectFile> createTLOF(const Triple &TT) {
   if (TT.isOSBinFormatMachO()) {
     if (TT.getArch() == Triple::x86_64)
       return std::make_unique<X86_64MachoTargetObjectFile>();
     return std::make_unique<TargetLoweringObjectFileMachO>();
   }

llvm/lib/Transforms/IPO/SampleProfileProbe.cpp

 using namespace llvm;
 #define DEBUG_TYPE "sample-profile-probe"
 
 STATISTIC(ArtificialDbgLine,
           "Number of probes that have an artificial debug line");
 
 SampleProfileProber::SampleProfileProber(Function &Func) : F(&Func) {
   BlockProbeIds.clear();
+  CallProbeIds.clear();
   LastProbeId = (uint32_t)PseudoProbeReservedId::Last;
   computeProbeIdForBlocks();
+  computeProbeIdForCallsites();
 }
 
 void SampleProfileProber::computeProbeIdForBlocks() {
   for (auto &BB : *F) {
     BlockProbeIds[&BB] = ++LastProbeId;
   }
 }
 
+void SampleProfileProber::computeProbeIdForCallsites() {
+  for (auto &BB : *F) {
+    for (auto &I : BB) {
+      if (!isa<CallBase>(I))
+        continue;
+      if (isa<IntrinsicInst>(&I))
+        continue;
+      CallProbeIds[&I] = ++LastProbeId;
+    }
+  }
+}
+
 uint32_t SampleProfileProber::getBlockId(const BasicBlock *BB) const {
   auto I = BlockProbeIds.find(const_cast<BasicBlock *>(BB));
   return I == BlockProbeIds.end() ? 0 : I->second;
 }
 
+uint32_t SampleProfileProber::getCallsiteId(const Instruction *Call) const {
+  auto Iter = CallProbeIds.find(const_cast<Instruction *>(Call));
+  return Iter == CallProbeIds.end() ? 0 : Iter->second;
+}
+
 void SampleProfileProber::instrumentOneFunc(Function &F, TargetMachine *TM) {
   Module *M = F.getParent();
   MDBuilder MDB(F.getContext());
   // Compute a GUID without considering the function's linkage type. This is
   // fine since function name is the only key in the profile database.
   uint64_t Guid = Function::getGUID(F.getName());
 
+  // Assign an artificial debug line to a probe that doesn't come with a real
+  // line. A probe not having a debug line will get an incomplete inline
+  // context. This will cause samples collected on the probe to be counted
+  // into the base profile instead of a context profile. The line number
+  // itself is not important though.
+  auto AssignDebugLoc = [&](Instruction *I) {
+    assert((isa<PseudoProbeInst>(I) || isa<CallBase>(I)) &&
+           "Expecting pseudo probe or call instructions");
+    if (!I->getDebugLoc()) {
+      if (auto *SP = F.getSubprogram()) {
+        auto DIL = DebugLoc::get(0, 0, SP);
+        I->setDebugLoc(DIL);
+        ArtificialDbgLine++;
+        LLVM_DEBUG({
+          dbgs() << "\nIn Function " << F.getName()
+                 << " Probe gets an artificial debug line\n";
+          I->dump();
+        });
+      }
+    }
+  };
+
   // Probe basic blocks.
   for (auto &I : BlockProbeIds) {
     BasicBlock *BB = I.first;
     uint32_t Index = I.second;
     // Insert a probe before an instruction with a valid debug line number which
     // will be assigned to the probe. The line number will be used later to
     // model the inline context when the probe is inlined into other functions.
     // Debug instructions, phi nodes and lifetime markers do not have an valid
for (auto &I : BlockProbeIds) {
     IRBuilder<> Builder(J);
     assert(Builder.GetInsertPoint() != BB->end() &&
            "Cannot get the probing point");
     Function *ProbeFn =
         llvm::Intrinsic::getDeclaration(M, Intrinsic::pseudoprobe);
     Value *Args[] = {Builder.getInt64(Guid), Builder.getInt64(Index),
                      Builder.getInt32(0)};
     auto *Probe = Builder.CreateCall(ProbeFn, Args);
-    // Assign an artificial debug line to a probe that doesn't come with a real
-    // line. A probe not having a debug line will get an incomplete inline
-    // context. This will cause samples collected on the probe to be counted
-    // into the base profile instead of a context profile. The line number
-    // itself is not important though.
-    if (!Probe->getDebugLoc()) {
-      if (auto *SP = F.getSubprogram()) {
-        auto DIL = DebugLoc::get(0, 0, SP);
-        Probe->setDebugLoc(DIL);
-        ArtificialDbgLine++;
-        LLVM_DEBUG(dbgs() << "\nIn Function " << F.getName() << " Probe "
-                          << Index << " gets an artificial debug line\n";);
-      }
+    AssignDebugLoc(Probe);
+
+  // Probe both direct calls and indirect calls. Direct calls are probed so that
+  // their probe ID can be used as an call site identifier to represent a
+  // calling context.
+  for (auto &I : CallProbeIds) {
+    auto *Call = I.first;
+    uint32_t Index = I.second;
+    uint32_t Type = cast<CallBase>(Call)->getCalledFunction()
+                        ? (uint32_t)PseudoProbeType::DirectCall
+                        : (uint32_t)PseudoProbeType::IndirectCall;
+    AssignDebugLoc(Call);
+    // Levarge the 32-bit discriminator field of debug data to store the ID and
+    // type of a callsite probe. This gets rid of the dependency on plumbing a
+    // customized metadata through the codegen pipeline.
+    uint32_t V = PseudoProbeDwarfDiscriminator::packProbeData(Index, Type);
+    if (auto DIL = Call->getDebugLoc()) {
+      DIL = DIL->cloneWithDiscriminator(V);
+      Call->setDebugLoc(DIL);
     }
   }
 }
 
 PreservedAnalyses SampleProfileProbePass::run(Module &M,
                                               ModuleAnalysisManager &AM) {
   for (auto &F : M) {
     if (F.isDeclaration())
       continue;
     SampleProfileProber ProbeManager(F);
     ProbeManager.instrumentOneFunc(F, TM);
   }
 
   return PreservedAnalyses::none();
 }

llvm/test/Transforms/SampleProfile/pseudo-probe-emit.ll

 ; REQUIRES: x86-registered-target
 ; RUN: opt < %s -passes=pseudo-probe -function-sections -S -o %t
 ; RUN: FileCheck %s < %t --check-prefix=CHECK-IL
-; RUN: llc %t -stop-after=instruction-select -o - | FileCheck %s --check-prefix=CHECK-MIR
+; RUN: llc %t -pseudo-probe-for-profiling -stop-after=pseudo-probe-inserter -o - | FileCheck %s --check-prefix=CHECK-MIR
 ;
 ;; Check the generation of pseudoprobe intrinsic call.
 
 define void @foo(i32 %x) !dbg !3 {
 bb0:
   %cmp = icmp eq i32 %x, 0
 ; CHECK-IL: call void @llvm.pseudoprobe(i64 [[#GUID:]], i64 1, i32 0), !dbg ![[#FAKELINE:]]
 ; CHECK-MIR: PSEUDO_PROBE [[#GUID:]], 1, 0, 0

 ; CHECK-MIR: PSEUDO_PROBE [[#GUID]], 4, 0, 0
   br label %bb3
 
 bb3:
 ; CHECK-IL: call void @llvm.pseudoprobe(i64 [[#GUID]], i64 4, i32 0), !dbg ![[#REALLINE:]]
   ret void, !dbg !12
 }
 
+declare void @bar(i32 %x) 
+
+define internal void @foo2(void (i32)* %f) !dbg !4 {
+entry:
+; CHECK-IL: call void @llvm.pseudoprobe(i64 [[#GUID2:]], i64 1, i32 0)
+; CHECK-MIR: PSEUDO_PROBE [[#GUID2:]], 1, 0, 0
+; Check pseudo_probe metadata attached to the indirect call instruction.
+; CHECK-IL: call void %f(i32 1), !dbg ![[#PROBE0:]]
+; CHECK-MIR: PSEUDO_PROBE [[#GUID2]], 2, 1, 0
+  call void %f(i32 1), !dbg !13
+; Check pseudo_probe metadata attached to the direct call instruction.
+; CHECK-IL: call void @bar(i32 1), !dbg ![[#PROBE1:]]
+; CHECK-MIR: PSEUDO_PROBE [[#GUID2]], 3, 2, 0
+  call void @bar(i32 1)
+  ret void
+}
+
 ; CHECK-IL: ![[#FOO:]] = distinct !DISubprogram(name: "foo"
 ; CHECK-IL: ![[#FAKELINE]] = !DILocation(line: 0, scope: ![[#FOO]])
 ; CHECK-IL: ![[#REALLINE]] = !DILocation(line: 2, scope: ![[#FOO]])
+; CHECK-IL: ![[#PROBE0]] = !DILocation(line: 2, column: 20, scope: ![[#SCOPE0:]])
+;; A discriminator of 67108887 which is 0x4000017 in hexdecimal, stands for a direct call probe
+;; with an index of 2.
+; CHECK-IL: ![[#SCOPE0]] = !DILexicalBlockFile(scope: ![[#]], file: ![[#]], discriminator: 67108887)
+; CHECK-IL: ![[#PROBE1]] = !DILocation(line: 0, scope: ![[#SCOPE1:]])
+;; A discriminator of 134217759 which is 0x800001f in hexdecimal, stands for a direct call probe
+;; with an index of 3.
+; CHECK-IL: ![[#SCOPE1]] = !DILexicalBlockFile(scope: ![[#]], file: ![[#]], discriminator: 134217759)
+
 
 !llvm.dbg.cu = !{!0}
 !llvm.module.flags = !{!9, !10}
 
 !0 = distinct !DICompileUnit(language: DW_LANG_C99, file: !1)
 !1 = !DIFile(filename: "test.c", directory: "")
 !2 = !{}
 !3 = distinct !DISubprogram(name: "foo", scope: !1, file: !1, line: 1, type: !5, unit: !0, retainedNodes: !2)
+!4 = distinct !DISubprogram(name: "foo2", scope: !1, file: !1, line: 2, type: !5, unit: !0, retainedNodes: !2)
 !5 = !DISubroutineType(types: !6)
 !6 = !{!7}
 !7 = !DIBasicType(name: "int", size: 32, align: 32, encoding: DW_ATE_signed)
 !9 = !{i32 2, !"Dwarf Version", i32 4}
 !10 = !{i32 2, !"Debug Info Version", i32 3}
 !11 = !{!"clang version 3.9.0"}
-!12 = !DILocation(line: 2, scope: !3)
+!12 = !DILocation(line: 2, scope: !3)
+!13 = !DILocation(line: 2, column: 20, scope: !4)