[SimplifyCFG] Don't speculatively execute BB[s] if they are predictably not taken

[lebedev.ri]

Summary

Same as D106650, but for FoldTwoEntryPHINode()

[spatel]

LGTM - as with the similar D106650, I think we could use 1-2 more tests with a metadata/branch prob just below getPredictableBranchThreshold(), so we verify that calculation works as expected.

[lebedev.ri]

In D106717#2903320, @spatel wrote:

LGTM - as with the similar D106650, I think we could use 1-2 more tests with a metadata/branch prob just below getPredictableBranchThreshold(), so we verify that calculation works as expected.

Will do, thank you for the review!

llvm/lib/Transforms/Utils/SimplifyCFG.cpp

if (isa<ConstantInt>(IfCond))
     return false;
 
   BasicBlock *DomBlock = DomBI->getParent();
   SmallVector<BasicBlock *, 2> IfBlocks;
   llvm::copy_if(
       PN->blocks(), std::back_inserter(IfBlocks), [](BasicBlock *IfBlock) {
         return cast<BranchInst>(IfBlock->getTerminator())->isUnconditional();
       });
-  assert(!IfBlocks.empty() && "Will have at least one block to speculate.");
+  assert((IfBlocks.size() == 1 || IfBlocks.size() == 2) &&
+         "Will have either one or two blocks to speculate.");
+
+  // If the branch is non-unpredictable, see if we either predictably jump to
+  // the merge bb (if we have only a single 'then' block), or if we predictably
+  // jump to one specific 'then' block (if we have two of them).
+  // It isn't beneficial to speculatively execute the code
+  // from the block that we know is predictably not entered.
+  if (!DomBI->getMetadata(LLVMContext::MD_unpredictable)) {
+    uint64_t TWeight, FWeight;
+    if (DomBI->extractProfMetadata(TWeight, FWeight) &&
+        (TWeight + FWeight) != 0) {
+      BranchProbability BITrueProb =
+          BranchProbability::getBranchProbability(TWeight, TWeight + FWeight);
+      BranchProbability Likely = TTI.getPredictableBranchThreshold();
+      BranchProbability BIFalseProb = BITrueProb.getCompl();
+      if (IfBlocks.size() == 1) {
+        BranchProbability BIBBProb =
+            DomBI->getSuccessor(0) == BB ? BITrueProb : BIFalseProb;
+        if (BIBBProb >= Likely)
+          return false;
+      } else {
+        if (BITrueProb >= Likely || BIFalseProb >= Likely)
+          return false;
+      }
+    }
+  }
 
   // Don't try to fold an unreachable block. For example, the phi node itself
   // can't be the candidate if-condition for a select that we want to form.
   if (auto *IfCondPhiInst = dyn_cast<PHINode>(IfCond))
     if (IfCondPhiInst->getParent() == BB)
       return false;
 
   // Okay, we found that we can merge this two-entry phi node into a select.

