[SimplifyCFG][profcheck] Synthesize profile for br (X == 0 | X == 1), T, F1 -> switch (llvm#161549)

We cannot calculate the weights of the switch precisely, but we do know the probability of the default branch. We then split equally the remaining probability over the rest of the cases. If we did nothing, the static estimation could be considerably poorer.



Issue llvm#147390

Author: mtrofin

llvm/lib/Transforms/Utils/SimplifyCFG.cpp

@@ -5152,14 +5152,18 @@ bool SimplifyCFGOpt::simplifyBranchOnICmpChain(BranchInst *BI,
   if (ExtraCase && Values.size() < 2)
     return false;
 
-  // TODO: Preserve branch weight metadata, similarly to how
-  // foldValueComparisonIntoPredecessors preserves it.
+  SmallVector<uint32_t> BranchWeights;
+  const bool HasProfile = !ProfcheckDisableMetadataFixes &&
+                          extractBranchWeights(*BI, BranchWeights);
 
   // Figure out which block is which destination.
   BasicBlock *DefaultBB = BI->getSuccessor(1);
   BasicBlock *EdgeBB = BI->getSuccessor(0);
-  if (!TrueWhenEqual)
+  if (!TrueWhenEqual) {
     std::swap(DefaultBB, EdgeBB);
+    if (HasProfile)
+      std::swap(BranchWeights[0], BranchWeights[1]);
+  }
 
   BasicBlock *BB = BI->getParent();
 
@@ -5190,10 +5194,11 @@ bool SimplifyCFGOpt::simplifyBranchOnICmpChain(BranchInst *BI,
     if (!isGuaranteedNotToBeUndefOrPoison(ExtraCase, AC, BI, nullptr))
       ExtraCase = Builder.CreateFreeze(ExtraCase);
 
-    if (TrueWhenEqual)
-      Builder.CreateCondBr(ExtraCase, EdgeBB, NewBB);
-    else
-      Builder.CreateCondBr(ExtraCase, NewBB, EdgeBB);
+    // We don't have any info about this condition.
+    auto *Br = TrueWhenEqual ? Builder.CreateCondBr(ExtraCase, EdgeBB, NewBB)
+                             : Builder.CreateCondBr(ExtraCase, NewBB, EdgeBB);
+    setExplicitlyUnknownBranchWeightsIfProfiled(*Br, *NewBB->getParent(),
+                                                DEBUG_TYPE);
 
     OldTI->eraseFromParent();
 
@@ -5220,6 +5225,17 @@ bool SimplifyCFGOpt::simplifyBranchOnICmpChain(BranchInst *BI,
 
   // Create the new switch instruction now.
   SwitchInst *New = Builder.CreateSwitch(CompVal, DefaultBB, Values.size());
+  if (HasProfile) {
+    // We know the weight of the default case. We don't know the weight of the
+    // other cases, but rather than completely lose profiling info, we split
+    // the remaining probability equally over them.
+    SmallVector<uint32_t> NewWeights(Values.size() + 1);
+    NewWeights[0] = BranchWeights[1]; // this is the default, and we swapped if
+                                      // TrueWhenEqual.
+    for (auto &V : drop_begin(NewWeights))
+      V = BranchWeights[0] / Values.size();
+    setBranchWeights(*New, NewWeights, /*IsExpected=*/false);
+  }
 
   // Add all of the 'cases' to the switch instruction.
   for (ConstantInt *Val : Values)

llvm/test/Transforms/SimplifyCFG/switch_create.ll

@@ -1,17 +1,17 @@
-; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --check-globals
 ; RUN: opt -S -passes=simplifycfg -simplifycfg-require-and-preserve-domtree=1 -switch-range-to-icmp < %s | FileCheck %s
 ; RUN: opt -S -data-layout="p:32:32-p1:16:16" -passes=simplifycfg -simplifycfg-require-and-preserve-domtree=1 -switch-range-to-icmp < %s | FileCheck -check-prefix=CHECK -check-prefix=DL %s
 
 declare void @foo1()
 
 declare void @foo2()
 
-define void @test1(i32 %V) {
+define void @test1(i32 %V) !prof !0 {
 ; CHECK-LABEL: @test1(
 ; CHECK-NEXT:    switch i32 [[V:%.*]], label [[F:%.*]] [
 ; CHECK-NEXT:      i32 17, label [[T:%.*]]
 ; CHECK-NEXT:      i32 4, label [[T]]
-; CHECK-NEXT:    ]
+; CHECK-NEXT:    ], !prof [[PROF1:![0-9]+]]
 ; CHECK:       common.ret:
 ; CHECK-NEXT:    ret void
 ; CHECK:       T:
@@ -24,7 +24,7 @@ define void @test1(i32 %V) {
   %C1 = icmp eq i32 %V, 4         ; <i1> [#uses=1]
   %C2 = icmp eq i32 %V, 17                ; <i1> [#uses=1]
   %CN = or i1 %C1, %C2            ; <i1> [#uses=1]
-  br i1 %CN, label %T, label %F
+  br i1 %CN, label %T, label %F, !prof !1
 T:              ; preds = %0
   call void @foo1( )
   ret void
@@ -116,12 +116,12 @@ F:              ; preds = %0
   ret void
 }
 
-define void @test2(i32 %V) {
+define void @test2(i32 %V) !prof !0 {
 ; CHECK-LABEL: @test2(
 ; CHECK-NEXT:    switch i32 [[V:%.*]], label [[T:%.*]] [
 ; CHECK-NEXT:      i32 17, label [[F:%.*]]
 ; CHECK-NEXT:      i32 4, label [[F]]
-; CHECK-NEXT:    ]
+; CHECK-NEXT:    ], !prof [[PROF2:![0-9]+]]
 ; CHECK:       common.ret:
 ; CHECK-NEXT:    ret void
 ; CHECK:       T:
@@ -134,7 +134,7 @@ define void @test2(i32 %V) {
   %C1 = icmp ne i32 %V, 4         ; <i1> [#uses=1]
   %C2 = icmp ne i32 %V, 17                ; <i1> [#uses=1]
   %CN = and i1 %C1, %C2           ; <i1> [#uses=1]
-  br i1 %CN, label %T, label %F
+  br i1 %CN, label %T, label %F, !prof !1
 T:              ; preds = %0
   call void @foo1( )
   ret void
@@ -1313,3 +1313,16 @@ if.then:
 if.end:
   ret void
 }
+
+!0 = !{!"function_entry_count", i32 100}
+!1 = !{!"branch_weights", i32 6, i32 10}
+;.
+; DL: attributes #[[ATTR0:[0-9]+]] = { noredzone nounwind ssp }
+; DL: attributes #[[ATTR1:[0-9]+]] = { nounwind }
+; DL: attributes #[[ATTR2]] = { noredzone nounwind }
+; DL: attributes #[[ATTR3]] = { noredzone }
+;.
+; DL: [[META0:![0-9]+]] = !{!"function_entry_count", i32 100}
+; DL: [[PROF1]] = !{!"branch_weights", i32 10, i32 3, i32 3}
+; DL: [[PROF2]] = !{!"branch_weights", i32 6, i32 5, i32 5}
+;.