[WebAssembly] Exception structure recovery in CFGStackify (experimental)

[aheejin]

Summary

This patch implements exception structure recovery routine in CFGStackify pass. This will be used to sort blocks and place TryBlock/EndBlock markers, but sorting and marker placing part is not implemented.

Experimental and not tested. Doesn't run anyway because of missing instruction selection. I'm uploading this just for record.
Do not commit.

[aheejin]

@dschuff: Oh, and this does not implement the special treatment of terminate.lpad.

[aheejin]

This revision was created just for the record and was not intended for landing.

lib/Target/WebAssembly/WebAssembly.h

 
 class WebAssemblyTargetMachine;
 class ModulePass;
 class FunctionPass;
 
 // LLVM IR passes.
 ModulePass *createWebAssemblyLowerEmscriptenEHSjLj(bool DoEH, bool DoSjLj);
 void initializeWebAssemblyLowerEmscriptenEHSjLjPass(PassRegistry &);
+void initializeMachineExceptionInfoPass(PassRegistry &);
 FunctionPass *createWebAssemblyOptimizeReturned();
 
 // ISel and immediate followup passes.
 FunctionPass *createWebAssemblyISelDag(WebAssemblyTargetMachine &TM,
                                        CodeGenOpt::Level OptLevel);
 FunctionPass *createWebAssemblyArgumentMove();
 FunctionPass *createWebAssemblySetP2AlignOperands();
 

lib/Target/WebAssembly/WebAssemblyCFGStackify.cpp

 /// Then, it inserts BLOCK and LOOP markers to mark the start of scopes, since
 /// scope boundaries serve as the labels for WebAssembly's control transfers.
 ///
 /// This is sufficient to convert arbitrary CFGs into a form that works on
 /// WebAssembly, provided that all loops are single-entry.
 ///
 //===----------------------------------------------------------------------===//
 
-#include "WebAssembly.h"
 #include "MCTargetDesc/WebAssemblyMCTargetDesc.h"
+#include "WebAssembly.h"
 #include "WebAssemblyMachineFunctionInfo.h"
 #include "WebAssemblySubtarget.h"
