[Profile[ Enhance expect lowering to handle correlated branches

[davidxl]

Summary

There are situations buitlin_expect is used on a set of conditions (&&, ||) where there is only one expect builtin call inserted. There is also situations where the result of builtin_expect is not fed into any branches (e.g, a ternary expression). In the above cases, the lowering phase fails to annotate the all related branches.

This patch enhances the lowering phase by doing 'backward' analysis with Phidefs. For examples, see the test cases.

[wmi]

Overall looks good. Some minor comments inlined.

lib/Transforms/Scalar/LowerExpectIntrinsic.cpp

112–119 ↗

(On Diff #98899)

BinaryOp should already be canonicalized here, so ConstantInt will be in Opnd1.

142 ↗

(On Diff #98899)

docminating --> dominating

159 ↗

(On Diff #98899)

wolk -> walk

162–164 ↗

(On Diff #98899)

isa<ConstantInt> (PhiOpnd) can be replaced by
ConstantInt *CI;
if (CI = dyn_cast<ConstantInt>(PhiOpnd))

test/Transforms/LowerExpectIntrinsic/phi_merge.ll

25 ↗

(On Diff #98899)

Maybe add the whole br statement? Otherwise we may check a !prof annotated to a different branch other than we expect.

[davidxl]

lib/Transforms/Scalar/LowerExpectIntrinsic.cpp

112–119 ↗

(On Diff #98899)

Done. added an assert.

[davidxl]

Address Wei's review feebback.

[chandlerc]

lib/Transforms/Scalar/LowerExpectIntrinsic.cpp

86–93 ↗

(On Diff #99208)

Please use a doxygen comment, and format it consistently with modern documentation comments in LLVM.

94 ↗

(On Diff #99208)

LLVM always uses static void not void static.

94 ↗

(On Diff #99208)

Use a reference since this cannot be null.

95 ↗

(On Diff #99208)

ArgValue seems redundant, just Arg?

99–100 ↗

(On Diff #99208)

This comment isn't really verbose enough to help th ereader out. Can you use proper prose and expand on it so that folks not familiar with the specific pattern you're thinking of ('copy chain'?) understand what is going on?

102 ↗

(On Diff #99208)

We write infinite loops more often as 'for (;;)', but I'm not sure this really needs to use that pattern...

113 ↗

(On Diff #99208)

You can't assert that IR is in canonical form. You can only assert that it would pass the verifier.

But you don't need to *handle* non-canonical form, you can just fail to add the metadata.

115–117 ↗

(On Diff #99208)

If you just want to test for zero, you can do that directly without zext-ing the constant...

134–135 ↗

(On Diff #99208)

You could just use one variable above (V or Phi) and write the loop to essentially be:

// Walk up through the operands until we find a PHI node.
while (!isa<PHINode>(V)) {
  ...
}

auto *PN = dyn_cast<PHINode>(V);
if (!PN)
  // Nothing to do if we cannot find a PHI node.
  return;

147–150 ↗

(On Diff #99208)

If the CFG isn't simplified, we need to walk more than one. If it is, we don't need to walk any. Do you have cases where this was necessary?

157 ↗

(On Diff #99208)

This comment isn't really helping me. From the code, I know you're walking the PHI operands, but I don't know why or what you're planning to do with them.

176 ↗

(On Diff #99208)

I think this is definitely a place where explaining more would help. The comment as is doesn't really help me read the code.

184–185 ↗

(On Diff #99208)

Why not just call BR->setMetadata twice above? The variable doesn't seem to save a lot here...

279–282 ↗

(On Diff #99208)

This doesn't erase a PHI node though... Maybe this comment just needs updating?

[chandlerc]

test/Transforms/LowerExpectIntrinsic/phi_merge.ll

14–24 ↗

(On Diff #99208)

Can you minimize the test cases rather than including all of the boilerplate introduced by Clang?

Similarly, no use variadic function calls here now that it is an LLVM IR test case.

[davidxl]

lib/Transforms/Scalar/LowerExpectIntrinsic.cpp

113 ↗

(On Diff #99208)

Removed.

115–117 ↗

(On Diff #99208)

This is not for constant but the value flow.

134–135 ↗

(On Diff #99208)

The loop is not greatly simplified and more readable.

147–150 ↗

(On Diff #99208)

The comment is stale. We need to handle diamond shaped control. See also phi_tern.ll for an example (where cfg is not simplified).

279–282 ↗

(On Diff #99208)

It is referring to the removal after this call.

[davidxl]

Address Chandler's feedbacks.

[chandlerc]

lib/Transforms/Scalar/LowerExpectIntrinsic.cpp

97 ↗

(On Diff #99222)

You don't check whether this is null... Either add a check or use cast?

100 ↗

(On Diff #99222)

nit: Some vertical whitespace before this comment block would make it easier to read IMO.

121–136 ↗

(On Diff #99222)

I think all of this can now be re-written using early return though, and I think that plus removing some unnecessary variables helps make the rest of the loop more readable:

BinaryOperator *BinOp = dyn_cast<BinaryOperator>(V)
if (!BinOp || BinOp->getOpcode() != Instruction::Xor)
  return;
  
ConstantInt *COperand = dyn_cast<ConstantInt>(BinOp->getOperand(1));
if (!COperand)
  return;

if (!COperand->isZero())
  ValueInverted = !ValueInverted;

V = BinOp->getOperand(0);

This also makes it more clear (at least to me) that there is a bug here. When the value isn't an i1 we can't assume that all non-zero xors negate. I think this will need to check that the constant operand is specifically either zero or all ones. Which ends up looking like:

ConstantInt *COperand = dyn_cast<ConstantInt>(BinOp->getOperand(1));
if (!COperand)
  return;

if (!COperand->isZero()) {
  // If this isn't a no-op copy, check that it is exactly a negation.
  if (!COperand->isMinusOne())
    return;

  ValueInverted = !ValueInverted;
}

V = BinOp->getOperand(0);

140 ↗

(On Diff #99222)

nit: I actually prefer the dyn_cast line right before the condition. Also, declare this here with auto *?

144–145 ↗

(On Diff #99222)

How do you want this to handle non-i1 expect calls?

149 ↗

(On Diff #99222)

Capturing by copy is pretty expensive. Given that this is called only from within this function, capturing by reference seems a bit more appropriate?

181–186 ↗

(On Diff #99222)

Reading this loop, I felt like this could end up setting the same metadata on a branch multiple times...

But reading this I feel like this isn't actually the case. But I'm also not exactly sure how this is working. I mean, I understand at a high level the cases you're trying to handle (I've read the test cases), but it really isn't clear how *this* code is effectively doing that.

Maybe an ascii-art diagram of the two cases you're trying to handle here would help? Maybe specifically talking about why we're looping over all of the phi's operands? (I had asked for the latter already...)

279–282 ↗

(On Diff #99208)

Ok, but you've written this comment right above code that erases the *call* and not any phi node, so it seems a bit confusing to write the comment this way. And erasing the phi node doesn't seem as relevant as erasing the call...

[davidxl]

lib/Transforms/Scalar/LowerExpectIntrinsic.cpp

144–145 ↗

(On Diff #99222)

As other cases, arbitrary value builtin_expect is not yet properly handled. One possible use of it is to annotate switch statement, but it is pretty uncommon.

I added a TODO there

181–186 ↗

(On Diff #99222)

I used a lambda function and I think it is clearer now.

[davidxl]

Address review feedback.

[davidxl]

Any more comments?

[chandlerc]

lib/Transforms/Scalar/LowerExpectIntrinsic.cpp

125 ↗

(On Diff #99234)

My suggestion was somewhat specifically to use early return here rather than break?

127 ↗

(On Diff #99234)

This comment isn't really helping. The code already is clearly handling XOR.

I would suggest something more like before you get the value as a binary operator saying:

// Lastly, we try to handle binary operators that either copy
// or invert the boolean value. These are canonicalized in LLVM
// to an `xor` instruction.

Once you have that, not sure you even need a comment here at all.

128–129 ↗

(On Diff #99234)

These variables don't seem to help anything. They're each only used once. Sink the initializer to the one usage below?

136–139 ↗

(On Diff #99234)

This looks like it is wrong when the xor operand is neither zero nor -1 (possibly due to being on a zero-extended value). We move V to the operand which might be the phi, and then break out.

As I had suggested, I find this easier to read when structured using early return and it also avoids the bug:

if (!CInt->isZero()) {
  if (!CInt->isMinusOne())
    return;

  ValueInverted = !ValueInverted;
}

That said, I would also personally only assign V to the operand after this block just to avoid confusion.

154 ↗

(On Diff #99234)

nit: BI is more commonly used for a branch instruction than BR in LLVM.

166 ↗

(On Diff #99234)

An old comment that didn't get addressed but is on a different line now:
"""
This comment isn't really helping me. From the code, I know you're walking the PHI operands, but I don't know why or what you're planning to do with them.
"""

175–177 ↗

(On Diff #99234)

We typically use 'FIXME'. I'd put this comment up on the computing of IsPhiValueNonZero as that's when this became confusing for me.

I also don't think the IsUnlikely variable is useful here. I would just sink it into the if.

The comment Not an interesting case: doesn't really tell the reader much. I think it would be more helpful to explain why only matched operands and values are meaningful to turn into metadata.

187–192 ↗

(On Diff #99234)

Minor prose suggestions:

"Figuring out which" -> "Test whether a particular"
"that leads" -> "leads"

"triangular shaped" should either be just "triangular" or "triangle shaped". I mildly prefer the former.

Ending this with a ':' doesn't make much sense to me.

Overall, the second sentence here doesn't help me as much. I would maybe suggest putting the explanation of the "how does this work" attached to the code inside the lambda. I'll suggest things there.

199 ↗

(On Diff #99234)

Maybe comment:

// If this successor is the incoming block for this Phi operand, then
// this successor does lead to the Phi.

However, I feel like this needs to check that OpndIncomingBB->getUniquePredecessor() returns non-null. Otherwise this might not be an actual diamond.

201–202 ↗

(On Diff #99234)

Maybe comment:

// Otherwise, if the edge is directly from the branch to the Phi,
// this successor is the one feeding this Phi operand.

97 ↗

(On Diff #99222)

This doesn't appear to be done.

100 ↗

(On Diff #99222)

nit: This is a blonk comment line rather than a blank line?

144–145 ↗

(On Diff #99222)

I don't think a TODO is really enough. We shouldn't put completely wrong probabilities on something just because we don't handle non-i1 expects. I'm fine if you just want to check at the start of this that the type is i1 and return otherwise.

279–282 ↗

(On Diff #99208)

Outstanding comment.

test/Transforms/LowerExpectIntrinsic/phi_merge.ll

12 ↗

(On Diff #99234)

CHECK-LABEL?

38 ↗

(On Diff #99234)

I typically put checks after the IR they are matching against.

53 ↗

(On Diff #99234)

CHECK-LABEL?

63 ↗

(On Diff #99234)

It would be good to make this a more precise check by first checking for the branch, and then checking that the profile data isn't there:

; CHECK: br i1 %tmp13,
; CHECK-NOT: !prof

14–24 ↗

(On Diff #99208)

Can still remove the '(...)' from the call instructions...

[davidxl]

lib/Transforms/Scalar/LowerExpectIntrinsic.cpp

136–139 ↗

(On Diff #99234)

I made some flow changes to match previous patterns. It is more readable now.

166 ↗

(On Diff #99234)

Added more comments.

175–177 ↗

(On Diff #99234)

Added more explanation.

199 ↗

(On Diff #99234)

I added the comment. There is no need to check unique predecessor here though. The existance of other predecessors does not invalidate the fact that any incoming edge to OpndIncomingBB is unlikely.

[davidxl]

Addressed review comments.

[davidxl]

Any more comments?

[chandlerc]

lib/Transforms/Scalar/LowerExpectIntrinsic.cpp

86–87 ↗

(On Diff #99663)

nit: wording might be clearer as:

Handler for PHINodes that define the value argument to an @llvm.expect call.

89–90 ↗

(On Diff #99663)

nit: no need to repeate the summary here. I think you can jump to If the operand ....

98 ↗

(On Diff #99663)

If you're just going to assert, you should use cast rather than dyn_cast -- it does exactly this, it asserts, but it does so more cheaply in non-asserts builds.

198 ↗

(On Diff #99663)

nit: ':' -> '.'.

200–203 ↗

(On Diff #99663)

I think this API is overall a bit confusing because we call it with integer indices, but it captures a specific branch and phi instruction...

Maybe, rather than getting the incoming block inside the lambda, why not pull the computation of OpndIncomingBB out of the lambda entirely? It doesn't change between the calls below.

Similarly, I feel like it would be a bit more readable to pass in the successor basic block and call BI->getSuccessor() in the call as that avoids documenting the magical 1 and 0 indices.

202 ↗

(On Diff #99663)

OperandIdx and SuccIdx would be better parameter names here I think.

318 ↗

(On Diff #99663)

Again, I think 'Before erasing the phi' is confusing here. This code doesn't erase the phi node. I think this should say 'Before erasing the llvm.expect call, ...' which can then make it easier to refer to the first argument in the rest of the sentence.

144–145 ↗

(On Diff #99222)

I'm still not seeing any response to my follow-up.

test/Transforms/LowerExpectIntrinsic/phi_merge.ll

14–24 ↗

(On Diff #99208)

Outstanding comment.

[davidxl]

lib/Transforms/Scalar/LowerExpectIntrinsic.cpp

200–203 ↗

(On Diff #99663)

I modified the function.

318 ↗

(On Diff #99663)

Ouch, I meant llvm.expect.

100 ↗

(On Diff #99222)

Removed the comment marker

144–145 ↗

(On Diff #99222)

I am not sure what your followup question is, but non-i1 case in the patch is handled correctly (but conservatively) -- basically no wrong probability will be emitted -- and it will early return as you expected (isUnlikely will be false). I enhanced the comment.

[davidxl]

Address Chandler's feedbacks.

[davidxl]

Added support to handle integer with any bit-width. New test cases are added to test int type other than i1.

[chandlerc]

some of minor nits a a small thing left... thanks for the handling and the nice test cases around larger integers!

lib/Transforms/Scalar/LowerExpectIntrinsic.cpp

96 ↗

(On Diff #101089)

You'll need to use the APInt so this can handle larger than 64-bit integers (or you'll need to bail out when the integer requires more than 64-bits). The APInt code should be pretty straight forward here though.

(also, add a test case (regardless of which direction you go) with 73 bit integer or some such?)

116–117 ↗

(On Diff #101089)

MathExtras.h provides maskTrailingOnes which implements this. However, if you end up moving to APInt it has a pretty rich interface for doing these kinds of operations. You may be able to just call .zext on it to model the zext.

122–124 ↗

(On Diff #101089)

It's a bit surprising that these two statements are in the opposite order as above... But anyways, doesn't this need to be a sign extend of the expected value?

181–188 ↗

(On Diff #101089)

FWIW, this comment now really helps me with the whole thing. Thanks!

203 ↗

(On Diff #101089)

Hoist this above the lambda, and thin skip the lambda argument? It doesn't vary between calls.

204–205 ↗

(On Diff #101089)

Skip the variables and just inline the successor access?

test/Transforms/LowerExpectIntrinsic/phi_merge.ll

119 ↗

(On Diff #101089)

Add a test (somewhere) for sext when the incoming value from the Phi node will get is -1? I'd probably use an i8 so its easy to write 255 as the value and sign extend it but expect that it is still 255. I think that test case would do the wrong thing currently? (see comment above, but maybe I've misread the code above.)

270–275 ↗

(On Diff #101089)

Note that in this test case and all the others you can strip off the attributes and the metadata here.

14–24 ↗

(On Diff #99208)

Still outstanding...

[davidxl]

lib/Transforms/Scalar/LowerExpectIntrinsic.cpp

96 ↗

(On Diff #101089)

Rewrote the code to use APInt. Also do a forward computation from the operand value using the recorded operations and compare with the expected value. This is more flexible which allows us to handle other operations in the future.

test/Transforms/LowerExpectIntrinsic/phi_merge.ll

119 ↗

(On Diff #101089)

Added two sext test case.

From i8 255 --> sext to i64 and compare with 255 -- does not match -- unlikely annotation emitted.
From i8 255 --> sext to i64 and compare with -1 --> match -- no annotation.

[davidxl]

Address review feebacks.

[chandlerc]

Cool! I like the replay logic. It's a little alarming that we will replay arbitrary number of operations on an arbitrary number of phi inputs, but I'm *really* not worried here and the simplicity and clarity of the logic resulting from simple replay is great.

The remaining comments are just superficial stuff. If they make sense to you, feel free to go ahead and commit with those changes. If something isn't clear or needs more discussion, just lemme know.

lib/Transforms/Scalar/LowerExpectIntrinsic.cpp

98–101 ↗

(On Diff #101150)

Probably want to update this comment to explain that we're now recording the operations so we can replay them.

135 ↗

(On Diff #101150)

Heh, for once, you can omit the return type! ;] Maybe add a brief comment explaining this lambda replays the operations on an APInt?

Or could name it something more explanatory like ApplyOperations

137 ↗

(On Diff #101150)

You can use a range based for loop:

for (Instruction *Op : llvm::reverse(Operations)) {

141–142 ↗

(On Diff #101150)

Just use cast here, and likely skip the variable.

186–191 ↗

(On Diff #101150)

I would sink the comparison into the condition. The variable 'IsUnlikely' isn't as helpful as the comment here, and its only used for the if.

[davidxl]

thanks for the review! I am pretty happy with the final version .

llvm/trunk/lib/Transforms/Scalar/LowerExpectIntrinsic.cpp

 //===- LowerExpectIntrinsic.cpp - Lower expect intrinsic ------------------===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 //
 // This pass lowers the 'expect' intrinsic to LLVM metadata.
 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/Transforms/Scalar/LowerExpectIntrinsic.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/Statistic.h"
+#include "llvm/ADT/iterator_range.h"
 #include "llvm/IR/BasicBlock.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/Intrinsics.h"
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/MDBuilder.h"
 #include "llvm/IR/Metadata.h"
static bool handleSwitchExpect(SwitchInst &SI) {
 
   SI.setMetadata(LLVMContext::MD_prof,
                  MDBuilder(CI->getContext()).createBranchWeights(Weights));
 
   SI.setCondition(ArgValue);
   return true;
 }
 
+/// Handler for PHINodes that define the value argument to an
+/// @llvm.expect call.
+///
+/// If the operand of the phi has a constant value and it 'contradicts'
+/// with the expected value of phi def, then the corresponding incoming
+/// edge of the phi is unlikely to be taken. Using that information,
+/// the branch probability info for the originating branch can be inferred.
+static void handlePhiDef(CallInst *Expect) {
+  Value &Arg = *Expect->getArgOperand(0);
+  ConstantInt *ExpectedValue = cast<ConstantInt>(Expect->getArgOperand(1));
+  const APInt &ExpectedPhiValue = ExpectedValue->getValue();
+
+  // Walk up in backward a list of instructions that
+  // have 'copy' semantics by 'stripping' the copies
+  // until a PHI node or an instruction of unknown kind
+  // is reached. Negation via xor is also handled.
+  //
+  //       C = PHI(...);
+  //       B = C;
+  //       A = B;
+  //       D = __builtin_expect(A, 0);
+  //
+  Value *V = &Arg;
+  SmallVector<Instruction *, 4> Operations;
+  while (!isa<PHINode>(V)) {
+    if (ZExtInst *ZExt = dyn_cast<ZExtInst>(V)) {
+      V = ZExt->getOperand(0);
+      Operations.push_back(ZExt);
+      continue;
+    }
+
+    if (SExtInst *SExt = dyn_cast<SExtInst>(V)) {
+      V = SExt->getOperand(0);
+      Operations.push_back(SExt);
+      continue;
+    }
+
+    BinaryOperator *BinOp = dyn_cast<BinaryOperator>(V);
+    if (!BinOp || BinOp->getOpcode() != Instruction::Xor)
+      return;
+
+    ConstantInt *CInt = dyn_cast<ConstantInt>(BinOp->getOperand(1));
+    if (!CInt)
+      return;
+
+    V = BinOp->getOperand(0);
+    Operations.push_back(BinOp);
+  }
+
+  // Executes the recorded operations on input 'Value'.
+  auto ApplyOperations = [&](const APInt &Value) {
+    APInt Result = Value;
+    for (auto Op : llvm::reverse(Operations)) {
+      switch (Op->getOpcode()) {
+      case Instruction::Xor:
+        Result ^= cast<ConstantInt>(Op->getOperand(1))->getValue();
+        break;
+      case Instruction::ZExt:
+        Result = Result.zext(Op->getType()->getIntegerBitWidth());
+        break;
+      case Instruction::SExt:
+        Result = Result.sext(Op->getType()->getIntegerBitWidth());
+        break;
+      default:
+        llvm_unreachable("Unexpected operation");
+      }
+    }
+    return Result;
+  };
+
+  auto *PhiDef = dyn_cast<PHINode>(V);
+
+  // Get the first dominating conditional branch of the operand
+  // i's incoming block.
+  auto GetDomConditional = [&](unsigned i) -> BranchInst * {
+    BasicBlock *BB = PhiDef->getIncomingBlock(i);
+    BranchInst *BI = dyn_cast<BranchInst>(BB->getTerminator());
+    if (BI && BI->isConditional())
+      return BI;
+    BB = BB->getSinglePredecessor();
+    if (!BB)
+      return nullptr;
+    BI = dyn_cast<BranchInst>(BB->getTerminator());
+    if (!BI || BI->isUnconditional())
+      return nullptr;
+    return BI;
+  };
+
+  // Now walk through all Phi operands to find phi oprerands with values
+  // conflicting with the expected phi output value. Any such operand
+  // indicates the incoming edge to that operand is unlikely.
+  for (unsigned i = 0, e = PhiDef->getNumIncomingValues(); i != e; ++i) {
+
+    Value *PhiOpnd = PhiDef->getIncomingValue(i);
+    ConstantInt *CI = dyn_cast<ConstantInt>(PhiOpnd);
+    if (!CI)
+      continue;
+
+    // Not an interesting case when IsUnlikely is false -- we can not infer
+    // anything useful when the operand value matches the expected phi
+    // output.
+    if (ExpectedPhiValue == ApplyOperations(CI->getValue()))
+      continue;
+
+    BranchInst *BI = GetDomConditional(i);
+    if (!BI)
+      continue;
+
+    MDBuilder MDB(PhiDef->getContext());
+
+    // There are two situations in which an operand of the PhiDef comes
+    // from a given successor of a branch instruction BI.
+    // 1) When the incoming block of the operand is the successor block;
+    // 2) When the incoming block is BI's enclosing block and the
+    // successor is the PhiDef's enclosing block.
+    //
+    // Returns true if the operand which comes from OpndIncomingBB
+    // comes from outgoing edge of BI that leads to Succ block.
+    auto *OpndIncomingBB = PhiDef->getIncomingBlock(i);
+    auto IsOpndComingFromSuccessor = [&](BasicBlock *Succ) {
+      if (OpndIncomingBB == Succ)
+        // If this successor is the incoming block for this
+        // Phi operand, then this successor does lead to the Phi.
+        return true;
+      if (OpndIncomingBB == BI->getParent() && Succ == PhiDef->getParent())
+        // Otherwise, if the edge is directly from the branch
+        // to the Phi, this successor is the one feeding this
+        // Phi operand.
+        return true;
+      return false;
+    };
+
+    if (IsOpndComingFromSuccessor(BI->getSuccessor(1)))
+      BI->setMetadata(
+          LLVMContext::MD_prof,
+          MDB.createBranchWeights(LikelyBranchWeight, UnlikelyBranchWeight));
+    else if (IsOpndComingFromSuccessor(BI->getSuccessor(0)))
+      BI->setMetadata(
+          LLVMContext::MD_prof,
+          MDB.createBranchWeights(UnlikelyBranchWeight, LikelyBranchWeight));
+  }
+}
+
 // Handle both BranchInst and SelectInst.
 template <class BrSelInst> static bool handleBrSelExpect(BrSelInst &BSI) {
 
   // Handle non-optimized IR code like:
   //   %expval = call i64 @llvm.expect.i64(i64 %conv1, i64 1)
   //   %tobool = icmp ne i64 %expval, 0
   //   br i1 %tobool, label %if.then, label %if.end
   //
for (auto BI = BB.rbegin(), BE = BB.rend(); BI != BE;) {
           if (handleBrSelExpect(*SI))
             ExpectIntrinsicsHandled++;
         }
         continue;
       }
 
       Function *Fn = CI->getCalledFunction();
       if (Fn && Fn->getIntrinsicID() == Intrinsic::expect) {
+        // Before erasing the llvm.expect, walk backward to find
+        // phi that define llvm.expect's first arg, and
+        // infer branch probability:
+        handlePhiDef(CI);
         Value *Exp = CI->getArgOperand(0);
         CI->replaceAllUsesWith(Exp);
         CI->eraseFromParent();
         Changed = true;
       }
     }
   }
 

llvm/trunk/test/Transforms/LowerExpectIntrinsic/phi_merge.ll

+; RUN: opt -lower-expect  -S -o - < %s | FileCheck %s
+; RUN: opt -S -passes='function(lower-expect)' < %s | FileCheck %s
+
+; The C case
+; if (__builtin_expect((x > goo() && y > hoo() && z > too()), 1)) 
+; For the above case, all 3 branches should be annotated.
+;
+; if (__builtin_expect((x > goo() && y > hoo() && z > too()), 0)) 
+; For the above case, we don't have enough information, so
+; only the last branch is annotated.
+
+define void @foo(i32 %arg, i32 %arg1, i32 %arg2, i32 %arg3) {
+; CHECK-LABEL: void @foo
+bb:
+  %tmp8 = call i32  @goo() 
+  %tmp9 = icmp sgt i32 %tmp8, %arg
+  br i1 %tmp9, label %bb10, label %bb18
+; CHECK: !prof [[WEIGHT:![0-9]+]]
+
+bb10:                                             ; preds = %bb
+  %tmp12 = call i32  @hoo()
+  %tmp13 = icmp sgt i32 %arg1, %tmp12
+  br i1 %tmp13, label %bb14, label %bb18
+; CHECK: br i1 %tmp13, {{.*}}!prof [[WEIGHT]]
+
+bb14:                                             ; preds = %bb10
+  %tmp16 = call i32  @too()
+  %tmp17 = icmp sgt i32 %arg2, %tmp16
+  br label %bb18
+
+bb18:                                             ; preds = %bb14, %bb10, %bb
+  %tmp19 = phi i1 [ false, %bb10 ], [ false, %bb ], [ %tmp17, %bb14 ]
+  %tmp20 = xor i1 %tmp19, true
+  %tmp21 = xor i1 %tmp20, true
+  %tmp22 = zext i1 %tmp21 to i32
+  %tmp23 = sext i32 %tmp22 to i64
+  %tmp24 = call i64 @llvm.expect.i64(i64 %tmp23, i64 1)
+  %tmp25 = icmp ne i64 %tmp24, 0
+  br i1 %tmp25, label %bb26, label %bb28
+; CHECK: br i1 %tmp25,{{.*}}!prof [[WEIGHT]]
+
+bb26:                                             ; preds = %bb18
+  %tmp27 = call i32  @goo()
+  br label %bb30
+
+bb28:                                             ; preds = %bb18
+  %tmp29 = call i32  @hoo()
+  br label %bb30
+
+bb30:                                             ; preds = %bb28, %bb26
+  ret void
+}
+
+define void @foo2(i32 %arg, i32 %arg1, i32 %arg2, i32 %arg3) {
+; CHECK-LABEL: void @foo2
+bb:
+  %tmp8 = call i32  @goo() 
+  %tmp9 = icmp sgt i32 %tmp8, %arg
+  br i1 %tmp9, label %bb10, label %bb18
+; CHECK:  br i1 %tmp9
+; CHECK-NOT: !prof
+
+bb10:                                             ; preds = %bb
+  %tmp12 = call i32  @hoo()
+  %tmp13 = icmp sgt i32 %arg1, %tmp12
+  br i1 %tmp13, label %bb14, label %bb18
+; CHECK: br i1 %tmp13
+; CHECK-NOT: !prof
+
+bb14:                                             ; preds = %bb10
+  %tmp16 = call i32 @too()
+  %tmp17 = icmp sgt i32 %arg2, %tmp16
+  br label %bb18
+
+bb18:                                             ; preds = %bb14, %bb10, %bb
+  %tmp19 = phi i1 [ false, %bb10 ], [ false, %bb ], [ %tmp17, %bb14 ]
+  %tmp20 = xor i1 %tmp19, true
+  %tmp21 = xor i1 %tmp20, true
+  %tmp22 = zext i1 %tmp21 to i32
+  %tmp23 = sext i32 %tmp22 to i64
+  %tmp24 = call i64 @llvm.expect.i64(i64 %tmp23, i64 0)
+  %tmp25 = icmp ne i64 %tmp24, 0
+  br i1 %tmp25, label %bb26, label %bb28
+; CHECK: br i1 %tmp25,{{.*}}!prof [[WEIGHT2:![0-9]+]]
+
+bb26:                                             ; preds = %bb18
+  %tmp27 = call i32 @goo()
+  br label %bb30
+
+bb28:                                             ; preds = %bb18
+  %tmp29 = call i32 @hoo()
+  br label %bb30
+
+bb30:                                             ; preds = %bb28, %bb26
+  ret void
+}
+
+define void @foo_i32(i32 %arg, i32 %arg1, i32 %arg2, i32 %arg3) {
+; CHECK-LABEL: void @foo_i32
+bb:
+  %tmp8 = call i32  @goo() 
+  %tmp9 = icmp sgt i32 %tmp8, %arg
+  br i1 %tmp9, label %bb10, label %bb18
+; CHECK: !prof [[WEIGHT]]
+
+bb10:                                             ; preds = %bb
+  %tmp12 = call i32 @hoo()
+  %tmp13 = icmp sgt i32 %arg1, %tmp12
+  br i1 %tmp13, label %bb14, label %bb18
+; CHECK: br i1 %tmp13, {{.*}}!prof [[WEIGHT]]
+
+bb14:                                             ; preds = %bb10
+  %tmp16 = call i32 @too()
+  %tmp17 = icmp sgt i32 %arg2, %tmp16
+  br label %bb18
+
+bb18:                                             ; preds = %bb14, %bb10, %bb
+  %tmp19 = phi i32 [ 5, %bb10 ], [ 5, %bb ], [ %tmp16, %bb14 ]
+  %tmp23 = sext i32 %tmp19 to i64
+  %tmp24 = call i64 @llvm.expect.i64(i64 %tmp23, i64 4)
+  %tmp25 = icmp ne i64 %tmp24, 0
+  br i1 %tmp25, label %bb26, label %bb28
+; CHECK: br i1 %tmp25,{{.*}}!prof [[WEIGHT]]
+
+bb26:                                             ; preds = %bb18
+  %tmp27 = call i32 @goo()
+  br label %bb30
+
+bb28:                                             ; preds = %bb18
+  %tmp29 = call i32 @hoo()
+  br label %bb30
+
+bb30:                                             ; preds = %bb28, %bb26
+  ret void
+}
+
+
+define void @foo_i32_not_unlikely(i32 %arg, i32 %arg1, i32 %arg2, i32 %arg3)  {
+; CHECK-LABEL: void @foo_i32_not_unlikely
+bb:
+  %tmp8 = call i32 @goo() 
+  %tmp9 = icmp sgt i32 %tmp8, %arg
+  br i1 %tmp9, label %bb10, label %bb18
+; CHECK: br i1 %tmp9
+; CHECK-NOT: !prof
+
+bb10:                                             ; preds = %bb
+  %tmp12 = call i32 @hoo()
+  %tmp13 = icmp sgt i32 %arg1, %tmp12
+  br i1 %tmp13, label %bb14, label %bb18
+; CHECK: br i1 %tmp13
+; CHECK-NOT: !prof
+
+bb14:                                             ; preds = %bb10
+  %tmp16 = call i32  @too()
+  %tmp17 = icmp sgt i32 %arg2, %tmp16
+  br label %bb18
+
+bb18:                                             ; preds = %bb14, %bb10, %bb
+  %tmp19 = phi i32 [ 4, %bb10 ], [ 4, %bb ], [ %tmp16, %bb14 ]
+  %tmp23 = sext i32 %tmp19 to i64
+  %tmp24 = call i64 @llvm.expect.i64(i64 %tmp23, i64 4)
+  %tmp25 = icmp ne i64 %tmp24, 0
+  br i1 %tmp25, label %bb26, label %bb28
+; CHECK: br i1 %tmp25,{{.*}}!prof [[WEIGHT]]
+
+bb26:                                             ; preds = %bb18
+  %tmp27 = call i32  @goo()
+  br label %bb30
+
+bb28:                                             ; preds = %bb18
+  %tmp29 = call i32 @hoo()
+  br label %bb30
+
+bb30:                                             ; preds = %bb28, %bb26
+  ret void
+}
+
+define void @foo_i32_xor(i32 %arg, i32 %arg1, i32 %arg2, i32 %arg3)  {
+; CHECK-LABEL: void @foo_i32_xor
+bb:
+  %tmp8 = call i32  @goo() 
+  %tmp9 = icmp sgt i32 %tmp8, %arg
+  br i1 %tmp9, label %bb10, label %bb18
+; CHECK: br i1 %tmp9,{{.*}}!prof [[WEIGHT]]
+
+bb10:                                             ; preds = %bb
+  %tmp12 = call i32  @hoo()
+  %tmp13 = icmp sgt i32 %arg1, %tmp12
+  br i1 %tmp13, label %bb14, label %bb18
+; CHECK: br i1 %tmp13,{{.*}}!prof [[WEIGHT]]
+
+bb14:                                             ; preds = %bb10
+  %tmp16 = call i32  @too()
+  %tmp17 = icmp sgt i32 %arg2, %tmp16
+  br label %bb18
+
+bb18:                                             ; preds = %bb14, %bb10, %bb
+  %tmp19 = phi i32 [ 6, %bb10 ], [ 6, %bb ], [ %tmp16, %bb14 ]
+  %tmp20 = xor i32 %tmp19, 3
+  %tmp23 = sext i32 %tmp20 to i64
+  %tmp24 = call i64 @llvm.expect.i64(i64 %tmp23, i64 4)
+  %tmp25 = icmp ne i64 %tmp24, 0
+  br i1 %tmp25, label %bb26, label %bb28
+; CHECK: br i1 %tmp25,{{.*}}!prof [[WEIGHT]]
+
+bb26:                                             ; preds = %bb18
+  %tmp27 = call i32 @goo()
+  br label %bb30
+
+bb28:                                             ; preds = %bb18
+  %tmp29 = call i32 @hoo()
+  br label %bb30
+bb30:                                             ; preds = %bb28, %bb26
+  ret void
+}
+
+define void @foo_i8_sext(i32 %arg, i32 %arg1, i8 %arg2, i32 %arg3)  {
+; CHECK-LABEL: void @foo_i8_sext
+bb:
+  %tmp8 = call i32  @goo() 
+  %tmp9 = icmp sgt i32 %tmp8, %arg
+  br i1 %tmp9, label %bb10, label %bb18
+; CHECK: br i1 %tmp9,{{.*}}!prof [[WEIGHT]]
+
+bb10:                                             ; preds = %bb
+  %tmp12 = call i32  @hoo()
+  %tmp13 = icmp sgt i32 %arg1, %tmp12
+  br i1 %tmp13, label %bb14, label %bb18
+; CHECK: br i1 %tmp13,{{.*}}!prof [[WEIGHT]]
+
+bb14:                                             ; preds = %bb10
+  %tmp16 = call i8  @too8()
+  %tmp17 = icmp sgt i8 %arg2, %tmp16
+  br label %bb18
+
+bb18:                                             ; preds = %bb14, %bb10, %bb
+  %tmp19 = phi i8 [ 255, %bb10 ], [ 255, %bb ], [ %tmp16, %bb14 ]
+  %tmp23 = sext i8 %tmp19 to i64
+; after sign extension, the operand value becomes -1 which does not match 255
+  %tmp24 = call i64 @llvm.expect.i64(i64 %tmp23, i64 255)
+  %tmp25 = icmp ne i64 %tmp24, 0
+  br i1 %tmp25, label %bb26, label %bb28
+; CHECK: br i1 %tmp25,{{.*}}!prof [[WEIGHT]]
+
+bb26:                                             ; preds = %bb18
+  %tmp27 = call i32 @goo()
+  br label %bb30
+
+bb28:                                             ; preds = %bb18
+  %tmp29 = call i32 @hoo()
+  br label %bb30
+bb30:                                             ; preds = %bb28, %bb26
+  ret void
+}
+
+define void @foo_i8_sext_not_unlikely(i32 %arg, i32 %arg1, i8 %arg2, i32 %arg3)  {
+; CHECK-LABEL: void @foo_i8_sext_not_unlikely
+bb:
+  %tmp8 = call i32  @goo() 
+  %tmp9 = icmp sgt i32 %tmp8, %arg
+  br i1 %tmp9, label %bb10, label %bb18
+; CHECK: br i1 %tmp9
+; CHECK-NOT: !prof
+
+bb10:                                             ; preds = %bb
+  %tmp12 = call i32  @hoo()
+  %tmp13 = icmp sgt i32 %arg1, %tmp12
+  br i1 %tmp13, label %bb14, label %bb18
+; CHECK: br i1 %tmp13
+; CHECK-NOT: !prof
+
+bb14:                                             ; preds = %bb10
+  %tmp16 = call i8  @too8()
+  %tmp17 = icmp sgt i8 %arg2, %tmp16
+  br label %bb18
+
+bb18:                                             ; preds = %bb14, %bb10, %bb
+  %tmp19 = phi i8 [ 255, %bb10 ], [ 255, %bb ], [ %tmp16, %bb14 ]
+  %tmp23 = sext i8 %tmp19 to i64
+; after sign extension, the operand value becomes -1 which matches -1
+  %tmp24 = call i64 @llvm.expect.i64(i64 %tmp23, i64 -1)
+  %tmp25 = icmp ne i64 %tmp24, 0
+  br i1 %tmp25, label %bb26, label %bb28
+; CHECK: br i1 %tmp25,{{.*}}!prof [[WEIGHT]]
+
+bb26:                                             ; preds = %bb18
+  %tmp27 = call i32 @goo()
+  br label %bb30
+
+bb28:                                             ; preds = %bb18
+  %tmp29 = call i32 @hoo()
+  br label %bb30
+bb30:                                             ; preds = %bb28, %bb26
+  ret void
+}
+
+
+define void @foo_i32_xor_not_unlikely(i32 %arg, i32 %arg1, i32 %arg2, i32 %arg3)  {
+; CHECK-LABEL: void @foo_i32_xor_not_unlikely
+bb:
+  %tmp8 = call i32 @goo() 
+  %tmp9 = icmp sgt i32 %tmp8, %arg
+  br i1 %tmp9, label %bb10, label %bb18
+; CHECK: br i1 %tmp9
+; CHECK-NOT: !prof
+
+bb10:                                             ; preds = %bb
+  %tmp12 = call i32  @hoo()
+  %tmp13 = icmp sgt i32 %arg1, %tmp12
+  br i1 %tmp13, label %bb14, label %bb18
+; CHECK: br i1 %tmp13
+; CHECK-NOT: !prof
+
+bb14:                                             ; preds = %bb10
+  %tmp16 = call i32 @too()
+  %tmp17 = icmp sgt i32 %arg2, %tmp16
+  br label %bb18
+
+bb18:                                             ; preds = %bb14, %bb10, %bb
+  %tmp19 = phi i32 [ 6, %bb10 ], [ 6, %bb ], [ %tmp16, %bb14 ]
+  %tmp20 = xor i32 %tmp19, 2
+  %tmp23 = sext i32 %tmp20 to i64
+  %tmp24 = call i64 @llvm.expect.i64(i64 %tmp23, i64 4)
+  %tmp25 = icmp ne i64 %tmp24, 0
+  br i1 %tmp25, label %bb26, label %bb28
+; CHECK: br i1 %tmp25,{{.*}}!prof [[WEIGHT]]
+
+bb26:                                             ; preds = %bb18
+  %tmp27 = call i32 @goo()
+  br label %bb30
+
+bb28:                                             ; preds = %bb18
+  %tmp29 = call i32  @hoo()
+  br label %bb30
+
+bb30:                                             ; preds = %bb28, %bb26
+  ret void
+}
+
+declare i32 @goo()
+
+declare i32 @hoo()
+
+declare i32 @too()
+
+declare i8 @too8()
+
+; Function Attrs: nounwind readnone
+declare i64 @llvm.expect.i64(i64, i64) 
+
+!llvm.ident = !{!0}
+
+!0 = !{!"clang version 5.0.0 (trunk 302965)"}
+; CHECK: [[WEIGHT]] = !{!"branch_weights", i32 2000, i32 1}
+; CHECK: [[WEIGHT2]] = !{!"branch_weights", i32 1, i32 2000}

llvm/trunk/test/Transforms/LowerExpectIntrinsic/phi_or.ll

+; RUN: opt -lower-expect  -S -o - < %s | FileCheck %s
+; RUN: opt -S -passes='function(lower-expect)' < %s | FileCheck %s
+; 
+; if (__builtin_expect((x > goo() || y > hoo()), 1)) {
+;  ..
+; }
+; For the above case, only the second branch should be
+; annotated.
+; if (__builtin_expect((x > goo() || y > hoo()), 0)) {
+;  ..
+; }
+; For the above case, two branches should be annotated.
+; Function Attrs: noinline nounwind uwtable
+define void @foo(i32 %arg, i32 %arg1, i32 %arg2, i32 %arg3)  {
+; CHECK-LABEL: void @foo
+bb:
+  %tmp8 = call i32 @goo()
+  %tmp9 = icmp slt i32 %arg, %tmp8
+  br i1 %tmp9, label %bb14, label %bb10
+; CHECK: br i1 %tmp9
+; CHECK-NOT: br i1 %tmp9{{.*}}!prof
+
+bb10:                                             ; preds = %bb
+  %tmp12 = call i32  @hoo()
+  %tmp13 = icmp sgt i32 %arg1, %tmp12
+  br label %bb14
+
+bb14:                                             ; preds = %bb10, %bb
+  %tmp15 = phi i1 [ true, %bb ], [ %tmp13, %bb10 ]
+  %tmp16 = zext i1 %tmp15 to i32
+  %tmp17 = sext i32 %tmp16 to i64
+  %expect = call i64 @llvm.expect.i64(i64 %tmp17, i64 1)
+  %tmp18 = icmp ne i64 %expect, 0
+  br i1 %tmp18, label %bb19, label %bb21
+; CHECK: br i1 %tmp18{{.*}}!prof [[WEIGHT:![0-9]+]]
+
+bb19:                                             ; preds = %bb14
+  %tmp20 = call i32 @goo()
+  br label %bb23
+
+bb21:                                             ; preds = %bb14
+  %tmp22 = call i32  @hoo()
+  br label %bb23
+
+bb23:                                             ; preds = %bb21, %bb19
+  ret void
+}
+
+define void @foo2(i32 %arg, i32 %arg1, i32 %arg2, i32 %arg3)  {
+; CHECK-LABEL: void @foo2
+bb:
+  %tmp = alloca i32, align 4
+  %tmp4 = alloca i32, align 4
+  %tmp5 = alloca i32, align 4
+  %tmp6 = alloca i32, align 4
+  store i32 %arg, i32* %tmp, align 4
+  store i32 %arg1, i32* %tmp4, align 4
+  store i32 %arg2, i32* %tmp5, align 4
+  store i32 %arg3, i32* %tmp6, align 4
+  %tmp7 = load i32, i32* %tmp, align 4
+  %tmp8 = call i32  @goo()
+  %tmp9 = icmp slt i32 %tmp7, %tmp8
+  br i1 %tmp9, label %bb14, label %bb10
+; CHECK: br i1 %tmp9{{.*}}!prof [[WEIGHT2:![0-9]+]]
+
+bb10:                                             ; preds = %bb
+  %tmp11 = load i32, i32* %tmp5, align 4
+  %tmp12 = call i32 @hoo()
+  %tmp13 = icmp sgt i32 %tmp11, %tmp12
+  br label %bb14
+
+bb14:                                             ; preds = %bb10, %bb
+  %tmp15 = phi i1 [ true, %bb ], [ %tmp13, %bb10 ]
+  %tmp16 = zext i1 %tmp15 to i32
+  %tmp17 = sext i32 %tmp16 to i64
+  %expect = call i64 @llvm.expect.i64(i64 %tmp17, i64 0)
+  %tmp18 = icmp ne i64 %expect, 0
+  br i1 %tmp18, label %bb19, label %bb21
+; CHECK: br i1 %tmp18{{.*}}!prof [[WEIGHT2]]
+
+bb19:                                             ; preds = %bb14
+  %tmp20 = call i32 @goo()
+  br label %bb23
+
+bb21:                                             ; preds = %bb14
+  %tmp22 = call i32 @hoo()
+  br label %bb23
+
+bb23:                                             ; preds = %bb21, %bb19
+  ret void
+}
+
+declare i32 @goo() 
+declare i32 @hoo() 
+declare i64 @llvm.expect.i64(i64, i64) 
+
+
+!llvm.ident = !{!0}
+
+
+!0 = !{!"clang version 5.0.0 (trunk 302965)"}
+; CHECK: [[WEIGHT]] = !{!"branch_weights", i32 2000, i32 1}
+; CHECK: [[WEIGHT2]] = !{!"branch_weights", i32 1, i32 2000}

llvm/trunk/test/Transforms/LowerExpectIntrinsic/phi_tern.ll

+; RUN: opt -lower-expect  -S -o - < %s | FileCheck %s
+; RUN: opt -S -passes='function(lower-expect)' < %s | FileCheck %s
+
+; return __builtin_expect((a > b ? 1, goo(), 0);
+;  
+; Function Attrs: noinline nounwind uwtable
+define i32 @foo(i32 %arg, i32 %arg1)  {
+; CHECK-LABEL: i32 @foo
+bb:
+  %tmp5 = icmp sgt i32 %arg, %arg1
+  br i1 %tmp5, label %bb9, label %bb7
+; CHECK: br i1 %tmp5{{.*}}!prof [[WEIGHT:![0-9]+]]
+
+bb7:                                              ; preds = %bb
+  %tmp8 = call i32 @goo()
+  br label %bb9
+
+bb9:                                              ; preds = %bb7, %bb9
+  %tmp10 = phi i32 [ 1, %bb ], [ %tmp8, %bb7 ]
+  %tmp11 = sext i32 %tmp10 to i64
+  %expect = call i64 @llvm.expect.i64(i64 %tmp11, i64 0)
+  %tmp12 = trunc i64 %expect to i32
+  ret i32 %tmp12
+}
+
+define i32 @foo2(i32 %arg, i32 %arg1)  {
+bb:
+  %tmp5 = icmp sgt i32 %arg, %arg1
+  br i1 %tmp5, label %bb6, label %bb7
+; CHECK: br i1 %tmp5{{.*}}!prof [[WEIGHT:![0-9]+]]
+
+bb6:                                              ; preds = %bb
+  br label %bb9
+
+bb7:                                              ; preds = %bb
+  %tmp8 = call i32 @goo()
+  br label %bb9
+
+bb9:                                              ; preds = %bb7, %bb6
+  %tmp10 = phi i32 [ 1, %bb6 ], [ %tmp8, %bb7 ]
+  %tmp11 = sext i32 %tmp10 to i64
+  %expect = call i64 @llvm.expect.i64(i64 %tmp11, i64 0)
+  %tmp12 = trunc i64 %expect to i32
+  ret i32 %tmp12
+}
+
+declare i32 @goo() 
+declare i64 @llvm.expect.i64(i64, i64) 
+
+
+
+!llvm.ident = !{!0}
+
+!0 = !{!"clang version 5.0.0 (trunk 302965)"}
+
+; CHECK: [[WEIGHT]] = !{!"branch_weights", i32 1, i32 2000}