1 //===- Transforms/Utils/SampleProfileInference.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 profile inference algorithm, profi.
11 //
12 //===----------------------------------------------------------------------===//
13 
14 #ifndef LLVM_TRANSFORMS_UTILS_SAMPLEPROFILEINFERENCE_H
15 #define LLVM_TRANSFORMS_UTILS_SAMPLEPROFILEINFERENCE_H
16 
17 #include "llvm/ADT/DenseMap.h"
18 #include "llvm/ADT/DepthFirstIterator.h"
19 #include "llvm/ADT/SmallVector.h"
20 
21 #include "llvm/IR/BasicBlock.h"
22 #include "llvm/IR/Instruction.h"
23 #include "llvm/IR/Instructions.h"
24 
25 namespace llvm {
26 
27 class Function;
28 class MachineBasicBlock;
29 class MachineFunction;
30 
31 namespace afdo_detail {
32 
33 template <class BlockT> struct TypeMap {};
34 template <> struct TypeMap<BasicBlock> {
35   using BasicBlockT = BasicBlock;
36   using FunctionT = Function;
37 };
38 template <> struct TypeMap<MachineBasicBlock> {
39   using BasicBlockT = MachineBasicBlock;
40   using FunctionT = MachineFunction;
41 };
42 
43 } // end namespace afdo_detail
44 
45 struct FlowJump;
46 
47 /// A wrapper of a binary basic block.
48 struct FlowBlock {
49   uint64_t Index;
50   uint64_t Weight{0};
51   bool UnknownWeight{false};
52   uint64_t Flow{0};
53   bool HasSelfEdge{false};
54   std::vector<FlowJump *> SuccJumps;
55   std::vector<FlowJump *> PredJumps;
56 
57   /// Check if it is the entry block in the function.
58   bool isEntry() const { return PredJumps.empty(); }
59 
60   /// Check if it is an exit block in the function.
61   bool isExit() const { return SuccJumps.empty(); }
62 };
63 
64 /// A wrapper of a jump between two basic blocks.
65 struct FlowJump {
66   uint64_t Source;
67   uint64_t Target;
68   uint64_t Flow{0};
69   bool IsUnlikely{false};
70 };
71 
72 /// A wrapper of binary function with basic blocks and jumps.
73 struct FlowFunction {
74   std::vector<FlowBlock> Blocks;
75   std::vector<FlowJump> Jumps;
76   /// The index of the entry block.
77   uint64_t Entry;
78 };
79 
80 void applyFlowInference(FlowFunction &Func);
81 
82 /// Sample profile inference pass.
83 template <typename BT> class SampleProfileInference {
84 public:
85   using BasicBlockT = typename afdo_detail::TypeMap<BT>::BasicBlockT;
86   using FunctionT = typename afdo_detail::TypeMap<BT>::FunctionT;
87   using Edge = std::pair<const BasicBlockT *, const BasicBlockT *>;
88   using BlockWeightMap = DenseMap<const BasicBlockT *, uint64_t>;
89   using EdgeWeightMap = DenseMap<Edge, uint64_t>;
90   using BlockEdgeMap =
91       DenseMap<const BasicBlockT *, SmallVector<const BasicBlockT *, 8>>;
92 
93   SampleProfileInference(FunctionT &F, BlockEdgeMap &Successors,
94                          BlockWeightMap &SampleBlockWeights)
95       : F(F), Successors(Successors), SampleBlockWeights(SampleBlockWeights) {}
96 
97   /// Apply the profile inference algorithm for a given function
98   void apply(BlockWeightMap &BlockWeights, EdgeWeightMap &EdgeWeights);
99 
100 private:
101   /// Try to infer branch probabilities mimicking implementation of
102   /// BranchProbabilityInfo. Unlikely taken branches are marked so that the
103   /// inference algorithm can avoid sending flow along corresponding edges.
104   void findUnlikelyJumps(const std::vector<const BasicBlockT *> &BasicBlocks,
105                          BlockEdgeMap &Successors, FlowFunction &Func);
106 
107   /// Determine whether the block is an exit in the CFG.
108   bool isExit(const BasicBlockT *BB);
109 
110   /// Function.
111   const FunctionT &F;
112 
113   /// Successors for each basic block in the CFG.
114   BlockEdgeMap &Successors;
115 
116   /// Map basic blocks to their sampled weights.
117   BlockWeightMap &SampleBlockWeights;
118 };
119 
120 template <typename BT>
121 void SampleProfileInference<BT>::apply(BlockWeightMap &BlockWeights,
122                                        EdgeWeightMap &EdgeWeights) {
123   // Find all forwards reachable blocks which the inference algorithm will be
124   // applied on.
125   df_iterator_default_set<const BasicBlockT *> Reachable;
126   for (auto *BB : depth_first_ext(&F, Reachable))
127     (void)BB /* Mark all reachable blocks */;
128 
129   // Find all backwards reachable blocks which the inference algorithm will be
130   // applied on.
131   df_iterator_default_set<const BasicBlockT *> InverseReachable;
132   for (const auto &BB : F) {
133     // An exit block is a block without any successors.
134     if (isExit(&BB)) {
135       for (auto *RBB : inverse_depth_first_ext(&BB, InverseReachable))
136         (void)RBB;
137     }
138   }
139 
140   // Keep a stable order for reachable blocks
141   DenseMap<const BasicBlockT *, uint64_t> BlockIndex;
142   std::vector<const BasicBlockT *> BasicBlocks;
143   BlockIndex.reserve(Reachable.size());
144   BasicBlocks.reserve(Reachable.size());
145   for (const auto &BB : F) {
146     if (Reachable.count(&BB) && InverseReachable.count(&BB)) {
147       BlockIndex[&BB] = BasicBlocks.size();
148       BasicBlocks.push_back(&BB);
149     }
150   }
151 
152   BlockWeights.clear();
153   EdgeWeights.clear();
154   bool HasSamples = false;
155   for (const auto *BB : BasicBlocks) {
156     auto It = SampleBlockWeights.find(BB);
157     if (It != SampleBlockWeights.end() && It->second > 0) {
158       HasSamples = true;
159       BlockWeights[BB] = It->second;
160     }
161   }
162   // Quit early for functions with a single block or ones w/o samples
163   if (BasicBlocks.size() <= 1 || !HasSamples) {
164     return;
165   }
166 
167   // Create necessary objects
168   FlowFunction Func;
169   Func.Blocks.reserve(BasicBlocks.size());
170   // Create FlowBlocks
171   for (const auto *BB : BasicBlocks) {
172     FlowBlock Block;
173     if (SampleBlockWeights.find(BB) != SampleBlockWeights.end()) {
174       Block.UnknownWeight = false;
175       Block.Weight = SampleBlockWeights[BB];
176     } else {
177       Block.UnknownWeight = true;
178       Block.Weight = 0;
179     }
180     Block.Index = Func.Blocks.size();
181     Func.Blocks.push_back(Block);
182   }
183   // Create FlowEdges
184   for (const auto *BB : BasicBlocks) {
185     for (auto *Succ : Successors[BB]) {
186       if (!BlockIndex.count(Succ))
187         continue;
188       FlowJump Jump;
189       Jump.Source = BlockIndex[BB];
190       Jump.Target = BlockIndex[Succ];
191       Func.Jumps.push_back(Jump);
192       if (BB == Succ) {
193         Func.Blocks[BlockIndex[BB]].HasSelfEdge = true;
194       }
195     }
196   }
197   for (auto &Jump : Func.Jumps) {
198     Func.Blocks[Jump.Source].SuccJumps.push_back(&Jump);
199     Func.Blocks[Jump.Target].PredJumps.push_back(&Jump);
200   }
201 
202   // Try to infer probabilities of jumps based on the content of basic block
203   findUnlikelyJumps(BasicBlocks, Successors, Func);
204 
205   // Find the entry block
206   for (size_t I = 0; I < Func.Blocks.size(); I++) {
207     if (Func.Blocks[I].isEntry()) {
208       Func.Entry = I;
209       break;
210     }
211   }
212 
213   // Create and apply the inference network model.
214   applyFlowInference(Func);
215 
216   // Extract the resulting weights from the control flow
217   // All weights are increased by one to avoid propagation errors introduced by
218   // zero weights.
219   for (const auto *BB : BasicBlocks) {
220     BlockWeights[BB] = Func.Blocks[BlockIndex[BB]].Flow;
221   }
222   for (auto &Jump : Func.Jumps) {
223     Edge E = std::make_pair(BasicBlocks[Jump.Source], BasicBlocks[Jump.Target]);
224     EdgeWeights[E] = Jump.Flow;
225   }
226 
227 #ifndef NDEBUG
228   // Unreachable blocks and edges should not have a weight.
229   for (auto &I : BlockWeights) {
230     assert(Reachable.contains(I.first));
231     assert(InverseReachable.contains(I.first));
232   }
233   for (auto &I : EdgeWeights) {
234     assert(Reachable.contains(I.first.first) &&
235            Reachable.contains(I.first.second));
236     assert(InverseReachable.contains(I.first.first) &&
237            InverseReachable.contains(I.first.second));
238   }
239 #endif
240 }
241 
242 template <typename BT>
243 inline void SampleProfileInference<BT>::findUnlikelyJumps(
244     const std::vector<const BasicBlockT *> &BasicBlocks,
245     BlockEdgeMap &Successors, FlowFunction &Func) {}
246 
247 template <>
248 inline void SampleProfileInference<BasicBlock>::findUnlikelyJumps(
249     const std::vector<const BasicBlockT *> &BasicBlocks,
250     BlockEdgeMap &Successors, FlowFunction &Func) {
251   for (auto &Jump : Func.Jumps) {
252     const auto *BB = BasicBlocks[Jump.Source];
253     const auto *Succ = BasicBlocks[Jump.Target];
254     const Instruction *TI = BB->getTerminator();
255     // Check if a block ends with InvokeInst and mark non-taken branch unlikely.
256     // In that case block Succ should be a landing pad
257     if (Successors[BB].size() == 2 && Successors[BB].back() == Succ) {
258       if (isa<InvokeInst>(TI)) {
259         Jump.IsUnlikely = true;
260       }
261     }
262     const Instruction *SuccTI = Succ->getTerminator();
263     // Check if the target block contains UnreachableInst and mark it unlikely
264     if (SuccTI->getNumSuccessors() == 0) {
265       if (isa<UnreachableInst>(SuccTI)) {
266         Jump.IsUnlikely = true;
267       }
268     }
269   }
270 }
271 
272 template <typename BT>
273 inline bool SampleProfileInference<BT>::isExit(const BasicBlockT *BB) {
274   return BB->succ_empty();
275 }
276 
277 template <>
278 inline bool SampleProfileInference<BasicBlock>::isExit(const BasicBlock *BB) {
279   return succ_empty(BB);
280 }
281 
282 } // end namespace llvm
283 #endif // LLVM_TRANSFORMS_UTILS_SAMPLEPROFILEINFERENCE_H
284