diff --git a/llvm/lib/Transforms/Utils/SimplifyCFG.cpp b/llvm/lib/Transforms/Utils/SimplifyCFG.cpp index 68b5b1a78a346..18187bcdedf09 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(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) diff --git a/llvm/test/Transforms/SimplifyCFG/switch_mask.ll b/llvm/test/Transforms/SimplifyCFG/switch_mask.ll new file mode 100644 index 0000000000000..123519bc69211 --- /dev/null +++ b/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 +}