1//=- WebAssemblyFixIrreducibleControlFlow.cpp - Fix irreducible control flow -//
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/// \file
10/// This file implements a pass that removes irreducible control flow.
11/// Irreducible control flow means multiple-entry loops, which this pass
12/// transforms to have a single entry.
13///
14/// Note that LLVM has a generic pass that lowers irreducible control flow, but
15/// it linearizes control flow, turning diamonds into two triangles, which is
16/// both unnecessary and undesirable for WebAssembly.
17///
18/// The big picture: We recursively process each "region", defined as a group
19/// of blocks with a single entry and no branches back to that entry. A region
20/// may be the entire function body, or the inner part of a loop, i.e., the
21/// loop's body without branches back to the loop entry. In each region we
22/// identify all the strongly-connected components (SCCs). We fix up multi-entry
23/// loops (SCCs) by adding a new block that can dispatch to each of the loop
24/// entries, based on the value of a label "helper" variable, and we replace
25/// direct branches to the entries with assignments to the label variable and a
26/// branch to the dispatch block. Then the dispatch block is the single entry in
27/// the loop containing the previous multiple entries. Each time we fix some
28/// irreducibility, we recalculate the SCCs. After ensuring all the SCCs in a
29/// region are reducible, we recurse into them. The total time complexity of
30/// this pass is roughly:
31/// O((NumBlocks + NumEdges) * (NumNestedLoops + NumIrreducibleLoops))
32///
33/// This pass is similar to what the Relooper [1] does. Both identify looping
34/// code that requires multiple entries, and resolve it in a similar way (in
35/// Relooper terminology, we implement a Multiple shape in a Loop shape). Note
36/// also that like the Relooper, we implement a "minimal" intervention: we only
37/// use the "label" helper for the blocks we absolutely must and no others. We
38/// also prioritize code size and do not duplicate code in order to resolve
39/// irreducibility. The graph algorithms for finding loops and entries and so
40/// forth are also similar to the Relooper. The main differences between this
41/// pass and the Relooper are:
42///
43/// * We just care about irreducibility, so we just look at loops.
44/// * The Relooper emits structured control flow (with ifs etc.), while we
45/// emit a CFG.
46///
47/// [1] Alon Zakai. 2011. Emscripten: an LLVM-to-JavaScript compiler. In
48/// Proceedings of the ACM international conference companion on Object oriented
49/// programming systems languages and applications companion (SPLASH '11). ACM,
50/// New York, NY, USA, 301-312. DOI=10.1145/2048147.2048224
51/// http://doi.acm.org/10.1145/2048147.2048224
52///
53//===----------------------------------------------------------------------===//
54
55#include "MCTargetDesc/WebAssemblyMCTargetDesc.h"
56#include "WebAssembly.h"
57#include "WebAssemblySubtarget.h"
58#include "llvm/ADT/SCCIterator.h"
59#include "llvm/CodeGen/MachineBasicBlock.h"
60#include "llvm/CodeGen/MachineFunction.h"
61#include "llvm/CodeGen/MachineFunctionAnalysisManager.h"
62#include "llvm/CodeGen/MachineFunctionPass.h"
63#include "llvm/CodeGen/MachineInstrBuilder.h"
64#include "llvm/CodeGen/MachinePassManager.h"
65#include "llvm/IR/Analysis.h"
66#include "llvm/Support/Debug.h"
67#include <limits>
68using namespace llvm;
69
70#define DEBUG_TYPE "wasm-fix-irreducible-control-flow"
71
72namespace {
73
74using BlockVector = SmallVector<MachineBasicBlock *, 4>;
75using BlockSet = SmallPtrSet<MachineBasicBlock *, 4>;
76
77static BlockVector getSortedEntries(const BlockSet &Entries) {
78 BlockVector SortedEntries(Entries.begin(), Entries.end());
79 llvm::sort(C&: SortedEntries,
80 Comp: [](const MachineBasicBlock *A, const MachineBasicBlock *B) {
81 auto ANum = A->getNumber();
82 auto BNum = B->getNumber();
83 return ANum < BNum;
84 });
85 return SortedEntries;
86}
87
88struct ReachabilityNode {
89 MachineBasicBlock *MBB;
90 SmallVector<ReachabilityNode *, 4> Succs;
91 unsigned SCCId = std::numeric_limits<unsigned>::max();
92};
93
94// Analyzes the SCC (strongly-connected component) structure in a region.
95// Ignores branches to blocks outside of the region, and ignores branches to the
96// region entry (for the case where the region is the inner part of a loop).
97class ReachabilityGraph {
98public:
99 ReachabilityGraph(MachineBasicBlock *Entry, const BlockSet &Blocks)
100 : Entry(Entry), Blocks(Blocks) {
101#ifndef NDEBUG
102 // The region must have a single entry.
103 for (auto *MBB : Blocks) {
104 if (MBB != Entry) {
105 for (auto *Pred : MBB->predecessors()) {
106 assert(inRegion(Pred));
107 }
108 }
109 }
110#endif
111 calculate();
112 }
113
114 // Get all blocks that are loop entries.
115 const BlockSet &getLoopEntries() const { return LoopEntries; }
116 const BlockSet &getLoopEntriesForSCC(unsigned SCCId) const {
117 return LoopEntriesBySCC[SCCId];
118 }
119
120 unsigned getSCCId(MachineBasicBlock *MBB) const {
121 return getNode(MBB)->SCCId;
122 }
123
124 friend struct GraphTraits<ReachabilityGraph *>;
125
126private:
127 MachineBasicBlock *Entry;
128 const BlockSet &Blocks;
129
130 BlockSet LoopEntries;
131 SmallVector<BlockSet, 0> LoopEntriesBySCC;
132
133 bool inRegion(MachineBasicBlock *MBB) const { return Blocks.count(Ptr: MBB); }
134
135 SmallVector<ReachabilityNode, 0> Nodes;
136 DenseMap<MachineBasicBlock *, ReachabilityNode *> MBBToNodeMap;
137
138 ReachabilityNode *getNode(MachineBasicBlock *MBB) const {
139 return MBBToNodeMap.at(Val: MBB);
140 }
141
142 void calculate();
143};
144} // end anonymous namespace
145
146namespace llvm {
147template <> struct GraphTraits<ReachabilityGraph *> {
148 using NodeRef = ReachabilityNode *;
149 using ChildIteratorType = SmallVectorImpl<NodeRef>::iterator;
150
151 static NodeRef getEntryNode(ReachabilityGraph *G) {
152 return G->getNode(MBB: G->Entry);
153 }
154
155 static inline ChildIteratorType child_begin(NodeRef N) {
156 return N->Succs.begin();
157 }
158
159 static inline ChildIteratorType child_end(NodeRef N) {
160 return N->Succs.end();
161 }
162};
163} // end namespace llvm
164
165namespace {
166
167void ReachabilityGraph::calculate() {
168 auto NumBlocks = Blocks.size();
169 Nodes.assign(NumElts: NumBlocks, Elt: {});
170
171 MBBToNodeMap.clear();
172 MBBToNodeMap.reserve(NumEntries: NumBlocks);
173
174 // Initialize mappings.
175 unsigned MBBIdx = 0;
176 for (auto *MBB : Blocks) {
177 auto &Node = Nodes[MBBIdx++];
178
179 Node.MBB = MBB;
180 MBBToNodeMap[MBB] = &Node;
181 }
182
183 // Add all relevant direct branches.
184 MBBIdx = 0;
185 for (auto *MBB : Blocks) {
186 auto &Node = Nodes[MBBIdx++];
187
188 for (auto *Succ : MBB->successors()) {
189 if (Succ != Entry && inRegion(MBB: Succ)) {
190 Node.Succs.push_back(Elt: getNode(MBB: Succ));
191 }
192 }
193 }
194
195 unsigned CurrSCCIdx = 0;
196 for (auto &SCC : make_range(x: scc_begin(G: this), y: scc_end(G: this))) {
197 LoopEntriesBySCC.push_back(Elt: {});
198 auto &SCCLoopEntries = LoopEntriesBySCC.back();
199
200 for (auto *Node : SCC) {
201 // Make sure nodes are only ever assigned one SCC
202 assert(Node->SCCId == std::numeric_limits<unsigned>::max());
203
204 Node->SCCId = CurrSCCIdx;
205 }
206
207 bool SelfLoop = false;
208 if (SCC.size() == 1) {
209 auto &Node = SCC[0];
210
211 for (auto *Succ : Node->Succs) {
212 if (Succ == Node) {
213 SelfLoop = true;
214 break;
215 }
216 }
217 }
218
219 // Blocks outside any (multi-block) loop will be isolated in their own
220 // single-element SCC. Thus blocks that are in a loop are those in
221 // multi-element SCCs or are self-looping.
222 if (SCC.size() > 1 || SelfLoop) {
223 // Find the loop entries - loop body blocks with predecessors outside
224 // their SCC
225 for (auto *Node : SCC) {
226 if (Node->MBB == Entry)
227 continue;
228
229 for (auto *Pred : Node->MBB->predecessors()) {
230 // This test is accurate despite not having assigned all nodes an SCC
231 // yet. We only care if a node has been assigned into this SCC or not.
232 if (getSCCId(MBB: Pred) != CurrSCCIdx) {
233 LoopEntries.insert(Ptr: Node->MBB);
234 SCCLoopEntries.insert(Ptr: Node->MBB);
235 }
236 }
237 }
238 }
239 ++CurrSCCIdx;
240 }
241
242#ifndef NDEBUG
243 // Make sure all nodes have been processed
244 for (auto &Node : Nodes) {
245 assert(Node.SCCId != std::numeric_limits<unsigned>::max());
246 }
247#endif
248}
249
250class WebAssemblyFixIrreducibleControlFlowLegacy final
251 : public MachineFunctionPass {
252 StringRef getPassName() const override {
253 return "WebAssembly Fix Irreducible Control Flow";
254 }
255
256 bool runOnMachineFunction(MachineFunction &MF) override;
257
258public:
259 static char ID; // Pass identification, replacement for typeid
260 WebAssemblyFixIrreducibleControlFlowLegacy() : MachineFunctionPass(ID) {}
261};
262
263// Given a set of entries to a single loop, create a single entry for that
264// loop by creating a dispatch block for them, routing control flow using
265// a helper variable. Also updates Blocks with any new blocks created, so
266// that we properly track all the blocks in the region. But this does not update
267// ReachabilityGraph; this will be updated in the caller of this function as
268// needed.
269void makeSingleEntryLoop(const BlockSet &Entries, BlockSet &Blocks,
270 MachineFunction &MF, const ReachabilityGraph &Graph) {
271 assert(Entries.size() >= 2);
272
273 // Sort the entries to ensure a deterministic build.
274 BlockVector SortedEntries = getSortedEntries(Entries);
275
276#ifndef NDEBUG
277 for (auto *Block : SortedEntries)
278 assert(Block->getNumber() != -1);
279 if (SortedEntries.size() > 1) {
280 for (auto I = SortedEntries.begin(), E = SortedEntries.end() - 1; I != E;
281 ++I) {
282 auto ANum = (*I)->getNumber();
283 auto BNum = (*(std::next(I)))->getNumber();
284 assert(ANum != BNum);
285 }
286 }
287#endif
288
289 // Create a dispatch block which will contain a jump table to the entries.
290 MachineBasicBlock *Dispatch = MF.CreateMachineBasicBlock();
291 MF.insert(MBBI: MF.end(), MBB: Dispatch);
292 Blocks.insert(Ptr: Dispatch);
293
294 // Add the jump table.
295 const auto &TII = *MF.getSubtarget<WebAssemblySubtarget>().getInstrInfo();
296 MachineInstrBuilder MIB =
297 BuildMI(BB: Dispatch, MIMD: DebugLoc(), MCID: TII.get(Opcode: WebAssembly::BR_TABLE_I32));
298
299 // Add the register which will be used to tell the jump table which block to
300 // jump to.
301 MachineRegisterInfo &MRI = MF.getRegInfo();
302 Register Reg = MRI.createVirtualRegister(RegClass: &WebAssembly::I32RegClass);
303 MIB.addReg(RegNo: Reg);
304
305 // Compute the indices in the superheader, one for each bad block, and
306 // add them as successors.
307 DenseMap<MachineBasicBlock *, unsigned> Indices;
308 for (auto *Entry : SortedEntries) {
309 auto Pair = Indices.try_emplace(Key: Entry);
310 assert(Pair.second);
311
312 unsigned Index = MIB.getInstr()->getNumExplicitOperands() - 1;
313 Pair.first->second = Index;
314
315 MIB.addMBB(MBB: Entry);
316 Dispatch->addSuccessor(Succ: Entry);
317 }
318
319 // Rewrite the problematic successors for every block that wants to reach
320 // the bad blocks. For simplicity, we just introduce a new block for every
321 // edge we need to rewrite. (Fancier things are possible.)
322
323 BlockVector AllPreds;
324 for (auto *Entry : SortedEntries) {
325 for (auto *Pred : Entry->predecessors()) {
326 if (Pred != Dispatch) {
327 AllPreds.push_back(Elt: Pred);
328 }
329 }
330 }
331
332 // This set stores predecessors within this loop.
333 DenseSet<MachineBasicBlock *> InLoop;
334 for (auto *Pred : AllPreds) {
335 auto PredSCCId = Graph.getSCCId(MBB: Pred);
336
337 for (auto *Entry : Pred->successors()) {
338 if (!Entries.count(Ptr: Entry))
339 continue;
340 if (Graph.getSCCId(MBB: Entry) == PredSCCId) {
341 InLoop.insert(V: Pred);
342 break;
343 }
344 }
345 }
346
347 // Record if each entry has a layout predecessor. This map stores
348 // <<loop entry, Predecessor is within the loop?>, layout predecessor>
349 DenseMap<PointerIntPair<MachineBasicBlock *, 1, bool>, MachineBasicBlock *>
350 EntryToLayoutPred;
351 for (auto *Pred : AllPreds) {
352 bool PredInLoop = InLoop.count(V: Pred);
353 for (auto *Entry : Pred->successors())
354 if (Entries.count(Ptr: Entry) && Pred->isLayoutSuccessor(MBB: Entry))
355 EntryToLayoutPred[{Entry, PredInLoop}] = Pred;
356 }
357
358 // We need to create at most two routing blocks per entry: one for
359 // predecessors outside the loop and one for predecessors inside the loop.
360 // This map stores
361 // <<loop entry, Predecessor is within the loop?>, routing block>
362 DenseMap<PointerIntPair<MachineBasicBlock *, 1, bool>, MachineBasicBlock *>
363 Map;
364 for (auto *Pred : AllPreds) {
365 bool PredInLoop = InLoop.count(V: Pred);
366 for (auto *Entry : Pred->successors()) {
367 if (!Entries.count(Ptr: Entry) || Map.count(Val: {Entry, PredInLoop}))
368 continue;
369 // If there exists a layout predecessor of this entry and this predecessor
370 // is not that, we rather create a routing block after that layout
371 // predecessor to save a branch.
372 if (auto *OtherPred = EntryToLayoutPred.lookup(Val: {Entry, PredInLoop}))
373 if (OtherPred != Pred)
374 continue;
375
376 // This is a successor we need to rewrite.
377 MachineBasicBlock *Routing = MF.CreateMachineBasicBlock();
378 MF.insert(MBBI: Pred->isLayoutSuccessor(MBB: Entry)
379 ? MachineFunction::iterator(Entry)
380 : MF.end(),
381 MBB: Routing);
382 Blocks.insert(Ptr: Routing);
383
384 // Set the jump table's register of the index of the block we wish to
385 // jump to, and jump to the jump table.
386 BuildMI(BB: Routing, MIMD: DebugLoc(), MCID: TII.get(Opcode: WebAssembly::CONST_I32), DestReg: Reg)
387 .addImm(Val: Indices[Entry]);
388 BuildMI(BB: Routing, MIMD: DebugLoc(), MCID: TII.get(Opcode: WebAssembly::BR)).addMBB(MBB: Dispatch);
389 Routing->addSuccessor(Succ: Dispatch);
390 Map[{Entry, PredInLoop}] = Routing;
391 }
392 }
393
394 for (auto *Pred : AllPreds) {
395 bool PredInLoop = InLoop.count(V: Pred);
396 // Remap the terminator operands and the successor list.
397 for (MachineInstr &Term : Pred->terminators())
398 for (auto &Op : Term.explicit_uses())
399 if (Op.isMBB() && Indices.count(Val: Op.getMBB()))
400 Op.setMBB(Map[{Op.getMBB(), PredInLoop}]);
401
402 for (auto *Succ : Pred->successors()) {
403 if (!Entries.count(Ptr: Succ))
404 continue;
405 auto *Routing = Map[{Succ, PredInLoop}];
406 Pred->replaceSuccessor(Old: Succ, New: Routing);
407 }
408 }
409
410 // Create a fake default label, because br_table requires one.
411 MIB.addMBB(MBB: MIB.getInstr()
412 ->getOperand(i: MIB.getInstr()->getNumExplicitOperands() - 1)
413 .getMBB());
414}
415
416bool processRegion(MachineBasicBlock *Entry, BlockSet &Blocks,
417 MachineFunction &MF) {
418 bool Changed = false;
419 // Remove irreducibility before processing child loops, which may take
420 // multiple iterations.
421 while (true) {
422 ReachabilityGraph Graph(Entry, Blocks);
423
424 bool FoundIrreducibility = false;
425
426 for (auto *LoopEntry : getSortedEntries(Entries: Graph.getLoopEntries())) {
427 // Find mutual entries - all entries which can reach this one, and
428 // are reached by it (that always includes LoopEntry itself). All mutual
429 // entries must be in the same SCC, so if we have more than one, then we
430 // have irreducible control flow.
431 //
432 // (Note that we need to sort the entries here, as otherwise the order can
433 // matter: being mutual is a symmetric relationship, and each set of
434 // mutuals will be handled properly no matter which we see first. However,
435 // there can be multiple disjoint sets of mutuals, and which we process
436 // first changes the output.)
437 //
438 // Note that irreducibility may involve inner loops, e.g. imagine A
439 // starts one loop, and it has B inside it which starts an inner loop.
440 // If we add a branch from all the way on the outside to B, then in a
441 // sense B is no longer an "inner" loop, semantically speaking. We will
442 // fix that irreducibility by adding a block that dispatches to either
443 // either A or B, so B will no longer be an inner loop in our output.
444 // (A fancier approach might try to keep it as such.)
445 //
446 // Note that we still need to recurse into inner loops later, to handle
447 // the case where the irreducibility is entirely nested - we would not
448 // be able to identify that at this point, since the enclosing loop is
449 // a group of blocks all of whom can reach each other. (We'll see the
450 // irreducibility after removing branches to the top of that enclosing
451 // loop.)
452 auto &MutualLoopEntries =
453 Graph.getLoopEntriesForSCC(SCCId: Graph.getSCCId(MBB: LoopEntry));
454
455 if (MutualLoopEntries.size() > 1) {
456 makeSingleEntryLoop(Entries: MutualLoopEntries, Blocks, MF, Graph);
457 FoundIrreducibility = true;
458 Changed = true;
459 break;
460 }
461 }
462
463 // Only go on to actually process the inner loops when we are done
464 // removing irreducible control flow and changing the graph. Modifying
465 // the graph as we go is possible, and that might let us avoid looking at
466 // the already-fixed loops again if we are careful, but all that is
467 // complex and bug-prone. Since irreducible loops are rare, just starting
468 // another iteration is best.
469 if (FoundIrreducibility) {
470 continue;
471 }
472
473 for (auto *LoopEntry : Graph.getLoopEntries()) {
474 BlockSet InnerBlocks;
475
476 auto EntrySCCId = Graph.getSCCId(MBB: LoopEntry);
477 for (auto *Block : Blocks) {
478 if (EntrySCCId == Graph.getSCCId(MBB: Block)) {
479 InnerBlocks.insert(Ptr: Block);
480 }
481 }
482
483 // Each of these calls to processRegion may change the graph, but are
484 // guaranteed not to interfere with each other. The only changes we make
485 // to the graph are to add blocks on the way to a loop entry. As the
486 // loops are disjoint, that means we may only alter branches that exit
487 // another loop, which are ignored when recursing into that other loop
488 // anyhow.
489 if (processRegion(Entry: LoopEntry, Blocks&: InnerBlocks, MF)) {
490 Changed = true;
491 }
492 }
493
494 return Changed;
495 }
496}
497
498} // end anonymous namespace
499
500char WebAssemblyFixIrreducibleControlFlowLegacy::ID = 0;
501INITIALIZE_PASS(WebAssemblyFixIrreducibleControlFlowLegacy, DEBUG_TYPE,
502 "Removes irreducible control flow", false, false)
503
504FunctionPass *llvm::createWebAssemblyFixIrreducibleControlFlowLegacyPass() {
505 return new WebAssemblyFixIrreducibleControlFlowLegacy();
506}
507
508// Test whether the given register has an ARGUMENT def.
509static bool hasArgumentDef(unsigned Reg, const MachineRegisterInfo &MRI) {
510 for (const auto &Def : MRI.def_instructions(Reg))
511 if (WebAssembly::isArgument(Opc: Def.getOpcode()))
512 return true;
513 return false;
514}
515
516// Add a register definition with IMPLICIT_DEFs for every register to cover for
517// register uses that don't have defs in every possible path.
518// TODO: This is fairly heavy-handed; find a better approach.
519static void addImplicitDefs(MachineFunction &MF) {
520 const MachineRegisterInfo &MRI = MF.getRegInfo();
521 const auto &TII = *MF.getSubtarget<WebAssemblySubtarget>().getInstrInfo();
522 MachineBasicBlock &Entry = *MF.begin();
523 for (unsigned I = 0, E = MRI.getNumVirtRegs(); I < E; ++I) {
524 Register Reg = Register::index2VirtReg(Index: I);
525
526 // Skip unused registers.
527 if (MRI.use_nodbg_empty(RegNo: Reg))
528 continue;
529
530 // Skip registers that have an ARGUMENT definition.
531 if (hasArgumentDef(Reg, MRI))
532 continue;
533
534 BuildMI(BB&: Entry, I: Entry.begin(), MIMD: DebugLoc(),
535 MCID: TII.get(Opcode: WebAssembly::IMPLICIT_DEF), DestReg: Reg);
536 }
537
538 // Move ARGUMENT_* instructions to the top of the entry block, so that their
539 // liveness reflects the fact that these really are live-in values.
540 for (MachineInstr &MI : llvm::make_early_inc_range(Range&: Entry)) {
541 if (WebAssembly::isArgument(Opc: MI.getOpcode())) {
542 MI.removeFromParent();
543 Entry.insert(I: Entry.begin(), MI: &MI);
544 }
545 }
546}
547
548static bool fixIrreducibleControlFlow(MachineFunction &MF) {
549 LLVM_DEBUG(dbgs() << "********** Fixing Irreducible Control Flow **********\n"
550 "********** Function: "
551 << MF.getName() << '\n');
552
553 // Start the recursive process on the entire function body.
554 BlockSet AllBlocks;
555 for (auto &MBB : MF) {
556 AllBlocks.insert(Ptr: &MBB);
557 }
558
559 if (LLVM_UNLIKELY(processRegion(&*MF.begin(), AllBlocks, MF))) {
560 // We rewrote part of the function; recompute relevant things.
561 MF.RenumberBlocks();
562 // Now we've inserted dispatch blocks, some register uses can have incoming
563 // paths without a def. For example, before this pass register %a was
564 // defined in BB1 and used in BB2, and there was only one path from BB1 and
565 // BB2. But if this pass inserts a dispatch block having multiple
566 // predecessors between the two BBs, now there are paths to BB2 without
567 // visiting BB1, and %a's use in BB2 is not dominated by its def. Adding
568 // IMPLICIT_DEFs to all regs is one simple way to fix it.
569 addImplicitDefs(MF);
570 return true;
571 }
572
573 return false;
574}
575
576bool WebAssemblyFixIrreducibleControlFlowLegacy::runOnMachineFunction(
577 MachineFunction &MF) {
578 return fixIrreducibleControlFlow(MF);
579}
580
581PreservedAnalyses WebAssemblyFixIrreducibleControlFlowPass::run(
582 MachineFunction &MF, MachineFunctionAnalysisManager &MFAM) {
583 return fixIrreducibleControlFlow(MF)
584 ? getMachineFunctionPassPreservedAnalyses()
585 : PreservedAnalyses::all();
586}
587