[SimplifyCFG] propagate branch metadata when creating select

[spatel]

Summary

Hopefully, the change itself is straightforward - plug another hole where we drop branch weight metadata while creating a select.

I'm curious about the lack of use of BranchProbability in SimplifyCFG though. Is there a reason we don't use that here (and therefore need duplicate functionality like "FitWeights()")?

[davidxl]

looks fine, but can you write another more direct test case for better checking if the profile update is correct or not.

[spatel]

Patch updated:
Minimize test case to show that existing metadata on a select isn't changed and the missing weights fixed by this patch are correct.

[davidxl]

test/Transforms/SimplifyCFG/preserve-branchweights.ll

422 ↗

(On Diff #55690)

If I read the code correctly, should the weights for select be 9:5 ? !11 does not look correct.

423 ↗

(On Diff #55690)

This one is correct.

[spatel]

test/Transforms/SimplifyCFG/preserve-branchweights.ll

422 ↗

(On Diff #55690)

I thought the weights for the select would be equal to the weights of the original branch because they have the same comparison condition (%cmpb), 3:5. How do you calculate 9:5?

[davidxl]

test/Transforms/SimplifyCFG/preserve-branchweights.ll

422 ↗

(On Diff #55690)

Should it be 3 : 5 * (1/3) which 9 : 5? When block 2 is entered, there is only 1/3 probability it reaches block3.

[spatel]

test/Transforms/SimplifyCFG/preserve-branchweights.ll

422 ↗

(On Diff #55690)

Aha - yes. Although it's the same comparison operand, we have changed the program structure - we are speculatively executing the select, so the original branch weights no longer apply.

This brings up my question from the original summary then - do you know if there's a reason not to use the BranchProbability class for calculating the new weights?

[davidxl]

test/Transforms/SimplifyCFG/preserve-branchweights.ll

422 ↗

(On Diff #55690)

BPI provides general APIs to access/read branch probability information that is either from profile data, or from static branch prediction.

Here the transformation just needs to update the profile data, using low level meta data interfaces is fine (Besides BPI does not have complete update APIs either).

[spatel]

Patch updated:
The branch weights of the first select are recalculated based on the new program structure.

[davidxl]

lgtm

lib/Transforms/Utils/SimplifyCFG.cpp

2882 ↗

(On Diff #56069)

The select is newly created instruction that does not exist before -- it sits at the same BB where the original phi, so it is not 'speculatively' executed.

The original PBI's branch weights do not apply because the select's 'logical' edges are incoming edges of the phi that is eliminated, not the out edges of PBI.

[spatel]

Thanks, David. I think that was the last of the branch --> select branch_weights drops in SimplifyCFG.
Still need to look at switch --> select and other transforms.

[vitalybuka]

Looks like a cause of
http://lab.llvm.org:8011/builders/sanitizer-x86_64-linux-bootstrap/builds/11620/steps/check-llvm%20msan/logs/stdio

[vitalybuka]

I am going to revert.

[vitalybuka]

Reverted in r268577

[spatel]

In D19674#421811, @vitalybuka wrote:

Reverted in r268577

Thanks, Vitaly. I'll try to get this debugged tomorrow.

[spatel]

In D19674#421885, @spatel wrote:

In D19674#421811, @vitalybuka wrote:

Reverted in r268577

Thanks, Vitaly. I'll try to get this debugged tomorrow.

I've tried to reproduce the error using sanitized builds of opt/llc as well as running under valgrind, but I don't see it. How do I find the exact steps used by the bot to induce the bug?

llvm/trunk/lib/Transforms/Utils/SimplifyCFG.cpp

static void GetBranchWeights(TerminatorInst *TI,
   if (BranchInst* BI = dyn_cast<BranchInst>(TI)) {
     assert(Weights.size() == 2);
     ICmpInst *ICI = cast<ICmpInst>(BI->getCondition());
     if (ICI->getPredicate() == ICmpInst::ICMP_EQ)
       std::swap(Weights.front(), Weights.back());
   }
 }
 
-/// Keep halving the weights until all can fit in uint32_t.
+/// Scale each weight so they all fit in uint32_t.
 static void FitWeights(MutableArrayRef<uint64_t> Weights) {
   uint64_t Max = *std::max_element(Weights.begin(), Weights.end());
   if (Max > UINT_MAX) {
     unsigned Offset = 32 - countLeadingZeros(Max);
     for (uint64_t &I : Weights)
       I >>= Offset;
   }
 }
static bool SimplifyCondBranchToCondBranch(BranchInst *PBI, BranchInst *BI,
   Value *Cond = Builder.CreateOr(PBICond, BICond, "brmerge");
 
   // Modify PBI to branch on the new condition to the new dests.
   PBI->setCondition(Cond);
   PBI->setSuccessor(0, CommonDest);
   PBI->setSuccessor(1, OtherDest);
 
   // Update branch weight for PBI.
+  MDBuilder MDB(BI->getContext());
   uint64_t PredTrueWeight, PredFalseWeight, SuccTrueWeight, SuccFalseWeight;
-  bool PredHasWeights =
-      PBI->extractProfMetadata(PredTrueWeight, PredFalseWeight);
-  bool SuccHasWeights =
-      BI->extractProfMetadata(SuccTrueWeight, SuccFalseWeight);
-  if (PredHasWeights && SuccHasWeights) {
-    uint64_t PredCommon = PBIOp ? PredFalseWeight : PredTrueWeight;
-    uint64_t PredOther = PBIOp ?PredTrueWeight : PredFalseWeight;
-    uint64_t SuccCommon = BIOp ? SuccFalseWeight : SuccTrueWeight;
-    uint64_t SuccOther = BIOp ? SuccTrueWeight : SuccFalseWeight;
+  uint64_t PredCommon, PredOther, SuccCommon, SuccOther;
+  bool HasWeights = PBI->extractProfMetadata(PredTrueWeight, PredFalseWeight);
+  if (HasWeights)
+    HasWeights = BI->extractProfMetadata(SuccTrueWeight, SuccFalseWeight);
+  if (HasWeights) {
+    PredCommon = PBIOp ? PredFalseWeight : PredTrueWeight;
+    PredOther = PBIOp ? PredTrueWeight : PredFalseWeight;
+    SuccCommon = BIOp ? SuccFalseWeight : SuccTrueWeight;
+    SuccOther = BIOp ? SuccTrueWeight : SuccFalseWeight;
     // The weight to CommonDest should be PredCommon * SuccTotal +
     //                                    PredOther * SuccCommon.
     // The weight to OtherDest should be PredOther * SuccOther.
     uint64_t NewWeights[2] = {PredCommon * (SuccCommon + SuccOther) +
                                   PredOther * SuccCommon,
                               PredOther * SuccOther};
-    // Halve the weights if any of them cannot fit in an uint32_t
     FitWeights(NewWeights);
 
     PBI->setMetadata(LLVMContext::MD_prof,
-                     MDBuilder(BI->getContext())
-                         .createBranchWeights(NewWeights[0], NewWeights[1]));
+                     MDB.createBranchWeights(NewWeights[0], NewWeights[1]));
   }
 
   // OtherDest may have phi nodes.  If so, add an entry from PBI's
   // block that are identical to the entries for BI's block.
   AddPredecessorToBlock(OtherDest, PBI->getParent(), BB);
 
   // We know that the CommonDest already had an edge from PBI to
   // it.  If it has PHIs though, the PHIs may have different
   // entries for BB and PBI's BB.  If so, insert a select to make
   // them agree.
   PHINode *PN;
   for (BasicBlock::iterator II = CommonDest->begin();
        (PN = dyn_cast<PHINode>(II)); ++II) {
     Value *BIV = PN->getIncomingValueForBlock(BB);
     unsigned PBBIdx = PN->getBasicBlockIndex(PBI->getParent());
     Value *PBIV = PN->getIncomingValue(PBBIdx);
     if (BIV != PBIV) {
       // Insert a select in PBI to pick the right value.
-      Value *NV = cast<SelectInst>
+      SelectInst *NV = cast<SelectInst>
         (Builder.CreateSelect(PBICond, PBIV, BIV, PBIV->getName() + ".mux"));
       PN->setIncomingValue(PBBIdx, NV);
+      // Although the select has the same condition as PBI, the original branch
+      // weights for PBI do not apply to the new select because the select's
+      // 'logical' edges are incoming edges of the phi that is eliminated, not
+      // the outgoing edges of PBI.
+      if (HasWeights) {
+        // The weight to PredCommonDest should be PredCommon * SuccTotal.
+        // The weight to PredOtherDest should be PredOther * SuccCommon.
+        uint64_t NewWeights[2] = {PredCommon * (SuccCommon + SuccOther),
+                                  PredOther * SuccCommon};
+
+        FitWeights(NewWeights);
+
+        NV->setMetadata(LLVMContext::MD_prof,
+                        MDB.createBranchWeights(NewWeights[0], NewWeights[1]));
+      }
     }
   }
 
   DEBUG(dbgs() << "INTO: " << *PBI->getParent());
   DEBUG(dbgs() << *PBI->getParent()->getParent());
 
   // This basic block is probably dead.  We know it has at least
   // one fewer predecessor.

llvm/trunk/test/Transforms/SimplifyCFG/preserve-branchweights.ll

 else:
   br label %return
 
 return:
   %retval.0 = phi i32 [ 1, %if ], [ 234, %else ]
   ret i32 %retval.0
 }
 
-; The 1st select should have branch weights equal to the 1st branch.
-; The 2nd select should have freshly calculated branch weights.
+; The selects should have freshly calculated branch weights.
 
 define i32 @SimplifyCondBranchToCondBranch(i1 %cmpa, i1 %cmpb) {
 ; CHECK-LABEL: @SimplifyCondBranchToCondBranch(
 ; CHECK-NEXT:  block1:
-; CHECK-NEXT:    [[BRMERGE:%.*]] = or i1 %cmpb, %cmpa
-; CHECK-NEXT:    [[DOTMUX:%.*]] = select i1 %cmpb, i32 0, i32 2
-; CHECK-NEXT:    [[OUTVAL:%.*]] = select i1 [[BRMERGE]], i32 [[DOTMUX]], i32 1, !prof !12
+; CHECK-NEXT:    [[BRMERGE:%.*]] = or i1 %cmpa, %cmpb
+; CHECK-NEXT:    [[DOTMUX:%.*]] = select i1 %cmpa, i32 0, i32 2, !prof !12
+; CHECK-NEXT:    [[OUTVAL:%.*]] = select i1 [[BRMERGE]], i32 [[DOTMUX]], i32 1, !prof !13
 ; CHECK-NEXT:    ret i32 [[OUTVAL]]
 ;
 block1:
-  br i1 %cmpb, label %block3, label %block2, !prof !0
+  br i1 %cmpa, label %block3, label %block2, !prof !13
 
 block2:
-  br i1 %cmpa, label %block3, label %exit, !prof !2
+  br i1 %cmpb, label %block3, label %exit, !prof !14
+
+block3:
+  %cowval = phi i32 [ 2, %block2 ], [ 0, %block1 ]
+  br label %exit
+
+exit:
+  %outval = phi i32 [ %cowval, %block3 ], [ 1, %block2 ]
+  ret i32 %outval
+}
+
+; Swap the operands of the compares to verify that the weights update correctly.
+
+define i32 @SimplifyCondBranchToCondBranchSwap(i1 %cmpa, i1 %cmpb) {
+; CHECK-LABEL: @SimplifyCondBranchToCondBranchSwap(
+; CHECK-NEXT:  block1:
+; CHECK-NEXT:    [[CMPA_NOT:%.*]] = xor i1 %cmpa, true
+; CHECK-NEXT:    [[CMPB_NOT:%.*]] = xor i1 %cmpb, true
+; CHECK-NEXT:    [[BRMERGE:%.*]] = or i1 [[CMPA_NOT]], [[CMPB_NOT]]
+; CHECK-NEXT:    [[DOTMUX:%.*]] = select i1 [[CMPA_NOT]], i32 0, i32 2, !prof !14
+; CHECK-NEXT:    [[OUTVAL:%.*]] = select i1 [[BRMERGE]], i32 [[DOTMUX]], i32 1, !prof !15
+; CHECK-NEXT:    ret i32 [[OUTVAL]]
+;
+block1:
+  br i1 %cmpa, label %block2, label %block3, !prof !13
+
+block2:
+  br i1 %cmpb, label %exit, label %block3, !prof !14
 
 block3:
   %cowval = phi i32 [ 2, %block2 ], [ 0, %block1 ]
   br label %exit
 
 exit:
   %outval = phi i32 [ %cowval, %block3 ], [ 1, %block2 ]
   ret i32 %outval
 }
 
 
 !0 = !{!"branch_weights", i32 3, i32 5}
 !1 = !{!"branch_weights", i32 1, i32 1}
 !2 = !{!"branch_weights", i32 1, i32 2}
 !3 = !{!"branch_weights", i32 4, i32 3, i32 2, i32 1}
 !4 = !{!"branch_weights", i32 4, i32 3, i32 2, i32 1}
 !5 = !{!"branch_weights", i32 7, i32 6, i32 5}
 !6 = !{!"branch_weights", i32 1, i32 3}
 !7 = !{!"branch_weights", i32 33, i32 9, i32 8, i32 7}
 !8 = !{!"branch_weights", i32 33, i32 9, i32 8}
 !9 = !{!"branch_weights", i32 7, i32 6}
 !10 = !{!"branch_weights", i32 672646, i32 21604207}
 !11 = !{!"branch_weights", i32 6960, i32 21597248}
 !12 = !{!"these_are_not_the_branch_weights_you_are_looking_for", i32 3, i32 5}
+!13 = !{!"branch_weights", i32 2, i32 3}
+!14 = !{!"branch_weights", i32 4, i32 7}
 
 ; CHECK: !0 = !{!"branch_weights", i32 5, i32 11}
 ; CHECK: !1 = !{!"branch_weights", i32 1, i32 5}
 ; CHECK: !2 = !{!"branch_weights", i32 7, i32 1, i32 2}
 ; CHECK: !3 = !{!"branch_weights", i32 49, i32 12, i32 24, i32 35}
 ; CHECK: !4 = !{!"branch_weights", i32 11, i32 5}
 ; CHECK: !5 = !{!"branch_weights", i32 17, i32 15}
 ; CHECK: !6 = !{!"branch_weights", i32 9, i32 7}
 ; CHECK: !7 = !{!"branch_weights", i32 17, i32 9, i32 8, i32 7, i32 17}
 ; CHECK: !8 = !{!"branch_weights", i32 24, i32 33}
 ; CHECK: !9 = !{!"branch_weights", i32 8, i32 33}
 ;; The false weight prints out as a negative integer here, but inside llvm, we
 ;; treat the weight as an unsigned integer.
 ; CHECK: !10 = !{!"branch_weights", i32 112017436, i32 -735157296}
 ; CHECK: !11 = !{!"branch_weights", i32 3, i32 5}
-; CHECK: !12 = !{!"branch_weights", i32 14, i32 10}
+; CHECK: !12 = !{!"branch_weights", i32 22, i32 12}
+; CHECK: !13 = !{!"branch_weights", i32 34, i32 21}
+; CHECK: !14 = !{!"branch_weights", i32 33, i32 14}
+; CHECK: !15 = !{!"branch_weights", i32 47, i32 8}