Skip to content

Commit

Permalink
[SimplifyCFG] Find the smallest table considering overflow in `switch…
Browse files Browse the repository at this point in the history
…ToLookupTable`
  • Loading branch information
DianQK committed Feb 5, 2024
1 parent 6a3b696 commit 4eaa728
Show file tree
Hide file tree
Showing 2 changed files with 71 additions and 51 deletions.
92 changes: 64 additions & 28 deletions llvm/lib/Transforms/Utils/SimplifyCFG.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -6516,18 +6516,20 @@ ShouldBuildLookupTable(SwitchInst *SI, uint64_t TableSize,
}

static bool ShouldUseSwitchConditionAsTableIndex(
ConstantInt &MinCaseVal, const ConstantInt &MaxCaseVal,
const ConstantInt &BeginCaseVal, const ConstantInt &EndCaseVal,
bool HasDefaultResults, const SmallDenseMap<PHINode *, Type *> &ResultTypes,
const DataLayout &DL, const TargetTransformInfo &TTI) {
if (MinCaseVal.isNullValue())
if (BeginCaseVal.isNullValue())
return true;
if (MinCaseVal.isNegative() ||
MaxCaseVal.getLimitedValue() == std::numeric_limits<uint64_t>::max() ||
if (BeginCaseVal.getValue().sge(EndCaseVal.getValue()))
return false;
if (BeginCaseVal.isNegative() ||
EndCaseVal.getLimitedValue() == std::numeric_limits<uint64_t>::max() ||
!HasDefaultResults)
return false;
return all_of(ResultTypes, [&](const auto &KV) {
return SwitchLookupTable::WouldFitInRegister(
DL, MaxCaseVal.getLimitedValue() + 1 /* TableSize */,
DL, EndCaseVal.getLimitedValue() + 1 /* TableSize */,
KV.second /* ResultType */);
});
}
Expand Down Expand Up @@ -6618,7 +6620,8 @@ static void reuseTableCompare(
/// lookup tables.
static bool SwitchToLookupTable(SwitchInst *SI, IRBuilder<> &Builder,
DomTreeUpdater *DTU, const DataLayout &DL,
const TargetTransformInfo &TTI) {
const TargetTransformInfo &TTI,
bool TryMinTableSize) {
assert(SI->getNumCases() > 1 && "Degenerate switch?");

BasicBlock *BB = SI->getParent();
Expand All @@ -6644,9 +6647,6 @@ static bool SwitchToLookupTable(SwitchInst *SI, IRBuilder<> &Builder,
// Figure out the corresponding result for each case value and phi node in the
// common destination, as well as the min and max case values.
assert(!SI->cases().empty());
SwitchInst::CaseIt CI = SI->case_begin();
ConstantInt *MinCaseVal = CI->getCaseValue();
ConstantInt *MaxCaseVal = CI->getCaseValue();

BasicBlock *CommonDest = nullptr;

Expand All @@ -6657,17 +6657,49 @@ static bool SwitchToLookupTable(SwitchInst *SI, IRBuilder<> &Builder,
SmallDenseMap<PHINode *, Type *> ResultTypes;
SmallVector<PHINode *, 4> PHIs;

for (SwitchInst::CaseIt E = SI->case_end(); CI != E; ++CI) {
ConstantInt *CaseVal = CI->getCaseValue();
if (CaseVal->getValue().slt(MinCaseVal->getValue()))
MinCaseVal = CaseVal;
if (CaseVal->getValue().sgt(MaxCaseVal->getValue()))
MaxCaseVal = CaseVal;
SmallVector<ConstantInt *, 8> CaseVals(llvm::map_range(
SI->cases(), [](const auto &C) { return C.getCaseValue(); }));

llvm::sort(CaseVals, [](const auto *L, const auto *R) {
return L->getValue().slt(R->getValue());
});
auto *CaseValIter = CaseVals.begin();
ConstantInt *BeginCaseVal = *CaseValIter;
ConstantInt *EndCaseVal = CaseVals.back();
bool RangeOverflow = false;
uint64_t MinTableSize = EndCaseVal->getValue()
.ssub_ov(BeginCaseVal->getValue(), RangeOverflow)
.getLimitedValue() +
1;
// If there is no overflow, then this must be the minimal table.
// The signed max-min can no longer build a lookup table, so return.
if (RangeOverflow && TryMinTableSize) {
// We consider cases where the starting to the endpoint will cross the
// signed max and min. For example, for the i8 range `[-128, -127, 126,
// 127]`, we choose from 126 to -127. The length of the lookup table is 4.
while (CaseValIter != CaseVals.end()) {
auto *CurCaseVal = *CaseValIter++;
if (CaseValIter == CaseVals.end())
break;
auto *NextCaseVal = *CaseValIter;
const auto &NextVal = NextCaseVal->getValue();
const auto &CurVal = CurCaseVal->getValue();
uint64_t RequireTableSize = (CurVal - NextVal).getLimitedValue() + 1;
if (RequireTableSize < MinTableSize) {
BeginCaseVal = NextCaseVal;
EndCaseVal = CurCaseVal;
MinTableSize = RequireTableSize;
}
}
}

for (const auto &CI : SI->cases()) {
ConstantInt *CaseVal = CI.getCaseValue();

// Resulting value at phi nodes for this case value.
using ResultsTy = SmallVector<std::pair<PHINode *, Constant *>, 4>;
ResultsTy Results;
if (!getCaseResults(SI, CaseVal, CI->getCaseSuccessor(), &CommonDest,
if (!getCaseResults(SI, CaseVal, CI.getCaseSuccessor(), &CommonDest,
Results, DL, TTI))
return false;

Expand Down Expand Up @@ -6702,13 +6734,12 @@ static bool SwitchToLookupTable(SwitchInst *SI, IRBuilder<> &Builder,
}

bool UseSwitchConditionAsTableIndex = ShouldUseSwitchConditionAsTableIndex(
*MinCaseVal, *MaxCaseVal, HasDefaultResults, ResultTypes, DL, TTI);
*BeginCaseVal, *EndCaseVal, HasDefaultResults, ResultTypes, DL, TTI);
uint64_t TableSize;
if (UseSwitchConditionAsTableIndex)
TableSize = MaxCaseVal->getLimitedValue() + 1;
TableSize = EndCaseVal->getLimitedValue() + 1;
else
TableSize =
(MaxCaseVal->getValue() - MinCaseVal->getValue()).getLimitedValue() + 1;
TableSize = MinTableSize;

bool TableHasHoles = (NumResults < TableSize);
bool NeedMask = (TableHasHoles && !HasDefaultResults);
Expand All @@ -6721,13 +6752,16 @@ static bool SwitchToLookupTable(SwitchInst *SI, IRBuilder<> &Builder,
}

if (!ShouldBuildLookupTable(SI, TableSize, TTI, DL, ResultTypes))
return false;
// When a signed max-min cannot construct a lookup table, try to find a
// range with a minimal lookup table.
return !TryMinTableSize &&
SwitchToLookupTable(SI, Builder, DTU, DL, TTI, true);

std::vector<DominatorTree::UpdateType> Updates;

// Compute the maximum table size representable by the integer type we are
// switching upon.
unsigned CaseSize = MinCaseVal->getType()->getPrimitiveSizeInBits();
unsigned CaseSize = BeginCaseVal->getType()->getPrimitiveSizeInBits();
uint64_t MaxTableSize = CaseSize > 63 ? UINT64_MAX : 1ULL << CaseSize;
assert(MaxTableSize >= TableSize &&
"It is impossible for a switch to have more entries than the max "
Expand All @@ -6749,15 +6783,17 @@ static bool SwitchToLookupTable(SwitchInst *SI, IRBuilder<> &Builder,
Value *TableIndex;
ConstantInt *TableIndexOffset;
if (UseSwitchConditionAsTableIndex) {
TableIndexOffset = ConstantInt::get(MaxCaseVal->getIntegerType(), 0);
TableIndexOffset = ConstantInt::get(EndCaseVal->getIntegerType(), 0);
TableIndex = SI->getCondition();
} else {
TableIndexOffset = MinCaseVal;
TableIndexOffset = BeginCaseVal;
// If the default is unreachable, all case values are s>= MinCaseVal. Then
// we can try to attach nsw.
bool MayWrap = true;
if (!DefaultIsReachable) {
APInt Res = MaxCaseVal->getValue().ssub_ov(MinCaseVal->getValue(), MayWrap);
if (!DefaultIsReachable &&
EndCaseVal->getValue().sge(BeginCaseVal->getValue())) {
APInt Res =
EndCaseVal->getValue().ssub_ov(BeginCaseVal->getValue(), MayWrap);
(void)Res;
}

Expand Down Expand Up @@ -6798,7 +6834,7 @@ static bool SwitchToLookupTable(SwitchInst *SI, IRBuilder<> &Builder,
// PHI value for the default case in case we're using a bit mask.
} else {
Value *Cmp = Builder.CreateICmpULT(
TableIndex, ConstantInt::get(MinCaseVal->getType(), TableSize));
TableIndex, ConstantInt::get(BeginCaseVal->getType(), TableSize));
RangeCheckBranch =
Builder.CreateCondBr(Cmp, LookupBB, SI->getDefaultDest());
if (DTU)
Expand Down Expand Up @@ -7112,7 +7148,7 @@ bool SimplifyCFGOpt::simplifySwitch(SwitchInst *SI, IRBuilder<> &Builder) {
// CVP. Therefore, only apply this transformation during late stages of the
// optimisation pipeline.
if (Options.ConvertSwitchToLookupTable &&
SwitchToLookupTable(SI, Builder, DTU, DL, TTI))
SwitchToLookupTable(SI, Builder, DTU, DL, TTI, false))
return requestResimplify();

if (simplifySwitchOfPowersOfTwo(SI, Builder, DL, TTI))
Expand Down
30 changes: 7 additions & 23 deletions llvm/test/Transforms/SimplifyCFG/X86/switch_to_lookup_table.ll
Original file line number Diff line number Diff line change
Expand Up @@ -120,31 +120,15 @@ return:
define i32 @f_i8_128(i8 %c) {
; CHECK-LABEL: @f_i8_128(
; CHECK-NEXT: entry:
; CHECK-NEXT: switch i8 [[C:%.*]], label [[SW_DEFAULT:%.*]] [
; CHECK-NEXT: i8 122, label [[RETURN:%.*]]
; CHECK-NEXT: i8 123, label [[SW_BB1:%.*]]
; CHECK-NEXT: i8 124, label [[SW_BB2:%.*]]
; CHECK-NEXT: i8 125, label [[SW_BB3:%.*]]
; CHECK-NEXT: i8 126, label [[SW_BB4:%.*]]
; CHECK-NEXT: i8 127, label [[SW_BB5:%.*]]
; CHECK-NEXT: i8 -128, label [[SW_BB6:%.*]]
; CHECK-NEXT: ]
; CHECK: sw.bb1:
; CHECK-NEXT: br label [[RETURN]]
; CHECK: sw.bb2:
; CHECK-NEXT: br label [[RETURN]]
; CHECK: sw.bb3:
; CHECK-NEXT: br label [[RETURN]]
; CHECK: sw.bb4:
; CHECK-NEXT: br label [[RETURN]]
; CHECK: sw.bb5:
; CHECK-NEXT: br label [[RETURN]]
; CHECK: sw.bb6:
; CHECK-NEXT: br label [[RETURN]]
; CHECK: sw.default:
; CHECK-NEXT: [[SWITCH_TABLEIDX:%.*]] = sub i8 [[C:%.*]], 122
; CHECK-NEXT: [[TMP0:%.*]] = icmp ult i8 [[SWITCH_TABLEIDX]], 7
; CHECK-NEXT: br i1 [[TMP0]], label [[SWITCH_LOOKUP:%.*]], label [[RETURN:%.*]]
; CHECK: switch.lookup:
; CHECK-NEXT: [[SWITCH_GEP:%.*]] = getelementptr inbounds [7 x i32], ptr @switch.table.f_i8_128, i32 0, i8 [[SWITCH_TABLEIDX]]
; CHECK-NEXT: [[SWITCH_LOAD:%.*]] = load i32, ptr [[SWITCH_GEP]], align 4
; CHECK-NEXT: br label [[RETURN]]
; CHECK: return:
; CHECK-NEXT: [[RETVAL_0:%.*]] = phi i32 [ 15, [[SW_DEFAULT]] ], [ 1, [[SW_BB6]] ], [ 62, [[SW_BB5]] ], [ 27, [[SW_BB4]] ], [ -1, [[SW_BB3]] ], [ 0, [[SW_BB2]] ], [ 123, [[SW_BB1]] ], [ 55, [[ENTRY:%.*]] ]
; CHECK-NEXT: [[RETVAL_0:%.*]] = phi i32 [ [[SWITCH_LOAD]], [[SWITCH_LOOKUP]] ], [ 15, [[ENTRY:%.*]] ]
; CHECK-NEXT: ret i32 [[RETVAL_0]]
;
entry:
Expand Down

0 comments on commit 4eaa728

Please sign in to comment.