1//===- IndirectBrExpandPass.cpp - Expand indirectbr to switch -------------===//
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/// \file
9///
10/// Implements an expansion pass to turn `indirectbr` instructions in the IR
11/// into `switch` instructions. This works by enumerating the basic blocks in
12/// a dense range of integers, replacing each `blockaddr` constant with the
13/// corresponding integer constant, and then building a switch that maps from
14/// the integers to the actual blocks. All of the indirectbr instructions in the
15/// function are redirected to this common switch.
16///
17/// While this is generically useful if a target is unable to codegen
18/// `indirectbr` natively, it is primarily useful when there is some desire to
19/// get the builtin non-jump-table lowering of a switch even when the input
20/// source contained an explicit indirect branch construct.
21///
22/// Note that it doesn't make any sense to enable this pass unless a target also
23/// disables jump-table lowering of switches. Doing that is likely to pessimize
24/// the code.
25///
26//===----------------------------------------------------------------------===//
27
28#include "llvm/ADT/STLExtras.h"
29#include "llvm/ADT/Sequence.h"
30#include "llvm/ADT/SmallVector.h"
31#include "llvm/Analysis/BlockFrequencyInfo.h"
32#include "llvm/Analysis/DomTreeUpdater.h"
33#include "llvm/Analysis/LazyBlockFrequencyInfo.h"
34#include "llvm/CodeGen/IndirectBrExpand.h"
35#include "llvm/CodeGen/TargetPassConfig.h"
36#include "llvm/CodeGen/TargetSubtargetInfo.h"
37#include "llvm/IR/BasicBlock.h"
38#include "llvm/IR/Constants.h"
39#include "llvm/IR/Dominators.h"
40#include "llvm/IR/Function.h"
41#include "llvm/IR/Instructions.h"
42#include "llvm/IR/ProfDataUtils.h"
43#include "llvm/InitializePasses.h"
44#include "llvm/Pass.h"
45#include "llvm/Support/CodeGen.h"
46#include "llvm/Support/ErrorHandling.h"
47#include "llvm/Support/ScaledNumber.h"
48#include "llvm/Target/TargetMachine.h"
49#include <optional>
50
51using namespace llvm;
52
53#define DEBUG_TYPE "indirectbr-expand"
54
55namespace llvm {
56extern cl::opt<bool> ProfcheckDisableMetadataFixes;
57} // namespace llvm
58
59namespace {
60
61class IndirectBrExpandLegacyPass : public FunctionPass {
62 CodeGenOptLevel OptLevel;
63
64public:
65 static char ID; // Pass identification, replacement for typeid
66
67 IndirectBrExpandLegacyPass(CodeGenOptLevel OptLevel)
68 : FunctionPass(ID), OptLevel(OptLevel) {}
69
70 IndirectBrExpandLegacyPass()
71 : IndirectBrExpandLegacyPass(CodeGenOptLevel::None) {}
72
73 void getAnalysisUsage(AnalysisUsage &AU) const override {
74 if (OptLevel != CodeGenOptLevel::None)
75 LazyBlockFrequencyInfoPass::getLazyBFIAnalysisUsage(AU);
76 AU.addPreserved<DominatorTreeWrapperPass>();
77 }
78
79 bool runOnFunction(Function &F) override;
80};
81
82} // end anonymous namespace
83
84static bool runImpl(Function &F, const TargetLowering *TLI, DomTreeUpdater *DTU,
85 function_ref<BlockFrequencyInfo *()> GetBFI,
86 bool PreserveProfile);
87
88PreservedAnalyses IndirectBrExpandPass::run(Function &F,
89 FunctionAnalysisManager &FAM) {
90 auto *STI = TM->getSubtargetImpl(F);
91 if (!STI->enableIndirectBrExpand())
92 return PreservedAnalyses::all();
93
94 auto *TLI = STI->getTargetLowering();
95 auto *DT = FAM.getCachedResult<DominatorTreeAnalysis>(IR&: F);
96 DomTreeUpdater DTU(DT, DomTreeUpdater::UpdateStrategy::Lazy);
97
98 bool Changed = runImpl(
99 F, TLI, DTU: DT ? &DTU : nullptr,
100 GetBFI: [&]() { return &FAM.getResult<BlockFrequencyAnalysis>(IR&: F); },
101 /*PreserveProfile=*/true);
102 if (!Changed)
103 return PreservedAnalyses::all();
104 PreservedAnalyses PA;
105 PA.preserve<DominatorTreeAnalysis>();
106 return PA;
107}
108
109char IndirectBrExpandLegacyPass::ID = 0;
110
111INITIALIZE_PASS_BEGIN(IndirectBrExpandLegacyPass, DEBUG_TYPE,
112 "Expand indirectbr instructions", false, false)
113INITIALIZE_PASS_DEPENDENCY(DominatorTreeWrapperPass)
114INITIALIZE_PASS_END(IndirectBrExpandLegacyPass, DEBUG_TYPE,
115 "Expand indirectbr instructions", false, false)
116
117FunctionPass *llvm::createIndirectBrExpandPass(CodeGenOptLevel OptLevel) {
118 return new IndirectBrExpandLegacyPass(OptLevel);
119}
120
121bool runImpl(Function &F, const TargetLowering *TLI, DomTreeUpdater *DTU,
122 function_ref<BlockFrequencyInfo *()> GetBFI,
123 bool PreserveProfile) {
124 auto &DL = F.getDataLayout();
125
126 SmallVector<IndirectBrInst *, 1> IndirectBrs;
127 SmallVector<uint64_t, 1> IndirectBrsBlockFrequencies;
128 SmallVector<uint64_t, 1> IndirectBrsBranchWeightSums;
129 bool SkipProfileUpdates = !PreserveProfile;
130 BlockFrequencyInfo *BFI = nullptr;
131
132 struct IndirectBrSuccessor {
133 // The index into the IndirectBrs, IndirectBrsBlockFrequencies, and
134 // IndirectBrsBranchWeightSums vectors.
135 size_t IndirectBrIndex = 0;
136 uint64_t SuccessorBranchWeight = 0;
137 };
138
139 // Set of all potential successors for indirectbr instructions.
140 DenseMap<const BasicBlock *, SmallVector<IndirectBrSuccessor>>
141 IndirectBrSuccToIndirectBr;
142
143 // Build a list of indirectbrs that we want to rewrite.
144 for (BasicBlock &BB : F)
145 if (auto *IBr = dyn_cast<IndirectBrInst>(Val: BB.getTerminator())) {
146 // Handle the degenerate case of no successors by replacing the indirectbr
147 // with unreachable as there is no successor available.
148 if (IBr->getNumSuccessors() == 0) {
149 (void)new UnreachableInst(F.getContext(), IBr->getIterator());
150 IBr->eraseFromParent();
151 continue;
152 }
153
154 IndirectBrs.push_back(Elt: IBr);
155 const size_t CurrentIndirectBrIndex = IndirectBrs.size() - 1;
156 for (const BasicBlock *SuccessorBB : IBr->successors())
157 IndirectBrSuccToIndirectBr.insert(KV: {SuccessorBB, {}});
158
159 if (SkipProfileUpdates)
160 continue;
161 if (!BFI)
162 BFI = GetBFI();
163 std::optional<uint64_t> BlockFrequency = BFI->getBlockProfileCount(BB: &BB);
164 if (!BlockFrequency.has_value()) {
165 SkipProfileUpdates = true;
166 continue;
167 }
168 IndirectBrsBlockFrequencies.push_back(Elt: *BlockFrequency);
169 SmallVector<uint32_t> IndirectBrBranchWeights;
170 bool HasBranchWeights =
171 extractBranchWeights(I: *IBr, Weights&: IndirectBrBranchWeights);
172 if (!HasBranchWeights) {
173 SkipProfileUpdates = true;
174 continue;
175 }
176 for (const auto [SuccessorBB, SuccessorBranchWeight] :
177 zip_equal(t: IBr->successors(), u&: IndirectBrBranchWeights))
178 IndirectBrSuccToIndirectBr[SuccessorBB].push_back(
179 Elt: {.IndirectBrIndex: CurrentIndirectBrIndex, .SuccessorBranchWeight: SuccessorBranchWeight});
180 IndirectBrsBranchWeightSums.push_back(Elt: sum_of(Range&: IndirectBrBranchWeights));
181 assert(IndirectBrsBranchWeightSums.size() == IndirectBrs.size() &&
182 "expected an identical number of blocks in both vectors");
183 }
184
185 if (IndirectBrs.empty())
186 return false;
187
188 // If we need to replace any indirectbrs we need to establish integer
189 // constants that will correspond to each of the basic blocks in the function
190 // whose address escapes. We do that here and rewrite all the blockaddress
191 // constants to just be those integer constants cast to a pointer type.
192 SmallVector<BasicBlock *, 4> BBs;
193 SmallVector<ScaledNumber<uint64_t>, 4> BBWeights;
194
195 for (BasicBlock &BB : F) {
196 // Skip blocks that aren't successors to an indirectbr we're going to
197 // rewrite.
198 auto IndirectBrSuccToIndirectBrIt = IndirectBrSuccToIndirectBr.find(Val: &BB);
199 if (IndirectBrSuccToIndirectBrIt == IndirectBrSuccToIndirectBr.end())
200 continue;
201
202 auto *BA = BlockAddress::lookup(BB: &BB);
203
204 // Skip if the constant was formed but ended up not being used (due to DCE
205 // or whatever).
206 if (!BA || !BA->isConstantUsed())
207 continue;
208
209 // Compute the index we want to use for this basic block. We can't use zero
210 // because null can be compared with block addresses.
211 int BBIndex = BBs.size() + 1;
212 BBs.push_back(Elt: &BB);
213
214 auto *ITy = cast<IntegerType>(Val: DL.getIntPtrType(BA->getType()));
215 ConstantInt *BBIndexC = ConstantInt::get(Ty: ITy, V: BBIndex);
216
217 // Now rewrite the blockaddress to an integer constant based on the index.
218 // FIXME: This part doesn't properly recognize other uses of blockaddress
219 // expressions, for instance, where they are used to pass labels to
220 // asm-goto. This part of the pass needs a rework.
221 BA->replaceAllUsesWith(V: ConstantExpr::getIntToPtr(C: BBIndexC, Ty: BA->getType()));
222
223 if (SkipProfileUpdates)
224 continue;
225 ScaledNumber<uint64_t> BranchWeightSumsProduct(1, 0);
226 for (uint64_t BranchWeightSum : IndirectBrsBranchWeightSums)
227 BranchWeightSumsProduct *= ScaledNumber<uint64_t>(BranchWeightSum, 0);
228 ScaledNumber<uint64_t> BlockWeight(0, 0);
229 for (const auto &[IndirectBrIndex, BlockBranchProbability] :
230 IndirectBrSuccToIndirectBrIt->second) {
231 // If the branch weight sum is zero, skip adding the block weight or
232 // otherwise we end up dividing by zero.
233 const uint64_t CurrentBranchWeightSum =
234 IndirectBrsBranchWeightSums[IndirectBrIndex];
235 if (CurrentBranchWeightSum == 0)
236 continue;
237 BlockWeight += ScaledNumber<uint64_t>(
238 IndirectBrsBlockFrequencies[IndirectBrIndex], 0) *
239 ScaledNumber<uint64_t>(BlockBranchProbability, 0) *
240 (BranchWeightSumsProduct /
241 ScaledNumber<uint64_t>(CurrentBranchWeightSum, 0));
242 }
243 BBWeights.push_back(Elt: BlockWeight);
244 }
245
246 if (BBs.empty()) {
247 // There are no blocks whose address is taken, so any indirectbr instruction
248 // cannot get a valid input and we can replace all of them with unreachable.
249 SmallVector<DominatorTree::UpdateType, 8> Updates;
250 if (DTU)
251 Updates.reserve(N: IndirectBrSuccToIndirectBr.size());
252 for (auto *IBr : IndirectBrs) {
253 if (DTU) {
254 for (BasicBlock *SuccBB : IBr->successors())
255 Updates.push_back(Elt: {DominatorTree::Delete, IBr->getParent(), SuccBB});
256 }
257 (void)new UnreachableInst(F.getContext(), IBr->getIterator());
258 IBr->eraseFromParent();
259 }
260 if (DTU) {
261 assert(Updates.size() == IndirectBrSuccToIndirectBr.size() &&
262 "Got unexpected update count.");
263 DTU->applyUpdates(Updates);
264 }
265 return true;
266 }
267
268 BasicBlock *SwitchBB;
269 Value *SwitchValue;
270
271 // Compute a common integer type across all the indirectbr instructions.
272 IntegerType *CommonITy = nullptr;
273 for (auto *IBr : IndirectBrs) {
274 auto *ITy =
275 cast<IntegerType>(Val: DL.getIntPtrType(IBr->getAddress()->getType()));
276 if (!CommonITy || ITy->getBitWidth() > CommonITy->getBitWidth())
277 CommonITy = ITy;
278 }
279
280 auto GetSwitchValue = [CommonITy](IndirectBrInst *IBr) {
281 return CastInst::CreatePointerCast(S: IBr->getAddress(), Ty: CommonITy,
282 Name: Twine(IBr->getAddress()->getName()) +
283 ".switch_cast",
284 InsertBefore: IBr->getIterator());
285 };
286
287 SmallVector<DominatorTree::UpdateType, 8> Updates;
288
289 if (IndirectBrs.size() == 1) {
290 // If we only have one indirectbr, we can just directly replace it within
291 // its block.
292 IndirectBrInst *IBr = IndirectBrs[0];
293 SwitchBB = IBr->getParent();
294 SwitchValue = GetSwitchValue(IBr);
295 if (DTU) {
296 Updates.reserve(N: IndirectBrSuccToIndirectBr.size());
297 for (BasicBlock *SuccBB : IBr->successors())
298 Updates.push_back(Elt: {DominatorTree::Delete, IBr->getParent(), SuccBB});
299 assert(Updates.size() == IndirectBrSuccToIndirectBr.size() &&
300 "Got unexpected update count.");
301 }
302 IBr->eraseFromParent();
303 } else {
304 // Otherwise we need to create a new block to hold the switch across BBs,
305 // jump to that block instead of each indirectbr, and phi together the
306 // values for the switch.
307 SwitchBB = BasicBlock::Create(Context&: F.getContext(), Name: "switch_bb", Parent: &F);
308 auto *SwitchPN = PHINode::Create(Ty: CommonITy, NumReservedValues: IndirectBrs.size(),
309 NameStr: "switch_value_phi", InsertBefore: SwitchBB);
310 SwitchValue = SwitchPN;
311
312 // Now replace the indirectbr instructions with direct branches to the
313 // switch block and fill out the PHI operands.
314 if (DTU)
315 Updates.reserve(N: IndirectBrs.size() +
316 2 * IndirectBrSuccToIndirectBr.size());
317 for (auto *IBr : IndirectBrs) {
318 SwitchPN->addIncoming(V: GetSwitchValue(IBr), BB: IBr->getParent());
319 UncondBrInst::Create(Target: SwitchBB, InsertBefore: IBr->getIterator());
320 if (DTU) {
321 Updates.push_back(Elt: {DominatorTree::Insert, IBr->getParent(), SwitchBB});
322 for (BasicBlock *SuccBB : IBr->successors())
323 Updates.push_back(Elt: {DominatorTree::Delete, IBr->getParent(), SuccBB});
324 }
325 IBr->eraseFromParent();
326 }
327 }
328
329 // Now build the switch in the block. The block will have no terminator
330 // already.
331 auto *SI = SwitchInst::Create(Value: SwitchValue, Default: BBs[0], NumCases: BBs.size(), InsertBefore: SwitchBB);
332
333 // Add a case for each block.
334 for (int i : llvm::seq<int>(Begin: 1, End: BBs.size()))
335 SI->addCase(OnVal: ConstantInt::get(Ty: CommonITy, V: i + 1), Dest: BBs[i]);
336
337 if (DTU) {
338 // If there were multiple indirectbr's, they may have common successors,
339 // but in the dominator tree, we only track unique edges.
340 SmallPtrSet<BasicBlock *, 8> UniqueSuccessors;
341 Updates.reserve(N: Updates.size() + BBs.size());
342 for (BasicBlock *BB : BBs) {
343 if (UniqueSuccessors.insert(Ptr: BB).second)
344 Updates.push_back(Elt: {DominatorTree::Insert, SwitchBB, BB});
345 }
346 DTU->applyUpdates(Updates);
347 }
348
349 if (SkipProfileUpdates || ProfcheckDisableMetadataFixes) {
350 setExplicitlyUnknownBranchWeightsIfProfiled(I&: *SI, DEBUG_TYPE);
351 return true;
352 }
353
354 // We need to convert the ScaledNumber weights (which might not be
355 // representable in 64 bits) back to normal 64 bit integers so we can apply
356 // them as metadata. They might not have the same scale though, so we find the
357 // max scale and then scale down any weights that have a scale less than the
358 // max scale. This ensures that all the weights have the same scale.
359 int16_t MaxScale = 0;
360 for (const ScaledNumber<uint64_t> &BBWeight : BBWeights)
361 MaxScale = std::max(a: MaxScale, b: BBWeight.getScale());
362 SmallVector<uint64_t, 4> ExtractedBBWeights;
363 ExtractedBBWeights.reserve(N: BBWeights.size());
364 for (ScaledNumber<uint64_t> &BBWeight : BBWeights) {
365 int16_t Shift = MaxScale - BBWeight.getScale();
366 assert(Shift >= 0 && "expected non-negative shift");
367 ExtractedBBWeights.push_back(Elt: BBWeight.getDigits() >> Shift);
368 }
369 setFittedBranchWeights(I&: *SI, Weights: ExtractedBBWeights, IsExpected: false);
370
371 return true;
372}
373
374bool IndirectBrExpandLegacyPass::runOnFunction(Function &F) {
375 auto *TPC = getAnalysisIfAvailable<TargetPassConfig>();
376 if (!TPC)
377 return false;
378
379 auto &TM = TPC->getTM<TargetMachine>();
380 auto &STI = *TM.getSubtargetImpl(F);
381 if (!STI.enableIndirectBrExpand())
382 return false;
383 auto *TLI = STI.getTargetLowering();
384
385 std::optional<DomTreeUpdater> DTU;
386 if (auto *DTWP = getAnalysisIfAvailable<DominatorTreeWrapperPass>())
387 DTU.emplace(args&: DTWP->getDomTree(), args: DomTreeUpdater::UpdateStrategy::Lazy);
388
389 return runImpl(
390 F, TLI, DTU: DTU ? &*DTU : nullptr,
391 GetBFI: [&]() { return &getAnalysis<LazyBlockFrequencyInfoPass>().getBFI(); },
392 PreserveProfile: OptLevel != CodeGenOptLevel::None);
393}
394