diff --git a/include/tvm/expr.h b/include/tvm/expr.h index ea578152899d..fc52421d903b 100644 --- a/include/tvm/expr.h +++ b/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; + using FType = NodeFunctor; TVM_DLL static FType& vtable(); }; diff --git a/include/tvm/ir_functor_ext.h b/include/tvm/ir_functor_ext.h index 54a5eff6846b..04ce7934ff2f 100644 --- a/include/tvm/ir_functor_ext.h +++ b/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 +#include + #include namespace tvm { @@ -104,7 +105,7 @@ template class ExprFunctor { private: using TSelf = ExprFunctor; - using FType = IRFunctor; + using FType = NodeFunctor; public: /*! \brief the result type of this functor */ @@ -213,7 +214,7 @@ template class StmtFunctor { private: using TSelf = StmtFunctor; - using FType = IRFunctor; + using FType = NodeFunctor; public: /*! \brief the result type of this functor */ diff --git a/include/tvm/ir_mutator.h b/include/tvm/ir_mutator.h index c910a48620c8..5460ae0f4ba9 100644 --- a/include/tvm/ir_mutator.h +++ b/include/tvm/ir_mutator.h @@ -28,7 +28,7 @@ #include #include "expr.h" #include "ir.h" -#include "tvm/node/ir_functor.h" +#include "tvm/node/functor.h" namespace tvm { namespace ir { @@ -36,13 +36,13 @@ 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; + using FMutateExpr = NodeFunctor; /*! \brief functor type of stmt mutation */ - using FMutateStmt = IRFunctor; + using FMutateStmt = NodeFunctor; /*! \return internal vtable of expr */ static FMutateExpr& vtable_expr(); // NOLINT(*) /*! \return internal stmt of expr */ diff --git a/include/tvm/ir_visitor.h b/include/tvm/ir_visitor.h index bebf94585ed6..b85cf233a42f 100644 --- a/include/tvm/ir_visitor.h +++ b/include/tvm/ir_visitor.h @@ -25,7 +25,7 @@ #define TVM_IR_VISITOR_H_ #include "ir.h" -#include "tvm/node/ir_functor.h" +#include "tvm/node/functor.h" namespace tvm { namespace ir { @@ -33,7 +33,7 @@ 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; + using FVisit = NodeFunctor; /*! \return internal vtable*/ static FVisit& vtable(); // overloadable visit function. diff --git a/include/tvm/node/functor.h b/include/tvm/node/functor.h new file mode 100644 index 000000000000..d56fb8dde799 --- /dev/null +++ b/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 +#include +#include + +#include +#include +#include + +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 tostr; + * tostr.set_dispatch([](const ObjectRef& op, std::string prefix) { + * return prefix + "Add"; + * }); + * tostr.set_dispatch([](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 +class NodeFunctor; + +template +class NodeFunctor { + private: + /*! \brief internal function pointer type */ + typedef R (*FPointer)(const ObjectRef&n, Args...); + /*! \brief refer to itself. */ + using TSelf = NodeFunctor; + /*! \brief internal function table */ + std::vector 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)...); + } + /*! + * \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 + 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 + 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; + * // 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([](const ObjectRef& ref, IRPrinter* p) { + * auto* n = static_cast(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_ diff --git a/include/tvm/node/ir_functor.h b/include/tvm/node/ir_functor.h deleted file mode 100644 index e902e8fb6d44..000000000000 --- a/include/tvm/node/ir_functor.h +++ /dev/null @@ -1,282 +0,0 @@ -/* - * 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/ir_functor.h - * \brief Defines the IRFunctor data structures. - */ -#ifndef TVM_NODE_IR_FUNCTOR_H_ -#define TVM_NODE_IR_FUNCTOR_H_ - -#include -#include -#include -#include -#include -#include -#include -#include "node.h" - -namespace tvm { -/*! - * \brief A dynamically dispatched functor on ObjectRef in the first argument. - * - * \code - * IRFunctor tostr; - * tostr.set_dispatch([](const Add* op, std::string prefix) { - * return prefix + "Add"; - * }); - * tostr.set_dispatch([](const IntImm* op) { - * 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 -class IRFunctor; - -template -class IRFunctor { - private: - using Function = std::function; - using TSelf = IRFunctor; - /*! \brief internal function table */ - std::vector 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. - */ - inline 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. - */ - inline R operator()(const ObjectRef& n, Args... args) const { - uint32_t type_index = n->type_index(); - CHECK(type_index < func_.size() && - func_[type_index] != nullptr) - << "IRFunctor calls un-registered function on type " - << n->GetTypeKey(); - return func_[type_index](n, std::forward(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 - inline TSelf& set_dispatch(Function 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 set the dispacher for type TNode - * This allows f to used detailed const Node pointer to replace ObjectRef - * - * \param f The function to be set. - * \tparam TNode the type of Node to be dispatched. - * \return reference to self. - */ - template - inline TSelf& set_dispatch(std::function f) { // NOLINT(*) - Function fun = [f](const ObjectRef& n, Args... args) { - return f(static_cast(n.get()), - std::forward(args)...); - }; - return this->set_dispatch(fun); - } - /*! - * \brief unset the dispacher for type TNode - * - * \tparam TNode the type of Node to be dispatched. - * \return reference to self. - */ - template - inline 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; - } -}; - -#if defined(__GNUC__) -#define TVM_ATTRIBUTE_UNUSED __attribute__((unused)) -#else -#define TVM_ATTRIBUTE_UNUSED -#endif - -/*! \brief helper macro to generate string concat */ -#define TVM_STR_CONCAT_(__x, __y) __x##__y -#define TVM_STR_CONCAT(__x, __y) TVM_STR_CONCAT_(__x, __y) - -#define TVM_REGISTER_VAR_DEF(ClsName) \ - static TVM_ATTRIBUTE_UNUSED auto & __make_functor ## _ ## ClsName - -/*! - * \brief Useful macro to set IRFunctor dispatch in a global static field. - * - * \code - * // Use IRFunctor to implement IRPrinter similar to Visitor Pattern. - * // vtable allows easy patch in 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 = IRFunctor; - * // 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([](const Add* n, IRPrinter* p) { - * 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 IRFunctor. - */ -#define TVM_STATIC_IR_FUNCTOR(ClsName, FField) \ - TVM_STR_CONCAT(TVM_REGISTER_VAR_DEF(ClsName), __COUNTER__) = \ - ClsName::FField() - - /*! - * \brief A container for a list of callbacks. All callbacks are invoked when - * the object is destructed. - */ -class IRFunctorCleanList { - public: - ~IRFunctorCleanList() { - for (auto &f : clean_items) { - f(); - } - } - - void append(std::function func) { - clean_items.push_back(func); - } - - private: - std::vector< std::function > clean_items; -}; - -/*! -* \brief A wrapper around IRFunctor that will record calls to set_dispatch -* and make a corresponding call to clear_dispatch when the last copy of -* the IRFunctorStaticRegistry is destructed. When assigned to a static variable, -* this can be used by NNVM and other libraries to unregister callbacks when -* the library is unloaded. This prevents crashes when the underlying IRFunctor -* is destructed as it will no longer contain std::function instances allocated -* by a library that has been unloaded. -*/ -template -class IRFunctorStaticRegistry; - -template -class IRFunctorStaticRegistry { - private: - IRFunctor *irf_; - std::shared_ptr free_list; - - using TSelf = IRFunctorStaticRegistry; - - public: - IRFunctorStaticRegistry(IRFunctor *irf) { - irf_ = irf; - free_list = std::make_shared(); - } - - template - inline TSelf& set_dispatch(std::function f) { // NOLINT(*) - irf_->template set_dispatch(f); - auto irf_copy = irf_; - free_list.get()->append([irf_copy] { - irf_copy->template clear_dispatch(); - }); - return *this; - } -}; - -/*! -* \brief Helper function for constructing an IRFunctorStaticRegistry. This allows -* the compiler to deduce the template types. -*/ -template -IRFunctorStaticRegistry MakeIRFunctorStaticRegistry( - IRFunctor *irf) { - return IRFunctorStaticRegistry(irf); -} - -#define TVM_AUTO_REGISTER_VAR_DEF(ClsName) \ - static TVM_ATTRIBUTE_UNUSED auto __make_functor ## _ ## ClsName - -/*! -* \brief Macro to set IRFunctor dispatch in a global static field using an IRFunctorStaticRegistry. -* Usage is exactly the same as TVM_STATIC_IR_FUNCTOR. Libraries should use this instead of -* TVM_STATIC_IR_FUNCTOR. -*/ -#define TVM_STATIC_IR_FUNCTOR_REGISTER(ClsName, FField) \ - TVM_STR_CONCAT(TVM_AUTO_REGISTER_VAR_DEF(ClsName), __COUNTER__) = \ - MakeIRFunctorStaticRegistry(&ClsName::FField()) - -} // namespace tvm -#endif // TVM_NODE_IR_FUNCTOR_H_ diff --git a/include/tvm/node/reflection.h b/include/tvm/node/reflection.h index e6caa443ab9c..35a8e1d4a657 100644 --- a/include/tvm/node/reflection.h +++ b/include/tvm/node/reflection.h @@ -48,20 +48,20 @@ using runtime::ObjectRef; * Each objects that wants reflection will need to implement * a VisitAttrs function and call visitor->Visit on each of its field. */ -class TVM_DLL AttrVisitor { +class AttrVisitor { public: //! \cond Doxygen_Suppress - virtual ~AttrVisitor() = default; - virtual void Visit(const char* key, double* value) = 0; - virtual void Visit(const char* key, int64_t* value) = 0; - virtual void Visit(const char* key, uint64_t* value) = 0; - virtual void Visit(const char* key, int* value) = 0; - virtual void Visit(const char* key, bool* value) = 0; - virtual void Visit(const char* key, std::string* value) = 0; - virtual void Visit(const char* key, void** value) = 0; - virtual void Visit(const char* key, DataType* value) = 0; - virtual void Visit(const char* key, runtime::NDArray* value) = 0; - virtual void Visit(const char* key, runtime::ObjectRef* value) = 0; + TVM_DLL virtual ~AttrVisitor() = default; + TVM_DLL virtual void Visit(const char* key, double* value) = 0; + TVM_DLL virtual void Visit(const char* key, int64_t* value) = 0; + TVM_DLL virtual void Visit(const char* key, uint64_t* value) = 0; + TVM_DLL virtual void Visit(const char* key, int* value) = 0; + TVM_DLL virtual void Visit(const char* key, bool* value) = 0; + TVM_DLL virtual void Visit(const char* key, std::string* value) = 0; + TVM_DLL virtual void Visit(const char* key, void** value) = 0; + TVM_DLL virtual void Visit(const char* key, DataType* value) = 0; + TVM_DLL virtual void Visit(const char* key, runtime::NDArray* value) = 0; + TVM_DLL virtual void Visit(const char* key, runtime::ObjectRef* value) = 0; template::value>::type> void Visit(const char* key, ENum* ptr) { @@ -93,13 +93,13 @@ class ReflectionVTable { * If this is not empty then FGlobalKey must be defined for the object. * \return The created function. */ - using FCreate = std::function(const std::string& global_key)>; + typedef ObjectPtr (*FCreate)(const std::string& global_key); /*! * \brief Global key function, only needed by global objects. * \param node The node pointer. * \return node The global key to the node. */ - using FGlobalKey = std::function; + typedef std::string (*FGlobalKey)(const Object* self); /*! * \brief Dispatch the VisitAttrs function. * \param self The pointer to the object. @@ -193,7 +193,7 @@ class ReflectionVTable::Registry { static DMLC_ATTRIBUTE_UNUSED ::tvm::ReflectionVTable::Registry & \ __make_Node ## _ ## TypeName ## __ = \ ::tvm::ReflectionVTable::Global()->Register() \ - .set_creator([](const std::string&) { \ + .set_creator([](const std::string&) -> ObjectPtr { \ return ::tvm::runtime::make_object(); \ }) diff --git a/include/tvm/relay/expr_functor.h b/include/tvm/relay/expr_functor.h index 8bc87a27f66f..722f73f03826 100644 --- a/include/tvm/relay/expr_functor.h +++ b/include/tvm/relay/expr_functor.h @@ -25,7 +25,7 @@ #ifndef TVM_RELAY_EXPR_FUNCTOR_H_ #define TVM_RELAY_EXPR_FUNCTOR_H_ -#include +#include #include #include #include @@ -66,7 +66,7 @@ template class ExprFunctor { private: using TSelf = ExprFunctor; - using FType = tvm::IRFunctor; + using FType = tvm::NodeFunctor; public: /*! \brief the result type of this functor */ diff --git a/include/tvm/relay/pattern_functor.h b/include/tvm/relay/pattern_functor.h index c15523cb25de..d84d43af82a7 100644 --- a/include/tvm/relay/pattern_functor.h +++ b/include/tvm/relay/pattern_functor.h @@ -25,7 +25,7 @@ #ifndef TVM_RELAY_PATTERN_FUNCTOR_H_ #define TVM_RELAY_PATTERN_FUNCTOR_H_ -#include +#include #include #include #include @@ -66,7 +66,7 @@ template class PatternFunctor { private: using TSelf = PatternFunctor; - using FType = tvm::IRFunctor; + using FType = tvm::NodeFunctor; public: /*! \brief the result type of this functor */ diff --git a/nnvm/src/compiler/compile_engine.cc b/nnvm/src/compiler/compile_engine.cc index 5ce78d1d58d6..cd84f923e27a 100644 --- a/nnvm/src/compiler/compile_engine.cc +++ b/nnvm/src/compiler/compile_engine.cc @@ -391,8 +391,9 @@ TVM_REGISTER_GLOBAL("nnvm.compiler.CacheItem2ScheduleArgs") TVM_REGISTER_NODE_TYPE(GraphFuncNode); TVM_REGISTER_NODE_TYPE(GraphCacheEntryNode); -TVM_STATIC_IR_FUNCTOR_REGISTER(IRPrinter, vtable) -.set_dispatch([](const GraphFuncNode *op, IRPrinter *p) { +TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) +.set_dispatch([](const ObjectRef& ref, IRPrinter* p) { + auto* op = static_cast(ref.get()); p->stream << "GraphFunc(name=" << op->func_name << ", addr=" << op << ")"; }); diff --git a/nnvm/src/compiler/graph_hash.cc b/nnvm/src/compiler/graph_hash.cc index b76f99fa58d3..bbbc0dbea885 100644 --- a/nnvm/src/compiler/graph_hash.cc +++ b/nnvm/src/compiler/graph_hash.cc @@ -6,9 +6,9 @@ * 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 @@ -101,8 +101,9 @@ GraphKey GraphKeyNode::make(Graph graph, return GraphKey(n); } -TVM_STATIC_IR_FUNCTOR_REGISTER(IRPrinter, vtable) -.set_dispatch([](const GraphKeyNode *op, IRPrinter *p) { +TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) +.set_dispatch([](const ObjectRef& ref, IRPrinter* p) { + auto* op = static_cast(ref.get()); p->stream << "GraphKeyNode("<< op << ")"; }); diff --git a/nnvm/src/compiler/graph_runtime.cc b/nnvm/src/compiler/graph_runtime.cc index d8ff3bf34bf8..a4b398cd41ea 100644 --- a/nnvm/src/compiler/graph_runtime.cc +++ b/nnvm/src/compiler/graph_runtime.cc @@ -30,6 +30,8 @@ namespace nnvm { namespace compiler { +using tvm::Object; +using tvm::ObjectPtr; using tvm::runtime::TVMArgs; using tvm::runtime::TVMRetValue; using tvm::runtime::PackedFunc; diff --git a/src/arithmetic/const_int_bound.cc b/src/arithmetic/const_int_bound.cc index 168486ee0018..d494a50d128d 100644 --- a/src/arithmetic/const_int_bound.cc +++ b/src/arithmetic/const_int_bound.cc @@ -53,7 +53,8 @@ inline void PrintBoundValue(std::ostream& os, int64_t val) { } TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) -.set_dispatch([](const ConstIntBoundNode* op, IRPrinter* p) { +.set_dispatch([](const ObjectRef& node, IRPrinter* p) { + auto* op = static_cast(node.get()); p->stream << "ConstIntBound["; PrintBoundValue(p->stream, op->min_value); p->stream << ','; diff --git a/src/arithmetic/int_set.cc b/src/arithmetic/int_set.cc index 409477578758..9f8effb6c612 100644 --- a/src/arithmetic/int_set.cc +++ b/src/arithmetic/int_set.cc @@ -810,7 +810,8 @@ IntSet EvalSet(Range r, TVM_REGISTER_NODE_TYPE(IntervalSetNode); TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) -.set_dispatch([](const IntervalSetNode *op, IRPrinter *p) { +.set_dispatch([](const ObjectRef& node, IRPrinter *p) { + auto* op = static_cast(node.get()); p->stream << "IntervalSet" << "[" << op->min_value << ", " << op->max_value << ']'; diff --git a/src/arithmetic/modular_set.cc b/src/arithmetic/modular_set.cc index 9e363e7cf99a..25c7391fd9c4 100644 --- a/src/arithmetic/modular_set.cc +++ b/src/arithmetic/modular_set.cc @@ -45,7 +45,8 @@ ModularSet::ModularSet(int64_t coeff, int64_t base) { } TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) -.set_dispatch([](const ModularSetNode *op, IRPrinter *p) { +.set_dispatch([](const ObjectRef& node, IRPrinter *p) { + auto* op = static_cast(node.get()); p->stream << "ModularSet(" << "coeff=" << op->coeff << ", base=" << op->base << ')'; diff --git a/src/codegen/build_module.cc b/src/codegen/build_module.cc index cfcb0607858f..3f279f8772df 100644 --- a/src/codegen/build_module.cc +++ b/src/codegen/build_module.cc @@ -37,8 +37,9 @@ TVM_REGISTER_NODE_TYPE(TargetNode); TVM_REGISTER_NODE_TYPE(GenericFuncNode); TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) -.set_dispatch([](const TargetNode *op, IRPrinter *p) { - p->stream << op->str(); +.set_dispatch([](const ObjectRef& node, IRPrinter *p) { + auto* op = static_cast(node.get()); + p->stream << op->str(); }); @@ -654,7 +655,8 @@ tvm::BuildConfig BuildConfig::Current() { TVM_REGISTER_NODE_TYPE(BuildConfigNode); TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) -.set_dispatch([](const BuildConfigNode *op, IRPrinter *p) { +.set_dispatch([](const ObjectRef& node, IRPrinter *p) { + auto* op = static_cast(node.get()); p->stream << "build_config("; p->stream << "data_alignment=" << op->data_alignment << ", "; p->stream << "offset_factor=" << op->offset_factor << ", "; diff --git a/src/lang/api_registry.cc b/src/lang/api_registry.cc index cd3d43b7dcf3..3c486767c724 100644 --- a/src/lang/api_registry.cc +++ b/src/lang/api_registry.cc @@ -26,11 +26,12 @@ namespace tvm { TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) -.set_dispatch([](const EnvFuncNode *op, IRPrinter *p) { +.set_dispatch([](const ObjectRef& node, IRPrinter *p) { + auto* op = static_cast(node.get()); p->stream << "EnvFunc(" << op->name << ")"; }); -NodePtr CreateEnvNode(const std::string& name) { +ObjectPtr CreateEnvNode(const std::string& name) { auto* f = runtime::Registry::Get(name); CHECK(f != nullptr) << "Cannot find global function \'" << name << '\''; NodePtr n = make_node(); @@ -62,7 +63,7 @@ TVM_REGISTER_API("_EnvFuncGetPackedFunc") TVM_REGISTER_NODE_TYPE(EnvFuncNode) .set_creator(CreateEnvNode) -.set_global_key([](const Object* n) { +.set_global_key([](const Object* n) -> std::string { return static_cast(n)->name; }); diff --git a/src/lang/attr_functor.h b/src/lang/attr_functor.h index b9391e4895b9..51b355e81df3 100644 --- a/src/lang/attr_functor.h +++ b/src/lang/attr_functor.h @@ -18,7 +18,6 @@ */ /*! - * Copyright (c) 2018 by Contributors * \file attr_functor.h * \brief A way to define arbitrary function signature * with dispatch on common attributes. @@ -31,6 +30,7 @@ #ifndef TVM_LANG_ATTR_FUNCTOR_H_ #define TVM_LANG_ATTR_FUNCTOR_H_ +#include #include namespace tvm { @@ -54,7 +54,7 @@ template class AttrFunctor { private: using TSelf = AttrFunctor; - using FType = tvm::IRFunctor; + using FType = tvm::NodeFunctor; public: /*! \brief the result type of this functor */ diff --git a/src/lang/attrs.cc b/src/lang/attrs.cc index a299e17996e0..0b036c31c8e1 100644 --- a/src/lang/attrs.cc +++ b/src/lang/attrs.cc @@ -61,7 +61,8 @@ Attrs DictAttrsNode::make(Map dict) { } TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) -.set_dispatch([](const DictAttrsNode *op, IRPrinter *p) { +.set_dispatch([](const ObjectRef& node, IRPrinter *p) { + auto* op = static_cast(node.get()); p->stream << op->dict; }); diff --git a/src/lang/buffer.cc b/src/lang/buffer.cc index 689b291ae2ed..bc14e2bc009f 100644 --- a/src/lang/buffer.cc +++ b/src/lang/buffer.cc @@ -452,7 +452,8 @@ Buffer BufferNode::make(Var data, } TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) -.set_dispatch([](const BufferNode *op, IRPrinter *p) { +.set_dispatch([](const ObjectRef& node, IRPrinter *p) { + auto* op = static_cast(node.get()); p->stream << "buffer(" << op->name << ", " << op << ")"; }); diff --git a/src/lang/channel.cc b/src/lang/channel.cc index 6746a3c5ecb5..c564d61477b7 100644 --- a/src/lang/channel.cc +++ b/src/lang/channel.cc @@ -6,9 +6,9 @@ * 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 @@ -33,7 +33,8 @@ Channel ChannelNode::make(Var handle_var, Type dtype) { } TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) -.set_dispatch([](const ChannelNode *op, IRPrinter *p) { +.set_dispatch([](const ObjectRef& node, IRPrinter *p) { + auto* op = static_cast(node.get()); p->stream << "channel(" << op->handle_var << ", " << op->dtype << ")"; }); diff --git a/src/lang/data_layout.cc b/src/lang/data_layout.cc index 3686d5f887b8..7c76e40bf01c 100644 --- a/src/lang/data_layout.cc +++ b/src/lang/data_layout.cc @@ -196,7 +196,8 @@ int32_t Layout::FactorOf(const LayoutAxis& axis) const { } TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) -.set_dispatch([](const LayoutNode* l, IRPrinter* p) { +.set_dispatch([](const ObjectRef& node, IRPrinter* p) { + auto* l = static_cast(node.get()); p->stream << "Layout(" << l->name << ")"; }); @@ -352,7 +353,8 @@ BijectiveLayout BijectiveLayoutNode::make(const Layout& src_layout, } TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) -.set_dispatch([](const BijectiveLayoutNode* b, IRPrinter* p) { +.set_dispatch([](const ObjectRef& node, IRPrinter* p) { + auto* b = static_cast(node.get()); p->stream << "BijectiveLayout(" << b->src_layout.name() << "->" << b->dst_layout.name() << ")"; }); diff --git a/src/lang/expr.cc b/src/lang/expr.cc index 31ade90dd587..6a69fdaa20c4 100644 --- a/src/lang/expr.cc +++ b/src/lang/expr.cc @@ -182,7 +182,8 @@ IRPrinter::FType& IRPrinter::vtable() { } TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) -.set_dispatch([](const IntImm *op, IRPrinter *p) { +.set_dispatch([](const ObjectRef& node, IRPrinter* p) { + auto* op = static_cast(node.get()); if (op->type == Int(32)) { p->stream << op->value; } else { @@ -191,7 +192,8 @@ TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) }); TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) -.set_dispatch([](const IterVarNode *op, IRPrinter *p) { +.set_dispatch([](const ObjectRef& node, IRPrinter* p) { + auto* op = static_cast(node.get()); p->stream << "iter_var("; if (op->var->name_hint.length() != 0) { p->stream << op->var->name_hint << ", "; @@ -206,7 +208,8 @@ TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) }); TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) -.set_dispatch([](const RangeNode* op, IRPrinter* p) { +.set_dispatch([](const ObjectRef& node, IRPrinter* p) { + auto* op = static_cast(node.get()); p->stream << "range(min=" << op->min << ", ext=" << op->extent << ')'; }); diff --git a/src/lang/ir.cc b/src/lang/ir.cc index 04e04aef455c..bb8401dae843 100644 --- a/src/lang/ir.cc +++ b/src/lang/ir.cc @@ -553,12 +553,14 @@ Stmt Evaluate::make(Expr value) { // Printers TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) -.set_dispatch([](const UIntImm* op, IRPrinter* p) { +.set_dispatch([](const ObjectRef& node, IRPrinter* p) { + auto* op = static_cast(node.get()); p->stream << "(" << op->type << ")" << op->value; }); TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) -.set_dispatch([](const FloatImm* op, IRPrinter* p) { +.set_dispatch([](const ObjectRef& node, IRPrinter* p) { + auto* op = static_cast(node.get()); auto& stream = p->stream; switch (op->type.bits()) { case 64: @@ -576,7 +578,8 @@ TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) }); TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) -.set_dispatch([](const StringImm* op, IRPrinter* p) { +.set_dispatch([](const ObjectRef& node, IRPrinter* p) { + auto* op = static_cast(node.get()); auto& stream = p->stream; stream << '"'; for (size_t i = 0; i < op->value.size(); ++i) { @@ -611,101 +614,116 @@ TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) }); TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) -.set_dispatch([](const Cast* op, IRPrinter* p) { +.set_dispatch([](const ObjectRef& node, IRPrinter* p) { + auto* op = static_cast(node.get()); p->stream << op->type << '('; p->Print(op->value); p->stream << ')'; }) -.set_dispatch([](const Variable* op, IRPrinter* p) { +.set_dispatch([](const ObjectRef& node, IRPrinter* p) { + auto* op = static_cast(node.get()); // omit the type // stream << op->name << "." << op->type; p->stream << op->name_hint; }) -.set_dispatch([](const Add* op, IRPrinter* p) { +.set_dispatch([](const ObjectRef& node, IRPrinter* p) { + auto* op = static_cast(node.get()); p->stream << '('; p->Print(op->a); p->stream << " + "; p->Print(op->b); p->stream << ')'; }) -.set_dispatch([](const Sub* op, IRPrinter* p) { +.set_dispatch([](const ObjectRef& node, IRPrinter* p) { + auto* op = static_cast(node.get()); p->stream << '('; p->Print(op->a); p->stream << " - "; p->Print(op->b); p->stream << ')'; }) -.set_dispatch([](const Mul* op, IRPrinter* p) { +.set_dispatch([](const ObjectRef& node, IRPrinter* p) { + auto* op = static_cast(node.get()); p->stream << '('; p->Print(op->a); p->stream << "*"; p->Print(op->b); p->stream << ')'; }) -.set_dispatch
([](const Div* op, IRPrinter* p) { +.set_dispatch
([](const ObjectRef& node, IRPrinter* p) { + auto* op = static_cast(node.get()); p->stream << '('; p->Print(op->a); p->stream << "/"; p->Print(op->b); p->stream << ')'; }) -.set_dispatch([](const Mod* op, IRPrinter* p) { +.set_dispatch([](const ObjectRef& node, IRPrinter* p) { + auto* op = static_cast(node.get()); p->stream << '('; p->Print(op->a); p->stream << " % "; p->Print(op->b); p->stream << ')'; }) -.set_dispatch([](const Min* op, IRPrinter* p) { +.set_dispatch([](const ObjectRef& node, IRPrinter* p) { + auto* op = static_cast(node.get()); p->stream << "min("; p->Print(op->a); p->stream << ", "; p->Print(op->b); p->stream << ")"; }) -.set_dispatch([](const Max* op, IRPrinter* p) { +.set_dispatch([](const ObjectRef& node, IRPrinter* p) { + auto* op = static_cast(node.get()); p->stream << "max("; p->Print(op->a); p->stream << ", "; p->Print(op->b); p->stream << ")"; }) -.set_dispatch([](const EQ* op, IRPrinter* p) { +.set_dispatch([](const ObjectRef& node, IRPrinter* p) { + auto* op = static_cast(node.get()); p->stream << '('; p->Print(op->a); p->stream << " == "; p->Print(op->b); p->stream << ')'; }) -.set_dispatch([](const NE* op, IRPrinter* p) { +.set_dispatch([](const ObjectRef& node, IRPrinter* p) { + auto* op = static_cast(node.get()); p->stream << '('; p->Print(op->a); p->stream << " != "; p->Print(op->b); p->stream << ')'; }) -.set_dispatch([](const LT* op, IRPrinter* p) { +.set_dispatch([](const ObjectRef& node, IRPrinter* p) { + auto* op = static_cast(node.get()); p->stream << '('; p->Print(op->a); p->stream << " < "; p->Print(op->b); p->stream << ')'; }) -.set_dispatch([](const LE* op, IRPrinter* p) { +.set_dispatch([](const ObjectRef& node, IRPrinter* p) { + auto* op = static_cast(node.get()); p->stream << '('; p->Print(op->a); p->stream << " <= "; p->Print(op->b); p->stream << ')'; }) -.set_dispatch([](const GT* op, IRPrinter* p) { +.set_dispatch([](const ObjectRef& node, IRPrinter* p) { + auto* op = static_cast(node.get()); p->stream << '('; p->Print(op->a); p->stream << " > "; p->Print(op->b); p->stream << ')'; }) -.set_dispatch([](const GE* op, IRPrinter* p) { +.set_dispatch([](const ObjectRef& node, IRPrinter* p) { + auto* op = static_cast(node.get()); p->stream << '('; p->Print(op->a); p->stream << " >= "; @@ -714,17 +732,20 @@ TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) }); TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) -.set_dispatch([](const FloorDiv* op, IRPrinter *p) { +.set_dispatch([](const ObjectRef& node, IRPrinter* p) { + auto* op = static_cast(node.get()); p->stream << "floordiv(" << op->a << ", " << op->b << ")"; }); TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) -.set_dispatch([](const FloorMod* op, IRPrinter *p) { +.set_dispatch([](const ObjectRef& node, IRPrinter* p) { + auto* op = static_cast(node.get()); p->stream << "floormod(" << op->a << ", " << op->b << ")"; }); TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) -.set_dispatch([](const And* op, IRPrinter* p) { +.set_dispatch([](const ObjectRef& node, IRPrinter* p) { + auto* op = static_cast(node.get()); p->stream << '('; p->Print(op->a); p->stream << " && "; @@ -733,7 +754,8 @@ TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) }); TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) -.set_dispatch([](const Or* op, IRPrinter* p) { +.set_dispatch([](const ObjectRef& node, IRPrinter* p) { + auto* op = static_cast(node.get()); p->stream << '('; p->Print(op->a); p->stream << " || "; @@ -742,13 +764,15 @@ TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) }); TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) -.set_dispatch([](const Not* op, IRPrinter* p) { +.set_dispatch([](const ObjectRef& node, IRPrinter* p) { + auto* op = static_cast(node.get()); p->stream << '!'; p->Print(op->a); }); TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) -.set_dispatch([](const ObjectRef& node, IRPrinter* p) { + auto* op = static_cast(node.get()); p->stream << "select("; p->Print(op->condition); p->stream << ", "; @@ -759,7 +783,8 @@ TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) }); TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) -.set_dispatch([](const Load* op, IRPrinter* p) { +.set_dispatch([](const ObjectRef& node, IRPrinter* p) { + auto* op = static_cast(node.get()); p->stream << op->buffer_var << "["; p->Print(op->index); p->stream << "]"; @@ -770,7 +795,8 @@ TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) }); TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) -.set_dispatch([](const Ramp* op, IRPrinter* p) { +.set_dispatch([](const ObjectRef& node, IRPrinter* p) { + auto* op = static_cast(node.get()); p->stream << "ramp("; p->Print(op->base); p->stream << ", "; @@ -779,14 +805,16 @@ TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) }); TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) -.set_dispatch([](const Broadcast* op, IRPrinter* p) { +.set_dispatch([](const ObjectRef& node, IRPrinter* p) { + auto* op = static_cast(node.get()); p->stream << "x" << op->lanes << "("; p->Print(op->value); p->stream << ")"; }); TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) -.set_dispatch([](const Call* op, IRPrinter* p) { +.set_dispatch([](const ObjectRef& node, IRPrinter* p) { + auto* op = static_cast(node.get()); p->stream << op->name << "("; for (size_t i = 0; i < op->args.size(); ++i) { p->Print(op->args[i]); @@ -798,7 +826,8 @@ TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) }); TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) -.set_dispatch([](const Let* op, IRPrinter* p) { +.set_dispatch([](const ObjectRef& node, IRPrinter* p) { + auto* op = static_cast(node.get()); p->stream << "(let " << op->var << " = "; p->Print(op->value); p->stream << " in "; @@ -807,7 +836,8 @@ TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) }); TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) -.set_dispatch([](const LetStmt* op, IRPrinter* p) { +.set_dispatch([](const ObjectRef& node, IRPrinter* p) { + auto* op = static_cast(node.get()); p->PrintIndent(); p->stream << "let " << op->var << " = "; p->Print(op->value); @@ -816,7 +846,8 @@ TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) }); TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) -.set_dispatch([](const AttrStmt* op, IRPrinter* p) { +.set_dispatch([](const ObjectRef& node, IRPrinter* p) { + auto* op = static_cast(node.get()); p->PrintIndent(); p->stream << "// attr ["; p->Print(op->node); @@ -828,7 +859,8 @@ TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) }); TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) -.set_dispatch([](const AssertStmt* op, IRPrinter* p) { +.set_dispatch([](const ObjectRef& node, IRPrinter* p) { + auto* op = static_cast(node.get()); p->PrintIndent(); p->stream << "assert("; p->Print(op->condition); @@ -839,7 +871,8 @@ TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) }); TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) -.set_dispatch([](const ProducerConsumer* op, IRPrinter* p) { +.set_dispatch([](const ObjectRef& node, IRPrinter* p) { + auto* op = static_cast(node.get()); if (op->is_producer) { p->PrintIndent(); p->stream << "produce " << op->func->func_name() << " {\n"; @@ -872,7 +905,8 @@ std::ostream &operator<<(std::ostream& out, ForType type) { // NOLINT(*) } TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) -.set_dispatch([](const For* op, IRPrinter* p) { +.set_dispatch([](const ObjectRef& node, IRPrinter* p) { + auto* op = static_cast(node.get()); p->PrintIndent(); p->stream << op->for_type << " (" << op->loop_var << ", "; p->Print(op->min); @@ -889,7 +923,8 @@ TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) }); TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) -.set_dispatch([](const Store* op, IRPrinter* p) { +.set_dispatch([](const ObjectRef& node, IRPrinter* p) { + auto* op = static_cast(node.get()); p->PrintIndent(); p->stream << op->buffer_var << "["; p->Print(op->index); @@ -903,7 +938,8 @@ TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) }); TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) -.set_dispatch([](const Provide* op, IRPrinter* p) { +.set_dispatch([](const ObjectRef& node, IRPrinter* p) { + auto* op = static_cast(node.get()); p->PrintIndent(); p->stream << op->func->func_name() << "("; for (size_t i = 0; i < op->args.size(); ++i) { @@ -920,7 +956,8 @@ TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) }); TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) -.set_dispatch([](const Allocate* op, IRPrinter* p) { +.set_dispatch([](const ObjectRef& node, IRPrinter* p) { + auto* op = static_cast(node.get()); p->PrintIndent(); p->stream << "allocate " << op->buffer_var << "[" << op->type; for (size_t i = 0; i < op->extents.size(); ++i) { @@ -937,14 +974,16 @@ TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) }); TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) -.set_dispatch([](const Free* op, IRPrinter* p) { +.set_dispatch([](const ObjectRef& node, IRPrinter* p) { + auto* op = static_cast(node.get()); p->PrintIndent(); p->stream << "free " << op->buffer_var; p->stream << '\n'; }); TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) -.set_dispatch([](const Realize* op, IRPrinter* p) { +.set_dispatch([](const ObjectRef& node, IRPrinter* p) { + auto* op = static_cast(node.get()); p->PrintIndent(); p->stream << "realize " << op->func->func_name() << "("; for (size_t i = 0; i < op->bounds.size(); ++i) { @@ -974,7 +1013,8 @@ TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) }); TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) -.set_dispatch([](const Prefetch* op, IRPrinter* p) { +.set_dispatch([](const ObjectRef& node, IRPrinter* p) { + auto* op = static_cast(node.get()); p->PrintIndent(); p->stream << "prefetch " << op->func->func_name() << "("; for (size_t i = 0; i < op->bounds.size(); ++i) { @@ -992,13 +1032,15 @@ TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) }); TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) -.set_dispatch([](const Block* op, IRPrinter* p) { +.set_dispatch([](const ObjectRef& node, IRPrinter* p) { + auto* op = static_cast(node.get()); p->Print(op->first); if (op->rest.defined()) p->Print(op->rest); }); TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) -.set_dispatch([](const IfThenElse* op, IRPrinter* p) { +.set_dispatch([](const ObjectRef& node, IRPrinter* p) { + auto* op = static_cast(node.get()); p->PrintIndent(); while (true) { p->stream << "if (" << op->condition << ") {\n"; @@ -1028,7 +1070,8 @@ TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) }); TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) -.set_dispatch([](const Evaluate* op, IRPrinter* p) { +.set_dispatch([](const ObjectRef& node, IRPrinter* p) { + auto* op = static_cast(node.get()); p->PrintIndent(); p->Print(op->value); p->stream << "\n"; @@ -1045,7 +1088,8 @@ void PrintList(const Array &exprs, IRPrinter* p) { } TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) -.set_dispatch([](const Shuffle* op, IRPrinter* p) { +.set_dispatch([](const ObjectRef& node, IRPrinter* p) { + auto* op = static_cast(node.get()); p->stream << "shuffle("; PrintList(op->vectors, p); p->stream << ", "; @@ -1055,7 +1099,8 @@ TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) // Container printer TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) -.set_dispatch([](const ArrayNode* op, IRPrinter* p) { +.set_dispatch([](const ObjectRef& node, IRPrinter* p) { + auto* op = static_cast(node.get()); p->stream << '['; for (size_t i = 0 ; i < op->data.size(); ++i) { if (i != 0) { @@ -1067,7 +1112,8 @@ TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) }); TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) -.set_dispatch([](const MapNode* op, IRPrinter* p) { +.set_dispatch([](const ObjectRef& node, IRPrinter* p) { + auto* op = static_cast(node.get()); p->stream << '{'; for (auto it = op->data.begin(); it != op->data.end(); ++it) { if (it != op->data.begin()) { @@ -1081,7 +1127,8 @@ TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) }); TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) -.set_dispatch([](const StrMapNode* op, IRPrinter* p) { +.set_dispatch([](const ObjectRef& node, IRPrinter* p) { + auto* op = static_cast(node.get()); p->stream << '{'; for (auto it = op->data.begin(); it != op->data.end(); ++it) { if (it != op->data.begin()) { @@ -1094,7 +1141,8 @@ TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) }); TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) -.set_dispatch([](const Reduce* op, IRPrinter* p) { +.set_dispatch([](const ObjectRef& node, IRPrinter* p) { + auto* op = static_cast(node.get()); p->stream << "reduce(combiner=" << op->combiner; p->stream << ", source=" << op->source; @@ -1105,7 +1153,8 @@ TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) }); TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) -.set_dispatch([](const CommReducerNode* op, IRPrinter* p) { +.set_dispatch([](const ObjectRef& node, IRPrinter* p) { + auto* op = static_cast(node.get()); p->stream << "comm_reducer(result=" << op->result << ", lhs=" << op->lhs << ", rhs=" << op->rhs @@ -1114,8 +1163,8 @@ TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) }); TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) -.set_dispatch([](const Any *op, IRPrinter *p) { - p->stream << "?"; +.set_dispatch([](const ObjectRef& node, IRPrinter* p) { + p->stream << "?"; }); TVM_REGISTER_NODE_TYPE(CommReducerNode); diff --git a/src/lang/lowered_func.cc b/src/lang/lowered_func.cc index 626b9f7d8d80..cb1ee0585a0c 100644 --- a/src/lang/lowered_func.cc +++ b/src/lang/lowered_func.cc @@ -6,9 +6,9 @@ * 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 @@ -26,7 +26,8 @@ namespace tvm { TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) -.set_dispatch([](const LoweredFuncNode *op, IRPrinter *p) { +.set_dispatch([](const ObjectRef& node, IRPrinter *p) { + auto* op = static_cast(node.get()); p->stream << "LoweredFunc(" << op->name << ", " << op << ")"; }); diff --git a/src/lang/target_info.cc b/src/lang/target_info.cc index 481a9269193b..8c45a19cf818 100644 --- a/src/lang/target_info.cc +++ b/src/lang/target_info.cc @@ -27,7 +27,8 @@ namespace tvm { TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) -.set_dispatch([](const MemoryInfoNode *op, IRPrinter *p) { +.set_dispatch([](const ObjectRef& node, IRPrinter *p) { + auto* op = static_cast(node.get()); p->stream << "mem-info(" << "unit_bits=" << op->unit_bits << ", " << "max_num_bits=" << op->max_num_bits << ", " diff --git a/src/lang/tensor.cc b/src/lang/tensor.cc index 1ac564293c28..db90e4e021ba 100644 --- a/src/lang/tensor.cc +++ b/src/lang/tensor.cc @@ -69,7 +69,8 @@ Tensor TensorNode::make(Array shape, } TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) -.set_dispatch([](const TensorNode *t, IRPrinter *p) { +.set_dispatch([](const ObjectRef& node, IRPrinter *p) { + auto* t = static_cast(node.get()); p->stream << "Tensor(shape=" << t->shape << ", op.name=" << t->op->name << ')'; }); @@ -100,8 +101,9 @@ TensorIntrin TensorIntrinNode::make(std::string name, } TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) -.set_dispatch([](const TensorIntrinNode *n, IRPrinter *p) { - p->stream << "TensorIntrin(name=" << n->name << ", " << n << ")"; +.set_dispatch([](const ObjectRef& node, IRPrinter* p) { + auto* op = static_cast(node.get()); + p->stream << "TensorIntrin(name=" << op->name << ", " << op << ")"; }); TVM_REGISTER_NODE_TYPE(TensorIntrinNode); @@ -124,7 +126,8 @@ TensorIntrinCall TensorIntrinCallNode::make(TensorIntrin intrin, } TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) -.set_dispatch([](const TensorIntrinCallNode *n, IRPrinter *p) { +.set_dispatch([](const ObjectRef& node, IRPrinter *p) { + auto* n = static_cast(node.get()); p->stream << "TensorIntrinCall(intrin=" << n->intrin << ", " << n << ")"; }); diff --git a/src/op/compute_op.cc b/src/op/compute_op.cc index 69589423b663..5f5d2d4f475b 100644 --- a/src/op/compute_op.cc +++ b/src/op/compute_op.cc @@ -40,7 +40,8 @@ namespace tvm { using namespace ir; TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) -.set_dispatch([](const ComputeOpNode *op, IRPrinter *p) { +.set_dispatch([](const ObjectRef& node, IRPrinter* p) { + auto* op = static_cast(node.get()); p->stream << "compute(" << op->name << ", " << op << ")"; }); diff --git a/src/op/extern_op.cc b/src/op/extern_op.cc index 9f3341536ddd..35fe469fbe16 100644 --- a/src/op/extern_op.cc +++ b/src/op/extern_op.cc @@ -31,7 +31,8 @@ namespace tvm { using namespace ir; // ExternOpNode TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) -.set_dispatch([](const ExternOpNode *op, IRPrinter *p) { +.set_dispatch([](const ObjectRef& node, IRPrinter* p) { + auto* op = static_cast(node.get()); p->stream << "extern(" << op->name << ", " << op << ")"; }); diff --git a/src/op/hybrid_op.cc b/src/op/hybrid_op.cc index e6a46fe19846..7a99ea10b74d 100644 --- a/src/op/hybrid_op.cc +++ b/src/op/hybrid_op.cc @@ -37,7 +37,8 @@ namespace tvm { using namespace ir; // HybridOpNode TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) -.set_dispatch([](const HybridOpNode *op, IRPrinter *p) { +.set_dispatch([](const ObjectRef& node, IRPrinter* p) { + auto* op = static_cast(node.get()); p->stream << "hybrid(" << op->name << ", " << op << ")"; }); diff --git a/src/op/placeholder_op.cc b/src/op/placeholder_op.cc index 97d01ca063f1..4d08fa316ae8 100644 --- a/src/op/placeholder_op.cc +++ b/src/op/placeholder_op.cc @@ -6,9 +6,9 @@ * 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 @@ -28,7 +28,8 @@ namespace tvm { // PlaceholderOpNode TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) -.set_dispatch([](const PlaceholderOpNode *op, IRPrinter *p) { +.set_dispatch([](const ObjectRef& node, IRPrinter* p) { + auto* op = static_cast(node.get()); p->stream << "placeholder(" << op->name << ", " << op << ")"; }); diff --git a/src/op/scan_op.cc b/src/op/scan_op.cc index 7b7a47c61c26..b02073b5357e 100644 --- a/src/op/scan_op.cc +++ b/src/op/scan_op.cc @@ -32,7 +32,8 @@ namespace tvm { using namespace ir; TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) -.set_dispatch([](const ScanOpNode *op, IRPrinter *p) { +.set_dispatch([](const ObjectRef& node, IRPrinter* p) { + auto* op = static_cast(node.get()); p->stream << "scan(" << op->name << ", " << op << ")"; }); TVM_REGISTER_NODE_TYPE(ScanOpNode); diff --git a/src/op/tensor_compute_op.cc b/src/op/tensor_compute_op.cc index d333461c14b5..6533b0e8657a 100644 --- a/src/op/tensor_compute_op.cc +++ b/src/op/tensor_compute_op.cc @@ -6,9 +6,9 @@ * 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 @@ -36,8 +36,8 @@ namespace tvm { using namespace ir; // TensorComputeOpNode TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) -.set_dispatch([](const TensorComputeOpNode *op, - IRPrinter *p) { +.set_dispatch([](const ObjectRef& node, IRPrinter* p) { + auto* op = static_cast(node.get()); p->stream << "tensor_compute_op(" << op->name << ", " << op << ")"; }); diff --git a/src/pass/ir_mutator.cc b/src/pass/ir_mutator.cc index fda12378a766..c8e46c9129a5 100644 --- a/src/pass/ir_mutator.cc +++ b/src/pass/ir_mutator.cc @@ -118,9 +118,9 @@ inline Array MutateIterVarArr(Array rdom, IRMutator *m) { // Mutate Stmt -#define DISPATCH_TO_MUTATE_STMT(OP) \ - set_dispatch([](const OP* op, const Stmt& s, IRMutator* m) { \ - return m->Mutate_(op, s); \ +#define DISPATCH_TO_MUTATE_STMT(OP) \ + set_dispatch([](const ObjectRef& node, const Stmt& s, IRMutator* m) { \ + return m->Mutate_(static_cast(node.get()), s); \ }) Stmt IRMutator::Mutate_(const AttrStmt* op, const Stmt& s) { @@ -344,9 +344,9 @@ TVM_STATIC_IR_FUNCTOR(IRMutator, vtable_stmt) // Mutate Expr -#define DISPATCH_TO_MUTATE_EXPR(OP) \ - set_dispatch([](const OP* op, const Expr& e, IRMutator* m) { \ - return m->Mutate_(op, e); \ +#define DISPATCH_TO_MUTATE_EXPR(OP) \ + set_dispatch([](const ObjectRef& node, const Expr& e, IRMutator* m) { \ + return m->Mutate_(static_cast(node.get()), e); \ }) Expr IRMutator::Mutate_(const Variable *op, const Expr& e) { diff --git a/src/pass/ir_visitor.cc b/src/pass/ir_visitor.cc index fde183e0c41a..38c8490c0560 100644 --- a/src/pass/ir_visitor.cc +++ b/src/pass/ir_visitor.cc @@ -237,9 +237,9 @@ DEFINE_OP_NO_VISIT_(UIntImm) DEFINE_OP_NO_VISIT_(FloatImm) DEFINE_OP_NO_VISIT_(StringImm) -#define DISPATCH_TO_VISIT(OP) \ - set_dispatch([](const OP* op, IRVisitor* v) { \ - v->Visit_(op); \ +#define DISPATCH_TO_VISIT(OP) \ + set_dispatch([](const ObjectRef& node, IRVisitor* v) { \ + v->Visit_(static_cast(node.get())); \ }) TVM_STATIC_IR_FUNCTOR(IRVisitor, vtable) diff --git a/src/relay/backend/graph_runtime_codegen.cc b/src/relay/backend/graph_runtime_codegen.cc index 7ec287b0e0a2..0342aa6ab1ba 100644 --- a/src/relay/backend/graph_runtime_codegen.cc +++ b/src/relay/backend/graph_runtime_codegen.cc @@ -24,7 +24,6 @@ #include #include -#include #include #include diff --git a/src/relay/backend/interpreter.cc b/src/relay/backend/interpreter.cc index 8c6daceedd5c..130fd4b8ce40 100644 --- a/src/relay/backend/interpreter.cc +++ b/src/relay/backend/interpreter.cc @@ -53,8 +53,9 @@ Closure ClosureNode::make(tvm::Map env, Function func) { TVM_REGISTER_API("relay._make.Closure") .set_body_typed(ClosureNode::make); -TVM_STATIC_IR_FUNCTOR_REGISTER(IRPrinter, vtable) -.set_dispatch([](const ClosureNode* node, tvm::IRPrinter* p) { +TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) +.set_dispatch([](const ObjectRef& ref, IRPrinter* p) { + auto* node = static_cast(ref.get()); p->stream << "ClosureNode(" << node->func << ", " << node->env << ")"; }); @@ -71,10 +72,11 @@ RecClosure RecClosureNode::make(Closure clos, Var bind) { TVM_REGISTER_API("relay._make.RecClosure") .set_body_typed(RecClosureNode::make); -TVM_STATIC_IR_FUNCTOR_REGISTER(IRPrinter, vtable) -.set_dispatch([](const RecClosureNode* node, tvm::IRPrinter* p) { - p->stream << "RecClosureNode(" << node->clos << ")"; - }); +TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) +.set_dispatch([](const ObjectRef& ref, IRPrinter* p) { + auto* node = static_cast(ref.get()); + p->stream << "RecClosureNode(" << node->clos << ")"; + }); TupleValue TupleValueNode::make(tvm::Array value) { NodePtr n = make_node(); @@ -85,8 +87,9 @@ TupleValue TupleValueNode::make(tvm::Array value) { TVM_REGISTER_API("relay._make.TupleValue") .set_body_typed(TupleValueNode::make); -TVM_STATIC_IR_FUNCTOR_REGISTER(IRPrinter, vtable) -.set_dispatch([](const TupleValueNode* node, tvm::IRPrinter* p) { +TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) +.set_dispatch([](const ObjectRef& ref, IRPrinter* p) { + auto* node = static_cast(ref.get()); p->stream << "TupleValueNode(" << node->fields << ")"; }); @@ -96,8 +99,9 @@ TensorValue TensorValueNode::make(runtime::NDArray data) { return TensorValue(n); } -TVM_STATIC_IR_FUNCTOR_REGISTER(IRPrinter, vtable) -.set_dispatch([](const TensorValueNode* node, tvm::IRPrinter* p) { +TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) +.set_dispatch([](const ObjectRef& ref, IRPrinter* p) { + auto* node = static_cast(ref.get()); auto to_str = GetPackedFunc("relay._tensor_value_repr"); std::string data_str = to_str(GetRef(node)); p->stream << "TensorValueNode(" << data_str << ")"; @@ -117,11 +121,11 @@ TVM_REGISTER_API("relay._make.RefValue") TVM_REGISTER_NODE_TYPE(RefValueNode); -TVM_STATIC_IR_FUNCTOR_REGISTER(IRPrinter, vtable) -.set_dispatch([](const RefValueNode* node, - tvm::IRPrinter* p) { - p->stream << "RefValueNode(" << node->value << ")"; - }); +TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) +.set_dispatch([](const ObjectRef& ref, IRPrinter* p) { + auto* node = static_cast(ref.get()); + p->stream << "RefValueNode(" << node->value << ")"; + }); ConstructorValue ConstructorValueNode::make(int32_t tag, tvm::Array fields, @@ -138,9 +142,9 @@ TVM_REGISTER_API("relay._make.ConstructorValue") TVM_REGISTER_NODE_TYPE(ConstructorValueNode); -TVM_STATIC_IR_FUNCTOR_REGISTER(IRPrinter, vtable) -.set_dispatch([](const ConstructorValueNode* node, - tvm::IRPrinter* p) { +TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) +.set_dispatch([](const ObjectRef& ref, IRPrinter* p) { + auto* node = static_cast(ref.get()); p->stream << "ConstructorValueNode(" << node->tag << "," << node->fields << ")"; }); diff --git a/src/relay/ir/adt.cc b/src/relay/ir/adt.cc index 12cebe5f5d3c..1f51ecc84fdc 100644 --- a/src/relay/ir/adt.cc +++ b/src/relay/ir/adt.cc @@ -37,9 +37,8 @@ TVM_REGISTER_NODE_TYPE(PatternWildcardNode); TVM_REGISTER_API("relay._make.PatternWildcard") .set_body_typed(PatternWildcardNode::make); -TVM_STATIC_IR_FUNCTOR_REGISTER(IRPrinter, vtable) -.set_dispatch([](const PatternWildcardNode* node, - tvm::IRPrinter* p) { +TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) +.set_dispatch([](const ObjectRef& ref, IRPrinter* p) { p->stream << "PatternWildcardNode()"; }); @@ -54,9 +53,9 @@ TVM_REGISTER_NODE_TYPE(PatternVarNode); TVM_REGISTER_API("relay._make.PatternVar") .set_body_typed(PatternVarNode::make); -TVM_STATIC_IR_FUNCTOR_REGISTER(IRPrinter, vtable) -.set_dispatch([](const PatternVarNode* node, - tvm::IRPrinter* p) { +TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) +.set_dispatch([](const ObjectRef& ref, IRPrinter* p) { + auto* node = static_cast(ref.get()); p->stream << "PatternVarNode(" << node->var << ")"; }); @@ -73,9 +72,9 @@ TVM_REGISTER_NODE_TYPE(PatternConstructorNode); TVM_REGISTER_API("relay._make.PatternConstructor") .set_body_typed(PatternConstructorNode::make); -TVM_STATIC_IR_FUNCTOR_REGISTER(IRPrinter, vtable) -.set_dispatch([](const PatternConstructorNode* node, - tvm::IRPrinter* p) { +TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) +.set_dispatch([](const ObjectRef& ref, IRPrinter* p) { + auto* node = static_cast(ref.get()); p->stream << "PatternConstructorNode(" << node->constructor << ", " << node->patterns << ")"; }); @@ -91,9 +90,9 @@ TVM_REGISTER_NODE_TYPE(PatternTupleNode); TVM_REGISTER_API("relay._make.PatternTuple") .set_body_typed(PatternTupleNode::make); -TVM_STATIC_IR_FUNCTOR_REGISTER(IRPrinter, vtable) -.set_dispatch([](const PatternTupleNode* node, - tvm::IRPrinter* p) { +TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) +.set_dispatch([](const ObjectRef& ref, IRPrinter* p) { + auto* node = static_cast(ref.get()); p->stream << "PatternTupleNode(" << node->patterns << ")"; }); @@ -112,9 +111,9 @@ TVM_REGISTER_NODE_TYPE(ConstructorNode); TVM_REGISTER_API("relay._make.Constructor") .set_body_typed(ConstructorNode::make); -TVM_STATIC_IR_FUNCTOR_REGISTER(IRPrinter, vtable) -.set_dispatch([](const ConstructorNode* node, - tvm::IRPrinter* p) { +TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) +.set_dispatch([](const ObjectRef& ref, IRPrinter* p) { + auto* node = static_cast(ref.get()); p->stream << "ConstructorNode(" << node->name_hint << ", " << node->inputs << ", " << node->belong_to << ")"; }); @@ -134,9 +133,9 @@ TVM_REGISTER_NODE_TYPE(TypeDataNode); TVM_REGISTER_API("relay._make.TypeData") .set_body_typed(TypeDataNode::make); -TVM_STATIC_IR_FUNCTOR_REGISTER(IRPrinter, vtable) -.set_dispatch([](const TypeDataNode* node, - tvm::IRPrinter* p) { +TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) +.set_dispatch([](const ObjectRef& ref, IRPrinter* p) { + auto* node = static_cast(ref.get()); p->stream << "TypeDataNode(" << node->header << ", " << node->type_vars << ", " << node->constructors << ")"; }); @@ -153,9 +152,9 @@ TVM_REGISTER_NODE_TYPE(ClauseNode); TVM_REGISTER_API("relay._make.Clause") .set_body_typed(ClauseNode::make); -TVM_STATIC_IR_FUNCTOR_REGISTER(IRPrinter, vtable) -.set_dispatch([](const ClauseNode* node, - tvm::IRPrinter* p) { +TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) +.set_dispatch([](const ObjectRef& ref, IRPrinter* p) { + auto* node = static_cast(ref.get()); p->stream << "ClauseNode(" << node->lhs << ", " << node->rhs << ")"; }); @@ -173,9 +172,9 @@ TVM_REGISTER_NODE_TYPE(MatchNode); TVM_REGISTER_API("relay._make.Match") .set_body_typed(MatchNode::make); -TVM_STATIC_IR_FUNCTOR_REGISTER(IRPrinter, vtable) -.set_dispatch([](const MatchNode* node, - tvm::IRPrinter* p) { +TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) +.set_dispatch([](const ObjectRef& ref, IRPrinter* p) { + auto* node = static_cast(ref.get()); p->stream << "MatchNode(" << node->data << ", " << node->clauses << ", " << node->complete << ")"; }); diff --git a/src/relay/ir/base.cc b/src/relay/ir/base.cc index 80f07904662f..3bc916d9a406 100644 --- a/src/relay/ir/base.cc +++ b/src/relay/ir/base.cc @@ -18,7 +18,6 @@ */ /*! - * Copyright (c) 2018 by Contributors * \file base.cc * \brief The core base types for Relay. */ @@ -31,7 +30,7 @@ namespace relay { using tvm::IRPrinter; using namespace tvm::runtime; -NodePtr GetSourceNameNode(const std::string& name) { +ObjectPtr GetSourceNameNode(const std::string& name) { // always return pointer as the reference can change as map re-allocate. // or use another level of indirection by creating a unique_ptr static std::unordered_map > source_map; @@ -54,8 +53,9 @@ SourceName SourceName::Get(const std::string& name) { TVM_REGISTER_API("relay._make.SourceName") .set_body_typed(SourceName::Get); -TVM_STATIC_IR_FUNCTOR_REGISTER(IRPrinter, vtable) -.set_dispatch([](const SourceNameNode* node, tvm::IRPrinter* p) { +TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) +.set_dispatch([](const ObjectRef& ref, tvm::IRPrinter* p) { + auto* node = static_cast(ref.get()); p->stream << "SourceName(" << node->name << ", " << node << ")"; }); @@ -78,8 +78,9 @@ TVM_REGISTER_NODE_TYPE(SpanNode); TVM_REGISTER_API("relay._make.Span") .set_body_typed(SpanNode::make); -TVM_STATIC_IR_FUNCTOR_REGISTER(IRPrinter, vtable) -.set_dispatch([](const SpanNode* node, tvm::IRPrinter* p) { +TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) +.set_dispatch([](const ObjectRef& ref, tvm::IRPrinter* p) { + auto* node = static_cast(ref.get()); p->stream << "SpanNode(" << node->source << ", " << node->lineno << ", " << node->col_offset << ")"; }); diff --git a/src/relay/ir/expr.cc b/src/relay/ir/expr.cc index c36b4c8566b8..3a4e87d8f629 100644 --- a/src/relay/ir/expr.cc +++ b/src/relay/ir/expr.cc @@ -18,7 +18,6 @@ */ /*! - * Copyright (c) 2018 by Contributors * \file src/tvm/ir/expr.cc * \brief The expression AST nodes of Relay. */ @@ -41,8 +40,9 @@ TVM_REGISTER_NODE_TYPE(ConstantNode); TVM_REGISTER_API("relay._make.Constant") .set_body_typed(ConstantNode::make); -TVM_STATIC_IR_FUNCTOR_REGISTER(IRPrinter, vtable) -.set_dispatch([](const ConstantNode* node, tvm::IRPrinter* p) { +TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) +.set_dispatch([](const ObjectRef& ref, IRPrinter* p) { + auto* node = static_cast(ref.get()); const PackedFunc* fprint = Registry::Get("relay._constant_repr"); CHECK(fprint) << "unable to find printing function for constants"; std::string data = (*fprint)(GetRef(node)); @@ -73,8 +73,9 @@ TVM_REGISTER_NODE_TYPE(TupleNode); TVM_REGISTER_API("relay._make.Tuple") .set_body_typed(TupleNode::make); -TVM_STATIC_IR_FUNCTOR_REGISTER(IRPrinter, vtable) -.set_dispatch([](const TupleNode* node, tvm::IRPrinter* p) { +TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) +.set_dispatch([](const ObjectRef& ref, IRPrinter* p) { + auto* node = static_cast(ref.get()); p->stream << "Tuple(" << node->fields << ")"; }); @@ -97,8 +98,9 @@ TVM_REGISTER_NODE_TYPE(VarNode); TVM_REGISTER_API("relay._make.Var") .set_body_typed(static_cast(VarNode::make)); -TVM_STATIC_IR_FUNCTOR_REGISTER(IRPrinter, vtable) -.set_dispatch([](const VarNode* node, tvm::IRPrinter* p) { +TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) +.set_dispatch([](const ObjectRef& ref, IRPrinter* p) { + auto* node = static_cast(ref.get()); p->stream << "Var(" << node->name_hint(); if (node->type_annotation.defined()) { p->stream << ", ty="; @@ -118,8 +120,9 @@ TVM_REGISTER_NODE_TYPE(GlobalVarNode); TVM_REGISTER_API("relay._make.GlobalVar") .set_body_typed(GlobalVarNode::make); -TVM_STATIC_IR_FUNCTOR_REGISTER(IRPrinter, vtable) -.set_dispatch([](const GlobalVarNode* node, tvm::IRPrinter* p) { +TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) +.set_dispatch([](const ObjectRef& ref, IRPrinter* p) { + auto* node = static_cast(ref.get()); p->stream << "GlobalVar(" << node->name_hint << ")"; }); @@ -217,12 +220,12 @@ TVM_REGISTER_NODE_TYPE(FunctionNode); TVM_REGISTER_API("relay._make.Function") .set_body_typed(FunctionNode::make); -TVM_STATIC_IR_FUNCTOR_REGISTER(IRPrinter, vtable) -.set_dispatch([](const FunctionNode* node, - tvm::IRPrinter* p) { - p->stream << "FunctionNode(" << node->params << ", " << node->ret_type - << ", " << node->body << ", " << node->type_params << ", " - << node->attrs << ")"; +TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) +.set_dispatch([](const ObjectRef& ref, IRPrinter* p) { + auto* node = static_cast(ref.get()); + p->stream << "FunctionNode(" << node->params << ", " << node->ret_type + << ", " << node->body << ", " << node->type_params << ", " + << node->attrs << ")"; }); Call CallNode::make(Expr op, Array args, Attrs attrs, @@ -240,11 +243,12 @@ TVM_REGISTER_NODE_TYPE(CallNode); TVM_REGISTER_API("relay._make.Call") .set_body_typed(CallNode::make); -TVM_STATIC_IR_FUNCTOR_REGISTER(IRPrinter, vtable) -.set_dispatch([](const CallNode* node, tvm::IRPrinter* p) { - p->stream << "CallNode(" << node->op << ", " << node->args << ", " - << node->attrs << ", " << node->type_args << ")"; -}); +TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) +.set_dispatch([](const ObjectRef& ref, IRPrinter* p) { + auto* node = static_cast(ref.get()); + p->stream << "CallNode(" << node->op << ", " << node->args << ", " + << node->attrs << ", " << node->type_args << ")"; + }); Let LetNode::make(Var var, Expr value, Expr body) { NodePtr n = make_node(); @@ -259,8 +263,9 @@ TVM_REGISTER_NODE_TYPE(LetNode); TVM_REGISTER_API("relay._make.Let") .set_body_typed(LetNode::make); -TVM_STATIC_IR_FUNCTOR_REGISTER(IRPrinter, vtable) -.set_dispatch([](const LetNode* node, tvm::IRPrinter* p) { +TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) +.set_dispatch([](const ObjectRef& ref, IRPrinter* p) { + auto* node = static_cast(ref.get()); p->stream << "LetNode(" << node->var << ", " << node->value << ", " << node->body << ")"; }); @@ -278,8 +283,9 @@ TVM_REGISTER_NODE_TYPE(IfNode); TVM_REGISTER_API("relay._make.If") .set_body_typed(IfNode::make); -TVM_STATIC_IR_FUNCTOR_REGISTER(IRPrinter, vtable) -.set_dispatch([](const IfNode* node, tvm::IRPrinter* p) { +TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) +.set_dispatch([](const ObjectRef& ref, IRPrinter* p) { + auto* node = static_cast(ref.get()); p->stream << "IfNode(" << node->cond << ", " << node->true_branch << ", " << node->false_branch << ")"; }); @@ -296,8 +302,9 @@ TVM_REGISTER_NODE_TYPE(TupleGetItemNode); TVM_REGISTER_API("relay._make.TupleGetItem") .set_body_typed(TupleGetItemNode::make); -TVM_STATIC_IR_FUNCTOR_REGISTER(IRPrinter, vtable) -.set_dispatch([](const TupleGetItemNode* node, tvm::IRPrinter* p) { +TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) +.set_dispatch([](const ObjectRef& ref, IRPrinter* p) { + auto* node = static_cast(ref.get()); p->stream << "TupleGetItemNode(" << node->tuple << ", " << node->index << ")"; }); @@ -312,8 +319,9 @@ TVM_REGISTER_NODE_TYPE(RefCreateNode); TVM_REGISTER_API("relay._make.RefCreate") .set_body_typed(RefCreateNode::make); -TVM_STATIC_IR_FUNCTOR_REGISTER(IRPrinter, vtable) -.set_dispatch([](const RefCreateNode* node, tvm::IRPrinter* p) { +TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) +.set_dispatch([](const ObjectRef& ref, IRPrinter* p) { + auto* node = static_cast(ref.get()); p->stream << "RefCreateNode(" << node->value << ")"; }); @@ -328,8 +336,9 @@ TVM_REGISTER_NODE_TYPE(RefReadNode); TVM_REGISTER_API("relay._make.RefRead") .set_body_typed(RefReadNode::make); -TVM_STATIC_IR_FUNCTOR_REGISTER(IRPrinter, vtable) -.set_dispatch([](const RefReadNode* node, tvm::IRPrinter* p) { +TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) +.set_dispatch([](const ObjectRef& ref, IRPrinter* p) { + auto* node = static_cast(ref.get()); p->stream << "RefReadNode(" << node->ref << ")"; }); @@ -345,8 +354,9 @@ TVM_REGISTER_NODE_TYPE(RefWriteNode); TVM_REGISTER_API("relay._make.RefWrite") .set_body_typed(RefWriteNode::make); -TVM_STATIC_IR_FUNCTOR_REGISTER(IRPrinter, vtable) -.set_dispatch([](const RefWriteNode* node, tvm::IRPrinter* p) { +TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) +.set_dispatch([](const ObjectRef& ref, IRPrinter* p) { + auto* node = static_cast(ref.get()); p->stream << "RefWriteNode(" << node->ref << ", " << node->value << ")"; }); diff --git a/src/relay/ir/module.cc b/src/relay/ir/module.cc index cd5b1e69f1d2..7c2e6b3f92ba 100644 --- a/src/relay/ir/module.cc +++ b/src/relay/ir/module.cc @@ -404,9 +404,9 @@ TVM_REGISTER_API("relay._module.Module_ImportFromStd") mod->ImportFromStd(path); });; -TVM_STATIC_IR_FUNCTOR_REGISTER(IRPrinter, vtable) -.set_dispatch( - [](const ModuleNode *node, tvm::IRPrinter *p) { +TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) +.set_dispatch([](const ObjectRef& ref, IRPrinter* p) { + auto* node = static_cast(ref.get()); p->stream << "ModuleNode( " << node->functions << ")"; }); diff --git a/src/relay/ir/op.cc b/src/relay/ir/op.cc index 7bfe41c05058..c4557ac16ad5 100644 --- a/src/relay/ir/op.cc +++ b/src/relay/ir/op.cc @@ -199,8 +199,9 @@ TVM_REGISTER_NODE_TYPE(OpNode) return static_cast(n)->name; }); -TVM_STATIC_IR_FUNCTOR_REGISTER(IRPrinter, vtable) -.set_dispatch([](const OpNode* node, tvm::IRPrinter* p) { +TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) +.set_dispatch([](const ObjectRef& ref, IRPrinter* p) { + auto* node = static_cast(ref.get()); p->stream << "Op(" << node->name << ")"; }); diff --git a/src/relay/ir/type.cc b/src/relay/ir/type.cc index 2604896c8605..471b36964179 100644 --- a/src/relay/ir/type.cc +++ b/src/relay/ir/type.cc @@ -18,7 +18,6 @@ */ /*! - * Copyright (c) 2018 by Contributors * \file src/tvm/ir/type.cc * \brief The type system AST nodes of Relay. */ @@ -58,9 +57,9 @@ TVM_REGISTER_NODE_TYPE(TensorTypeNode); TVM_REGISTER_API("relay._make.TensorType") .set_body_typed(TensorTypeNode::make); -TVM_STATIC_IR_FUNCTOR_REGISTER(IRPrinter, vtable) -.set_dispatch([](const TensorTypeNode* node, - tvm::IRPrinter* p) { +TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) +.set_dispatch([](const ObjectRef& ref, IRPrinter* p) { + auto* node = static_cast(ref.get()); p->stream << "TensorType(" << node->shape << ", " << node->dtype << ")"; }); @@ -78,9 +77,9 @@ TVM_REGISTER_API("relay._make.TypeVar") return TypeVarNode::make(name, static_cast(kind)); }); -TVM_STATIC_IR_FUNCTOR_REGISTER(IRPrinter, vtable) -.set_dispatch([](const TypeVarNode* node, - tvm::IRPrinter* p) { +TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) +.set_dispatch([](const ObjectRef& ref, IRPrinter* p) { + auto* node = static_cast(ref.get()); p->stream << "TypeVarNode(" << node->var->name_hint << ", " << node->kind << ")"; }); @@ -99,9 +98,9 @@ TVM_REGISTER_API("relay._make.GlobalTypeVar") return GlobalTypeVarNode::make(name, static_cast(kind)); }); -TVM_STATIC_IR_FUNCTOR_REGISTER(IRPrinter, vtable) -.set_dispatch([](const GlobalTypeVarNode *node, - tvm::IRPrinter *p) { +TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) +.set_dispatch([](const ObjectRef& ref, IRPrinter* p) { + auto* node = static_cast(ref.get()); p->stream << "GlobalTypeVarNode(" << node->var->name_hint << ", " << node->kind << ")"; }); @@ -118,9 +117,9 @@ TVM_REGISTER_NODE_TYPE(TypeCallNode); TVM_REGISTER_API("relay._make.TypeCall") .set_body_typed(TypeCallNode::make); -TVM_STATIC_IR_FUNCTOR_REGISTER(IRPrinter, vtable) -.set_dispatch([](const TypeCallNode* node, - tvm::IRPrinter* p) { +TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) +.set_dispatch([](const ObjectRef& ref, IRPrinter* p) { + auto* node = static_cast(ref.get()); p->stream << "TypeCallNode(" << node->func << ", " << node->args << ")"; }); @@ -138,12 +137,11 @@ TVM_REGISTER_API("relay._make.IncompleteType") return IncompleteTypeNode::make(static_cast(kind)); }); -TVM_STATIC_IR_FUNCTOR_REGISTER(IRPrinter, vtable) -.set_dispatch( - [](const IncompleteTypeNode* node, - tvm::IRPrinter* p) { - p->stream << "IncompleteTypeNode(" << node->kind << ", " << node << ")"; - }); +TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) +.set_dispatch([](const ObjectRef& ref, IRPrinter* p) { + auto* node = static_cast(ref.get()); + p->stream << "IncompleteTypeNode(" << node->kind << ", " << node << ")"; + }); FuncType FuncTypeNode::make(tvm::Array arg_types, Type ret_type, @@ -162,9 +160,9 @@ TVM_REGISTER_NODE_TYPE(FuncTypeNode); TVM_REGISTER_API("relay._make.FuncType") .set_body_typed(FuncTypeNode::make); -TVM_STATIC_IR_FUNCTOR_REGISTER(IRPrinter, vtable) -.set_dispatch([](const FuncTypeNode* node, - tvm::IRPrinter* p) { +TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) +.set_dispatch([](const ObjectRef& ref, IRPrinter* p) { + auto* node = static_cast(ref.get()); p->stream << "FuncTypeNode(" << node->type_params << ", " << node->arg_types << ", " << node->ret_type << ", " << node->type_constraints << ")"; @@ -187,8 +185,9 @@ TVM_REGISTER_NODE_TYPE(TypeRelationNode); TVM_REGISTER_API("relay._make.TypeRelation") .set_body_typed(TypeRelationNode::make); -TVM_STATIC_IR_FUNCTOR_REGISTER(IRPrinter, vtable) -.set_dispatch([](const TypeRelationNode* node, tvm::IRPrinter* p) { +TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) +.set_dispatch([](const ObjectRef& ref, IRPrinter* p) { + auto* node = static_cast(ref.get()); p->stream << "TypeRelationNode(" << node->func->name << ", " << node->args << ")"; @@ -205,9 +204,9 @@ TVM_REGISTER_NODE_TYPE(TupleTypeNode); TVM_REGISTER_API("relay._make.TupleType") .set_body_typed(TupleTypeNode::make); -TVM_STATIC_IR_FUNCTOR_REGISTER(IRPrinter, vtable) -.set_dispatch([](const TupleTypeNode* node, - tvm::IRPrinter* p) { +TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) +.set_dispatch([](const ObjectRef& ref, IRPrinter* p) { + auto* node = static_cast(ref.get()); p->stream << "TupleTypeNode(" << node->fields << ")"; }); @@ -222,9 +221,9 @@ TVM_REGISTER_API("relay._make.RefType") TVM_REGISTER_NODE_TYPE(RefTypeNode); -TVM_STATIC_IR_FUNCTOR_REGISTER(IRPrinter, vtable) -.set_dispatch([](const RefTypeNode* node, - tvm::IRPrinter* p) { +TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) +.set_dispatch([](const ObjectRef& ref, IRPrinter* p) { + auto* node = static_cast(ref.get()); p->stream << "RefTypeNode(" << node->value << ")"; }); diff --git a/src/relay/ir/type_functor.h b/src/relay/ir/type_functor.h index bd9e649a3b4f..67c139185ebf 100644 --- a/src/relay/ir/type_functor.h +++ b/src/relay/ir/type_functor.h @@ -18,14 +18,13 @@ */ /*! - * Copyright (c) 2018 by Contributors * \file type_functor.h * \brief A way to defined arbitrary function signature with dispatch on types. */ #ifndef TVM_RELAY_IR_TYPE_FUNCTOR_H_ #define TVM_RELAY_IR_TYPE_FUNCTOR_H_ -#include +#include #include #include #include @@ -54,7 +53,7 @@ template class TypeFunctor { private: using TSelf = TypeFunctor; - using FType = tvm::IRFunctor; + using FType = tvm::NodeFunctor; public: /*! \brief the result type of this functor */ diff --git a/src/relay/pass/pass_manager.cc b/src/relay/pass/pass_manager.cc index d2688620b0c3..d9b8d72eb417 100644 --- a/src/relay/pass/pass_manager.cc +++ b/src/relay/pass/pass_manager.cc @@ -449,9 +449,9 @@ TVM_REGISTER_API("relay._transform.Info") *ret = pass->Info(); }); -TVM_STATIC_IR_FUNCTOR_REGISTER(IRPrinter, vtable) -.set_dispatch([](const PassInfoNode* node, - tvm::IRPrinter* p) { +TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) +.set_dispatch([](const ObjectRef& ref, tvm::IRPrinter* p) { + auto* node = static_cast(ref.get()); p->stream << "The meta data of the pass: "; p->stream << "pass name: " << node->name; p->stream << "opt_level: " << node->opt_level; @@ -475,9 +475,9 @@ TVM_REGISTER_API("relay._transform.RunPass") *ret = pass(mod); }); -TVM_STATIC_IR_FUNCTOR_REGISTER(IRPrinter, vtable) -.set_dispatch([](const ModulePassNode* node, - tvm::IRPrinter* p) { +TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) +.set_dispatch([](const ObjectRef& ref, IRPrinter* p) { + auto* node = static_cast(ref.get()); const PassInfo info = node->Info(); p->stream << "Run Module pass: " << info->name << " at the optimization level " << info->opt_level; @@ -488,9 +488,9 @@ TVM_REGISTER_NODE_TYPE(FunctionPassNode); TVM_REGISTER_API("relay._transform.MakeFunctionPass") .set_body_typed(FunctionPassNode::make); -TVM_STATIC_IR_FUNCTOR_REGISTER(IRPrinter, vtable) -.set_dispatch([](const FunctionPassNode* node, - tvm::IRPrinter* p) { +TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) +.set_dispatch([](const ObjectRef& ref, IRPrinter* p) { + auto* node = static_cast(ref.get()); const PassInfo info = node->Info(); p->stream << "Run Function pass: " << info->name << " at the optimization level " << info->opt_level; @@ -508,9 +508,9 @@ TVM_REGISTER_API("relay._transform.Sequential") *ret = Sequential(passes, pass_info); }); -TVM_STATIC_IR_FUNCTOR_REGISTER(IRPrinter, vtable) -.set_dispatch([](const SequentialNode* node, - tvm::IRPrinter* p) { +TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) +.set_dispatch([](const ObjectRef& ref, IRPrinter* p) { + auto* node = static_cast(ref.get()); const PassInfo info = node->Info(); p->stream << "Run Sequential pass: " << info->name << " at the optimization level " << info->opt_level << ". "; @@ -538,9 +538,9 @@ TVM_REGISTER_API("relay._transform.PassContext") *ret = pctx; }); -TVM_STATIC_IR_FUNCTOR_REGISTER(IRPrinter, vtable) -.set_dispatch([](const PassContextNode* node, - tvm::IRPrinter* p) { +TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) +.set_dispatch([](const ObjectRef& ref, IRPrinter* p) { + auto* node = static_cast(ref.get()); p->stream << "Pass context information: " << "\n"; p->stream << "\topt_level: " << node->opt_level << "\n"; p->stream << "\tfallback device: " diff --git a/src/relay/pass/quantize/quantize.cc b/src/relay/pass/quantize/quantize.cc index d564d2e76dbe..2793577cfee2 100644 --- a/src/relay/pass/quantize/quantize.cc +++ b/src/relay/pass/quantize/quantize.cc @@ -117,7 +117,8 @@ QConfig& QConfig::Current() { TVM_REGISTER_NODE_TYPE(QConfigNode); TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) -.set_dispatch([](const QConfigNode *op, IRPrinter *p) { +.set_dispatch([](const ObjectRef& ref, IRPrinter* p) { + auto* op = static_cast(ref.get()); p->stream << "qconfig("; p->stream << "nbit_input=" << op->nbit_input << ", "; p->stream << "nbit_weight=" << op->nbit_weight << ", "; diff --git a/src/schedule/schedule_lang.cc b/src/schedule/schedule_lang.cc index 407729df8038..54503fc4a8b8 100644 --- a/src/schedule/schedule_lang.cc +++ b/src/schedule/schedule_lang.cc @@ -800,17 +800,20 @@ TVM_REGISTER_NODE_TYPE(ScheduleNode); // Printer TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) -.set_dispatch([](const StageNode *op, IRPrinter *p) { +.set_dispatch([](const ObjectRef& node, IRPrinter* p) { + auto* op = static_cast(node.get()); if (op->op.defined()) { p->stream << "stage(" << op->origin_op->name << ", " << op << ")"; } else { p->stream << "group-stage(" << op << ")"; } }) -.set_dispatch([](const IterVarAttrNode *op, IRPrinter *p) { +.set_dispatch([](const ObjectRef& node, IRPrinter* p) { + auto* op = static_cast(node.get()); p->stream << IterVarType2String(op->iter_type); }) -.set_dispatch([](const SplitNode *op, IRPrinter *p) { +.set_dispatch([](const ObjectRef& node, IRPrinter* p) { + auto* op = static_cast(node.get()); p->stream << "split(parent="; p->Print(op->parent); p->stream << ", outer="; @@ -819,7 +822,8 @@ TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) p->Print(op->inner); p->stream << ')'; }) -.set_dispatch([](const FuseNode *op, IRPrinter *p) { +.set_dispatch([](const ObjectRef& node, IRPrinter* p) { + auto* op = static_cast(node.get()); p->stream << "split("; p->stream << "outer="; p->Print(op->outer); @@ -829,7 +833,8 @@ TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) p->Print(op->fused); p->stream << ')'; }) -.set_dispatch([](const RebaseNode *op, IRPrinter *p) { +.set_dispatch([](const ObjectRef& node, IRPrinter* p) { + auto* op = static_cast(node.get()); p->stream << "rebase("; p->stream << "parent="; p->Print(op->parent); @@ -837,12 +842,14 @@ TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable) p->Print(op->rebased); p->stream << ')'; }) -.set_dispatch([](const SingletonNode *op, IRPrinter *p) { +.set_dispatch([](const ObjectRef& node, IRPrinter* p) { + auto* op = static_cast(node.get()); p->stream << "singleton("; p->Print(op->iter); p->stream << ')'; }) -.set_dispatch([](const ScheduleNode *op, IRPrinter *p) { +.set_dispatch([](const ObjectRef& node, IRPrinter* p) { + auto* op = static_cast(node.get()); p->stream << "schedule(" << op << ")"; }); } // namespace tvm diff --git a/tests/cpp/attrs_test.cc b/tests/cpp/attrs_test.cc index 038ad6f171ec..9ccb9c96166d 100644 --- a/tests/cpp/attrs_test.cc +++ b/tests/cpp/attrs_test.cc @@ -6,9 +6,9 @@ * 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 @@ -53,7 +53,7 @@ struct TestAttrs : public AttrsNode { TEST(Attrs, Basic) { using namespace tvm; using namespace tvm::test; - std::shared_ptr n = std::make_shared(); + ObjectPtr n = make_object(); try { n->InitBySeq("axis", 10); LOG(FATAL) << "bad"; diff --git a/tests/cpp/ir_functor_test.cc b/tests/cpp/ir_functor_test.cc index fef43f97d3c3..5636958a5b26 100644 --- a/tests/cpp/ir_functor_test.cc +++ b/tests/cpp/ir_functor_test.cc @@ -21,7 +21,7 @@ #include #include #include -#include +#include #include TEST(IRF, Basic) { @@ -30,12 +30,12 @@ TEST(IRF, Basic) { Var x("x"); auto z = x + 1; - IRFunctor f; + NodeFunctor f; LOG(INFO) << "x"; - f.set_dispatch([](const Variable* n, int b) { + f.set_dispatch([](const ObjectRef& n, int b) { return b; }); - f.set_dispatch([](const Add* n, int b) { + f.set_dispatch([](const ObjectRef& n, int b) { return b + 2; }); CHECK_EQ(f(x, 2), 2); diff --git a/tests/cpp/ir_mutator_test.cc b/tests/cpp/ir_mutator_test.cc index 30972e762314..1b3296da17c3 100644 --- a/tests/cpp/ir_mutator_test.cc +++ b/tests/cpp/ir_mutator_test.cc @@ -45,7 +45,7 @@ IRMutator::FMutateExpr &IRVar2Const::vtable_expr() { // NOLINT(*) } TVM_STATIC_IR_FUNCTOR(IRVar2Const, vtable_expr) -.set_dispatch([](const Variable* op, const Expr &e, IRMutator* m) { +.set_dispatch([](const ObjectRef& ref, const Expr &e, IRMutator* m) { IRVar2Const* vm = static_cast(m); if (e.same_as(vm->var)) { return Expr(IntImm::make(Int(32), vm->int_val));