addressing comments

Author: mkatanbaf

python/tvm/contrib/micro/meta_schedule/local_builder_micro.py

@@ -14,7 +14,7 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
-"""Local builder for microTVM projects.that compile on the local host"""
+"""Local builder for microTVM projects that compile on the local host"""
 
 import os
 import tempfile
@@ -34,10 +34,10 @@
 def get_micro_local_builder():
     """Return micro-compatible Builder for meta schedule."""
 
-    def micro_build(
+    def _micro_build(
         mod: IRModule, target: Target, _params: Optional[Dict[str, NDArray]]
     ) -> OperatorModule:
-        """build function for micro targets.
+        """Build function for micro targets.
 
         Parameters
         ----------
@@ -54,7 +54,8 @@ def micro_build(
             The built Module.
         """
 
-        # Note: changing the global symbol is necessary for micro targets,
+        # Note: tvm_build assigns "global_symbol" to the name of generated C function
+        # changing it is necessary for micro targets,
         # since the generated projects already include a main function.
         prim_func = mod["main"].with_attr("global_symbol", "default_function")
         mod = IRModule({"main": prim_func})
@@ -63,8 +64,8 @@ def micro_build(
         rt_mod = tvm_build(mod, target=target, runtime=runtime)
         return rt_mod
 
-    def micro_export(mod: OperatorModule) -> str:
-        """export function for micro targets.
+    def _micro_export(mod: OperatorModule) -> str:
+        """Export function for micro targets.
 
         Parameters
         ----------
@@ -80,4 +81,4 @@ def micro_export(mod: OperatorModule) -> str:
         micro.export_model_library_format(mod, artifact_path)
         return artifact_path
 
-    return LocalBuilder(f_build=micro_build, f_export=micro_export)
+    return LocalBuilder(f_build=_micro_build, f_export=_micro_export)

python/tvm/contrib/micro/meta_schedule/rpc_runner_micro.py

@@ -27,10 +27,7 @@
 from tvm.rpc.server import Server
 from tvm.rpc.tracker import Tracker
 from tvm.meta_schedule.logging import get_logger
-from tvm.meta_schedule.utils import (
-    cpu_count,
-    derived_object,
-)
+from tvm.meta_schedule.utils import cpu_count, derived_object
 from tvm.meta_schedule.runner.config import EvaluatorConfig, RPCConfig
 from tvm.meta_schedule.runner import PyRunner, RunnerFuture, RunnerInput
 from tvm.meta_schedule.runner.rpc_runner import RPCRunnerFuture
@@ -159,51 +156,44 @@ def _worker_func(
 def get_rpc_runner_micro(
     platform,
     options,
-    session_timeout_sec: int = 300,
-    number: int = 3,
-    repeat: int = 1,
-    min_repeat_ms: int = 100,
+    rpc_config: RPCConfig = None,
+    evaluator_config: EvaluatorConfig = None,
+    session_timeout_sec=300,
 ):
     """Parameters
     ----------
     platform: str
         The platform used for project generation.
     project_options: dict
         The options for the generated micro project.
+    rpc_config: RPCConfig
+        The rpc configuration.
+    evaluator_config: EvaluatorConfig
+        The evaluator configuration.
     session_timeout_sec: int
         The session timeout. if the number of candidates sent to runner is larger
         than the runner workers, increase the timeout.
-    number: int
-        The number of times to run the evaluator function for taking average.
-        We call these runs as one `repeat` of measurement.
-    repeat: int
-        The number of times to repeat the measurement.
-        In total, the function will be invoked (1 + number x repeat) times,
-        where the first one is warm up and will be discarded.
-        The returned result contains `repeat` costs,
-        each of which is an average of `number` costs.
-    min_repeat_ms: int
-        Minimum repeat time in ms. if the execution latency is too short,
-        increase the number of runs to the given time (in ms) to reduce the measurement error.
     """
-    tracker_host = "127.0.0.1"
-    tracker_port = 9000
-    tracker_key = "$local$device$%d" % tracker_port
-    rpc_config = RPCConfig(
-        tracker_host=tracker_host,
-        tracker_port=tracker_port,
-        tracker_key=tracker_key,
-        session_priority=0,
-        session_timeout_sec=session_timeout_sec,
-    )
-    rpc_config = RPCConfig._normalized(rpc_config)
-    tracker_port_end = 10000
-    evaluator_config = EvaluatorConfig(
-        number=number,
-        repeat=repeat,
-        min_repeat_ms=min_repeat_ms,
-        enable_cpu_cache_flush=False,
-    )
+    if rpc_config is None:
+        tracker_host = "127.0.0.1"
+        tracker_port = 9000
+        tracker_key = "$local$device$%d" % tracker_port
+        rpc_config = RPCConfig(
+            tracker_host=tracker_host,
+            tracker_port=tracker_port,
+            tracker_key=tracker_key,
+            session_priority=0,
+            session_timeout_sec=session_timeout_sec,
+        )
+    tracker_port_end = rpc_config.tracker_port + 1000
+
+    if evaluator_config is None:
+        evaluator_config = EvaluatorConfig(
+            number=3,
+            repeat=1,
+            min_repeat_ms=100,
+            enable_cpu_cache_flush=False,
+        )
 
     tracker = Tracker(
         port=rpc_config.tracker_port,

python/tvm/contrib/micro/meta_schedule/test_autotune_ms.py

@@ -105,7 +105,7 @@ def test_micro_tuning_with_meta_schedule(platform, options):
         platform=platform, options=options, session_timeout_sec=120
     ) as runner:
         with ms.Profiler() as profiler:
-            db: tvm.runtime.Module = ms.relay_integration.tune_relay(
+            db: ms.Database = ms.relay_integration.tune_relay(
                 mod=mod,
                 params=params,
                 target=target,

src/meta_schedule/schedule_rule/schedule_rule.cc

@@ -261,10 +261,11 @@ Array<ScheduleRule> ScheduleRule::DefaultMicro() {
           /*require_injective=*/true,
           /*require_ordered=*/true,
           /*disallow_op=*/Array<String>{"tir.exp"}),
-      ScheduleRule::MultiLevelTilingWideVector(
-          /*structure=*/"SRSRS",
-          /*vector_length_in_bits=*/1024,
-          /*max_innermost_factor=*/Integer(128),
+      ScheduleRule::MultiLevelTiling(
+          /*structure=*/"SSRSRS",
+          /*tile_binds=*/NullOpt,
+          /*max_innermost_factor=*/Integer(64),
+          /*vector_load_lens=*/NullOpt,
           /*reuse_read=*/NullOpt,
           /*reuse_write=*/
           Map<String, ObjectRef>{{"req", String("may")},