[Dominators] Introduce DomTreeNodeTraits to allow customization. (NFC)

[fhahn]

Summary

This patch introduces DomTreeNodeTraits for customization. Clients can implement
DomTreeNodeTraitsCustom to provide custom ParentPtr, getEntryNode and getParent.
There's also a default specialization if DomTreeNodeTraitsCustom is not implemented,
that assume a Function-like NodeT. This is what is used for the existing DominatorTree
and MachineDominatorTree.

The main motivation for this patch is using DominatorTreeBase across all
regions of a VPlan, see D140513.

[kuhar]

llvm/include/llvm/Support/GenericDomTree.h

514–515

Consider moving the graph traits type to the class definition (together with other typedefs) or to a new trait like in D141260.

[fhahn]

Implement and use DomTreeNodeTraits for customization. Clients can implement DomTreeNodeTraitsCustom to provide custom ParentPtr, getEntryNode and getParent. There's also a default specialization if DomTreeNodeTraitsCustom is not implemented, that assume a Function-like NodeT. This is what is used for the existing DominatorTree and MachineDominatorTree.

[fhahn]

Fix build failure.

[kuhar]

llvm/include/llvm/Support/GenericDomTree.h

240–244

Do we need this? I'd think one could just directly specialize DomTreeNodeTraits for their type without the need for the second template parameter and DomTreeNodeTraitsCustom:

namespace llvm {
template <>
struct DomTreeNodeTraits<MyNodeType> {
  ...
};
}

[kuhar]

llvm/include/llvm/Support/GenericDomTree.h

236–237

nit: In the default implementation, it also seems reasonable to me to go through the related graph traits methods. But since both will end up calling these, either way (i.e., direct calls vs. graph traits) seems fine to me.

[fhahn]

Remove DomTreeNodeTraitsCustom and more complicated specialization.

[fhahn]

llvm/include/llvm/Support/GenericDomTree.h

236–237

The current graph traits only specify getEntryNode(NodePtr) and not getEntryNode(ParentPtr), which is the main reason for doing the direct call here.

240–244

Yeah, good point! My original thought. was to not require DomTreeNodeTraitsCustom to specify everything (e.g. NodeType), but it makes things more complicated than necessary.

I remove DomTreeNodeTraitsCustom and updated the DomTreeNodeTraits specialization in VPlanDominatorTree.h to specify everything.

[fhahn]

rebase to re-trigger precommit tests before merging

llvm/include/llvm/Support/GenericDomTree.h

 //===- GenericDomTree.h - Generic dominator trees for graphs ----*- 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
 //
 //===----------------------------------------------------------------------===//
 /// \file
 ///
 /// This file defines a set of templates that efficiently compute a dominator
 /// tree over a generic graph. This is used typically in LLVM for fast
 /// dominance queries on the CFG, but is fully generic w.r.t. the underlying
 /// graph types.
 ///
 /// Unlike ADT/* graph algorithms, generic dominator tree has more requirements
 /// on the graph's NodeRef. The NodeRef should be a pointer and,
-/// NodeRef->getParent() must return the parent node that is also a pointer.
+/// either NodeRef->getParent() must return the parent node that is also a
+/// pointer or DomTreeNodeTraitsCustom needs to be implemented.
 ///
 /// FIXME: Maybe GenericDomTree needs a TreeTraits, instead of GraphTraits.
 ///
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_SUPPORT_GENERICDOMTREE_H
 #define LLVM_SUPPORT_GENERICDOMTREE_H
 
void ApplyUpdates(DomTreeT &DT,
                             DomTreeT::IsPostDominator> &PreViewCFG,
                   GraphDiff<typename DomTreeT::NodePtr,
                             DomTreeT::IsPostDominator> *PostViewCFG);
 
 template <typename DomTreeT>
 bool Verify(const DomTreeT &DT, typename DomTreeT::VerificationLevel VL);
 }  // namespace DomTreeBuilder
 
