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