1 //===- Transforms/IPO/SampleProfileProbe.h ----------*- C++ -*-===//
2 //
3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4 // See https://llvm.org/LICENSE.txt for license information.
5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6 //
7 //===----------------------------------------------------------------------===//
8 //
9 /// \file
10 /// This file provides the interface for the pseudo probe implementation for
11 /// AutoFDO.
12 //
13 //===----------------------------------------------------------------------===//
14 
15 #ifndef LLVM_TRANSFORMS_IPO_SAMPLEPROFILEPROBE_H
16 #define LLVM_TRANSFORMS_IPO_SAMPLEPROFILEPROBE_H
17 
18 #include "llvm/ADT/DenseMap.h"
19 #include "llvm/IR/PassManager.h"
20 #include "llvm/IR/PseudoProbe.h"
21 #include "llvm/Target/TargetMachine.h"
22 #include <unordered_map>
23 
24 namespace llvm {
25 
26 class Module;
27 
28 using BlockIdMap = std::unordered_map<BasicBlock *, uint32_t>;
29 using InstructionIdMap = std::unordered_map<Instruction *, uint32_t>;
30 
31 enum class PseudoProbeReservedId { Invalid = 0, Last = Invalid };
32 
33 
34 /// Sample profile pseudo prober.
35 ///
36 /// Insert pseudo probes for block sampling and value sampling.
37 class SampleProfileProber {
38 public:
39   // Give an empty module id when the prober is not used for instrumentation.
40   SampleProfileProber(Function &F);
41   void instrumentOneFunc(Function &F, TargetMachine *TM);
42 
43 private:
getFunction()44   Function *getFunction() const { return F; }
45   uint32_t getBlockId(const BasicBlock *BB) const;
46   uint32_t getCallsiteId(const Instruction *Call) const;
47   void computeProbeIdForBlocks();
48   void computeProbeIdForCallsites();
49 
50   Function *F;
51 
52   /// Map basic blocks to the their pseudo probe ids.
53   BlockIdMap BlockProbeIds;
54 
55   /// Map indirect calls to the their pseudo probe ids.
56   InstructionIdMap CallProbeIds;
57 
58   /// The ID of the last probe, Can be used to number a new probe.
59   uint32_t LastProbeId;
60 };
61 
62 class SampleProfileProbePass : public PassInfoMixin<SampleProfileProbePass> {
63   TargetMachine *TM;
64 
65 public:
SampleProfileProbePass(TargetMachine * TM)66   SampleProfileProbePass(TargetMachine *TM) : TM(TM) {}
67   PreservedAnalyses run(Module &M, ModuleAnalysisManager &AM);
68 };
69 
70 } // end namespace llvm
71 #endif // LLVM_TRANSFORMS_IPO_SAMPLEPROFILEPROBE_H
72