[JumpThreading] Preserve profile metadata during select unfolding

[rp]

Summary

Jump threading can replace select and unconditional branch with
conditional branch, but when doing so loses profile information.

This destructive transform can eventually lead to a performance
degradation due to folding of branches in
shouldFoldCondBranchesToCommonDestination as branch probabilities
are no longer known.

[rp]

rebase

[rp]

Ping

[mkazantsev]

llvm/test/Transforms/JumpThreading/select-unfold-prof.ll

2 ↗

(On Diff #476840)

1. Please use -passes=jump-threading, old PM is being deprecated
2. Please precommit the test and rebase on top of it

[rp]

Removed the test that was copied from select.ll

The checks were added into existing tests in select.ll instead
as other transformations there also throw away the prof attribute
and will need fixes.

[mkazantsev]

llvm/lib/Transforms/Scalar/JumpThreading.cpp

2762

Should this also be under HasProfileData?

[rp]

llvm/lib/Transforms/Scalar/JumpThreading.cpp

2762

The code is not very consistent, for example updatePredecessorProfileMetadata updates MD_prof without checking for HasProfile.

As IR with MD_prof attached to instruction yet without MD_prof attached to function will pass validation, I've concluded that it would not be polite to just ignore the attribute.

[mkazantsev]

Ok, looks pretty straightforward. LG.

[rp]

Rebased

[rp]

In D138132#3981252, @mkazantsev wrote:

Ok, looks pretty straightforward. LG.

Thanks for review. I've rebased the changes, but I need someone to land them for me as this is my first time. Could you perhaps be the one to merge this?

[mkazantsev]

Yup, I'll do it.

[mkazantsev]

This comment was removed by mkazantsev.

[chapuni]

llvm/test/Transforms/JumpThreading/select.ll

2–9

Does it require asserts?

[lebedev.ri]

llvm/test/Transforms/JumpThreading/select.ll

4–9

You should just check the metadata (--check-globals)

[cmtice]

This commit breaks the test llvm/test/Transforms/JumpThreading/select.ll if you build the compiler optimized I.e. with '-DCMAKE_BUILD_TYPE=Release'. I tested this by building both the version just before this commit and just after this commit with -DCMAKE_BUILD_TYPE=Release, then running the test:

$ ./bin/llvm-lit -s -v ../llvm/test/Transforms/JumpThreading/select.ll

The version built just before this commit passed.:

$ ./bin/llvm-lit -s -v ../llvm/test/Transforms/JumpThreading/select.ll

Testing Time: 0.10s

Passed: 1

The one built just after this commit failed:

$ ./bin/llvm-lit -s -v ../llvm/test/Transforms/JumpThreading/select.ll
FAIL: LLVM :: Transforms/JumpThreading/select.ll (1 of 1)

• TEST 'LLVM :: Transforms/JumpThreading/select.ll' FAILED ****

Script:

: 'RUN: at line 2'; /usr/local/google3/cmtice/new-llvm-project/build-opt/bin/opt -S -passes=jump-threading -debug-only=branch-prob < /usr/local/google3/cmtice/new-llvm-project/llvm/test/Transforms/JumpThreading/select.ll 2>&1 | /usr/local/google3/cmtice/new-llvm-project/build-opt/bin/FileCheck /usr/local/google3/cmtice/new-llvm-project/llvm/test/Transforms/JumpThreading/select.ll

Exit Code: 1

Command Output (stderr):

/usr/local/google3/cmtice/new-llvm-project/llvm/test/Transforms/JumpThreading/select.ll:4:16: error: CHECK-LABEL: expected string not found in input
; CHECK-LABEL: ---- Branch Probability Info : unfold1 ----

^

<stdin>:1:1: note: scanning from here
opt: Unknown command line argument '-debug-only=branch-prob'. Try: '/usr/local/google3/cmtice/new-llvm-project/build-opt/bin/opt --help'
^
<stdin>:1:43: note: possible intended match here
opt: Unknown command line argument '-debug-only=branch-prob'. Try: '/usr/local/google3/cmtice/new-llvm-project/build-opt/bin/opt --help'

^

Input file: <stdin>
Check file: /usr/local/google3/cmtice/new-llvm-project/llvm/test/Transforms/JumpThreading/select.ll

-dump-input=help explains the following input dump.

Input was:
<<<<<<

1: opt: Unknown command line argument '-debug-only=branch-prob'. Try: '/usr/local/google3/cmtice/new-llvm-project/build-opt/bin/opt --help'

label:4'0 X~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ error: no match found
label:4'1 ? possible intended match

2: opt: Did you mean '--debug-pass=branch-prob'?

label:4'0 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

---

---

Failed Tests (1):

LLVM :: Transforms/JumpThreading/select.ll

Testing Time: 0.52s

Failed: 1

$

[cmtice]

Never mind; I just noticed that https://github.com/llvm/llvm-project/commit/230dba0ad30e826407e783d19234c8b9d3373e2a fixes this issue.

[gribozavr2]

Unfortunately, this change causes assertion failures and I'm going to revert it.

Addition in this line can overflow: BP.emplace_back(BranchProbability(TrueWeight, TrueWeight + FalseWeight));. As a result, the constructor of BranchProbability() can trigger an assertion.

Here's an example in a debugger:

frame #9: 0x000055556f3d47ea clang`llvm::JumpThreadingPass::unfoldSelectInstr(this=0x0000068f36045618, Pred=0x0000068f360e55a0, BB=0x0000068f3613d1a0, SI=0x0000068f36b8a6a0, SIUse=0x0000068f397a53d0, Idx=1) at JumpThreading.cpp:2797:23
   2794     if (extractBranchWeights(*SI, TrueWeight, FalseWeight) &&
   2795         (TrueWeight + FalseWeight) != 0) {
   2796       SmallVector<BranchProbability, 2> BP;
-> 2797       BP.emplace_back(BranchProbability(TrueWeight, TrueWeight + FalseWeight));
   2798       BP.emplace_back(BranchProbability(FalseWeight, TrueWeight + FalseWeight));
   2799       BPI->setEdgeProbability(Pred, BP);
   2800     }
(lldb) p TrueWeight
(uint64_t) $0 = 2147483648
(lldb) p FalseWeight
(uint64_t) $1 = 2147483648
(lldb) down
frame #8: 0x0000555571acc3a6 clang`llvm::BranchProbability::BranchProbability(this=0x00007fffffff0d80, Numerator=2147483648, Denominator=0) at BranchProbability.cpp:41:3
   38   #endif
   39
   40   BranchProbability::BranchProbability(uint32_t Numerator, uint32_t Denominator) {
-> 41     assert(Denominator > 0 && "Denominator cannot be 0!");
   42     assert(Numerator <= Denominator && "Probability cannot be bigger than 1!");
   43     if (Denominator == D)
   44       N = Numerator;
(lldb) p Denominator
(uint32_t) $2 = 0

Unfortunately, I don't have a reduced IR example, but hopefully this should be enough to debug and fix.

[gribozavr2]

Actually I do have a reduced example. Apply this change to the test that you added in this patch to reproduce the assertion failure:

diff --git a/llvm/test/Transforms/JumpThreading/select.ll b/llvm/test/Transforms/JumpThreading/select.ll
index 4c7f3ea7d277..cdaae23694c4 100644
--- a/llvm/test/Transforms/JumpThreading/select.ll
+++ b/llvm/test/Transforms/JumpThreading/select.ll
@@ -660,5 +660,5 @@ if.end:
   ret i32 %v1
 }

-!0 = !{!"branch_weights", i32 1, i32 9}
+!0 = !{!"branch_weights", i32 2147483648, i32 2147483648}
 !1 = !{!"function_entry_count", i64 1984}

[rp]

1. using getBranchProbability instead of BranchProbability ctor to avoid overflow
2. branch_weights changed to i64, values are UINT_MAX/4 and 3*UINT_MAX/4
3. added REQUIRES: asserts
4. regenerated checks with --check-globals

[rp]

In D138132#4039260, @gribozavr2 wrote:

Actually I do have a reduced example. Apply this change to the test that you added in this patch to reproduce the assertion failure

Thank you very much, I've updated branch weights to be overflowing and 25/75% instead of 50/50%. If I may ask, how did you get the overflows, is there some automation that triggers it?

[rp]

llvm/test/Transforms/JumpThreading/select.ll

4–9

Is it possible to check edge probabilities with metadata? As when I hard-code this code into probability calculation:

BP.emplace_back(BranchProbability::getBranchProbability(1, 10));
      BP.emplace_back(BranchProbability::getBranchProbability(9, 10));

The probabilities in those checks will be reported as 10/90%, yet the prof metadata will stay as it was - 25/75%.

[gribozavr2]

In D138132#4039546, @rp wrote:

If I may ask, how did you get the overflows, is there some automation that triggers it?

I got this crash while running Clang on a real-world project internally at Google.

[rp]

run clang-format

[mkazantsev]

LG since the overflow is now fixed.

[mkazantsev]

llvm/test/Transforms/JumpThreading/select.ll

21

Please regenerate test.

llvm/lib/Transforms/Scalar/JumpThreading.cpp

void JumpThreadingPass::unfoldSelectInstr(BasicBlock *Pred, BasicBlock *BB,
   BasicBlock *NewBB = BasicBlock::Create(BB->getContext(), "select.unfold",
                                          BB->getParent(), BB);
   // Move the unconditional branch to NewBB.
   PredTerm->removeFromParent();
   NewBB->getInstList().insert(NewBB->end(), PredTerm);
   // Create a conditional branch and update PHI nodes.
   auto *BI = BranchInst::Create(NewBB, BB, SI->getCondition(), Pred);
   BI->applyMergedLocation(PredTerm->getDebugLoc(), SI->getDebugLoc());
+  BI->copyMetadata(*SI, {LLVMContext::MD_prof});
   SIUse->setIncomingValue(Idx, SI->getFalseValue());
   SIUse->addIncoming(SI->getTrueValue(), NewBB);
+  // Set the block frequency of NewBB.
+  if (HasProfileData) {
+    uint64_t TrueWeight, FalseWeight;
+    if (extractBranchWeights(*SI, TrueWeight, FalseWeight) &&
+        (TrueWeight + FalseWeight) != 0) {
+      SmallVector<BranchProbability, 2> BP;
+      BP.emplace_back(BranchProbability(TrueWeight, TrueWeight + FalseWeight));
+      BP.emplace_back(BranchProbability(FalseWeight, TrueWeight + FalseWeight));
+      BPI->setEdgeProbability(Pred, BP);
+    }
+
+    auto NewBBFreq =
+        BFI->getBlockFreq(Pred) * BPI->getEdgeProbability(Pred, NewBB);
+    BFI->setBlockFreq(NewBB, NewBBFreq.getFrequency());
+  }
 
   // The select is now dead.
   SI->eraseFromParent();
   DTU->applyUpdatesPermissive({{DominatorTree::Insert, NewBB, BB},
                                {DominatorTree::Insert, Pred, NewBB}});
 
   // Update any other PHI nodes in BB.
   for (BasicBlock::iterator BI = BB->begin();

llvm/test/Transforms/JumpThreading/select.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
-; RUN: opt -S -passes=jump-threading < %s | FileCheck %s
+; RUN: opt -S -passes=jump-threading -debug-only=branch-prob < %s 2>&1 | FileCheck %s
+
+; CHECK-LABEL:  ---- Branch Probability Info : unfold1 ----
+; CHECK:       set edge cond.false -> 0 successor probability to 0x0ccccccd / 0x80000000 = 10.00%
+; CHECK:       set edge cond.false -> 1 successor probability to 0x73333333 / 0x80000000 = 90.00%
+; CHECK-LABEL:  ---- Branch Probability Info : unfold2 ----
+; CHECK:       set edge cond.false -> 0 successor probability to 0x0ccccccd / 0x80000000 = 10.00%
+; CHECK:       set edge cond.false -> 1 successor probability to 0x73333333 / 0x80000000 = 90.00%
 
 declare void @foo()
 declare void @bar()
 declare void @baz()
 declare void @quux()
 
 
 ; Jump threading of branch with select as condition.
 ; Mostly theoretical since instruction combining simplifies all selects of
 ; booleans where at least one operand is true/false/undef.
 
 define void @test_br(i1 %cond, i1 %value) nounwind {
 ; CHECK-LABEL: @test_br(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    br i1 [[COND:%.*]], label [[L1:%.*]], label [[L0:%.*]]
 ; CHECK:       L0:
 ; CHECK-NEXT:    call void @baz()
 ; CHECK-NEXT:    [[EXPR:%.*]] = select i1 [[COND]], i1 true, i1 [[VALUE:%.*]]
 ; CHECK-NEXT:    br i1 [[EXPR]], label [[L1]], label [[L2:%.*]]
 ; CHECK:       L1:

 L3:
   store i32 10000, ptr %status, align 4
   br label %L4
 
 L4:
   ret void
 }
 
-define void @unfold1(double %x, double %y) nounwind {
+define void @unfold1(double %x, double %y) nounwind !prof !1 {
 ; CHECK-LABEL: @unfold1(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[SUB:%.*]] = fsub double [[X:%.*]], [[Y:%.*]]
 ; CHECK-NEXT:    [[CMP:%.*]] = fcmp ogt double [[SUB]], 1.000000e+01
 ; CHECK-NEXT:    br i1 [[CMP]], label [[COND_END4:%.*]], label [[COND_FALSE:%.*]]
 ; CHECK:       cond.false:
 ; CHECK-NEXT:    [[ADD:%.*]] = fadd double [[X]], [[Y]]
 ; CHECK-NEXT:    [[CMP1:%.*]] = fcmp ogt double [[ADD]], 1.000000e+01
-; CHECK-NEXT:    br i1 [[CMP1]], label [[COND_END4]], label [[IF_THEN:%.*]]
+; CHECK-NEXT:    br i1 [[CMP1]], label [[COND_END4]], label [[IF_THEN:%.*]], !prof [[PROF1:![0-9]+]]
 ; CHECK:       cond.end4:
 ; CHECK-NEXT:    [[COND5:%.*]] = phi double [ [[SUB]], [[ENTRY:%.*]] ], [ [[ADD]], [[COND_FALSE]] ]
 ; CHECK-NEXT:    [[CMP6:%.*]] = fcmp oeq double [[COND5]], 0.000000e+00
 ; CHECK-NEXT:    br i1 [[CMP6]], label [[IF_THEN]], label [[IF_END:%.*]]
 ; CHECK:       if.then:
 ; CHECK-NEXT:    call void @foo()
 ; CHECK-NEXT:    br label [[IF_END]]
 ; CHECK:       if.end:
 ; CHECK-NEXT:    ret void
 ;
 entry:
   %sub = fsub double %x, %y
   %cmp = fcmp ogt double %sub, 1.000000e+01
   br i1 %cmp, label %cond.end4, label %cond.false
 
 cond.false:                                       ; preds = %entry
   %add = fadd double %x, %y
   %cmp1 = fcmp ogt double %add, 1.000000e+01
-  %add. = select i1 %cmp1, double %add, double 0.000000e+00
+  %add. = select i1 %cmp1, double %add, double 0.000000e+00, !prof !0
   br label %cond.end4
 
 cond.end4:                                        ; preds = %entry, %cond.false
   %cond5 = phi double [ %add., %cond.false ], [ %sub, %entry ]
   %cmp6 = fcmp oeq double %cond5, 0.000000e+00
   br i1 %cmp6, label %if.then, label %if.end
 
 if.then:                                          ; preds = %cond.end4
   call void @foo()
   br label %if.end
 
 if.end:                                           ; preds = %if.then, %cond.end4
   ret void
 
 }
 
 
-define void @unfold2(i32 %x, i32 %y) nounwind {
+define void @unfold2(i32 %x, i32 %y) nounwind !prof !1 {
 ; CHECK-LABEL: @unfold2(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[SUB:%.*]] = sub nsw i32 [[X:%.*]], [[Y:%.*]]
 ; CHECK-NEXT:    [[CMP:%.*]] = icmp sgt i32 [[SUB]], 10
 ; CHECK-NEXT:    br i1 [[CMP]], label [[IF_END:%.*]], label [[COND_FALSE:%.*]]
 ; CHECK:       cond.false:
 ; CHECK-NEXT:    [[ADD:%.*]] = add nsw i32 [[X]], [[Y]]
 ; CHECK-NEXT:    [[CMP1:%.*]] = icmp sgt i32 [[ADD]], 10
-; CHECK-NEXT:    br i1 [[CMP1]], label [[IF_THEN:%.*]], label [[COND_END4:%.*]]
+; CHECK-NEXT:    br i1 [[CMP1]], label [[IF_THEN:%.*]], label [[COND_END4:%.*]], !prof [[PROF1:![0-9]+]]
 ; CHECK:       cond.end4:
 ; CHECK-NEXT:    [[COND5:%.*]] = phi i32 [ [[ADD]], [[COND_FALSE]] ]
 ; CHECK-NEXT:    [[CMP6:%.*]] = icmp eq i32 [[COND5]], 0
 ; CHECK-NEXT:    br i1 [[CMP6]], label [[IF_THEN]], label [[IF_END]]
 ; CHECK:       if.then:
 ; CHECK-NEXT:    call void @foo()
 ; CHECK-NEXT:    br label [[IF_END]]
 ; CHECK:       if.end:
 ; CHECK-NEXT:    ret void
 ;
 entry:
   %sub = sub nsw i32 %x, %y
   %cmp = icmp sgt i32 %sub, 10
   br i1 %cmp, label %cond.end4, label %cond.false
 
 cond.false:                                       ; preds = %entry
   %add = add nsw i32 %x, %y
   %cmp1 = icmp sgt i32 %add, 10
-  %add. = select i1 %cmp1, i32 0, i32 %add
+  %add. = select i1 %cmp1, i32 0, i32 %add, !prof !0
   br label %cond.end4
 
 cond.end4:                                        ; preds = %entry, %cond.false
   %cond5 = phi i32 [ %add., %cond.false ], [ %sub, %entry ]
   %cmp6 = icmp eq i32 %cond5, 0
   br i1 %cmp6, label %if.then, label %if.end
 
 if.then:                                          ; preds = %cond.end4

 if.else:
   br label %if.end
 
 if.end:
   %v = phi i1 [ %x, %if.else ], [ false, %entry ]
   %v1 = select i1 %v, i32 %s, i32 42
   ret i32 %v1
 }
+
+!0 = !{!"branch_weights", i32 1, i32 9}
+!1 = !{!"function_entry_count", i64 1984}