[llvm] [SimplifyCFG] Fix uint32_t overflow in cbranch to cbranch merge prevention check. (PR #72329)

[Valery Pykhtin via llvm-commits]

https://github.com/vpykhtin created https://github.com/llvm/llvm-project/pull/72329

This fixes https://github.com/llvm/llvm-project/issues/72323.

Resulted from https://github.com/llvm/llvm-project/commit/f054947c0da99ec8b3c4bb043e5225672420a313

>From 2cfd9cf7fa1e1a457fc6ddc6a466b1dfdc0e95af Mon Sep 17 00:00:00 2001
From: Valery Pykhtin <valery.pykhtin at gmail.com>
Date: Wed, 15 Nov 2023 01:17:05 +0100
Subject: [PATCH 1/2] test before fix

---
 .../SimplifyCFG/branch-cond-dont-merge.ll     | 19 +++++++++++++++++++
 1 file changed, 19 insertions(+)

diff --git a/llvm/test/Transforms/SimplifyCFG/branch-cond-dont-merge.ll b/llvm/test/Transforms/SimplifyCFG/branch-cond-dont-merge.ll
index 5c21f163826ee25..3e924dc15bbe123 100644
--- a/llvm/test/Transforms/SimplifyCFG/branch-cond-dont-merge.ll
+++ b/llvm/test/Transforms/SimplifyCFG/branch-cond-dont-merge.ll
@@ -79,6 +79,25 @@ exit:
   ret void
 }
 
+define void @uint32_overflow_test(i1 %arg, i1 %arg1) {
+; CHECK-LABEL: @uint32_overflow_test(
+; CHECK-NEXT:  bb:
+; CHECK-NEXT:    ret void
+;
+bb:
+  br i1 %arg, label %bb4, label %bb2, !prof !3
+
+bb2:
+  br i1 %arg1, label %bb4, label %bb3
+
+bb3:
+  br label %bb4
+
+bb4:
+  ret void
+}
+
 !0 = !{!"branch_weights", i32 1, i32 1000}
 !1 = !{!"branch_weights", i32 1000, i32 1}
 !2 = !{!"branch_weights", i32 3, i32 2}
+!3 = !{!"branch_weights", i32 -258677585, i32 -1212131848}

>From 5ffc278fcc5af769629a7e1a03f51021cde0412f Mon Sep 17 00:00:00 2001
From: Valery Pykhtin <valery.pykhtin at gmail.com>
Date: Wed, 15 Nov 2023 01:18:39 +0100
Subject: [PATCH 2/2] fix

---
 llvm/lib/Transforms/Utils/SimplifyCFG.cpp                  | 4 ++--
 llvm/test/Transforms/SimplifyCFG/branch-cond-dont-merge.ll | 1 +
 2 files changed, 3 insertions(+), 2 deletions(-)

diff --git a/llvm/lib/Transforms/Utils/SimplifyCFG.cpp b/llvm/lib/Transforms/Utils/SimplifyCFG.cpp
index 6009558efca06af..f15ab6858401fc6 100644
--- a/llvm/lib/Transforms/Utils/SimplifyCFG.cpp
+++ b/llvm/lib/Transforms/Utils/SimplifyCFG.cpp
@@ -4351,10 +4351,10 @@ static bool SimplifyCondBranchToCondBranch(BranchInst *PBI, BranchInst *BI,
   SmallVector<uint32_t, 2> PredWeights;
   if (!PBI->getMetadata(LLVMContext::MD_unpredictable) &&
       extractBranchWeights(*PBI, PredWeights) &&
-      (PredWeights[0] + PredWeights[1]) != 0) {
+      ((uint64_t)PredWeights[0] + PredWeights[1]) != 0) {
 
     BranchProbability CommonDestProb = BranchProbability::getBranchProbability(
-        PredWeights[PBIOp], PredWeights[0] + PredWeights[1]);
+        PredWeights[PBIOp], (uint64_t)PredWeights[0] + PredWeights[1]);
 
     BranchProbability Likely = TTI.getPredictableBranchThreshold();
     if (CommonDestProb >= Likely)
diff --git a/llvm/test/Transforms/SimplifyCFG/branch-cond-dont-merge.ll b/llvm/test/Transforms/SimplifyCFG/branch-cond-dont-merge.ll
index 3e924dc15bbe123..0f540007b242e29 100644
--- a/llvm/test/Transforms/SimplifyCFG/branch-cond-dont-merge.ll
+++ b/llvm/test/Transforms/SimplifyCFG/branch-cond-dont-merge.ll
@@ -82,6 +82,7 @@ exit:
 define void @uint32_overflow_test(i1 %arg, i1 %arg1) {
 ; CHECK-LABEL: @uint32_overflow_test(
 ; CHECK-NEXT:  bb:
+; CHECK-NEXT:    [[BRMERGE:%.*]] = select i1 [[ARG:%.*]], i1 true, i1 [[ARG1:%.*]]
 ; CHECK-NEXT:    ret void
 ;
 bb: