contrib/llvm/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 /// \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   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 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 }