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 "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 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() << printMBBReference(*MBB) << " has index " << Index << "\n");
 209
 210     Pair.first->second = Index;
 211     for (auto Pred : MBB->predecessors())
 212       RewriteSuccs.insert(Pred);
 213
 214     MIB.addMBB(MBB);
 215     Dispatch->addSuccessor(MBB);
 216
 217     MetaBlock Meta(MBB);
 218     for (auto *Succ : Meta.successors())
 219       if (Succ != Header && (!Loop || Loop->contains(Succ)))
 220         SuccWorklist.push_back(Succ);
 221   }
 222
 223   // Rewrite the problematic successors for every block in RewriteSuccs.
 224   // For simplicity, we just introduce a new block for every edge we need to
 225   // rewrite. Fancier things are possible.
 226   for (MachineBasicBlock *MBB : RewriteSuccs) {
 227     DenseMap<MachineBasicBlock *, MachineBasicBlock *> Map;
 228     for (auto *Succ : MBB->successors()) {
 229       if (!Indices.count(Succ))
 230         continue;
 231
 232       MachineBasicBlock *Split = MF.CreateMachineBasicBlock();
 233       MF.insert(MBB->isLayoutSuccessor(Succ) ? MachineFunction::iterator(Succ)
 234                                              : MF.end(),
 235                 Split);
 236       MLI.changeLoopFor(Split, Loop);
 237
 238       // Set the jump table's register of the index of the block we wish to
 239       // jump to, and jump to the jump table.
 240       BuildMI(*Split, Split->end(), DebugLoc(), TII.get(WebAssembly::CONST_I32),
 241               Reg)
 242           .addImm(Indices[Succ]);
 243       BuildMI(*Split, Split->end(), DebugLoc(), TII.get(WebAssembly::BR))
 244           .addMBB(Dispatch);
 245       Split->addSuccessor(Dispatch);
 246       Map[Succ] = Split;
 247     }
 248     // Remap the terminator operands and the successor list.
 249     for (MachineInstr &Term : MBB->terminators())
 250       for (auto &Op : Term.explicit_uses())
 251         if (Op.isMBB() && Indices.count(Op.getMBB()))
 252           Op.setMBB(Map[Op.getMBB()]);
 253     for (auto Rewrite : Map)
 254       MBB->replaceSuccessor(Rewrite.first, Rewrite.second);
 255   }
 256
 257   // Create a fake default label, because br_table requires one.
 258   MIB.addMBB(MIB.getInstr()
 259                  ->getOperand(MIB.getInstr()->getNumExplicitOperands() - 1)
 260                  .getMBB());
 261
 262   return true;
 263 }
 264
 265 bool WebAssemblyFixIrreducibleControlFlow::runOnMachineFunction(
 266     MachineFunction &MF) {
 267   DEBUG(dbgs() << "********** Fixing Irreducible Control Flow **********\n"
 268                   "********** Function: "
 269                << MF.getName() << '\n');
 270
 271   bool Changed = false;
 272   auto &MLI = getAnalysis<MachineLoopInfo>();
 273
 274   // Visit the function body, which is identified as a null loop.
 275   Changed |= VisitLoop(MF, MLI, nullptr);
 276
 277   // Visit all the loops.
 278   SmallVector<MachineLoop *, 8> Worklist(MLI.begin(), MLI.end());
 279   while (!Worklist.empty()) {
 280     MachineLoop *CurLoop = Worklist.pop_back_val();
 281     Worklist.append(CurLoop->begin(), CurLoop->end());
 282     Changed |= VisitLoop(MF, MLI, CurLoop);
 283   }
 284
 285   // If we made any changes, completely recompute everything.
 286   if (LLVM_UNLIKELY(Changed)) {
 287     DEBUG(dbgs() << "Recomputing dominators and loops.\n");
 288     MF.getRegInfo().invalidateLiveness();
 289     MF.RenumberBlocks();
 290     getAnalysis<MachineDominatorTree>().runOnMachineFunction(MF);
 291     MLI.runOnMachineFunction(MF);
 292   }
 293
 294   return Changed;
 295 }