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

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}
+;.