[RUNTIME][OBJECT] Introduce static slots for common objects.

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<typename TTypeNode>
   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);
 };
 

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);
 };
 

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);
 };
 

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);
 };
 

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);
 };
 

include/tvm/runtime/container.h

@@ -200,8 +200,8 @@ class ADTObj : public Object, public InplaceArrayBase<ADTObj, ObjectRef> {
   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);
 

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:

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.

include/tvm/runtime/packed_func.h

@@ -1268,6 +1268,8 @@ struct unpack_call_dispatcher<void, 0, index, F> {
 
 template<typename R, int nargs, typename F>
 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<R, nargs, 0, F>::run(f, args, rv);
 }
 

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);
 };
 

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);
 };
 

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);
 };
 

python/tvm/relay/quantize/quantize.py

@@ -121,7 +121,7 @@ def __enter__(self):
         return self
 
     def __exit__(self, ptype, value, trace):
-        _quantize._ExitQConfigScope(self)
+        _quantize._ExitQConfigScope()
 
     def __setattr__(self, name, value):
         if name in QConfig._node_defaults:

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.
 

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.
 

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);
 };
 

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

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<int> 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<int64_t>(ObjectHash()(obj));
 });
+
+TVM_REGISTER_GLOBAL("runtime.DumpTypeTable")
+.set_body_typed([](int min_child_count) {
+  TypeContext::Global()->Dump(min_child_count);
+});
 }  // namespace runtime
 }  // namespace tvm
 

src/target/codegen.cc

@@ -47,6 +47,7 @@ runtime::Module Build(IRModule mod, const Target& target) {
   if (BuildConfig::Current()->disable_assert) {
     mod = tir::transform::SkipAssert()(mod);
   }
+
   std::string build_f_name = "target.build." + target->target_name;
   // the build function.
   const PackedFunc* bf = runtime::Registry::Get(build_f_name);

src/target/opt/build_cuda_on.cc

@@ -127,7 +127,7 @@ std::string NVRTCCompile(const std::string& code, bool include_path = false) {
   return ptx;
 }
 
-runtime::Module BuildCUDA(IRModule mod) {
+runtime::Module BuildCUDA(IRModule mod, std::string target) {
   using tvm::runtime::Registry;
   bool output_ssa = false;
   CodeGenCUDA cg;

src/target/source/codegen_opencl.cc

@@ -238,7 +238,7 @@ void CodeGenOpenCL::VisitExpr_(const FloatImmNode *op, std::ostream& os) { // NO
   }
 }
 
-runtime::Module BuildOpenCL(IRModule mod) {
+runtime::Module BuildOpenCL(IRModule mod, std::string target) {
   using tvm::runtime::Registry;
   bool output_ssa = false;
   CodeGenOpenCL cg;

src/target/source/codegen_opengl.cc

@@ -289,7 +289,7 @@ void CodeGenOpenGL::VisitStmt_(const EvaluateNode* op) {
   this->stream << GetVarID(buffer) << " = " << PrintExpr(value) << ";\n";
 }
 
-runtime::Module BuildOpenGL(IRModule mod) {
+runtime::Module BuildOpenGL(IRModule mod, std::string target) {
   bool output_ssa = false;
   CodeGenOpenGL cg;
   cg.Init(output_ssa);

src/target/spirv/build_vulkan.cc

@@ -70,7 +70,7 @@ class SPIRVTools {
   spv_context ctx_;
 };
 
-runtime::Module BuildSPIRV(IRModule mod) {
+runtime::Module BuildSPIRV(IRModule mod, std::string target) {
   using tvm::runtime::Registry;
   using tvm::runtime::VulkanShader;
 

src/target/stackvm/codegen_stackvm.cc

@@ -517,7 +517,7 @@ void CodeGenStackVM::VisitExpr_(const LetNode* op) {
   this->Push(op->body);
 }
 
-runtime::Module BuildStackVM(const IRModule& mod) {
+runtime::Module BuildStackVM(const IRModule& mod, const std::string& target) {
   std::unordered_map<std::string, StackVM> fmap;
   std::string entry_func;