lib/Target/WebAssembly/WebAssemblyFixIrreducibleControlFlow.cpp

   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 /// 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 "MCTargetDesc/WebAssemblyMCTargetDesc.h"
  30 #include "WebAssembly.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   StringRef 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 INITIALIZE_PASS(WebAssemblyFixIrreducibleControlFlow, DEBUG_TYPE,
  75                 "Removes irreducible control flow", false, false)
  76
  77 FunctionPass *llvm::createWebAssemblyFixIrreducibleControlFlow() {
  78   return new WebAssemblyFixIrreducibleControlFlow();
  79 }
  80
  81 namespace {
  82
  83 /// A utility for walking the blocks of a loop, handling a nested inner
  84 /// loop as a monolithic conceptual block.
  85 class MetaBlock {
  86   MachineBasicBlock *Block;
  87   SmallVector<MachineBasicBlock *, 2> Preds;
  88   SmallVector<MachineBasicBlock *, 2> Succs;
  89
  90 public:
  91   explicit MetaBlock(MachineBasicBlock *MBB)
  92       : Block(MBB), Preds(MBB->pred_begin(), MBB->pred_end()),
  93         Succs(MBB->succ_begin(), MBB->succ_end()) {}
  94
  95   explicit MetaBlock(MachineLoop *Loop) : Block(Loop->getHeader()) {
  96     Loop->getExitBlocks(Succs);
  97     for (MachineBasicBlock *Pred : Block->predecessors())
  98       if (!Loop->contains(Pred))
  99         Preds.push_back(Pred);
 100   }
 101
 102   MachineBasicBlock *getBlock() const { return Block; }
 103
 104   const SmallVectorImpl<MachineBasicBlock *> &predecessors() const {
 105     return Preds;
 106   }
 107   const SmallVectorImpl<MachineBasicBlock *> &successors() const {
 108     return Succs;
 109   }
 110
 111   bool operator==(const MetaBlock &MBB) { return Block == MBB.Block; }
 112   bool operator!=(const MetaBlock &MBB) { return Block != MBB.Block; }
 113 };
 114
 115 class SuccessorList final : public MetaBlock {
 116   size_t Index;
 117   size_t Num;
 118
 119 public:
 120   explicit SuccessorList(MachineBasicBlock *MBB)
 121       : MetaBlock(MBB), Index(0), Num(successors().size()) {}
 122
 123   explicit SuccessorList(MachineLoop *Loop)
 124       : MetaBlock(Loop), Index(0), Num(successors().size()) {}
 125
 126   bool HasNext() const { return Index != Num; }
 127
 128   MachineBasicBlock *Next() {
 129     assert(HasNext());
 130     return successors()[Index++];
 131   }
 132 };
 133
 134 } // end anonymous namespace
 135
 136 bool WebAssemblyFixIrreducibleControlFlow::VisitLoop(MachineFunction &MF,
 137                                                      MachineLoopInfo &MLI,
 138                                                      MachineLoop *Loop) {
 139   MachineBasicBlock *Header = Loop ? Loop->getHeader() : &*MF.begin();
 140   SetVector<MachineBasicBlock *> RewriteSuccs;
 141
 142   // DFS through Loop's body, looking for irreducible control flow. Loop is
 143   // natural, and we stay in its body, and we treat any nested loops
 144   // monolithically, so any cycles we encounter indicate irreducibility.
 145   SmallPtrSet<MachineBasicBlock *, 8> OnStack;
 146   SmallPtrSet<MachineBasicBlock *, 8> Visited;
 147   SmallVector<SuccessorList, 4> LoopWorklist;
 148   LoopWorklist.push_back(SuccessorList(Header));
 149   OnStack.insert(Header);
 150   Visited.insert(Header);
 151   while (!LoopWorklist.empty()) {
 152     SuccessorList &Top = LoopWorklist.back();
 153     if (Top.HasNext()) {
 154       MachineBasicBlock *Next = Top.Next();
 155       if (Next == Header || (Loop && !Loop->contains(Next)))
 156         continue;
 157       if (LLVM_LIKELY(OnStack.insert(Next).second)) {
 158         if (!Visited.insert(Next).second) {
 159           OnStack.erase(Next);
 160           continue;
 161         }
 162         MachineLoop *InnerLoop = MLI.getLoopFor(Next);
 163         if (InnerLoop != Loop)
 164           LoopWorklist.push_back(SuccessorList(InnerLoop));
 165         else
 166           LoopWorklist.push_back(SuccessorList(Next));
 167       } else {
 168         RewriteSuccs.insert(Top.getBlock());
 169       }
 170       continue;
 171     }
 172     OnStack.erase(Top.getBlock());
 173     LoopWorklist.pop_back();
 174   }
 175
 176   // Most likely, we didn't find any irreducible control flow.
 177   if (LLVM_LIKELY(RewriteSuccs.empty()))
 178     return false;
 179
 180   DEBUG(dbgs() << "Irreducible control flow detected!\n");
 181
 182   // Ok. We have irreducible control flow! Create a dispatch block which will
 183   // contains a jump table to any block in the problematic set of blocks.
 184   MachineBasicBlock *Dispatch = MF.CreateMachineBasicBlock();
 185   MF.insert(MF.end(), Dispatch);
 186   MLI.changeLoopFor(Dispatch, Loop);
 187
 188   // Add the jump table.
 189   const auto &TII = *MF.getSubtarget<WebAssemblySubtarget>().getInstrInfo();
 190   MachineInstrBuilder MIB = BuildMI(*Dispatch, Dispatch->end(), DebugLoc(),
 191                                     TII.get(WebAssembly::BR_TABLE_I32));
 192
 193   // Add the register which will be used to tell the jump table which block to
 194   // jump to.
 195   MachineRegisterInfo &MRI = MF.getRegInfo();
 196   unsigned Reg = MRI.createVirtualRegister(&WebAssembly::I32RegClass);
 197   MIB.addReg(Reg);
 198
 199   // Collect all the blocks which need to have their successors rewritten,
 200   // add the successors to the jump table, and remember their index.
 201   DenseMap<MachineBasicBlock *, unsigned> Indices;
 202   SmallVector<MachineBasicBlock *, 4> SuccWorklist(RewriteSuccs.begin(),
 203                                                    RewriteSuccs.end());
 204   while (!SuccWorklist.empty()) {
 205     MachineBasicBlock *MBB = SuccWorklist.pop_back_val();
 206     auto Pair = Indices.insert(std::make_pair(MBB, 0));
 207     if (!Pair.second)
 208       continue;
 209
 210     unsigned Index = MIB.getInstr()->getNumExplicitOperands() - 1;
 211     DEBUG(dbgs() << printMBBReference(*MBB) << " has index " << Index << "\n");
 212
 213     Pair.first->second = Index;
 214     for (auto Pred : MBB->predecessors())
 215       RewriteSuccs.insert(Pred);
 216
 217     MIB.addMBB(MBB);
 218     Dispatch->addSuccessor(MBB);
 219
 220     MetaBlock Meta(MBB);
 221     for (auto *Succ : Meta.successors())
 222       if (Succ != Header && (!Loop || Loop->contains(Succ)))
 223         SuccWorklist.push_back(Succ);
 224   }
 225
 226   // Rewrite the problematic successors for every block in RewriteSuccs.
 227   // For simplicity, we just introduce a new block for every edge we need to
 228   // rewrite. Fancier things are possible.
 229   for (MachineBasicBlock *MBB : RewriteSuccs) {
 230     DenseMap<MachineBasicBlock *, MachineBasicBlock *> Map;
 231     for (auto *Succ : MBB->successors()) {
 232       if (!Indices.count(Succ))
 233         continue;
 234
 235       MachineBasicBlock *Split = MF.CreateMachineBasicBlock();
 236       MF.insert(MBB->isLayoutSuccessor(Succ) ? MachineFunction::iterator(Succ)
 237                                              : MF.end(),
 238                 Split);
 239       MLI.changeLoopFor(Split, Loop);
 240
 241       // Set the jump table's register of the index of the block we wish to
 242       // jump to, and jump to the jump table.
 243       BuildMI(*Split, Split->end(), DebugLoc(), TII.get(WebAssembly::CONST_I32),
 244               Reg)
 245           .addImm(Indices[Succ]);
 246       BuildMI(*Split, Split->end(), DebugLoc(), TII.get(WebAssembly::BR))
 247           .addMBB(Dispatch);
 248       Split->addSuccessor(Dispatch);
 249       Map[Succ] = Split;
 250     }
 251     // Remap the terminator operands and the successor list.
 252     for (MachineInstr &Term : MBB->terminators())
 253       for (auto &Op : Term.explicit_uses())
 254         if (Op.isMBB() && Indices.count(Op.getMBB()))
 255           Op.setMBB(Map[Op.getMBB()]);
 256     for (auto Rewrite : Map)
 257       MBB->replaceSuccessor(Rewrite.first, Rewrite.second);
 258   }
 259
 260   // Create a fake default label, because br_table requires one.
 261   MIB.addMBB(MIB.getInstr()
 262                  ->getOperand(MIB.getInstr()->getNumExplicitOperands() - 1)
 263                  .getMBB());
 264
 265   return true;
 266 }
 267
 268 bool WebAssemblyFixIrreducibleControlFlow::runOnMachineFunction(
 269     MachineFunction &MF) {
 270   DEBUG(dbgs() << "********** Fixing Irreducible Control Flow **********\n"
 271                   "********** Function: "
 272                << MF.getName() << '\n');
 273
 274   bool Changed = false;
 275   auto &MLI = getAnalysis<MachineLoopInfo>();
 276
 277   // Visit the function body, which is identified as a null loop.
 278   Changed |= VisitLoop(MF, MLI, nullptr);
 279
 280   // Visit all the loops.
 281   SmallVector<MachineLoop *, 8> Worklist(MLI.begin(), MLI.end());
 282   while (!Worklist.empty()) {
 283     MachineLoop *CurLoop = Worklist.pop_back_val();
 284     Worklist.append(CurLoop->begin(), CurLoop->end());
 285     Changed |= VisitLoop(MF, MLI, CurLoop);
 286   }
 287
 288   // If we made any changes, completely recompute everything.
 289   if (LLVM_UNLIKELY(Changed)) {
 290     DEBUG(dbgs() << "Recomputing dominators and loops.\n");
 291     MF.getRegInfo().invalidateLiveness();
 292     MF.RenumberBlocks();
 293     getAnalysis<MachineDominatorTree>().runOnMachineFunction(MF);
 294     MLI.runOnMachineFunction(MF);
 295   }
 296
 297   return Changed;
 298 }