[FFI][ABI] ABI Updates to for future metadata and complex ordering (#18254)

This PR updates the ABI to enable potential future need
for getting metadata from a dynamically loaded module.

Orders the current static object into simple objects that have C
ABI and more complex one that may need c++.

These items changes ABI to be future compact before we freeze.

Author: tqchen

ffi/include/tvm/ffi/c_api.h

@@ -126,19 +126,22 @@ typedef enum {
   kTVMFFIError = 67,
   /*! \brief Function object. */
   kTVMFFIFunction = 68,
-  /*! \brief Array object. */
-  kTVMFFIArray = 69,
-  /*! \brief Map object. */
-  kTVMFFIMap = 70,
   /*!
    * \brief Shape object, layout = { TVMFFIObject, { const int64_t*, size_t }, ... }
    */
-  kTVMFFIShape = 71,
+  kTVMFFIShape = 69,
   /*!
    * \brief NDArray object, layout = { TVMFFIObject, DLTensor, ... }
    */
-  kTVMFFINDArray = 72,
-  /*! \brief Runtime module object. */
+  kTVMFFINDArray = 70,
+  /*! \brief Array object. */
+  kTVMFFIArray = 71,
+  //----------------------------------------------------------------
+  // more complex objects
+  //----------------------------------------------------------------
+  /*! \brief Map object. */
+  kTVMFFIMap = 72,
+  /*! \brief Runtime dynamic loaded module object. */
   kTVMFFIModule = 73,
   kTVMFFIStaticObjectEnd,
   // [Section] Dynamic Boxed: [kTVMFFIDynObjectBegin, +oo)
@@ -763,11 +766,11 @@ typedef struct TVMFFITypeInfo {
  *
  * \param name The name of the function.
  * \param f The function to be registered.
- * \param override Whether allow override already registered function.
+ * \param allow_override Whether allow override already registered function.
  * \return 0 when success, nonzero when failure happens
  */
 TVM_FFI_DLL int TVMFFIFunctionSetGlobal(const TVMFFIByteArray* name, TVMFFIObjectHandle f,
-                                        int override);
+                                        int allow_override);
 
 /*!
  * \brief Register the function to runtime's global table with method info.
@@ -780,7 +783,7 @@ TVM_FFI_DLL int TVMFFIFunctionSetGlobal(const TVMFFIByteArray* name, TVMFFIObjec
  * \return 0 when success, nonzero when failure happens
  */
 TVM_FFI_DLL int TVMFFIFunctionSetGlobalFromMethodInfo(const TVMFFIMethodInfo* method_info,
-                                                      int override);
+                                                      int allow_override);
 
 /*!
  * \brief Register type field information for runtime reflection.

ffi/include/tvm/ffi/extra/module.h

@@ -68,6 +68,12 @@ class TVM_FFI_EXTRA_CXX_API ModuleObj : public Object {
    * \return True if the module implements the function, false otherwise.
    */
   virtual bool ImplementsFunction(const String& name) { return GetFunction(name).defined(); }
+  /*!
+   * \brief Get the metadata of the function, if available.
+   * \param name The name of the function.
+   * \return The metadata stored in json string format.
+   */
+  virtual Optional<String> GetFunctionMetadata(const String& name) { return std::nullopt; }
   /*!
    * \brief Write the current module to file with given format (for further compilation).
    *
@@ -121,6 +127,12 @@ class TVM_FFI_EXTRA_CXX_API ModuleObj : public Object {
    * \return True if the module implements the function, false otherwise.
    */
   bool ImplementsFunction(const String& name, bool query_imports);
+  /*!
+   * \brief Get the function metadata of the function if available.
+   * \param name The name of the function.
+   * \return The function metadata of the function in json format.
+   */
+  Optional<String> GetFunctionMetadata(const String& name, bool query_imports);
   /*!
    * \brief Get the imports of the module.
    * \return The imports of the module.
@@ -215,6 +227,8 @@ namespace symbol {
 constexpr const char* tvm_ffi_library_ctx = "__tvm_ffi_library_ctx";
 /*! \brief Global variable to store binary data alongside a library module. */
 constexpr const char* tvm_ffi_library_bin = "__tvm_ffi_library_bin";
+/*! \brief Optional metadata prefix of a symbol. */
+constexpr const char* tvm_ffi_metadata_prefix = "__tvm_ffi_metadata_";
 /*! \brief Default entry function of a library module. */
 constexpr const char* tvm_ffi_main = "__tvm_ffi_main__";
 }  // namespace symbol

ffi/pyproject.toml

