wehub-resource-sync
512 lines
19 KiB
Python
512 lines
19 KiB
Python
# 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()
|