[MetaSchedule][M4b] Misc improvement of the Measurer

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<Map<String, runtime::NDArray>> params;
 
   void VisitAttrs(tvm::AttrVisitor* v) {
     v->Visit("mod", &mod);
     v->Visit("target", &target);
+    v->Visit("params", &params);
   }
 
   static constexpr const char* _type_key = "meta_schedule.BuilderInput";
@@ -52,8 +55,10 @@ class BuilderInput : public runtime::ObjectRef {
    * \brief Constructor of BuilderInput.
    * \param mod The IRModule to be built.
    * \param target The target to be built for.
+   * \param params The optional parameters used for build
    */
-  TVM_DLL explicit BuilderInput(IRModule mod, Target target);
+  TVM_DLL explicit BuilderInput(IRModule mod, Target target,
+                                Optional<Map<String, runtime::NDArray>> params = NullOpt);
   TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(BuilderInput, runtime::ObjectRef, BuilderInputNode);
 };
 

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,
         )
 
 

python/tvm/meta_schedule/builder/local_builder.py

@@ -15,20 +15,36 @@
 # 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
 from ..utils import cpu_count, get_global_func_with_default_on_worker
 from .builder import BuilderInput, BuilderResult, PyBuilder
 
 
+logger = logging.getLogger(__name__)  # pylint: disable=invalid-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
     -------

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__)  # pylint: disable=invalid-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,

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__)  # pylint: disable=invalid-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,

python/tvm/meta_schedule/utils.py

@@ -32,6 +32,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
 
@@ -57,7 +61,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(

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<Map<String, runtime::NDArray>> params) {
   ObjectPtr<BuilderInputNode> n = make_object<BuilderInputNode>();
   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<Map<String, runtime::NDArray>> params) -> BuilderInput {
+      return BuilderInput(mod, target, params);
     });
 
 TVM_REGISTER_GLOBAL("meta_schedule.BuilderResult")

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(