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