[JumpThreading] Copy profile metadata on select unfolding

[rp]

Summary

When select is unfolded, copy profile metadata attached
to the select to a newly created branch

[rp]

comment text

[rp]

This comment was removed by rp.

[rp]

This patch addresses losing of the profile attribute during jump-threading, later causing not ideal block merge due to missing profile in shouldFoldCondBranchesToCommonDestination() where blocks are merged in case branch probability is unknown (and it is unknown as it was not propagated from select to br).

[ebrevnov]

We should update BPI\BFI in a similar way as it is done in updateBlockFreqAndEdgeWeight.

[rp]

Added setEdgeProbability/setBlockFreq

[rp]

clang-format

[ebrevnov]

Looks OK, except mentioned nit.

llvm/lib/Transforms/Scalar/JumpThreading.cpp

2768

You should check sum of TrueWeight and FalseWeight is not 0

[rp]

check TrueWeight and FalseWeight for division by 0

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

 ; 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 !0
 ; 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 !0
 ; 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 2146410443, i32 1073205}
+!1 = !{!"function_entry_count", i64 1984}