def optimize(self, mod, target=None, target_host=None, params=None):
"""Helper method that optimizes a Relay module via VM.
Parameters
----------
mod : tvm.IRModule
target : any multi-target like object, see Target.canon_multi_target
For homogeneous compilation, the unique build target.
For heterogeneous compilation, a dictionary or list of possible build targets.
target_host : any target-like object, see Target.canon_target
Host compilation target, if target is device.
params : dict of str to NDArray
Input parameters to the graph that do not change
during inference time. Used for constant folding.
Returns
-------
mod : tvm.IRModule
The optimized relay module.
params : dict
The parameters of the final module.
"""
raw_targets = Target.canon_multi_target_and_host(target, target_host)
if params:
self.set_params(params)
return self._optimize(mod, raw_targets), self.get_params()
def test_canon_multi_target_and_host_6():
"""Test `canon_target_and_host` by using TVM Objects"""
cuda_device_type = tvm.device("cuda").device_type
target = {cuda_device_type: Target(target="cuda", host="llvm")}
host = None
raw_targets_1 = Target.canon_multi_target_and_host(target, host)
assert len(raw_targets_1) == 1
assert raw_targets_1[0].kind.name == "cuda"
assert raw_targets_1[0].host.kind.name == "llvm"
target = {cuda_device_type: Target(tvm.runtime.container.String("cuda"))}
host = Target(tvm.runtime.container.String("llvm"))
target = tvm.runtime.convert(target)
assert isinstance(target, tvm.ir.container.Map)
raw_targets_2 = Target.canon_multi_target_and_host(target, host)
assert len(raw_targets_2) == 1
assert raw_targets_2[0].kind.name == "cuda"
assert raw_targets_2[0].host.kind.name == "llvm"
def optimize(mod, target=None, params=None):
"""Helper function that optimizes a Relay module.
Parameters
----------
mod : :py:class:`~tvm.IRModule`
The module to build. Using relay.Function is deprecated.
target : None, or any multi-target like object, see Target.canon_multi_target
For homogeneous compilation, the unique build target.
For heterogeneous compilation, a dictionary or list of possible build targets.
Defaults to the current target in the environment if None.
params : dict of str to NDArray
Input parameters to the graph that do not change
during inference time. Used for constant folding.
Returns
-------
mod : :py:class:`~tvm.IRModule`
The optimized relay module.
params : dict
The parameters of the final graph.
"""
if not isinstance(mod, (IRModule, _function.Function)):
raise ValueError("Type of input parameter mod must be tvm.IRModule")
if isinstance(mod, _function.Function):
if params:
mod = bind_params_by_name(mod, params)
mod = IRModule.from_expr(mod)
warnings.warn(
"Please use input parameter mod (tvm.IRModule) "
"instead of deprecated parameter func (tvm.relay.function.Function)",
DeprecationWarning,
)
raw_targets = Target.canon_multi_target_and_host(
Target.target_or_current(target))
# If current dispatch context is fallback context (the default root context),
# then load pre-tuned parameters from TopHub
if isinstance(autotvm.DispatchContext.current, autotvm.FallbackContext):
tophub_context = autotvm.tophub.context(raw_targets)
else:
tophub_context = autotvm.utils.EmptyContext()
with tophub_context:
bld_mod = BuildModule()
mod, params = bld_mod.optimize(mod, target=raw_targets, params=params)
return mod, params
def lower(self, mod, target=None, target_host=None):
"""Lower the module to VM bytecode.
Parameters
----------
mod : tvm.IRModule
The Relay module to build.
target : any multi-target like object, see Target.canon_multi_target
For homogeneous compilation, the unique build target.
For heterogeneous compilation, a dictionary or list of possible build targets.
target_host : any target-like object, see Target.canon_target
Host compilation target, if target is device.
"""
raw_targets = Target.canon_multi_target_and_host(target, target_host)
tophub_context = self._tophub_context(raw_targets)
with tophub_context:
self._lower(mod, raw_targets)
def optimize(self, mod, target=None, target_host=None, params=None):
"""
Parameters
----------
mod : :py:class:`~tvm.IRModule`
The IR module to build.
target : any multi-target like object, see Target.canon_multi_target.
For homogeneous compilation, the unique build target.
For heterogeneous compilation, a dictionary or list of possible build targets.
target_host : None, or any target-like object, see Target.canon_target
Host compilation target, if target is device.
params : dict of str to NDArray
Input parameters to the graph that do not change
during inference time. Used for constant folding.
Returns
-------
mod : :py:class:`~tvm.IRModule`
The optimized relay module.
params : dict
The parameters of the final graph.
"""
raw_targets = Target.canon_multi_target_and_host(target, target_host)
# Setup the params.
if params:
self._set_params(params)
mod = self._optimize(mod, raw_targets)
# Get artifacts
params = self.get_params()
return mod, params
def test_canon_multi_target_and_host_5():
raw_targets = Target.canon_multi_target_and_host("cuda", "llvm")
assert len(raw_targets) == 1
assert raw_targets[0].kind.name == "cuda"
assert raw_targets[0].host.kind.name == "llvm"
def test_canon_multi_target_and_host_3():
raw_targets = Target.canon_multi_target_and_host(["llvm", "cuda"])
assert len(raw_targets) == 2
assert raw_targets[0].kind.name == "llvm"
assert raw_targets[1].kind.name == "cuda"
def test_canon_multi_target_and_host_2():
raw_targets = Target.canon_multi_target_and_host({1: "llvm", 2: "cuda"})
assert len(raw_targets) == 2
assert raw_targets[0].kind.name == "llvm"
assert raw_targets[1].kind.name == "cuda"
def test_canon_multi_target_and_host_1():
raw_targets = Target.canon_multi_target_and_host({"kind": "llvm"})
assert len(raw_targets) == 1
assert raw_targets[0].kind.name == "llvm"
def test_canon_multi_target_and_host_0():
with pytest.raises(AssertionError):
Target.canon_multi_target_and_host(None)
def _setup(self, mod, target):
raw_targets = Target.canon_multi_target_and_host(target)
self._init(mod, raw_targets)
def build(
self,
mod,
target=None,
target_host=None,
executor=Executor("graph"),
runtime=Runtime("cpp"),
workspace_memory_pools=None,
params=None,
mod_name=None,
):
"""
Parameters
----------
mod : :py:class:`~tvm.IRModule`
The IRModule to build.
target : any multi-target like object, see Target.canon_multi_target
For homogeneous compilation, the unique build target.
For heterogeneous compilation, a dictionary or list of possible build targets.
target_host : None, or any target-like object, see Target.canon_target
Host compilation target, if target is device.
When TVM compiles device specific program such as CUDA,
we also need host(CPU) side code to interact with the driver
to setup the dimensions and parameters correctly.
target_host is used to specify the host side codegen target.
By default, llvm is used if it is enabled,
otherwise a stackvm interpreter is used.
executor : Optional[Executor]
The executor configuration with which to build the model.
Defaults to "graph" if no executor specified.
runtime : Optional[Runtime]
Runtime configuration to use when building the model.
Defaults to "cpp" if no runtime specified.
workspace_memory_pools : Optional[WorkspaceMemoryPools]
The object that contains an Array of PoolInfo objects
that hold properties of workspace pools that could be
used by the inference.
params : dict of str to NDArray
Input parameters to the graph that do not change
during inference time. Used for constant folding.
mod_name: Optional[str]
The module name we will build
Returns
-------
graph_json : str
The json string that can be accepted by graph executor.
mod : tvm.Module
The module containing necessary libraries.
params : dict
The parameters of the final graph.
"""
raw_targets = Target.canon_multi_target_and_host(target, target_host)
# Setup the params.
if params:
self._set_params(params)
# Build the IR module. If auto_scheduler is not enabled,
# then use the TOPI-defined schedule.
use_auto_scheduler = PassContext.current().config.get(
"relay.backend.use_auto_scheduler", False)
# Turn off AutoTVM config not found warnings if auto_scheduler is enabled.
old_autotvm_silent = autotvm.GLOBAL_SCOPE.silent
autotvm.GLOBAL_SCOPE.silent = use_auto_scheduler or old_autotvm_silent
mod_name = mangle_module_name(mod_name)
self._build(mod, raw_targets, executor, runtime,
workspace_memory_pools, mod_name)
autotvm.GLOBAL_SCOPE.silent = old_autotvm_silent
# Get artifacts
mod = self.get_module()
params = self.get_params()
executor_config = self.get_graph_json() if str(
executor) == "graph" else None
return executor_config, mod, params
def build(
ir_mod,
target=None,
target_host=None,
executor=Executor("graph"),
runtime=Runtime("cpp"),
workspace_memory_pools=None,
params=None,
mod_name="default",
):
# fmt: off
# pylint: disable=line-too-long
"""Helper function that builds a Relay function to run on TVM graph executor.
Parameters
----------
ir_mod : :py:class:`~tvm.IRModule`
The IR module to build. Using relay.Function is deprecated.
target : None, or any multi-target like object, see Target.canon_multi_target
For homogeneous compilation, the unique build target.
For heterogeneous compilation, a dictionary or list of possible build targets.
Defaults to the current target in the environment if None.
target_host : None, or any target like object, see Target.canon_target
Host compilation target, if target is device.
executor : Optional[Executor]
The executor configuration with which to build the model.
Defaults to "graph" if no executor specified.
runtime : Optional[Runtime]
Runtime configuration to use when building the model.
Defaults to "cpp" if no runtime specified.
workspace_memory_pools : Optional[WorkspaceMemoryPools]
The object that contains an Array of PoolInfo objects
that hold properties of workspace pools that could be
used by the inference.
params : dict of str to NDArray
Input parameters to the graph that do not change
during inference time. Used for constant folding.
mod_name: Optional[str]
The module name we will build
Returns
-------
factory_module : tvm.relay.backend.executor_factory.ExecutorFactoryModule
The runtime factory for the TVM graph executor.
"""
# pylint: enable=line-too-long
# fmt: on
if not isinstance(ir_mod, (IRModule, _function.Function)):
raise ValueError("Type of input parameter mod must be tvm.IRModule")
if isinstance(ir_mod, _function.Function):
if params:
ir_mod = bind_params_by_name(ir_mod, params)
ir_mod = IRModule.from_expr(ir_mod)
warnings.warn(
"Please use input parameter mod (tvm.IRModule) "
"instead of deprecated parameter mod (tvm.relay.function.Function)",
DeprecationWarning,
)
raw_targets = Target.canon_multi_target_and_host(
Target.target_or_current(target), target_host)
assert len(raw_targets) > 0
target_host = raw_targets[0].host
# All of this logic is to raise deprecation warnings for various parameters
# TODO(Mousius) Remove these after some time
deprecated_params_target = target_host or list(raw_targets)[0]
deprecated_executor, deprecated_runtime = _reconstruct_from_deprecated_options(
deprecated_params_target)
if deprecated_executor:
executor = deprecated_executor
if deprecated_runtime:
runtime = deprecated_runtime
# If current dispatch context is fallback context (the default root context),
# then load pre-tuned parameters from TopHub
if isinstance(autotvm.DispatchContext.current, autotvm.FallbackContext):
tophub_context = autotvm.tophub.context(list(raw_targets))
else:
tophub_context = autotvm.utils.EmptyContext()
with tophub_context:
bld_mod = BuildModule()
graph_json, runtime_mod, params = bld_mod.build(
mod=ir_mod,
target=raw_targets,
params=params,
executor=executor,
runtime=runtime,
workspace_memory_pools=workspace_memory_pools,
mod_name=mod_name,
)
func_metadata = bld_mod.get_function_metadata()
devices = bld_mod.get_devices()
lowered_ir_mods = bld_mod.get_irmodule()
executor_codegen_metadata = bld_mod.get_executor_codegen_metadata()
if str(executor) == "aot":
executor_factory = _executor_factory.AOTExecutorFactoryModule(
ir_mod,
lowered_ir_mods,
raw_targets,
executor,
runtime,
runtime_mod,
mod_name,
params,
func_metadata,
executor_codegen_metadata,
devices,
)
elif str(executor) == "graph":
executor_factory = _executor_factory.GraphExecutorFactoryModule(
ir_mod,
raw_targets,
executor,
graph_json,
runtime_mod,
mod_name,
params,
func_metadata,
)
else:
assert False, "Executor " + executor + " not supported"
return executor_factory
