From e53341d50c69c66214412e32a69ece97adc27eee Mon Sep 17 00:00:00 2001 From: Junru Shao Date: Thu, 16 Dec 2021 09:47:56 -0800 Subject: [PATCH] [MetaScheduler] Misc improvement of the Measurer --- include/tvm/meta_schedule/builder.h | 6 ++- python/tvm/meta_schedule/builder/builder.py | 16 ++++++-- .../meta_schedule/builder/local_builder.py | 37 ++++++++++++++++--- .../tvm/meta_schedule/runner/local_runner.py | 7 +++- python/tvm/meta_schedule/runner/rpc_runner.py | 7 +++- python/tvm/meta_schedule/utils.py | 6 ++- src/meta_schedule/builder/builder.cc | 9 +++-- .../unittest/test_meta_schedule_builder.py | 4 +- 8 files changed, 74 insertions(+), 18 deletions(-) diff --git a/include/tvm/meta_schedule/builder.h b/include/tvm/meta_schedule/builder.h index b809843f4157..6029959d81e2 100644 --- a/include/tvm/meta_schedule/builder.h +++ b/include/tvm/meta_schedule/builder.h @@ -32,10 +32,13 @@ class BuilderInputNode : public runtime::Object { IRModule mod; /*! \brief The target to be built for. */ Target target; + /*! \brief The optional parameters used for build */ + Optional> params; void VisitAttrs(tvm::AttrVisitor* v) { v->Visit("mod", &mod); v->Visit("target", &target); + v->Visit("params", ¶ms); } static constexpr const char* _type_key = "meta_schedule.BuilderInput"; @@ -53,7 +56,8 @@ class BuilderInput : public runtime::ObjectRef { * \param mod The IRModule to be built. * \param target The target to be built for. */ - TVM_DLL explicit BuilderInput(IRModule mod, Target target); + TVM_DLL explicit BuilderInput(IRModule mod, Target target, + Optional> params = NullOpt); TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(BuilderInput, runtime::ObjectRef, BuilderInputNode); }; diff --git a/python/tvm/meta_schedule/builder/builder.py b/python/tvm/meta_schedule/builder/builder.py index 381051e85f55..ef5d3ca130a7 100644 --- a/python/tvm/meta_schedule/builder/builder.py +++ b/python/tvm/meta_schedule/builder/builder.py @@ -15,11 +15,11 @@ # specific language governing permissions and limitations # under the License. """Meta Schedule builders that translate IRModule to runtime.Module, and then export""" -from typing import List, Optional +from typing import Dict, List, Optional from tvm._ffi import register_object from tvm.ir import IRModule -from tvm.runtime import Object +from tvm.runtime import NDArray, Object from tvm.target import Target from .. import _ffi_api @@ -36,12 +36,19 @@ class BuilderInput(Object): The IRModule to be built. target : Target The target to be built for. + params: Optional[Dict[str, NDArray]] + The parameters for Relay build module """ mod: IRModule target: Target - def __init__(self, mod: IRModule, target: Target) -> None: + def __init__( + self, + mod: IRModule, + target: Target, + params: Optional[Dict[str, NDArray]] = None, + ) -> None: """Constructor. Parameters @@ -50,11 +57,14 @@ def __init__(self, mod: IRModule, target: Target) -> None: The IRModule to be built. target : Target The target to be built for. + params: Optional[Dict[str, NDArray]] + The parameters for Relay build module """ self.__init_handle_by_constructor__( _ffi_api.BuilderInput, # type: ignore # pylint: disable=no-member mod, target, + params, ) diff --git a/python/tvm/meta_schedule/builder/local_builder.py b/python/tvm/meta_schedule/builder/local_builder.py index 5146ac4e03b9..cac958b50e13 100644 --- a/python/tvm/meta_schedule/builder/local_builder.py +++ b/python/tvm/meta_schedule/builder/local_builder.py @@ -15,13 +15,14 @@ # specific language governing permissions and limitations # under the License. """Local builder that compile on the local host""" +import logging import os import tempfile -from typing import Callable, List, Optional, Union +from typing import Callable, Dict, List, Optional, Union from tvm._ffi import register_func from tvm.ir import IRModule -from tvm.runtime import Module +from tvm.runtime import Module, NDArray, load_param_dict, save_param_dict from tvm.target import Target from ...contrib.popen_pool import MapResult, PopenPoolExecutor, StatusKind @@ -29,6 +30,21 @@ from .builder import BuilderInput, BuilderResult, PyBuilder +logger = logging.getLogger(__name__) + + +def _serialize_params(params: Optional[Dict[str, NDArray]]) -> Optional[bytearray]: + if params is None: + return None + return save_param_dict(params) + + +def _deserialize_params(params: Optional[bytearray]) -> Optional[Dict[str, NDArray]]: + if params is None: + return None + return load_param_dict(params) + + class LocalBuilder(PyBuilder): """A builder that builds the given input on local host. @@ -52,7 +68,11 @@ class LocalBuilder(PyBuilder): .. code-block:: python - def default_build(mod: IRModule, target: Target) -> Module: + def default_build( + mod: IRModule, + target: Target, + params: Optional[Dict[str, NDArray]] + ) -> Module: ... T_EXPORT : typing._GenericAlias @@ -71,7 +91,7 @@ def default_export(mod: Module) -> str: please send the registration logic via initializer. """ - T_BUILD = Callable[[IRModule, Target], Module] + T_BUILD = Callable[[IRModule, Target, Optional[Dict[str, NDArray]]], Module] T_EXPORT = Callable[[Module], str] pool: PopenPoolExecutor @@ -110,6 +130,7 @@ def __init__( if max_workers is None: max_workers = cpu_count() + logger.info("LocalBuilder: max_workers = %d", max_workers) self.pool = PopenPoolExecutor( max_workers=max_workers, @@ -134,6 +155,7 @@ def build(self, build_inputs: List[BuilderInput]) -> List[BuilderResult]: self.f_export, build_input.mod, build_input.target, + _serialize_params(build_input.params), ) for build_input in build_inputs ], @@ -172,6 +194,7 @@ def _worker_func( _f_export: Union[None, str, T_EXPORT], mod: IRModule, target: Target, + params: Optional[bytearray], ) -> str: # Step 0. Get the registered functions f_build: LocalBuilder.T_BUILD = get_global_func_with_default_on_worker( @@ -183,14 +206,14 @@ def _worker_func( default_export, ) # Step 1. Build the IRModule - rt_mod: Module = f_build(mod, target) + rt_mod: Module = f_build(mod, target, _deserialize_params(params)) # Step 2. Export the Module artifact_path: str = f_export(rt_mod) return artifact_path @register_func("meta_schedule.builder.default_build") -def default_build(mod: IRModule, target: Target) -> Module: +def default_build(mod: IRModule, target: Target, _params: Optional[Dict[str, NDArray]]) -> Module: """Default build function. Parameters @@ -199,6 +222,8 @@ def default_build(mod: IRModule, target: Target) -> Module: The IRModule to be built. target : Target The target to be built. + _params : Optional[Dict[str, NDArray]] + The parameters to be used for the build. Must be None. Returns ------- diff --git a/python/tvm/meta_schedule/runner/local_runner.py b/python/tvm/meta_schedule/runner/local_runner.py index caa266f97eb3..6af403905cb4 100644 --- a/python/tvm/meta_schedule/runner/local_runner.py +++ b/python/tvm/meta_schedule/runner/local_runner.py @@ -16,7 +16,9 @@ # under the License. """Local Runner""" from contextlib import contextmanager +import logging from typing import Callable, List, Optional, Union + import tvm from ...contrib.popen_pool import PopenPoolExecutor @@ -25,12 +27,14 @@ from .config import EvaluatorConfig from .runner import PyRunner, RunnerFuture, RunnerInput, RunnerResult from .utils import ( - T_ARG_INFO_JSON_OBJ_LIST, T_ARGUMENT_LIST, + T_ARG_INFO_JSON_OBJ_LIST, alloc_argument_common, run_evaluator_common, ) +logger = logging.getLogger(__name__) + class LocalRunnerFuture(RunnerFuture): """Local based runner future @@ -214,6 +218,7 @@ def __init__( self.f_run_evaluator = f_run_evaluator self.f_cleanup = f_cleanup + logger.info("LocalRunner: max_workers = 1") self.pool = PopenPoolExecutor( max_workers=1, # one local worker timeout=timeout_sec, diff --git a/python/tvm/meta_schedule/runner/rpc_runner.py b/python/tvm/meta_schedule/runner/rpc_runner.py index 3ba1c1dccf5f..b38448f4b65a 100644 --- a/python/tvm/meta_schedule/runner/rpc_runner.py +++ b/python/tvm/meta_schedule/runner/rpc_runner.py @@ -17,6 +17,7 @@ """RPC Runner""" import concurrent.futures from contextlib import contextmanager +import logging import os.path as osp from typing import Callable, List, Optional, Union @@ -31,13 +32,16 @@ from .config import EvaluatorConfig, RPCConfig from .runner import PyRunner, RunnerFuture, RunnerInput, RunnerResult from .utils import ( - T_ARG_INFO_JSON_OBJ_LIST, T_ARGUMENT_LIST, + T_ARG_INFO_JSON_OBJ_LIST, alloc_argument_common, run_evaluator_common, ) +logger = logging.getLogger(__name__) + + class RPCRunnerFuture(RunnerFuture): """RPC based runner future @@ -275,6 +279,7 @@ def __init__( self.f_alloc_argument = f_alloc_argument self.f_run_evaluator = f_run_evaluator self.f_cleanup = f_cleanup + logger.info("RPCRunner: max_workers = %d", max_workers) self.pool = PopenPoolExecutor( max_workers=max_workers, timeout=rpc_config.session_timeout_sec, diff --git a/python/tvm/meta_schedule/utils.py b/python/tvm/meta_schedule/utils.py index a9ef514543f8..f3574a223d32 100644 --- a/python/tvm/meta_schedule/utils.py +++ b/python/tvm/meta_schedule/utils.py @@ -31,6 +31,10 @@ @register_func("meta_schedule.cpu_count") +def _cpu_count_impl(logical: bool = True) -> int: + return psutil.cpu_count(logical=logical) or 1 + + def cpu_count(logical: bool = True) -> int: """Return the number of logical or physical CPUs in the system @@ -56,7 +60,7 @@ def cpu_count(logical: bool = True) -> int: Setting these variables may interfere the host-side search with profiling of generated kernels when measuring locally. """ - return psutil.cpu_count(logical=logical) or 1 + return _cpu_count_impl(logical) def get_global_func_with_default_on_worker( diff --git a/src/meta_schedule/builder/builder.cc b/src/meta_schedule/builder/builder.cc index fb63b7e65332..25adfbfdfa1d 100644 --- a/src/meta_schedule/builder/builder.cc +++ b/src/meta_schedule/builder/builder.cc @@ -23,10 +23,12 @@ namespace meta_schedule { /******** Constructors ********/ -BuilderInput::BuilderInput(IRModule mod, Target target) { +BuilderInput::BuilderInput(IRModule mod, Target target, + Optional> params) { ObjectPtr n = make_object(); n->mod = std::move(mod); n->target = std::move(target); + n->params = std::move(params); data_ = std::move(n); } @@ -51,8 +53,9 @@ TVM_REGISTER_OBJECT_TYPE(BuilderNode); TVM_REGISTER_NODE_TYPE(PyBuilderNode); TVM_REGISTER_GLOBAL("meta_schedule.BuilderInput") - .set_body_typed([](IRModule mod, Target target) -> BuilderInput { - return BuilderInput(mod, target); + .set_body_typed([](IRModule mod, Target target, + Optional> params) -> BuilderInput { + return BuilderInput(mod, target, params); }); TVM_REGISTER_GLOBAL("meta_schedule.BuilderResult") diff --git a/tests/python/unittest/test_meta_schedule_builder.py b/tests/python/unittest/test_meta_schedule_builder.py index fb3fa135a9b8..af95ea57e34a 100644 --- a/tests/python/unittest/test_meta_schedule_builder.py +++ b/tests/python/unittest/test_meta_schedule_builder.py @@ -163,7 +163,7 @@ def test_meta_schedule_error_handle_build_func(): def initializer(): @register_func("meta_schedule.builder.test_build") - def test_build(mod: Module, target: Target) -> None: # pylint: disable=unused-variable + def test_build(mod: Module, target: Target, _) -> None: # pylint: disable=unused-variable raise ValueError("Builder intended Test Error (build func).") builder = LocalBuilder(f_build="meta_schedule.builder.test_build", initializer=initializer) @@ -201,7 +201,7 @@ def test_meta_schedule_error_handle_time_out(): def initializer(): @register_func("meta_schedule.builder.test_time_out") - def timeout_build(mod, target): # pylint: disable=unused-argument, unused-variable + def timeout_build(mod, target, _): # pylint: disable=unused-argument, unused-variable time.sleep(2) builder = LocalBuilder(