diff --git a/src/cfg/wto.h b/src/cfg/wto.h index 742ca6f88b0..339b837a36a 100644 --- a/src/cfg/wto.h +++ b/src/cfg/wto.h @@ -81,7 +81,6 @@ #include #include -#include #include #include "cfg/domtree.h" @@ -94,18 +93,25 @@ namespace wasm { template struct WeakTopologicalOrdering { static constexpr Index NoIndex = Index(-1); - struct Cycle; - using Element = std::variant; - using List = std::vector; - - struct Cycle { - List elems; - - BasicBlock* head() const { return std::get(elems.front()); } - bool operator==(const Cycle& other) const { return elems == other.elems; } + // `entries` is the flattened sequence of blocks in weak topological order, + // plus markers for the end of cycles. + // + // For normal block entries: + // - cycleTarget == NoIndex + // - block is the corresponding basic block + // + // For end of cycle entries: + // - cycleTarget is the index in `entries` of the head of the cycle, i.e. + // where we will go next if we need to process the cycle again. + // - block is the basic block at the head of the cycle, used to check whether + // we need to process the cycle again (by checking contents.inQueue). + // + struct Entry { + BasicBlock* block; + Index cycleTarget; }; - List elems; + std::vector entries; WeakTopologicalOrdering(std::vector>& blocks); }; @@ -239,27 +245,26 @@ WeakTopologicalOrdering::WeakTopologicalOrdering( } } - // Traverse the linked lists of children, materializing them as WTO elements. + // Traverse the linked lists of children, materializing them as WTO entries. // Loop depth should be limited, so doing this recursively should be fine. If // it ever causes an issue, we can un-recurse this. - // TODO: Flatten the WTO into a single contiguous vector of entries with cycle - // jump targets to avoid per-cycle vector allocations and recursion. - auto buildList = [&](auto& self, Index firstChild, List& out) -> void { + entries.reserve(numBlocks * 2); + auto emitList = [&](auto& self, Index firstChild) -> void { for (Index curr = firstChild; curr != NoIndex; curr = nodes[curr].nextSibling) { auto* block = blocks[curr].get(); if (nodes[curr].isLoopHeader) { - Cycle cycle; - cycle.elems.emplace_back(block); - self(self, nodes[curr].firstChild, cycle.elems); - out.emplace_back(std::move(cycle)); + Index startPc = entries.size(); + entries.push_back({block, NoIndex}); + self(self, nodes[curr].firstChild); + entries.push_back({block, startPc}); } else { - out.emplace_back(block); + entries.push_back({block, NoIndex}); } } }; - buildList(buildList, topFirstChild, elems); + emitList(emitList, topFirstChild); } // Given a CFG in reverse postorder (e.g. from cfg-traversal), run a forward @@ -304,34 +309,28 @@ template struct WTOWorklist { } return; } - // Iterate through each element in the current cycle's list (or the - // top-level list), which will be in reverse postorder. Visit those that are - // in the queue, which may push later elements to the queue. When there is a - // nested cycle, repeatedly visit it recursively until it stabilizes before - // continuing on. We could un-recurse this, but the loop depth is expected - // to be acceptably small. + // Iterate through the flattened WTO entries in reverse postorder. Visit + // blocks that are in the queue, which may push later blocks or loop headers + // to the queue. At the end of a cycle, jump back to the cycle header entry + // if the header was re-queued so the cycle repeats until it stabilizes. WeakTopologicalOrdering wto(cfg.basicBlocks); - auto evalList = - [&](auto& self, - const typename WeakTopologicalOrdering::List& list) - -> void { - for (const auto& elem : list) { - if (auto* block = std::get_if(&elem)) { - if ((*block)->contents.inQueue) { - (*block)->contents.inQueue = false; - visit(*block); - } - } else { - const auto& cycle = - std::get::Cycle>(elem); - BasicBlock* head = cycle.head(); - do { - self(self, cycle.elems); - } while (head->contents.inQueue); + const auto& entries = wto.entries; + Index pc = 0; + Index end = entries.size(); + while (pc < end) { + const auto& entry = entries[pc]; + if (entry.cycleTarget == WeakTopologicalOrdering::NoIndex) { + if (entry.block->contents.inQueue) { + entry.block->contents.inQueue = false; + visit(entry.block); } + ++pc; + } else if (entry.block->contents.inQueue) { + pc = entry.cycleTarget; + } else { + ++pc; } - }; - evalList(evalList, wto.elems); + } } }; diff --git a/test/gtest/wto.cpp b/test/gtest/wto.cpp index d08e65dfb6e..79e1f63acd9 100644 --- a/test/gtest/wto.cpp +++ b/test/gtest/wto.cpp @@ -135,18 +135,31 @@ std::ostream& operator<<(std::ostream& os, const WTOList& list) { using BasicBlock = TestCFG::BasicBlock; -WTOList toIndexWTO(const WeakTopologicalOrdering::List& src) { - WTOList dst; - for (const auto& elem : src) { - if (auto* b = std::get_if(&elem)) { - dst.emplace_back((*b)->contents.index); +WTOList toIndexWTO(const WeakTopologicalOrdering& wto) { + constexpr Index NoIndex = WeakTopologicalOrdering::NoIndex; + std::vector> stack; + for (Index pc = 0; pc < wto.entries.size(); ++pc) { + const auto& entry = wto.entries[pc]; + if (entry.cycleTarget == NoIndex) { + stack.emplace_back(pc, WTOElem(entry.block->contents.index)); } else { - const auto& cycle = - std::get::Cycle>(elem); - EXPECT_EQ(cycle.head(), std::get(cycle.elems.front())); - dst.emplace_back(WTOCycle(toIndexWTO(cycle.elems))); + EXPECT_EQ(entry.block, wto.entries[entry.cycleTarget].block); + size_t start = stack.size(); + while (stack[start - 1].first != entry.cycleTarget) { + --start; + } + WTOList cycleElems; + for (size_t i = start - 1; i < stack.size(); ++i) { + cycleElems.push_back(std::move(stack[i].second)); + } + stack.erase(stack.begin() + start, stack.end()); + stack.back() = {NoIndex, WTOElem(WTOCycle(std::move(cycleElems)))}; } } + WTOList dst; + for (auto& [_, elem] : stack) { + dst.push_back(std::move(elem)); + } return dst; } @@ -209,8 +222,8 @@ void verifyWTOInvariants(const TestCFG& cfg, const WTOList& wto) { WTOList getWTO(TestCFG& cfg) { WeakTopologicalOrdering wto(cfg.basicBlocks); - EXPECT_EQ(wto.elems, wto.elems); - auto list = toIndexWTO(wto.elems); + auto list = toIndexWTO(wto); + EXPECT_EQ(list, list); verifyWTOInvariants(cfg, list); return list; }