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