[llvm] [SimplifyCFG] Delete the unnecessary range check for small mask operation (PR #65835)

[via llvm-commits]

https://github.com/vfdff updated https://github.com/llvm/llvm-project/pull/65835

>From 78e41596876e37df558a6fb2befe2d5061b8fbd1 Mon Sep 17 00:00:00 2001
From: zhongyunde 00443407 <zhongyunde at huawei.com>
Date: Fri, 1 Sep 2023 23:05:45 -0400
Subject: [PATCH] [SimplifyCFG] Delete the unnecessary range check for small
 mask operation

When the small mask value little than 64, we can eliminate the checking
for upper limit of the range by enlarge the lookup table size to its next
pow2 value. Similar to https://godbolt.org/z/TcshM956f
```
bool f(unsigned x) {
    switch (x % 8) {
        case 0: return 1;
        case 1: return 0;
        case 2: return 0;
        case 3: return 1;
        case 4: return 1;
        case 5: return 0;
        case 6: return 1;

        // This would remove the range check: case 7: return 0;
    }
    return 0;
}
```
Use WouldFitInRegister instead of fitsInLegalInteger to support
more result type beside bool.

Fixes https://github.com/llvm/llvm-project/issues/65120
---
 llvm/lib/Transforms/Utils/SimplifyCFG.cpp     | 16 +++++-
 .../Transforms/SimplifyCFG/switch_mask.ll     | 57 +++++++++++++++++++
 2 files changed, 71 insertions(+), 2 deletions(-)
 create mode 100644 llvm/test/Transforms/SimplifyCFG/switch_mask.ll

diff --git a/llvm/lib/Transforms/Utils/SimplifyCFG.cpp b/llvm/lib/Transforms/Utils/SimplifyCFG.cpp
index 35fead111aa9666..08c25cf2618f96d 100644
--- a/llvm/lib/Transforms/Utils/SimplifyCFG.cpp
+++ b/llvm/lib/Transforms/Utils/SimplifyCFG.cpp
@@ -6598,9 +6598,8 @@ static bool SwitchToLookupTable(SwitchInst *SI, IRBuilder<> &Builder,
   // If the default destination is unreachable, or if the lookup table covers
   // all values of the conditional variable, branch directly to the lookup table
   // BB. Otherwise, check that the condition is within the case range.
-  const bool DefaultIsReachable =
+  bool DefaultIsReachable =
       !isa<UnreachableInst>(SI->getDefaultDest()->getFirstNonPHIOrDbg());
-  const bool GeneratingCoveredLookupTable = (MaxTableSize == TableSize);
 
   // Create the BB that does the lookups.
   Module &Mod = *CommonDest->getParent()->getParent();
@@ -6631,6 +6630,19 @@ static bool SwitchToLookupTable(SwitchInst *SI, IRBuilder<> &Builder,
 
   BranchInst *RangeCheckBranch = nullptr;
 
+  // Grow the table to cover all possible index values to avoid the range check.
+  if (UseSwitchConditionAsTableIndex) {
+    ConstantRange CR = computeConstantRange(TableIndex, /* ForSigned */ false);
+    if (WouldFitInRegister(DL, CR.getUpper().getLimitedValue(),
+                           PHIs[0]->getType())) {
+      // The default branch is unreachable when we enlarge the lookup table.
+      // Adjust DefaultIsReachable to reuse code path.
+      TableSize = CR.getUpper().getZExtValue();
+      DefaultIsReachable = false;
+    }
+  }
+
+  const bool GeneratingCoveredLookupTable = (MaxTableSize == TableSize);
   if (!DefaultIsReachable || GeneratingCoveredLookupTable) {
     Builder.CreateBr(LookupBB);
     if (DTU)
diff --git a/llvm/test/Transforms/SimplifyCFG/switch_mask.ll b/llvm/test/Transforms/SimplifyCFG/switch_mask.ll
new file mode 100644
index 000000000000000..454814ec7070d42
--- /dev/null
+++ b/llvm/test/Transforms/SimplifyCFG/switch_mask.ll
@@ -0,0 +1,57 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt -passes=simplifycfg --switch-to-lookup -S < %s | FileCheck %s
+
+target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128"
+
+; https://alive2.llvm.org/ce/z/tuxLhJ
+define i1 @switch_lookup_with_small_mask(i64 %x) {
+; CHECK-LABEL: @switch_lookup_with_small_mask(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[AND:%.*]] = and i64 [[X:%.*]], 15
+; CHECK-NEXT:    [[SWITCH_CAST:%.*]] = trunc i64 [[AND]] to i16
+; CHECK-NEXT:    [[SWITCH_SHIFTAMT:%.*]] = mul nuw nsw i16 [[SWITCH_CAST]], 1
+; CHECK-NEXT:    [[SWITCH_DOWNSHIFT:%.*]] = lshr i16 1030, [[SWITCH_SHIFTAMT]]
+; CHECK-NEXT:    [[SWITCH_MASKED:%.*]] = trunc i16 [[SWITCH_DOWNSHIFT]] to i1
+; CHECK-NEXT:    ret i1 [[SWITCH_MASKED]]
+;
+entry:
+  %and = and i64 %x, 15
+  switch i64 %and, label %default [
+  i64 10, label %lor.end
+  i64 1, label %lor.end
+  i64 2, label %lor.end
+  ]
+
+default:                                          ; preds = %entry
+  br label %lor.end
+
+lor.end:                                          ; preds = %entry, %entry, %entry, %default
+  %0 = phi i1 [ true, %entry ], [ false, %default ], [ true, %entry ], [ true, %entry ]
+  ret i1 %0
+}
+
+define i1 @switch_lookup_with_small_urem(i64 %x) {
+; CHECK-LABEL: @switch_lookup_with_small_urem(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[AND:%.*]] = urem i64 [[X:%.*]], 42
+; CHECK-NEXT:    [[SWITCH_CAST:%.*]] = trunc i64 [[AND]] to i42
+; CHECK-NEXT:    [[SWITCH_SHIFTAMT:%.*]] = mul nuw nsw i42 [[SWITCH_CAST]], 1
+; CHECK-NEXT:    [[SWITCH_DOWNSHIFT:%.*]] = lshr i42 1030, [[SWITCH_SHIFTAMT]]
+; CHECK-NEXT:    [[SWITCH_MASKED:%.*]] = trunc i42 [[SWITCH_DOWNSHIFT]] to i1
+; CHECK-NEXT:    ret i1 [[SWITCH_MASKED]]
+;
+entry:
+  %and = urem i64 %x, 42
+  switch i64 %and, label %default [
+  i64 10, label %lor.end
+  i64 1, label %lor.end
+  i64 2, label %lor.end
+  ]
+
+default:                                          ; preds = %entry
+  br label %lor.end
+
+lor.end:                                          ; preds = %entry, %entry, %entry, %default
+  %0 = phi i1 [ true, %entry ], [ false, %default ], [ true, %entry ], [ true, %entry ]
+  ret i1 %0
+}