# 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. # ruff: noqa: F401 # pylint: disable=invalid-name, unused-import, import-outside-toplevel, inconsistent-return-statements """Runtime Module namespace.""" import json import os import struct from collections.abc import Sequence import numpy as np from tvm_ffi import ( Module as _Module, ) from tvm_ffi import libinfo as tvm_ffi_libinfo from tvm_ffi import ( load_module as _load_module, ) from tvm_ffi import ( register_object as _register_object, ) from tvm_ffi import ( system_lib, ) import tvm.libinfo from tvm.base import _RUNTIME_ONLY from . import _ffi_api class BenchmarkResult: """Runtimes from benchmarking""" def __init__(self, results: Sequence[float]): """Construct a new BenchmarkResult from a sequence of runtimes. Parameters ---------- results : Sequence[float] Raw times from benchmarking Attributes ---------- min : float Minimum runtime in seconds of all results. mean : float Mean runtime in seconds of all results. If py:meth:`Module.time_evaluator` or `benchmark` is called with `number` > 0, then each result is already the mean of a `number` of runtimes, so this becomes the mean of means. median : float Median runtime in seconds of all results. If py:meth:`Module.time_evaluator` is called with `number` > 0, then each result is already the mean of a `number` of runtimes, so this becomes the median of means. max : float Maximum runtime in seconds of all results. If py:meth:`Module.time_evaluator` is called with `number` > 0, then each result is already the mean of a `number` of runtimes, so this becomes the maximum of those means. std : float Standard deviation in seconds of runtimes. If py:meth:`Module.time_evaluator` is called with `number` > 0, then each result is already the mean of a `number` of runtimes, so this becomes the standard deviation of means. results : Sequence[float] The collected runtimes (in seconds). This may be a series of mean runtimes if py:meth:`Module.time_evaluator` or `benchmark` was run with `number` > 1. """ self.results = results self.mean = np.mean(self.results) self.std = np.std(self.results) self.median = np.median(self.results) self.min = np.min(self.results) self.max = np.max(self.results) def __repr__(self): return ( f"BenchmarkResult(min={self.min}, mean={self.mean}, median={self.median}, " f"max={self.max}, std={self.std}, results={self.results})" ) def __str__(self): return ( f"Execution time summary:\n" f"{'mean (ms)':^12} {'median (ms)':^12} {'max (ms)':^12} " f"{'min (ms)':^12} {'std (ms)':^12}\n" f"{self.mean * 1000:^12.4f} {self.median * 1000:^12.4f} {self.max * 1000:^12.4f} " f"{self.min * 1000:^12.4f} {self.std * 1000:^12.4f}" " " ) # override the Module class in ffi.Module @_register_object("ffi.Module") class Module(_Module): """Runtime Module.""" 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): """Collect all compilation exportable modules from the import tree.""" 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 tvmjs as _tvmjs from tvm.support import cc as _cc from tvm.support import tar as _tar from tvm.support import utils as _utils 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 = 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 = 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 is None and hasattr(fcompile, "get_target_triple"): triple = fcompile.get_target_triple() assert triple, "Target triple should not be empty" llvm_target = json.dumps({"kind": "llvm", "mtriple": triple.strip()}) if getattr(fcompile, "need_system_lib", False) and not is_system_lib: raise ValueError(f"{fcompile!s} 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 _ffi_api.RuntimeEnabled("llvm") and llvm_target: 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, 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] default_include_paths = [ tvm.libinfo.find_include_path(), tvm_ffi_libinfo.find_include_path(), tvm_ffi_libinfo.find_dlpack_include_path(), ] opts = options + ["-I" + path for path in default_include_paths] kwargs.update({"options": opts}) 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. 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.dlpack_device_type(), dev.index, 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): """Load module from file. Parameters ---------- path : str The path to the module file. Returns ------- module : runtime.Module The loaded module Note ---- This function will automatically call cc.create_shared if the path is in format .o or .tar """ if os.path.isfile(path): path = os.path.realpath(path) else: raise ValueError(f"cannot find file {path}") # High level handling for .o and .tar file. # We support this to be consistent with RPC module load. if path.endswith(".o"): # Extra dependencies during runtime. from tvm.support import cc as _cc _cc.create_shared(path + ".so", path) path += ".so" elif path.endswith(".tar"): # Extra dependencies during runtime. from tvm.support import cc as _cc from tvm.support import tar as _tar from tvm.support import utils as _utils tar_temp = _utils.tempdir(custom_path=path.replace(".tar", "")) _tar.untar(path, tar_temp.temp_dir) files = [tar_temp.relpath(x) for x in tar_temp.listdir()] _cc.create_shared(path + ".so", files) path += ".so" # Redirect to the load API return _load_module(path) def load_static_library(path, func_names): """Load the .o library at path which implements functions with func_names. Unlike the generic load_module the result will remain as a static_library and will not be relinked on-the-fly into a .so library.""" return _ffi_api.ModuleLoadStaticLibrary(path, func_names) def enabled(target): """Whether module runtime is enabled for target Parameters ---------- target : str or Dict[str, Any] or tvm.target.Target The target device type. Returns ------- enabled : bool Whether runtime is enabled. Examples -------- The following code checks if gpu is enabled. >>> tvm.runtime.enabled("gpu") """ if isinstance(target, dict): target = target.get("kind", "") elif hasattr(target, "kind"): target = target.kind.name return _ffi_api.RuntimeEnabled(target) def num_threads() -> int: """Get the number of threads in use by the TVM runtime. Returns ------- int Number of threads in use. """ return _ffi_api.NumThreads()