[RUNTIME] Refactor extension type handling, now it is header only (#924)

* [RUNTIME] Refactor extension type handling, now it is header only

apps/extension/src/tvm_ext.cc

@@ -22,13 +22,10 @@ struct extension_class_info<tvm_ext::IntVector> {
 }  // namespace tvm
 }  // namespace runtime
 
-
-namespace tvm_ext {
-
 using namespace tvm;
 using namespace tvm::runtime;
 
-TVM_REGISTER_EXT_TYPE(IntVector);
+namespace tvm_ext {
 
 TVM_REGISTER_GLOBAL("tvm_ext.ivec_create")
 .set_body([](TVMArgs args, TVMRetValue *rv) {
@@ -66,3 +63,18 @@ TVM_REGISTER_GLOBAL("device_api.ext_dev")
     *rv = (*tvm::runtime::Registry::Get("device_api.cpu"))();
   });
 }  // namespace tvm_ext
+
+// This callback approach allows extension allows tvm to extract
+// This way can be helpful when we want to use a header only
+// minimum version of TVM Runtime.
+extern "C" int TVMExtDeclare(TVMFunctionHandle pregister) {
+  const PackedFunc& fregister =
+      *static_cast<PackedFunc*>(pregister);
+  auto mul = [](TVMArgs args, TVMRetValue *rv) {
+    int x = args[0];
+    int y = args[1];
+    *rv = x * y;
+  };
+  fregister("mul", PackedFunc(mul));
+  return 0;
+}

apps/extension/tests/test_ext.py

@@ -44,8 +44,14 @@ def test_ext_vec():
 
     tvm.convert(ivec_cb)(ivec)
 
+def test_extract_ext():
+    fdict = tvm.extract_ext_funcs(tvm_ext._LIB.TVMExtDeclare)
+    assert fdict["mul"](3, 4) == 12
+
+
 if __name__ == "__main__":
     test_ext_dev()
     test_ext_vec()
     test_bind_add()
     test_sym_add()
+    test_extract_ext()

include/tvm/runtime/c_runtime_api.h

@@ -24,6 +24,13 @@
 #define TVM_EXTERN_C
 #endif
 
+// Macros to do weak linking
+#ifdef _MSC_VER
+#define TVM_WEAK __declspec(selectany)
+#else
+#define TVM_WEAK __attribute__((weak))
+#endif
+
 #ifdef __EMSCRIPTEN__
 #include <emscripten/emscripten.h>
 #define TVM_DLL EMSCRIPTEN_KEEPALIVE
@@ -314,6 +321,17 @@ typedef int (*TVMPackedCFunc)(
 typedef void (*TVMPackedCFuncFinalizer)(void* resource_handle);
 
 /*!
+ * \brief Signature for extension function declarer.
+ *
+ *  TVM call this function to get the extension functions
+ *  The declarer will call register_func to register function and their name.
+ *
+ * \param resource_func_handle The register function
+ * \return 0 if success, -1 if failure happens
+ */
+typedef int (*TVMExtensionFuncDeclarer)(TVMFunctionHandle register_func_handle);
+
+/*!
  * \brief Wrap a TVMPackedCFunc to become a FunctionHandle.
  *
  * The resource_handle will be managed by TVM API, until the function is no longer used.

include/tvm/runtime/module.h

@@ -38,8 +38,14 @@ class Module {
    * \param query_imports Whether also query dependency modules.
    * \return The result function.
    *  This function will return PackedFunc(nullptr) if function do not exist.
+   * \note Implemented in packed_func.cc
    */
-  TVM_DLL PackedFunc GetFunction(const std::string& name, bool query_imports = false);
+  inline PackedFunc GetFunction(const std::string& name, bool query_imports = false);
+  /*! \return internal container */
+  inline ModuleNode* operator->();
+  /*! \return internal container */
+  inline const ModuleNode* operator->() const;
+  // The following functions requires link with runtime.
   /*!
    * \brief Import another module into this module.
    * \param other The module to be imported.
@@ -57,10 +63,6 @@ class Module {
    */
   TVM_DLL static Module LoadFromFile(const std::string& file_name,
                                      const std::string& format = "");
-  /*! \return internal container */
-  inline ModuleNode* operator->();
-  /*! \return internal container */
-  inline const ModuleNode* operator->() const;
 
  private:
   std::shared_ptr<ModuleNode> node_;

include/tvm/runtime/packed_func.h

@@ -183,31 +183,17 @@ struct extension_class_info {
 };
 
 /*!
- * \brief Runtime function table about extension type.
+ * \brief Capsule structure holding extension types
+ *  Capsule is self-contained and include
+ *  all the information to clone and destroy the type.
  */
-class ExtTypeVTable {
- public:
+struct TVMExtTypeCapsule {
+  /*! \brief The pointer to the object */
+  void* ptr;
   /*! \brief function to be called to delete a handle */
   void (*destroy)(void* handle);
   /*! \brief function to be called when clone a handle */
   void* (*clone)(void* handle);
-  /*!
-   * \brief Register type
-   * \tparam T The type to be register.
-   * \return The registered vtable.
-   */
-  template <typename T>
-  static inline ExtTypeVTable* Register_();
-  /*!
-   * \brief Get a vtable based on type code.
-   * \param type_code The type code
-   * \return The registered vtable.
-   */
-  TVM_DLL static ExtTypeVTable* Get(int type_code);
-
- private:
-  // Internal registration function.
-  TVM_DLL static ExtTypeVTable* RegisterInternal(int type_code, const ExtTypeVTable& vt);
 };
 
 /*!
@@ -255,8 +241,9 @@ class TVMPODValue_ {
   }
   template<typename TExtension>
   const TExtension& AsExtension() const {
-    CHECK_LT(type_code_, kExtEnd);
-    return static_cast<TExtension*>(value_.v_handle)[0];
+    CHECK_EQ(type_code_, extension_class_info<TExtension>::code);
+    return static_cast<TExtension*>(
+        static_cast<TVMExtTypeCapsule*>(value_.v_handle)->ptr)[0];
   }
   int type_code() const {
     return type_code_;
@@ -488,14 +475,6 @@ class TVMRetValue : public TVMPODValue_ {
     this->Assign(other);
     return *this;
   }
-  template<typename T,
-           typename = typename std::enable_if<
-             extension_class_info<T>::code != 0>::type>
-  TVMRetValue& operator=(const T& other) {
-    this->SwitchToClass<T>(
-        extension_class_info<T>::code, other);
-    return *this;
-  }
   /*!
    * \brief Move the value back to front-end via C API.
    *  This marks the current container as null.
@@ -521,6 +500,11 @@ class TVMRetValue : public TVMPODValue_ {
           type_code_ != kStr) << "TVMRetValue.value can only be used for POD data";
     return value_;
   }
+  // assign extension
+  template<typename T,
+           typename = typename std::enable_if<
+             extension_class_info<T>::code != 0>::type>
+  inline TVMRetValue& operator=(const T& other);
   // NodeRef related extenstions: in tvm/packed_func_ext.h
   template<typename T,
            typename = typename std::enable_if<
@@ -564,11 +548,9 @@ class TVMRetValue : public TVMPODValue_ {
           SwitchToPOD(other.type_code());
           value_ = other.value_;
         } else {
-          this->Clear();
-          type_code_ = other.type_code();
-          value_.v_handle =
-              (*(ExtTypeVTable::Get(other.type_code())->clone))(
-                  other.value().v_handle);
+          TVMExtTypeCapsule cap = *other.template ptr<TVMExtTypeCapsule>();
+          cap.ptr = cap.clone(cap.ptr);
+          SwitchToClass<TVMExtTypeCapsule>(other.type_code(), cap);
         }
         break;
       }
@@ -600,7 +582,9 @@ class TVMRetValue : public TVMPODValue_ {
       case kNodeHandle: delete ptr<std::shared_ptr<Node> >(); break;
     }
     if (type_code_ > kExtBegin) {
-      (*(ExtTypeVTable::Get(type_code_)->destroy))(value_.v_handle);
+      TVMExtTypeCapsule *cap = ptr<TVMExtTypeCapsule>();
+      cap->destroy(cap->ptr);
+      delete cap;
     }
     type_code_ = kNull;
   }
@@ -716,8 +700,10 @@ inline void for_each(const F& f, Args&&... args) {  // NOLINT(*)
 /* \brief argument settter to PackedFunc */
 class TVMArgsSetter {
  public:
-  TVMArgsSetter(TVMValue* values, int* type_codes)
-      : values_(values), type_codes_(type_codes) {}
+  TVMArgsSetter(TVMValue* values,
+                int* type_codes,
+                TVMExtTypeCapsule* exts)
+      : values_(values), type_codes_(type_codes), exts_(exts) {}
   // setters for POD types
   template<typename T,
            typename = typename std::enable_if<
@@ -807,15 +793,21 @@ class TVMArgsSetter {
   TVMValue* values_;
   /*! \brief The type code fields */
   int* type_codes_;
+  /*! \brief Temporary storage for extension types */
+  TVMExtTypeCapsule* exts_;
 };
 
 template<typename... Args>
 inline TVMRetValue PackedFunc::operator()(Args&& ...args) const {
   const int kNumArgs = sizeof...(Args);
+  // Compiler will remove an static array when it is not touched.
   const int kArraySize = kNumArgs > 0 ? kNumArgs : 1;
   TVMValue values[kArraySize];
   int type_codes[kArraySize];
-  detail::for_each(TVMArgsSetter(values, type_codes),
+  // If the function call does not contain extension type,
+  // exts will get optimized away by compiler.
+  TVMExtTypeCapsule exts[kArraySize];
+  detail::for_each(TVMArgsSetter(values, type_codes, exts),
                    std::forward<Args>(args)...);
   TVMRetValue rv;
   body_(TVMArgs(values, type_codes, kNumArgs), &rv);
@@ -853,14 +845,6 @@ inline TVMRetValue::operator T() const {
       ::Apply(this);
 }
 
-template<typename T, typename>
-inline void TVMArgsSetter::operator()(size_t i, const T& value) const {
-  static_assert(extension_class_info<T>::code != 0,
-                "Need to have extesion code");
-  type_codes_[i] = extension_class_info<T>::code;
-  values_[i].v_handle = const_cast<T*>(&value);
-}
-
 // extension type handling
 template<typename T>
 struct ExtTypeInfo {
@@ -872,16 +856,42 @@ struct ExtTypeInfo {
   }
 };
 
-template<typename T>
-inline ExtTypeVTable* ExtTypeVTable::Register_() {
-  const int code = extension_class_info<T>::code;
-  static_assert(code != 0,
-                "require extension_class_info traits to be declared with non-zero code");
-  ExtTypeVTable vt;
-  vt.clone = ExtTypeInfo<T>::clone;
-  vt.destroy = ExtTypeInfo<T>::destroy;
-  return ExtTypeVTable::RegisterInternal(code, vt);
+template<typename T, typename>
+inline TVMRetValue& TVMRetValue::operator=(const T& other) {
+  TVMExtTypeCapsule cap;
+  cap.clone = ExtTypeInfo<T>::clone;
+  cap.destroy = ExtTypeInfo<T>::destroy;
+  cap.ptr = new T(other);
+  SwitchToClass<TVMExtTypeCapsule>(
+      extension_class_info<T>::code, cap);
+  return *this;
+}
+
+template<typename T, typename>
+inline void TVMArgsSetter::operator()(size_t i, const T& value) const {
+  static_assert(extension_class_info<T>::code != 0,
+                "Need to have extesion code");
+  type_codes_[i] = extension_class_info<T>::code;
+  exts_[i].clone = ExtTypeInfo<T>::clone;
+  exts_[i].destroy = ExtTypeInfo<T>::destroy;
+  exts_[i].ptr = const_cast<T*>(&value);
+  values_[i].v_handle = &exts_[i];
 }
+
+// Implement Module::GetFunction
+// Put implementation in this file so we have seen the PackedFunc
+inline PackedFunc Module::GetFunction(const std::string& name, bool query_imports) {
+  PackedFunc pf = node_->GetFunction(name, node_);
+  if (pf != nullptr) return pf;
+  if (query_imports) {
+    for (const Module& m : node_->imports_) {
+      pf = m.node_->GetFunction(name, m.node_);
+      if (pf != nullptr) return pf;
+    }
+  }
+  return pf;
+}
+
 }  // namespace runtime
 }  // namespace tvm
 #endif  // TVM_RUNTIME_PACKED_FUNC_H_

python/tvm/_ffi/function.py

@@ -234,6 +234,31 @@ def list_global_func_names():
     return fnames
 
 
+def extract_ext_funcs(finit):
+    """
+    Extract the extension PackedFuncs from a C module.
+
+    Parameters
+    ----------
+    finit : ctypes function
+        a ctypes that takes signature of TVMExtensionDeclarer
+
+    Returns
+    -------
+    fdict : dict of str to Function
+        The extracted functions
+    """
+    fdict = {}
+    def _list(name, func):
+        fdict[name] = func
+    myf = convert_to_tvm_func(_list)
+    ret = finit(myf.handle)
+    _ = myf
+    if ret != 0:
+        raise RuntimeError("cannot initialize with %s" % finit)
+    return fdict
+
+
 def _get_api(f):
     flocal = f
     flocal.is_global = True

python/tvm/api.py

@@ -8,7 +8,7 @@ from ._ffi.base import string_types
 from ._ffi.node import register_node, NodeBase
 from ._ffi.node import convert_to_node as _convert_to_node
 from ._ffi.function import Function
-from ._ffi.function import _init_api, register_func, get_global_func
+from ._ffi.function import _init_api, register_func, get_global_func, extract_ext_funcs
 from ._ffi.function import convert_to_tvm_func as _convert_tvm_func
 from ._ffi.runtime_ctypes import TVMType
 from . import _api_internal

python/tvm/build_module.py

@@ -23,16 +23,16 @@ from . import target as _target
 from . import make
 
 class DumpIR(object):
-    """Dump IR for each pass.
-       With it, you can dump ir just like gcc/llvm.
-
-       How to use:
-       -----------
-       .. code-block:: python
+    """
+    Dump IR for each pass.
+    With it, you can dump ir just like gcc/llvm.
 
-          with tvm.build_config(dump_pass_ir=True)
-              run()
+    How to use:
+    -----------
+    .. code-block:: python
 
+    with tvm.build_config(dump_pass_ir=True)
+        run()
     """
     scope_level = 0
     def __init__(self):
@@ -40,9 +40,9 @@ class DumpIR(object):
         self._recover_list = []
 
     def decorate(self, func):
-        ''' decorate the pass function'''
+        """ decorate the pass function"""
         def dump(*args, **kwargs):
-            '''dump function'''
+            """dump function"""
             retv = func(*args, **kwargs)
             if not isinstance(retv, (_stmt.Stmt, container.LoweredFunc, container.Array)):
                 return retv
@@ -59,7 +59,7 @@ class DumpIR(object):
         return dump
 
     def decorate_irpass(self):
-        '''decorate ir_pass and ScheduleOps'''
+        """decorate ir_pass and ScheduleOps"""
         self._old_sgpass = schedule.ScheduleOps
         schedule.ScheduleOps = self.decorate(schedule.ScheduleOps)
         vset = vars(ir_pass)
@@ -71,7 +71,7 @@ class DumpIR(object):
             vset[k] = self.decorate(v) if isinstance(v, types.FunctionType) else v
 
     def decorate_custompass(self):
-        ''' decorate add_lower_pass pass in BuildConfig'''
+        """ decorate add_lower_pass pass in BuildConfig"""
         cfg = BuildConfig.current
         self._old_custom_pass = cfg.add_lower_pass
         custom_pass = cfg.add_lower_pass if cfg.add_lower_pass else []
@@ -79,7 +79,7 @@ class DumpIR(object):
         BuildConfig.current.add_lower_pass = pass_list
 
     def enter(self):
-        '''only decorate outermost nest'''
+        """only decorate outermost nest"""
         if DumpIR.scope_level > 0:
             return
         self.decorate_irpass()
@@ -88,7 +88,7 @@ class DumpIR(object):
         DumpIR.scope_level += 1
 
     def exit(self):
-        '''recover outermost nest'''
+        """recover outermost nest"""
         if DumpIR.scope_level > 1:
             return
         # recover decorated functions
@@ -163,6 +163,7 @@ class BuildConfig(NodeBase):
                 "'%s' object cannot set attribute '%s'" % (str(type(self)), name))
         return super(BuildConfig, self).__setattr__(name, value)
 
+
 def build_config(**kwargs):
     """Configure the build behavior by setting config variables.
 
@@ -226,6 +227,7 @@ def build_config(**kwargs):
             setattr(config, k, kwargs[k])
     return config
 
+
 if not _RUNTIME_ONLY:
     # BuildConfig is not available in tvm_runtime
     BuildConfig.current = build_config()
@@ -352,8 +354,10 @@ def lower(sch,
         stmt = f(stmt)
     if simple_mode:
         return stmt
+
     return ir_pass.MakeAPI(stmt, name, arg_list, 0, cfg.restricted_func)
 
+
 def build(sch,
           args=None,
           target=None,

src/runtime/c_runtime_api.cc

@@ -347,6 +347,14 @@ int TVMFuncCreateFromCFunc(TVMPackedCFunc func,
   API_END();
 }
 
+int TVMExtTypeFree(void* handle, int type_code) {
+  API_BEGIN();
+  TVMExtTypeCapsule* cap = static_cast<TVMExtTypeCapsule*>(handle);
+  cap->destroy(cap->ptr);
+  delete cap;
+  API_END();
+}
+
 int TVMArrayAlloc(const tvm_index_t* shape,
                   int ndim,
                   int dtype_code,

src/runtime/module.cc

@@ -13,19 +13,6 @@
 namespace tvm {
 namespace runtime {
 
-PackedFunc Module::GetFunction(
-    const std::string& name, bool query_imports) {
-  PackedFunc pf = node_->GetFunction(name, node_);
-  if (pf != nullptr) return pf;
-  if (query_imports) {
-    for (const Module& m : node_->imports_) {
-      pf = m.node_->GetFunction(name, m.node_);
-      if (pf != nullptr) return pf;
-    }
-  }
-  return pf;
-}
-
 void Module::Import(Module other) {
   // specially handle rpc
   if (!std::strcmp((*this)->type_key(), "rpc")) {

src/runtime/registry.cc

@@ -22,15 +22,10 @@ struct Registry::Manager {
   // and the resource can become invalid because of indeterminstic order of destruction.
   // The resources will only be recycled during program exit.
   std::unordered_map<std::string, Registry*> fmap;
-  // vtable for extension type
-  std::array<ExtTypeVTable, kExtEnd> ext_vtable;
   // mutex
   std::mutex mutex;
 
   Manager() {
-    for (auto& x : ext_vtable) {
-      x.destroy = nullptr;
-    }
   }
 
   static Manager* Global() {
@@ -88,24 +83,6 @@ std::vector<std::string> Registry::ListNames() {
   return keys;
 }
 
-ExtTypeVTable* ExtTypeVTable::Get(int type_code) {
-  CHECK(type_code > kExtBegin && type_code < kExtEnd);
-  Registry::Manager* m = Registry::Manager::Global();
-  ExtTypeVTable* vt = &(m->ext_vtable[type_code]);
-  CHECK(vt->destroy != nullptr)
-      << "Extension type not registered";
-  return vt;
-}
-
-ExtTypeVTable* ExtTypeVTable::RegisterInternal(
-    int type_code, const ExtTypeVTable& vt) {
-  CHECK(type_code > kExtBegin && type_code < kExtEnd);
-  Registry::Manager* m = Registry::Manager::Global();
-  std::lock_guard<std::mutex>(m->mutex);
-  ExtTypeVTable* pvt = &(m->ext_vtable[type_code]);
-  pvt[0] = vt;
-  return pvt;
-}
 }  // namespace runtime
 }  // namespace tvm
 
@@ -120,12 +97,6 @@ struct TVMFuncThreadLocalEntry {
 /*! \brief Thread local store that can be used to hold return values. */
 typedef dmlc::ThreadLocalStore<TVMFuncThreadLocalEntry> TVMFuncThreadLocalStore;
 
-int TVMExtTypeFree(void* handle, int type_code) {
-  API_BEGIN();
-  tvm::runtime::ExtTypeVTable::Get(type_code)->destroy(handle);
-  API_END();
-}
-
 int TVMFuncRegisterGlobal(
     const char* name, TVMFunctionHandle f, int override) {
   API_BEGIN();

tests/cpp/packed_func_test.cc

@@ -126,9 +126,6 @@ struct extension_class_info<test::IntVector> {
 }  // runtime
 }  // tvm
 
-// do registration, this need to be in cc file
-TVM_REGISTER_EXT_TYPE(test::IntVector);
-
 TEST(PackedFunc, ExtensionType) {
   using namespace tvm;
   using namespace tvm::runtime;

tests/scripts/task_python_integration.sh

@@ -6,6 +6,7 @@ rm -rf python/tvm/*.pyc python/tvm/*/*.pyc
 
 # Test TVM
 make cython || exit -1
+make cython3 || exit -1
 
 # Test extern package package
 cd apps/extension

topi/src/topi.cc

@@ -54,8 +54,6 @@ namespace topi {
 using namespace tvm;
 using namespace tvm::runtime;
 
-TVM_REGISTER_EXT_TYPE(tvm::Target);
-
 /*! \brief Canonicalize an argument that may be Array<Expr> or int to Array<Expr> */
 Array<Expr> ArrayOrInt(TVMArgValue arg) {
   if (arg.type_code() == kDLInt || arg.type_code() == kDLUInt) {