1 //=- WebAssemblyFixIrreducibleControlFlow.cpp - Fix irreducible control flow -// 2 // 3 // The LLVM Compiler Infrastructure 4 // 5 // This file is distributed under the University of Illinois Open Source 6 // License. See LICENSE.TXT for details. 7 // 8 //===----------------------------------------------------------------------===// 9 /// 10 /// \file 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. 15 /// 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. 19 /// 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. 26 /// 27 //===----------------------------------------------------------------------===// 28 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" 44 using namespace llvm; 45 46 #define DEBUG_TYPE "wasm-fix-irreducible-control-flow" 47 48 namespace { 49 class WebAssemblyFixIrreducibleControlFlow final : public MachineFunctionPass { 50 const char *getPassName() const override { 51 return "WebAssembly Fix Irreducible Control Flow"; 52 } 53 54 void getAnalysisUsage(AnalysisUsage &AU) const override { 55 AU.setPreservesCFG(); 56 AU.addRequired<MachineDominatorTree>(); 57 AU.addPreserved<MachineDominatorTree>(); 58 AU.addRequired<MachineLoopInfo>(); 59 AU.addPreserved<MachineLoopInfo>(); 60 MachineFunctionPass::getAnalysisUsage(AU); 61 } 62 63 bool runOnMachineFunction(MachineFunction &MF) override; 64 65 bool VisitLoop(MachineFunction &MF, MachineLoopInfo &MLI, MachineLoop *Loop); 66 67 public: 68 static char ID; // Pass identification, replacement for typeid 69 WebAssemblyFixIrreducibleControlFlow() : MachineFunctionPass(ID) {} 70 }; 71 } // end anonymous namespace 72 73 char WebAssemblyFixIrreducibleControlFlow::ID = 0; 74 FunctionPass *llvm::createWebAssemblyFixIrreducibleControlFlow() { 75 return new WebAssemblyFixIrreducibleControlFlow(); 76 } 77 78 namespace { 79 80 /// A utility for walking the blocks of a loop, handling a nested inner 81 /// loop as a monolithic conceptual block. 82 class MetaBlock { 83 MachineBasicBlock *Block; 84 SmallVector<MachineBasicBlock *, 2> Preds; 85 SmallVector<MachineBasicBlock *, 2> Succs; 86 87 public: 88 explicit MetaBlock(MachineBasicBlock *MBB) 89 : Block(MBB), Preds(MBB->pred_begin(), MBB->pred_end()), 90 Succs(MBB->succ_begin(), MBB->succ_end()) {} 91 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); 97 } 98 99 MachineBasicBlock *getBlock() const { return Block; } 100 101 const SmallVectorImpl<MachineBasicBlock *> &predecessors() const { 102 return Preds; 103 } 104 const SmallVectorImpl<MachineBasicBlock *> &successors() const { 105 return Succs; 106 } 107 108 bool operator==(const MetaBlock &MBB) { return Block == MBB.Block; } 109 bool operator!=(const MetaBlock &MBB) { return Block != MBB.Block; } 110 }; 111 112 class SuccessorList final : public MetaBlock { 113 size_t Index; 114 size_t Num; 115 116 public: 117 explicit SuccessorList(MachineBasicBlock *MBB) 118 : MetaBlock(MBB), Index(0), Num(successors().size()) {} 119 120 explicit SuccessorList(MachineLoop *Loop) 121 : MetaBlock(Loop), Index(0), Num(successors().size()) {} 122 123 bool HasNext() const { return Index != Num; } 124 125 MachineBasicBlock *Next() { 126 assert(HasNext()); 127 return successors()[Index++]; 128 } 129 }; 130 131 } // end anonymous namespace 132 133 bool WebAssemblyFixIrreducibleControlFlow::VisitLoop(MachineFunction &MF, 134 MachineLoopInfo &MLI, 135 MachineLoop *Loop) { 136 MachineBasicBlock *Header = Loop ? Loop->getHeader() : &*MF.begin(); 137 SetVector<MachineBasicBlock *> RewriteSuccs; 138 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(); 150 if (Top.HasNext()) { 151 MachineBasicBlock *Next = Top.Next(); 152 if (Next == Header || (Loop && !Loop->contains(Next))) 153 continue; 154 if (LLVM_LIKELY(OnStack.insert(Next).second)) { 155 if (!Visited.insert(Next).second) { 156 OnStack.erase(Next); 157 continue; 158 } 159 MachineLoop *InnerLoop = MLI.getLoopFor(Next); 160 if (InnerLoop != Loop) 161 LoopWorklist.push_back(SuccessorList(InnerLoop)); 162 else 163 LoopWorklist.push_back(SuccessorList(Next)); 164 } else { 165 RewriteSuccs.insert(Top.getBlock()); 166 } 167 continue; 168 } 169 OnStack.erase(Top.getBlock()); 170 LoopWorklist.pop_back(); 171 } 172 173 // Most likely, we didn't find any irreducible control flow. 174 if (LLVM_LIKELY(RewriteSuccs.empty())) 175 return false; 176 177 DEBUG(dbgs() << "Irreducible control flow detected!\n"); 178 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); 184 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)); 189 190 // Add the register which will be used to tell the jump table which block to 191 // jump to. 192 MachineRegisterInfo &MRI = MF.getRegInfo(); 193 unsigned Reg = MRI.createVirtualRegister(&WebAssembly::I32RegClass); 194 MIB.addReg(Reg); 195 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(), 200 RewriteSuccs.end()); 201 while (!SuccWorklist.empty()) { 202 MachineBasicBlock *MBB = SuccWorklist.pop_back_val(); 203 auto Pair = Indices.insert(std::make_pair(MBB, 0)); 204 if (!Pair.second) 205 continue; 206 207 unsigned Index = MIB.getInstr()->getNumExplicitOperands() - 1; 208 DEBUG(dbgs() << "MBB#" << MBB->getNumber() << " has index " << Index 209 << "\n"); 210 211 Pair.first->second = Index; 212 for (auto Pred : MBB->predecessors()) 213 RewriteSuccs.insert(Pred); 214 215 MIB.addMBB(MBB); 216 Dispatch->addSuccessor(MBB); 217 218 MetaBlock Meta(MBB); 219 for (auto *Succ : Meta.successors()) 220 if (Succ != Header && (!Loop || Loop->contains(Succ))) 221 SuccWorklist.push_back(Succ); 222 } 223 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)) 231 continue; 232 233 MachineBasicBlock *Split = MF.CreateMachineBasicBlock(); 234 MF.insert(MBB->isLayoutSuccessor(Succ) ? MachineFunction::iterator(Succ) 235 : MF.end(), 236 Split); 237 MLI.changeLoopFor(Split, Loop); 238 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), 242 Reg) 243 .addImm(Indices[Succ]); 244 BuildMI(*Split, Split->end(), DebugLoc(), TII.get(WebAssembly::BR)) 245 .addMBB(Dispatch); 246 Split->addSuccessor(Dispatch); 247 Map[Succ] = Split; 248 } 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); 256 } 257 258 // Create a fake default label, because br_table requires one. 259 MIB.addMBB(MIB.getInstr() 260 ->getOperand(MIB.getInstr()->getNumExplicitOperands() - 1) 261 .getMBB()); 262 263 return true; 264 } 265 266 bool WebAssemblyFixIrreducibleControlFlow::runOnMachineFunction( 267 MachineFunction &MF) { 268 DEBUG(dbgs() << "********** Fixing Irreducible Control Flow **********\n" 269 "********** Function: " 270 << MF.getName() << '\n'); 271 272 bool Changed = false; 273 auto &MLI = getAnalysis<MachineLoopInfo>(); 274 275 // Visit the function body, which is identified as a null loop. 276 Changed |= VisitLoop(MF, MLI, nullptr); 277 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); 284 } 285 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(); 290 MF.RenumberBlocks(); 291 getAnalysis<MachineDominatorTree>().runOnMachineFunction(MF); 292 MLI.runOnMachineFunction(MF); 293 } 294 295 return Changed; 296 } 297