1 //=- WebAssemblyFixIrreducibleControlFlow.cpp - Fix irreducible control flow -//
3 // The LLVM Compiler Infrastructure
5 // This file is distributed under the University of Illinois Open Source
6 // License. See LICENSE.TXT for details.
8 //===----------------------------------------------------------------------===//
11 /// \brief This file implements a pass that transforms irreducible control flow
12 /// into reducible control flow. Irreducible control flow means multiple-entry
13 /// loops; they appear as CFG cycles that are not recorded in MachineLoopInfo
14 /// due to being unnatural.
16 /// Note that LLVM has a generic pass that lowers irreducible control flow, but
17 /// it linearizes control flow, turning diamonds into two triangles, which is
18 /// both unnecessary and undesirable for WebAssembly.
20 /// TODO: The transformation implemented here handles all irreducible control
21 /// flow, without exponential code-size expansion, though it does so by creating
22 /// inefficient code in many cases. Ideally, we should add other
23 /// transformations, including code-duplicating cases, which can be more
24 /// efficient in common cases, and they can fall back to this conservative
25 /// implementation as needed.
27 //===----------------------------------------------------------------------===//
29 #include "WebAssembly.h"
30 #include "MCTargetDesc/WebAssemblyMCTargetDesc.h"
31 #include "WebAssemblyMachineFunctionInfo.h"
32 #include "WebAssemblySubtarget.h"
33 #include "llvm/ADT/PriorityQueue.h"
34 #include "llvm/ADT/SCCIterator.h"
35 #include "llvm/ADT/SetVector.h"
36 #include "llvm/CodeGen/MachineDominators.h"
37 #include "llvm/CodeGen/MachineFunction.h"
38 #include "llvm/CodeGen/MachineInstrBuilder.h"
39 #include "llvm/CodeGen/MachineLoopInfo.h"
40 #include "llvm/CodeGen/MachineRegisterInfo.h"
41 #include "llvm/CodeGen/Passes.h"
42 #include "llvm/Support/Debug.h"
43 #include "llvm/Support/raw_ostream.h"
46 #define DEBUG_TYPE "wasm-fix-irreducible-control-flow"
49 class WebAssemblyFixIrreducibleControlFlow final : public MachineFunctionPass {
50 StringRef getPassName() const override {
51 return "WebAssembly Fix Irreducible Control Flow";
54 void getAnalysisUsage(AnalysisUsage &AU) const override {
56 AU.addRequired<MachineDominatorTree>();
57 AU.addPreserved<MachineDominatorTree>();
58 AU.addRequired<MachineLoopInfo>();
59 AU.addPreserved<MachineLoopInfo>();
60 MachineFunctionPass::getAnalysisUsage(AU);
63 bool runOnMachineFunction(MachineFunction &MF) override;
65 bool VisitLoop(MachineFunction &MF, MachineLoopInfo &MLI, MachineLoop *Loop);
68 static char ID; // Pass identification, replacement for typeid
69 WebAssemblyFixIrreducibleControlFlow() : MachineFunctionPass(ID) {}
71 } // end anonymous namespace
73 char WebAssemblyFixIrreducibleControlFlow::ID = 0;
74 FunctionPass *llvm::createWebAssemblyFixIrreducibleControlFlow() {
75 return new WebAssemblyFixIrreducibleControlFlow();
80 /// A utility for walking the blocks of a loop, handling a nested inner
81 /// loop as a monolithic conceptual block.
83 MachineBasicBlock *Block;
84 SmallVector<MachineBasicBlock *, 2> Preds;
85 SmallVector<MachineBasicBlock *, 2> Succs;
88 explicit MetaBlock(MachineBasicBlock *MBB)
89 : Block(MBB), Preds(MBB->pred_begin(), MBB->pred_end()),
90 Succs(MBB->succ_begin(), MBB->succ_end()) {}
92 explicit MetaBlock(MachineLoop *Loop) : Block(Loop->getHeader()) {
93 Loop->getExitBlocks(Succs);
94 for (MachineBasicBlock *Pred : Block->predecessors())
95 if (!Loop->contains(Pred))
96 Preds.push_back(Pred);
99 MachineBasicBlock *getBlock() const { return Block; }
101 const SmallVectorImpl<MachineBasicBlock *> &predecessors() const {
104 const SmallVectorImpl<MachineBasicBlock *> &successors() const {
108 bool operator==(const MetaBlock &MBB) { return Block == MBB.Block; }
109 bool operator!=(const MetaBlock &MBB) { return Block != MBB.Block; }
112 class SuccessorList final : public MetaBlock {
117 explicit SuccessorList(MachineBasicBlock *MBB)
118 : MetaBlock(MBB), Index(0), Num(successors().size()) {}
120 explicit SuccessorList(MachineLoop *Loop)
121 : MetaBlock(Loop), Index(0), Num(successors().size()) {}
123 bool HasNext() const { return Index != Num; }
125 MachineBasicBlock *Next() {
127 return successors()[Index++];
131 } // end anonymous namespace
133 bool WebAssemblyFixIrreducibleControlFlow::VisitLoop(MachineFunction &MF,
134 MachineLoopInfo &MLI,
136 MachineBasicBlock *Header = Loop ? Loop->getHeader() : &*MF.begin();
137 SetVector<MachineBasicBlock *> RewriteSuccs;
139 // DFS through Loop's body, looking for for irreducible control flow. Loop is
140 // natural, and we stay in its body, and we treat any nested loops
141 // monolithically, so any cycles we encounter indicate irreducibility.
142 SmallPtrSet<MachineBasicBlock *, 8> OnStack;
143 SmallPtrSet<MachineBasicBlock *, 8> Visited;
144 SmallVector<SuccessorList, 4> LoopWorklist;
145 LoopWorklist.push_back(SuccessorList(Header));
146 OnStack.insert(Header);
147 Visited.insert(Header);
148 while (!LoopWorklist.empty()) {
149 SuccessorList &Top = LoopWorklist.back();
151 MachineBasicBlock *Next = Top.Next();
152 if (Next == Header || (Loop && !Loop->contains(Next)))
154 if (LLVM_LIKELY(OnStack.insert(Next).second)) {
155 if (!Visited.insert(Next).second) {
159 MachineLoop *InnerLoop = MLI.getLoopFor(Next);
160 if (InnerLoop != Loop)
161 LoopWorklist.push_back(SuccessorList(InnerLoop));
163 LoopWorklist.push_back(SuccessorList(Next));
165 RewriteSuccs.insert(Top.getBlock());
169 OnStack.erase(Top.getBlock());
170 LoopWorklist.pop_back();
173 // Most likely, we didn't find any irreducible control flow.
174 if (LLVM_LIKELY(RewriteSuccs.empty()))
177 DEBUG(dbgs() << "Irreducible control flow detected!\n");
179 // Ok. We have irreducible control flow! Create a dispatch block which will
180 // contains a jump table to any block in the problematic set of blocks.
181 MachineBasicBlock *Dispatch = MF.CreateMachineBasicBlock();
182 MF.insert(MF.end(), Dispatch);
183 MLI.changeLoopFor(Dispatch, Loop);
185 // Add the jump table.
186 const auto &TII = *MF.getSubtarget<WebAssemblySubtarget>().getInstrInfo();
187 MachineInstrBuilder MIB = BuildMI(*Dispatch, Dispatch->end(), DebugLoc(),
188 TII.get(WebAssembly::BR_TABLE_I32));
190 // Add the register which will be used to tell the jump table which block to
192 MachineRegisterInfo &MRI = MF.getRegInfo();
193 unsigned Reg = MRI.createVirtualRegister(&WebAssembly::I32RegClass);
196 // Collect all the blocks which need to have their successors rewritten,
197 // add the successors to the jump table, and remember their index.
198 DenseMap<MachineBasicBlock *, unsigned> Indices;
199 SmallVector<MachineBasicBlock *, 4> SuccWorklist(RewriteSuccs.begin(),
201 while (!SuccWorklist.empty()) {
202 MachineBasicBlock *MBB = SuccWorklist.pop_back_val();
203 auto Pair = Indices.insert(std::make_pair(MBB, 0));
207 unsigned Index = MIB.getInstr()->getNumExplicitOperands() - 1;
208 DEBUG(dbgs() << "MBB#" << MBB->getNumber() << " has index " << Index
211 Pair.first->second = Index;
212 for (auto Pred : MBB->predecessors())
213 RewriteSuccs.insert(Pred);
216 Dispatch->addSuccessor(MBB);
219 for (auto *Succ : Meta.successors())
220 if (Succ != Header && (!Loop || Loop->contains(Succ)))
221 SuccWorklist.push_back(Succ);
224 // Rewrite the problematic successors for every block in RewriteSuccs.
225 // For simplicity, we just introduce a new block for every edge we need to
226 // rewrite. Fancier things are possible.
227 for (MachineBasicBlock *MBB : RewriteSuccs) {
228 DenseMap<MachineBasicBlock *, MachineBasicBlock *> Map;
229 for (auto *Succ : MBB->successors()) {
230 if (!Indices.count(Succ))
233 MachineBasicBlock *Split = MF.CreateMachineBasicBlock();
234 MF.insert(MBB->isLayoutSuccessor(Succ) ? MachineFunction::iterator(Succ)
237 MLI.changeLoopFor(Split, Loop);
239 // Set the jump table's register of the index of the block we wish to
240 // jump to, and jump to the jump table.
241 BuildMI(*Split, Split->end(), DebugLoc(), TII.get(WebAssembly::CONST_I32),
243 .addImm(Indices[Succ]);
244 BuildMI(*Split, Split->end(), DebugLoc(), TII.get(WebAssembly::BR))
246 Split->addSuccessor(Dispatch);
249 // Remap the terminator operands and the successor list.
250 for (MachineInstr &Term : MBB->terminators())
251 for (auto &Op : Term.explicit_uses())
252 if (Op.isMBB() && Indices.count(Op.getMBB()))
253 Op.setMBB(Map[Op.getMBB()]);
254 for (auto Rewrite : Map)
255 MBB->replaceSuccessor(Rewrite.first, Rewrite.second);
258 // Create a fake default label, because br_table requires one.
259 MIB.addMBB(MIB.getInstr()
260 ->getOperand(MIB.getInstr()->getNumExplicitOperands() - 1)
266 bool WebAssemblyFixIrreducibleControlFlow::runOnMachineFunction(
267 MachineFunction &MF) {
268 DEBUG(dbgs() << "********** Fixing Irreducible Control Flow **********\n"
269 "********** Function: "
270 << MF.getName() << '\n');
272 bool Changed = false;
273 auto &MLI = getAnalysis<MachineLoopInfo>();
275 // Visit the function body, which is identified as a null loop.
276 Changed |= VisitLoop(MF, MLI, nullptr);
278 // Visit all the loops.
279 SmallVector<MachineLoop *, 8> Worklist(MLI.begin(), MLI.end());
280 while (!Worklist.empty()) {
281 MachineLoop *CurLoop = Worklist.pop_back_val();
282 Worklist.append(CurLoop->begin(), CurLoop->end());
283 Changed |= VisitLoop(MF, MLI, CurLoop);
286 // If we made any changes, completely recompute everything.
287 if (LLVM_UNLIKELY(Changed)) {
288 DEBUG(dbgs() << "Recomputing dominators and loops.\n");
289 MF.getRegInfo().invalidateLiveness();
291 getAnalysis<MachineDominatorTree>().runOnMachineFunction(MF);
292 MLI.runOnMachineFunction(MF);