llvm/lib/Transforms/Scalar/LowerExpectIntrinsic.cpp

   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
  29 using namespace llvm;
  30
  31 #define DEBUG_TYPE "lower-expect-intrinsic"
  32
  33 STATISTIC(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().
  48 static cl::opt<uint32_t> LikelyBranchWeight(
  49     "likely-branch-weight", cl::Hidden, cl::init(2000),
  50     cl::desc("Weight of the branch likely to be taken (default = 2000)"));
  51 static cl::opt<uint32_t> UnlikelyBranchWeight(
  52     "unlikely-branch-weight", cl::Hidden, cl::init(1),
  53     cl::desc("Weight of the branch unlikely to be taken (default = 1)"));
  54
  55 static std::tuple<uint32_t, uint32_t>
  56 getBranchWeight(Intrinsic::ID IntrinsicID, CallInst *CI, int BranchCount) {
  57   if (IntrinsicID == Intrinsic::expect) {
  58     // __builtin_expect
  59     return std::make_tuple(LikelyBranchWeight.getValue(),
  60                            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>(CI->getArgOperand(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((TrueProb * (double)(INT32_MAX - 1)) + 1.0);
  71     uint32_t UnlikelyBW = ceil((FalseProb * (double)(INT32_MAX - 1)) + 1.0);
  72     return std::make_tuple(LikelyBW, UnlikelyBW);
  73   }
  74 }
  75
  76 static bool handleSwitchExpect(SwitchInst &SI) {
  77   CallInst *CI = dyn_cast<CallInst>(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(0);
  87   ConstantInt *ExpectedValue = dyn_cast<ConstantInt>(CI->getArgOperand(1));
  88   if (!ExpectedValue)
  89     return false;
  90
  91   SwitchInst::CaseHandle Case = *SI.findCaseValue(ExpectedValue);
  92   unsigned n = SI.getNumCases(); // +1 for default case.
  93   uint32_t LikelyBranchWeightVal, UnlikelyBranchWeightVal;
  94   std::tie(LikelyBranchWeightVal, UnlikelyBranchWeightVal) =
  95       getBranchWeight(Fn->getIntrinsicID(), CI, 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(SI, Weights, /*IsFrontend=*/true);
 103
 104   SI.setCondition(ArgValue);
 105   setBranchWeights(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.
 116 static void handlePhiDef(CallInst *Expect) {
 117   Value &Arg = *Expect->getArgOperand(0);
 118   ConstantInt *ExpectedValue = dyn_cast<ConstantInt>(Expect->getArgOperand(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>(Expect->getArgOperand(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>(V)) {
 147     if (ZExtInst *ZExt = dyn_cast<ZExtInst>(V)) {
 148       V = ZExt->getOperand(0);
 149       Operations.push_back(ZExt);
 150       continue;
 151     }
 152
 153     if (SExtInst *SExt = dyn_cast<SExtInst>(V)) {
 154       V = SExt->getOperand(0);
 155       Operations.push_back(SExt);
 156       continue;
 157     }
 158
 159     BinaryOperator *BinOp = dyn_cast<BinaryOperator>(V);
 160     if (!BinOp || BinOp->getOpcode() != Instruction::Xor)
 161       return;
 162
 163     ConstantInt *CInt = dyn_cast<ConstantInt>(BinOp->getOperand(1));
 164     if (!CInt)
 165       return;
 166
 167     V = BinOp->getOperand(0);
 168     Operations.push_back(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(Operations)) {
 175       switch (Op->getOpcode()) {
 176       case Instruction::Xor:
 177         Result ^= cast<ConstantInt>(Op->getOperand(1))->getValue();
 178         break;
 179       case Instruction::ZExt:
 180         Result = Result.zext(Op->getType()->getIntegerBitWidth());
 181         break;
 182       case Instruction::SExt:
 183         Result = Result.sext(Op->getType()->getIntegerBitWidth());
 184         break;
 185       default:
 186         llvm_unreachable("Unexpected operation");
 187       }
 188     }
 189     return Result;
 190   };
 191
 192   auto *PhiDef = cast<PHINode>(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>(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>(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>(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(LikelyBranchWeightVal, UnlikelyBranchWeightVal) = getBranchWeight(
 258         Expect->getCalledFunction()->getIntrinsicID(), Expect, 2);
 259     if (!ExpectedValueIsLikely)
 260       std::swap(LikelyBranchWeightVal, UnlikelyBranchWeightVal);
 261
 262     if (IsOpndComingFromSuccessor(BI->getSuccessor(1)))
 263       BI->setMetadata(LLVMContext::MD_prof,
 264                       MDB.createBranchWeights(LikelyBranchWeightVal,
 265                                               UnlikelyBranchWeightVal,
 266                                               /*IsExpected=*/true));
 267     else if (IsOpndComingFromSuccessor(BI->getSuccessor(0)))
 268       BI->setMetadata(LLVMContext::MD_prof,
 269                       MDB.createBranchWeights(UnlikelyBranchWeightVal,
 270                                               LikelyBranchWeightVal,
 271                                               /*IsExpected=*/true));
 272   }
 273 }
 274
 275 // Handle both BranchInst and SelectInst.
 276 template <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>(CmpI->getOperand(1));
 301     if (!CmpConstOperand)
 302       return false;
 303     CI = dyn_cast<CallInst>(CmpI->getOperand(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(0);
 322   ConstantInt *ExpectedValue = dyn_cast<ConstantInt>(CI->getArgOperand(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(LikelyBranchWeightVal, UnlikelyBranchWeightVal) =
 331       getBranchWeight(Fn->getIntrinsicID(), CI, 2);
 332
 333   SmallVector<uint32_t, 4> ExpectedWeights;
 334   if ((ExpectedValue->getZExtValue() == ValueComparedTo) ==
 335       (Predicate == CmpInst::ICMP_EQ)) {
 336     Node = MDB.createBranchWeights(
 337         LikelyBranchWeightVal, UnlikelyBranchWeightVal, /*IsExpected=*/true);
 338     ExpectedWeights = {LikelyBranchWeightVal, UnlikelyBranchWeightVal};
 339   } else {
 340     Node = MDB.createBranchWeights(UnlikelyBranchWeightVal,
 341                                    LikelyBranchWeightVal, /*IsExpected=*/true);
 342     ExpectedWeights = {UnlikelyBranchWeightVal, LikelyBranchWeightVal};
 343   }
 344
 345   if (CmpI)
 346     CmpI->setOperand(0, ArgValue);
 347   else
 348     BSI.setCondition(ArgValue);
 349
 350   misexpect::checkFrontendInstrumentation(BSI, ExpectedWeights);
 351
 352   BSI.setMetadata(LLVMContext::MD_prof, Node);
 353
 354   return true;
 355 }
 356
 357 static bool handleBranchExpect(BranchInst &BI) {
 358   if (BI.isUnconditional())
 359     return false;
 360
 361   return handleBrSelExpect<BranchInst>(BI);
 362 }
 363
 364 static 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>(BB.getTerminator())) {
 370       if (handleBranchExpect(*BI))
 371         ExpectIntrinsicsHandled++;
 372     } else if (SwitchInst *SI = dyn_cast<SwitchInst>(BB.getTerminator())) {
 373       if (handleSwitchExpect(*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(llvm::reverse(BB))) {
 381       CallInst *CI = dyn_cast<CallInst>(&Inst);
 382       if (!CI) {
 383         if (SelectInst *SI = dyn_cast<SelectInst>(&Inst)) {
 384           if (handleBrSelExpect(*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(CI);
 397         Value *Exp = CI->getArgOperand(0);
 398         CI->replaceAllUsesWith(Exp);
 399         CI->eraseFromParent();
 400         Changed = true;
 401       }
 402     }
 403   }
 404
 405   return Changed;
 406 }
 407
 408 PreservedAnalyses LowerExpectIntrinsicPass::run(Function &F,
 409                                                 FunctionAnalysisManager &) {
 410   if (lowerExpectIntrinsic(F))
 411     return PreservedAnalyses::none();
 412
 413   return PreservedAnalyses::all();
 414 }