[llvm] [SimplifyCFG] Fix Missing Case Computation for Arbitrary Width Integers (PR #77831)

[Qiongsi Wu via llvm-commits]

https://github.com/qiongsiwu created https://github.com/llvm/llvm-project/pull/77831

https://github.com/llvm/llvm-project/pull/76669 taught SimplifyCFG to handle switches when `default` has only one case. When the `switch`'s condition is wider than 64 bit, the current implementation can calculate the wrong default value. For example, when the condition is an `i128` and the cases are `1 * 2^110`, `2 * 2^110` and `3 * 2^110` and the default case is `0 * 2^110`, the default value should be `i128 0`, but https://github.com/llvm/llvm-project/pull/76669 computes `UINT64_MAX` as the default value. This PR computes the default with `APInt` instead of `uint64_t`. 

>From bf3bbc4ef3ae2afb5a2681f62c348c3d237a3c26 Mon Sep 17 00:00:00 2001
From: Qiongsi Wu <qwu at ibm.com>
Date: Thu, 11 Jan 2024 15:45:24 -0500
Subject: [PATCH] Fix missing case computation for wide integers.

---
 llvm/lib/Transforms/Utils/SimplifyCFG.cpp     |  7 ++---
 .../SimplifyCFG/switch-dead-default.ll        | 26 +++++++++++++++++++
 2 files changed, 30 insertions(+), 3 deletions(-)

diff --git a/llvm/lib/Transforms/Utils/SimplifyCFG.cpp b/llvm/lib/Transforms/Utils/SimplifyCFG.cpp
index 7515e539e7fb78..47bda070ebc0a5 100644
--- a/llvm/lib/Transforms/Utils/SimplifyCFG.cpp
+++ b/llvm/lib/Transforms/Utils/SimplifyCFG.cpp
@@ -5623,11 +5623,12 @@ static bool eliminateDeadSwitchCases(SwitchInst *SI, DomTreeUpdater *DTU,
     // optimization, such as lookup tables.
     if (SI->getNumCases() == AllNumCases - 1) {
       assert(NumUnknownBits > 1 && "Should be canonicalized to a branch");
-      uint64_t MissingCaseVal = 0;
+      IntegerType *CondType = cast<IntegerType>(Cond->getType());
+      APInt MissingCaseVal = APInt::getZero(CondType->getBitWidth());
       for (const auto &Case : SI->cases())
-        MissingCaseVal ^= Case.getCaseValue()->getValue().getLimitedValue();
+        MissingCaseVal ^= Case.getCaseValue()->getValue();
       auto *MissingCase =
-          cast<ConstantInt>(ConstantInt::get(Cond->getType(), MissingCaseVal));
+          ConstantInt::get(SI->getModule()->getContext(), MissingCaseVal);
       SwitchInstProfUpdateWrapper SIW(*SI);
       SIW.addCase(MissingCase, SI->getDefaultDest(), SIW.getSuccessorWeight(0));
       createUnreachableSwitchDefault(SI, DTU, /*RemoveOrigDefaultBlock*/ false);
diff --git a/llvm/test/Transforms/SimplifyCFG/switch-dead-default.ll b/llvm/test/Transforms/SimplifyCFG/switch-dead-default.ll
index e30a535c523237..4e47d4751ab15d 100644
--- a/llvm/test/Transforms/SimplifyCFG/switch-dead-default.ll
+++ b/llvm/test/Transforms/SimplifyCFG/switch-dead-default.ll
@@ -2,6 +2,7 @@
 ; RUN: opt %s -S -passes='simplifycfg<switch-to-lookup>' -simplifycfg-require-and-preserve-domtree=1 -switch-range-to-icmp | FileCheck %s
 
 target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128"
+
 declare void @foo(i32)
 
 define void @test(i1 %a) {
@@ -313,6 +314,31 @@ default:
 
 declare void @llvm.assume(i1)
 
+define zeroext i1 @test8(i128 %a) {
+; CHECK-LABEL: define zeroext i1 @test8(
+; CHECK-SAME: i128 [[A:%.*]]) {
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[TMP0:%.*]] = and i128 [[A]], 3894222643901120721397872246915072
+; CHECK-NEXT:    [[SWITCH:%.*]] = icmp ult i128 [[TMP0]], 1
+; CHECK-NEXT:    [[SPEC_SELECT:%.*]] = select i1 [[SWITCH]], i1 false, i1 true
+; CHECK-NEXT:    ret i1 [[SPEC_SELECT]]
+;
+entry:
+  %0 = and i128 %a, 3894222643901120721397872246915072
+  switch i128 %0, label %lor.rhs [
+  i128 1298074214633706907132624082305024, label %lor.end
+  i128 2596148429267413814265248164610048, label %lor.end
+  i128 3894222643901120721397872246915072, label %lor.end
+  ]
+
+lor.rhs:                                          ; preds = %entry
+  br label %lor.end
+
+lor.end:                                          ; preds = %entry, %entry, %entry, %lor.rhs
+  %1 = phi i1 [ true, %entry ], [ false, %lor.rhs ], [ true, %entry ], [ true, %entry ]
+  ret i1 %1
+}
+
 !0 = !{!"branch_weights", i32 8, i32 4, i32 2, i32 1}
 ;.
 ; CHECK: [[PROF0]] = !{!"branch_weights", i32 0, i32 4, i32 2, i32 1, i32 8}