@@ -17,7 +17,7 @@
 
 [project]
 name = "apache-tvm-ffi"
-version = "0.1.0a2"
+version = "0.1.0a3"
 description = "tvm ffi"
 
 authors = [{ name = "TVM FFI team" }]

ffi/python/tvm_ffi/cython/base.pxi

@@ -48,12 +48,13 @@ cdef extern from "tvm/ffi/c_api.h":
         kTVMFFIBytes = 66
         kTVMFFIError = 67
         kTVMFFIFunction = 68
-        kTVMFFIArray = 69
-        kTVMFFIMap = 70
-        kTVMFFIShape = 71
-        kTVMFFINDArray = 72
+        kTVMFFIShape = 69
+        kTVMFFINDArray = 70
+        kTVMFFIArray = 71
+        kTVMFFIMap = 72
         kTVMFFIModule = 73
 
+
     ctypedef void* TVMFFIObjectHandle
 
     ctypedef struct DLDataType:

ffi/scripts/run_tests.sh

@@ -17,7 +17,9 @@
 # under the License.
 set -euxo pipefail
 
-BUILD_TYPE=Release
+BUILD_TYPE=RelWithDebugInfo
+
+rm -rf build/CMakeCache.txt
 
 cmake -G Ninja -S . -B build -DTVM_FFI_BUILD_TESTS=ON -DCMAKE_BUILD_TYPE=${BUILD_TYPE} \
   -DCMAKE_EXPORT_COMPILE_COMMANDS=ON -DCMAKE_CXX_COMPILER_LAUNCHER=ccache

ffi/src/ffi/extra/module.cc

@@ -44,6 +44,20 @@ Optional<Function> ModuleObj::GetFunction(const String& name, bool query_imports
   return std::nullopt;
 }
 
+Optional<String> ModuleObj::GetFunctionMetadata(const String& name, bool query_imports) {
+  if (auto opt_metadata = this->GetFunctionMetadata(name)) {
+    return opt_metadata;
+  }
+  if (query_imports) {
+    for (const Any& import : imports_) {
+      if (auto opt_metadata = import.cast<Module>()->GetFunctionMetadata(name, query_imports)) {
+        return *opt_metadata;
+      }
+    }
+  }
+  return std::nullopt;
+}
+
 void ModuleObj::ImportModule(const Module& other) {
   std::unordered_set<const ModuleObj*> visited{other.operator->()};
   std::vector<const ModuleObj*> stack{other.operator->()};
@@ -115,6 +129,10 @@ TVM_FFI_STATIC_INIT_BLOCK({
                   [](Module mod, String name, bool query_imports) {
                     return mod->ImplementsFunction(name, query_imports);
                   })
+      .def_method("ffi.ModuleGetFunctionMetadata",
+                  [](Module mod, String name, bool query_imports) {
+                    return mod->GetFunctionMetadata(name, query_imports);
+                  })
       .def_method("ffi.ModuleGetFunction",
                   [](Module mod, String name, bool query_imports) {
                     return mod->GetFunction(name, query_imports);

jvm/core/src/main/java/org/apache/tvm/TypeIndex.java

@@ -36,9 +36,9 @@ public class TypeIndex {
   public static final int kTVMFFIBytes = 66;
   public static final int kTVMFFIError = 67;
   public static final int kTVMFFIFunction = 68;
-  public static final int kTVMFFIArray = 69;
-  public static final int kTVMFFIMap = 70;
-  public static final int kTVMFFIShape = 71;
-  public static final int kTVMFFINDArray = 72;
+  public static final int kTVMFFIShape = 70;
+  public static final int kTVMFFINDArray = 71;
+  public static final int kTVMFFIArray = 72;
+  public static final int kTVMFFIMap = 73;
   public static final int kTVMFFIModule = 73;
 }

web/src/ctypes.ts

@@ -97,16 +97,16 @@ export const enum TypeIndex {
   kTVMFFIFunction = 68,
   /*! \brief Array object. */
   kTVMFFIArray = 69,
-  /*! \brief Map object. */
-  kTVMFFIMap = 70,
   /*!
    * \brief Shape object, layout = { TVMFFIObject, { const int64_t*, size_t }, ... }
    */
-  kTVMFFIShape = 71,
+  kTVMFFIShape = 70,
   /*!
    * \brief NDArray object, layout = { TVMFFIObject, DLTensor, ... }
    */
-  kTVMFFINDArray = 72,
+  kTVMFFINDArray = 71,
+  /*! \brief Map object. */
+  kTVMFFIMap = 72,
   /*! \brief Runtime module object. */
   kTVMFFIModule = 73,
 }