[SimplifyCFG] Delete the unnecessary range check for small mask operation

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,25 @@ 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);
+    // Grow the table shouldn't have any size impact by checking
+    // WouldFitInRegister.
+    // TODO: Consider growing the table also when it doesn't fit in a register
+    // if no optsize is specified.
+    if (all_of(ResultTypes, [&](const auto &KV) {
+          return SwitchLookupTable::WouldFitInRegister(
+              DL, CR.getUpper().getLimitedValue(), KV.second /* ResultType */);
+        })) {
+      // 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)

llvm/test/Transforms/SimplifyCFG/switch_mask.ll

@@ -0,0 +1,105 @@
+; 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_i1(i64 %x) {
+; CHECK-LABEL: @switch_lookup_with_small_i1(
+; 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
+}
+
+; https://godbolt.org/z/sjbjorKon
+define i8 @switch_lookup_with_small_i8(i64 %x) {
+; CHECK-LABEL: @switch_lookup_with_small_i8(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[REM:%.*]] = urem i64 [[X:%.*]], 5
+; CHECK-NEXT:    [[SWITCH_CAST:%.*]] = trunc i64 [[REM]] to i40
+; CHECK-NEXT:    [[SWITCH_SHIFTAMT:%.*]] = mul nuw nsw i40 [[SWITCH_CAST]], 8
+; CHECK-NEXT:    [[SWITCH_DOWNSHIFT:%.*]] = lshr i40 460303, [[SWITCH_SHIFTAMT]]
+; CHECK-NEXT:    [[SWITCH_MASKED:%.*]] = trunc i40 [[SWITCH_DOWNSHIFT]] to i8
+; CHECK-NEXT:    ret i8 [[SWITCH_MASKED]]
+;
+entry:
+  %rem = urem i64 %x, 5
+  switch i64 %rem, label %default [
+  i64 0, label %sw.bb0
+  i64 1, label %sw.bb1
+  i64 2, label %sw.bb2
+  ]
+
+sw.bb0:                                           ; preds = %entry
+  br label %lor.end
+
+sw.bb1:                                           ; preds = %entry
+  br label %lor.end
+
+sw.bb2:                                           ; preds = %entry
+  br label %lor.end
+
+default:                                          ; preds = %entry
+  br label %lor.end
+
+lor.end:
+  %0 = phi i8 [ 15, %sw.bb0 ], [ 6, %sw.bb1 ], [ 7, %sw.bb2 ], [ 0, %default ]
+  ret i8 %0
+}
+
+; Negative test: Table size would not fit the register.
+define i8 @switch_lookup_with_small_i8_negative(i64 %x) {
+; CHECK-LABEL: @switch_lookup_with_small_i8_negative(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[REM:%.*]] = urem i64 [[X:%.*]], 9
+; CHECK-NEXT:    [[TMP0:%.*]] = icmp ult i64 [[REM]], 3
+; CHECK-NEXT:    [[SWITCH_CAST:%.*]] = trunc i64 [[REM]] to i24
+; CHECK-NEXT:    [[SWITCH_SHIFTAMT:%.*]] = mul nuw nsw i24 [[SWITCH_CAST]], 8
+; CHECK-NEXT:    [[SWITCH_DOWNSHIFT:%.*]] = lshr i24 460303, [[SWITCH_SHIFTAMT]]
+; CHECK-NEXT:    [[SWITCH_MASKED:%.*]] = trunc i24 [[SWITCH_DOWNSHIFT]] to i8
+; CHECK-NEXT:    [[TMP1:%.*]] = select i1 [[TMP0]], i8 [[SWITCH_MASKED]], i8 0
+; CHECK-NEXT:    ret i8 [[TMP1]]
+;
+entry:
+  %rem = urem i64 %x, 9                           ; 9 * 8 = 72 > 64, not fit the register
+  switch i64 %rem, label %default [
+  i64 0, label %sw.bb0
+  i64 1, label %sw.bb1
+  i64 2, label %sw.bb2
+  ]
+
+sw.bb0:                                           ; preds = %entry
+  br label %lor.end
+
+sw.bb1:                                           ; preds = %entry
+  br label %lor.end
+
+sw.bb2:                                           ; preds = %entry
+  br label %lor.end
+
+default:                                          ; preds = %entry
+  br label %lor.end
+
+lor.end:
+  %0 = phi i8 [ 15, %sw.bb0 ], [ 6, %sw.bb1 ], [ 7, %sw.bb2 ], [ 0, %default ]
+  ret i8 %0
+}