This is an automated email from the ASF dual-hosted git repository. tqchen pushed a commit to branch ffi-module in repository https://gitbox.apache.org/repos/asf/tvm.git
commit fe1d6384e05d513a346b7182e2b1b6312ff2ef99 Author: tqchen <[email protected]> AuthorDate: Sun Aug 17 17:46:51 2025 -0400 [FFI][REFACTOR] Establish ffi.Module in python This PR refactors and establishes ffi.Module under the python tvm ffi api. Also moves export_library method to executable so it aligns more with compiled artifact. --- python/tvm/ffi/__init__.py | 5 + python/tvm/ffi/module.py | 257 ++++++++++++ python/tvm/relax/vm_build.py | 6 +- python/tvm/runtime/executable.py | 209 ++++++++-- python/tvm/runtime/module.py | 443 +-------------------- tests/python/runtime/test_runtime_module_export.py | 2 +- 6 files changed, 459 insertions(+), 463 deletions(-) diff --git a/python/tvm/ffi/__init__.py b/python/tvm/ffi/__init__.py index e615e22a0c..801a8d2989 100644 --- a/python/tvm/ffi/__init__.py +++ b/python/tvm/ffi/__init__.py @@ -30,6 +30,7 @@ from .ndarray import Device, device from .ndarray import cpu, cuda, rocm, opencl, metal, vpi, vulkan, ext_dev, hexagon, webgpu from .ndarray import from_dlpack, NDArray, Shape from .container import Array, Map +from .module import Module, ModulePropertyMask, system_lib, load_module from . import serialization from . import access_path from . import testing @@ -71,4 +72,8 @@ __all__ = [ "testing", "access_path", "serialization", + "Module", + "ModulePropertyMask", + "system_lib", + "load_module", ] diff --git a/python/tvm/ffi/module.py b/python/tvm/ffi/module.py new file mode 100644 index 0000000000..b7189594d4 --- /dev/null +++ b/python/tvm/ffi/module.py @@ -0,0 +1,257 @@ +# 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. +"""Module related objects and functions.""" + +from enum import IntEnum +from . import _ffi_api + +from . import core +from .registry import register_object + +__all__ = ["Module", "ModulePropertyMask", "system_lib", "load_module"] + + +class ModulePropertyMask(IntEnum): + """Runtime Module Property Mask.""" + + BINARY_SERIALIZABLE = 0b001 + RUNNABLE = 0b010 + COMPILATION_EXPORTABLE = 0b100 + + +@register_object("ffi.Module") +class Module(core.Object): + """Runtime Module.""" + + def __new__(cls): + instance = super(Module, cls).__new__(cls) # pylint: disable=no-value-for-parameter + instance.entry_name = "__tvm_ffi_main__" + instance._entry = None + return instance + + @property + def entry_func(self): + """Get the entry function + + Returns + ------- + f : tvm.ffi.Function + The entry function if exist + """ + if self._entry: + return self._entry + self._entry = self.get_function("__tvm_ffi_main__") + return self._entry + + @property + def kind(self): + """Get type key of the module.""" + return _ffi_api.ModuleGetKind(self) + + @property + def imports(self): + """Get imported modules + + Returns + ---------- + modules : list of Module + The module + """ + return self.imports_ + + def implements_function(self, name, query_imports=False): + """Returns True if the module has a definition for the global function with name. Note + that has_function(name) does not imply get_function(name) is non-null since the module + may be, eg, a CSourceModule which cannot supply a packed-func implementation of the function + without further compilation. However, get_function(name) non null should always imply + has_function(name). + + Parameters + ---------- + name : str + The name of the function + + query_imports : bool + Whether to also query modules imported by this module. + + Returns + ------- + b : Bool + True if module (or one of its imports) has a definition for name. + """ + return _ffi_api.ModuleImplementsFunction(self, name, query_imports) + + def get_function(self, name, query_imports=False): + """Get function from the module. + + Parameters + ---------- + name : str + The name of the function + + query_imports : bool + Whether also query modules imported by this module. + + Returns + ------- + f : tvm.ffi.Function + The result function. + """ + func = _ffi_api.ModuleGetFunction(self, name, query_imports) + if func is None: + raise AttributeError(f"Module has no function '{name}'") + return func + + def import_module(self, module): + """Add module to the import list of current one. + + Parameters + ---------- + module : tvm.runtime.Module + The other module. + """ + _ffi_api.ModuleImportModule(self, module) + + def __getitem__(self, name): + if not isinstance(name, str): + raise ValueError("Can only take string as function name") + return self.get_function(name) + + def __call__(self, *args): + if self._entry: + return self._entry(*args) + # pylint: disable=not-callable + return self.entry_func(*args) + + def inspect_source(self, fmt=""): + """Get source code from module, if available. + + Parameters + ---------- + fmt : str, optional + The specified format. + + Returns + ------- + source : str + The result source code. + """ + return _ffi_api.ModuleInspectSource(self, fmt) + + def get_write_formats(self): + """Get the format of the module.""" + return _ffi_api.ModuleGetWriteFormats(self) + + def get_property_mask(self): + """Get the runtime module property mask. The mapping is stated in ModulePropertyMask. + + Returns + ------- + mask : int + Bitmask of runtime module property + """ + return _ffi_api.ModuleGetPropertyMask(self) + + def is_binary_serializable(self): + """Module 'binary serializable', save_to_bytes is supported. + + Returns + ------- + b : Bool + True if the module is binary serializable. + """ + return (self.get_property_mask() & ModulePropertyMask.BINARY_SERIALIZABLE) != 0 + + def is_runnable(self): + """Module 'runnable', get_function is supported. + + Returns + ------- + b : Bool + True if the module is runnable. + """ + return (self.get_property_mask() & ModulePropertyMask.RUNNABLE) != 0 + + def is_compilation_exportable(self): + """Module 'compilation exportable', write_to_file is supported for object or source. + + Returns + ------- + b : Bool + True if the module is compilation exportable. + """ + return (self.get_property_mask() & ModulePropertyMask.COMPILATION_EXPORTABLE) != 0 + + def clear_imports(self): + """Remove all imports of the module.""" + _ffi_api.ModuleClearImports(self) + + def write_to_file(self, file_name, fmt=""): + """Write the current module to file. + + Parameters + ---------- + file_name : str + The name of the file. + fmt : str + The format of the file. + + See Also + -------- + runtime.Module.export_library : export the module to shared library. + """ + _ffi_api.ModuleWriteToFile(self, file_name, fmt) + + +def system_lib(symbol_prefix=""): + """Get system-wide library module singleton. + + System lib is a global module that contains self register functions in startup. + Unlike normal dso modules which need to be loaded explicitly. + It is useful in environments where dynamic loading api like dlopen is banned. + + The system lib is intended to be linked and loaded during the entire life-cyle of the program. + If you want dynamic loading features, use dso modules instead. + + Parameters + ---------- + symbol_prefix: Optional[str] + Optional symbol prefix that can be used for search. When we lookup a symbol + symbol_prefix + name will first be searched, then the name without symbol_prefix. + + Returns + ------- + module : runtime.Module + The system-wide library module. + """ + return _ffi_api.SystemLib(symbol_prefix) + + +def load_module(path): + """Load module from file. + + Parameters + ---------- + path : str + The path to the module file. + + Returns + ------- + module : ffi.Module + The loaded module + """ + return _ffi_api.ModuleLoadFromFile(path) diff --git a/python/tvm/relax/vm_build.py b/python/tvm/relax/vm_build.py index f6db61af61..cf8cd86330 100644 --- a/python/tvm/relax/vm_build.py +++ b/python/tvm/relax/vm_build.py @@ -99,6 +99,10 @@ def _auto_attach_system_lib_prefix( return tir_mod +def _is_device_module(mod: tvm.runtime.Module) -> bool: + return mod.kind in ["cuda", "opencl", "metal", "hip", "vulkan", "webgpu"] + + def _vmlink( builder: "relax.ExecBuilder", target: Optional[Union[str, tvm.target.Target]], @@ -153,7 +157,7 @@ def _vmlink( tir_mod = _auto_attach_system_lib_prefix(tir_mod, target, system_lib) lib = tvm.tir.build(tir_mod, target=target, pipeline=tir_pipeline) for ext_mod in ext_libs: - if ext_mod.is_device_module(): + if _is_device_module(ext_mod): tir_ext_libs.append(ext_mod) else: relax_ext_libs.append(ext_mod) diff --git a/python/tvm/runtime/executable.py b/python/tvm/runtime/executable.py index 51f0a772e4..2a95bc40ff 100644 --- a/python/tvm/runtime/executable.py +++ b/python/tvm/runtime/executable.py @@ -17,10 +17,15 @@ # pylint: disable=invalid-name, no-member """Executable object for TVM Runtime""" +import os from typing import Any, Callable, Dict, List, Optional, Union import tvm + +from tvm.base import _RUNTIME_ONLY +from tvm.libinfo import find_include_path from tvm.contrib import utils as _utils +from . import _ffi_api from . import PackedFunc, Module @@ -94,7 +99,7 @@ class Executable: return x.kind in ("c", "static_library") # pylint:disable = protected-access - not_runnable_list = self.mod._collect_from_import_tree(_not_runnable) + not_runnable_list = self._collect_from_import_tree(_not_runnable) # everything is runnable, directly return mod. if len(not_runnable_list) == 0: @@ -105,20 +110,63 @@ class Executable: # by collecting the link and allow export_library skip those modules. workspace_dir = _utils.tempdir() dso_path = workspace_dir.relpath("exported.so") - self.mod.export_library(dso_path, fcompile=fcompile, addons=addons, **kwargs) + self.export_library(dso_path, fcompile=fcompile, addons=addons, **kwargs) self._jitted_mod = tvm.runtime.load_module(dso_path) return self._jitted_mod + def _collect_from_import_tree(self, filter_func): + """Helper function to collect modules from the tree matching a filter_func, then return it. + + Parameters + ---------- + filter_func : Callable[[Module], bool] + A function which is invoked for each Module discovered in the import tree (including + self). + + Returns + ------- + list[Module] : + A list of matching Module. + """ + visited, stack, dso_modules = set(), [], [] + # append root module + visited.add(self.mod) + stack.append(self.mod) + while stack: + module = stack.pop() + assert ( + module.is_compilation_exportable() or module.is_binary_serializable() + ), f"Module {module.kind} should be either dso exportable or binary serializable." + + if filter_func(module): + dso_modules.append(module) + for m in module.imports: + if m not in visited: + visited.add(m) + stack.append(m) + return dso_modules + + def _collect_dso_modules(self): + return self._collect_from_import_tree(lambda m: m.is_compilation_exportable()) + def export_library( self, - file_name: str, + file_name, *, - fcompile: Optional[Union[str, Callable[[str, List[str], Dict[str, Any]], None]]] = None, - addons: Optional[List[str]] = None, - workspace_dir: Optional[str] = None, + fcompile=None, + fpack_imports=None, + addons=None, + workspace_dir=None, **kwargs, - ) -> Any: - """Export the executable to a library which can then be loaded back. + ): + """ + Export the module and all imported modules into a single device library. + + This function only works on host LLVM modules, other runtime::Module + subclasses will work with this API but they must support implement + the save and load mechanisms of modules completely including saving + from streams and files. This will pack your non-shared library module + into a single shared library which can later be loaded by TVM. Parameters ---------- @@ -127,9 +175,25 @@ class Executable: fcompile : function(target, file_list, kwargs), optional The compilation function to use create the final library object during + export. - addons : list of str, optional - Additional object files to link against. + For example, when fcompile=_cc.create_shared, or when it is not supplied but + module is "llvm," this is used to link all produced artifacts + into a final dynamic library. + + This behavior is controlled by the type of object exported. + If fcompile has attribute object_format, will compile host library + to that format. Otherwise, will use default format "o". + + fpack_imports: function(mod: runtime.Module, is_system_lib: bool, symbol_prefix: str, + workspace_dir: str) -> str + Function used to pack imported modules from `mod` into a file suitable for passing + to fcompile as an input file. The result can be a C source, or an .o object file, + or any other file that the fcompile function can handle. The function returns the + name of the created file. + + If not provided, the imported modules will be serialized either via packing to an + LLVM module, or to a C source file. workspace_dir : str, optional The path of the directory used to create the intermediate @@ -144,22 +208,111 @@ class Executable: result of fcompile() : unknown, optional If the compilation function returns an artifact it would be returned via export_library, if any. - - Examples - -------- - .. code:: python - - ex = tvm.compile(mod, target) - # export the library - ex.export_library("exported.so") - - # load it back for future uses. - rt_mod = tvm.runtime.load_module("exported.so") """ - return self.mod.export_library( - file_name=file_name, - fcompile=fcompile, - addons=addons, - workspace_dir=workspace_dir, - **kwargs, - ) + # NOTE: this function depends on contrib library features + # which are only available in when TVM function is available. + if _RUNTIME_ONLY: + raise RuntimeError("Cannot call export_library in runtime only mode") + + # Extra dependencies during runtime. + from pathlib import Path + from tvm.contrib import cc as _cc, tar as _tar, utils as _utils, tvmjs as _tvmjs + + if isinstance(file_name, Path): + file_name = str(file_name) + + modules = self._collect_dso_modules() + if workspace_dir is None: + temp = _utils.tempdir() + workspace_dir = temp.temp_dir + files = addons if addons else [] + is_system_lib = False + has_c_module = False + system_lib_prefix = None + llvm_target_string = None + global_object_format = "o" + + def get_source_format_from_module(module): + for fmt in module.get_write_formats(): + if fmt in ["c", "cc", "cpp", "cu"]: + return fmt + raise ValueError(f"Module {module.kind} does not exporting to c, cc, cpp or cu.") + + for index, module in enumerate(modules): + if fcompile is not None and hasattr(fcompile, "object_format"): + if module.kind == "c": + object_format = get_source_format_from_module(module) + has_c_module = True + else: + global_object_format = object_format = fcompile.object_format + else: + if module.kind == "c": + if len(module.get_write_formats()) > 0: + object_format = get_source_format_from_module(module) + else: + object_format = "c" + if "cc" in kwargs: + if kwargs["cc"] == "nvcc": + object_format = "cu" + has_c_module = True + else: + assert module.is_compilation_exportable() + global_object_format = object_format = "o" + + path_obj = os.path.join(workspace_dir, f"lib{index}.{object_format}") + module.write_to_file(path_obj) + files.append(path_obj) + if module.kind == "llvm": + is_system_lib = module.get_function("__tvm_is_system_module")() + llvm_target_string = module.get_function("_get_target_string")() + system_lib_prefix = module.get_function("__tvm_get_system_lib_prefix")() + + if not fcompile: + if file_name.endswith(".tar"): + fcompile = _tar.tar + elif file_name.endswith(".wasm"): + fcompile = _tvmjs.create_tvmjs_wasm + else: + fcompile = _cc.create_shared + + if llvm_target_string is None and hasattr(fcompile, "get_target_triple"): + triple = fcompile.get_target_triple() + assert triple, "Target triple should not be empty" + llvm_target_string = "llvm -mtriple " + triple + + if getattr(fcompile, "need_system_lib", False) and not is_system_lib: + raise ValueError(f"{str(fcompile)} need --system-lib option") + + if self.mod.imports: + pack_lib_prefix = system_lib_prefix if system_lib_prefix else "" + + if fpack_imports is not None: + path_out = fpack_imports(self.mod, is_system_lib, pack_lib_prefix, workspace_dir) + files.append(path_out) + elif _ffi_api.RuntimeEnabled("llvm") and llvm_target_string: + path_obj = os.path.join( + workspace_dir, f"{pack_lib_prefix}devc.{global_object_format}" + ) + m = _ffi_api.ModulePackImportsToLLVM( + self.mod, is_system_lib, llvm_target_string, pack_lib_prefix + ) + m.write_to_file(path_obj) + files.append(path_obj) + else: + path_cc = os.path.join(workspace_dir, f"{pack_lib_prefix}devc.c") + with open(path_cc, "w") as f: + f.write(_ffi_api.ModulePackImportsToC(self.mod, is_system_lib, pack_lib_prefix)) + files.append(path_cc) + + # The imports could contain a c module but the object format could be tar + # Thus, it would not recognize the following include paths as options + # which are there assuming a c compiler is the fcompile. + if has_c_module and not file_name.endswith(".tar"): + options = [] + if "options" in kwargs: + opts = kwargs["options"] + options = opts if isinstance(opts, (list, tuple)) else [opts] + opts = options + ["-I" + path for path in find_include_path()] + kwargs.update({"options": opts}) + + return fcompile(file_name, files, **kwargs) diff --git a/python/tvm/runtime/module.py b/python/tvm/runtime/module.py index 3925c24365..e445371cc3 100644 --- a/python/tvm/runtime/module.py +++ b/python/tvm/runtime/module.py @@ -18,17 +18,17 @@ # pylint: disable=invalid-name, unused-import, import-outside-toplevel, inconsistent-return-statements """Runtime Module namespace.""" import os -import ctypes import struct from typing import Sequence import numpy as np -import tvm.ffi -from tvm.base import _RUNTIME_ONLY -from tvm.libinfo import find_include_path - from . import _ffi_api -from ..ffi import _ffi_api as _mod_ffi_api +from ..ffi import ( + Module as _Module, + load_module as _load_module, + register_object as _register_object, +) +from ..ffi import system_lib class BenchmarkResult: @@ -90,203 +90,9 @@ class BenchmarkResult: ) -class ModulePropertyMask(object): - """Runtime Module Property Mask.""" - - BINARY_SERIALIZABLE = 0b001 - RUNNABLE = 0b010 - COMPILATION_EXPORTABLE = 0b100 - - [email protected]_object("ffi.Module") -class Module(tvm.ffi.Object): - """Runtime Module.""" - - def __new__(cls): - instance = super(Module, cls).__new__(cls) # pylint: disable=no-value-for-parameter - instance.entry_name = "__tvm_ffi_main__" - instance._entry = None - return instance - - @property - def entry_func(self): - """Get the entry function - - Returns - ------- - f : tvm.runtime.PackedFunc - The entry function if exist - """ - if self._entry: - return self._entry - self._entry = self.get_function("__tvm_ffi_main__") - return self._entry - - @property - def kind(self): - """Get type key of the module.""" - return _mod_ffi_api.ModuleGetKind(self) - - @property - def imports(self): - """Get imported modules - - Returns - ---------- - modules : list of Module - The module - """ - return self.imports_ - - def implements_function(self, name, query_imports=False): - """Returns True if the module has a definition for the global function with name. Note - that has_function(name) does not imply get_function(name) is non-null since the module - may be, eg, a CSourceModule which cannot supply a packed-func implementation of the function - without further compilation. However, get_function(name) non null should always imply - has_function(name). - - Parameters - ---------- - name : str - The name of the function - - query_imports : bool - Whether to also query modules imported by this module. - - Returns - ------- - b : Bool - True if module (or one of its imports) has a definition for name. - """ - return _mod_ffi_api.ModuleImplementsFunction(self, name, query_imports) - - def get_function(self, name, query_imports=False): - """Get function from the module. - - Parameters - ---------- - name : str - The name of the function - - query_imports : bool - Whether also query modules imported by this module. - - Returns - ------- - f : tvm.runtime.PackedFunc - The result function. - """ - func = _mod_ffi_api.ModuleGetFunction(self, name, query_imports) - if func is None: - raise AttributeError(f"Module has no function '{name}'") - return func - - def import_module(self, module): - """Add module to the import list of current one. - - Parameters - ---------- - module : tvm.runtime.Module - The other module. - """ - _mod_ffi_api.ModuleImportModule(self, module) - - def __getitem__(self, name): - if not isinstance(name, str): - raise ValueError("Can only take string as function name") - return self.get_function(name) - - def __call__(self, *args): - if self._entry: - return self._entry(*args) - # pylint: disable=not-callable - return self.entry_func(*args) - - def inspect_source(self, fmt=""): - """Get source code from module, if available. - - Parameters - ---------- - fmt : str, optional - The specified format. - - Returns - ------- - source : str - The result source code. - """ - return _mod_ffi_api.ModuleInspectSource(self, fmt) - - def get_write_formats(self): - """Get the format of the module.""" - return _mod_ffi_api.ModuleGetWriteFormats(self) - - def get_property_mask(self): - """Get the runtime module property mask. The mapping is stated in ModulePropertyMask. - - Returns - ------- - mask : int - Bitmask of runtime module property - """ - return _mod_ffi_api.ModuleGetPropertyMask(self) - - def is_binary_serializable(self): - """Returns true if module is 'binary serializable', ie can be serialzed into binary - stream and loaded back to the runtime module. - - Returns - ------- - b : Bool - True if the module is binary serializable. - """ - return (self.get_property_mask() & ModulePropertyMask.BINARY_SERIALIZABLE) != 0 - - def is_runnable(self): - """Returns true if module is 'runnable'. ie can be executed without any extra - compilation/linking steps. - - Returns - ------- - b : Bool - True if the module is runnable. - """ - return (self.get_property_mask() & ModulePropertyMask.RUNNABLE) != 0 - - def is_device_module(self): - return self.kind in ["cuda", "opencl", "metal", "hip", "vulkan", "webgpu"] - - def is_compilation_exportable(self): - """Returns true if module is 'compilation exportable', ie can be included in result of - export_library by the external compiler directly. - - Returns - ------- - b : Bool - True if the module is compilation exportable. - """ - return (self.get_property_mask() & ModulePropertyMask.COMPILATION_EXPORTABLE) != 0 - - def clear_imports(self): - """Remove all imports of the module.""" - _mod_ffi_api.ModuleClearImports(self) - - def write_to_file(self, file_name, fmt=""): - """Write the current module to file. - - Parameters - ---------- - file_name : str - The name of the file. - fmt : str - The format of the file. - - See Also - -------- - runtime.Module.export_library : export the module to shared library. - """ - _mod_ffi_api.ModuleWriteToFile(self, file_name, fmt) - +# override the Module class in ffi.Module +@_register_object("ffi.Module") +class Module(_Module): def time_evaluator( self, func_name, @@ -385,235 +191,6 @@ class Module(tvm.ffi.Object): except NameError: raise NameError("time_evaluator is only supported when RPC is enabled") - def _collect_from_import_tree(self, filter_func): - """Helper function to collect modules from the tree matching a filter_func, then return it. - - Parameters - ---------- - filter_func : Callable[[Module], bool] - A function which is invoked for each Module discovered in the import tree (including - self). - - Returns - ------- - list[Module] : - A list of matching Module. - """ - visited, stack, dso_modules = set(), [], [] - # append root module - visited.add(self) - stack.append(self) - while stack: - module = stack.pop() - assert ( - module.is_compilation_exportable() or module.is_binary_serializable() - ), f"Module {module.kind} should be either dso exportable or binary serializable." - - if filter_func(module): - dso_modules.append(module) - for m in module.imports: - if m not in visited: - visited.add(m) - stack.append(m) - return dso_modules - - def _collect_dso_modules(self): - return self._collect_from_import_tree(lambda m: m.is_compilation_exportable()) - - def export_library( - self, - file_name, - *, - fcompile=None, - fpack_imports=None, - addons=None, - workspace_dir=None, - **kwargs, - ): - """ - Export the module and all imported modules into a single device library. - - This function only works on host LLVM modules, other runtime::Module - subclasses will work with this API but they must support implement - the save and load mechanisms of modules completely including saving - from streams and files. This will pack your non-shared library module - into a single shared library which can later be loaded by TVM. - - Parameters - ---------- - file_name : str - The name of the shared library. - - fcompile : function(target, file_list, kwargs), optional - The compilation function to use create the final library object during - export. - - For example, when fcompile=_cc.create_shared, or when it is not supplied but - module is "llvm," this is used to link all produced artifacts - into a final dynamic library. - - This behavior is controlled by the type of object exported. - If fcompile has attribute object_format, will compile host library - to that format. Otherwise, will use default format "o". - - fpack_imports: function(mod: runtime.Module, is_system_lib: bool, symbol_prefix: str, - workspace_dir: str) -> str - Function used to pack imported modules from `mod` into a file suitable for passing - to fcompile as an input file. The result can be a C source, or an .o object file, - or any other file that the fcompile function can handle. The function returns the - name of the created file. - - If not provided, the imported modules will be serialized either via packing to an - LLVM module, or to a C source file. - - workspace_dir : str, optional - The path of the directory used to create the intermediate - artifacts when exporting the module. - If this is not provided a temporary dir will be created. - - kwargs : dict, optional - Additional arguments passed to fcompile - - Returns - ------- - result of fcompile() : unknown, optional - If the compilation function returns an artifact it would be returned via - export_library, if any. - """ - # NOTE: this function depends on contrib library features - # which are only available in when TVM function is available. - if _RUNTIME_ONLY: - raise RuntimeError("Cannot call export_library in runtime only mode") - # Extra dependencies during runtime. - from pathlib import Path - from tvm.contrib import cc as _cc, tar as _tar, utils as _utils, tvmjs as _tvmjs - - if isinstance(file_name, Path): - file_name = str(file_name) - - modules = self._collect_dso_modules() - if workspace_dir is None: - temp = _utils.tempdir() - workspace_dir = temp.temp_dir - files = addons if addons else [] - is_system_lib = False - has_c_module = False - system_lib_prefix = None - llvm_target_string = None - global_object_format = "o" - - def get_source_format_from_module(module): - for fmt in module.get_write_formats(): - if fmt in ["c", "cc", "cpp", "cu"]: - return fmt - raise ValueError(f"Module {module.kind} does not exporting to c, cc, cpp or cu.") - - for index, module in enumerate(modules): - if fcompile is not None and hasattr(fcompile, "object_format"): - if module.kind == "c": - object_format = get_source_format_from_module(module) - has_c_module = True - else: - global_object_format = object_format = fcompile.object_format - else: - if module.kind == "c": - if len(module.get_write_formats()) > 0: - object_format = get_source_format_from_module(module) - else: - object_format = "c" - if "cc" in kwargs: - if kwargs["cc"] == "nvcc": - object_format = "cu" - has_c_module = True - else: - assert module.is_compilation_exportable() - global_object_format = object_format = "o" - - path_obj = os.path.join(workspace_dir, f"lib{index}.{object_format}") - module.write_to_file(path_obj) - files.append(path_obj) - if module.kind == "llvm": - is_system_lib = module.get_function("__tvm_is_system_module")() - llvm_target_string = module.get_function("_get_target_string")() - system_lib_prefix = module.get_function("__tvm_get_system_lib_prefix")() - - if not fcompile: - if file_name.endswith(".tar"): - fcompile = _tar.tar - elif file_name.endswith(".wasm"): - fcompile = _tvmjs.create_tvmjs_wasm - else: - fcompile = _cc.create_shared - - if llvm_target_string is None and hasattr(fcompile, "get_target_triple"): - triple = fcompile.get_target_triple() - assert triple, "Target triple should not be empty" - llvm_target_string = "llvm -mtriple " + triple - - if getattr(fcompile, "need_system_lib", False) and not is_system_lib: - raise ValueError(f"{str(fcompile)} need --system-lib option") - - if self.imports: - pack_lib_prefix = system_lib_prefix if system_lib_prefix else "" - - if fpack_imports is not None: - path_out = fpack_imports(self, is_system_lib, pack_lib_prefix, workspace_dir) - files.append(path_out) - elif enabled("llvm") and llvm_target_string: - path_obj = os.path.join( - workspace_dir, f"{pack_lib_prefix}devc.{global_object_format}" - ) - m = _ffi_api.ModulePackImportsToLLVM( - self, is_system_lib, llvm_target_string, pack_lib_prefix - ) - m.write_to_file(path_obj) - files.append(path_obj) - else: - path_cc = os.path.join(workspace_dir, f"{pack_lib_prefix}devc.c") - with open(path_cc, "w") as f: - f.write(_ffi_api.ModulePackImportsToC(self, is_system_lib, pack_lib_prefix)) - files.append(path_cc) - - # The imports could contain a c module but the object format could be tar - # Thus, it would not recognize the following include paths as options - # which are there assuming a c compiler is the fcompile. - if has_c_module and not file_name.endswith(".tar"): - options = [] - if "options" in kwargs: - opts = kwargs["options"] - options = opts if isinstance(opts, (list, tuple)) else [opts] - opts = options + ["-I" + path for path in find_include_path()] - kwargs.update({"options": opts}) - - return fcompile(file_name, files, **kwargs) - - -def system_lib(symbol_prefix=""): - """Get system-wide library module singleton. - - System lib is a global module that contains self register functions in startup. - Unlike normal dso modules which need to be loaded explicitly. - It is useful in environments where dynamic loading api like dlopen is banned. - - To build system lib function, simply specify target option ```llvm --system-lib``` - The system lib will be available as long as the result code is linked by the program. - - The system lib is intended to be linked and loaded during the entire life-cyle of the program. - If you want dynamic loading features, use dso modules instead. - - Parameters - ---------- - symbol_prefix: Optional[str] - Optional symbol prefix that can be used for search. When we lookup a symbol - symbol_prefix + name will first be searched, then the name without symbol_prefix. - - Returns - ------- - module : runtime.Module - The system-wide library module. - """ - return _mod_ffi_api.SystemLib(symbol_prefix) - def load_module(path): """Load module from file. @@ -656,7 +233,7 @@ def load_module(path): _cc.create_shared(path + ".so", files) path += ".so" # Redirect to the load API - return _mod_ffi_api.ModuleLoadFromFile(path) + return _load_module(path) def load_static_library(path, func_names): diff --git a/tests/python/runtime/test_runtime_module_export.py b/tests/python/runtime/test_runtime_module_export.py index 0db1fa93dc..168d839dbc 100644 --- a/tests/python/runtime/test_runtime_module_export.py +++ b/tests/python/runtime/test_runtime_module_export.py @@ -54,7 +54,7 @@ def test_import_static_library(): # Import mod1 as a static library into mod0 and compile to its own DSO. mod0.import_module(mod1_o) mod0_dso_path = temp.relpath("mod0.so") - mod0.export_library(mod0_dso_path) + tvm.runtime.Executable(mod0).export_library(mod0_dso_path) # The imported mod1 is statically linked into mod0. loaded_lib = tvm.runtime.load_module(mod0_dso_path)
