[NODE][REFACTOR] Rename IRFunctor->NodeFunctor, use func pointer

include/tvm/expr.h

@@ -32,7 +32,7 @@
 #include "dtype.h"
 #include "node/node.h"
 #include "node/container.h"
-#include "node/ir_functor.h"
+#include "node/functor.h"
 #include "runtime/c_runtime_api.h"
 
 namespace tvm {
@@ -487,7 +487,7 @@ class IRPrinter {
   /*! \brief Print indent to the stream */
   TVM_DLL void PrintIndent();
   // Allow registration to be printer.
-  using FType = IRFunctor<void(const ObjectRef&, IRPrinter *)>;
+  using FType = NodeFunctor<void(const ObjectRef&, IRPrinter *)>;
   TVM_DLL static FType& vtable();
 };
 

include/tvm/ir_functor_ext.h

@@ -24,8 +24,9 @@
 #ifndef TVM_IR_FUNCTOR_EXT_H_
 #define TVM_IR_FUNCTOR_EXT_H_
 
-#include "tvm/node/ir_functor.h"
-#include "ir.h"
+#include <tvm/node/functor.h>
+#include <tvm/ir.h>
+
 #include <utility>
 
 namespace tvm {
@@ -104,7 +105,7 @@ template<typename R, typename ...Args>
 class ExprFunctor<R(const Expr& n, Args...)> {
  private:
   using TSelf = ExprFunctor<R(const Expr& n, Args...)>;
-  using FType = IRFunctor<R(const ObjectRef& n, TSelf* self, Args...)>;
+  using FType = NodeFunctor<R(const ObjectRef& n, TSelf* self, Args...)>;
 
  public:
   /*! \brief the result type of this functor */
@@ -213,7 +214,7 @@ template<typename R, typename ...Args>
 class StmtFunctor<R(const Stmt& n, Args... args)> {
  private:
   using TSelf = StmtFunctor<R(const Stmt& n, Args... args)>;
-  using FType = IRFunctor<R(const ObjectRef& n, TSelf* self, Args... args)>;
+  using FType = NodeFunctor<R(const ObjectRef& n, TSelf* self, Args... args)>;
 
  public:
   /*! \brief the result type of this functor */

include/tvm/ir_mutator.h

@@ -28,21 +28,21 @@
 #include <utility>
 #include "expr.h"
 #include "ir.h"
-#include "tvm/node/ir_functor.h"
+#include "tvm/node/functor.h"
 
 namespace tvm {
 namespace ir {
 /*!
  * \brief a base class for mutator to iterative mutate the IR
  *
  *  This IRMutator is implemented via Visitor Pattern.
- *  Also you can implement via IRFunctor.
+ *  Also you can implement via NodeFunctor.
  *  This enables easy extensions of possible new Node.
  *  It also makes changing return types easier.
  *
  * \note If you want to return a different type other than Expr and Stmt,
  *       Simply following the same pattern as IRMutator and create a seperate class.
- * \sa IRFunctor
+ * \sa NodeFunctor
  */
 class TVM_DLL IRMutator {
  public:
@@ -65,9 +65,9 @@ class TVM_DLL IRMutator {
   /*! \brief destructor */
   virtual ~IRMutator() {}
   /*! \brief functor type of expr mutation */
-  using FMutateExpr = IRFunctor<Expr(const ObjectRef&, const Expr&, IRMutator*)>;
+  using FMutateExpr = NodeFunctor<Expr(const ObjectRef&, const Expr&, IRMutator*)>;
   /*! \brief functor type of stmt mutation */
-  using FMutateStmt = IRFunctor<Stmt(const ObjectRef&, const Stmt&, IRMutator*)>;
+  using FMutateStmt = NodeFunctor<Stmt(const ObjectRef&, const Stmt&, IRMutator*)>;
   /*! \return internal vtable of expr */
   static FMutateExpr& vtable_expr();  // NOLINT(*)
   /*! \return internal stmt of expr */

include/tvm/ir_visitor.h

@@ -25,15 +25,15 @@
 #define TVM_IR_VISITOR_H_
 
 #include "ir.h"
-#include "tvm/node/ir_functor.h"
+#include "tvm/node/functor.h"
 
 namespace tvm {
 namespace ir {
 
 /*!
  * \brief a base class for visitor to iterative traverse the IR
  *
- *  This IRVisitor is implemented via IRFunctor
+ *  This IRVisitor is implemented via NodeFunctor
  *  This enables extensions of possible new Node.
  *
  * \sa ExprFunctor, StmtFunctor, PostOrderVisit
@@ -94,7 +94,7 @@ class TVM_DLL IRVisitor {
   /*! \brief destructor */
   virtual ~IRVisitor() {}
   /*! \brief functor type of visitor */
-  using FVisit = IRFunctor<void(const ObjectRef&, IRVisitor*)>;
+  using FVisit = NodeFunctor<void(const ObjectRef&, IRVisitor*)>;
   /*! \return internal vtable*/
   static FVisit& vtable();
   // overloadable visit function.

include/tvm/node/functor.h

@@ -0,0 +1,181 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+/*!
+ * \file tvm/node/functor.h
+ * \brief Defines the Functor data structures.
+ */
+#ifndef TVM_NODE_FUNCTOR_H_
+#define TVM_NODE_FUNCTOR_H_
+
+#include <dmlc/logging.h>
+#include <tvm/runtime/registry.h>
+#include <tvm/node/node.h>
+
+#include <vector>
+#include <type_traits>
+#include <utility>
+
+namespace tvm {
+/*!
+ * \brief A dynamically dispatched functor on the type of the first argument.
+ *
+ * This is a class that is useful to construct polymorphic dispatching
+ * base on the AST/IR node's type.
+ *
+ * \code
+ *   NodeFunctor<std::string (const ObjectRef& n, std::string prefix)> tostr;
+ *   tostr.set_dispatch<Add>([](const ObjectRef& op, std::string prefix) {
+ *     return prefix + "Add";
+ *   });
+ *   tostr.set_dispatch<IntImm>([](const ObjectRef& op, std::string prefix) {
+ *     return prefix + "IntImm"
+ *   });
+ *
+ *   Expr x = make_const(1);
+ *   Expr y = x + x;
+ *   // dispatch to IntImm, outputs "MyIntImm"
+ *   LOG(INFO) << tostr(x, "My");
+ *   // dispatch to IntImm, outputs "MyAdd"
+ *   LOG(INFO) << tostr(y, "My");
+ * \endcode
+ *
+ * \tparam FType function signiture
+ *  This type if only defined for FType with function signature
+ */
+template<typename FType>
+class NodeFunctor;
+
+template<typename R, typename ...Args>
+class NodeFunctor<R(const ObjectRef& n, Args...)> {
+ private:
+  /*! \brief internal function pointer type */
+  typedef R (*FPointer)(const ObjectRef&n, Args...);
+  /*! \brief refer to itself. */
+  using TSelf = NodeFunctor<R (const ObjectRef& n, Args...)>;
+  /*! \brief internal function table */
+  std::vector<FPointer> func_;
+
+ public:
+  /*! \brief the result type of this functor */
+  using result_type = R;
+  /*!
+   * \brief Whether the functor can dispatch the corresponding Node
+   * \param n The node to be dispatched
+   * \return Whether dispatching function is registered for n's type.
+   */
+  bool can_dispatch(const ObjectRef& n) const {
+    uint32_t type_index = n->type_index();
+    return type_index < func_.size() && func_[type_index] != nullptr;
+  }
+  /*!
+   * \brief invoke the functor, dispatch on type of n
+   * \param n The Node argument
+   * \param args The additional arguments
+   * \return The result.
+   */
+  R operator()(const ObjectRef& n, Args... args) const {
+    CHECK(can_dispatch(n))
+        << "NodeFunctor calls un-registered function on type "
+        << n->GetTypeKey();
+    return (*func_[n->type_index()])(n, std::forward<Args>(args)...);
+  }
+  /*!
+   * \brief set the dispacher for type TNode
+   * \param f The function to be set.
+   * \tparam TNode the type of Node to be dispatched.
+   * \return reference to self.
+   */
+  template<typename TNode>
+  TSelf& set_dispatch(FPointer f) {  // NOLINT(*)
+    uint32_t tindex = TNode::RuntimeTypeIndex();
+    if (func_.size() <= tindex) {
+      func_.resize(tindex + 1, nullptr);
+    }
+    CHECK(func_[tindex] == nullptr)
+        << "Dispatch for " << TNode::_type_key
+        << " is already set";
+    func_[tindex] = f;
+    return *this;
+  }
+  /*!
+  * \brief unset the dispacher for type TNode
+  *
+  * \tparam TNode the type of Node to be dispatched.
+  * \return reference to self.
+  */
+  template<typename TNode>
+  TSelf& clear_dispatch() {  // NOLINT(*)
+    uint32_t tindex = TNode::RuntimeTypeIndex();
+    CHECK_LT(tindex, func_.size())
+        << "clear_dispatch: index out of range";
+    func_[tindex] = nullptr;
+    return *this;
+  }
+};
+
+
+#define TVM_REG_FUNC_VAR_DEF(ClsName)                                 \
+  static TVM_ATTRIBUTE_UNUSED auto & __make_functor ## _ ## ClsName
+
+/*!
+ * \brief Useful macro to set NodeFunctor dispatch in a global static field.
+ *
+ * \code
+ *  // Use NodeFunctor to implement IRPrinter similar to Visitor Pattern.
+ *  // vtable allows easy patch of new Node types, without changing
+ *  // interface of IRPrinter.
+ *
+ *  class IRPrinter {
+ *   public:
+ *    std::ostream& stream;
+ *    // the dispatch function.
+ *    void print(Expr e) {
+ *      const static FType& f = *vtable();
+ *      f(e, this);
+ *    }
+ *
+ *    using FType = NodeFunctor<void (const ObjectRef&, IRPrinter *)>;
+ *    // function to return global function table
+ *    static FType& vtable();
+ *  };
+ *
+ *  // in cpp/cc file
+ *  IRPrinter::FType& IRPrinter::vtable() { // NOLINT(*)
+ *    static FType inst; return inst;
+ *  }
+ *
+ *  TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable)
+ *  .set_dispatch<Add>([](const ObjectRef& ref, IRPrinter* p) {
+ *    auto* n = static_cast<const Add*>(ref.get());
+ *    p->print(n->a);
+ *    p->stream << '+'
+ *    p->print(n->b);
+ *  });
+ *
+ *
+ * \endcode
+ *
+ * \param ClsName The name of the class
+ * \param FField The static function that returns a singleton of NodeFunctor.
+ */
+#define TVM_STATIC_IR_FUNCTOR(ClsName, FField)                       \
+  TVM_STR_CONCAT(TVM_REG_FUNC_VAR_DEF(ClsName), __COUNTER__)  =      \
+      ClsName::FField()
+}  // namespace tvm
+#endif  // TVM_NODE_FUNCTOR_H_