[SPIR-V] add convergence region analysis (#78456)

This new analysis returns a hierarchical view of the convergence regions
in the given function.
This will allow our passes to query which basic block belongs to which
convergence region, and structurize the code in consequence.

Definition
----------

A convergence region is a CFG with:
 - a single entry node.
 - one or multiple exit nodes (different from LLVM's regions).
 - one back-edge
 - zero or more subregions.

Excluding sub-regions nodes, the nodes of a region can only reference a
single convergence token. A subregion uses a different convergence
token.

Algorithm
---------

This algorithm assumes all loops are in the Simplify form.

Create an initial convergence region for the whole function.
  - the convergence token is the function entry token.
  - the entry is the function entrypoint.
- Exits are all the basic blocks terminating with a return instruction.

Take the function CFG, and process it in DAG order (ignoring
back-edges). If a basic block is a loop header:
 - Create a new region.
- The parent region is the parent's loop region if any, otherwise, the
top level region.
   - The region blocks are all the blocks belonging to this loop.
- For each loop exit: - visit the rest of the CFG in DAG order (ignore
back-edges). - if the region's convergence token is found, add all the
blocks dominated by the exit from which the token is reachable to the
region.
   - continue the algorithm with the loop headers successors.

llvm/lib/Target/SPIRV/Analysis/CMakeLists.txt

@@ -0,0 +1,10 @@
+add_llvm_component_library(LLVMSPIRVAnalysis
+  SPIRVConvergenceRegionAnalysis.cpp
+
+  LINK_COMPONENTS
+  Core
+  Support
+
+  ADD_TO_COMPONENT
+  SPIRV
+  )

llvm/lib/Target/SPIRV/Analysis/SPIRVConvergenceRegionAnalysis.cpp

@@ -0,0 +1,350 @@
+//===- ConvergenceRegionAnalysis.h -----------------------------*- C++ -*--===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// The analysis determines the convergence region for each basic block of
+// the module, and provides a tree-like structure describing the region
+// hierarchy.
+//
+//===----------------------------------------------------------------------===//
+
+#include "SPIRVConvergenceRegionAnalysis.h"
+#include "llvm/Analysis/LoopInfo.h"
+#include "llvm/IR/Dominators.h"
+#include "llvm/IR/IntrinsicInst.h"
+#include "llvm/InitializePasses.h"
+#include "llvm/Transforms/Utils/LoopSimplify.h"
+#include <optional>
+#include <queue>
+
+#define DEBUG_TYPE "spirv-convergence-region-analysis"
+
+using namespace llvm;
+
+namespace llvm {
+void initializeSPIRVConvergenceRegionAnalysisWrapperPassPass(PassRegistry &);
+} // namespace llvm
+
+INITIALIZE_PASS_BEGIN(SPIRVConvergenceRegionAnalysisWrapperPass,
+                      "convergence-region",
+                      "SPIRV convergence regions analysis", true, true)
+INITIALIZE_PASS_DEPENDENCY(LoopSimplify)
+INITIALIZE_PASS_DEPENDENCY(DominatorTreeWrapperPass)
+INITIALIZE_PASS_DEPENDENCY(LoopInfoWrapperPass)
+INITIALIZE_PASS_END(SPIRVConvergenceRegionAnalysisWrapperPass,
+                    "convergence-region", "SPIRV convergence regions analysis",
+                    true, true)
+
+namespace llvm {
+namespace SPIRV {
+namespace {
+
+template <typename BasicBlockType, typename IntrinsicInstType>
+std::optional<IntrinsicInstType *>
+getConvergenceTokenInternal(BasicBlockType *BB) {
+  static_assert(std::is_const_v<IntrinsicInstType> ==
+                    std::is_const_v<BasicBlockType>,
+                "Constness must match between input and output.");
+  static_assert(std::is_same_v<BasicBlock, std::remove_const_t<BasicBlockType>>,
+                "Input must be a basic block.");
+  static_assert(
+      std::is_same_v<IntrinsicInst, std::remove_const_t<IntrinsicInstType>>,
+      "Output type must be an intrinsic instruction.");
+
+  for (auto &I : *BB) {
+    if (auto *II = dyn_cast<IntrinsicInst>(&I)) {
+      switch (II->getIntrinsicID()) {
+      case Intrinsic::experimental_convergence_entry:
+      case Intrinsic::experimental_convergence_loop:
+        return II;
+      case Intrinsic::experimental_convergence_anchor: {
+        auto Bundle = II->getOperandBundle(LLVMContext::OB_convergencectrl);
+        assert(Bundle->Inputs.size() == 1 &&
+               Bundle->Inputs[0]->getType()->isTokenTy());
+        auto TII = dyn_cast<IntrinsicInst>(Bundle->Inputs[0].get());
+        assert(TII != nullptr);
+        return TII;
+      }
+      }
+    }
+
+    if (auto *CI = dyn_cast<CallInst>(&I)) {
+      auto OB = CI->getOperandBundle(LLVMContext::OB_convergencectrl);
+      if (!OB.has_value())
+        continue;
+      return dyn_cast<IntrinsicInst>(OB.value().Inputs[0]);
+    }
+  }
+
+  return std::nullopt;
+}
+
+// Given a ConvergenceRegion tree with |Start| as its root, finds the smallest
+// region |Entry| belongs to. If |Entry| does not belong to the region defined
+// by |Start|, this function returns |nullptr|.
+ConvergenceRegion *findParentRegion(ConvergenceRegion *Start,
+                                    BasicBlock *Entry) {
+  ConvergenceRegion *Candidate = nullptr;
+  ConvergenceRegion *NextCandidate = Start;
+
+  while (Candidate != NextCandidate && NextCandidate != nullptr) {
+    Candidate = NextCandidate;
+    NextCandidate = nullptr;
+
+    // End of the search, we can return.
+    if (Candidate->Children.size() == 0)
+      return Candidate;
+
+    for (auto *Child : Candidate->Children) {
+      if (Child->Blocks.count(Entry) != 0) {
+        NextCandidate = Child;
+        break;
+      }
+    }
+  }
+
+  return Candidate;
+}
+
+} // anonymous namespace
+
+std::optional<IntrinsicInst *> getConvergenceToken(BasicBlock *BB) {
+  return getConvergenceTokenInternal<BasicBlock, IntrinsicInst>(BB);
+}
+
+std::optional<const IntrinsicInst *> getConvergenceToken(const BasicBlock *BB) {
+  return getConvergenceTokenInternal<const BasicBlock, const IntrinsicInst>(BB);
+}
+
+ConvergenceRegion::ConvergenceRegion(DominatorTree &DT, LoopInfo &LI,
+                                     Function &F)
+    : DT(DT), LI(LI), Parent(nullptr) {
+  Entry = &F.getEntryBlock();
+  ConvergenceToken = getConvergenceToken(Entry);
+  for (auto &B : F) {
+    Blocks.insert(&B);
+    if (isa<ReturnInst>(B.getTerminator()))
+      Exits.insert(&B);
+  }
+}
+
+ConvergenceRegion::ConvergenceRegion(
+    DominatorTree &DT, LoopInfo &LI,
+    std::optional<IntrinsicInst *> ConvergenceToken, BasicBlock *Entry,
+    SmallPtrSet<BasicBlock *, 8> &&Blocks, SmallPtrSet<BasicBlock *, 2> &&Exits)
+    : DT(DT), LI(LI), ConvergenceToken(ConvergenceToken), Entry(Entry),
+      Exits(std::move(Exits)), Blocks(std::move(Blocks)) {
+  for (auto *BB : this->Exits)
+    assert(this->Blocks.count(BB) != 0);
+  assert(this->Blocks.count(this->Entry) != 0);
+}
+
+void ConvergenceRegion::releaseMemory() {
+  // Parent memory is owned by the parent.
+  Parent = nullptr;
+  for (auto *Child : Children) {
+    Child->releaseMemory();
+    delete Child;
+  }
+  Children.resize(0);
+}
+
+void ConvergenceRegion::dump(const unsigned IndentSize) const {
+  const std::string Indent(IndentSize, '\t');
+  dbgs() << Indent << this << ": {\n";
+  dbgs() << Indent << "	Parent: " << Parent << "\n";
+
+  if (ConvergenceToken.value_or(nullptr)) {
+    dbgs() << Indent
+           << "	ConvergenceToken: " << ConvergenceToken.value()->getName()
+           << "\n";
+  }
+
+  if (Entry->getName() != "")
+    dbgs() << Indent << "	Entry: " << Entry->getName() << "\n";
+  else
+    dbgs() << Indent << "	Entry: " << Entry << "\n";
+
+  dbgs() << Indent << "	Exits: { ";
+  for (const auto &Exit : Exits) {
+    if (Exit->getName() != "")
+      dbgs() << Exit->getName() << ", ";
+    else
+      dbgs() << Exit << ", ";
+  }
+  dbgs() << "	}\n";
+
+  dbgs() << Indent << "	Blocks: { ";
+  for (const auto &Block : Blocks) {
+    if (Block->getName() != "")
+      dbgs() << Block->getName() << ", ";
+    else
+      dbgs() << Block << ", ";
+  }
+  dbgs() << "	}\n";
+
+  dbgs() << Indent << "	Children: {\n";
+  for (const auto Child : Children)
+    Child->dump(IndentSize + 2);
+  dbgs() << Indent << "	}\n";
+
+  dbgs() << Indent << "}\n";
+}
+
+class ConvergenceRegionAnalyzer {
+
+public:
+  ConvergenceRegionAnalyzer(Function &F, DominatorTree &DT, LoopInfo &LI)
+      : DT(DT), LI(LI), F(F) {}
+
+private:
+  bool isBackEdge(const BasicBlock *From, const BasicBlock *To) const {
+    assert(From != To && "From == To. This is awkward.");
+
+    // We only handle loop in the simplified form. This means:
+    // - a single back-edge, a single latch.
+    // - meaning the back-edge target can only be the loop header.
+    // - meaning the From can only be the loop latch.
+    if (!LI.isLoopHeader(To))
+      return false;
+
+    auto *L = LI.getLoopFor(To);
+    if (L->contains(From) && L->isLoopLatch(From))
+      return true;
+
+    return false;
+  }
+
+  std::unordered_set<BasicBlock *>
+  findPathsToMatch(LoopInfo &LI, BasicBlock *From,
+                   std::function<bool(const BasicBlock *)> isMatch) const {
+    std::unordered_set<BasicBlock *> Output;
+
+    if (isMatch(From))
+      Output.insert(From);
+
+    auto *Terminator = From->getTerminator();
+    for (unsigned i = 0; i < Terminator->getNumSuccessors(); ++i) {
+      auto *To = Terminator->getSuccessor(i);
+      if (isBackEdge(From, To))
+        continue;
+
+      auto ChildSet = findPathsToMatch(LI, To, isMatch);
+      if (ChildSet.size() == 0)
+        continue;
+
+      Output.insert(ChildSet.begin(), ChildSet.end());
+      Output.insert(From);
+      if (LI.isLoopHeader(From)) {
+        auto *L = LI.getLoopFor(From);
+        for (auto *BB : L->getBlocks()) {
+          Output.insert(BB);
+        }
+      }
+    }
+
+    return Output;
+  }
+
+  SmallPtrSet<BasicBlock *, 2>
+  findExitNodes(const SmallPtrSetImpl<BasicBlock *> &RegionBlocks) {
+    SmallPtrSet<BasicBlock *, 2> Exits;
+
+    for (auto *B : RegionBlocks) {
+      auto *Terminator = B->getTerminator();
+      for (unsigned i = 0; i < Terminator->getNumSuccessors(); ++i) {
+        auto *Child = Terminator->getSuccessor(i);
+        if (RegionBlocks.count(Child) == 0)
+          Exits.insert(B);
+      }
+    }
+
+    return Exits;
+  }
+
+public:
+  ConvergenceRegionInfo analyze() {
+    ConvergenceRegion *TopLevelRegion = new ConvergenceRegion(DT, LI, F);
+    std::queue<Loop *> ToProcess;
+    for (auto *L : LI.getLoopsInPreorder())
+      ToProcess.push(L);
+
+    while (ToProcess.size() != 0) {
+      auto *L = ToProcess.front();
+      ToProcess.pop();
+      assert(L->isLoopSimplifyForm());
+
+      auto CT = getConvergenceToken(L->getHeader());
+      SmallPtrSet<BasicBlock *, 8> RegionBlocks(L->block_begin(),
+                                                L->block_end());
+      SmallVector<BasicBlock *> LoopExits;
+      L->getExitingBlocks(LoopExits);
+      if (CT.has_value()) {
+        for (auto *Exit : LoopExits) {
+          auto N = findPathsToMatch(LI, Exit, [&CT](const BasicBlock *block) {
+            auto Token = getConvergenceToken(block);
+            if (Token == std::nullopt)
+              return false;
+            return Token.value() == CT.value();
+          });
+          RegionBlocks.insert(N.begin(), N.end());
+        }
+      }
+
+      auto RegionExits = findExitNodes(RegionBlocks);
+      ConvergenceRegion *Region = new ConvergenceRegion(
+          DT, LI, CT, L->getHeader(), std::move(RegionBlocks),
+          std::move(RegionExits));
+      Region->Parent = findParentRegion(TopLevelRegion, Region->Entry);
+      assert(Region->Parent != nullptr && "This is impossible.");
+      Region->Parent->Children.push_back(Region);
+    }
+
+    return ConvergenceRegionInfo(TopLevelRegion);
+  }
+
+private:
+  DominatorTree &DT;
+  LoopInfo &LI;
+  Function &F;
+};
+
+ConvergenceRegionInfo getConvergenceRegions(Function &F, DominatorTree &DT,
+                                            LoopInfo &LI) {
+  ConvergenceRegionAnalyzer Analyzer(F, DT, LI);
+  return Analyzer.analyze();
+}
+
+} // namespace SPIRV
+
+char SPIRVConvergenceRegionAnalysisWrapperPass::ID = 0;
+
+SPIRVConvergenceRegionAnalysisWrapperPass::
+    SPIRVConvergenceRegionAnalysisWrapperPass()
+    : FunctionPass(ID) {}
+
+bool SPIRVConvergenceRegionAnalysisWrapperPass::runOnFunction(Function &F) {
+  DominatorTree &DT = getAnalysis<DominatorTreeWrapperPass>().getDomTree();
+  LoopInfo &LI = getAnalysis<LoopInfoWrapperPass>().getLoopInfo();
+
+  CRI = SPIRV::getConvergenceRegions(F, DT, LI);
+  // Nothing was modified.
+  return false;
+}
+
+SPIRVConvergenceRegionAnalysis::Result
+SPIRVConvergenceRegionAnalysis::run(Function &F, FunctionAnalysisManager &AM) {
+  Result CRI;
+  auto &DT = AM.getResult<DominatorTreeAnalysis>(F);
+  auto &LI = AM.getResult<LoopAnalysis>(F);
+  CRI = SPIRV::getConvergenceRegions(F, DT, LI);
+  return CRI;
+}
+
+AnalysisKey SPIRVConvergenceRegionAnalysis::Key;
+
+} // namespace llvm