[FFI] Formalize ffi.Module

apps/android_rpc/app/src/main/jni/tvm_runtime.h

@@ -37,20 +37,21 @@
 #include "../ffi/src/ffi/container.cc"
 #include "../ffi/src/ffi/dtype.cc"
 #include "../ffi/src/ffi/error.cc"
+#include "../ffi/src/ffi/extra/library_module.cc"
+#include "../ffi/src/ffi/extra/library_module_dynamic_lib.cc"
+#include "../ffi/src/ffi/extra/library_module_system_lib.cc"
+#include "../ffi/src/ffi/extra/module.cc"
 #include "../ffi/src/ffi/function.cc"
 #include "../ffi/src/ffi/ndarray.cc"
 #include "../ffi/src/ffi/object.cc"
 #include "../ffi/src/ffi/testing.cc"
 #include "../ffi/src/ffi/traceback.cc"
 #include "../src/runtime/cpu_device_api.cc"
 #include "../src/runtime/device_api.cc"
-#include "../src/runtime/dso_library.cc"
 #include "../src/runtime/file_utils.cc"
-#include "../src/runtime/library_module.cc"
 #include "../src/runtime/logging.cc"
 #include "../src/runtime/memory/memory_manager.cc"
 #include "../src/runtime/minrpc/minrpc_logger.cc"
-#include "../src/runtime/module.cc"
 #include "../src/runtime/ndarray.cc"
 #include "../src/runtime/profiling.cc"
 #include "../src/runtime/registry.cc"
@@ -62,7 +63,6 @@
 #include "../src/runtime/rpc/rpc_server_env.cc"
 #include "../src/runtime/rpc/rpc_session.cc"
 #include "../src/runtime/rpc/rpc_socket_impl.cc"
-#include "../src/runtime/system_library.cc"
 #include "../src/runtime/thread_pool.cc"
 #include "../src/runtime/threading_backend.cc"
 #include "../src/runtime/workspace_pool.cc"

apps/cpp_rpc/rpc_env.cc

