This is an automated email from the ASF dual-hosted git repository. tqchen pushed a commit to branch ffi-compat-module-create-signatures in repository https://gitbox.apache.org/repos/asf/tvm.git
commit deb878107a2e4a3ceae954e888a09b2e512b5548 Author: tqchen <[email protected]> AuthorDate: Sun Apr 26 20:00:55 2026 +0000 [REFACTOR] Use FFI types in runtime inline module-create wrapper signatures This PR cleans up the public inline wrapper API for runtime backend module creators introduced in #19447. The wrappers previously used a mix of std::unordered_map / std::vector / std::string and ffi::* types, with conversion glue inside the wrapper body. Changes: - ConstLoaderModuleCreate: both parameters change from std::unordered_map<std::string, T> to ffi::Map<ffi::String, T>; the conversion loops in the wrapper body are removed (12 lines deleted). - VulkanModuleCreate: source parameter changes from std::string to ffi::String; SPIRVShader smap stays std::unordered_map because SPIRVShader is not an FFI-storable type. - OpenCLModuleCreate (SPIRV overload): spirv_text parameter changes from const std::string& to ffi::String; same SPIRVShader constraint applies. - metal_module.h: remove stale std-library includes (unordered_map, vector, memory, string) that were no longer needed after #19447. - codegen_vm.cc (LinkModules): callsite updated to build ffi::Map<ffi::String, ffi::Array<ffi::String>> directly instead of std::unordered_map<std::string, std::vector<std::string>>, and passes the existing ffi::Map<ffi::String, Tensor> params directly to ConstLoaderModuleCreate without copying. Internal CreateImpl functions and the compiler-side SourceModuleCreate family are unchanged. --- src/relax/backend/vm/codegen_vm.cc | 18 ++++-------------- src/runtime/const_loader_module.h | 23 +++-------------------- src/runtime/metal/metal_module.h | 5 ----- src/runtime/opencl/opencl_module.h | 8 +++----- src/runtime/vulkan/vulkan_module.h | 4 ++-- 5 files changed, 12 insertions(+), 46 deletions(-) diff --git a/src/relax/backend/vm/codegen_vm.cc b/src/relax/backend/vm/codegen_vm.cc index 13dc02fde4..bf29556768 100644 --- a/src/relax/backend/vm/codegen_vm.cc +++ b/src/relax/backend/vm/codegen_vm.cc @@ -465,30 +465,20 @@ void LinkModules(ObjectPtr<VMExecutable> exec, const ffi::Map<ffi::String, runti const tvm::ffi::Module& lib, const ffi::Array<ffi::Module>& ext_libs) { // query if we need const loader for ext_modules // Wrap all submodules in the initialization wrapper. - std::unordered_map<std::string, std::vector<std::string>> const_vars_by_symbol; + ffi::Map<ffi::String, ffi::Array<ffi::String>> const_vars_by_symbol; for (tvm::ffi::Module mod : ext_libs) { auto pf_sym = mod->GetFunction("get_symbol"); auto pf_var = mod->GetFunction("get_const_vars"); - std::vector<std::string> symbol_const_vars; if (pf_sym.has_value() && pf_var.has_value()) { ffi::String symbol = (*pf_sym)().cast<ffi::String>(); ffi::Array<ffi::String> variables = (*pf_var)().cast<ffi::Array<ffi::String>>(); - for (size_t i = 0; i < variables.size(); i++) { - symbol_const_vars.push_back(variables[i].operator std::string()); - } - TVM_FFI_ICHECK_EQ(const_vars_by_symbol.count(symbol), 0U) - << "Found duplicated symbol: " << symbol; - const_vars_by_symbol[symbol] = symbol_const_vars; + TVM_FFI_ICHECK(!const_vars_by_symbol.count(symbol)) << "Found duplicated symbol: " << symbol; + const_vars_by_symbol.Set(symbol, variables); } } if (!const_vars_by_symbol.empty() || !params.empty()) { // need runtime const information, run link const loader - std::unordered_map<std::string, runtime::Tensor> const_var_tensor; - for (const auto& [name, param] : params) { - const_var_tensor[name] = param; - } - ffi::Module const_loader_mod = - runtime::ConstLoaderModuleCreate(const_var_tensor, const_vars_by_symbol); + ffi::Module const_loader_mod = runtime::ConstLoaderModuleCreate(params, const_vars_by_symbol); const_loader_mod->ImportModule(lib); for (const auto& it : ext_libs) { const_loader_mod->ImportModule(it); diff --git a/src/runtime/const_loader_module.h b/src/runtime/const_loader_module.h index c97232016d..2e2ad97f76 100644 --- a/src/runtime/const_loader_module.h +++ b/src/runtime/const_loader_module.h @@ -32,10 +32,6 @@ #include <tvm/runtime/base.h> #include <tvm/runtime/tensor.h> -#include <string> -#include <unordered_map> -#include <vector> - namespace tvm { namespace runtime { @@ -52,26 +48,13 @@ namespace runtime { * The creator is always available (ConstLoaderModule is a runtime-universal module). */ inline ffi::Module ConstLoaderModuleCreate( - const std::unordered_map<std::string, Tensor>& const_var_tensor, - const std::unordered_map<std::string, std::vector<std::string>>& const_vars_by_symbol) { + ffi::Map<ffi::String, Tensor> const_var_tensor, + ffi::Map<ffi::String, ffi::Array<ffi::String>> const_vars_by_symbol) { static const auto fcreate = ffi::Function::GetGlobal("ffi.Module.create.const_loader"); TVM_FFI_CHECK(fcreate.has_value(), RuntimeError) << "ffi.Module.create.const_loader is not registered in runtime. " << "Ensure libtvm_runtime is loaded."; - // Convert to FFI-compatible types. - ffi::Map<ffi::String, Tensor> ffi_const_var_tensor; - for (const auto& kv : const_var_tensor) { - ffi_const_var_tensor.Set(kv.first, kv.second); - } - ffi::Map<ffi::String, ffi::Array<ffi::String>> ffi_const_vars_by_symbol; - for (const auto& kv : const_vars_by_symbol) { - ffi::Array<ffi::String> vars; - for (const auto& v : kv.second) { - vars.push_back(ffi::String(v)); - } - ffi_const_vars_by_symbol.Set(kv.first, vars); - } - return (*fcreate)(ffi_const_var_tensor, ffi_const_vars_by_symbol).cast<ffi::Module>(); + return (*fcreate)(const_var_tensor, const_vars_by_symbol).cast<ffi::Module>(); } } // namespace runtime diff --git a/src/runtime/metal/metal_module.h b/src/runtime/metal/metal_module.h index fe9454f674..3f4b3965ad 100644 --- a/src/runtime/metal/metal_module.h +++ b/src/runtime/metal/metal_module.h @@ -28,11 +28,6 @@ #include <tvm/ffi/extra/module.h> #include <tvm/ffi/function.h> -#include <memory> -#include <string> -#include <unordered_map> -#include <vector> - #include "../metadata.h" namespace tvm { diff --git a/src/runtime/opencl/opencl_module.h b/src/runtime/opencl/opencl_module.h index 6697badd48..ff8e16c3a1 100644 --- a/src/runtime/opencl/opencl_module.h +++ b/src/runtime/opencl/opencl_module.h @@ -28,10 +28,8 @@ #include <tvm/ffi/function.h> #include <tvm/runtime/base.h> -#include <memory> #include <string> #include <unordered_map> -#include <vector> #include "../../support/bytes_io.h" #include "../metadata.h" @@ -73,8 +71,8 @@ inline ffi::Module OpenCLModuleCreate(ffi::String data, ffi::String fmt, * registered the creator. */ inline ffi::Module OpenCLModuleCreate( - const std::unordered_map<std::string, spirv::SPIRVShader>& shaders, - const std::string& spirv_text, ffi::Map<ffi::String, FunctionInfo> fmap) { + const std::unordered_map<std::string, spirv::SPIRVShader>& shaders, ffi::String spirv_text, + ffi::Map<ffi::String, FunctionInfo> fmap) { static const auto fcreate = ffi::Function::GetGlobal("ffi.Module.create.opencl.spirv"); TVM_FFI_CHECK(fcreate.has_value(), RuntimeError) << "ffi.Module.create.opencl.spirv is not registered in runtime. " @@ -87,7 +85,7 @@ inline ffi::Module OpenCLModuleCreate( strm.Write(kv.second); shader_bytes.Set(kv.first, ffi::Bytes(std::move(buf))); } - return (*fcreate)(shader_bytes, ffi::String(spirv_text), fmap).cast<ffi::Module>(); + return (*fcreate)(shader_bytes, spirv_text, fmap).cast<ffi::Module>(); } } // namespace runtime } // namespace tvm diff --git a/src/runtime/vulkan/vulkan_module.h b/src/runtime/vulkan/vulkan_module.h index 87df473753..fb3bff73dc 100644 --- a/src/runtime/vulkan/vulkan_module.h +++ b/src/runtime/vulkan/vulkan_module.h @@ -50,7 +50,7 @@ namespace vulkan { */ inline ffi::Module VulkanModuleCreate(std::unordered_map<std::string, SPIRVShader> smap, ffi::Map<ffi::String, FunctionInfo> fmap, - std::string source) { + ffi::String source) { static const auto fcreate = ffi::Function::GetGlobal("ffi.Module.create.vulkan"); TVM_FFI_CHECK(fcreate.has_value(), RuntimeError) << "ffi.Module.create.vulkan is not registered in runtime. " @@ -63,7 +63,7 @@ inline ffi::Module VulkanModuleCreate(std::unordered_map<std::string, SPIRVShade strm.Write(kv.second); shader_bytes.Set(kv.first, ffi::Bytes(std::move(buf))); } - return (*fcreate)(shader_bytes, fmap, ffi::String(source)).cast<ffi::Module>(); + return (*fcreate)(shader_bytes, fmap, source).cast<ffi::Module>(); } } // namespace vulkan
