move partial-order and sorting to utils

The partial ordering visitor is used for block sorting and
structurizing, and the sorting is required for both structurizing
and normal SPIR-V emission.

The reason why we sort after structurizing, then all target in prepare
functions is that LLVM requires the MIR to be correct after each pass.
So if we don't sort after the structurizer, then we fail this
requirement. But structurizing is only required for vulkan env, meaning
we also need a path to sort non-vulkan SPIR-V, as the sorting is also
required in the general SPIR-V spec.

Signed-off-by: Nathan Gauër <[email protected]>

Author: Keenuts

llvm/lib/Target/SPIRV/SPIRVPrepareFunctions.cpp

@@ -546,8 +546,10 @@ SPIRVPrepareFunctions::removeAggregateTypesFromSignature(Function *F) {
 
 bool SPIRVPrepareFunctions::runOnModule(Module &M) {
   bool Changed = false;
-  for (Function &F : M)
+  for (Function &F : M) {
     Changed |= substituteIntrinsicCalls(&F);
+    Changed |= sortBlocks(F);
+  }
 
   std::vector<Function *> FuncsWorklist;
   for (auto &F : M)

llvm/lib/Target/SPIRV/SPIRVStructurizer.cpp

@@ -29,6 +29,7 @@
 #include "llvm/Transforms/Utils/LowerMemIntrinsics.h"
 #include <queue>
 #include <stack>
+#include <unordered_set>
 
 using namespace llvm;
 using namespace SPIRV;
@@ -42,184 +43,6 @@ namespace {
 using BlockSet = std::unordered_set<BasicBlock *>;
 using Edge = std::pair<BasicBlock *, BasicBlock *>;
 
-// This class implements a partial ordering visitor, which visits a cyclic graph
-// in natural topological-like ordering. Topological ordering is not defined for
-// directed graphs with cycles, so this assumes cycles are a single node, and
-// ignores back-edges. The cycle is visited from the entry in the same
-// topological-like ordering.
-//
-// This means once we visit a node, we know all the possible ancestors have been
-// visited.
-//
-// clang-format off
-//
-// Given this graph:
-//
-//     ,-> B -\
-// A -+        +---> D ----> E -> F -> G -> H
-//     `-> C -/      ^                 |
-//                   +-----------------+
-//
-// Visit order is:
-//  A, [B, C in any order], D, E, F, G, H
-//
-// clang-format on
-//
-// Changing the function CFG between the construction of the visitor and
-// visiting is undefined. The visitor can be reused, but if the CFG is updated,
-// the visitor must be rebuilt.
-class PartialOrderingVisitor {
-  DomTreeBuilder::BBDomTree DT;
-  LoopInfo LI;
-  BlockSet Visited = {};
-
-  struct OrderInfo {
-    size_t Rank;
-    size_t TraversalIndex;
-  };
-
-  using BlockToOrderInfoMap = std::unordered_map<BasicBlock *, OrderInfo>;
-  BlockToOrderInfoMap BlockToOrder;
-
-  // std::unordered_map<BasicBlock *, std::pair<size_t, size_t>> B2R = {};
-  std::vector<BasicBlock *> Order = {};
-
-  // Get all basic-blocks reachable from Start.
-  BlockSet getReachableFrom(BasicBlock *Start) {
-    std::queue<BasicBlock *> ToVisit;
-    ToVisit.push(Start);
-
-    BlockSet Output;
-    while (ToVisit.size() != 0) {
-      BasicBlock *BB = ToVisit.front();
-      ToVisit.pop();
-
-      if (Output.count(BB) != 0)
-        continue;
-      Output.insert(BB);
-
-      for (BasicBlock *Successor : successors(BB)) {
-        if (DT.dominates(Successor, BB))
-          continue;
-        ToVisit.push(Successor);
-      }
-    }
-
-    return Output;
-  }
-
-  size_t visit(BasicBlock *BB, size_t Rank) {
-    if (Visited.count(BB) != 0)
-      return Rank;
-
-    Loop *L = LI.getLoopFor(BB);
-    const bool isLoopHeader = LI.isLoopHeader(BB);
-
-    if (BlockToOrder.count(BB) == 0) {
-      OrderInfo Info = {Rank, Visited.size()};
-      BlockToOrder.emplace(BB, Info);
-    } else {
-      BlockToOrder[BB].Rank = std::max(BlockToOrder[BB].Rank, Rank);
-    }
-
-    for (BasicBlock *Predecessor : predecessors(BB)) {
-      if (isLoopHeader && L->contains(Predecessor)) {
-        continue;
-      }
-
-      if (BlockToOrder.count(Predecessor) == 0) {
-        return Rank;
-      }
-    }
-
-    Visited.insert(BB);
-
-    SmallVector<BasicBlock *, 2> OtherSuccessors;
-    SmallVector<BasicBlock *, 2> LoopSuccessors;
-
-    for (BasicBlock *Successor : successors(BB)) {
-      // Ignoring back-edges.
-      if (DT.dominates(Successor, BB))
-        continue;
-
-      if (isLoopHeader && L->contains(Successor)) {
-        LoopSuccessors.push_back(Successor);
-      } else
-        OtherSuccessors.push_back(Successor);
-    }
-
-    for (BasicBlock *BB : LoopSuccessors)
-      Rank = std::max(Rank, visit(BB, Rank + 1));
-
-    size_t OutputRank = Rank;
-    for (BasicBlock *Item : OtherSuccessors)
-      OutputRank = std::max(OutputRank, visit(Item, Rank + 1));
-    return OutputRank;
-  };
-
-public:
-  // Build the visitor to operate on the function F.
-  PartialOrderingVisitor(Function &F) {
-    DT.recalculate(F);
-    LI = LoopInfo(DT);
-
-    visit(&*F.begin(), 0);
-
-    Order.reserve(F.size());
-    for (auto &[BB, Info] : BlockToOrder)
-      Order.emplace_back(BB);
-
-    std::sort(
-        Order.begin(), Order.end(),
-        [&](const auto &LHS, const auto &RHS) { return compare(LHS, RHS); });
-  }
-
-  bool compare(const BasicBlock *LHS, const BasicBlock *RHS) const {
-    const OrderInfo &InfoLHS = BlockToOrder.at(const_cast<BasicBlock *>(LHS));
-    const OrderInfo &InfoRHS = BlockToOrder.at(const_cast<BasicBlock *>(RHS));
-    if (InfoLHS.Rank != InfoRHS.Rank)
-      return InfoLHS.Rank < InfoRHS.Rank;
-    return InfoLHS.TraversalIndex < InfoRHS.TraversalIndex;
-  }
-
-  // Visit the function starting from the basic block |Start|, and calling |Op|
-  // on each visited BB. This traversal ignores back-edges, meaning this won't
-  // visit a node to which |Start| is not an ancestor.
-  // If Op returns |true|, the visitor continues. If |Op| returns false, the
-  // visitor will stop at that rank. This means if 2 nodes share the same rank,
-  // and Op returns false when visiting the first, the second will be visited
-  // afterwards. But none of their successors will.
-  void partialOrderVisit(BasicBlock &Start,
-                         std::function<bool(BasicBlock *)> Op) {
-    BlockSet Reachable = getReachableFrom(&Start);
-    assert(BlockToOrder.count(&Start) != 0);
-
-    // Skipping blocks with a rank inferior to |Start|'s rank.
-    auto It = Order.begin();
-    while (It != Order.end() && *It != &Start)
-      ++It;
-
-    // This is unexpected. Worst case |Start| is the last block,
-    // so It should point to the last block, not past-end.
-    assert(It != Order.end());
-
-    // By default, there is no rank limit. Setting it to the maximum value.
-    std::optional<size_t> EndRank = std::nullopt;
-    for (; It != Order.end(); ++It) {
-      if (EndRank.has_value() && BlockToOrder[*It].Rank > *EndRank)
-        break;
-
-      if (Reachable.count(*It) == 0) {
-        continue;
-      }
-
-      if (!Op(*It)) {
-        EndRank = BlockToOrder[*It].Rank;
-      }
-    }
-  }
-};
-
 // Helper function to do a partial order visit from the block |Start|, calling
 // |Op| on each visited node.
 void partialOrderVisit(BasicBlock &Start,
@@ -1318,44 +1141,6 @@ class SPIRVStructurizer : public FunctionPass {
     return Modified;
   }
 
-  // Sort blocks in a partial ordering, so each block is after all its
-  // dominators. This should match both the SPIR-V and the MIR requirements.
-  bool sortBlocks(Function &F) {
-    if (F.size() == 0)
-      return false;
-
-    bool Modified = false;
-
-    std::vector<BasicBlock *> Order;
-    Order.reserve(F.size());
-
-    PartialOrderingVisitor Visitor(F);
-    Visitor.partialOrderVisit(*F.begin(), [&Order](BasicBlock *Block) {
-      Order.push_back(Block);
-      return true;
-    });
-
-    assert(&*F.begin() == Order[0]);
-    BasicBlock *LastBlock = &*F.begin();
-    for (BasicBlock *BB : Order) {
-      if (BB != LastBlock && &*LastBlock->getNextNode() != BB) {
-        Modified = true;
-        BB->moveAfter(LastBlock);
-      }
-      LastBlock = BB;
-    }
-#if 0
-    for (auto It = Order.begin() + 1; It != Order.end(); ++It) {
-      if (*It != &*LastBlock->getNextNode()) {
-        Modified = true;
-        (*It)->moveAfter(LastBlock);
-      }
-      LastBlock = *It;
-    }
-#endif
-    return Modified;
-  }
-
 public:
   static char ID;