def with_minrpc(compile_func, server="posix_popen_server", runtime="libtvm"):
"""Attach the compiler function with minrpc related options.
Parameters
----------
compile_func : Union[str, Callable[[str, str, Optional[str]], None]]
The compilation function to decorate.
server : str
The server type.
runtime : str
The runtime library.
Returns
-------
fcompile : function
The return compilation.
"""
server_path = find_minrpc_server_libpath(server)
runtime_path = libinfo.find_lib_path([runtime, runtime + ".so", runtime + ".dylib"])[0]
runtime_dir = os.path.abspath(os.path.dirname(runtime_path))
options = ["-std=c++14"]
# Make sure the rpath to the libtvm is set so we can do local tests.
# Note that however, this approach won't work on remote.
# Always recommend to to link statically.
options += ["-Wl,-rpath=" + runtime_dir]
options += ["-I" + path for path in libinfo.find_include_path()]
fcompile = cc.cross_compiler(
compile_func, options=options, add_files=[server_path, runtime_path]
)
fcompile.__name__ = "with_minrpc"
fcompile.need_system_lib = True
return fcompile
def check_device(device):
ctx = tvm.context(device, 0)
if not ctx.exist:
print("Skip because %s is not enabled" % device)
return
temp = util.tempdir()
name = "myadd_%s" % device
if sys.platform == "darwin" or sys.platform.startswith('linux'):
f = tvm.build(s, [A, B], device, "llvm -system-lib", name=name)
elif sys.platform == "win32":
f = tvm.build(s, [A, B], device, "llvm", name=name)
else:
raise ValueError("Unsupported platform")
path_dso = temp.relpath("dev_lib.so")
# test cross compiler function
f.export_library(path_dso, cc.cross_compiler("g++"))
f1 = tvm.runtime.load_module(path_dso)
a = tvm.nd.array(np.random.uniform(size=1024).astype(A.dtype), ctx)
b = tvm.nd.array(np.zeros(1024, dtype=A.dtype), ctx)
f1(a, b)
np.testing.assert_equal(b.asnumpy(), a.asnumpy() + 1)
if sys.platform != "win32":
f2 = tvm.runtime.system_lib()
f2[name](a, b)
np.testing.assert_equal(b.asnumpy(), a.asnumpy() + 1)
def create_aot_shared(so_name: Union[str, pathlib.Path],
files,
hexagon_arch: str,
options=None):
"""Export Hexagon AOT module."""
options = options or []
if not os.access(str(HEXAGON_CLANG_PLUS), os.X_OK):
raise Exception('The Clang++ "' + str(HEXAGON_CLANG_PLUS) +
'" does not exist or is not executable.')
if not HEXAGON_TOOLCHAIN:
raise Exception(
" The environment variable HEXAGON_TOOLCHAIN is unset. Please export "
+ "HEXAGON_TOOLCHAIN in your environment.")
if not HEXAGON_SDK_PATH:
raise Exception(
" The environment variable HEXAGON_SDK_PATH is unset. Please export "
+ "HEXAGON_SDK_PATH in your environment.")
# The AOT C codegen uses TVM runtime functions
# (e.g. TVMBackendAllocWorkspace) directly. On Hexagon these calls
# should be made using functions pointers provided as __TVM*
# variables in the provided context. This workaround allows the
# the TVM runtime symbols to be visible to the compiled shared
# library.
#
# This workaround can be removed when AOT codegen can be done with
# LLVM codegen.
workaround_link_flags = os.environ.get("HEXAGON_SHARED_LINK_FLAGS")
if workaround_link_flags:
options.extend(workaround_link_flags.split())
tvm_dir = pathlib.Path(os.path.dirname(
os.path.realpath(__file__))) / ".." / ".." / ".." / ".."
compute_arch = f"compute{hexagon_arch}"
compile_options = [
f"-O3",
f"-I{tvm_dir / 'include'}",
f"-I{tvm_dir / '3rdparty' / 'dlpack' / 'include'}",
f"-I{tvm_dir / '3rdparty' / 'dmlc-core' / 'include'}",
f"-I{pathlib.Path(HEXAGON_SDK_PATH) / 'rtos' / 'qurt' / compute_arch / 'include'/ 'posix'}",
f"-I{pathlib.Path(HEXAGON_SDK_PATH) / 'rtos' / 'qurt' / compute_arch / 'include' / 'qurt'}",
f"-DDMLC_USE_LOGGING_LIBRARY=<tvm/runtime/logging.h>",
f"-D_MACH_I32=int",
]
# For debugging
for path in HEXAGON_SDK_INCLUDE_DIRS:
compile_options.append(f"-I{str(path)}")
cross_compile = cc.cross_compiler(compile_func=hexagon_clang_plus())
cross_compile.output_format = "o"
c_files = [str(file) for file in files]
cross_compile(str(so_name), c_files, options=compile_options + options)
def save_module(module_path, graph, lib, params, cross=None):
"""
Create a tarball containing the generated TVM graph,
exported library and parameters
Parameters
----------
module_path : str
path to the target tar.gz file to be created,
including the file name
graph : str
A JSON-serialized TVM execution graph.
lib : tvm.module.Module
A TVM module containing the compiled functions.
params : dict
The parameters (weights) for the TVM module.
cross : str or callable object, optional
Function that performs the actual compilation
"""
lib_name = "mod.so"
graph_name = "mod.json"
param_name = "mod.params"
temp = util.tempdir()
path_lib = temp.relpath(lib_name)
if not cross:
logger.debug("exporting library to %s", path_lib)
lib.export_library(path_lib)
else:
logger.debug("exporting library to %s , using cross compiler %s",
path_lib, cross)
lib.export_library(path_lib, cc.cross_compiler(cross))
with open(temp.relpath(graph_name), "w") as graph_file:
logger.debug("writing graph to file to %s", graph_file.name)
graph_file.write(graph)
with open(temp.relpath(param_name), "wb") as params_file:
logger.debug("writing params to file to %s", params_file.name)
params_file.write(relay.save_param_dict(params))
logger.debug("saving module as tar file to %s", module_path)
with tarfile.open(module_path, "w") as tar:
tar.add(path_lib, lib_name)
tar.add(temp.relpath(graph_name), graph_name)
tar.add(temp.relpath(param_name), param_name)
def create_aot_shared(so_name: Union[str, pathlib.Path],
files,
hexagon_arch: str,
options=None):
"""Export Hexagon AOT module."""
options = options or []
if not os.access(str(HEXAGON_CLANG_PLUS), os.X_OK):
raise Exception('The Clang++ "' + str(HEXAGON_CLANG_PLUS) +
'" does not exist or is not executable.')
if not HEXAGON_TOOLCHAIN:
raise Exception(
" The environment variable HEXAGON_TOOLCHAIN is unset. Please export "
+ "HEXAGON_TOOLCHAIN in your environment.")
if not HEXAGON_SDK_PATH:
raise Exception(
" The environment variable HEXAGON_SDK_PATH is unset. Please export "
+ "HEXAGON_SDK_PATH in your environment.")
tvm_dir = pathlib.Path(os.path.dirname(
os.path.realpath(__file__))) / ".." / ".." / ".." / ".."
compute_arch = f"compute{hexagon_arch}"
compile_options = [
f"-O3",
f"-I{tvm_dir / 'include'}",
f"-I{tvm_dir / '3rdparty' / 'dlpack' / 'include'}",
f"-I{tvm_dir / '3rdparty' / 'dmlc-core' / 'include'}",
f"-I{pathlib.Path(HEXAGON_SDK_PATH) / 'rtos' / 'qurt' / compute_arch / 'include'/ 'posix'}",
f"-I{pathlib.Path(HEXAGON_SDK_PATH) / 'rtos' / 'qurt' / compute_arch / 'include' / 'qurt'}",
f"-DDMLC_USE_LOGGING_LIBRARY=<tvm/runtime/logging.h>",
f"-D_MACH_I32=int",
]
# For debugging
for path in HEXAGON_SDK_INCLUDE_DIRS:
compile_options.append(f"-I{str(path)}")
cross_compile = cc.cross_compiler(compile_func=hexagon_clang_plus())
cross_compile.output_format = "o"
c_files = [str(file) for file in files]
cross_compile(str(so_name), c_files, options=compile_options + options)
def cross_compiler(dev_config, lib_type):
"""Create a cross-compile function that wraps `create_lib` for a `Binutil` instance.
For use in `tvm.module.Module.export_library`.
Parameters
----------
dev_config : Dict[str, Any]
MicroTVM config dict for the target device
lib_type : micro.LibType
whether to compile a MicroTVM runtime or operator library
Return
------
func : Callable[[str, str, Optional[str]], None]
cross compile function taking a destination path for the object file
and a path for the input source file.
Example
--------
.. code-block:: python
c_mod = ... # some module generated with "c" as the target
fcompile = tvm.micro.cross_compiler(dev_config, LibType.OPERATOR)
c_mod.export_library("dev_lib.obj", fcompile=fcompile)
"""
dev_funcs = tvm.micro.device.get_device_funcs(dev_config['device_id'])
create_micro_lib = dev_funcs['create_micro_lib']
def compile_func(obj_path, src_path, **kwargs):
if isinstance(obj_path, list):
obj_path = obj_path[0]
if isinstance(src_path, list):
src_path = src_path[0]
create_micro_lib(obj_path, src_path, lib_type,
kwargs.get("options", None))
return _cc.cross_compiler(compile_func, output_format="obj")
def cross_compiler(toolchain_prefix, include_dev_lib_header=True):
"""Creates a cross compile function that wraps `create_micro_lib`.
For use in `tvm.module.Module.export_library`.
Parameters
----------
toolchain_prefix : str
toolchain prefix to be used
include_dev_lib_header : Optional[bool]
whether to include the device library header containing definitions of
library functions.
Return
------
func : Callable[[str, str, Optional[str]], None]
cross compile function taking a destination path for the object file
and a path for the input source file.
Example
--------
.. code-block:: python
c_mod = ... # some module generated with "c" as the target
fcompile = tvm.micro.cross_compiler(toolchain_prefix="")
c_mod.export_library("dev_lib.obj", fcompile=fcompile)
"""
def compile_func(obj_path, src_path, **kwargs):
if isinstance(obj_path, list):
obj_path = obj_path[0]
if isinstance(src_path, list):
src_path = src_path[0]
create_micro_lib(obj_path, src_path, toolchain_prefix,
kwargs.get("options", None), include_dev_lib_header)
return _cc.cross_compiler(compile_func)
def cross_compiler(dev_config,
lib_type,
lib_src_paths=None,
lib_headers=None,
lib_include_paths=None):
"""Create a cross compile function that wraps `create_lib` for a `Binutil` instance.
For use in `tvm.runtime.Module.export_library`.
Parameters
----------
create_micro_lib : func
function for creating MicroTVM libraries for a specific device (e.g.,
`tvm.micro.device.get_device_funcs('arm.stm32f746xx')['create_micro_lib']`)
lib_type : micro.LibType
whether to compile a MicroTVM runtime or operator library
lib_src_paths: TODO
TODO
lib_headers: TODO
e.g., `['cmsis_gcc.h', 'arm_math.h']`
lib_include_paths: TODO
TODO
Return
------
func : Callable[[str, str, Optional[str]], None]
cross compile function taking a destination path for the object file
and a path for the input source file.
Example
--------
.. code-block:: python
c_mod = ... # some module generated with "c" as the target
fcompile = tvm.micro.cross_compiler(dev_config, LibType.OPERATOR)
c_mod.export_library('dev_lib.obj', fcompile=fcompile)
"""
assert (lib_headers is None) == (lib_include_paths is None), \
"must specify both `lib_headers` and `lib_include_paths` or neither"
if lib_src_paths is None:
lib_src_paths = []
if lib_include_paths is None:
lib_include_paths = []
include_options = []
for include_path in lib_include_paths:
include_options.append("-I")
include_options.append(include_path)
create_micro_lib = tvm.micro.device.get_device_funcs(
dev_config["device_id"])["create_micro_lib"]
mem_layout = dev_config["mem_layout"]
def compile_func(obj_path, src_path, **kwargs):
if isinstance(obj_path, list):
obj_path = obj_path[0]
if isinstance(src_path, list):
src_path = src_path[0]
options = kwargs.get("options", [])
options += include_options
# check that workspace allocations don't exceed available workspace memory
with open(src_path) as f:
src_contents = f.read()
max_ws_usage = _calc_max_workspace_usage(src_contents)
available_mem = mem_layout["workspace"]["size"]
if max_ws_usage > available_mem:
raise RuntimeError(
f"workspace allocations in library ({max_ws_usage}) "
f"exceed available memory ({available_mem})")
# inject headers into new source path, if requested
if lib_headers:
headers_to_inject = "\n".join(
map(lambda s: f"#include <{s}>", lib_headers)) + "\n"
new_src_contents = headers_to_inject + src_contents
tmp_dir = _util.tempdir()
src_path = tmp_dir.relpath(os.path.basename(src_path))
with open(src_path, "w") as f:
f.write(new_src_contents)
create_micro_lib(obj_path,
src_path,
lib_type,
options,
lib_src_paths=lib_src_paths)
return _cc.cross_compiler(compile_func, output_format="obj")
