From 928b54b8ebe3747108a62251584dfdd7d883a62b Mon Sep 17 00:00:00 2001 From: tqchen Date: Thu, 23 Apr 2020 10:20:51 -0700 Subject: [PATCH] [RUNTIME][OBJECT] Introduce static slots for common objects. The _type_child_slots can be used to enable quick type checking optimization by checking the whether the type index is within the bound. This PR enables these static slots: - Introduce a static assert to avoid the scenario when a developer forget to _type_child_slots when the field is set for the type's parent. - Revamp and assign static type index to common runtime objects - Add a DumpTypeTable call to allow developer monitor the current situation of type table and offers suggestions for the slots(ideally the slots equals the number of children so there is no overflow. --- include/tvm/ir/expr.h | 7 +++-- include/tvm/ir/function.h | 1 + include/tvm/ir/tensor_type.h | 1 + include/tvm/ir/type.h | 2 ++ include/tvm/relay/expr.h | 1 + include/tvm/runtime/container.h | 6 ++--- include/tvm/runtime/ndarray.h | 4 +-- include/tvm/runtime/object.h | 43 ++++++++++++++++++++++--------- include/tvm/runtime/packed_func.h | 2 ++ include/tvm/runtime/vm.h | 4 +-- include/tvm/tir/stmt.h | 1 + include/tvm/tir/var.h | 1 + python/tvm/runtime/container.py | 2 +- python/tvm/runtime/ndarray.py | 2 +- src/arith/canonical_simplify.cc | 1 + src/runtime/module.cc | 2 ++ src/runtime/object.cc | 34 +++++++++++++++++++++--- 17 files changed, 87 insertions(+), 27 deletions(-) diff --git a/include/tvm/ir/expr.h b/include/tvm/ir/expr.h index 6630bf3ded20e..fba35a9193f94 100644 --- a/include/tvm/ir/expr.h +++ b/include/tvm/ir/expr.h @@ -42,9 +42,10 @@ namespace tvm { */ class BaseExprNode : public Object { public: - static constexpr const char* _type_key = "Expr"; + static constexpr const char* _type_key = "BaseExpr"; static constexpr const bool _type_has_method_sequal_reduce = true; static constexpr const bool _type_has_method_shash_reduce = true; + static constexpr const uint32_t _type_child_slots = 58; TVM_DECLARE_BASE_OBJECT_INFO(BaseExprNode, Object); }; @@ -88,6 +89,7 @@ class PrimExprNode : public BaseExprNode { DataType dtype; static constexpr const char* _type_key = "PrimExpr"; + static constexpr const uint32_t _type_child_slots = 34; TVM_DECLARE_BASE_OBJECT_INFO(PrimExprNode, BaseExprNode); }; @@ -161,7 +163,8 @@ class RelayExprNode : public BaseExprNode { template inline const TTypeNode* type_as() const; - static constexpr const char* _type_key = "relay.Expr"; + static constexpr const char* _type_key = "RelayExpr"; + static constexpr const uint32_t _type_child_slots = 22; TVM_DECLARE_BASE_OBJECT_INFO(RelayExprNode, BaseExprNode); }; diff --git a/include/tvm/ir/function.h b/include/tvm/ir/function.h index d55656f34b006..b4a9ed0d89645 100644 --- a/include/tvm/ir/function.h +++ b/include/tvm/ir/function.h @@ -140,6 +140,7 @@ class BaseFuncNode : public RelayExprNode { } static constexpr const char* _type_key = "BaseFunc"; + static constexpr const uint32_t _type_child_slots = 2; TVM_DECLARE_BASE_OBJECT_INFO(BaseFuncNode, RelayExprNode); }; diff --git a/include/tvm/ir/tensor_type.h b/include/tvm/ir/tensor_type.h index e993cd9afaccf..489ea64409006 100644 --- a/include/tvm/ir/tensor_type.h +++ b/include/tvm/ir/tensor_type.h @@ -36,6 +36,7 @@ namespace tvm { class BaseTensorTypeNode : public TypeNode { public: static constexpr const char* _type_key = "relay.BaseTensorType"; + static constexpr const uint32_t _type_child_slots = 1; TVM_DECLARE_BASE_OBJECT_INFO(BaseTensorTypeNode, TypeNode); }; diff --git a/include/tvm/ir/type.h b/include/tvm/ir/type.h index 0e65758a2e1c9..0ef03c42c03e3 100644 --- a/include/tvm/ir/type.h +++ b/include/tvm/ir/type.h @@ -81,6 +81,7 @@ class TypeNode : public Object { static constexpr const char* _type_key = "Type"; static constexpr const bool _type_has_method_sequal_reduce = true; static constexpr const bool _type_has_method_shash_reduce = true; + static constexpr const uint32_t _type_child_slots = 14; TVM_DECLARE_BASE_OBJECT_INFO(TypeNode, Object); }; @@ -391,6 +392,7 @@ inline bool IsVoidType(const Type& type) { class TypeConstraintNode : public TypeNode { public: static constexpr const char* _type_key = "TypeConstraint"; + static constexpr const uint32_t _type_child_slots = 1; TVM_DECLARE_BASE_OBJECT_INFO(TypeConstraintNode, TypeNode); }; diff --git a/include/tvm/relay/expr.h b/include/tvm/relay/expr.h index fe240c30e4717..8c5026050d91f 100644 --- a/include/tvm/relay/expr.h +++ b/include/tvm/relay/expr.h @@ -630,6 +630,7 @@ class TempExprNode : public ExprNode { static constexpr const char* _type_key = "relay.TempExpr"; static constexpr const bool _type_has_method_sequal_reduce = false; static constexpr const bool _type_has_method_shash_reduce = false; + static constexpr const uint32_t _type_child_slots = 0; TVM_DECLARE_BASE_OBJECT_INFO(TempExprNode, ExprNode); }; diff --git a/include/tvm/runtime/container.h b/include/tvm/runtime/container.h index 7d08613af2153..cdb92ba6779a1 100644 --- a/include/tvm/runtime/container.h +++ b/include/tvm/runtime/container.h @@ -200,8 +200,8 @@ class ADTObj : public Object, public InplaceArrayBase { uint32_t size; // The fields of the structure follows directly in memory. - static constexpr const uint32_t _type_index = TypeIndex::kVMADT; - static constexpr const char* _type_key = "vm.ADT"; + static constexpr const uint32_t _type_index = TypeIndex::kRuntimeADT; + static constexpr const char* _type_key = "runtime.ADT"; TVM_DECLARE_FINAL_OBJECT_INFO(ADTObj, Object); private: @@ -314,7 +314,7 @@ class StringObj : public Object { /*! \brief The length of the string object. */ uint64_t size; - static constexpr const uint32_t _type_index = TypeIndex::kDynamic; + static constexpr const uint32_t _type_index = TypeIndex::kRuntimeString; static constexpr const char* _type_key = "runtime.String"; TVM_DECLARE_FINAL_OBJECT_INFO(StringObj, Object); diff --git a/include/tvm/runtime/ndarray.h b/include/tvm/runtime/ndarray.h index 17f81a2a8b680..33f27f49fbbe0 100644 --- a/include/tvm/runtime/ndarray.h +++ b/include/tvm/runtime/ndarray.h @@ -288,10 +288,10 @@ class NDArray::Container : using Object::IncRef; // Information for object protocol. - static constexpr const uint32_t _type_index = TypeIndex::kDynamic; + static constexpr const uint32_t _type_index = TypeIndex::kRuntimeNDArray; static constexpr const uint32_t _type_child_slots = 0; static constexpr const uint32_t _type_child_slots_can_overflow = true; - static constexpr const char* _type_key = "NDArray"; + static constexpr const char* _type_key = "runtime.NDArray"; TVM_DECLARE_BASE_OBJECT_INFO(NDArray::Container, Object); protected: diff --git a/include/tvm/runtime/object.h b/include/tvm/runtime/object.h index edca925baeb0d..764dcdf3640c1 100644 --- a/include/tvm/runtime/object.h +++ b/include/tvm/runtime/object.h @@ -46,17 +46,31 @@ namespace tvm { namespace runtime { -/*! \brief list of the type index. */ -enum TypeIndex { - /*! \brief Root object type. */ - kRoot = 0, - kClosure = 1, - kVMADT = 2, - kRuntimeModule = 3, - kStaticIndexEnd, - /*! \brief Type index is allocated during runtime. */ - kDynamic = kStaticIndexEnd -}; +/*! + * \brief Namespace for the list of type index. + * \note Use struct so that we have to use TypeIndex::ENumName to refer to + * the constant, but still able to use enum. + */ +struct TypeIndex { + enum { + /*! \brief Root object type. */ + kRoot = 0, + // Standard static index assignments, + // Frontends can take benefit of these constants. + /*! \brief runtime::Module. */ + kRuntimeModule = 1, + /*! \brief runtime::NDArray. */ + kRuntimeNDArray = 2, + /*! \brief runtime::String. */ + kRuntimeString = 3, + // static assignments that may subject to change. + kRuntimeClosure, + kRuntimeADT, + kStaticIndexEnd, + /*! \brief Type index is allocated during runtime. */ + kDynamic = kStaticIndexEnd + }; +}; // namespace TypeIndex /*! * \brief base class of all object containers. @@ -198,7 +212,7 @@ class Object { using RefCounterType = int32_t; #endif - static constexpr const char* _type_key = "Object"; + static constexpr const char* _type_key = "runtime.Object"; static uint32_t _GetOrAllocRuntimeTypeIndex() { return TypeIndex::kRoot; @@ -675,6 +689,10 @@ struct ObjectEqual { #define TVM_DECLARE_BASE_OBJECT_INFO(TypeName, ParentType) \ static_assert(!ParentType::_type_final, "ParentObj maked as final"); \ static uint32_t RuntimeTypeIndex() { \ + static_assert(TypeName::_type_child_slots == 0 || \ + ParentType::_type_child_slots == 0 || \ + TypeName::_type_child_slots < ParentType::_type_child_slots, \ + "Need to set _type_child_slots when parent specifies it."); \ if (TypeName::_type_index != ::tvm::runtime::TypeIndex::kDynamic) { \ return TypeName::_type_index; \ } \ @@ -690,6 +708,7 @@ struct ObjectEqual { return tidx; \ } \ + /*! * \brief helper macro to declare type information in a final class. * \param TypeName The name of the current type. diff --git a/include/tvm/runtime/packed_func.h b/include/tvm/runtime/packed_func.h index 3d5a7e865303a..c56600730eda3 100644 --- a/include/tvm/runtime/packed_func.h +++ b/include/tvm/runtime/packed_func.h @@ -1268,6 +1268,8 @@ struct unpack_call_dispatcher { template inline void unpack_call(const F& f, const TVMArgs& args, TVMRetValue* rv) { + CHECK_EQ(nargs, args.size()) + << "Expect " << nargs << " arguments but get " << args.size(); unpack_call_dispatcher::run(f, args, rv); } diff --git a/include/tvm/runtime/vm.h b/include/tvm/runtime/vm.h index 43c222d0994aa..44d58987954c6 100644 --- a/include/tvm/runtime/vm.h +++ b/include/tvm/runtime/vm.h @@ -44,8 +44,8 @@ namespace vm { */ class ClosureObj : public Object { public: - static constexpr const uint32_t _type_index = TypeIndex::kClosure; - static constexpr const char* _type_key = "Closure"; + static constexpr const uint32_t _type_index = TypeIndex::kRuntimeClosure; + static constexpr const char* _type_key = "runtime.Closure"; TVM_DECLARE_BASE_OBJECT_INFO(ClosureObj, Object); }; diff --git a/include/tvm/tir/stmt.h b/include/tvm/tir/stmt.h index 20c2d009b93cb..aed8b5c77ae58 100644 --- a/include/tvm/tir/stmt.h +++ b/include/tvm/tir/stmt.h @@ -40,6 +40,7 @@ class StmtNode : public Object { static constexpr const char* _type_key = "Stmt"; static constexpr const bool _type_has_method_sequal_reduce = true; static constexpr const bool _type_has_method_shash_reduce = true; + static constexpr const uint32_t _type_child_slots = 15; TVM_DECLARE_BASE_OBJECT_INFO(StmtNode, Object); }; diff --git a/include/tvm/tir/var.h b/include/tvm/tir/var.h index 19c904a1230f9..bb73bf03d88b5 100644 --- a/include/tvm/tir/var.h +++ b/include/tvm/tir/var.h @@ -78,6 +78,7 @@ class VarNode : public PrimExprNode { } static constexpr const char* _type_key = "tir.Var"; + static constexpr const uint32_t _type_child_slots = 1; TVM_DECLARE_BASE_OBJECT_INFO(VarNode, PrimExprNode); }; diff --git a/python/tvm/runtime/container.py b/python/tvm/runtime/container.py index a719dcd4eaf06..69652bea1dc56 100644 --- a/python/tvm/runtime/container.py +++ b/python/tvm/runtime/container.py @@ -60,7 +60,7 @@ def getitem_helper(obj, elem_getter, length, idx): return elem_getter(obj, idx) -@tvm._ffi.register_object("vm.ADT") +@tvm._ffi.register_object("runtime.ADT") class ADT(Object): """Algebatic data type(ADT) object. diff --git a/python/tvm/runtime/ndarray.py b/python/tvm/runtime/ndarray.py index ee7ab7b5d11fe..10bbb6ef54c29 100644 --- a/python/tvm/runtime/ndarray.py +++ b/python/tvm/runtime/ndarray.py @@ -36,7 +36,7 @@ from tvm._ffi._ctypes.ndarray import NDArrayBase -@tvm._ffi.register_object +@tvm._ffi.register_object("runtime.NDArray") class NDArray(NDArrayBase): """Lightweight NDArray class of TVM runtime. diff --git a/src/arith/canonical_simplify.cc b/src/arith/canonical_simplify.cc index 7a6e772c2935d..2bb0189890a6a 100644 --- a/src/arith/canonical_simplify.cc +++ b/src/arith/canonical_simplify.cc @@ -57,6 +57,7 @@ class CanonicalExprNode : public PrimExprNode { } static constexpr const char* _type_key = "arith.CanonicalExpr"; + static constexpr const uint32_t _type_child_slots = 2; TVM_DECLARE_BASE_OBJECT_INFO(CanonicalExprNode, PrimExprNode); }; diff --git a/src/runtime/module.cc b/src/runtime/module.cc index ab2cb67cfa4d1..d2ed7ff9e2b7b 100644 --- a/src/runtime/module.cc +++ b/src/runtime/module.cc @@ -188,5 +188,7 @@ TVM_REGISTER_GLOBAL("runtime.ModuleSaveToFile") .set_body_typed([](Module mod, std::string name, std::string fmt) { mod->SaveToFile(name, fmt); }); + +TVM_REGISTER_OBJECT_TYPE(ModuleNode); } // namespace runtime } // namespace tvm diff --git a/src/runtime/object.cc b/src/runtime/object.cc index 0d85b9dab42cc..03012006bd797 100644 --- a/src/runtime/object.cc +++ b/src/runtime/object.cc @@ -86,7 +86,8 @@ class TypeContext { return it->second; } // try to allocate from parent's type table. - CHECK_LT(parent_tindex, type_table_.size()); + CHECK_LT(parent_tindex, type_table_.size()) + << " skey= " << skey << "static_index=" << static_tindex; TypeInfo& pinfo = type_table_[parent_tindex]; CHECK_EQ(pinfo.index, parent_tindex); @@ -108,7 +109,7 @@ class TypeContext { << " between " << type_table_[allocated_tindex].name << " and " << skey; - } else if (pinfo.allocated_slots + num_slots < pinfo.num_slots) { + } else if (pinfo.allocated_slots + num_slots <= pinfo.num_slots) { // allocate the slot from parent's reserved pool allocated_tindex = parent_tindex + pinfo.allocated_slots; // update parent's state @@ -119,8 +120,8 @@ class TypeContext { // allocate new entries. allocated_tindex = type_counter_; type_counter_ += num_slots; - CHECK_LE(type_table_.size(), allocated_tindex); - type_table_.resize(allocated_tindex + 1, TypeInfo()); + CHECK_LE(type_table_.size(), type_counter_); + type_table_.resize(type_counter_, TypeInfo()); } CHECK_GT(allocated_tindex, parent_tindex); // initialize the slot. @@ -161,6 +162,25 @@ class TypeContext { return it->second; } + void Dump(int min_children_count) { + std::vector num_children(type_table_.size(), 0); + // reverse accumulation so we can get total counts in a bottom-up manner. + for (auto it = type_table_.rbegin(); it != type_table_.rend(); ++it) { + if (it->index != 0) { + num_children[it->parent_index] += num_children[it->index] + 1; + } + } + + for (const auto& info : type_table_) { + if (info.index != 0 && num_children[info.index] >= min_children_count) { + std::cerr <<'[' << info.index << "] "<< info.name + << "\tparent=" << type_table_[info.parent_index].name + << "\tnum_child_slots=" << info.num_slots - 1 + << "\tnum_children=" << num_children[info.index] << std::endl; + } + } + } + static TypeContext* Global() { static TypeContext inst; return &inst; @@ -169,6 +189,7 @@ class TypeContext { private: TypeContext() { type_table_.resize(TypeIndex::kStaticIndexEnd, TypeInfo()); + type_table_[0].name = "runtime.Object"; } // mutex to avoid registration from multiple threads. std::mutex mutex_; @@ -208,6 +229,11 @@ TVM_REGISTER_GLOBAL("runtime.ObjectHash") .set_body_typed([](ObjectRef obj) { return static_cast(ObjectHash()(obj)); }); + +TVM_REGISTER_GLOBAL("runtime.DumpTypeTable") +.set_body_typed([](int min_child_count) { + TypeContext::Global()->Dump(min_child_count); +}); } // namespace runtime } // namespace tvm