[REFACTOR] Use FFI types in runtime inline module-create wrapper signatures#19449
Conversation
…atures This PR cleans up the public inline wrapper API for runtime backend module creators introduced in apache#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 apache#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.
There was a problem hiding this comment.
Code Review
This pull request refactors the backend and runtime modules to use TVM FFI container types (ffi::Map, ffi::Array, ffi::String) instead of standard C++ containers, simplifying the code by removing manual conversion logic. Review feedback warns that removing standard library headers from the const loader module may cause compilation failures in the corresponding source file. Additionally, suggestions were made to explicitly include tvm/ffi/string.h and to pass FFI containers by const reference to avoid unnecessary atomic reference counting overhead.
I am having trouble creating individual review comments. Click here to see my feedback.
src/runtime/const_loader_module.h (35-37)
Removing these standard library includes from the header will cause compilation failures in src/runtime/const_loader_module.cc. That file relies heavily on std::string, std::unordered_map, and std::vector (e.g., in the ConstLoaderModuleObj definition and various serialization methods) but does not include them itself. Please add these includes to src/runtime/const_loader_module.cc before removing them from this header.
src/runtime/const_loader_module.h (32)
This header uses ffi::String in the ConstLoaderModuleCreate signature but does not explicitly include <tvm/ffi/string.h>. While it may be pulled in transitively by other FFI headers, it is best practice to include it explicitly to ensure the header is self-contained.
#include <tvm/ffi/string.h>
#include <tvm/runtime/base.h>
src/runtime/const_loader_module.h (51-52)
FFI container types like ffi::Map and ffi::Array are ObjectRef handles. Passing them by value causes an atomic reference count increment and decrement. For input parameters that are not moved from, it is more efficient and consistent with other parts of the codebase (like LinkModules in codegen_vm.cc) to pass them by const &.
const ffi::Map<ffi::String, Tensor>& const_var_tensor,
const ffi::Map<ffi::String, ffi::Array<ffi::String>>& const_vars_by_symbol) {
src/runtime/opencl/opencl_module.h (74-75)
To avoid unnecessary atomic reference count increments, ffi::String and ffi::Map should be passed by const & when used as input parameters.
const std::unordered_map<std::string, spirv::SPIRVShader>& shaders, const ffi::String& spirv_text,
const ffi::Map<ffi::String, FunctionInfo>& fmap) {
src/runtime/vulkan/vulkan_module.h (53)
To avoid unnecessary atomic reference count increments, ffi::String should be passed by const & when used as an input parameter.
const ffi::String& source) {
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 <tvm/ffi/string.h>` 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.
Summary
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::stringandffi::*types with conversion glue inside the wrapper body.Changes
ConstLoaderModuleCreate: both parameters change fromstd::unordered_map<std::string, T>toffi::Map<ffi::String, T>; the conversion loops in the wrapper body are removed (net −12 lines of glue code). The callsite incodegen_vm.ccis updated to buildffi::Mapdirectly and passes the existingffi::Map<ffi::String, Tensor>params argument through without a copy.VulkanModuleCreate:sourceparameter changes fromstd::stringtoffi::String. TheSPIRVShadersmap remainsstd::unordered_mapbecauseSPIRVShaderis a plain C++ struct, not FFI-storable.OpenCLModuleCreate(SPIRV overload):spirv_textchanges fromconst std::string&toffi::String. SameSPIRVShaderconstraint applies.metal_module.h: removes stale<unordered_map>,<vector>,<memory>,<string>includes left over from before [REFACTOR] Isolate backend module creation via ffi.Module.create.<kind> registry #19447.Internal
*ModuleCreateImplfunctions and the compiler-sideSourceModuleCreatefamily are unchanged.