[VPlan] Add VPPartialDominatorTree via optional NodeTrait on DominatorTreeBase

llvm/include/llvm/ADT/GenericSSAContext.h

@@ -17,11 +17,11 @@
 #ifndef LLVM_ADT_GENERICSSACONTEXT_H
 #define LLVM_ADT_GENERICSSACONTEXT_H
 
+#include "llvm/Support/GenericDomTree.h"
 #include "llvm/Support/Printable.h"
 
 namespace llvm {
 
-template <typename, bool> class DominatorTreeBase;
 template <typename> class SmallVectorImpl;
 
 namespace Intrinsic {

llvm/include/llvm/Support/GenericDomTree.h

@@ -41,7 +41,7 @@
 
 namespace llvm {
 
-template <typename NodeT, bool IsPostDom>
+template <typename, bool, typename>
 class DominatorTreeBase;
 
 namespace DomTreeBuilder {
@@ -52,10 +52,10 @@ struct SemiNCAInfo;
 /// Base class for the actual dominator tree node.
 template <class NodeT> class DomTreeNodeBase {
   friend class PostDominatorTree;
-  friend class DominatorTreeBase<NodeT, false>;
-  friend class DominatorTreeBase<NodeT, true>;
-  friend struct DomTreeBuilder::SemiNCAInfo<DominatorTreeBase<NodeT, false>>;
-  friend struct DomTreeBuilder::SemiNCAInfo<DominatorTreeBase<NodeT, true>>;
+  template <class NT, bool IsPD, class Trait>
+  friend class DominatorTreeBase;
+  template <class DomTreeT>
+  friend struct DomTreeBuilder::SemiNCAInfo;
 
   NodeT *TheBB;
   DomTreeNodeBase *IDom;
@@ -232,15 +232,20 @@ template <typename NodeT> struct DomTreeNodeTraits {
 ///
 /// This class is a generic template over graph nodes. It is instantiated for
 /// various graphs in the LLVM IR or in the code generator.
-template <typename NodeT, bool IsPostDom>
+///
+/// An optional NodeTrait template parameter allows customizing the traits used
+/// for the dominator tree, e.g. to use a custom entry node or parent type.
+/// This enables non-pointer NodeT types where NodePtr from NodeTrait is a
+/// pointer type different from NodeT*.
+template <typename NodeT, bool IsPostDom,
+          typename NodeTrait = DomTreeNodeTraits<NodeT>>
 class DominatorTreeBase {
  public:
-  static_assert(std::is_pointer_v<typename GraphTraits<NodeT *>::NodeRef>,
-                "Currently DominatorTreeBase supports only pointer nodes");
-  using NodeTrait = DomTreeNodeTraits<NodeT>;
   using NodeType = typename NodeTrait::NodeType;
   using NodePtr = typename NodeTrait::NodePtr;
   using ParentPtr = typename NodeTrait::ParentPtr;
+  static_assert(std::is_pointer_v<NodePtr>,
+                "DominatorTreeBase requires NodePtr to be a pointer type");
   static_assert(std::is_pointer_v<ParentPtr>,
                 "Currently NodeT's parent must be a pointer type");
   using ParentType = std::remove_pointer_t<ParentPtr>;
@@ -398,7 +403,8 @@ class DominatorTreeBase {
   /// may (but is not required to) be null for a forward (backwards)
   /// statically unreachable block.
   DomTreeNodeBase<NodeT> *getNode(const NodeT *BB) const {
-    assert((!BB || Parent == NodeTrait::getParent(const_cast<NodeT *>(BB))) &&
+    assert((!BB || !NodeTrait::getParent(const_cast<NodeT *>(BB)) ||
+            Parent == NodeTrait::getParent(const_cast<NodeT *>(BB))) &&
            "cannot get DomTreeNode of block with different parent");
     if (auto Idx = getNodeIndex(BB); Idx && *Idx < DomTreeNodes.size())
       return DomTreeNodes[*Idx].get();
@@ -517,16 +523,16 @@ class DominatorTreeBase {
   /// must have tree nodes.
   NodeT *findNearestCommonDominator(NodeT *A, NodeT *B) const {
     assert(A && B && "Pointers are not valid");
-    assert(NodeTrait::getParent(A) == NodeTrait::getParent(B) &&
+    assert((!NodeTrait::getParent(A) || !NodeTrait::getParent(B) ||
+            NodeTrait::getParent(A) == NodeTrait::getParent(B)) &&
            "Two blocks are not in same function");
 
     // If either A or B is a entry block then it is nearest common dominator
     // (for forward-dominators).
     if (!isPostDominator()) {
-      NodeT &Entry =
-          *DomTreeNodeTraits<NodeT>::getEntryNode(NodeTrait::getParent(A));
-      if (A == &Entry || B == &Entry)
-        return &Entry;
+      NodeT *Entry = NodeTrait::getEntryNode(Parent);
+      if (A == Entry || B == Entry)
+        return Entry;
     }
 
     DomTreeNodeBase<NodeT> *NodeA = getNode(A);
@@ -651,8 +657,8 @@ class DominatorTreeBase {
   void insertEdge(NodeT *From, NodeT *To) {
     assert(From);
     assert(To);
-    assert(NodeTrait::getParent(From) == Parent);
-    assert(NodeTrait::getParent(To) == Parent);
+    assert(!NodeTrait::getParent(From) || NodeTrait::getParent(From) == Parent);
+    assert(!NodeTrait::getParent(To) || NodeTrait::getParent(To) == Parent);
     DomTreeBuilder::InsertEdge(*this, From, To);
   }
 
@@ -669,8 +675,8 @@ class DominatorTreeBase {
   void deleteEdge(NodeT *From, NodeT *To) {
     assert(From);
     assert(To);
-    assert(NodeTrait::getParent(From) == Parent);
-    assert(NodeTrait::getParent(To) == Parent);
+    assert(!NodeTrait::getParent(From) || NodeTrait::getParent(From) == Parent);
+    assert(!NodeTrait::getParent(To) || NodeTrait::getParent(To) == Parent);
     DomTreeBuilder::DeleteEdge(*this, From, To);
   }
 
@@ -1023,16 +1029,16 @@ using PostDomTreeBase = DominatorTreeBase<T, true>;
 
 // These two functions are declared out of line as a workaround for building
 // with old (< r147295) versions of clang because of pr11642.
-template <typename NodeT, bool IsPostDom>
-bool DominatorTreeBase<NodeT, IsPostDom>::dominates(const NodeT *A,
-                                                    const NodeT *B) const {
+template <typename NodeT, bool IsPostDom, typename NodeTrait>
+bool DominatorTreeBase<NodeT, IsPostDom, NodeTrait>::dominates(
+    const NodeT *A, const NodeT *B) const {
   if (A == B)
     return true;
 
   return dominates(getNode(A), getNode(B));
 }
-template <typename NodeT, bool IsPostDom>
-bool DominatorTreeBase<NodeT, IsPostDom>::properlyDominates(
+template <typename NodeT, bool IsPostDom, typename NodeTrait>
+bool DominatorTreeBase<NodeT, IsPostDom, NodeTrait>::properlyDominates(
     const NodeT *A, const NodeT *B) const {
   if (A == B)
     return false;

llvm/lib/Transforms/Vectorize/VPlanDominatorTree.h

@@ -17,6 +17,7 @@
 
 #include "VPlan.h"
 #include "VPlanCFG.h"
+#include "llvm/ADT/DepthFirstIterator.h"
 #include "llvm/ADT/GraphTraits.h"
 #include "llvm/IR/Dominators.h"
 #include "llvm/Support/GenericDomTree.h"
@@ -58,5 +59,68 @@ template <>
 struct GraphTraits<const VPDomTreeNode *>
     : public DomTreeGraphTraitsBase<const VPDomTreeNode,
                                     VPDomTreeNode::const_iterator> {};
+
+/// A lightweight wrapper representing a subgraph of a VPlan rooted at a
+/// chosen entry block. Used as the parent type for VPPartialDominatorTree so
+/// that the entry can be any VPBlockBase, not necessarily the VPlan entry.
+struct VPBlockSubgraph {
+  VPBlockBase *Entry;
+};
+
+/// GraphTraits for VPBlockSubgraph* – traversal starts at Entry using the
+/// same successor-iterator as VPBlockBase*.
+template <> struct GraphTraits<VPBlockSubgraph *> {
+  using GraphRef = VPBlockSubgraph *;
+  using NodeRef = VPBlockBase *;
+  using nodes_iterator = df_iterator<NodeRef>;
+  using ChildIteratorType = VPAllSuccessorsIterator<VPBlockBase *>;
+
+  static NodeRef getEntryNode(GraphRef SG) { return SG->Entry; }
+
+  static ChildIteratorType child_begin(NodeRef N) {
+    return ChildIteratorType(N);
+  }
+  static ChildIteratorType child_end(NodeRef N) {
+    return ChildIteratorType::end(N);
+  }
+
+  static nodes_iterator nodes_begin(GraphRef SG) {
+    return nodes_iterator::begin(SG->Entry);
+  }
+  static nodes_iterator nodes_end(GraphRef SG) {
+    return nodes_iterator::end(SG->Entry);
+  }
+};
+
+/// DomTreeNodeTraits for a partial VPlan dominator tree. The parent is a
+/// VPBlockSubgraph* carrying the chosen entry block. getParent() returns
+/// nullptr because there is no single VPBlockBase field that stores which
+/// subgraph a block belongs to; the assertions in DominatorTreeBase are
+/// relaxed to treat a null return as "don't check".
+struct VPPartialDomTreeTraits {
+  using NodeType = VPBlockBase;
+  using NodePtr = VPBlockBase *;
+  using ParentPtr = VPBlockSubgraph *;
+
+  static NodePtr getEntryNode(ParentPtr SG) { return SG->Entry; }
+  /// Returns nullptr to opt out of the parent-equality assertion in
+  /// DominatorTreeBase; the tree is still correct because construction only
+  /// visits nodes reachable from the chosen entry.
+  static ParentPtr getParent(NodePtr) { return nullptr; }
+};
+
+/// A dominator tree for a subgraph of a VPlan whose entry can be any
+/// VPBlockBase, not necessarily the VPlan's own entry block.
+class VPPartialDominatorTree
+    : public DominatorTreeBase<VPBlockBase, false, VPPartialDomTreeTraits> {
+  using Base = DominatorTreeBase<VPBlockBase, false, VPPartialDomTreeTraits>;
+  VPBlockSubgraph SubgraphInfo;
+
+public:
+  explicit VPPartialDominatorTree(VPBlockBase *Entry) : SubgraphInfo{Entry} {
+    recalculate(SubgraphInfo);
+  }
+};
+
 } // namespace llvm
 #endif // LLVM_TRANSFORMS_VECTORIZE_VPLANDOMINATORTREE_H