[LLVM] refactor GenericSSAContext and its specializations

Fix the GenericSSAContext template so that it actually declares all the
necessary typenames and the methods that must be implemented by its
specializations SSAContext and MachineSSAContext.

Reviewed By: arsenm

Differential Revision: https://reviews.llvm.org/D156288

Author: ssahasra

llvm/include/llvm/ADT/GenericCycleImpl.h

@@ -354,11 +354,11 @@ template <typename ContextT> void GenericCycleInfo<ContextT>::clear() {
 template <typename ContextT>
 void GenericCycleInfo<ContextT>::compute(FunctionT &F) {
   GenericCycleInfoCompute<ContextT> Compute(*this);
-  Context.setFunction(F);
+  Context = ContextT(&F);
 
   LLVM_DEBUG(errs() << "Computing cycles for function: " << F.getName()
                     << "\n");
-  Compute.run(ContextT::getEntryBlock(F));
+  Compute.run(&F.front());
 
   assert(validateTree());
 }

llvm/include/llvm/ADT/GenericCycleInfo.h

@@ -256,7 +256,7 @@ template <typename ContextT> class GenericCycleInfo {
   void clear();
   void compute(FunctionT &F);
 
-  FunctionT *getFunction() const { return Context.getFunction(); }
+  const FunctionT *getFunction() const { return Context.getFunction(); }
   const ContextT &getSSAContext() const { return Context; }
 
   CycleT *getCycle(const BlockT *Block) const;

llvm/include/llvm/ADT/GenericSSAContext.h

@@ -21,87 +21,78 @@
 
 namespace llvm {
 
+template <typename, bool> class DominatorTreeBase;
+template <typename> class SmallVectorImpl;
+
+// Specializations of this template should provide the types used by the
+// template GenericSSAContext below.
+template <typename _FunctionT> struct GenericSSATraits;
+
+// Ideally this should have been a stateless traits class. But the print methods
+// for Machine IR need access to the owning function. So we track that state in
+// the template itself.
+//
+// We use FunctionT as a template argument and not GenericSSATraits to allow
+// forward declarations using well-known typenames.
 template <typename _FunctionT> class GenericSSAContext {
-public:
-  // Specializations should provide the following types that are similar to how
-  // LLVM IR is structured:
+  using SSATraits = GenericSSATraits<_FunctionT>;
+  const typename SSATraits::FunctionT *F;
 
+public:
   // The smallest unit of the IR is a ValueT. The SSA context uses a ValueRefT,
   // which is a pointer to a ValueT, since Machine IR does not have the
   // equivalent of a ValueT.
-  //
-  // using ValueRefT = ...
-  //
+  using ValueRefT = typename SSATraits::ValueRefT;
+
   // The ConstValueRefT is needed to work with "const Value *", where const
   // needs to bind to the pointee and not the pointer.
-  //
-  // using ConstValueRefT = ...
-  //
-  // The null value for ValueRefT.
-  //
-  // static constexpr ValueRefT ValueRefNull;
+  using ConstValueRefT = typename SSATraits::ConstValueRefT;
+
+  // The null value for ValueRefT. For LLVM IR and MIR, this is simply the
+  // default constructed value.
+  static constexpr ValueRefT *ValueRefNull = {};
 
   // An InstructionT usually defines one or more ValueT objects.
-  //
-  // using InstructionT = ... must be a subclass of Value
+  using InstructionT = typename SSATraits::InstructionT;
 
   // A UseT represents a data-edge from the defining instruction to the using
   // instruction.
-  //
-  // using UseT = ...
+  using UseT = typename SSATraits::UseT;
 
   // A BlockT is a sequence of InstructionT, and forms a node of the CFG. It
   // has global methods predecessors() and successors() that return
   // the list of incoming CFG edges and outgoing CFG edges
   // respectively.
-  //
-  // using BlockT = ...
+  using BlockT = typename SSATraits::BlockT;
 
   // A FunctionT represents a CFG along with arguments and return values. It is
   // the smallest complete unit of code in a Module.
-  //
-  // The compiler produces an error here if this class is implicitly
-  // specialized due to an instantiation. An explicit specialization
-  // of this template needs to be added before the instantiation point
-  // indicated by the compiler.
-  using FunctionT = typename _FunctionT::invalidTemplateInstanceError;
+  using FunctionT = typename SSATraits::FunctionT;
 
   // A dominator tree provides the dominance relation between basic blocks in
   // a given funciton.
-  //
-  // using DominatorTreeT = ...
-
-  // Initialize the SSA context with information about the FunctionT being
-  // processed.
-  //
-  // void setFunction(FunctionT &function);
-  // FunctionT* getFunction() const;
-
-  // Every FunctionT has a unique BlockT marked as its entry.
-  //
-  // static BlockT* getEntryBlock(FunctionT &F);
-
-  // Methods to examine basic blocks and values
-  //
-  // static void appendBlockDefs(SmallVectorImpl<ValueRefT> &defs,
-  //                             BlockT &block);
-  // static void appendBlockDefs(SmallVectorImpl<const ValueRefT> &defs,
-  //                             const BlockT &block);
-
-  // static void appendBlockTerms(SmallVectorImpl<InstructionT *> &terms,
-  //                              BlockT &block);
-  // static void appendBlockTerms(SmallVectorImpl<const InstructionT *> &terms,
-  //                              const BlockT &block);
-  //
-  // static bool comesBefore(const InstructionT *lhs, const InstructionT *rhs);
-  // static bool isConstantOrUndefValuePhi(const InstructionT &Instr);
-  // const BlockT *getDefBlock(const ValueRefT value) const;
-
-  // Methods to print various objects.
-  //
-  // Printable print(BlockT *block) const;
-  // Printable print(InstructionT *inst) const;
-  // Printable print(ValueRefT value) const;
+  using DominatorTreeT = DominatorTreeBase<BlockT, false>;
+
+  GenericSSAContext() = default;
+  GenericSSAContext(const FunctionT *F) : F(F) {}
+
+  const FunctionT *getFunction() const { return F; }
+
+  static void appendBlockDefs(SmallVectorImpl<ValueRefT> &defs, BlockT &block);
+  static void appendBlockDefs(SmallVectorImpl<ConstValueRefT> &defs,
+                              const BlockT &block);
+
+  static void appendBlockTerms(SmallVectorImpl<InstructionT *> &terms,
+                               BlockT &block);
+  static void appendBlockTerms(SmallVectorImpl<const InstructionT *> &terms,
+                               const BlockT &block);
+
+  static bool isConstantOrUndefValuePhi(const InstructionT &Instr);
+  const BlockT *getDefBlock(ConstValueRefT value) const;
+
+  Printable print(const BlockT *block) const;
+  Printable print(const InstructionT *inst) const;
+  Printable print(ConstValueRefT value) const;
 };
 } // namespace llvm
 

llvm/include/llvm/ADT/GenericUniformityImpl.h

@@ -129,11 +129,11 @@ template <typename ContextT> class ModifiedPostOrder {
   const ContextT &Context;
 
   void computeCyclePO(const CycleInfoT &CI, const CycleT *Cycle,
-                      SmallPtrSetImpl<BlockT *> &Finalized);
+                      SmallPtrSetImpl<const BlockT *> &Finalized);
 
-  void computeStackPO(SmallVectorImpl<BlockT *> &Stack, const CycleInfoT &CI,
-                      const CycleT *Cycle,
-                      SmallPtrSetImpl<BlockT *> &Finalized);
+  void computeStackPO(SmallVectorImpl<const BlockT *> &Stack,
+                      const CycleInfoT &CI, const CycleT *Cycle,
+                      SmallPtrSetImpl<const BlockT *> &Finalized);
 };
 
 template <typename> class DivergencePropagator;
@@ -342,11 +342,10 @@ template <typename ContextT> class GenericUniformityAnalysisImpl {
       typename SyncDependenceAnalysisT::DivergenceDescriptor;
   using BlockLabelMapT = typename SyncDependenceAnalysisT::BlockLabelMap;
 
-  GenericUniformityAnalysisImpl(const FunctionT &F, const DominatorTreeT &DT,
-                                const CycleInfoT &CI,
+  GenericUniformityAnalysisImpl(const DominatorTreeT &DT, const CycleInfoT &CI,
                                 const TargetTransformInfo *TTI)
-      : Context(CI.getSSAContext()), F(F), CI(CI), TTI(TTI), DT(DT),
-        SDA(Context, DT, CI) {}
+      : Context(CI.getSSAContext()), F(*Context.getFunction()), CI(CI),
+        TTI(TTI), DT(DT), SDA(Context, DT, CI) {}
 
   void initialize();
 
@@ -1135,10 +1134,9 @@ bool GenericUniformityAnalysisImpl<ContextT>::isAlwaysUniform(
 
 template <typename ContextT>
 GenericUniformityInfo<ContextT>::GenericUniformityInfo(
-    FunctionT &Func, const DominatorTreeT &DT, const CycleInfoT &CI,
-    const TargetTransformInfo *TTI)
-    : F(&Func) {
-  DA.reset(new ImplT{Func, DT, CI, TTI});
+    const DominatorTreeT &DT, const CycleInfoT &CI,
+    const TargetTransformInfo *TTI) {
+  DA.reset(new ImplT{DT, CI, TTI});
 }
 
 template <typename ContextT>
@@ -1214,6 +1212,12 @@ bool GenericUniformityInfo<ContextT>::hasDivergence() const {
   return DA->hasDivergence();
 }
 
+template <typename ContextT>
+const typename ContextT::FunctionT &
+GenericUniformityInfo<ContextT>::getFunction() const {
+  return DA->getFunction();
+}
+
 /// Whether \p V is divergent at its definition.
 template <typename ContextT>
 bool GenericUniformityInfo<ContextT>::isDivergent(ConstValueRefT V) const {
@@ -1243,8 +1247,8 @@ void GenericUniformityInfo<ContextT>::print(raw_ostream &out) const {
 
 template <typename ContextT>
 void llvm::ModifiedPostOrder<ContextT>::computeStackPO(
-    SmallVectorImpl<BlockT *> &Stack, const CycleInfoT &CI, const CycleT *Cycle,
-    SmallPtrSetImpl<BlockT *> &Finalized) {
+    SmallVectorImpl<const BlockT *> &Stack, const CycleInfoT &CI,
+    const CycleT *Cycle, SmallPtrSetImpl<const BlockT *> &Finalized) {
   LLVM_DEBUG(dbgs() << "inside computeStackPO\n");
   while (!Stack.empty()) {
     auto *NextBB = Stack.back();
@@ -1313,9 +1317,9 @@ void llvm::ModifiedPostOrder<ContextT>::computeStackPO(
 template <typename ContextT>
 void ModifiedPostOrder<ContextT>::computeCyclePO(
     const CycleInfoT &CI, const CycleT *Cycle,
-    SmallPtrSetImpl<BlockT *> &Finalized) {
+    SmallPtrSetImpl<const BlockT *> &Finalized) {
   LLVM_DEBUG(dbgs() << "inside computeCyclePO\n");
-  SmallVector<BlockT *> Stack;
+  SmallVector<const BlockT *> Stack;
   auto *CycleHeader = Cycle->getHeader();
 
   LLVM_DEBUG(dbgs() << "  noted header: "
@@ -1352,11 +1356,11 @@ void ModifiedPostOrder<ContextT>::computeCyclePO(
 /// \brief Generically compute the modified post order.
 template <typename ContextT>
 void llvm::ModifiedPostOrder<ContextT>::compute(const CycleInfoT &CI) {
-  SmallPtrSet<BlockT *, 32> Finalized;
-  SmallVector<BlockT *> Stack;
+  SmallPtrSet<const BlockT *, 32> Finalized;
+  SmallVector<const BlockT *> Stack;
   auto *F = CI.getFunction();
   Stack.reserve(24); // FIXME made-up number
-  Stack.push_back(GraphTraits<FunctionT *>::getEntryNode(F));
+  Stack.push_back(&F->front());
   computeStackPO(Stack, CI, nullptr, Finalized);
 }
 

llvm/include/llvm/ADT/GenericUniformityInfo.h

@@ -40,8 +40,7 @@ template <typename ContextT> class GenericUniformityInfo {
   using CycleInfoT = GenericCycleInfo<ContextT>;
   using CycleT = typename CycleInfoT::CycleT;
 
-  GenericUniformityInfo(FunctionT &F, const DominatorTreeT &DT,
-                        const CycleInfoT &CI,
+  GenericUniformityInfo(const DominatorTreeT &DT, const CycleInfoT &CI,
                         const TargetTransformInfo *TTI = nullptr);
   GenericUniformityInfo() = default;
   GenericUniformityInfo(GenericUniformityInfo &&) = default;
@@ -56,7 +55,7 @@ template <typename ContextT> class GenericUniformityInfo {
   bool hasDivergence() const;
 
   /// The GPU kernel this analysis result is for
-  const FunctionT &getFunction() const { return *F; }
+  const FunctionT &getFunction() const;
 
   /// Whether \p V is divergent at its definition.
   bool isDivergent(ConstValueRefT V) const;
@@ -82,7 +81,6 @@ template <typename ContextT> class GenericUniformityInfo {
 private:
   using ImplT = GenericUniformityAnalysisImpl<ContextT>;
 
-  FunctionT *F;
   std::unique_ptr<ImplT, GenericUniformityAnalysisImplDeleter<ImplT>> DA;
 
   GenericUniformityInfo(const GenericUniformityInfo &) = delete;

llvm/include/llvm/CodeGen/MachineSSAContext.h

@@ -15,6 +15,7 @@
 #ifndef LLVM_CODEGEN_MACHINESSACONTEXT_H
 #define LLVM_CODEGEN_MACHINESSACONTEXT_H
 
+#include "llvm/ADT/GenericSSAContext.h"
 #include "llvm/CodeGen/MachineBasicBlock.h"
 #include "llvm/Support/Printable.h"
 
@@ -23,8 +24,6 @@ class MachineRegisterInfo;
 class MachineInstr;
 class MachineFunction;
 class Register;
-template <typename _FunctionT> class GenericSSAContext;
-template <typename, bool> class DominatorTreeBase;
 
 inline unsigned succ_size(const MachineBasicBlock *BB) {
   return BB->succ_size();
@@ -34,37 +33,13 @@ inline unsigned pred_size(const MachineBasicBlock *BB) {
 }
 inline auto instrs(const MachineBasicBlock &BB) { return BB.instrs(); }
 
-template <> class GenericSSAContext<MachineFunction> {
-  const MachineRegisterInfo *RegInfo = nullptr;
-  MachineFunction *MF = nullptr;
-
-public:
+template <> struct GenericSSATraits<MachineFunction> {
   using BlockT = MachineBasicBlock;
   using FunctionT = MachineFunction;
   using InstructionT = MachineInstr;
   using ValueRefT = Register;
   using ConstValueRefT = Register;
   using UseT = MachineOperand;
-  using DominatorTreeT = DominatorTreeBase<BlockT, false>;
-
-  static constexpr Register ValueRefNull = 0;
-
-  void setFunction(MachineFunction &Fn);
-  MachineFunction *getFunction() const { return MF; }
-
-  static MachineBasicBlock *getEntryBlock(MachineFunction &F);
-  static void appendBlockDefs(SmallVectorImpl<Register> &defs,
-                              const MachineBasicBlock &block);
-  static void appendBlockTerms(SmallVectorImpl<MachineInstr *> &terms,
-                               MachineBasicBlock &block);
-  static void appendBlockTerms(SmallVectorImpl<const MachineInstr *> &terms,
-                               const MachineBasicBlock &block);
-  MachineBasicBlock *getDefBlock(Register) const;
-  static bool isConstantOrUndefValuePhi(const MachineInstr &Phi);
-
-  Printable print(const MachineBasicBlock *Block) const;
-  Printable print(const MachineInstr *Inst) const;
-  Printable print(Register Value) const;
 };
 
 using MachineSSAContext = GenericSSAContext<MachineFunction>;