[REFACTOR] Phase out include/tvm/runtime/module.h

apps/cpp_rpc/rpc_env.cc

@@ -20,9 +20,9 @@
  * \file rpc_env.cc
  * \brief Server environment of the RPC.
  */
+#include <tvm/ffi/extra/module.h>
 #include <tvm/ffi/function.h>
 #include <tvm/runtime/logging.h>
-#include <tvm/runtime/module.h>
 
 #include <cerrno>
 #ifndef _WIN32

apps/hexagon_launcher/launcher_core.h

@@ -22,9 +22,9 @@
 
 #include <dlpack/dlpack.h>
 #include <tvm/ffi/extra/json.h>
+#include <tvm/ffi/extra/module.h>
 #include <tvm/ffi/function.h>
 #include <tvm/runtime/data_type.h>
-#include <tvm/runtime/module.h>
 #include <tvm/runtime/tensor.h>
 
 #include <string>
@@ -81,10 +81,10 @@ struct OutputConfig {
 };
 
 struct Model {
-  Model(tvm::runtime::Module executor, tvm::runtime::Module module, std::string json);
+  Model(tvm::ffi::Module executor, tvm::ffi::Module module, std::string json);
 
-  tvm::runtime::Module model_executor;
-  tvm::runtime::Module graph_module;
+  tvm::ffi::Module model_executor;
+  tvm::ffi::Module graph_module;
   std::string graph_json;
 
   static tvm::Device device() { return tvm::Device{static_cast<DLDeviceType>(kDLHexagon), 0}; }
@@ -121,13 +121,13 @@ bool write_output_config(const std::string& file_name, OutputConfig* output_conf
 
 void reset_device_api();
 
-tvm::runtime::Module load_module(const std::string& file_name);
+tvm::ffi::Module load_module(const std::string& file_name);
 
 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);
+const tvm::ffi::Function get_module_func(tvm::ffi::Module module, const std::string& name);
 
-tvm::runtime::Module create_aot_executor(tvm::runtime::Module factory_module, tvm::Device device);
-tvm::runtime::Module create_graph_executor(const std::string& graph_json,
-                                           tvm::runtime::Module graph_module, tvm::Device device);
+tvm::ffi::Module create_aot_executor(tvm::ffi::Module factory_module, tvm::Device device);
+tvm::ffi::Module create_graph_executor(const std::string& graph_json, tvm::ffi::Module graph_module,
+                                       tvm::Device device);
 
 #endif  // TVM_RUNTIME_HEXAGON_LAUNCHER_LAUNCHER_CORE_H_

docs/arch/runtime.rst

@@ -169,7 +169,7 @@ TVM defines the compiled object as `Module`_.
 The user can get the compiled function from Module as PackedFunc.
 The generated compiled code can dynamically get function from Module in runtime. It caches the function handle in the first call and reuses in subsequent calls. We use this to link device code and callback into any PackedFunc(e.g., python) from generated code.
 
-.. _Module: https://github.com/apache/tvm/blob/main/include/tvm/runtime/module.h
+.. _Module: https://github.com/apache/tvm/blob/main/3rdparty/tvm-ffi/include/tvm/ffi/extra/module.h
 
 The ModuleNode is an abstract class that can be implemented by each type of device.
 So far we support modules for CUDA, Metal, OpenCL and loading dynamic shared libraries. This abstraction makes introduction

include/tvm/runtime/device_api.h

@@ -26,6 +26,7 @@
 
 #include <tvm/ffi/any.h>
 #include <tvm/ffi/optional.h>
+#include <tvm/ffi/string.h>
 #include <tvm/runtime/base.h>
 #include <tvm/runtime/logging.h>
 
@@ -404,6 +405,23 @@ inline Device AddRPCSessionMask(Device dev, int session_table_index) {
   return dev;
 }
 
+/*!
+ * \brief Check if runtime module is enabled for target.
+ * \param target The target module name.
+ * \return Whether runtime is enabled.
+ */
+TVM_DLL bool RuntimeEnabled(const ffi::String& target);
+
+/*! \brief namespace for constant symbols */
+namespace symbol {
+/*! \brief global function to set device */
+constexpr const char* tvm_set_device = "__tvm_set_device";
+/*! \brief Auxiliary counter to global barrier. */
+constexpr const char* tvm_global_barrier_state = "__tvm_global_barrier_state";
+/*! \brief Prepare the global barrier before kernels that uses global barrier. */
+constexpr const char* tvm_prepare_global_barrier = "__tvm_prepare_global_barrier";
+}  // namespace symbol
+
 }  // namespace runtime
 }  // namespace tvm
 

include/tvm/runtime/disco/builtin.h

@@ -19,8 +19,8 @@
 #ifndef TVM_RUNTIME_DISCO_BUILTIN_H_
 #define TVM_RUNTIME_DISCO_BUILTIN_H_
 
+#include <tvm/ffi/extra/module.h>
 #include <tvm/runtime/data_type.h>
-#include <tvm/runtime/module.h>
 #include <tvm/runtime/tensor.h>
 
 #include <string>

