|
| 1 | +//===- JumpTableToSwitch.cpp ----------------------------------------------===// |
| 2 | +// |
| 3 | +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. |
| 4 | +// See https://llvm.org/LICENSE.txt for license information. |
| 5 | +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception |
| 6 | +// |
| 7 | +//===----------------------------------------------------------------------===// |
| 8 | + |
| 9 | +#include "llvm/Transforms/Scalar/JumpTableToSwitch.h" |
| 10 | +#include "llvm/ADT/SmallVector.h" |
| 11 | +#include "llvm/Analysis/ConstantFolding.h" |
| 12 | +#include "llvm/Analysis/DomTreeUpdater.h" |
| 13 | +#include "llvm/Analysis/OptimizationRemarkEmitter.h" |
| 14 | +#include "llvm/Analysis/PostDominators.h" |
| 15 | +#include "llvm/IR/IRBuilder.h" |
| 16 | +#include "llvm/Support/CommandLine.h" |
| 17 | +#include "llvm/Support/Debug.h" |
| 18 | +#include "llvm/Transforms/Utils/BasicBlockUtils.h" |
| 19 | + |
| 20 | +using namespace llvm; |
| 21 | + |
| 22 | +// TODO: Consider adding a cost model for profitability analysis of this |
| 23 | +// transformation. Currently we replace a jump table with a switch if all the |
| 24 | +// functions in the jump table are smaller than the provided threshold. |
| 25 | +static cl::opt<unsigned> |
| 26 | + JumpTableSizeThreshold("jump-table-to-switch-size-threshold", cl::Hidden, |
| 27 | + cl::desc("Only split jump tables with size less or " |
| 28 | + "equal than JumpTableSizeThreshold."), |
| 29 | + cl::init(10)); |
| 30 | + |
| 31 | +static cl::opt<unsigned> FunctionSizeThreshold( |
| 32 | + "jump-table-to-switch-function-size-threshold", cl::Hidden, |
| 33 | + cl::desc("Only split jump tables containing functions whose sizes are less " |
| 34 | + "or equal than this threshold."), |
| 35 | + cl::init(50)); |
| 36 | + |
| 37 | +#define DEBUG_TYPE "jump-table-to-switch" |
| 38 | + |
| 39 | +namespace { |
| 40 | +struct JumpTableTy { |
| 41 | + Value *Index; |
| 42 | + SmallVector<Function *, 10> Funcs; |
| 43 | +}; |
| 44 | +} // anonymous namespace |
| 45 | + |
| 46 | +static std::optional<JumpTableTy> parseJumpTable(GetElementPtrInst *GEP, |
| 47 | + PointerType *PtrTy) { |
| 48 | + Constant *Ptr = dyn_cast<Constant>(GEP->getPointerOperand()); |
| 49 | + if (!Ptr) |
| 50 | + return std::nullopt; |
| 51 | + |
| 52 | + GlobalVariable *GV = dyn_cast<GlobalVariable>(Ptr); |
| 53 | + if (!GV || !GV->isConstant() || !GV->hasDefinitiveInitializer()) |
| 54 | + return std::nullopt; |
| 55 | + |
| 56 | + Function &F = *GEP->getParent()->getParent(); |
| 57 | + const DataLayout &DL = F.getParent()->getDataLayout(); |
| 58 | + const unsigned BitWidth = |
| 59 | + DL.getIndexSizeInBits(GEP->getPointerAddressSpace()); |
| 60 | + MapVector<Value *, APInt> VariableOffsets; |
| 61 | + APInt ConstantOffset(BitWidth, 0); |
| 62 | + if (!GEP->collectOffset(DL, BitWidth, VariableOffsets, ConstantOffset)) |
| 63 | + return std::nullopt; |
| 64 | + if (VariableOffsets.size() != 1) |
| 65 | + return std::nullopt; |
| 66 | + // TODO: consider supporting more general patterns |
| 67 | + if (!ConstantOffset.isZero()) |
| 68 | + return std::nullopt; |
| 69 | + APInt StrideBytes = VariableOffsets.front().second; |
| 70 | + const uint64_t JumpTableSizeBytes = DL.getTypeAllocSize(GV->getValueType()); |
| 71 | + if (JumpTableSizeBytes % StrideBytes.getZExtValue() != 0) |
| 72 | + return std::nullopt; |
| 73 | + const uint64_t N = JumpTableSizeBytes / StrideBytes.getZExtValue(); |
| 74 | + if (N > JumpTableSizeThreshold) |
| 75 | + return std::nullopt; |
| 76 | + |
| 77 | + JumpTableTy JumpTable; |
| 78 | + JumpTable.Index = VariableOffsets.front().first; |
| 79 | + JumpTable.Funcs.reserve(N); |
| 80 | + for (uint64_t Index = 0; Index < N; ++Index) { |
| 81 | + // ConstantOffset is zero. |
| 82 | + APInt Offset = Index * StrideBytes; |
| 83 | + Constant *C = ConstantFoldLoadFromConst( |
| 84 | + cast<Constant>(GV->getInitializer()), PtrTy, Offset, DL); |
| 85 | + auto *Func = dyn_cast_or_null<Function>(C); |
| 86 | + if (!Func || Func->isDeclaration() || |
| 87 | + Func->getInstructionCount() > FunctionSizeThreshold) |
| 88 | + return std::nullopt; |
| 89 | + JumpTable.Funcs.push_back(Func); |
| 90 | + } |
| 91 | + return JumpTable; |
| 92 | +} |
| 93 | + |
| 94 | +static BasicBlock *expandToSwitch(CallBase *CB, const JumpTableTy &JT, |
| 95 | + DomTreeUpdater &DTU, |
| 96 | + OptimizationRemarkEmitter &ORE) { |
| 97 | + const bool IsVoid = CB->getType() == Type::getVoidTy(CB->getContext()); |
| 98 | + |
| 99 | + SmallVector<DominatorTree::UpdateType, 8> DTUpdates; |
| 100 | + BasicBlock *BB = CB->getParent(); |
| 101 | + BasicBlock *Tail = SplitBlock(BB, CB, &DTU, nullptr, nullptr, |
| 102 | + BB->getName() + Twine(".tail")); |
| 103 | + DTUpdates.push_back({DominatorTree::Delete, BB, Tail}); |
| 104 | + BB->getTerminator()->eraseFromParent(); |
| 105 | + |
| 106 | + Function &F = *BB->getParent(); |
| 107 | + BasicBlock *BBUnreachable = BasicBlock::Create( |
| 108 | + F.getContext(), "default.switch.case.unreachable", &F, Tail); |
| 109 | + IRBuilder<> BuilderUnreachable(BBUnreachable); |
| 110 | + BuilderUnreachable.CreateUnreachable(); |
| 111 | + |
| 112 | + IRBuilder<> Builder(BB); |
| 113 | + SwitchInst *Switch = Builder.CreateSwitch(JT.Index, BBUnreachable); |
| 114 | + DTUpdates.push_back({DominatorTree::Insert, BB, BBUnreachable}); |
| 115 | + |
| 116 | + IRBuilder<> BuilderTail(CB); |
| 117 | + PHINode *PHI = |
| 118 | + IsVoid ? nullptr : BuilderTail.CreatePHI(CB->getType(), JT.Funcs.size()); |
| 119 | + |
| 120 | + for (auto [Index, Func] : llvm::enumerate(JT.Funcs)) { |
| 121 | + BasicBlock *B = BasicBlock::Create(Func->getContext(), |
| 122 | + "call." + Twine(Index), &F, Tail); |
| 123 | + DTUpdates.push_back({DominatorTree::Insert, BB, B}); |
| 124 | + DTUpdates.push_back({DominatorTree::Insert, B, Tail}); |
| 125 | + |
| 126 | + CallBase *Call = cast<CallBase>(CB->clone()); |
| 127 | + Call->setCalledFunction(Func); |
| 128 | + Call->insertInto(B, B->end()); |
| 129 | + Switch->addCase( |
| 130 | + cast<ConstantInt>(ConstantInt::get(JT.Index->getType(), Index)), B); |
| 131 | + BranchInst::Create(Tail, B); |
| 132 | + if (PHI) |
| 133 | + PHI->addIncoming(Call, B); |
| 134 | + } |
| 135 | + DTU.applyUpdates(DTUpdates); |
| 136 | + ORE.emit([&]() { |
| 137 | + return OptimizationRemark(DEBUG_TYPE, "ReplacedJumpTableWithSwitch", CB) |
| 138 | + << "expanded indirect call into switch"; |
| 139 | + }); |
| 140 | + if (PHI) |
| 141 | + CB->replaceAllUsesWith(PHI); |
| 142 | + CB->eraseFromParent(); |
| 143 | + return Tail; |
| 144 | +} |
| 145 | + |
| 146 | +PreservedAnalyses JumpTableToSwitchPass::run(Function &F, |
| 147 | + FunctionAnalysisManager &AM) { |
| 148 | + OptimizationRemarkEmitter &ORE = |
| 149 | + AM.getResult<OptimizationRemarkEmitterAnalysis>(F); |
| 150 | + DominatorTree *DT = AM.getCachedResult<DominatorTreeAnalysis>(F); |
| 151 | + PostDominatorTree *PDT = AM. getCachedResult<PostDominatorTreeAnalysis>(F); |
| 152 | + DomTreeUpdater DTU(DT, PDT, DomTreeUpdater::UpdateStrategy::Lazy); |
| 153 | + bool Changed = false; |
| 154 | + for (BasicBlock &BB : make_early_inc_range(F)) { |
| 155 | + BasicBlock *CurrentBB = &BB; |
| 156 | + while (CurrentBB) { |
| 157 | + BasicBlock *SplittedOutTail = nullptr; |
| 158 | + for (Instruction &I : make_early_inc_range(*CurrentBB)) { |
| 159 | + auto *Call = dyn_cast<CallInst>(&I); |
| 160 | + if (!Call || Call->getCalledFunction() || Call->isMustTailCall()) |
| 161 | + continue; |
| 162 | + auto *L = dyn_cast<LoadInst>(Call->getCalledOperand()); |
| 163 | + // Skip atomic or volatile loads. |
| 164 | + if (!L || !L->isSimple()) |
| 165 | + continue; |
| 166 | + auto *GEP = dyn_cast<GetElementPtrInst>(L->getPointerOperand()); |
| 167 | + if (!GEP) |
| 168 | + continue; |
| 169 | + auto *PtrTy = dyn_cast<PointerType>(L-> getType()); |
| 170 | + assert(PtrTy && "call operand must be a pointer"); |
| 171 | + std::optional<JumpTableTy> JumpTable = parseJumpTable(GEP, PtrTy); |
| 172 | + if (!JumpTable) |
| 173 | + continue; |
| 174 | + SplittedOutTail = expandToSwitch(Call, *JumpTable, DTU, ORE); |
| 175 | + Changed = true; |
| 176 | + break; |
| 177 | + } |
| 178 | + CurrentBB = SplittedOutTail ? SplittedOutTail : nullptr; |
| 179 | + } |
| 180 | + } |
| 181 | + |
| 182 | + if (!Changed) |
| 183 | + return PreservedAnalyses::all(); |
| 184 | + |
| 185 | + PreservedAnalyses PA; |
| 186 | + if (DT) |
| 187 | + PA.preserve<DominatorTreeAnalysis>(); |
| 188 | + if (PDT) |
| 189 | + PA. preserve<PostDominatorTreeAnalysis>(); |
| 190 | + return PA; |
| 191 | +} |
0 commit comments