llvm/test/Transforms/PGOProfile/chr.ll

 define i64 @test_chr_22(i1 %i, i64* %j, i64 %v0) !prof !14 {
 ; CHECK-LABEL: @test_chr_22(
 ; CHECK-NEXT:  bb0:
 ; CHECK-NEXT:    [[REASS_ADD:%.*]] = shl i64 [[V0:%.*]], 1
 ; CHECK-NEXT:    [[V2:%.*]] = add i64 [[REASS_ADD]], 3
 ; CHECK-NEXT:    [[C1:%.*]] = icmp slt i64 [[V2]], 100
 ; CHECK-NEXT:    br i1 [[C1]], label [[BB0_SPLIT:%.*]], label [[BB0_SPLIT_NONCHR:%.*]], !prof [[PROF15]]
 ; CHECK:       common.ret:
-; CHECK-NEXT:    ret i64 99
+; CHECK-NEXT:    [[COMMON_RET_OP:%.*]] = phi i64 [ 99, [[BB0_SPLIT]] ], [ 99, [[BB0_SPLIT_NONCHR]] ]
+; CHECK-NEXT:    ret i64 [[COMMON_RET_OP]]
 ; CHECK:       bb0.split:
 ; CHECK-NEXT:    [[V299:%.*]] = mul i64 [[V2]], 7860086430977039991
 ; CHECK-NEXT:    store i64 [[V299]], i64* [[J:%.*]], align 4
 ; CHECK-NEXT:    br label [[COMMON_RET:%.*]]
 ; CHECK:       bb0.split.nonchr:
 ; CHECK-NEXT:    [[V299_NONCHR:%.*]] = mul i64 [[V2]], 7860086430977039991
 ; CHECK-NEXT:    store i64 [[V299_NONCHR]], i64* [[J]], align 4
 ; CHECK-NEXT:    br label [[COMMON_RET]]

llvm/test/Transforms/SimplifyCFG/fold-two-entry-phi-node-with-one-block-profmd.ll

end:
   ret i32 %res
 }
 
 define i32 @predictably_nontaken(i32 %a, i32 %b, i32 %c, i32 %d) {
 ; CHECK-LABEL: @predictably_nontaken(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    call void @sideeffect0()
 ; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i32 [[A:%.*]], [[B:%.*]]
+; CHECK-NEXT:    br i1 [[CMP]], label [[END:%.*]], label [[COND_TRUE:%.*]], !prof [[PROF0]]
+; CHECK:       cond.true:
 ; CHECK-NEXT:    [[V0:%.*]] = add i32 [[C:%.*]], [[D:%.*]]
-; CHECK-NEXT:    [[RES:%.*]] = select i1 [[CMP]], i32 0, i32 [[V0]], !prof [[PROF0]]
+; CHECK-NEXT:    br label [[END]]
+; CHECK:       end:
+; CHECK-NEXT:    [[RES:%.*]] = phi i32 [ [[V0]], [[COND_TRUE]] ], [ 0, [[ENTRY:%.*]] ]
 ; CHECK-NEXT:    call void @sideeffect1()
 ; CHECK-NEXT:    ret i32 [[RES]]
 ;
 entry:
   call void @sideeffect0()
   %cmp = icmp eq i32 %a, %b
   br i1 %cmp, label %end, label %cond.true, !prof !0 ; likely branches to %end
 

llvm/test/Transforms/SimplifyCFG/fold-two-entry-phi-node-with-two-blocks-profmd.ll

end:
   ret i32 %res
 }
 
 define i32 @predictably_taken(i32 %a, i32 %b, i32 %c, i32 %d) {
 ; CHECK-LABEL: @predictably_taken(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    call void @sideeffect0()
 ; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i32 [[A:%.*]], [[B:%.*]]
+; CHECK-NEXT:    br i1 [[CMP]], label [[COND_TRUE:%.*]], label [[COND_FALSE:%.*]], !prof [[PROF0:![0-9]+]]
+; CHECK:       cond.true:
 ; CHECK-NEXT:    [[V0:%.*]] = add i32 [[C:%.*]], [[D:%.*]]
+; CHECK-NEXT:    br label [[END:%.*]]
+; CHECK:       cond.false:
 ; CHECK-NEXT:    [[V1:%.*]] = sub i32 [[C]], [[D]]
-; CHECK-NEXT:    [[RES:%.*]] = select i1 [[CMP]], i32 [[V0]], i32 [[V1]], !prof [[PROF0:![0-9]+]]
+; CHECK-NEXT:    br label [[END]]
+; CHECK:       end:
+; CHECK-NEXT:    [[RES:%.*]] = phi i32 [ [[V0]], [[COND_TRUE]] ], [ [[V1]], [[COND_FALSE]] ]
 ; CHECK-NEXT:    call void @sideeffect1()
 ; CHECK-NEXT:    ret i32 [[RES]]
 ;
 entry:
   call void @sideeffect0()
   %cmp = icmp eq i32 %a, %b
   br i1 %cmp, label %cond.true, label %cond.false, !prof !0 ; likely branches to %cond.true
 
end:
   ret i32 %res
 }
 
 define i32 @predictably_nontaken(i32 %a, i32 %b, i32 %c, i32 %d) {
 ; CHECK-LABEL: @predictably_nontaken(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    call void @sideeffect0()
 ; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i32 [[A:%.*]], [[B:%.*]]
+; CHECK-NEXT:    br i1 [[CMP]], label [[COND_FALSE:%.*]], label [[COND_TRUE:%.*]], !prof [[PROF0]]
+; CHECK:       cond.true:
 ; CHECK-NEXT:    [[V0:%.*]] = add i32 [[C:%.*]], [[D:%.*]]
+; CHECK-NEXT:    br label [[END:%.*]]
+; CHECK:       cond.false:
 ; CHECK-NEXT:    [[V1:%.*]] = sub i32 [[C]], [[D]]
-; CHECK-NEXT:    [[RES:%.*]] = select i1 [[CMP]], i32 [[V1]], i32 [[V0]], !prof [[PROF0]]
+; CHECK-NEXT:    br label [[END]]
+; CHECK:       end:
+; CHECK-NEXT:    [[RES:%.*]] = phi i32 [ [[V0]], [[COND_TRUE]] ], [ [[V1]], [[COND_FALSE]] ]
 ; CHECK-NEXT:    call void @sideeffect1()
 ; CHECK-NEXT:    ret i32 [[RES]]
 ;
 entry:
   call void @sideeffect0()
   %cmp = icmp eq i32 %a, %b
   br i1 %cmp, label %cond.false, label %cond.true, !prof !0 ; likely branches to %cond.false