From deb878107a2e4a3ceae954e888a09b2e512b5548 Mon Sep 17 00:00:00 2001 From: tqchen Date: Sun, 26 Apr 2026 20:00:55 +0000 Subject: [PATCH 1/2] [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 to ffi::Map; 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> directly instead of std::unordered_map>, and passes the existing ffi::Map 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 13dc02fde4a1..bf29556768c5 100644 --- a/src/relax/backend/vm/codegen_vm.cc +++ b/src/relax/backend/vm/codegen_vm.cc @@ -465,30 +465,20 @@ void LinkModules(ObjectPtr exec, const ffi::Map& ext_libs) { // query if we need const loader for ext_modules // Wrap all submodules in the initialization wrapper. - std::unordered_map> const_vars_by_symbol; + ffi::Map> 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 symbol_const_vars; if (pf_sym.has_value() && pf_var.has_value()) { ffi::String symbol = (*pf_sym)().cast(); ffi::Array variables = (*pf_var)().cast>(); - 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 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 c97232016d8a..2e2ad97f7662 100644 --- a/src/runtime/const_loader_module.h +++ b/src/runtime/const_loader_module.h @@ -32,10 +32,6 @@ #include #include -#include -#include -#include - 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& const_var_tensor, - const std::unordered_map>& const_vars_by_symbol) { + ffi::Map const_var_tensor, + ffi::Map> 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_const_var_tensor; - for (const auto& kv : const_var_tensor) { - ffi_const_var_tensor.Set(kv.first, kv.second); - } - ffi::Map> ffi_const_vars_by_symbol; - for (const auto& kv : const_vars_by_symbol) { - ffi::Array 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(); + return (*fcreate)(const_var_tensor, const_vars_by_symbol).cast(); } } // namespace runtime diff --git a/src/runtime/metal/metal_module.h b/src/runtime/metal/metal_module.h index fe9454f674d1..3f4b3965adc5 100644 --- a/src/runtime/metal/metal_module.h +++ b/src/runtime/metal/metal_module.h @@ -28,11 +28,6 @@ #include #include -#include -#include -#include -#include - #include "../metadata.h" namespace tvm { diff --git a/src/runtime/opencl/opencl_module.h b/src/runtime/opencl/opencl_module.h index 6697badd4885..ff8e16c3a172 100644 --- a/src/runtime/opencl/opencl_module.h +++ b/src/runtime/opencl/opencl_module.h @@ -28,10 +28,8 @@ #include #include -#include #include #include -#include #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& shaders, - const std::string& spirv_text, ffi::Map fmap) { + const std::unordered_map& shaders, ffi::String spirv_text, + ffi::Map 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(); + return (*fcreate)(shader_bytes, spirv_text, fmap).cast(); } } // namespace runtime } // namespace tvm diff --git a/src/runtime/vulkan/vulkan_module.h b/src/runtime/vulkan/vulkan_module.h index 87df473753d4..fb3bff73dc96 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 smap, ffi::Map 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(); + return (*fcreate)(shader_bytes, fmap, source).cast(); } } // namespace vulkan From ed11eb15db55fe54f4d2b4aa91b7bfc852574775 Mon Sep 17 00:00:00 2001 From: tqchen Date: Sun, 26 Apr 2026 20:09:09 +0000 Subject: [PATCH 2/2] [REFACTOR] Address Gemini review on #19449 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Four review points: 1. (high) Add std::string / std::unordered_map / std::vector includes to const_loader_module.cc — they were transitively pulled in by the old const_loader_module.h before it dropped them, and the .cc relies on them directly. 2. (medium) Add explicit `#include ` to const_loader_module.h so the header is self-contained. 3. (medium) Pass `ffi::Map` / `ffi::Array` / `ffi::String` parameters by `const &` in the inline wrappers — these are `ObjectRef` handles, pass-by-value triggers an atomic refcount inc/dec. Changed: const_loader_module.h (both params), opencl_module.h (spirv_text + fmap on the SPIRV overload), vulkan_module.h (smap, fmap, source). 4. (medium) Same const-ref consistency applied to the SPIRVShader smap (still std::unordered_map since SPIRVShader is not FFI-storable). Local verify: cpptest 128 / 128 pass; build clean with HIDE_PRIVATE_SYMBOLS=ON. --- src/runtime/const_loader_module.cc | 3 +++ src/runtime/const_loader_module.h | 5 +++-- src/runtime/opencl/opencl_module.h | 4 ++-- src/runtime/vulkan/vulkan_module.h | 6 +++--- 4 files changed, 11 insertions(+), 7 deletions(-) diff --git a/src/runtime/const_loader_module.cc b/src/runtime/const_loader_module.cc index 006c1f1e1acd..aaaeb9737e2d 100644 --- a/src/runtime/const_loader_module.cc +++ b/src/runtime/const_loader_module.cc @@ -39,6 +39,9 @@ #include #include +#include +#include +#include #include "../support/bytes_io.h" diff --git a/src/runtime/const_loader_module.h b/src/runtime/const_loader_module.h index 2e2ad97f7662..6722785cc950 100644 --- a/src/runtime/const_loader_module.h +++ b/src/runtime/const_loader_module.h @@ -29,6 +29,7 @@ #include #include #include +#include #include #include @@ -48,8 +49,8 @@ namespace runtime { * The creator is always available (ConstLoaderModule is a runtime-universal module). */ inline ffi::Module ConstLoaderModuleCreate( - ffi::Map const_var_tensor, - ffi::Map> const_vars_by_symbol) { + const ffi::Map& const_var_tensor, + const ffi::Map>& 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. " diff --git a/src/runtime/opencl/opencl_module.h b/src/runtime/opencl/opencl_module.h index ff8e16c3a172..9d16ea9231b2 100644 --- a/src/runtime/opencl/opencl_module.h +++ b/src/runtime/opencl/opencl_module.h @@ -71,8 +71,8 @@ inline ffi::Module OpenCLModuleCreate(ffi::String data, ffi::String fmt, * registered the creator. */ inline ffi::Module OpenCLModuleCreate( - const std::unordered_map& shaders, ffi::String spirv_text, - ffi::Map fmap) { + const std::unordered_map& shaders, + const ffi::String& spirv_text, const ffi::Map& 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. " diff --git a/src/runtime/vulkan/vulkan_module.h b/src/runtime/vulkan/vulkan_module.h index fb3bff73dc96..d8fdda4d9251 100644 --- a/src/runtime/vulkan/vulkan_module.h +++ b/src/runtime/vulkan/vulkan_module.h @@ -48,9 +48,9 @@ namespace vulkan { * and rehydrated on the runtime side. * Requires libtvm_runtime built with USE_VULKAN=ON to have registered the creator. */ -inline ffi::Module VulkanModuleCreate(std::unordered_map smap, - ffi::Map fmap, - ffi::String source) { +inline ffi::Module VulkanModuleCreate(const std::unordered_map& smap, + const ffi::Map& fmap, + const 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. "