[SimplifyCFG][profcheck] Handle branch weights in simplifySwitchLookup

[mtrofin]

The switch becomes a conditional branch, one edge going to what was the default target of the switch, the other to a BB that performs a lookup in a table. The branch weights are accurately determinable from the ones of the switch.

Issue #147390

[mtrofin]

• [profcheck] Update exclusion list to reflect fixes #161943
• [SimplifyCFG][profcheck] Profile propagation for indirectbr #161747
• [SimplifyCFG][profcheck] Handle branch weights in simplifySwitchLookup #161739 👈 (View in Graphite)
• [SimplifyCFG][profcheck] Synthesize profile for br (X == 0 | X == 1), T, F1 -> switch #161549
• [SimplifyCFG][PGO] Reuse existing setBranchWeights #160629
• [profcheck][SimplifyCFG] Propagate !prof from switch to select #159645
• [profcheck] Option to inject distinct small weights #159644
• main

This stack of pull requests is managed by Graphite. Learn more about stacking.

[llvmbot]

@llvm/pr-subscribers-llvm-transforms

Author: Mircea Trofin (mtrofin)

Changes

---

Full diff: https://github.com/llvm/llvm-project/pull/161739.diff

2 Files Affected:

• (modified) llvm/lib/Transforms/Utils/SimplifyCFG.cpp (+18-4)
• (modified) llvm/test/Transforms/SimplifyCFG/rangereduce.ll (+20-4)

diff --git a/llvm/lib/Transforms/Utils/SimplifyCFG.cpp b/llvm/lib/Transforms/Utils/SimplifyCFG.cpp
index 63f4b2e030b69..fa3ac273b39f9 100644
--- a/llvm/lib/Transforms/Utils/SimplifyCFG.cpp
+++ b/llvm/lib/Transforms/Utils/SimplifyCFG.cpp
@@ -7319,19 +7319,33 @@ static bool simplifySwitchLookup(SwitchInst *SI, IRBuilder<> &Builder,
   if (DTU)
     Updates.push_back({DominatorTree::Insert, LookupBB, CommonDest});
 
+  SmallVector<uint32_t> BranchWeights;
+  const bool HasBranchWeights = RangeCheckBranch &&
+                                !ProfcheckDisableMetadataFixes &&
+                                extractBranchWeights(*SI, BranchWeights);
+  uint64_t ToLookupWeight = 0;
+  uint64_t ToDefaultWeight = 0;
+
   // Remove the switch.
   SmallPtrSet<BasicBlock *, 8> RemovedSuccessors;
-  for (unsigned i = 0, e = SI->getNumSuccessors(); i < e; ++i) {
-    BasicBlock *Succ = SI->getSuccessor(i);
+  for (unsigned I = 0, E = SI->getNumSuccessors(); I < E; ++I) {
+    BasicBlock *Succ = SI->getSuccessor(I);
 
-    if (Succ == SI->getDefaultDest())
+    if (Succ == SI->getDefaultDest()) {
+      if (HasBranchWeights)
+        ToDefaultWeight += BranchWeights[I];
       continue;
+    }
     Succ->removePredecessor(BB);
     if (DTU && RemovedSuccessors.insert(Succ).second)
       Updates.push_back({DominatorTree::Delete, BB, Succ});
+    if (HasBranchWeights)
+      ToLookupWeight += BranchWeights[I];
   }
   SI->eraseFromParent();
-
+  if (HasBranchWeights)
+    setFittedBranchWeights(*RangeCheckBranch, {ToLookupWeight, ToDefaultWeight},
+                           /*IsExpected=*/false);
   if (DTU)
     DTU->applyUpdates(Updates);
 
diff --git a/llvm/test/Transforms/SimplifyCFG/rangereduce.ll b/llvm/test/Transforms/SimplifyCFG/rangereduce.ll
index 17d65a4d4fa5e..d1fba91d1e505 100644
--- a/llvm/test/Transforms/SimplifyCFG/rangereduce.ll
+++ b/llvm/test/Transforms/SimplifyCFG/rangereduce.ll
@@ -1,15 +1,22 @@
-; 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-to-lookup -S | FileCheck %s
 ; RUN: opt < %s -passes='simplifycfg<switch-to-lookup>' -S | FileCheck %s
 
 target datalayout = "e-n32"
 
-define i32 @test1(i32 %a) {
+;.
+; CHECK: @switch.table.test1 = private unnamed_addr constant [4 x i32] [i32 11984, i32 1143, i32 99783, i32 99783], align 4
+; CHECK: @switch.table.test3 = private unnamed_addr constant [3 x i32] [i32 11984, i32 1143, i32 99783], align 4
+; CHECK: @switch.table.test6 = private unnamed_addr constant [4 x i32] [i32 99783, i32 99783, i32 1143, i32 11984], align 4
+; CHECK: @switch.table.test8 = private unnamed_addr constant [5 x i32] [i32 11984, i32 1143, i32 99783, i32 8867, i32 99783], align 4
+; CHECK: @switch.table.test9 = private unnamed_addr constant [8 x i32] [i32 99783, i32 8867, i32 99783, i32 8867, i32 8867, i32 8867, i32 11984, i32 1143], align 4
+;.
+define i32 @test1(i32 %a) !prof !0 {
 ; CHECK-LABEL: @test1(
 ; CHECK-NEXT:    [[TMP1:%.*]] = sub i32 [[A:%.*]], 97
 ; CHECK-NEXT:    [[TMP2:%.*]] = call i32 @llvm.fshl.i32(i32 [[TMP1]], i32 [[TMP1]], i32 30)
 ; CHECK-NEXT:    [[TMP3:%.*]] = icmp ult i32 [[TMP2]], 4
-; CHECK-NEXT:    br i1 [[TMP3]], label [[SWITCH_LOOKUP:%.*]], label [[COMMON_RET:%.*]]
+; CHECK-NEXT:    br i1 [[TMP3]], label [[SWITCH_LOOKUP:%.*]], label [[COMMON_RET:%.*]], !prof [[PROF1:![0-9]+]]
 ; CHECK:       switch.lookup:
 ; CHECK-NEXT:    [[TMP4:%.*]] = zext nneg i32 [[TMP2]] to i64
 ; CHECK-NEXT:    [[SWITCH_GEP:%.*]] = getelementptr inbounds [4 x i32], ptr @switch.table.test1, i64 0, i64 [[TMP4]]
@@ -24,7 +31,7 @@ define i32 @test1(i32 %a) {
   i32 101, label %two
   i32 105, label %three
   i32 109, label %three
-  ]
+  ], !prof !1
 
 def:
   ret i32 8867
@@ -310,3 +317,12 @@ three:
   ret i32 99783
 }
 
+!0 = !{!"function_entry_count", i32 100}
+!1 = !{!"branch_weights", i32 5, i32 7, i32 11, i32 13, i32 17}
+;.
+; CHECK: attributes #[[ATTR0:[0-9]+]] = { optsize }
+; CHECK: attributes #[[ATTR1:[0-9]+]] = { nocallback nofree nosync nounwind speculatable willreturn memory(none) }
+;.
+; CHECK: [[META0:![0-9]+]] = !{!"function_entry_count", i32 100}
+; CHECK: [[PROF1]] = !{!"branch_weights", i32 48, i32 5}
+;.

[mtrofin]

Merge activity

• Oct 4, 4:03 AM UTC: A user started a stack merge that includes this pull request via Graphite.
• Oct 4, 4:05 AM UTC: @mtrofin merged this pull request with Graphite.