[SimplifyCFG] switch: Do Not Transform the Default Case if the Condition is Too Wide

[qiongsiwu]

#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. This PR skips cases where the condition is too wide.

[llvmbot]

@llvm/pr-subscribers-llvm-transforms

Author: Qiongsi Wu (qiongsiwu)

Changes

#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 #76669 computes UINT64_MAX as the default value. This PR computes the default with APInt instead of uint64_t.

---

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

2 Files Affected:

• (modified) llvm/lib/Transforms/Utils/SimplifyCFG.cpp (+4-3)
• (modified) llvm/test/Transforms/SimplifyCFG/switch-dead-default.ll (+26)

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}

[dtcxzyw]

I don't think it is worth handing switch i128 cases.
We can early exit if the bit width of condition value is larger than 64:

  auto *CondTy = cast<IntegerType>(SI->getCondition()->getType());
  if (CondTy->getIntegerBitWidth() > 64 ||
      !DL.fitsInLegalInteger(CondTy->getIntegerBitWidth()))
    return false;

[qiongsiwu]

I don't think it is worth handing switch i128 cases. We can early exit if the bit width of condition value is larger than 64:

  auto *CondTy = cast<IntegerType>(SI->getCondition()->getType());
  if (CondTy->getIntegerBitWidth() > 64 ||
      !DL.fitsInLegalInteger(CondTy->getIntegerBitWidth()))
    return false;

Thanks for the fast response @dtcxzyw ! I can certainly implement the early escape. Could you help me understand why it is not worth handling conditions of wider widths?

[dtcxzyw]

I don't think it is worth handing switch i128 cases. We can early exit if the bit width of condition value is larger than 64:

  auto *CondTy = cast<IntegerType>(SI->getCondition()->getType());
  if (CondTy->getIntegerBitWidth() > 64 ||
      !DL.fitsInLegalInteger(CondTy->getIntegerBitWidth()))
    return false;

Thanks for the fast response @dtcxzyw ! I can certainly implement the early escape. Could you help me understand why it is not worth handling conditions of wider widths?

• switch i128 is rare in real code.
• It is impossible to convert it into jump table.

[qiongsiwu]

* `switch i128` is rare in real code.

* It is impossible to convert it into jump table.

Sounds good! The PR is revised. I kept the test case to "document" the behaviour, but I can take it out if it is not necessary.

[DianQK]

LGTM. Thanks!

[qiongsiwu]

Thanks @DianQK !

In light of #76669 (comment), I will not merge this PR at the moment, since #76669 may get reverted or require more work.

[nikic]

@qiongsiwu Please do merge it now. If necessary, this change will be reverted together with the base commit.

[qiongsiwu]

@qiongsiwu Please do merge it now. If necessary, this change will be reverted together with the base commit.

Thanks for chiming in @nikic ! Merging this PR now.