diff --git a/python/tvm/ffi/__init__.py b/python/tvm/ffi/__init__.py index e615e22a0cbc..801a8d298906 100644 --- a/python/tvm/ffi/__init__.py +++ b/python/tvm/ffi/__init__.py @@ -30,6 +30,7 @@ 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 @@ "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 000000000000..0895b317c1d4 --- /dev/null +++ b/python/tvm/ffi/module.py @@ -0,0 +1,258 @@ +# 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.""" +# pylint: disable=invalid-name + +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 f6db61af61d2..cf8cd863309a 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 51f0a772e403..47c46959be28 100644 --- a/python/tvm/runtime/executable.py +++ b/python/tvm/runtime/executable.py @@ -17,9 +17,10 @@ # pylint: disable=invalid-name, no-member """Executable object for TVM Runtime""" -from typing import Any, Callable, Dict, List, Optional, Union +from typing import Any, Callable, Dict, List, Optional import tvm + from tvm.contrib import utils as _utils from . import PackedFunc, Module @@ -105,20 +106,27 @@ def _not_runnable(x): # 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 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, + 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,6 +135,15 @@ def export_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". addons : list of str, optional Additional object files to link against. @@ -144,20 +161,9 @@ def export_library( 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, + file_name, fcompile=fcompile, addons=addons, workspace_dir=workspace_dir, diff --git a/python/tvm/runtime/module.py b/python/tvm/runtime/module.py index 3925c24365d5..9cbc06708bd0 100644 --- a/python/tvm/runtime/module.py +++ b/python/tvm/runtime/module.py @@ -18,17 +18,20 @@ # 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 +import numpy as np 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, + system_lib, +) class BenchmarkResult: @@ -90,301 +93,11 @@ def __str__(self): ) -class ModulePropertyMask(object): - """Runtime Module Property Mask.""" - - BINARY_SERIALIZABLE = 0b001 - RUNNABLE = 0b010 - COMPILATION_EXPORTABLE = 0b100 - - -@tvm.ffi.register_object("ffi.Module") -class Module(tvm.ffi.Object): +# override the Module class in ffi.Module +@_register_object("ffi.Module") +class Module(_Module): """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) - - def time_evaluator( - self, - func_name, - dev, - number=10, - repeat=1, - min_repeat_ms=0, - limit_zero_time_iterations=100, - cooldown_interval_ms=0, - repeats_to_cooldown=1, - cache_flush_bytes=0, - f_preproc="", - ): - """Get an evaluator that measures time cost of running function. - - Parameters - ---------- - func_name: str - The name of the function in the module. - - dev: Device - The device we should run this function on. - - number: int - The number of times to run this function for taking average. - We call these runs as one `repeat` of measurement. - - repeat: int, optional - The number of times to repeat the measurement. - In total, the function will be invoked (1 + number x repeat) times, - where the first one is warm up and will be discarded. - The returned result contains `repeat` costs, - each of which is an average of `number` costs. - - min_repeat_ms: int, optional - The minimum duration of one `repeat` in milliseconds. - By default, one `repeat` contains `number` runs. If this parameter is set, - the parameters `number` will be dynamically adjusted to meet the - minimum duration requirement of one `repeat`. - i.e., When the run time of one `repeat` falls below this time, the `number` parameter - will be automatically increased. - - limit_zero_time_iterations: int, optional - The maximum number of repeats when measured time is equal to 0. - It helps to avoid hanging during measurements. - - cooldown_interval_ms: int, optional - The cooldown interval in milliseconds between the number of repeats defined by - `repeats_to_cooldown`. - - repeats_to_cooldown: int, optional - The number of repeats before the cooldown is activated. - - cache_flush_bytes: int, optional - The number of bytes to flush from the cache before each repeat. - - f_preproc: str, optional - The preprocess function name we want to execute before executing the time evaluator. - - Note - ---- - The function will be invoked (1 + number x repeat) times, - with the first call discarded in case there is lazy initialization. - - Returns - ------- - ftimer : function - The function that takes same argument as func and returns a BenchmarkResult. - The ProfileResult reports `repeat` time costs in seconds. - """ - try: - feval = _ffi_api.RPCTimeEvaluator( - self, - func_name, - dev.device_type, - dev.device_id, - number, - repeat, - min_repeat_ms, - limit_zero_time_iterations, - cooldown_interval_ms, - repeats_to_cooldown, - cache_flush_bytes, - f_preproc, - ) - - def evaluator(*args): - """Internal wrapped evaluator.""" - # Wrap feval so we can add more stats in future. - blob = feval(*args) - fmt = "@" + ("d" * repeat) - results = struct.unpack(fmt, blob) - return BenchmarkResult(results) - - return evaluator - 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. @@ -418,6 +131,7 @@ def _collect_from_import_tree(self, filter_func): return dso_modules def _collect_dso_modules(self): + """Collect all compilation exportable modules from the import tree.""" return self._collect_from_import_tree(lambda m: m.is_compilation_exportable()) def export_library( @@ -484,9 +198,10 @@ def export_library( # 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 + from tvm.contrib import cc as _cc, tar as _tar, tvmjs as _tvmjs, utils as _utils if isinstance(file_name, Path): file_name = str(file_name) @@ -559,7 +274,7 @@ def get_source_format_from_module(module): 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: + elif _ffi_api.RuntimeEnabled("llvm") and llvm_target_string: path_obj = os.path.join( workspace_dir, f"{pack_lib_prefix}devc.{global_object_format}" ) @@ -587,32 +302,103 @@ def get_source_format_from_module(module): return fcompile(file_name, files, **kwargs) + def time_evaluator( + self, + func_name, + dev, + number=10, + repeat=1, + min_repeat_ms=0, + limit_zero_time_iterations=100, + cooldown_interval_ms=0, + repeats_to_cooldown=1, + cache_flush_bytes=0, + f_preproc="", + ): + """Get an evaluator that measures time cost of running function. + + Parameters + ---------- + func_name: str + The name of the function in the module. + + dev: Device + The device we should run this function on. + + number: int + The number of times to run this function for taking average. + We call these runs as one `repeat` of measurement. + + repeat: int, optional + The number of times to repeat the measurement. + In total, the function will be invoked (1 + number x repeat) times, + where the first one is warm up and will be discarded. + The returned result contains `repeat` costs, + each of which is an average of `number` costs. + + min_repeat_ms: int, optional + The minimum duration of one `repeat` in milliseconds. + By default, one `repeat` contains `number` runs. If this parameter is set, + the parameters `number` will be dynamically adjusted to meet the + minimum duration requirement of one `repeat`. + i.e., When the run time of one `repeat` falls below this time, the `number` parameter + will be automatically increased. -def system_lib(symbol_prefix=""): - """Get system-wide library module singleton. + limit_zero_time_iterations: int, optional + The maximum number of repeats when measured time is equal to 0. + It helps to avoid hanging during measurements. - 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. + cooldown_interval_ms: int, optional + The cooldown interval in milliseconds between the number of repeats defined by + `repeats_to_cooldown`. - 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. + repeats_to_cooldown: int, optional + The number of repeats before the cooldown is activated. - 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. + cache_flush_bytes: int, optional + The number of bytes to flush from the cache before each repeat. - 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. + f_preproc: str, optional + The preprocess function name we want to execute before executing the time evaluator. - Returns - ------- - module : runtime.Module - The system-wide library module. - """ - return _mod_ffi_api.SystemLib(symbol_prefix) + Note + ---- + The function will be invoked (1 + number x repeat) times, + with the first call discarded in case there is lazy initialization. + + Returns + ------- + ftimer : function + The function that takes same argument as func and returns a BenchmarkResult. + The ProfileResult reports `repeat` time costs in seconds. + """ + try: + feval = _ffi_api.RPCTimeEvaluator( + self, + func_name, + dev.device_type, + dev.device_id, + number, + repeat, + min_repeat_ms, + limit_zero_time_iterations, + cooldown_interval_ms, + repeats_to_cooldown, + cache_flush_bytes, + f_preproc, + ) + + def evaluator(*args): + """Internal wrapped evaluator.""" + # Wrap feval so we can add more stats in future. + blob = feval(*args) + fmt = "@" + ("d" * repeat) + results = struct.unpack(fmt, blob) + return BenchmarkResult(results) + + return evaluator + except NameError: + raise NameError("time_evaluator is only supported when RPC is enabled") def load_module(path): @@ -656,7 +442,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 0db1fa93dc2a..168d839dbc7d 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)