[SPIR-V] Fix block sorting with irreducible CFG

Block sorting was assuming reducible CFG. Meaning we always
had a best node to continue with. Irreducible CFG makes breaks
this assumption, so the algorithm looped indefinitely because
no node was a valid candidate.

Fixes #116692

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

Author: Keenuts

llvm/lib/Target/SPIRV/SPIRVUtils.cpp

@@ -519,14 +519,18 @@ bool PartialOrderingVisitor::CanBeVisited(BasicBlock *BB) const {
 }
 
 size_t PartialOrderingVisitor::GetNodeRank(BasicBlock *BB) const {
-  size_t result = 0;
+  auto It = BlockToOrder.find(BB);
+  if (It != BlockToOrder.end())
+    return It->second.Rank;
 
+  size_t result = 0;
   for (BasicBlock *P : predecessors(BB)) {
     // Ignore back-edges.
     if (DT.dominates(BB, P))
       continue;
 
     auto Iterator = BlockToOrder.end();
+
     Loop *L = LI.getLoopFor(P);
     BasicBlock *Latch = L ? L->getLoopLatch() : nullptr;
 
@@ -541,7 +545,11 @@ size_t PartialOrderingVisitor::GetNodeRank(BasicBlock *BB) const {
       Iterator = BlockToOrder.find(Latch);
     }
 
-    assert(Iterator != BlockToOrder.end());
+    // This block hasn't been ranked yet. Ignoring.
+    // This doesn't happen often, but when dealing with irreducible CFG, we have
+    // to rank nodes without knowing the rank of all their predecessors.
+    if (Iterator == BlockToOrder.end())
+      continue;
     result = std::max(result, Iterator->second.Rank + 1);
   }
 
@@ -552,15 +560,27 @@ size_t PartialOrderingVisitor::visit(BasicBlock *BB, size_t Unused) {
   ToVisit.push(BB);
   Queued.insert(BB);
 
+  // When the graph is irreducible, we can end up in a case where each
+  // node has a predecessor we haven't ranked yet.
+  // When such case arise, we have to pick a node to continue.
+  // This index is used to determine when we looped through all candidates.
+  // Each time a candidate is processed, this counter is reset.
+  // If the index is larger than the queue size, it means we looped.
+  size_t QueueIndex = 0;
+
   while (ToVisit.size() != 0) {
     BasicBlock *BB = ToVisit.front();
     ToVisit.pop();
 
-    if (!CanBeVisited(BB)) {
+    // Either the node is a candidate, or we looped already, and this is
+    // the first node we tried.
+    if (!CanBeVisited(BB) && QueueIndex <= ToVisit.size()) {
       ToVisit.push(BB);
+      QueueIndex++;
       continue;
     }
 
+    QueueIndex = 0;
     size_t Rank = GetNodeRank(BB);
     OrderInfo Info = {Rank, BlockToOrder.size()};
     BlockToOrder.emplace(BB, Info);

llvm/lib/Target/SPIRV/SPIRVUtils.h

@@ -84,10 +84,11 @@ class PartialOrderingVisitor {
   // Visits |BB| with the current rank being |Rank|.
   size_t visit(BasicBlock *BB, size_t Rank);
 
-  size_t GetNodeRank(BasicBlock *BB) const;
   bool CanBeVisited(BasicBlock *BB) const;
 
 public:
+  size_t GetNodeRank(BasicBlock *BB) const;
+
   // Build the visitor to operate on the function F.
   PartialOrderingVisitor(Function &F);
 

llvm/test/CodeGen/SPIRV/structurizer/cf.if.nested.ll

@@ -34,28 +34,28 @@
 ; CHECK:    %[[#bb30:]] = OpLabel
 ; CHECK:                  OpSelectionMerge %[[#bb31:]] None
 ; CHECK:                  OpBranchConditional %[[#]] %[[#bb32:]] %[[#bb33:]]
-; CHECK:    %[[#bb32:]] = OpLabel
+; CHECK:     %[[#bb32]] = OpLabel
 ; CHECK:                  OpSelectionMerge %[[#bb34:]] None
-; CHECK:                  OpBranchConditional %[[#]] %[[#bb35:]] %[[#bb34:]]
-; CHECK:    %[[#bb33:]] = OpLabel
+; CHECK:                  OpBranchConditional %[[#]] %[[#bb35:]] %[[#bb34]]
+; CHECK:     %[[#bb33]] = OpLabel
 ; CHECK:                  OpSelectionMerge %[[#bb36:]] None
 ; CHECK:                  OpBranchConditional %[[#]] %[[#bb37:]] %[[#bb38:]]
-; CHECK:    %[[#bb35:]] = OpLabel
-; CHECK:                  OpBranch %[[#bb34:]]
-; CHECK:    %[[#bb37:]] = OpLabel
-; CHECK:                  OpBranch %[[#bb36:]]
-; CHECK:    %[[#bb38:]] = OpLabel
+; CHECK:     %[[#bb35]] = OpLabel
+; CHECK:                  OpBranch %[[#bb34]]
+; CHECK:     %[[#bb37]] = OpLabel
+; CHECK:                  OpBranch %[[#bb36]]
+; CHECK:     %[[#bb38]] = OpLabel
 ; CHECK:                  OpSelectionMerge %[[#bb39:]] None
-; CHECK:                  OpBranchConditional %[[#]] %[[#bb40:]] %[[#bb39:]]
-; CHECK:    %[[#bb34:]] = OpLabel
-; CHECK:                  OpBranch %[[#bb31:]]
-; CHECK:    %[[#bb40:]] = OpLabel
-; CHECK:                  OpBranch %[[#bb39:]]
-; CHECK:    %[[#bb39:]] = OpLabel
-; CHECK:                  OpBranch %[[#bb36:]]
-; CHECK:    %[[#bb36:]] = OpLabel
-; CHECK:                  OpBranch %[[#bb31:]]
-; CHECK:    %[[#bb31:]] = OpLabel
+; CHECK:                  OpBranchConditional %[[#]] %[[#bb40:]] %[[#bb39]]
+; CHECK:     %[[#bb34]] = OpLabel
+; CHECK:                  OpBranch %[[#bb31]]
+; CHECK:     %[[#bb40]] = OpLabel
+; CHECK:                  OpBranch %[[#bb39]]
+; CHECK:     %[[#bb39]] = OpLabel
+; CHECK:                  OpBranch %[[#bb36]]
+; CHECK:     %[[#bb36]] = OpLabel
+; CHECK:                  OpBranch %[[#bb31]]
+; CHECK:     %[[#bb31]] = OpLabel
 ; CHECK:                  OpReturnValue %[[#]]
 ; CHECK:                  OpFunctionEnd
 ; CHECK: %[[#func_26:]] = OpFunction %[[#void:]] DontInline %[[#]]