[JumpThreading] Propagate branch hint metadata in more cases

[inouehrs]

Summary

updatePredecessorProfileMetadata propagates branch metadata to keep the information after jump-threading eliminates the branch with metadata.
This patch extends updatePredecessorProfileMetadata in two ways:

1. The current updatePredecessorProfileMetadata propagates metadata only when the boolean value used in the conditional branch directly comes from comes from previous blocks. E.g. it supports conditional branch such as if (__builtin_expect(foo(), 0)).

This patch allows one additional icmp node before the conditional branch to catch common cases. E.g. this patch supports if (__builtin_expect(bar() == nullptr, 0)) as well.

2. This patch enhances GetPredOutEdge to find the branch to be associated with the metadata in general control flow case. The original GetPredOutEdge lambda only allows the simplest cases.

[nemanjai]

lib/Transforms/Scalar/JumpThreading.cpp

103

Maybe a comment mentioning why this default value (i.e. there's a bunch of mentions of 1:2000 identifying cold blocks, but this is set to 1000).

515

Please add an inline comment about the requirement for the MD having 3 operands.

520

Explicit comparisons against NULL seem at odds with other null checks added here.

537

I think nullptr is a more common way to initialize a null pointer. Same goes for other uses of NULL.

552

In general, I would recommend that a function with output parameters only set those output parameters on success. If there's a compelling reason to set WeightHot and WeightCold when returning nullptr, I think the comment at the top of the function should mention that behaviour.

562

Nit: please be careful with line length.

565

The logic in the function seems straightforward enough - looks like the cold region is all blocks dominated by ColdBlock and blocks post-dominated by ColdBlock (unless I'm misreading something). Then we set the branch weight on branches into that region.
Just seems like we should just be able to use the Dominators/PostDominators analyses for this kind of thing, but I very well may be missing something. In any case, if the existing analyses can't be used (even if it is an issue of analyses dependency) I think you should add a comment as to why that is.

1617

s/upword/upward

1618

This formatting looks odd. I almost missed the else.

[inouehrs]

Updated based on Nemanja's comments. Also rebased to the latest code base.

[inouehrs]

lib/Transforms/Scalar/JumpThreading.cpp

103

For the default threshold, I modify to use 2000 (instead of 1000 in the previous patch) as the default value based on the _builtin_expect handling.

515

Added comment.

520

Thank you. I changed the style.

552

I modified the code to set WeightHot and WeightCold only when it returns the non-null pointer as the return value.

562

Fixed.

565

I added the comments. I use my simple implementation of dominance analysis because the CFG may be already modified if we have multiple branches with biased branch weights.
Also, a method with branch weight metadata is relatively rare and so I avoid requesting dominance analysis only for this purpose.

1617

Oops. Fixed.

1618

Fixed the style.

[inouehrs]

Gentle ping. I appreciate any suggestions.

[inouehrs]

• rebased to the latest source tree.
• fixed coding style issues

[inouehrs]

Gentle ping.
Thank you.

[inouehrs]

• rebased to the latest source tree.
• fixed build break

[inouehrs]

• rebased to the latest codebase

[inouehrs]

• fixed build break with the latest source

[davidxl]

I have not looked at this patch in details, but from the test case, it seems that the problem should be handled by the updatePredecessorProfileMetadata () method. Have you checked why the existing update does not kick in?

[inouehrs]

@davidxl Thank you for the comment. I have never realized your patch D36864 added similar functionality.
I will investigate updatePredecessorProfileMetadata.

[inouehrs]

• re-implemented by piggybacking on https://reviews.llvm.org/D36864

[inouehrs]

@davidxl I reimplemented functionality of my patch by enhancing your code. I appreciate any comments. Thanks!

[inouehrs]

• rebased to the latest tree and tested again
• minor touch up

[inouehrs]

• rebase to the latest
• code cleanup

lib/Transforms/Scalar/JumpThreading.cpp

ImplicationSearchThreshold(
   "jump-threading-implication-search-threshold",
   cl::desc("The number of predecessors to search for a stronger "
            "condition to use to thread over a weaker condition"),
   cl::init(3), cl::Hidden);
 
 static cl::opt<bool> PrintLVIAfterJumpThreading(
     "print-lvi-after-jump-threading",
     cl::desc("Print the LazyValueInfo cache after JumpThreading"), cl::init(false),
     cl::Hidden);
 
 namespace {
 
   /// This pass performs 'jump threading', which looks at blocks that have
   /// multiple predecessors and multiple successors.  If one or more of the
   /// predecessors of the block can be proven to always jump to one of the
   /// successors, we forward the edge from the predecessor to the successor by
   /// duplicating the contents of this block.

 FunctionPass *llvm::createJumpThreadingPass(int Threshold) {
   return new JumpThreading(Threshold);
 }
 
 JumpThreadingPass::JumpThreadingPass(int T) {
   BBDupThreshold = (T == -1) ? BBDuplicateThreshold : unsigned(T);
 }
 
+// Returns the single outgoing edge of the dominating predecessor block
+// that leads to the PhiNode's incoming block:
+static std::pair<BasicBlock *, BasicBlock *>
+GetPredOutEdge(BasicBlock *IncomingBB, BasicBlock *PhiBB) {
+  // If the incoming edge is conditional branch, we just select it.
+  BranchInst *PredBr = dyn_cast<BranchInst>(IncomingBB->getTerminator());
+  if (PredBr && PredBr->isConditional())
+    return {IncomingBB, PhiBB};
+
+  // We identify the BBs post-dominated by IncomingBB.
+  // We do not use the existing DominatorTree analyses since the CFG may be
+  // already changed if we have multiple branches with branch hint.
+  SmallVector<BasicBlock *, 16> PostDominatedBBs;
+  PostDominatedBBs.push_back(IncomingBB);
+  unsigned Idx = 0;
+  while (Idx < PostDominatedBBs.size()) {
+    BasicBlock *CBB = PostDominatedBBs[Idx++];
+    for (BasicBlock *PrevBB : predecessors(CBB)) {
+      if (is_contained(PostDominatedBBs, PrevBB))
+        continue;
+
+      // If all the successors are known to be post-dominated by IncomingBB,
+      // this block is also post-dominated.
+      if (all_of(successors(PrevBB), [&](BasicBlock *SBB) {
+            return is_contained(PostDominatedBBs, SBB);
+          }))
+        PostDominatedBBs.push_back(PrevBB);
+    }
+  }
+
+  BasicBlock *PredBB = nullptr, *SuccBB = nullptr;
+  for (BasicBlock *CBB : PostDominatedBBs) {
+    for (BasicBlock *PrevBB : predecessors(CBB)) {
+      // If a previous block is not post-dominated by IncomingBB,
+      // this block is at the dominance frontier.
+      if (is_contained(PostDominatedBBs, PrevBB))
+        continue;
+
+      BranchInst *BI = dyn_cast<BranchInst>(PrevBB->getTerminator());
+      if (!BI || !BI->isConditional())
+        continue;
+
+      // We found a CFG edge that leads to the incoming block.
+      // If we already found another, we return nullptr.
+      if (PredBB)
+        return {nullptr, nullptr};
+      PredBB = PrevBB;
+      SuccBB = CBB;
+    }
+  }
+  return {PredBB, SuccBB};
+}
+
 // Update branch probability information according to conditional
 // branch probablity. This is usually made possible for cloned branches
 // in inline instances by the context specific profile in the caller.
 // For instance,
 //
 //  [Block PredBB]
 //  [Branch PredBr]
 //  if (t) {

 //
 //  which gives us:
 //     P(A) is less than P(cond == true), i.e.
 //     P(t == true) <= P(cond == true)
 //
 //  In other words, if we know P(cond == true) is unlikely, we know
 //  that P(t == true) is also unlikely.
 //
-static void updatePredecessorProfileMetadata(PHINode *PN, BasicBlock *BB) {
+static void
+updatePredecessorProfileMetadata(Instruction *CondInst, BasicBlock *BB) {
+  PHINode *PN = dyn_cast<PHINode>(CondInst);
+  ICmpInst *CmpNode = nullptr;
+  if (!PN) {
+    // We allow phi + icmp with an immediate + br as well as phi + br.
+    CmpNode = dyn_cast<ICmpInst>(CondInst);
+    if (CmpNode && isa<ConstantInt>(CmpNode->getOperand(1)))
+      PN = dyn_cast<PHINode>(CmpNode->getOperand(0));
+
+    if (!PN)
+      return;
+  }
+
   BranchInst *CondBr = dyn_cast<BranchInst>(BB->getTerminator());
   if (!CondBr)
     return;
 
   BranchProbability BP;
   uint64_t TrueWeight, FalseWeight;
   if (!CondBr->extractProfMetadata(TrueWeight, FalseWeight))
     return;
 
-  // Returns the outgoing edge of the dominating predecessor block
-  // that leads to the PhiNode's incoming block:
-  auto GetPredOutEdge =
-      [](BasicBlock *IncomingBB,
-         BasicBlock *PhiBB) -> std::pair<BasicBlock *, BasicBlock *> {
-    auto *PredBB = IncomingBB;
-    auto *SuccBB = PhiBB;
-    while (true) {
-      BranchInst *PredBr = dyn_cast<BranchInst>(PredBB->getTerminator());
-      if (PredBr && PredBr->isConditional())
-        return {PredBB, SuccBB};
-      auto *SinglePredBB = PredBB->getSinglePredecessor();
-      if (!SinglePredBB)
-        return {nullptr, nullptr};
-      SuccBB = PredBB;
-      PredBB = SinglePredBB;
-    }
-  };
-
   for (unsigned i = 0, e = PN->getNumIncomingValues(); i != e; ++i) {
     Value *PhiOpnd = PN->getIncomingValue(i);
     ConstantInt *CI = dyn_cast<ConstantInt>(PhiOpnd);
 
-    if (!CI || !CI->getType()->isIntegerTy(1))
+    if (!CI || !CI->getType()->isIntegerTy())
+      continue;
+
+    bool IsTrueEdge = false;
+    if (!CmpNode) {
+      if (!CI->getType()->isIntegerTy(1))
         continue;
 
-    BP = (CI->isOne() ? BranchProbability::getBranchProbability(
+      IsTrueEdge = CI->isOne();
+    }
+    else {
+      APInt Imm = (dyn_cast<ConstantInt>(CmpNode->getOperand(1)))->getValue();
+      switch(CmpNode->getPredicate()) {
+        case ICmpInst::ICMP_EQ:
+          IsTrueEdge = (CI->getValue() == Imm);
+          break;
+        case ICmpInst::ICMP_NE:
+          IsTrueEdge = (CI->getValue() != Imm);
+          break;
+        case ICmpInst::ICMP_ULT:
+          IsTrueEdge = CI->getValue().ult(Imm);
+          break;
+        case ICmpInst::ICMP_SLT:
+          IsTrueEdge = CI->getValue().slt(Imm);
+          break;
+        case ICmpInst::ICMP_UGT:
+          IsTrueEdge = CI->getValue().ugt(Imm);
+          break;
+        case ICmpInst::ICMP_SGT:
+          IsTrueEdge = CI->getValue().sgt(Imm);
+          break;
+        default: llvm_unreachable("Unhandled icmp opcode!");
+      }
+    }
+    BP = (IsTrueEdge ? BranchProbability::getBranchProbability(
                             TrueWeight, TrueWeight + FalseWeight)
-                      : BranchProbability::getBranchProbability(
                             FalseWeight, TrueWeight + FalseWeight));
 
     auto PredOutEdge = GetPredOutEdge(PN->getIncomingBlock(i), BB);
     if (!PredOutEdge.first)
       return;
 
     BasicBlock *PredBB = PredOutEdge.first;
     BranchInst *PredBr = cast<BranchInst>(PredBB->getTerminator());

 // guard/assume.
 static void ReplaceFoldableUses(Instruction *Cond, Value *ToVal) {
   assert(Cond->getType() == ToVal->getType());
   auto *BB = Cond->getParent();
   // We can unconditionally replace all uses in non-local blocks (i.e. uses
   // strictly dominated by BB), since LVI information is true from the
   // terminator of BB.
   replaceNonLocalUsesWith(Cond, ToVal);
   for (Instruction &I : reverse(*BB)) {
     // Reached the Cond whose uses we are trying to replace, so there are no
     // more uses.
     if (&I == Cond)
       break;
     // We only replace uses in instructions that are guaranteed to reach the end
     // of BB, where we know Cond is ToVal.
     if (!isGuaranteedToTransferExecutionToSuccessor(&I))
       break;
     I.replaceUsesOfWith(Cond, ToVal);
   }
   if (Cond->use_empty() && !Cond->mayHaveSideEffects())
     Cond->eraseFromParent();
 }
 
 /// Return the cost of duplicating a piece of this block from first non-phi
 /// and before StopAt instruction to thread across it. Stop scanning the block
 /// when exceeding the threshold. If duplication is impossible, returns ~0U.
 static unsigned getJumpThreadDuplicationCost(BasicBlock *BB,
                                              Instruction *StopAt,
                                              unsigned Threshold) {
   assert(StopAt->getParent() == BB && "Not an instruction from proper BB?");
   /// Ignore PHI nodes, these will be flattened when duplication happens.
   BasicBlock::const_iterator I(BB->getFirstNonPHI());
 
   // FIXME: THREADING will delete values that are just used to compute the
   // branch, so they shouldn't count against the duplication cost.
 
   unsigned Bonus = 0;
   if (BB->getTerminator() == StopAt) {
     // Threading through a switch statement is particularly profitable.  If this
     // block ends in a switch, decrease its cost to make it more likely to
     // happen.
     if (isa<SwitchInst>(StopAt))
       Bonus = 6;
 
     // The same holds for indirect branches, but slightly more so.
     if (isa<IndirectBrInst>(StopAt))
       Bonus = 8;
   }
 
   // Bump the threshold up so the early exit from the loop doesn't skip the
   // terminator-based Size adjustment at the end.
   Threshold += Bonus;
 
   // Sum up the cost of each instruction until we get to the terminator.  Don't
   // include the terminator because the copy won't include it.
   unsigned Size = 0;
   for (; &*I != StopAt; ++I) {
 
     // Stop scanning the block if we've reached the threshold.
     if (Size > Threshold)
       return Size;
 
     // Debugger intrinsics don't incur code size.
     if (isa<DbgInfoIntrinsic>(I)) continue;
 
     // If this is a pointer->pointer bitcast, it is free.
     if (isa<BitCastInst>(I) && I->getType()->isPointerTy())
bool JumpThreadingPass::ProcessBlock(BasicBlock *BB) {
 
   // TODO: There are other places where load PRE would be profitable, such as
   // more complex comparisons.
   if (LoadInst *LoadI = dyn_cast<LoadInst>(SimplifyValue))
     if (SimplifyPartiallyRedundantLoad(LoadI))
       return true;
 
   // Before threading, try to propagate profile data backwards:
-  if (PHINode *PN = dyn_cast<PHINode>(CondInst))
-    if (PN->getParent() == BB && isa<BranchInst>(BB->getTerminator()))
-      updatePredecessorProfileMetadata(PN, BB);
+  if (CondInst->getParent() == BB && isa<BranchInst>(BB->getTerminator()))
+    updatePredecessorProfileMetadata(CondInst, BB);
 
   // Handle a variety of cases where we are branching on something derived from
   // a PHI node in the current block.  If we can prove that any predecessors
   // compute a predictable value based on a PHI node, thread those predecessors.
   if (ProcessThreadableEdges(CondInst, BB, Preference, Terminator))
     return true;
 
   // If this is an otherwise-unfoldable branch on a phi node in the current
FindMostPopularDest(BasicBlock *BB,
   if (!SamePopularity.empty()) {
     SamePopularity.push_back(MostPopularDest);
     TerminatorInst *TI = BB->getTerminator();
     for (unsigned i = 0; ; ++i) {
       assert(i != TI->getNumSuccessors() && "Didn't find any successor!");
 
       if (!is_contained(SamePopularity, TI->getSuccessor(i)))
         continue;
 
       MostPopularDest = TI->getSuccessor(i);
       break;
     }
   }
 
   // Okay, we have finally picked the most popular destination.
   return MostPopularDest;
 }
 

test/Transforms/JumpThreading/branch-metadata.ll

+; RUN: opt < %s -jump-threading -S | FileCheck %s
+
+; This test examines how jump-threading handles branch weights metadata
+; if the branch to be eliminated has the metadata.
+
+define i32 @func(i32 %flag) {
+; CHECK-LABEL: @func
+; CHECK: br i1 %tobool.i, label %if.hot, label %if.cold, !prof !0
+; CHECK-DAG: if.hot:
+; CHECK-DAG: if.cold:
+entry:
+  %tobool.i = icmp eq i32 %flag, 0
+  br i1 %tobool.i, label %to.be.eliminated, label %if.then
+
+if.then:
+  br label %to.be.eliminated
+
+to.be.eliminated:
+  %retval.0.i = phi i32 [ 0, %if.then ], [ 1, %entry ]
+  %cmp = icmp eq i32 %retval.0.i, 0
+  br i1 %cmp, label %if.cold, label %if.hot, !prof !1
+
+if.cold:
+  call void @cold_func()
+  br label %return
+
+if.hot:
+  call void @hot_func()
+  br label %return
+
+return:
+  %retval.0 = phi i32 [ 0, %if.cold ], [ 1, %if.hot ]
+  ret i32 %retval.0
+}
+
+define i32 @func2(i32 %flag) {
+; CHECK-LABEL: @func2
+; CHECK: br i1 %tobool.i, label %if.cold, label %if.hot, !prof !1
+; CHECK-DAG: if.hot:
+; CHECK-DAG: if.cold:
+entry:
+  %tobool.i = icmp eq i32 %flag, 0
+  br i1 %tobool.i, label %to.be.eliminated, label %if.then
+
+if.then:
+  br label %to.be.eliminated
+
+to.be.eliminated:
+  %retval.0.i = phi i32 [ 0, %if.then ], [ 1, %entry ]
+  %cmp = icmp eq i32 %retval.0.i, 0
+  br i1 %cmp, label %if.hot, label %if.cold, !prof !0
+
+if.cold:
+  call void @cold_func()
+  br label %return
+
+if.hot:
+  call void @hot_func()
+  br label %return
+
+return:
+  %retval.0 = phi i32 [ 0, %if.cold ], [ 1, %if.hot ]
+  ret i32 %retval.0
+}
+
+
+define signext i32 @func3(i32 signext %flag) {
+; This is a test with more complicated CFG.
+; The following is the equivalent C code.
+; We want to set branch weights for "if (flag & 1)" and "if (flag & 2)".
+;
+; inline int bar(int flag) {
+;   if (flag & 1) {
+;     cold_func();
+;     return 0;
+;   }
+;   if (flag & 2) {
+;     if (__builtin_expect(flag & 4, 0)) cold_func();
+;     cold_func();
+;     return 0;
+;   }
+;   if (__builtin_expect(flag & 8, 0)) cold_func();
+;   return 1;
+; }
+;
+; int func3(int flag) {
+;   if (__builtin_expect(0 == bar(flag), 0)) {
+;     cold_func();
+;     return 0;
+;   }
+;   hot_func();
+;   return 1;
+; }
+
+; CHECK-LABEL: @func3
+; CHECK: br i1 %tobool.i, label %if.end.i, label %if.then.i, !prof !0
+; CHECK: if.then.i:
+; CHECK: br i1 %tobool2.i, label %if.end8.i, label %if.then3.i, !prof !0
+; CHECK: if.then3.i:
+
+entry:
+  %and.i = and i32 %flag, 1
+  %tobool.i = icmp eq i32 %and.i, 0
+  br i1 %tobool.i, label %if.end.i, label %if.then.i
+
+if.then.i:
+  call void @cold_func()
+  br label %bar.exit
+
+if.end.i:
+  %and1.i = and i32 %flag, 2
+  %tobool2.i = icmp eq i32 %and1.i, 0
+  br i1 %tobool2.i, label %if.end8.i, label %if.then3.i
+
+if.then3.i:
+  %and4.i = and i32 %flag, 4
+  %tobool5.i = icmp eq i32 %and4.i, 0
+  br i1 %tobool5.i, label %if.end7.i, label %if.then6.i, !prof !0
+
+if.then6.i:
+  call void @cold_func()
+  br label %if.end7.i
+
+if.end7.i:
+  call void @cold_func()
+  br label %bar.exit
+
+if.end8.i:
+  %and9.i = and i32 %flag, 8
+  %tobool12.i = icmp eq i32 %and9.i, 0
+  br i1 %tobool12.i, label %bar.exit, label %if.then13.i, !prof !0
+
+if.then13.i:
+  call void @cold_func()
+  br label %bar.exit
+
+bar.exit:
+  %retval.0.i = phi i32 [ 0, %if.then.i ], [ 0, %if.end7.i ], [ 1, %if.end8.i ], [ 1, %if.then13.i ]
+  %cmp = icmp eq i32 %retval.0.i, 0
+  br i1 %cmp, label %if.then, label %if.end, !prof !1
+
+if.then:
+  call void @cold_func()
+  br label %return
+
+if.end:
+  call void @hot_func()
+  br label %return
+
+return:
+  %retval.0 = phi i32 [ 0, %if.then ], [ 1, %if.end ]
+  ret i32 %retval.0
+}
+
+; CHECK-DAG: !0 = !{!"branch_weights", i32 2146410443, i32 1073205}
+; CHECK-DAG: !1 = !{!"branch_weights", i32 1073205, i32 2146410443}
+
+
+declare void @hot_func()
+declare void @cold_func()
+!0 = !{!"branch_weights", i32 2000, i32 1}
+!1 = !{!"branch_weights", i32 1, i32 2000}