include/tvm/runtime/profiling.h

@@ -26,10 +26,10 @@
 
 #include <tvm/ffi/container/array.h>
 #include <tvm/ffi/container/map.h>
+#include <tvm/ffi/extra/module.h>
 #include <tvm/ffi/function.h>
 #include <tvm/runtime/base.h>
 #include <tvm/runtime/device_api.h>
-#include <tvm/runtime/module.h>
 #include <tvm/runtime/object.h>
 #include <tvm/runtime/tensor.h>
 

include/tvm/runtime/vm/executable.h

@@ -23,8 +23,8 @@
 #ifndef TVM_RUNTIME_VM_EXECUTABLE_H_
 #define TVM_RUNTIME_VM_EXECUTABLE_H_
 
+#include <tvm/ffi/extra/module.h>
 #include <tvm/ffi/function.h>
-#include <tvm/runtime/module.h>
 #include <tvm/runtime/object.h>
 #include <tvm/support/io.h>
 #include <tvm/support/serializer.h>
@@ -166,14 +166,14 @@ class VMExecutable : public ffi::ModuleObj {
 
   virtual ~VMExecutable() {}
 
-  TVM_MODULE_VTABLE_BEGIN("relax.VMExecutable");
-  TVM_MODULE_VTABLE_ENTRY("stats", &VMExecutable::Stats);
-  TVM_MODULE_VTABLE_ENTRY("as_text", &VMExecutable::AsText);
-  TVM_MODULE_VTABLE_ENTRY("as_python", &VMExecutable::AsPython);
-  TVM_MODULE_VTABLE_ENTRY("vm_load_executable", &VMExecutable::VMLoadExecutable);
-  TVM_MODULE_VTABLE_ENTRY("vm_profiler_load_executable", &VMExecutable::VMProfilerLoadExecutable);
-  TVM_MODULE_VTABLE_ENTRY("has_function", &VMExecutable::HasFunction);
-  TVM_MODULE_VTABLE_END();
+  /*! \brief Module type key. */
+  const char* kind() const final;
+  /*!
+   * \brief Look up an exported function by name.
+   * \param name The function name.
+   * \return The function if found, otherwise std::nullopt.
+   */
+  ffi::Optional<ffi::Function> GetFunction(const ffi::String& name) override;
 
  private:
   /*!

include/tvm/target/codegen.h

@@ -24,8 +24,8 @@
 #ifndef TVM_TARGET_CODEGEN_H_
 #define TVM_TARGET_CODEGEN_H_
 
+#include <tvm/ffi/extra/module.h>
 #include <tvm/ir/module.h>
-#include <tvm/runtime/module.h>
 #include <tvm/target/target.h>
 #include <tvm/tirx/expr.h>
 

src/ir/op.cc

@@ -25,7 +25,6 @@
 #include <tvm/ffi/reflection/registry.h>
 #include <tvm/ir/op.h>
 #include <tvm/ir/type.h>
-#include <tvm/runtime/module.h>
 #include <tvm/tirx/op_attr_types.h>
 
 #include <memory>

src/relax/backend/contrib/cutlass/codegen.cc

@@ -21,13 +21,13 @@
  * \file src/relax/backend/contrib/cutlass/codegen.cc
  * \brief Implementation of the CUTLASS code generator for Relax.
  */
+#include <tvm/ffi/extra/module.h>
 #include <tvm/ffi/reflection/registry.h>
 #include <tvm/ir/module.h>
 #include <tvm/ir/name_supply.h>
 #include <tvm/relax/analysis.h>
 #include <tvm/relax/attrs/nn.h>
 #include <tvm/relax/type.h>
-#include <tvm/runtime/module.h>
 
 #include <memory>
 #include <string>

src/relax/transform/fold_constant.cc

@@ -17,14 +17,14 @@
  * under the License.
  */
 
+#include <tvm/ffi/extra/module.h>
 #include <tvm/ffi/reflection/registry.h>
 #include <tvm/ir/function.h>
 #include <tvm/relax/analysis.h>
 #include <tvm/relax/expr_functor.h>
 #include <tvm/relax/op_attr_types.h>
 #include <tvm/relax/transform.h>
 #include <tvm/relax/type.h>
-#include <tvm/runtime/module.h>
 #include <tvm/tirx/function.h>
 #include <tvm/tirx/op.h>
 

src/relax/transform/run_codegen.cc

@@ -23,11 +23,11 @@
  * \brief Run codegen for annotated relax functions.
  */
 
+#include <tvm/ffi/extra/module.h>
 #include <tvm/ffi/reflection/registry.h>
 #include <tvm/relax/analysis.h>
 #include <tvm/relax/expr_functor.h>
 #include <tvm/relax/transform.h>
-#include <tvm/runtime/module.h>
 
 #include "../../support/ordered_set.h"
 #include "utils.h"

src/runtime/cuda/cuda_module.cc

@@ -27,6 +27,7 @@
 #include <tvm/ffi/extra/c_env_api.h>
 #include <tvm/ffi/function.h>
 #include <tvm/ffi/reflection/registry.h>
+#include <tvm/runtime/device_api.h>
 #include <tvm/support/io.h>
 
 #include <array>

src/runtime/cuda/cuda_module.h

@@ -24,7 +24,7 @@
 #ifndef TVM_RUNTIME_CUDA_CUDA_MODULE_H_
 #define TVM_RUNTIME_CUDA_CUDA_MODULE_H_
 
-#include <tvm/runtime/module.h>
+#include <tvm/ffi/extra/module.h>
 
 #include <memory>
 #include <string>

src/runtime/device_api.cc

@@ -31,7 +31,6 @@
 #include <tvm/runtime/base.h>
 #include <tvm/runtime/c_backend_api.h>
 #include <tvm/runtime/device_api.h>
-#include <tvm/runtime/module.h>
 
 #include <algorithm>
 #include <array>

src/runtime/hexagon/hexagon_buffer.cc

@@ -18,8 +18,6 @@
  */
 #include "hexagon_buffer.h"
 
-#include <tvm/runtime/module.h>
-
 #include <algorithm>
 #include <string>
 #include <utility>

src/runtime/metadata.h

@@ -27,8 +27,8 @@
 #include <tvm/ffi/container/array.h>
 #include <tvm/ffi/container/map.h>
 #include <tvm/ffi/extra/json.h>
+#include <tvm/ffi/extra/module.h>
 #include <tvm/ffi/function.h>
-#include <tvm/runtime/module.h>
 #include <tvm/runtime/tensor.h>
 #include <tvm/support/io.h>
 #include <tvm/support/serializer.h>

src/runtime/module.cc

@@ -25,7 +25,7 @@
 #include <tvm/ffi/function.h>
 #include <tvm/ffi/reflection/registry.h>
 #include <tvm/runtime/c_backend_api.h>
-#include <tvm/runtime/module.h>
+#include <tvm/runtime/device_api.h>
 
 #include <cstring>
 #include <unordered_set>

src/runtime/rpc/rpc_session.h

@@ -24,9 +24,9 @@
 #ifndef TVM_RUNTIME_RPC_RPC_SESSION_H_
 #define TVM_RUNTIME_RPC_RPC_SESSION_H_
 
+#include <tvm/ffi/extra/module.h>
 #include <tvm/ffi/function.h>
 #include <tvm/runtime/device_api.h>
-#include <tvm/runtime/module.h>
 #include <tvm/runtime/object.h>
 
 #include <functional>

src/runtime/static_library.cc

@@ -24,10 +24,10 @@
  */
 #include "./static_library.h"
 
+#include <tvm/ffi/extra/module.h>
 #include <tvm/ffi/function.h>
 #include <tvm/ffi/memory.h>
 #include <tvm/ffi/reflection/registry.h>
-#include <tvm/runtime/module.h>
 #include <tvm/support/io.h>
 #include <tvm/support/serializer.h>
 

src/runtime/static_library.h

@@ -27,8 +27,8 @@
 #define TVM_RUNTIME_STATIC_LIBRARY_H_
 
 #include <tvm/ffi/container/array.h>
+#include <tvm/ffi/extra/module.h>
 #include <tvm/runtime/logging.h>
-#include <tvm/runtime/module.h>
 
 #include <array>
 #include <memory>

src/runtime/vm/executable.cc

@@ -32,6 +32,7 @@
 
 #include "../../support/bytes_io.h"
 #include "../file_utils.h"
+#include "./module_utils.h"
 
 namespace tvm {
 namespace runtime {
@@ -45,6 +46,21 @@ constexpr uint64_t kTVMVMBytecodeMagicV2 = 0xD225DE2F4214151E;
   TVM_FFI_ICHECK(val) << "Invalid VM file format in the " << section << " section." \
                       << "\n";
 
+const char* VMExecutable::kind() const { return "relax.VMExecutable"; }
+
+ffi::Optional<ffi::Function> VMExecutable::GetFunction(const ffi::String& _name) {
+  using SelfPtr = std::remove_cv_t<decltype(this)>;
+  ::tvm::ffi::ObjectPtr<::tvm::ffi::Object> _self =
+      ::tvm::ffi::GetObjectPtr<::tvm::ffi::Object>(this);
+  TVM_MODULE_VTABLE_ENTRY("stats", &VMExecutable::Stats);
+  TVM_MODULE_VTABLE_ENTRY("as_text", &VMExecutable::AsText);
+  TVM_MODULE_VTABLE_ENTRY("as_python", &VMExecutable::AsPython);
+  TVM_MODULE_VTABLE_ENTRY("vm_load_executable", &VMExecutable::VMLoadExecutable);
+  TVM_MODULE_VTABLE_ENTRY("vm_profiler_load_executable", &VMExecutable::VMProfilerLoadExecutable);
+  TVM_MODULE_VTABLE_ENTRY("has_function", &VMExecutable::HasFunction);
+  return std::nullopt;
+}
+
 std::string VMExecutable::Stats() const {
   std::ostringstream oss;
   oss << "Relax VM executable statistics:" << std::endl;