[CSSPGO] Pseudo probe instrumentation pass

[hoy]

Summary

This change introduces a pseudo probe instrumentation pass for block instrumentation. Please refer to https://reviews.llvm.org/D86193 for the whole story.

Given the following LLVM IR:

define internal void @foo2(i32 %x, void (i32)* %f) !dbg !4 {
bb0:
  %cmp = icmp eq i32 %x, 0
   br i1 %cmp, label %bb1, label %bb2
bb1:
   br label %bb3
bb2:
   br label %bb3
bb3:
   ret void
}

The instrumented IR will look like below. Note that each llvm.pseudoprobe intrinsic call represents a pseudo probe at a block, of which the first parameter is the GUID of the probe’s owner function and the second parameter is the probe’s ID.

define internal void @foo2(i32 %x, void (i32)* %f) !dbg !4 {
bb0:
   %cmp = icmp eq i32 %x, 0
   call void @llvm.pseudoprobe(i64 837061429793323041, i64 1)
   br i1 %cmp, label %bb1, label %bb2
bb1:                                             
   call void @llvm.pseudoprobe(i64 837061429793323041, i64 2)
   br label %bb3
bb2:                                              
   call void @llvm.pseudoprobe(i64 837061429793323041, i64 3)
   br label %bb3
bb3:                                              
   call void @llvm.pseudoprobe(i64 837061429793323041, i64 4)
   ret void
}

[hoy]

Updating D86499: [CSSPGO] Pseudo probe instrumentation pass

[wmi]

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

37

Will the class be used in files other than SampleProfileProbe.cpp? If not, it is better to be placed in SampleProfileProbe.cpp.

llvm/lib/Transforms/IPO/SampleProfileProbe.cpp

74–77

Can we add a stat support here to see how frequent the probe cannot find an instruction with valid dbgline? Even if no instruction with valid dbgline can be found, is it possible to create a dbgline for it and set the line number to 0 or to the function start (I assume line number is of no use. Only the inline context matters here)?

[hoy]

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

37

Yes, it will be used outside SampleProfileProbe.cpp. It'll later on introduce a field storing the name of IL metadata that'll be used by later probe emission pass.

llvm/lib/Transforms/IPO/SampleProfileProbe.cpp

74–77

This is a very good suggestion, thanks!

The prober pass had not been placed in the very beginning of the pipeline initially and we saw quite some instructions without a real debug line number. It is now placed in the beginning but with some statistics I was still seeing that happened to a few instructions. Giving an artificial line number to such probes will help.

[hoy]

Updating D86499: [CSSPGO] Pseudo probe instrumentation pass

[hoy]

Updating D86499: [CSSPGO] Pseudo probe instrumentation pass

Corresponding changes to support IR/MIR intrinsic attributes.

[wmi]

LGTM.

[hoy]

Updating D86499: [CSSPGO] Pseudo probe instrumentation pass

Fixing clang-tidy issues.

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

