[llvm] [SimplifyCFG] Fix branch-weight overflow when folding switch case into default (PR #212520)

[via llvm-commits]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-llvm-transforms

Author: Alok Kumar Sharma (alokkrsharma)

<details>
<summary>Changes</summary>

The overflow guard added in #<!-- -->178964 for merging a switch case's weight into the default's calls fitWeights() to downscale the weights, but the returned value was never assigned back to Weights, so the downscale didn't take effect. As a result, a default+case sum that overflows uint64_t would wrap around, turning a dominant default weight into a small, incorrect one.

Assign fitWeights()'s result back to Weights before the addition so the downscale takes effect as intended.

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


2 Files Affected:

- (modified) llvm/lib/Transforms/Utils/Local.cpp (+4-2) 
- (added) llvm/test/Transforms/SimplifyCFG/switch-default-fold-weight-overflow.ll (+45) 


``````````diff
diff --git a/llvm/lib/Transforms/Utils/Local.cpp b/llvm/lib/Transforms/Utils/Local.cpp
index b17740c0bc192..643ad0fb24267 100644
--- a/llvm/lib/Transforms/Utils/Local.cpp
+++ b/llvm/lib/Transforms/Utils/Local.cpp
@@ -228,8 +228,10 @@ bool llvm::ConstantFoldTerminator(BasicBlock *BB, bool DeleteDeadConditions,
           unsigned Idx = It->getCaseIndex();
 
           // Check for and prevent uint64_t overflow by reducing branch weights.
-          if (Weights[0] > UINT64_MAX - Weights[Idx + 1])
-            fitWeights(Weights);
+          if (Weights[0] > UINT64_MAX - Weights[Idx + 1]) {
+            SmallVector<uint32_t> Fitted = fitWeights(Weights);
+            Weights.assign(Fitted.begin(), Fitted.end());
+          }
 
           Weights[0] += Weights[Idx + 1];
           // Remove weight for this case.
diff --git a/llvm/test/Transforms/SimplifyCFG/switch-default-fold-weight-overflow.ll b/llvm/test/Transforms/SimplifyCFG/switch-default-fold-weight-overflow.ll
new file mode 100644
index 0000000000000..91311343e5572
--- /dev/null
+++ b/llvm/test/Transforms/SimplifyCFG/switch-default-fold-weight-overflow.ll
@@ -0,0 +1,45 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 6
+; RUN: opt < %s -passes=simplifycfg -S | FileCheck %s
+
+; A case whose destination is the same as the default gets folded into the
+; default, adding its weight to the default's. If default + case overflows
+; uint64_t, the weights must be downscaled via fitWeights() before the
+; addition; otherwise it wraps around and the default's huge weight turns
+; into a tiny, wrong one.
+;
+; default = 2^64 - 100, folded case = 200 (default + case wraps around to
+; 100 as a uint64_t), surviving case = 1000. Correctly scaled, default stays
+; by far the biggest weight. If the downscale isn't applied, it collapses to
+; 100 -- smaller than the surviving case's 1000, which is the wrong relative
+; order.
+
+define void @switch_default_fold_overflow(i32 %x, ptr %p) {
+; CHECK-LABEL: define void @switch_default_fold_overflow(
+; CHECK-SAME: i32 [[X:%.*]], ptr [[P:%.*]]) {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    [[COND:%.*]] = icmp eq i32 [[X]], 1
+; CHECK-NEXT:    br i1 [[COND]], label %[[OTHER:.*]], label %[[DEFAULT:.*]], !prof [[PROF0:![0-9]+]]
+; CHECK:       [[DEFAULT]]:
+; CHECK-NEXT:    ret void
+; CHECK:       [[OTHER]]:
+; CHECK-NEXT:    store i32 0, ptr [[P]], align 4
+; CHECK-NEXT:    br label %[[DEFAULT]]
+;
+entry:
+  switch i32 %x, label %default [
+  i32 0, label %default
+  i32 1, label %other
+  ], !prof !0
+
+default:
+  ret void
+
+other:
+  store i32 0, ptr %p
+  br label %default
+}
+
+!0 = !{!"branch_weights", i64 18446744073709551516, i64 200, i64 1000}
+;.
+; CHECK: [[PROF0]] = !{!"branch_weights", i32 0, i32 -1}
+;.

``````````

</details>


https://github.com/llvm/llvm-project/pull/212520