@@ -147,7 +147,7 @@ RPCEnv::RPCEnv(const std::string& wd) {
                              std::string file_name = this->GetPath(path);
                              file_name = BuildSharedLibrary(file_name);
                              LOG(INFO) << "Load module from " << file_name << " ...";
-                             return Module::LoadFromFile(file_name, "");
+                             return ffi::Module::LoadFromFile(file_name);
                            }));
 
   ffi::Function::SetGlobal("tvm.rpc.server.download_linked_module",

apps/hexagon_launcher/launcher_core.cc

@@ -144,7 +144,7 @@ const tvm::ffi::Function get_runtime_func(const std::string& name) {
 }
 
 const tvm::ffi::Function get_module_func(tvm::runtime::Module module, const std::string& name) {
-  return module.GetFunction(name, false);
+  return module->GetFunction(name, false).value_or(tvm::ffi::Function());
 }
 
 void reset_device_api() {
@@ -153,7 +153,7 @@ void reset_device_api() {
 }
 
 tvm::runtime::Module load_module(const std::string& file_name) {
-  static const tvm::ffi::Function loader = get_runtime_func("runtime.module.loadfile_hexagon");
+  static const tvm::ffi::Function loader = get_runtime_func("ffi.Module.load_from_file.hexagon");
   tvm::ffi::Any rv = loader(file_name);
   if (rv.type_code() == kTVMModuleHandle) {
     ICHECK_EQ(rv.type_code(), kTVMModuleHandle)

apps/ios_rpc/tvmrpc/TVMRuntime.mm

@@ -33,7 +33,7 @@
 
 #if defined(USE_CUSTOM_DSO_LOADER) && USE_CUSTOM_DSO_LOADER == 1
 // internal TVM header to achieve Library class
-#include <../../../src/runtime/library_module.h>
+#include <../../../ffi/src/ffi/extra/library_module.h>
 #include <custom_dlfcn.h>
 #endif
 
@@ -70,7 +70,7 @@ void LogMessageImpl(const std::string& file, int lineno, int level, const std::s
           NSBundle* bundle = [NSBundle mainBundle];
           base = [[bundle privateFrameworksPath] stringByAppendingPathComponent:@"tvm"];
 
-          if (tvm::ffi::Function::GetGlobal("runtime.module.loadfile_dylib_custom")) {
+          if (tvm::ffi::Function::GetGlobal("ffi.Module.load_from_file.dylib_custom")) {
             // Custom dso loader is present. Will use it.
             base = NSTemporaryDirectory();
             fmt = "dylib_custom";
@@ -114,11 +114,11 @@ void Init(const std::string& name) {
 // Add UnsignedDSOLoader plugin in global registry
 TVM_FFI_STATIC_INIT_BLOCK({
   namespace refl = tvm::ffi::reflection;
-  refl::GlobalDef().def_packed("runtime.module.loadfile_dylib_custom",
+  refl::GlobalDef().def_packed("ffi.Module.load_from_file.dylib_custom",
                                [](ffi::PackedArgs args, ffi::Any* rv) {
                                  auto n = make_object<UnsignedDSOLoader>();
                                  n->Init(args[0]);
-                                 *rv = CreateModuleFromLibrary(n);
+                                 *rv = tvm::ffi::CreateLibraryModule(n);
                                });
 });
 

ffi/CMakeLists.txt

@@ -69,6 +69,10 @@ if (TVM_FFI_USE_EXTRA_CXX_API)
     "${CMAKE_CURRENT_SOURCE_DIR}/src/ffi/extra/json_writer.cc"
     "${CMAKE_CURRENT_SOURCE_DIR}/src/ffi/extra/serialization.cc"
     "${CMAKE_CURRENT_SOURCE_DIR}/src/ffi/extra/reflection_extra.cc"
+    "${CMAKE_CURRENT_SOURCE_DIR}/src/ffi/extra/module.cc"
+    "${CMAKE_CURRENT_SOURCE_DIR}/src/ffi/extra/library_module.cc"
+    "${CMAKE_CURRENT_SOURCE_DIR}/src/ffi/extra/library_module_system_lib.cc"
+    "${CMAKE_CURRENT_SOURCE_DIR}/src/ffi/extra/library_module_dynamic_lib.cc"
   )
 endif()
 

ffi/include/tvm/ffi/extra/c_env_api.h

@@ -0,0 +1,70 @@
+/*
+ * 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/ffi/extra/c_env_api.h
+ * \brief Extra environment API.
+ */
+#ifndef TVM_FFI_EXTRA_C_ENV_API_H_
+#define TVM_FFI_EXTRA_C_ENV_API_H_
+
+#include <tvm/ffi/c_api.h>
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+/*!
+ * \brief FFI function to lookup a function from a module's imports.
+ *
+ * This is a helper function that is used by generated code.
+ *
+ * \param library_ctx The library context module handle.
+ * \param func_name The name of the function.
+ * \param out The result function.
+ * \note The returned function is a weak reference that is cached/owned by the module.
+ * \return 0 when no error is thrown, -1 when failure happens
+ */
+TVM_FFI_DLL int TVMFFIEnvLookupFromImports(TVMFFIObjectHandle library_ctx, const char* func_name,
+                                           TVMFFIObjectHandle* out);
+
+/*
+ * \brief Register a symbol value that will be initialized when a library with the symbol is loaded.
+ *
+ * This function can be used to make context functions to be available in the library
+ * module that wants to avoid an explicit link dependency
+ *
+ * \param name The name of the symbol.
+ * \param symbol The symbol to register.
+ * \return 0 when success, nonzero when failure happens
+ */
+TVM_FFI_DLL int TVMFFIEnvRegisterContextSymbol(const char* name, void* symbol);
+
+/*!
+ * \brief Register a symbol that will be initialized when a system library is loaded.
+ *
+ * \param name The name of the symbol.
+ * \param symbol The symbol to register.
+ * \return 0 when success, nonzero when failure happens
+ */
+TVM_FFI_DLL int TVMFFIEnvRegisterSystemLibSymbol(const char* name, void* symbol);
+
+#ifdef __cplusplus
+}  // extern "C"
+#endif
+#endif  // TVM_FFI_EXTRA_C_ENV_API_H_

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

@@ -0,0 +1,224 @@
+/*
+ * 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/ffi/module.h
+ * \brief A managed dynamic module in the TVM FFI.
+ */
+#ifndef TVM_FFI_EXTRA_MODULE_H_
+#define TVM_FFI_EXTRA_MODULE_H_
+
+#include <tvm/ffi/container/array.h>
+#include <tvm/ffi/container/map.h>
+#include <tvm/ffi/extra/base.h>
+#include <tvm/ffi/function.h>
+
+namespace tvm {
+namespace ffi {
+
+// forward declare Module
+class Module;
+
+/*!
+ * \brief A module that can dynamically load ffi::Functions or exportable source code.
+ */
+class TVM_FFI_EXTRA_CXX_API ModuleObj : public Object {
+ public:
+  /*!
+   * \return The per module type key.
+   * \note This key is used to for serializing custom modules.
+   */
+  virtual const char* kind() const = 0;
+  /*!
+   * \brief Get the property mask of the module.
+   * \return The property mask of the module.
+   *
+   * \sa Module::ModulePropertyMask
+   */
+  virtual int GetPropertyMask() const { return 0b000; }
+  /*!
+   * \brief Get a ffi::Function from the module.
+   * \param name The name of the function.
+   * \return The function.
+   */
+  virtual Optional<Function> GetFunction(const String& name) = 0;
+  /*!
+   * \brief Returns true if this module has a definition for a function of \p name.
+   *
+   * Note that even if this function returns true the corresponding \p GetFunction result
+   * may be nullptr if the function is not yet callable without further compilation.
+   *
+   * The default implementation just checks if \p GetFunction is non-null.
+   * \param name The name of the function.
+   * \return True if the module implements the function, false otherwise.
+   */
+  virtual bool ImplementsFunction(const String& name) { return GetFunction(name).defined(); }
+  /*!
+   * \brief Write the current module to file with given format (for further compilation).
+   *
+   * \param file_name The file to be saved to.
+   * \param format The format of the file.
+   *
+   * \note This function is mainly used by modules that
+   */
+  virtual void WriteToFile(const String& file_name, const String& format) const {
+    TVM_FFI_THROW(RuntimeError) << "Module[" << kind() << "] does not support WriteToFile";
+  }
+  /*!
+   * \brief Get the possible write formats of the module, when available.
+   * \return Possible write formats when available.
+   */
+  virtual Array<String> GetWriteFormats() const { return Array<String>(); }
+  /*!
+   * \brief Serialize the the module to bytes.
+   * \return The serialized module.
+   */
+  virtual Bytes SaveToBytes() const {
+    TVM_FFI_THROW(RuntimeError) << "Module[" << kind() << "] does not support SaveToBytes";
+    TVM_FFI_UNREACHABLE();
+  }
+  /*!
+   * \brief Get the source code of module, when available.
+   * \param format Format of the source code, can be empty by default.
+   * \return Possible source code when available, or empty string if not available.
+   */
+  virtual String InspectSource(const String& format = "") const { return String(); }
+  /*!
+   * \brief Import another module.
+   * \param other The module to import.
+   */
+  virtual void ImportModule(const Module& other);
+  /*!
+   * \brief Clear all imported modules.
+   */
+  virtual void ClearImports();
+  /*!
+   * \brief Overloaded function to optionally query from imports.
+   * \param name The name of the function.
+   * \param query_imports Whether to query imported modules.
+   * \return The function.
+   */
+  Optional<Function> GetFunction(const String& name, bool query_imports);
+  /*!
+   * \brief Overloaded function to optionally query from imports.
+   * \param name The name of the function.
+   * \param query_imports Whether to query imported modules.
+   * \return True if the module implements the function, false otherwise.
+   */
+  bool ImplementsFunction(const String& name, bool query_imports);
+  /*!
+   * \brief Get the imports of the module.
+   * \return The imports of the module.
+   * \note Note the signature is not part of the public API.
+   */
+  const Array<Any>& imports() const { return this->imports_; }
+
+  struct InternalUnsafe;
+
+  static constexpr const int32_t _type_index = TypeIndex::kTVMFFIModule;
+  static constexpr const char* _type_key = StaticTypeKey::kTVMFFIModule;
+  static const constexpr bool _type_final = true;
+  TVM_FFI_DECLARE_STATIC_OBJECT_INFO(ModuleObj, Object);
+
+ protected:
+  friend struct InternalUnsafe;
+
+  /*!
+   * \brief The modules that this module depends on.
+   * \note Use ObjectRef to avoid circular dep on Module.
+   */
+  Array<Any> imports_;
+
+ private:
+  /*!
+   * \brief cache used by TVMFFIModuleLookupFromImports
+   */
+  Map<String, ffi::Function> import_lookup_cache_;
+};
+
+/*!
+ * \brief Reference to module object.
+ */
+class Module : public ObjectRef {
+ public:
+  /*!
+   * \brief Property of ffi::Module
+   */
+  enum ModulePropertyMask : int {
+    /*!
+     * \brief The module can be serialized to bytes.
+     *
+     * This prooperty indicates that module implements SaveToBytes.
+     * The system also registers a GlobalDef function
+     * `ffi.Module.load_from_bytes.<kind>` with signature (Bytes) -> Module.
+     */
+    kBinarySerializable = 0b001,
+    /*!
+     * \brief The module can directly get runnable functions.
+     *
+     * This property indicates that module implements GetFunction that returns
+     * runnable ffi::Functions.
+     */
+    kRunnable = 0b010,
+    /*!
+     * \brief The module can be exported to a object file or source file that then be compiled.
+     *
+     * This property indicates that module implements WriteToFile with a given format
+     * that can be queried by GetLibExportFormat.
+     *
+     * Examples include modules that can be exported to .o, .cc, .cu files.
+     *
+     * Such modules can be exported, compiled and loaded back as a dynamic library module.
+     */
+    kCompilationExportable = 0b100
+  };
+
+  /*!
+   * \brief Load a module from file.
+   * \param file_name The name of the host function module.
+   * \param format The format of the file.
+   * \note This function won't load the import relationship.
+   *  Re-create import relationship by calling Import.
+   */
+  TVM_FFI_EXTRA_CXX_API static Module LoadFromFile(const String& file_name);
+  /*
+   * \brief Query context symbols that is registered via TVMEnvRegisterSymbols.
+   * \param callback The callback to be called with the symbol name and address.
+   * \note This helper can be used to implement custom Module that needs to access context symbols.
+   */
+  TVM_FFI_EXTRA_CXX_API static void VisitContextSymbols(
+      const ffi::TypedFunction<void(String, void*)>& callback);
+
+  TVM_FFI_DEFINE_MUTABLE_NOTNULLABLE_OBJECT_REF_METHODS(Module, ObjectRef, ModuleObj);
+};
+
+/*
+ * \brief Symbols for library module.
+ */
+namespace symbol {
+/*! \brief Global variable to store context pointer for a library module. */
+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 Default entry function of a library module. */
+constexpr const char* tvm_ffi_main = "__tvm_ffi_main__";
+}  // namespace symbol
+}  // namespace ffi
+}  // namespace tvm
+
+#endif  // TVM_FFI_EXTRA_MODULE_H_