[Profile] backward propagate profile data in jump-threading

[davidxl]

Summary

If a branch condition is defined by a function call result, the profile information for the branch may be (effectively) lost after inlinling.

For instance,

 bool cond() {
      if (a)
        return true;
      if (b)
        return true;
     return false;
 }

void test() {
    if (__builtin_expect(cond(), false) ) {
         bar();  // cold
     }
   ..
 }

After inlining and jump threading, the first branch in the inline instance does not have profile data -- this will make the block for bar() look hot.

In this patch, the problem is fixed in jump threading by 'backward' propagate the BP to the predecessor given the conditional probability discovered.

[chandlerc]

Really, really nice find here. A bunch of comments below, but they're all pretty minor on the whole.

lib/Transforms/Scalar/JumpThreading.cpp

159

Would be good to connect this comment to the arguments and behavior of this function.

(Don't get me wrong, this comment is *excellent* for explaining why we can do this back propagation and how to compute it! I just also want comments that explain how this routine fits into the model.)

160

Update -> update

175

You can use return {IncomingBB, PhiBB}; here and elsewhere to return. You'll have to specify the return type in the lambda, but then you won't even need to do that for nullptrs...

176–177

Why just one? You could do:

auto *PredBB = IncomingBB;
while (auto *SinglePredBB = PredBB->getSinglePredecessor())

PredBB = SinglePredBB;

at the top fo the function, and have a single set of logic to handle the whole thing?

189–192

Can simplify as:

if (!CI || !CI->getType()->isIntegerTy(1))
  continue;

194

This code only handle i1 PHIs that directly feed conditional branches, right?

Provided that's so, I would assert that CI is a 1 bit integer here, and then test whether it the null value or not rather than zexting and extracting it.

200

No need to compare with nullptr, pointers are testable directly:

if (PredOutEdge.first)

207

"we" -> "We"

I wonder, since what we have is fundamentally an upper bound, would it make sense to go ahead and apply this even when there is existing profile info?

Even for static hints, I can see this happening. For example, the pre-inlining code might have some relatively weak hint that is, after inlining, bounded by a powerful hint.

Or is your concern doing this without measuring the impact for PGO? If that's why, I would make that more clear -- we *intend* to always use this to refine weights, but PGO requires more careful tuning to make sure this doesn't regress anything.

Otherwise, I'd just go ahead and implement this as clamping any existing weights.

208–209

"With PGO, this can be turned on too to refine the profile with" -> "With PGO, this can be used to refine even existing profile data with"

213

I would remind the reader that this is because all we have is an upper bound.

[davidxl]

lib/Transforms/Scalar/JumpThreading.cpp

159

Annotated the comments with variable names in the code.

194

There is an early check of one bit type. I am not clear what is the suggestion about (regarding extracting value).

207

Yes, the latter . I have seen issues with PGO so without more testing, I will leave this for nonPGO case for now.

[davidxl]

Addressed Chandler's review feedbacks

[chandlerc]

Awesome. LGTM with the minor things below addressed. Let me know if any of this isn't make sense.

lib/Transforms/Scalar/JumpThreading.cpp

163

nit: spurious blank line...

176

nit: clang-format to fix line length

194

Essentially...

assert(CI.getBitWidth() == 1 && "We only handle i1 constants here.");
BP = CI->isOne() ? ... : ...;

You could also use CI->isZero(), but one seems more clear.

208–209

This wording improvement for the comment would still be helpful I think.

[davidxl]

Fix an obvious bug.

[davidxl]

Also improved test cases to cover the path that exposed the bug.

[chandlerc]

Cool! With the changes below, LGTM.

Any further testing you think this needs before relanding it? (I don't really have good ideas other than to patch it in internally and build some large things with PGO and ThinLTO)

lib/Transforms/Scalar/JumpThreading.cpp

184–188

I think this will be easier to read inverted:

auto *SinglePredBB = PredBB->getSinglePredecessor();
if (!SinglePredBB)
  return {nullptr, nullptr};

SuccBB = PredBB;
PredBB = SinglePredBB;

189

nit: When infinitely looping, I find for (;;) { a shorter and still easy to read idiom.

190–191

No need for this once the loop above returns directly when exiting.

[davidxl]

Addressed review feedback.

Tested with clang self build with PGO.

[LukeZhuang]

lib/Transforms/Scalar/JumpThreading.cpp

206

Hi David, may I ask why this is a return instead of continue here? If a phi node is like %cond = phi [true, %bb1], [true, %bb2], [%tmp, %bb3], then if bb1 is not "qualified", then it will not propagate to bb2. While it will do if bb1 and bb2 is swapped. May I ask is that expected?
And I have the same question for the newly added checker here: https://github.com/llvm/llvm-project/blob/main/llvm/lib/Transforms/Scalar/JumpThreading.cpp#L280.
Thank you very much!

[davidxl]

lib/Transforms/Scalar/JumpThreading.cpp

206

could be an oversight. If you have a test case to demonstrate the missed update, please send a patch.

lib/Transforms/Scalar/JumpThreading.cpp

 
 // Public interface to the Jump Threading pass
 FunctionPass *llvm::createJumpThreadingPass(int Threshold) { return new JumpThreading(Threshold); }
 
 JumpThreadingPass::JumpThreadingPass(int T) {
   BBDupThreshold = (T == -1) ? BBDuplicateThreshold : unsigned(T);
 }
 
-/// runOnFunction - Top level algorithm.
+// 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) {
+//     Block A;
+//  } else {
+//     Block B;
+//  }
+//
+//  [Block BB]
+//  cond = PN([true, %A], [..., %B]); // PHI node
+//  [Branch CondBr]
+//  if (cond) {
+//    ...  // P(cond == true) = 1%
+//  }
+//
+//  Here we know that when block A is taken, cond must be true, which means
+//      P(cond == true | A) = 1
+//
+//  Given that P(cond == true) = P(cond == true | A) * P(A) +
+//                               P(cond == true | B) * P(B)
+//  we get
+//     P(cond == true ) = P(A) + P(cond == true | B) * P(B)
+//
+//  which gives us:
+//     P(A) <= P(c == true), i.e.
+//     P(t == true) <= P(cond == true)
+//
+//  In other words, if we know P(cond == true), we know that P(t == true)
+//  can not be greater than 1%.
+static void updatePredecessorProfileMetadata(PHINode *PN, BasicBlock *BB) {
+  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;
+    for (;;) {
+      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))
+      continue;
+
+    BP = (CI->isOne() ? 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());
+
+    uint64_t PredTrueWeight, PredFalseWeight;
+    // FIXME: We currently only set the profile data when it is missing.
+    // With PGO, this can be used to refine even existing profile data with
+    // context information. This needs to be done after more performance
+    // testing.
+    if (PredBr->extractProfMetadata(PredTrueWeight, PredFalseWeight))
+      continue;
+
+    // We can not infer anything useful when BP >= 50%, because BP is the
+    // upper bound probability value.
+    if (BP >= BranchProbability(50, 100))
+      continue;
+
+    SmallVector<uint32_t, 2> Weights;
+    if (PredBr->getSuccessor(0) == PredOutEdge.second) {
+      Weights.push_back(BP.getNumerator());
+      Weights.push_back(BP.getCompl().getNumerator());
+    } else {
+      Weights.push_back(BP.getCompl().getNumerator());
+      Weights.push_back(BP.getNumerator());
+    }
+    PredBr->setMetadata(LLVMContext::MD_prof,
+                        MDBuilder(PredBr->getParent()->getContext())
+                            .createBranchWeights(Weights));
+  }
+}
+
+/// runOnFunction - Toplevel algorithm.
 ///
 bool JumpThreading::runOnFunction(Function &F) {
   if (skipFunction(F))
     return false;
   auto TLI = &getAnalysis<TargetLibraryInfoWrapperPass>().getTLI();
   auto LVI = &getAnalysis<LazyValueInfoWrapperPass>().getLVI();
   auto AA = &getAnalysis<AAResultsWrapperPass>().getAAResults();
   std::unique_ptr<BlockFrequencyInfo> BFI;

       SimplifyValue = CondCmp->getOperand(0);
 
   // TODO: There are other places where load PRE would be profitable, such as
   // more complex comparisons.
   if (LoadInst *LI = dyn_cast<LoadInst>(SimplifyValue))
     if (SimplifyPartiallyRedundantLoad(LI))
       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);
+
   // 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

test/Transforms/JumpThreading/threading_prof1.ll

+; RUN: opt -jump-threading -S < %s | FileCheck %s
+; RUN: opt -passes=jump-threading -S < %s | FileCheck %s
+
+define void @test() {
+; CHECK-LABEL: @test()
+bb:
+  %tmp = call i32 @a()
+  %tmp1 = icmp eq i32 %tmp, 1
+  br i1 %tmp1, label %bb5, label %bb2
+; CHECK: br i1 %tmp1,{{.*}} !prof ![[PROF1:[0-9]+]]
+
+bb2:                                              ; preds = %bb
+  %tmp3 = call i32 @b()
+  %tmp4 = icmp ne i32 %tmp3, 1
+  br label %bb5
+; CHECK: br i1 %tmp4, {{.*}} !prof ![[PROF2:[0-9]+]]
+
+bb5:                                              ; preds = %bb2, %bb
+  %tmp6 = phi i1 [ false, %bb ], [ %tmp4, %bb2 ]
+  br i1 %tmp6, label %bb8, label %bb7, !prof !0
+
+bb7:                                              ; preds = %bb5
+  call void @bar()
+  br label %bb8
+
+bb8:                                              ; preds = %bb7, %bb5
+  ret void
+}
+
+define void @test_single_pred1() {
+; CHECK-LABEL: @test_single_pred1()
+bb:
+  %tmp = call i32 @a()
+  %tmp1 = icmp eq i32 %tmp, 1
+  br i1 %tmp1, label %bb5_1, label %bb2
+; CHECK: br i1 %tmp1,{{.*}} !prof ![[PROF1:[0-9]+]]
+
+bb5_1:                                             
+  br label %bb5;
+
+bb2:                                              
+  %tmp3 = call i32 @b()
+  %tmp4 = icmp ne i32 %tmp3, 1
+  br label %bb5
+; CHECK: br i1 %tmp4, {{.*}} !prof ![[PROF2:[0-9]+]]
+
+bb5:                                             
+  %tmp6 = phi i1 [ false, %bb5_1 ], [ %tmp4, %bb2 ]
+  br i1 %tmp6, label %bb8, label %bb7, !prof !0
+
+bb7:                                            
+  call void @bar()
+  br label %bb8
+
+bb8:                                           
+  ret void
+}
+
+define void @test_single_pred2() {
+; CHECK-LABEL: @test_single_pred2()
+bb:
+  %tmp = call i32 @a()
+  %tmp1 = icmp eq i32 %tmp, 1
+  br i1 %tmp1, label %bb5_1, label %bb2
+; CHECK: br i1 %tmp1,{{.*}} !prof ![[PROF1:[0-9]+]]
+
+bb5_1:                                             
+  br label %bb5_2;
+
+bb5_2:                                             
+  br label %bb5;
+
+bb2:                          
+  %tmp3 = call i32 @b()
+  %tmp4 = icmp ne i32 %tmp3, 1
+  br label %bb5
+; CHECK: br i1 %tmp4, {{.*}} !prof ![[PROF2:[0-9]+]]
+
+bb5:                         
+  %tmp6 = phi i1 [ false, %bb5_2 ], [ %tmp4, %bb2 ]
+  br i1 %tmp6, label %bb8, label %bb7, !prof !0
+
+bb7:                        
+  call void @bar()
+  br label %bb8
+
+bb8:                       
+  ret void
+}
+
+declare void @bar()
+
+declare i32 @a()
+
+declare i32 @b()
+
+!0 = !{!"branch_weights", i32 2146410443, i32 1073205}
+;CHECK: ![[PROF1]] = !{!"branch_weights", i32 1073205, i32 2146410443}
+;CHECK: ![[PROF2]] = !{!"branch_weights", i32 2146410443, i32 1073205}

test/Transforms/JumpThreading/threading_prof2.ll

+; RUN: opt -jump-threading -S < %s | FileCheck %s
+; RUN: opt -passes=jump-threading -S < %s | FileCheck %s
+define void @test() {
+bb:
+  %tmp = call i32 @a()
+  %tmp1 = icmp eq i32 %tmp, 1
+  br i1 %tmp1, label %bb5, label %bb2
+; CHECK: br i1 %tmp1,{{.*}} !prof ![[PROF1:[0-9]+]]
+
+bb2:                                              
+  %tmp3 = call i32 @b()
+  %tmp4 = icmp ne i32 %tmp3, 1
+  br label %bb5
+; CHECK: br i1 %tmp4, {{.*}} !prof ![[PROF2:[0-9]+]]
+
+bb5:                                             
+  %tmp6 = phi i1 [ false, %bb ], [ %tmp4, %bb2 ]
+  br i1 %tmp6, label %bb8, label %bb7, !prof !0
+
+bb7:                                            
+  call void @bar()
+  br label %bb9
+
+bb8: 
+  call void @foo()
+  br label %bb9
+
+bb9:                                           
+  ret void
+}
+
+declare void @bar()
+
+declare void @foo()
+
+declare i32 @a()
+
+declare i32 @b()
+
+!0 = !{!"branch_weights", i32 2146410443, i32 1073205}
+;CHECK: ![[PROF1]] = !{!"branch_weights", i32 1073205, i32 2146410443}
+;CHECK: ![[PROF2]] = !{!"branch_weights", i32 2146410443, i32 1073205}