+template <typename NodeT> struct DomTreeNodeTraitsCustom {};
+
+/// Default DomTreeNode traits for used when DomTreeTraitsCustom is not defined
+/// for NodeT. The default implementation assume a Function-like NodeT.
+template <typename NodeT, typename = void> struct DomTreeNodeTraits {
+  using NodeType = NodeT;
+  using NodePtr = NodeT *;
+  using ParentPtr = decltype(std::declval<NodePtr>()->getParent());
+  static_assert(std::is_pointer<ParentPtr>::value,
+                "Currently NodeT's parent must be a pointer type");
+  using ParentType = std::remove_pointer_t<ParentPtr>;
+
+  static NodeT *getEntryNode(ParentPtr Parent) { return &Parent->front(); }
+  static ParentPtr getParent(NodePtr BB) { return BB->getParent(); }
+};
+
+/// DomTreeNode traits specialization to use DomTreeTraitsCustom if it is
+/// defined for NodeT.
+template <typename NodeT>
+struct DomTreeNodeTraits<
+    NodeT, std::void_t<typename DomTreeNodeTraitsCustom<NodeT>::ParentPtrTy>> {
+  using NodeType = NodeT;
+  using NodePtr = NodeT *;
+  using ParentPtr = typename DomTreeNodeTraitsCustom<NodeT>::ParentPtrTy;
+  static_assert(std::is_pointer<ParentPtr>::value,
+                "Currently NodeT's parent must be a pointer type");
+  using ParentType = std::remove_pointer_t<ParentPtr>;
+
+  static NodePtr getEntryNode(ParentPtr Parent) {
+    return DomTreeNodeTraitsCustom<NodeT>::getEntryNode(Parent);
+  }
+  static ParentPtr getParent(NodePtr BB) {
+    return DomTreeNodeTraitsCustom<NodeT>::getParent(BB);
+  }
+};
+
 /// Core dominator tree base class.
 ///
 /// 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>
 class DominatorTreeBase {
  public:
   static_assert(std::is_pointer<typename GraphTraits<NodeT *>::NodeRef>::value,
                 "Currently DominatorTreeBase supports only pointer nodes");
-  using NodeType = NodeT;
-  using NodePtr = NodeT *;
-  using ParentPtr = decltype(std::declval<NodeT *>()->getParent());
+  using NodeTrait = DomTreeNodeTraits<NodeT>;
+  using NodeType = typename NodeTrait::NodeType;
+  using NodePtr = typename NodeTrait::NodePtr;
+  using ParentPtr = typename NodeTrait::ParentPtr;
   static_assert(std::is_pointer<ParentPtr>::value,
                 "Currently NodeT's parent must be a pointer type");
   using ParentType = std::remove_pointer_t<ParentPtr>;
   static constexpr bool IsPostDominator = IsPostDom;
 
   using UpdateType = cfg::Update<NodePtr>;
   using UpdateKind = cfg::UpdateKind;
   static constexpr UpdateKind Insert = UpdateKind::Insert;
NodeT *getRoot() const {
     assert(this->Roots.size() == 1 && "Should always have entry node!");
     return this->Roots[0];
   }
 
   /// Find nearest common dominator basic block for basic block A and B. A and B
   /// must have tree nodes.
   NodeT *findNearestCommonDominator(NodeT *A, NodeT *B) const {
     assert(A && B && "Pointers are not valid");
-    assert(A->getParent() == B->getParent() &&
+    assert(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 = A->getParent()->front();
+      NodeT &Entry =
+          *DomTreeNodeTraits<NodeT>::getEntryNode(NodeTrait::getParent(A));
       if (A == &Entry || B == &Entry)
         return &Entry;
     }
 
     DomTreeNodeBase<NodeT> *NodeA = getNode(A);
     DomTreeNodeBase<NodeT> *NodeB = getNode(B);
     assert(NodeA && "A must be in the tree");
     assert(NodeB && "B must be in the tree");
#endif
   /// tree doesn't know about.
   ///
   /// Note that for postdominators it automatically takes care of inserting
   /// a reverse edge internally (so there's no need to swap the parameters).
   ///
   void insertEdge(NodeT *From, NodeT *To) {
     assert(From);
     assert(To);
-    assert(From->getParent() == Parent);
-    assert(To->getParent() == Parent);
+    assert(NodeTrait::getParent(From) == Parent);
+    assert(NodeTrait::getParent(To) == Parent);
     DomTreeBuilder::InsertEdge(*this, From, To);
   }
 
   /// Inform the dominator tree about a CFG edge deletion and update the tree.
   ///
   /// This function has to be called just after making the update on the actual
   /// CFG. An internal functions checks if the edge doesn't exist in the CFG in
   /// DEBUG mode. There cannot be any other updates that the
   /// dominator tree doesn't know about.
   ///
   /// Note that for postdominators it automatically takes care of deleting
   /// a reverse edge internally (so there's no need to swap the parameters).
   ///
   void deleteEdge(NodeT *From, NodeT *To) {
     assert(From);
     assert(To);
-    assert(From->getParent() == Parent);
-    assert(To->getParent() == Parent);
+    assert(NodeTrait::getParent(From) == Parent);
+    assert(NodeTrait::getParent(To) == Parent);
     DomTreeBuilder::DeleteEdge(*this, From, To);
   }
 
   /// Add a new node to the dominator tree information.
   ///
   /// This creates a new node as a child of DomBB dominator node, linking it
   /// into the children list of the immediate dominator.
   ///

llvm/lib/Transforms/Vectorize/VPlan.h

public:
   /// EntryBlock must have no predecessors.
   void setEntry(VPBlockBase *EntryBlock) {
     assert(EntryBlock->getPredecessors().empty() &&
            "Entry block cannot have predecessors.");
     Entry = EntryBlock;
     EntryBlock->setParent(this);
   }
 
-  // FIXME: DominatorTreeBase is doing 'A->getParent()->front()'. 'front' is a
-  // specific interface of llvm::Function, instead of using
-  // GraphTraints::getEntryNode. We should add a new template parameter to
-  // DominatorTreeBase representing the Graph type.
-  VPBlockBase &front() const { return *Entry; }
-
   const VPBlockBase *getExiting() const { return Exiting; }
   VPBlockBase *getExiting() { return Exiting; }
 
   /// Set \p ExitingBlock as the exiting VPBlockBase of this VPRegionBlock. \p
   /// ExitingBlock must have no successors.
   void setExiting(VPBlockBase *ExitingBlock) {
     assert(ExitingBlock->getSuccessors().empty() &&
            "Exit block cannot have successors.");

llvm/lib/Transforms/Vectorize/VPlanCFG.h

 
 // The following set of template specializations implement GraphTraits to treat
 // any VPBlockBase as a node in a graph of VPBlockBases. It's important to note
 // that VPBlockBase traits don't recurse into VPRegioBlocks, i.e., if the
 // VPBlockBase is a VPRegionBlock, this specialization provides access to its
 // successors/predecessors but not to the blocks inside the region.
 
 template <> struct GraphTraits<VPBlockBase *> {
+  using ParentPtrTy = VPRegionBlock *;
   using NodeRef = VPBlockBase *;
   using ChildIteratorType = SmallVectorImpl<VPBlockBase *>::iterator;
 
   static NodeRef getEntryNode(NodeRef N) { return N; }
+  static NodeRef getEntryNode(VPRegionBlock *R) { return R->getEntry(); }
 
   static inline ChildIteratorType child_begin(NodeRef N) {
     return N->getSuccessors().begin();
   }
 
   static inline ChildIteratorType child_end(NodeRef N) {
     return N->getSuccessors().end();
   }

llvm/lib/Transforms/Vectorize/VPlanDominatorTree.h

 
 #include "VPlan.h"
 #include "VPlanCFG.h"
 #include "llvm/ADT/GraphTraits.h"
 #include "llvm/IR/Dominators.h"
 
 namespace llvm {
 
+template <> struct DomTreeNodeTraitsCustom<VPBlockBase> {
+  using ParentPtrTy = VPRegionBlock *;
+
+  static VPBlockBase *getEntryNode(ParentPtrTy Parent) {
+    return Parent->getEntry();
+  }
+
+  static ParentPtrTy getParent(VPBlockBase *B) { return B->getParent(); }
+};
+
+///
 /// Template specialization of the standard LLVM dominator tree utility for
 /// VPBlockBases.
 using VPDominatorTree = DomTreeBase<VPBlockBase>;
 
 using VPDomTreeNode = DomTreeNodeBase<VPBlockBase>;
 
 /// Template specializations of GraphTraits for VPDomTreeNode.
 template <>