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) -> BranchInst * {
197 BasicBlock *BB = PhiDef->getIncomingBlock(i);
198 BranchInst *BI = dyn_cast<BranchInst>(Val: BB->getTerminator());
199 if (BI && BI->isConditional())
200 return BI;
201 BB = BB->getSinglePredecessor();
202 if (!BB)
203 return nullptr;
204 BI = dyn_cast<BranchInst>(Val: BB->getTerminator());
205 if (!BI || BI->isUnconditional())
206 return nullptr;
207 return BI;
208 };
209
210 // Now walk through all Phi operands to find phi oprerands with values
211 // conflicting with the expected phi output value. Any such operand
212 // indicates the incoming edge to that operand is unlikely.
213 for (unsigned i = 0, e = PhiDef->getNumIncomingValues(); i != e; ++i) {
214
215 Value *PhiOpnd = PhiDef->getIncomingValue(i);
216 ConstantInt *CI = dyn_cast<ConstantInt>(Val: PhiOpnd);
217 if (!CI)
218 continue;
219
220 // Not an interesting case when IsUnlikely is false -- we can not infer
221 // anything useful when:
222 // (1) We expect some phi output and the operand value matches it, or
223 // (2) We don't expect some phi output (i.e. the "ExpectedValue" has low
224 // probability) and the operand value doesn't match that.
225 const APInt &CurrentPhiValue = ApplyOperations(CI->getValue());
226 if (ExpectedValueIsLikely == (ExpectedPhiValue == CurrentPhiValue))
227 continue;
228
229 BranchInst *BI = GetDomConditional(i);
230 if (!BI)
231 continue;
232
233 MDBuilder MDB(PhiDef->getContext());
234
235 // There are two situations in which an operand of the PhiDef comes
236 // from a given successor of a branch instruction BI.
237 // 1) When the incoming block of the operand is the successor block;
238 // 2) When the incoming block is BI's enclosing block and the
239 // successor is the PhiDef's enclosing block.
240 //
241 // Returns true if the operand which comes from OpndIncomingBB
242 // comes from outgoing edge of BI that leads to Succ block.
243 auto *OpndIncomingBB = PhiDef->getIncomingBlock(i);
244 auto IsOpndComingFromSuccessor = [&](BasicBlock *Succ) {
245 if (OpndIncomingBB == Succ)
246 // If this successor is the incoming block for this
247 // Phi operand, then this successor does lead to the Phi.
248 return true;
249 if (OpndIncomingBB == BI->getParent() && Succ == PhiDef->getParent())
250 // Otherwise, if the edge is directly from the branch
251 // to the Phi, this successor is the one feeding this
252 // Phi operand.
253 return true;
254 return false;
255 };
256 uint32_t LikelyBranchWeightVal, UnlikelyBranchWeightVal;
257 std::tie(args&: LikelyBranchWeightVal, args&: UnlikelyBranchWeightVal) = getBranchWeight(
258 IntrinsicID: Expect->getCalledFunction()->getIntrinsicID(), CI: Expect, BranchCount: 2);
259 if (!ExpectedValueIsLikely)
260 std::swap(a&: LikelyBranchWeightVal, b&: UnlikelyBranchWeightVal);
261
262 if (IsOpndComingFromSuccessor(BI->getSuccessor(i: 1)))
263 BI->setMetadata(KindID: LLVMContext::MD_prof,
264 Node: MDB.createBranchWeights(TrueWeight: LikelyBranchWeightVal,
265 FalseWeight: UnlikelyBranchWeightVal,
266 /*IsExpected=*/true));
267 else if (IsOpndComingFromSuccessor(BI->getSuccessor(i: 0)))
268 BI->setMetadata(KindID: LLVMContext::MD_prof,
269 Node: MDB.createBranchWeights(TrueWeight: UnlikelyBranchWeightVal,
270 FalseWeight: LikelyBranchWeightVal,
271 /*IsExpected=*/true));
272 }
273}
274
275// Handle both BranchInst and SelectInst.
276template <class BrSelInst> static bool handleBrSelExpect(BrSelInst &BSI) {
277
278 // Handle non-optimized IR code like:
279 // %expval = call i64 @llvm.expect.i64(i64 %conv1, i64 1)
280 // %tobool = icmp ne i64 %expval, 0
281 // br i1 %tobool, label %if.then, label %if.end
282 //
283 // Or the following simpler case:
284 // %expval = call i1 @llvm.expect.i1(i1 %cmp, i1 1)
285 // br i1 %expval, label %if.then, label %if.end
286
287 CallInst *CI;
288
289 ICmpInst *CmpI = dyn_cast<ICmpInst>(BSI.getCondition());
290 CmpInst::Predicate Predicate;
291 ConstantInt *CmpConstOperand = nullptr;
292 if (!CmpI) {
293 CI = dyn_cast<CallInst>(BSI.getCondition());
294 Predicate = CmpInst::ICMP_NE;
295 } else {
296 Predicate = CmpI->getPredicate();
297 if (Predicate != CmpInst::ICMP_NE && Predicate != CmpInst::ICMP_EQ)
298 return false;
299
300 CmpConstOperand = dyn_cast<ConstantInt>(Val: CmpI->getOperand(i_nocapture: 1));
301 if (!CmpConstOperand)
302 return false;
303 CI = dyn_cast<CallInst>(Val: CmpI->getOperand(i_nocapture: 0));
304 }
305
306 if (!CI)
307 return false;
308
309 uint64_t ValueComparedTo = 0;
310 if (CmpConstOperand) {
311 if (CmpConstOperand->getBitWidth() > 64)
312 return false;
313 ValueComparedTo = CmpConstOperand->getZExtValue();
314 }
315
316 Function *Fn = CI->getCalledFunction();
317 if (!Fn || (Fn->getIntrinsicID() != Intrinsic::expect &&
318 Fn->getIntrinsicID() != Intrinsic::expect_with_probability))
319 return false;
320
321 Value *ArgValue = CI->getArgOperand(i: 0);
322 ConstantInt *ExpectedValue = dyn_cast<ConstantInt>(Val: CI->getArgOperand(i: 1));
323 if (!ExpectedValue)
324 return false;
325
326 MDBuilder MDB(CI->getContext());
327 MDNode *Node;
328
329 uint32_t LikelyBranchWeightVal, UnlikelyBranchWeightVal;
330 std::tie(args&: LikelyBranchWeightVal, args&: UnlikelyBranchWeightVal) =
331 getBranchWeight(IntrinsicID: Fn->getIntrinsicID(), CI, BranchCount: 2);
332
333 SmallVector<uint32_t, 4> ExpectedWeights;
334 if ((ExpectedValue->getZExtValue() == ValueComparedTo) ==
335 (Predicate == CmpInst::ICMP_EQ)) {
336 Node = MDB.createBranchWeights(
337 TrueWeight: LikelyBranchWeightVal, FalseWeight: UnlikelyBranchWeightVal, /*IsExpected=*/true);
338 ExpectedWeights = {LikelyBranchWeightVal, UnlikelyBranchWeightVal};
339 } else {
340 Node = MDB.createBranchWeights(TrueWeight: UnlikelyBranchWeightVal,
341 FalseWeight: LikelyBranchWeightVal, /*IsExpected=*/true);
342 ExpectedWeights = {UnlikelyBranchWeightVal, LikelyBranchWeightVal};
343 }
344
345 if (CmpI)
346 CmpI->setOperand(i_nocapture: 0, Val_nocapture: ArgValue);
347 else
348 BSI.setCondition(ArgValue);
349
350 misexpect::checkFrontendInstrumentation(I&: BSI, ExpectedWeights);
351
352 BSI.setMetadata(LLVMContext::MD_prof, Node);
353
354 return true;
355}
356
357static bool handleBranchExpect(BranchInst &BI) {
358 if (BI.isUnconditional())
359 return false;
360
361 return handleBrSelExpect<BranchInst>(BSI&: BI);
362}
363
364static bool lowerExpectIntrinsic(Function &F) {
365 bool Changed = false;
366
367 for (BasicBlock &BB : F) {
368 // Create "block_weights" metadata.
369 if (BranchInst *BI = dyn_cast<BranchInst>(Val: BB.getTerminator())) {
370 if (handleBranchExpect(BI&: *BI))
371 ExpectIntrinsicsHandled++;
372 } else if (SwitchInst *SI = dyn_cast<SwitchInst>(Val: BB.getTerminator())) {
373 if (handleSwitchExpect(SI&: *SI))
374 ExpectIntrinsicsHandled++;
375 }
376
377 // Remove llvm.expect intrinsics. Iterate backwards in order
378 // to process select instructions before the intrinsic gets
379 // removed.
380 for (Instruction &Inst : llvm::make_early_inc_range(Range: llvm::reverse(C&: BB))) {
381 CallInst *CI = dyn_cast<CallInst>(Val: &Inst);
382 if (!CI) {
383 if (SelectInst *SI = dyn_cast<SelectInst>(Val: &Inst)) {
384 if (handleBrSelExpect(BSI&: *SI))
385 ExpectIntrinsicsHandled++;
386 }
387 continue;
388 }
389
390 Function *Fn = CI->getCalledFunction();
391 if (Fn && (Fn->getIntrinsicID() == Intrinsic::expect ||
392 Fn->getIntrinsicID() == Intrinsic::expect_with_probability)) {
393 // Before erasing the llvm.expect, walk backward to find
394 // phi that define llvm.expect's first arg, and
395 // infer branch probability:
396 handlePhiDef(Expect: CI);
397 Value *Exp = CI->getArgOperand(i: 0);
398 CI->replaceAllUsesWith(V: Exp);
399 CI->eraseFromParent();
400 Changed = true;
401 }
402 }
403 }
404
405 return Changed;
406}
407
408PreservedAnalyses LowerExpectIntrinsicPass::run(Function &F,
409 FunctionAnalysisManager &) {
410 if (lowerExpectIntrinsic(F))
411 return PreservedAnalyses::none();
412
413 return PreservedAnalyses::all();
414}
415