+#include "llvm/ADT/PostOrderIterator.h"
 #include "llvm/ADT/PriorityQueue.h"
 #include "llvm/ADT/SetVector.h"
 #include "llvm/CodeGen/MachineDominators.h"
 #include "llvm/CodeGen/MachineFunction.h"
 #include "llvm/CodeGen/MachineInstrBuilder.h"
 #include "llvm/CodeGen/MachineLoopInfo.h"
 #include "llvm/CodeGen/MachineRegisterInfo.h"
 #include "llvm/CodeGen/Passes.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/raw_ostream.h"
 using namespace llvm;
 
 #define DEBUG_TYPE "wasm-cfg-stackify"
 
 namespace {
+
+class MachineException {
+  MachineBasicBlock *TopInvoke;
+  MachineBasicBlock *LPad;
+  MachineBasicBlock *Cont; // not included in exception
+  MachineException *ParentException;
+
+  std::vector<MachineException *> SubExceptions;
+  std::vector<MachineBasicBlock *> Blocks;
+  SmallPtrSet<const MachineBasicBlock *, 8> DenseBlockSet;
+  bool IsLPadShared;
+
+  MachineException(const MachineException &) = delete;
+  const MachineException &operator=(const MachineException &) = delete;
+
+public:
+  MachineException(MachineBasicBlock *TopInvoke, MachineBasicBlock *LPad,
+                   MachineBasicBlock *Cont, bool IsLPadShared = false)
+      : TopInvoke(TopInvoke), LPad(LPad), Cont(Cont), ParentException(nullptr),
+        IsLPadShared(IsLPadShared) {}
+  ~MachineException() {
+    for (size_t i = 0, e = SubExceptions.size(); i != e; ++i)
+      delete SubExceptions[i];
+  }
+
+  MachineBasicBlock *getTopInvoke() const { return TopInvoke; }
+  MachineBasicBlock *getLandingPad() const { return LPad; }
+  MachineBasicBlock *getCont() const { return Cont; }
+  MachineException *getParentException() const { return ParentException; }
+  bool isLandingPad(MachineBasicBlock *MBB) const { return MBB == LPad; }
+  void setParentException(MachineException *ME) { ParentException = ME; }
+  bool isLandingPadShared() { return IsLPadShared; }
+
+  bool contains(const MachineException *ME) const {
+    if (ME == this)
+      return true;
+    if (!ME)
+      return false;
+    return contains(ME->getParentException());
+  }
+  bool contains(const MachineBasicBlock *MBB) const {
+    return DenseBlockSet.count(MBB);
+  }
+
+  const std::vector<MachineException *> &getSubExceptions() const {
+    return SubExceptions;
+  }
+  std::vector<MachineException *> &getSubExceptionsVector() {
+    return SubExceptions;
+  }
+  typedef typename std::vector<MachineException *>::const_iterator iterator;
+  iterator begin() const { return SubExceptions.begin(); }
+  iterator end() const { return SubExceptions.end(); }
+
+  const std::vector<MachineBasicBlock *> &getBlocks() const { return Blocks; }
+
+  void reverseBlock(unsigned from) {
+    std::reverse(Blocks.begin() + from, Blocks.end());
+  }
+  void reserveBlocks(unsigned size) { Blocks.reserve(size); }
+
+  void addBlockEntry(MachineBasicBlock *MBB) {
+    Blocks.push_back(MBB);
+    DenseBlockSet.insert(MBB);
+  }
+
+  // Return the nesting level. An outermost one has depth 1.
+  unsigned getExceptionDepth() const {
+    unsigned D = 1;
+    for (const MachineException *CurException = ParentException; CurException;
+         CurException = CurException->ParentException)
+      ++D;
+    return D;
+  }
+  void print(raw_ostream &OS, unsigned Depth = 0, bool Verbose = false) const;
+  void dump() const;
+};
+
+class MachineExceptionInfo : public MachineFunctionPass {
+  DenseMap<const MachineBasicBlock *, MachineException *> BBMap;
+  std::vector<MachineException *> TopLevelExceptions;
+  MachineExceptionInfo &operator=(const MachineExceptionInfo &) = delete;
+  MachineExceptionInfo(const MachineExceptionInfo &) = delete;
+
+public:
+  static char ID;
+  MachineExceptionInfo() : MachineFunctionPass(ID) {}
+  ~MachineExceptionInfo() { releaseMemory(); }
+
+  bool runOnMachineFunction(MachineFunction &) override {
+    releaseMemory();
+    analyze(getAnalysis<MachineDominatorTree>());
+    return false;
+  }
+  void getAnalysisUsage(AnalysisUsage &AU) const override {
+    AU.setPreservesAll();
+    AU.addRequired<MachineDominatorTree>();
+    MachineFunctionPass::getAnalysisUsage(AU);
+  }
+
+  void releaseMemory() override {
+    BBMap.clear();
+    for (auto *ME : TopLevelExceptions)
+      delete ME;
+    TopLevelExceptions.clear();
+  }
+
+  bool empty() const { return TopLevelExceptions.empty(); }
+  MachineException *getExceptionFor(const MachineBasicBlock *MBB) const {
+    return BBMap.lookup(MBB);
+  }
+
+  void changeExceptionFor(MachineBasicBlock *MBB, MachineException *ME) {
+    if (!ME) {
+      BBMap.erase(MBB);
+      return;
+    }
+    BBMap[MBB] = ME;
+  }
+
+  void addTopLevelException(MachineException *New) {
+    assert(!New->getParentException() && "Exception already in subexceptions!");
+    TopLevelExceptions.push_back(New);
+  }
+
+  // Create the exception forest
+  void analyze(MachineDominatorTree &DomTree);
+  void insertIntoException(MachineBasicBlock *MBB);
+
+  void print(raw_ostream &OS, const Module *M = nullptr) const override;
+};
+
+} // end anonymous namespace
+
+namespace llvm {
+void initializeMachineExceptionInfoPass(PassRegistry &);
+extern char &MachineExceptionInfoID;
+}
+
+char MachineExceptionInfo::ID = 0;
+INITIALIZE_PASS_BEGIN(MachineExceptionInfo, "wasm-exception-info",
+                      "WebAssembly Machine Exception Information", true, true)
+INITIALIZE_PASS_DEPENDENCY(MachineDominatorTree)
+INITIALIZE_PASS_END(MachineExceptionInfo, "wasm-exception-info",
+                    "WebAssembly Machine Exception Information", true, true)
+char &llvm::MachineExceptionInfoID = MachineExceptionInfo::ID;
+
+void MachineException::print(raw_ostream &OS, unsigned Depth,
+                             bool Verbose) const {
+  OS.indent(Depth * 2) << "Exception at depth " << getExceptionDepth() <<
+                       (IsLPadShared ? "(lpad shared)" : "") << " containing:";
+
+  MachineBasicBlock *TI = getTopInvoke();
+  for (unsigned i = 0; i < getBlocks().size(); ++i) {
+    MachineBasicBlock *MBB = getBlocks()[i];
+    if (!Verbose) {
+      if (i)
+        OS << ",";
+      MBB->printAsOperand(OS, false);
+    } else
+      OS << "\n";
+
+    if (MBB == TI)
+      OS << "<top_invoke>";
+    if (isLandingPad(MBB))
+      OS << "<landingpad>";
+    if (Verbose)
+      MBB->print(OS);
+  }
+  OS << "\n";
+
+  for (iterator I = begin(), E = end(); I != E; ++I)
+    (*I)->print(OS, Depth + 2);
+}
+
+raw_ostream &operator<<(raw_ostream &OS, const MachineException &ME) {
+  ME.print(OS);
+  return OS;
+}
+
+#if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP)
+LLVM_DUMP_METHOD void MachineException::dump() const { print(dbgs()); }
+#endif
+
+// Discover a subexception
+static void discoverAndMapSubException(MachineException *ME,
+                                       MachineExceptionInfo *MEI,
+                                       MachineDominatorTree &DomTree) {
+  unsigned NumBlocks = 0;
+  unsigned NumSubExceptions = 0;
+
+  // Perform a backward CFG traversal using a worklist.
+  std::vector<MachineBasicBlock *> ReverseCFGWorklist(
+      ME->getCont()->pred_begin(), ME->getCont()->pred_end());
+  while (!ReverseCFGWorklist.empty()) {
+    MachineBasicBlock *PredBB = ReverseCFGWorklist.back();
+    ReverseCFGWorklist.pop_back();
+    MachineException *SubException = MEI->getExceptionFor(PredBB);
+    if (!SubException) {
+      if (!DomTree.isReachableFromEntry(PredBB))
+        continue;
+
+      // This is an undiscovered block. Map it to the current exception.
+      MEI->changeExceptionFor(PredBB, ME);
+      ++NumBlocks;
+      if (PredBB == ME->getTopInvoke())
+        continue;
+      // Push all block predecessors on the worklist.
+      ReverseCFGWorklist.insert(ReverseCFGWorklist.end(), PredBB->pred_begin(),
+                                PredBB->pred_end());
+
+    } else {
+      // This is a discovered block. Find its outermost discovered loop.
+      while (MachineException *Parent = SubException->getParentException())
+        SubException = Parent;
+
+      // If it is already discovered to be a subexception of this exception,
+      // continue.
+      if (SubException == ME)
+        continue;
+
+      // Discover a subexception of this exception.
+      SubException->setParentException(ME);
+      ++NumSubExceptions;
+      NumBlocks += SubException->getBlocks().capacity();
+      PredBB = SubException->getTopInvoke();
+      // Continue traversal along predecessor within this subexception tree
+      // itself. Note that a predecessor may directly reach another subexception
+      // that is not yet discovered to be a subexception of this exception,
+      // which we must traverse.
+      for (MachineBasicBlock *MBB : PredBB->predecessors())
+        if (MEI->getExceptionFor(MBB) != SubException)
+          ReverseCFGWorklist.push_back(MBB);
+    }
+  }
+
+  ME->getSubExceptionsVector().reserve(NumSubExceptions);
+  ME->reserveBlocks(NumBlocks);
+}
+
+// FIXME just to compile
+static MachineBasicBlock *getNormalDest(const MachineBasicBlock *MBB) {
+  return nullptr;
+}
+
+// Add a single block to its ancestor exceptions in postrorder. If the block is
+// a subexception's top invoke block, add the subexception to its parent in
+// postorder, then reverse the block and subexception vectors of the now
+// complete subexception to achieve RPO.
+void MachineExceptionInfo::insertIntoException(MachineBasicBlock *MBB) {
+  MachineException *SubException = getExceptionFor(MBB);
+  if (SubException && MBB == SubException->getTopInvoke()) {
+    // We reach this point once per subexception after processing all the blocks
+    // in the subexception.
+    if (SubException->getParentException())
+      SubException->getParentException()->getSubExceptionsVector().push_back(
+          SubException);
+    else
+      addTopLevelException(SubException);
+
+    // For convenience, Blocks and SubExceptions are inserted in postorder.
+    // Reverse the lists, except for the exception header, which is always at
+    // the beginning.
+    SubException->reverseBlock(1);
+    std::reverse(SubException->getSubExceptionsVector().begin(),
+                 SubException->getSubExceptionsVector().end());
+
+    SubException = SubException->getParentException();
+  }
+  for (; SubException; SubException = SubException->getParentException())
+    SubException->addBlockEntry(MBB);
+}
+
+// Do a post-order dominator traversal of CFG. For each landing pad we visit
+// Find its top invoke block and cont. block.  From the cont. block, do backward
+// CFG traversal until we visit the top invoke block. Only visit nodes dominated
+// by the top invoke block.
+// (The below is done in discoverAndMapSubException)
+// For each block we visit If it is already marked as ‘discovered’, register the
+// exception as a subexception of the current exception. Continue the backward
+// traversal from the top invoke block of the subexception.  If it is not marked
+// as ‘discovered’, mark it so and map it to the current exception.
+void MachineExceptionInfo::analyze(MachineDominatorTree &DomTree) {
+  // Postorder traversal of the dominator tree.
+  DomTree.dump();
+  for (auto DomNode : post_order(&DomTree)) {
+    MachineBasicBlock *LPad = DomNode->getBlock();
+    if (!LPad->isEHPad())
+      continue;
+
+    // If this landingpad is not shared between two exception structures, there
+    // should be a single Invoke BB that dominates this landingpad.
+    MachineBasicBlock *SingleTopInvoke = nullptr;
+    for (MachineBasicBlock *Pred : LPad->predecessors()) {
+      if (DomTree.dominates(Pred, LPad)) {
+        SingleTopInvoke = Pred;
+        break;
+      }
+    }
+    bool IsLPadShared = SingleTopInvoke != nullptr;
+
+    SmallVector<MachineBasicBlock *, 4> TopInvokes;
+
+    // Group invoke BBs into exception blocks
+    for (MachineBasicBlock *Pred : LPad->predecessors()) {
+      bool IsNewHeader = true;
+      for (size_t i = 0, e = TopInvokes.size(); i < e; ++i) {
+        if (DomTree.dominates(Pred, TopInvokes[i])) {
+          TopInvokes[i] = Pred;
+          IsNewHeader = false;
+        } else if (DomTree.dominates(TopInvokes[i], Pred)) {
+          IsNewHeader = false;
+        }
+      }
+      if (IsNewHeader)
+        TopInvokes.push_back(Pred);
+    }
+
+    // Perform a backward CFG traversal to discover and map blocks in these
+    // exception(s).
+    for (MachineBasicBlock *TopInvoke : TopInvokes) {
+      // Find try.cont block
+      MachineBasicBlock *Cont = TopInvoke;
+      while (LPad->isPredecessor(Cont))
+        Cont = getNormalDest(Cont); // FIXME this function doesn't exist
+      MachineException *ME =
+          new MachineException(TopInvoke, LPad, Cont, IsLPadShared);
+      discoverAndMapSubException(ME, this, DomTree);
+    }
+  }
+  // Perform a single forward CFG traversal to populate block and subexception
+  // vectors for all exceptions.
+  for (auto DomNode : post_order(&DomTree))
+    insertIntoException(DomNode->getBlock());
+}
+
+void MachineExceptionInfo::print(raw_ostream &OS, const Module *) const {
+  for (unsigned i = 0; i < TopLevelExceptions.size(); ++i)
+    TopLevelExceptions[i]->print(OS);
+}
+
+namespace {
+
 class WebAssemblyCFGStackify final : public MachineFunctionPass {
   const char *getPassName() const override {
     return "WebAssembly CFG Stackify";
   }
 
   void getAnalysisUsage(AnalysisUsage &AU) const override {
     AU.setPreservesCFG();
     AU.addRequired<MachineDominatorTree>();
     AU.addPreserved<MachineDominatorTree>();
     AU.addRequired<MachineLoopInfo>();
     AU.addPreserved<MachineLoopInfo>();
+    AU.addRequired<MachineExceptionInfo>();
+    AU.addPreserved<MachineExceptionInfo>();
     MachineFunctionPass::getAnalysisUsage(AU);
   }
 
   bool runOnMachineFunction(MachineFunction &MF) override;
 
 public:
   static char ID; // Pass identification, replacement for typeid
   WebAssemblyCFGStackify() : MachineFunctionPass(ID) {}

 };
 }
 
 /// Sort the blocks, taking special care to make sure that loops are not
 /// interrupted by blocks not dominated by their header.
 /// TODO: There are many opportunities for improving the heuristics here.
 /// Explore them.
 static void SortBlocks(MachineFunction &MF, const MachineLoopInfo &MLI,
+                       const MachineExceptionInfo &MEI,
                        const MachineDominatorTree &MDT) {
   // Prepare for a topological sort: Record the number of predecessors each
   // block has, ignoring loop backedges.
   MF.RenumberBlocks();
   SmallVector<unsigned, 16> NumPredsLeft(MF.getNumBlockIDs(), 0);
   for (MachineBasicBlock &MBB : MF) {
     unsigned N = MBB.pred_size();
     if (MachineLoop *L = MLI.getLoopFor(&MBB))

 }
 
 bool WebAssemblyCFGStackify::runOnMachineFunction(MachineFunction &MF) {
   DEBUG(dbgs() << "********** CFG Stackifying **********\n"
                   "********** Function: "
                << MF.getName() << '\n');
 
   const auto &MLI = getAnalysis<MachineLoopInfo>();
+  const auto &MEI = getAnalysis<MachineExceptionInfo>();
   auto &MDT = getAnalysis<MachineDominatorTree>();
   // Liveness is not tracked for EXPR_STACK physreg.
   const auto &TII = *MF.getSubtarget<WebAssemblySubtarget>().getInstrInfo();
   WebAssemblyFunctionInfo &MFI = *MF.getInfo<WebAssemblyFunctionInfo>();
   MF.getRegInfo().invalidateLiveness();
 
-  // Sort the blocks, with contiguous loops.
-  SortBlocks(MF, MLI, MDT);
+  // Sort the blocks, with contiguous loops and exceptions.
+  SortBlocks(MF, MLI, MEI, MDT);
 
   // Place the BLOCK and LOOP markers to indicate the beginnings of scopes.
   PlaceMarkers(MF, MLI, TII, MDT, MFI);
 
   return true;
 }

lib/Target/WebAssembly/WebAssemblyTargetMachine.cpp

static cl::opt<bool> EnableEmSjLj(
     cl::init(false));
 
 extern "C" void LLVMInitializeWebAssemblyTarget() {
   // Register the target.
   RegisterTargetMachine<WebAssemblyTargetMachine> X(TheWebAssemblyTarget32);
   RegisterTargetMachine<WebAssemblyTargetMachine> Y(TheWebAssemblyTarget64);
 
   // Register exception handling pass to opt
-  initializeWebAssemblyLowerEmscriptenEHSjLjPass(
-      *PassRegistry::getPassRegistry());
+  auto PR = PassRegistry::getPassRegistry();
+  initializeWebAssemblyLowerEmscriptenEHSjLjPass(*PR);
+  initializeMachineExceptionInfoPass(*PR);
 }
 
 //===----------------------------------------------------------------------===//
 // WebAssembly Lowering public interface.
 //===----------------------------------------------------------------------===//
 
 static Reloc::Model getEffectiveRelocModel(Optional<Reloc::Model> RM) {
   if (!RM.hasValue())