# 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.""" # tvm-ffi-stubgen(begin): import-section # fmt: off # isort: off from __future__ import annotations from typing import TYPE_CHECKING if TYPE_CHECKING: from collections.abc import Sequence from typing import Any # isort: on # fmt: on # tvm-ffi-stubgen(end) import json from enum import IntEnum from os import PathLike, fspath from typing import ClassVar, cast from . import _ffi_api, core from .registry import register_object __all__ = ["Module", "ModulePropertyMask", "load_module", "system_lib"] class ModulePropertyMask(IntEnum): """Runtime Module Property Mask.""" BINARY_SERIALIZABLE = 0b001 RUNNABLE = 0b010 COMPILATION_EXPORTABLE = 0b100 @register_object("ffi.Module") class Module(core.Object): """Module container for dynamically loaded Module. Examples -------- .. code-block:: python import tvm_ffi # Load the module from a shared library mod = tvm_ffi.load_module("path/to/library.so") # Call exported function mod.func_name(*args) # Query function metadata (type signature) metadata = mod.get_function_metadata("func_name") # Query function documentation (if available) doc = mod.get_function_doc("func_name") See Also -------- :py:func:`tvm_ffi.load_module` :py:meth:`get_function_metadata` :py:meth:`get_function_doc` Notes ----- If you load a module within a local scope, be careful when any called function creates and returns an object. The memory deallocation routines are part of the library's code. If the module is unloaded before the object is destroyed, the deleter may call an invalid address. Keep the module loaded until all returned objects are deleted. You can safely use returned objects inside a nested function that finishes before the module goes out of scope. When possible, consider keeping the module alive in a long-lived/global scope (for example, in a global state) to avoid premature unloading. .. code-block:: python def bad_pattern(x): # Bad: unload order of `tensor` and `mod` is not guaranteed mod = tvm_ffi.load_module("path/to/library.so") # ... do something with the tensor tensor = mod.func_create_and_return_tensor(x) def good_pattern(x): # Good: `tensor` is freed before `mod` goes out of scope mod = tvm_ffi.load_module("path/to/library.so") def run_some_tests(): tensor = mod.func_create_and_return_tensor(x) # ... do something with the tensor run_some_tests() """ # tvm-ffi-stubgen(begin): object/ffi.Module # fmt: off imports_: Sequence[Any] # fmt: on # tvm-ffi-stubgen(end) entry_name: ClassVar[str] = "main" # constant for entry function name @property def kind(self) -> str: """Get type key of the module.""" return _ffi_api.ModuleGetKind(self) @property def imports(self) -> list[Module]: """Get imported modules. Returns ------- modules The module """ return self.imports_ # type: ignore[return-value] def implements_function(self, name: str, query_imports: bool = False) -> bool: """Return True if the module defines a global function. Notes ----- ``implements_function(name)`` does not guarantee that ``get_function(name)`` will return a callable, because some module kinds (e.g. a source-only module) may not provide a packed function implementation until further compilation occurs. However, a non-null result from ``get_function(name)`` should imply the module implements the function. Parameters ---------- name The name of the function query_imports Whether to also query modules imported by this module. Returns ------- True if module (or one of its imports) has a definition for name. """ return _ffi_api.ModuleImplementsFunction(self, name, query_imports) def __getattr__(self, name: str) -> core.Function: """Accessor to allow getting functions as attributes.""" try: func = self.get_function(name) self.__dict__[name] = func return func except AttributeError: raise AttributeError(f"Module has no function '{name}'") def get_function(self, name: str, query_imports: bool = False) -> core.Function: """Get function from the module. Parameters ---------- name The name of the function query_imports Whether to also query modules imported by this module. Returns ------- The result function. """ func = _ffi_api.ModuleGetFunction(self, name, query_imports) func = cast(core.Function, func) if func is None: raise AttributeError(f"Module has no function '{name}'") return func def get_function_metadata( self, name: str, query_imports: bool = False ) -> dict[str, Any] | None: """Get metadata for a function exported from the module. This retrieves metadata for functions exported via c:macro:`TVM_FFI_DLL_EXPORT_TYPED_FUNC` and when c:macro:`TVM_FFI_DLL_EXPORT_INCLUDE_METADATA` is on, which includes type schema information. Parameters ---------- name The name of the function query_imports Whether to also query modules imported by this module. Returns ------- metadata A dictionary containing function metadata. The ``type_schema`` field encodes the callable signature. Examples -------- .. code-block:: python import tvm_ffi from tvm_ffi.core import TypeSchema import json mod = tvm_ffi.load_module("add_one_cpu.so") metadata = mod.get_function_metadata("add_one_cpu") schema = TypeSchema.from_json_str(metadata["type_schema"]) print(schema) # Shows function signature See Also -------- :py:func:`tvm_ffi.get_global_func_metadata` Get metadata for global registry functions. """ metadata_str = _ffi_api.ModuleGetFunctionMetadata(self, name, query_imports) if metadata_str is None: return None return json.loads(metadata_str) def get_function_doc(self, name: str, query_imports: bool = False) -> str | None: """Get documentation string for a function exported from the module. This retrieves documentation for functions exported via c:macro:`TVM_FFI_DLL_EXPORT_TYPED_FUNC_DOC`. Parameters ---------- name The name of the function query_imports Whether to also query modules imported by this module. Returns ------- doc : str or None The documentation string if available, None otherwise. Examples -------- .. code-block:: python import tvm_ffi mod = tvm_ffi.load_module("mylib.so") doc = mod.get_function_doc("process_batch") if doc: print(doc) See Also -------- :py:meth:`get_function_metadata` Get metadata including type schema. """ doc_str = _ffi_api.ModuleGetFunctionDoc(self, name, query_imports) return doc_str if doc_str else None def import_module(self, module: Module) -> None: """Add module to the import list of current one. Parameters ---------- module The other module. """ _ffi_api.ModuleImportModule(self, module) def __getitem__(self, name: str) -> core.Function: """Return function by name using item access (module["func"]).""" if not isinstance(name, str): raise ValueError("Can only take string as function name") return self.get_function(name) def __call__(self, *args: Any) -> Any: """Call the module's entry function (`main`).""" # pylint: disable=not-callable return self.main(*args) def inspect_source(self, fmt: str = "") -> str: """Get source code from module, if available. Parameters ---------- fmt The specified format. Returns ------- The result source code. """ return _ffi_api.ModuleInspectSource(self, fmt) def get_write_formats(self) -> Sequence[str]: """Get the format of the module.""" return _ffi_api.ModuleGetWriteFormats(self) def get_property_mask(self) -> int: """Get the runtime module property mask. The mapping is stated in ModulePropertyMask. Returns ------- Bitmask of runtime module property """ return _ffi_api.ModuleGetPropertyMask(self) def is_binary_serializable(self) -> bool: """Return whether the module is binary serializable (supports save_to_bytes). Returns ------- True if the module is binary serializable. """ return (self.get_property_mask() & ModulePropertyMask.BINARY_SERIALIZABLE) != 0 def is_runnable(self) -> bool: """Return whether the module is runnable (supports get_function). Returns ------- True if the module is runnable. """ return (self.get_property_mask() & ModulePropertyMask.RUNNABLE) != 0 def is_compilation_exportable(self) -> bool: """Return whether the module is compilation exportable. write_to_file is supported for object or source. Returns ------- True if the module is compilation exportable. """ return (self.get_property_mask() & ModulePropertyMask.COMPILATION_EXPORTABLE) != 0 def clear_imports(self) -> None: """Remove all imports of the module.""" _ffi_api.ModuleClearImports(self) def write_to_file(self, file_name: str, fmt: str = "") -> None: """Write the current module to file. Parameters ---------- file_name The name of the file. fmt The format of the file. See Also -------- tvm.runtime.Module.export_library : export the module to shared library. """ _ffi_api.ModuleWriteToFile(self, file_name, fmt) def system_lib(symbol_prefix: str = "") -> Module: """Get system-wide library module singleton with functions prefixed by ``__tvm_ffi_{symbol_prefix}``. The library module contains symbols that are registered via :cpp:func:`TVMFFIEnvModRegisterSystemLibSymbol`. .. note:: The system lib is intended to be statically linked and loaded during the entire lifecycle of the program. If you want dynamic loading features, use DSO modules instead. Parameters ---------- symbol_prefix 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 ------- The system-wide library module. Examples -------- Register the function ``test_symbol_add_one`` in C++ with the name ``__tvm_ffi_test_symbol_add_one`` via :cpp:func:`TVMFFIEnvModRegisterSystemLibSymbol`. .. code-block:: cpp // A function to be registered in the system lib static int test_symbol_add_one(void*, const TVMFFIAny* args, int32_t num_args, TVMFFIAny* ret) { TVM_FFI_SAFE_CALL_BEGIN(); TVM_FFI_CHECK(num_args == 1, "Expected 1 argument, but got: " + std::to_string(num_args)); int64_t x = reinterpret_cast(args)[0].cast(); reinterpret_cast(ret)[0] = x + 1; TVM_FFI_SAFE_CALL_END(); } // Register the function with name `test_symbol_add_one` prefixed by `__tvm_ffi_` int _ = TVMFFIEnvModRegisterSystemLibSymbol("__tvm_ffi_testing.add_one", reinterpret_cast(test_symbol_add_one)); Look up and call the function from Python: .. code-block:: python import tvm_ffi mod: tvm_ffi.Module = tvm_ffi.system_lib( "testing." ) # symbols prefixed with `__tvm_ffi_testing.` func: tvm_ffi.Function = mod["add_one"] # looks up `__tvm_ffi_testing.add_one` assert func(10) == 11 """ return _ffi_api.SystemLib(symbol_prefix) def load_module(path: str | PathLike, keep_module_alive: bool = True) -> Module: """Load module from file. Parameters ---------- path The path to the module file. keep_module_alive Whether to keep the module alive. If True, the module will be kept alive for the duration of the program until libtvm_ffi.so is unloaded. Returns ------- The loaded module Examples -------- .. code-block:: python import tvm_ffi from pathlib import Path # Works with string paths mod = tvm_ffi.load_module("path/to/module.so") # Also works with pathlib.Path objects mod = tvm_ffi.load_module(Path("path/to/module.so")) mod.func_name(*args) See Also -------- :py:class:`tvm_ffi.Module` """ path = fspath(path) mod = _ffi_api.ModuleLoadFromFile(path) if keep_module_alive: _ffi_api.ModuleGlobalsAdd(mod) return mod