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