+//===- Transforms/IPO/SampleProfileProbe.h ----------*- 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
+//
+//===----------------------------------------------------------------------===//
+//
+/// \file
+/// This file provides the interface for the pseudo probe implementation for
+/// AutoFDO.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_TRANSFORMS_IPO_SAMPLEPROFILEPROBE_H
+#define LLVM_TRANSFORMS_IPO_SAMPLEPROFILEPROBE_H
+
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/IR/PassManager.h"
+#include "llvm/Target/TargetMachine.h"
+#include <unordered_map>
+
+namespace llvm {
+
+class Module;
+
+using BlockIdMap = std::unordered_map<BasicBlock *, 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;
+  void computeProbeIdForBlocks();
+
+  Function *F;
+
+  /// Map basic blocks to the their pseudo probe ids.
+  BlockIdMap BlockProbeIds;
+
+  /// 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/Passes/PassBuilder.cpp

 #include "llvm/Transforms/IPO/Internalize.h"
 #include "llvm/Transforms/IPO/LoopExtractor.h"
 #include "llvm/Transforms/IPO/LowerTypeTests.h"
 #include "llvm/Transforms/IPO/MergeFunctions.h"
 #include "llvm/Transforms/IPO/OpenMPOpt.h"
 #include "llvm/Transforms/IPO/PartialInlining.h"
 #include "llvm/Transforms/IPO/SCCP.h"
 #include "llvm/Transforms/IPO/SampleProfile.h"
+#include "llvm/Transforms/IPO/SampleProfileProbe.h"
 #include "llvm/Transforms/IPO/StripDeadPrototypes.h"
 #include "llvm/Transforms/IPO/StripSymbols.h"
 #include "llvm/Transforms/IPO/SyntheticCountsPropagation.h"
 #include "llvm/Transforms/IPO/WholeProgramDevirt.h"
 #include "llvm/Transforms/InstCombine/InstCombine.h"
 #include "llvm/Transforms/Instrumentation.h"
 #include "llvm/Transforms/Instrumentation/AddressSanitizer.h"
 #include "llvm/Transforms/Instrumentation/BoundsChecking.h"

llvm/lib/Passes/PassRegistry.def

 MODULE_PASS("sample-profile", SampleProfileLoaderPass())
 MODULE_PASS("scc-oz-module-inliner",
   buildInlinerPipeline(OptimizationLevel::Oz, ThinLTOPhase::None))
 MODULE_PASS("loop-extract-single", LoopExtractorPass(1))
 MODULE_PASS("oz-module-optimizer",
   buildModuleOptimizationPipeline(OptimizationLevel::Oz, /*LTOPreLink*/false))
 MODULE_PASS("strip", StripSymbolsPass())
 MODULE_PASS("strip-dead-debug-info", StripDeadDebugInfoPass())
+MODULE_PASS("pseudo-probe", SampleProfileProbePass(TM))
 MODULE_PASS("strip-dead-prototypes", StripDeadPrototypesPass())
 MODULE_PASS("strip-debug-declare", StripDebugDeclarePass())
 MODULE_PASS("strip-nondebug", StripNonDebugSymbolsPass())
 MODULE_PASS("strip-nonlinetable-debuginfo", StripNonLineTableDebugInfoPass())
 MODULE_PASS("synthetic-counts-propagation", SyntheticCountsPropagation())
 MODULE_PASS("unique-internal-linkage-names", UniqueInternalLinkageNamesPass())
 MODULE_PASS("verify", VerifierPass())
 MODULE_PASS("wholeprogramdevirt", WholeProgramDevirtPass())

llvm/lib/Transforms/IPO/CMakeLists.txt

add_llvm_component_library(LLVMipo
   LoopExtractor.cpp
   LowerTypeTests.cpp
   MergeFunctions.cpp
   OpenMPOpt.cpp
   PartialInlining.cpp
   PassManagerBuilder.cpp
   PruneEH.cpp
   SampleProfile.cpp
+  SampleProfileProbe.cpp
   SCCP.cpp
   StripDeadPrototypes.cpp
   StripSymbols.cpp
   SyntheticCountsPropagation.cpp
   ThinLTOBitcodeWriter.cpp
   WholeProgramDevirt.cpp
 
   ADDITIONAL_HEADER_DIRS

llvm/lib/Transforms/IPO/SampleProfileProbe.cpp

+//===- SampleProfileProbe.cpp - Pseudo probe Instrumentation -------------===//
+//
+// 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 SampleProfileProber transformation.
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/Transforms/IPO/SampleProfileProbe.h"
+#include "llvm/ADT/Statistic.h"
+#include "llvm/Analysis/TargetLibraryInfo.h"
+#include "llvm/IR/BasicBlock.h"
+#include "llvm/IR/CFG.h"
+#include "llvm/IR/Constant.h"
+#include "llvm/IR/Constants.h"
+#include "llvm/IR/DebugInfoMetadata.h"
+#include "llvm/IR/GlobalValue.h"
+#include "llvm/IR/GlobalVariable.h"
+#include "llvm/IR/IRBuilder.h"
+#include "llvm/IR/Instruction.h"
+#include "llvm/IR/MDBuilder.h"
+#include "llvm/ProfileData/SampleProf.h"
+#include "llvm/Support/CRC.h"
+#include "llvm/Transforms/Instrumentation.h"
+#include "llvm/Transforms/Utils/ModuleUtils.h"
+#include <vector>
+
+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();
+  LastProbeId = (uint32_t)PseudoProbeReservedId::Last;
+  computeProbeIdForBlocks();
+}
+
+void SampleProfileProber::computeProbeIdForBlocks() {
+  for (auto &BB : *F) {
+    BlockProbeIds[&BB] = ++LastProbeId;
+  }
+}
+
+uint32_t SampleProfileProber::getBlockId(const BasicBlock *BB) const {
+  auto I = BlockProbeIds.find(const_cast<BasicBlock *>(BB));
+  return I == BlockProbeIds.end() ? 0 : I->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());
+
+  // 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
+    // line number. Real instructions generated by optimizations may not come
+    // with a line number either.
+    auto HasValidDbgLine = [](Instruction *J) {
+      return !isa<PHINode>(J) && !isa<DbgInfoIntrinsic>(J) &&
+             !J->isLifetimeStartOrEnd() && J->getDebugLoc();
+    };
+
+    Instruction *J = &*BB->getFirstInsertionPt();
+    while (J != BB->getTerminator() && !HasValidDbgLine(J)) {
+      J = J->getNextNode();
+    }
+
+    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";);
+      }
+    }
+  }
+}
+
+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

+; 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
+;
+;; 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
+  br i1 %cmp, label %bb1, label %bb2
+
+bb1:
+; CHECK-IL: call void @llvm.pseudoprobe(i64 [[#GUID:]], i64 2, i32 0), !dbg ![[#FAKELINE]]
+; CHECK-MIR: PSEUDO_PROBE [[#GUID]], 3, 0, 0
+; CHECK-MIR: PSEUDO_PROBE [[#GUID]], 4, 0, 0
+  br label %bb3
+
+bb2:
+; CHECK-IL: call void @llvm.pseudoprobe(i64 [[#GUID:]], i64 3, i32 0), !dbg ![[#FAKELINE]]
+; CHECK-MIR: PSEUDO_PROBE [[#GUID]], 2, 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
+}
+
+; CHECK-IL: ![[#FOO:]] = distinct !DISubprogram(name: "foo"
+; CHECK-IL: ![[#FAKELINE]] = !DILocation(line: 0, scope: ![[#FOO]])
+; CHECK-IL: ![[#REALLINE]] = !DILocation(line: 2, scope: ![[#FOO]])
+
+!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)
+!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)
 No newline at end of file