1//===- LowerExpectIntrinsic.cpp - Lower expect intrinsic ------------------===//
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// This pass lowers the 'expect' intrinsic to LLVM metadata.
10//
11//===----------------------------------------------------------------------===//
12
13#include "llvm/Transforms/Scalar/LowerExpectIntrinsic.h"
14#include "llvm/ADT/SmallVector.h"
15#include "llvm/ADT/Statistic.h"
16#include "llvm/IR/BasicBlock.h"
17#include "llvm/IR/Constants.h"
18#include "llvm/IR/Function.h"
19#include "llvm/IR/Instructions.h"
20#include "llvm/IR/Intrinsics.h"
21#include "llvm/IR/LLVMContext.h"
22#include "llvm/IR/MDBuilder.h"
23#include "llvm/IR/ProfDataUtils.h"
24#include "llvm/Support/CommandLine.h"
25#include "llvm/Transforms/Utils/MisExpect.h"
26
27#include <cmath>
28
29using namespace llvm;
30
31#define DEBUG_TYPE "lower-expect-intrinsic"
32
33STATISTIC(ExpectIntrinsicsHandled,
34 "Number of 'expect' intrinsic instructions handled");
35
36// These default values are chosen to represent an extremely skewed outcome for
37// a condition, but they leave some room for interpretation by later passes.
38//
39// If the documentation for __builtin_expect() was made explicit that it should
40// only be used in extreme cases, we could make this ratio higher. As it stands,
41// programmers may be using __builtin_expect() / llvm.expect to annotate that a
42// branch is likely or unlikely to be taken.
43
44// WARNING: these values are internal implementation detail of the pass.
45// They should not be exposed to the outside of the pass, front-end codegen
46// should emit @llvm.expect intrinsics instead of using these weights directly.
47// Transforms should use TargetTransformInfo's getPredictableBranchThreshold().
48static cl::opt<uint32_t> LikelyBranchWeight(
49 "likely-branch-weight", cl::Hidden, cl::init(Val: 2000),
50 cl::desc("Weight of the branch likely to be taken (default = 2000)"));
51static cl::opt<uint32_t> UnlikelyBranchWeight(
52 "unlikely-branch-weight", cl::Hidden, cl::init(Val: 1),
53 cl::desc("Weight of the branch unlikely to be taken (default = 1)"));
54
55static std::tuple<uint32_t, uint32_t>
56getBranchWeight(Intrinsic::ID IntrinsicID, CallInst *CI, int BranchCount) {
57 if (IntrinsicID == Intrinsic::expect) {
58 // __builtin_expect
59 return std::make_tuple(args&: LikelyBranchWeight.getValue(),
60 args&: UnlikelyBranchWeight.getValue());
61 } else {
62 // __builtin_expect_with_probability
63 assert(CI->getNumOperands() >= 3 &&
64 "expect with probability must have 3 arguments");
65 auto *Confidence = cast<ConstantFP>(Val: CI->getArgOperand(i: 2));
66 double TrueProb = Confidence->getValueAPF().convertToDouble();
67 assert((TrueProb >= 0.0 && TrueProb <= 1.0) &&
68 "probability value must be in the range [0.0, 1.0]");
69 double FalseProb = (1.0 - TrueProb) / (BranchCount - 1);
70 uint32_t LikelyBW = ceil(x: (TrueProb * (double)(INT32_MAX - 1)) + 1.0);
71 uint32_t UnlikelyBW = ceil(x: (FalseProb * (double)(INT32_MAX - 1)) + 1.0);
72 return std::make_tuple(args&: LikelyBW, args&: UnlikelyBW);
73 }
74}
75
76static bool handleSwitchExpect(SwitchInst &SI) {
77 CallInst *CI = dyn_cast<CallInst>(Val: SI.getCondition());
78 if (!CI)
79 return false;
80
81 Function *Fn = CI->getCalledFunction();
82 if (!Fn || (Fn->getIntrinsicID() != Intrinsic::expect &&
83 Fn->getIntrinsicID() != Intrinsic::expect_with_probability))
84 return false;
85
86 Value *ArgValue = CI->getArgOperand(i: 0);
87 ConstantInt *ExpectedValue = dyn_cast<ConstantInt>(Val: CI->getArgOperand(i: 1));
88 if (!ExpectedValue)
89 return false;
90
91 SwitchInst::CaseHandle Case = *SI.findCaseValue(C: ExpectedValue);
92 unsigned n = SI.getNumCases(); // +1 for default case.
93 uint32_t LikelyBranchWeightVal, UnlikelyBranchWeightVal;
94 std::tie(args&: LikelyBranchWeightVal, args&: UnlikelyBranchWeightVal) =
95 getBranchWeight(IntrinsicID: Fn->getIntrinsicID(), CI, BranchCount: n + 1);
96
97 SmallVector<uint32_t, 16> Weights(n + 1, UnlikelyBranchWeightVal);
98
99 uint64_t Index = (Case == *SI.case_default()) ? 0 : Case.getCaseIndex() + 1;
100 Weights[Index] = LikelyBranchWeightVal;
101
102 misexpect::checkExpectAnnotations(I: SI, ExistingWeights: Weights, /*IsFrontend=*/true);
103
104 SI.setCondition(ArgValue);
105 setBranchWeights(I&: SI, Weights, /*IsExpected=*/true);
106 return true;
107}
108
109/// Handler for PHINodes that define the value argument to an
110/// @llvm.expect call.
111///
112/// If the operand of the phi has a constant value and it 'contradicts'
113/// with the expected value of phi def, then the corresponding incoming
114/// edge of the phi is unlikely to be taken. Using that information,
115/// the branch probability info for the originating branch can be inferred.
116static void handlePhiDef(CallInst *Expect) {
117 Value &Arg = *Expect->getArgOperand(i: 0);
118 ConstantInt *ExpectedValue = dyn_cast<ConstantInt>(Val: Expect->getArgOperand(i: 1));
119 if (!ExpectedValue)
120 return;
121 const APInt &ExpectedPhiValue = ExpectedValue->getValue();
122 bool ExpectedValueIsLikely = true;
123 Function *Fn = Expect->getCalledFunction();
124 // If the function is expect_with_probability, then we need to take the
125 // probability into consideration. For example, in
126 // expect.with.probability.i64(i64 %a, i64 1, double 0.0), the
127 // "ExpectedValue" 1 is unlikely. This affects probability propagation later.
128 if (Fn->getIntrinsicID() == Intrinsic::expect_with_probability) {
129 auto *Confidence = cast<ConstantFP>(Val: Expect->getArgOperand(i: 2));
130 double TrueProb = Confidence->getValueAPF().convertToDouble();
131 ExpectedValueIsLikely = (TrueProb > 0.5);
132 }
133
134 // Walk up in backward a list of instructions that
135 // have 'copy' semantics by 'stripping' the copies
136 // until a PHI node or an instruction of unknown kind
137 // is reached. Negation via xor is also handled.
138 //
139 // C = PHI(...);
140 // B = C;
141 // A = B;
142 // D = __builtin_expect(A, 0);
143 //
144 Value *V = &Arg;
145 SmallVector<Instruction *, 4> Operations;
146 while (!isa<PHINode>(Val: V)) {
147 if (ZExtInst *ZExt = dyn_cast<ZExtInst>(Val: V)) {
148 V = ZExt->getOperand(i_nocapture: 0);
149 Operations.push_back(Elt: ZExt);
150 continue;
151 }
152
153 if (SExtInst *SExt = dyn_cast<SExtInst>(Val: V)) {
154 V = SExt->getOperand(i_nocapture: 0);
155 Operations.push_back(Elt: SExt);
156 continue;
157 }
158
159 BinaryOperator *BinOp = dyn_cast<BinaryOperator>(Val: V);
160 if (!BinOp || BinOp->getOpcode() != Instruction::Xor)
161 return;
162
163 ConstantInt *CInt = dyn_cast<ConstantInt>(Val: BinOp->getOperand(i_nocapture: 1));
164 if (!CInt)
165 return;
166
167 V = BinOp->getOperand(i_nocapture: 0);
168 Operations.push_back(Elt: BinOp);
169 }
170
171 // Executes the recorded operations on input 'Value'.
172 auto ApplyOperations = [&](const APInt &Value) {
173 APInt Result = Value;
174 for (auto *Op : llvm::reverse(C&: Operations)) {
175 switch (Op->getOpcode()) {
176 case Instruction::Xor:
177 Result ^= cast<ConstantInt>(Val: Op->getOperand(i: 1))->getValue();
178 break;
179 case Instruction::ZExt:
180 Result = Result.zext(width: Op->getType()->getIntegerBitWidth());
181 break;
182 case Instruction::SExt:
183 Result = Result.sext(width: Op->getType()->getIntegerBitWidth());
184 break;
185 default:
186 llvm_unreachable("Unexpected operation");
187 }
188 }
189 return Result;
190 };
191
192 auto *PhiDef = cast<PHINode>(Val: V);
193
194 // Get the first dominating conditional branch of the operand
195 // i's incoming block.
196 auto GetDomConditional = [&](unsigned i) -> CondBrInst * {
197 BasicBlock *BB = PhiDef->getIncomingBlock(i);
198 if (CondBrInst *BI = dyn_cast<CondBrInst>(Val: BB->getTerminator()))
199 return BI;
200 BB = BB->getSinglePredecessor();
201 if (!BB)
202 return nullptr;
203 return dyn_cast<CondBrInst>(Val: BB->getTerminator());
204 };
205
206 // Now walk through all Phi operands to find phi oprerands with values
207 // conflicting with the expected phi output value. Any such operand
208 // indicates the incoming edge to that operand is unlikely.
209 for (unsigned i = 0, e = PhiDef->getNumIncomingValues(); i != e; ++i) {
210
211 Value *PhiOpnd = PhiDef->getIncomingValue(i);
212 ConstantInt *CI = dyn_cast<ConstantInt>(Val: PhiOpnd);
213 if (!CI)
214 continue;
215
216 // Not an interesting case when IsUnlikely is false -- we can not infer
217 // anything useful when:
218 // (1) We expect some phi output and the operand value matches it, or
219 // (2) We don't expect some phi output (i.e. the "ExpectedValue" has low
220 // probability) and the operand value doesn't match that.
221 const APInt &CurrentPhiValue = ApplyOperations(CI->getValue());
222 if (ExpectedValueIsLikely == (ExpectedPhiValue == CurrentPhiValue))
223 continue;
224
225 CondBrInst *BI = GetDomConditional(i);
226 if (!BI)
227 continue;
228
229 MDBuilder MDB(PhiDef->getContext());
230
231 // There are two situations in which an operand of the PhiDef comes
232 // from a given successor of a branch instruction BI.
233 // 1) When the incoming block of the operand is the successor block;
234 // 2) When the incoming block is BI's enclosing block and the
235 // successor is the PhiDef's enclosing block.
236 //
237 // Returns true if the operand which comes from OpndIncomingBB
238 // comes from outgoing edge of BI that leads to Succ block.
239 auto *OpndIncomingBB = PhiDef->getIncomingBlock(i);
240 auto IsOpndComingFromSuccessor = [&](BasicBlock *Succ) {
241 if (OpndIncomingBB == Succ)
242 // If this successor is the incoming block for this
243 // Phi operand, then this successor does lead to the Phi.
244 return true;
245 if (OpndIncomingBB == BI->getParent() && Succ == PhiDef->getParent())
246 // Otherwise, if the edge is directly from the branch
247 // to the Phi, this successor is the one feeding this
248 // Phi operand.
249 return true;
250 return false;
251 };
252 uint32_t LikelyBranchWeightVal, UnlikelyBranchWeightVal;
253 std::tie(args&: LikelyBranchWeightVal, args&: UnlikelyBranchWeightVal) = getBranchWeight(
254 IntrinsicID: Expect->getCalledFunction()->getIntrinsicID(), CI: Expect, BranchCount: 2);
255 if (!ExpectedValueIsLikely)
256 std::swap(a&: LikelyBranchWeightVal, b&: UnlikelyBranchWeightVal);
257
258 if (IsOpndComingFromSuccessor(BI->getSuccessor(i: 1)))
259 BI->setMetadata(KindID: LLVMContext::MD_prof,
260 Node: MDB.createBranchWeights(TrueWeight: LikelyBranchWeightVal,
261 FalseWeight: UnlikelyBranchWeightVal,
262 /*IsExpected=*/true));
263 else if (IsOpndComingFromSuccessor(BI->getSuccessor(i: 0)))
264 BI->setMetadata(KindID: LLVMContext::MD_prof,
265 Node: MDB.createBranchWeights(TrueWeight: UnlikelyBranchWeightVal,
266 FalseWeight: LikelyBranchWeightVal,
267 /*IsExpected=*/true));
268 }
269}
270
271// Handle both CondBrInst and SelectInst.
272template <class BrSelInst> static bool handleBrSelExpect(BrSelInst &BSI) {
273
274 // Handle non-optimized IR code like:
275 // %expval = call i64 @llvm.expect.i64(i64 %conv1, i64 1)
276 // %tobool = icmp ne i64 %expval, 0
277 // br i1 %tobool, label %if.then, label %if.end
278 //
279 // Or the following simpler case:
280 // %expval = call i1 @llvm.expect.i1(i1 %cmp, i1 1)
281 // br i1 %expval, label %if.then, label %if.end
282
283 CallInst *CI;
284
285 ICmpInst *CmpI = dyn_cast<ICmpInst>(BSI.getCondition());
286 CmpInst::Predicate Predicate;
287 ConstantInt *CmpConstOperand = nullptr;
288 if (!CmpI) {
289 CI = dyn_cast<CallInst>(BSI.getCondition());
290 Predicate = CmpInst::ICMP_NE;
291 } else {
292 Predicate = CmpI->getPredicate();
293 if (Predicate != CmpInst::ICMP_NE && Predicate != CmpInst::ICMP_EQ)
294 return false;
295
296 CmpConstOperand = dyn_cast<ConstantInt>(Val: CmpI->getOperand(i_nocapture: 1));
297 if (!CmpConstOperand)
298 return false;
299 CI = dyn_cast<CallInst>(Val: CmpI->getOperand(i_nocapture: 0));
300 }
301
302 if (!CI)
303 return false;
304
305 uint64_t ValueComparedTo = 0;
306 if (CmpConstOperand) {
307 if (CmpConstOperand->getBitWidth() > 64)
308 return false;
309 ValueComparedTo = CmpConstOperand->getZExtValue();
310 }
311
312 Function *Fn = CI->getCalledFunction();
313 if (!Fn || (Fn->getIntrinsicID() != Intrinsic::expect &&
314 Fn->getIntrinsicID() != Intrinsic::expect_with_probability))
315 return false;
316
317 Value *ArgValue = CI->getArgOperand(i: 0);
318 ConstantInt *ExpectedValue = dyn_cast<ConstantInt>(Val: CI->getArgOperand(i: 1));
319 if (!ExpectedValue)
320 return false;
321
322 MDBuilder MDB(CI->getContext());
323 MDNode *Node;
324
325 uint32_t LikelyBranchWeightVal, UnlikelyBranchWeightVal;
326 std::tie(args&: LikelyBranchWeightVal, args&: UnlikelyBranchWeightVal) =
327 getBranchWeight(IntrinsicID: Fn->getIntrinsicID(), CI, BranchCount: 2);
328
329 SmallVector<uint32_t, 4> ExpectedWeights;
330 if ((ExpectedValue->getZExtValue() == ValueComparedTo) ==
331 (Predicate == CmpInst::ICMP_EQ)) {
332 Node = MDB.createBranchWeights(
333 TrueWeight: LikelyBranchWeightVal, FalseWeight: UnlikelyBranchWeightVal, /*IsExpected=*/true);
334 ExpectedWeights = {LikelyBranchWeightVal, UnlikelyBranchWeightVal};
335 } else {
336 Node = MDB.createBranchWeights(TrueWeight: UnlikelyBranchWeightVal,
337 FalseWeight: LikelyBranchWeightVal, /*IsExpected=*/true);
338 ExpectedWeights = {UnlikelyBranchWeightVal, LikelyBranchWeightVal};
339 }
340
341 if (CmpI)
342 CmpI->setOperand(i_nocapture: 0, Val_nocapture: ArgValue);
343 else
344 BSI.setCondition(ArgValue);
345
346 misexpect::checkFrontendInstrumentation(I: BSI, ExpectedWeights);
347
348 BSI.setMetadata(LLVMContext::MD_prof, Node);
349
350 return true;
351}
352
353static bool lowerExpectIntrinsic(Function &F) {
354 bool Changed = false;
355
356 for (BasicBlock &BB : F) {
357 // Create "block_weights" metadata.
358 if (CondBrInst *BI = dyn_cast<CondBrInst>(Val: BB.getTerminator())) {
359 if (handleBrSelExpect<CondBrInst>(BSI&: *BI))
360 ExpectIntrinsicsHandled++;
361 } else if (SwitchInst *SI = dyn_cast<SwitchInst>(Val: BB.getTerminator())) {
362 if (handleSwitchExpect(SI&: *SI))
363 ExpectIntrinsicsHandled++;
364 }
365
366 // Remove llvm.expect intrinsics. Iterate backwards in order
367 // to process select instructions before the intrinsic gets
368 // removed.
369 for (Instruction &Inst : llvm::make_early_inc_range(Range: llvm::reverse(C&: BB))) {
370 CallInst *CI = dyn_cast<CallInst>(Val: &Inst);
371 if (!CI) {
372 if (SelectInst *SI = dyn_cast<SelectInst>(Val: &Inst)) {
373 if (handleBrSelExpect(BSI&: *SI))
374 ExpectIntrinsicsHandled++;
375 }
376 continue;
377 }
378
379 Function *Fn = CI->getCalledFunction();
380 if (Fn && (Fn->getIntrinsicID() == Intrinsic::expect ||
381 Fn->getIntrinsicID() == Intrinsic::expect_with_probability)) {
382 // Before erasing the llvm.expect, walk backward to find
383 // phi that define llvm.expect's first arg, and
384 // infer branch probability:
385 handlePhiDef(Expect: CI);
386 Value *Exp = CI->getArgOperand(i: 0);
387 CI->replaceAllUsesWith(V: Exp);
388 CI->eraseFromParent();
389 Changed = true;
390 }
391 }
392 }
393
394 return Changed;
395}
396
397PreservedAnalyses LowerExpectIntrinsicPass::run(Function &F,
398 FunctionAnalysisManager &) {
399 if (lowerExpectIntrinsic(F))
400 return PreservedAnalyses::none();
401
402 return PreservedAnalyses::all();
403}
404