From 8549f6395b6772014c43ce182fc71cef304cc250 Mon Sep 17 00:00:00 2001
From: Giuseppe Rossini <giuseppe.rossini@arm.com>
Date: Fri, 27 Nov 2020 18:28:23 +0000
Subject: [PATCH 1/8] Add autoscheduler support to tvmc

- Add an autoschedule module to tvmc
- Extract common tuning option between autotuner and autoscheduler
- Add testing
---
 python/tvm/driver/tvmc/__init__.py            |   1 +
 python/tvm/driver/tvmc/autoscheduler.py       | 217 ++++++++++++++++++
 python/tvm/driver/tvmc/autotuner.py           |  85 +------
 python/tvm/driver/tvmc/common.py              |  87 +++++++
 python/tvm/driver/tvmc/compiler.py            |  30 ++-
 .../python/driver/tvmc/test_autoscheduler.py  | 101 ++++++++
 6 files changed, 433 insertions(+), 88 deletions(-)
 create mode 100644 python/tvm/driver/tvmc/autoscheduler.py
 create mode 100644 tests/python/driver/tvmc/test_autoscheduler.py

diff --git a/python/tvm/driver/tvmc/__init__.py b/python/tvm/driver/tvmc/__init__.py
index d96a725877eb..37344af42d42 100644
--- a/python/tvm/driver/tvmc/__init__.py
+++ b/python/tvm/driver/tvmc/__init__.py
@@ -21,3 +21,4 @@
 from . import autotuner
 from . import compiler
 from . import runner
+from . import autoscheduler
diff --git a/python/tvm/driver/tvmc/autoscheduler.py b/python/tvm/driver/tvmc/autoscheduler.py
new file mode 100644
index 000000000000..d8772f1ed54d
--- /dev/null
+++ b/python/tvm/driver/tvmc/autoscheduler.py
@@ -0,0 +1,217 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from docutils.nodes import target
+
+"""
+Provides support to auto-tuning networks using AutoTVM.
+"""
+import os.path
+import logging
+import time
+
+from urllib.parse import urlparse
+
+from tvm import auto_scheduler
+from tvm.auto_scheduler.auto_schedule import HardwareParams
+
+from . import common, frontends
+from .common import TVMCException, add_tuning_options
+from .main import register_parser
+
+
+# pylint: disable=invalid-name
+logger = logging.getLogger("TVMC")
+
+
+@register_parser
+def add_autoscheduler_parser(subparsers):
+    """ Include parser for 'autoschedule' subcommand """
+    parser = subparsers.add_parser("autoschedule", help="auto-schedule a model")
+    parser.set_defaults(func=drive_autoschedule)
+    add_tuning_options(parser)
+
+    parser.add_argument(
+        "--cache-line-bytes",
+        default=64,
+        help="the size of cache line in bytes",
+    )
+    parser.add_argument(
+        "--num-cores",
+        default=4,
+        help="the number of device cores",
+    )
+    parser.add_argument(
+        "--vector-unit-bytes",
+        default=16,
+        help="the width of vector units in bytes",
+    )
+    parser.add_argument(
+        "--model-format",
+        choices=frontends.get_frontend_names(),
+        help="specify input model format",
+    )
+
+
+def drive_autoschedule(args):
+    """Invoke auto-tuning with command line arguments
+
+    Parameters
+    ----------
+    args: argparse.Namespace
+        Arguments from command line parser.
+    """
+
+    # extra arguments validation before importing the model, so that obvious errors
+    # are pointed in advance.
+    if args.rpc_tracker:
+        parsed_url = urlparse("//%s" % args.rpc_tracker)
+        rpc_hostname = parsed_url.hostname
+        rpc_port = parsed_url.port or 9090
+        logger.info("RPC tracker hostname: %s", rpc_hostname)
+        logger.info("RPC tracker port: %s", rpc_port)
+
+        if not args.rpc_key:
+            raise common.TVMCException(
+                "need to provide an RPC tracker key (--rpc-key) for remote tuning"
+            )
+
+    target = common.target_from_cli(args.target)
+    mod, params = frontends.load_model(args.FILE, args.model_format)
+
+    # min_repeat_ms should be:
+    # a. the value provided by the user, if any, or
+    # b. 0ms in case target is "cpu"; otherwise 1000ms
+    if args.min_repeat_ms is not None:
+        min_repeat_ms = args.min_repeat_ms
+    else:
+        min_repeat_ms = 0 if target.keys[0] == "cpu" else 1000
+        logger.debug("Default --min-repeat-ms for this target is %s", min_repeat_ms)
+
+    if args.rpc_tracker:
+
+        runner = auto_scheduler.RPCRunner(
+            key=args.rpc_key,
+            host=rpc_hostname,
+            port=rpc_port,
+            number=args.number,
+            repeat=args.repeat,
+            n_parallel=args.parallel,
+            timeout=args.timeout,
+            min_repeat_ms=min_repeat_ms,
+        )
+    else:
+        logger.info("starting localhost tuning")
+        runner = auto_scheduler.LocalRunner(
+            number=args.number,
+            repeat=args.repeat,
+            timeout=args.timeout,
+            min_repeat_ms=min_repeat_ms,
+        )
+
+    # Create the autoscheduler tuning options
+    tuning_options = auto_scheduler.TuningOptions(
+        num_measure_trials=args.trials,
+        measure_callbacks=[auto_scheduler.RecordToFile(args.output)],
+        runner=runner,
+        builder="local",
+        early_stopping=args.early_stopping,
+    )
+
+    # Specify hardware parameters
+    hardware_params = HardwareParams(args.num_cores, args.vector_unit_bytes, args.cache_line_bytes)
+
+    # Extract the tasks from the model
+    tasks, weights = get_tuning_tasks(
+        mod, params, target, target_host, args.desired_layout, hardware_params
+    )
+
+    # Schedule the tasks (i.e., produce a schedule for each task)
+    schedule_tasks(
+        mod,
+        params,
+        target,
+        target_host=args.target_host,
+        tuning_records=args.tuning_records,
+        tuning_options=tuning_options,
+        alter_layout=args.desired_layout,
+    )
+
+
+def get_tuning_tasks(
+    mod, params, target, target_host=None, alter_layout=None, hardware_params=None
+):
+    """Get the tuning tasks for a given relay module.
+
+    Parameters
+    ----------
+    mod : tvm.relay.Module
+        The relay module from which to extract tuning tasks.
+    params : dict
+        The params for the relay module.
+    target : tvm.target.Target
+        The compilation target.
+    target_host : str, optional
+        The compilation target for the host.
+    alter_layout : str, optional
+        The layout to convert the graph to. Note, the convert layout
+        pass doesn't currently guarantee the whole of the graph will
+        be converted to the chosen layout.
+
+    Returns
+    -------
+    tasks : list of autotvm.Tasks
+        list of tasks to be tuned
+    weights : List[int]
+        the weight (i.e. the number of appearance) of extracted tasks
+    """
+    if alter_layout:
+        mod = common.convert_graph_layout(mod, alter_layout)
+
+    # Extract the tasks
+    tasks, task_weights = auto_scheduler.extract_tasks(
+        mod["main"], params, target=target, target_host=target_host, hardware_params=hardware_params
+    )
+
+    return tasks, task_weights
+
+
+def schedule_tasks(
+    tasks,
+    task_weights,
+    tuning_options,
+    tuning_records=None,
+):
+    """Generate the schedules for the different tasks (i.e., subgraphs) contained in the module.
+    Store the schedules in a json file that will be used later by the compiler.
+
+    Parameters
+    ----------
+    tasks : list
+        A list of autotvm.Tasks to tune.
+    task_weights : list
+        The weight (i.e. the number of appearance) of extracted tasks
+    tuning_records : str, optional
+        The json file used to preload the autoscheduler
+    tuning_options:
+        The options of tuning
+    """
+
+    # Create the scheduler
+    tuner = auto_scheduler.TaskScheduler(tasks, task_weights, load_log_file=tuning_records)
+
+    # Tune the tasks
+    tuner.tune(tuning_options)
diff --git a/python/tvm/driver/tvmc/autotuner.py b/python/tvm/driver/tvmc/autotuner.py
index 53c8f3bdc43d..dcbc63778973 100644
--- a/python/tvm/driver/tvmc/autotuner.py
+++ b/python/tvm/driver/tvmc/autotuner.py
@@ -30,7 +30,7 @@
 from tvm.autotvm.tuner import XGBTuner
 
 from . import common, frontends
-from .common import TVMCException
+from .common import TVMCException, add_tuning_options
 from .main import register_parser
 
 
@@ -44,99 +44,18 @@ def add_tune_parser(subparsers):
 
     parser = subparsers.add_parser("tune", help="auto-tune a model")
     parser.set_defaults(func=drive_tune)
-    parser.add_argument(
-        "--early-stopping",
-        type=int,
-        help="minimum number of trials before early stopping",
-    )
-
-    # There is some extra processing required to define the actual default value
-    # for --min-repeat-ms. This is done in `drive_tune`.
-    parser.add_argument(
-        "--min-repeat-ms",
-        default=None,
-        type=int,
-        help="minimum time to run each trial, in milliseconds. "
-        "Defaults to 0 on x86 and 1000 on all other targets",
-    )
+    add_tuning_options(parser)
     parser.add_argument(
         "--model-format",
         choices=frontends.get_frontend_names(),
         help="specify input model format",
     )
-    parser.add_argument(
-        "--number",
-        default=10,
-        type=int,
-        help="number of runs a single repeat is made of. "
-        "The final number of tuning executions is: "
-        "(1 + number * repeat)",
-    )
-    parser.add_argument(
-        "-o",
-        "--output",
-        required=True,
-        help="output file to store the tuning records for the tuning process",
-    )
-    parser.add_argument(
-        "--parallel",
-        default=4,
-        type=int,
-        help="the maximum number of parallel devices to use when tuning",
-    )
-    parser.add_argument(
-        "--repeat",
-        type=int,
-        default=1,
-        help="how many times to repeat each measurement",
-    )
-    parser.add_argument(
-        "--rpc-key",
-        help="the RPC tracker key of the target device. Required when --rpc-tracker is provided.",
-    )
-    parser.add_argument(
-        "--rpc-tracker",
-        help="hostname (required) and port (optional, defaults to 9090) of the RPC tracker, "
-        "e.g. '192.168.0.100:9999'",
-    )
-    parser.add_argument(
-        "--target",
-        help="compilation target as plain string, inline JSON or path to a JSON file",
-        required=True,
-    )
-    parser.add_argument(
-        "--target-host",
-        help="the host compilation target, defaults to 'llvm'",
-        default="llvm",
-    )
-    parser.add_argument("--timeout", default=10, help="compilation timeout, in seconds")
-    parser.add_argument(
-        "--trials",
-        type=int,
-        default=1000,
-        help="the maximum number of tuning trials to perform",
-    )
     parser.add_argument(
         "--tuner",
         choices=["ga", "gridsearch", "random", "xgb", "xgb_knob", "xgb-rank"],
         default="xgb",
         help="type of tuner to use",
     )
-    parser.add_argument(
-        "--tuning-records",
-        metavar="PATH",
-        help="path to an auto-tuning log file by AutoTVM.",
-    )
-    parser.add_argument(
-        "--desired-layout",
-        choices=["NCHW", "NHWC"],
-        default=None,
-        help="change the data layout of the whole graph",
-    )
-    # TODO (@leandron) This is a path to a physical file, but
-    #     can be improved in future to add integration with a modelzoo
-    #     or URL, for example.
-    parser.add_argument("FILE", help="path to the input model file")
 
 
 def drive_tune(args):
diff --git a/python/tvm/driver/tvmc/common.py b/python/tvm/driver/tvmc/common.py
index 9db22f3f3390..afcf3afe0ce0 100644
--- a/python/tvm/driver/tvmc/common.py
+++ b/python/tvm/driver/tvmc/common.py
@@ -36,6 +36,93 @@ class TVMCException(Exception):
     """TVMC Exception"""
 
 
+def add_tuning_options(parser):
+    """ Add common tuning options for the autotuner/autoscheduler."""
+
+    parser.add_argument(
+        "--early-stopping",
+        type=int,
+        help="minimum number of trials before early stopping",
+    )
+
+    # There is some extra processing required to define the actual default value
+    # for --min-repeat-ms. This is done in `drive_tune`.
+    parser.add_argument(
+        "--min-repeat-ms",
+        default=None,
+        type=int,
+        help="minimum time to run each trial, in milliseconds. "
+        "Defaults to 0 on x86 and 1000 on all other targets",
+    )
+    parser.add_argument(
+        "--number",
+        default=10,
+        type=int,
+        help="number of runs a single repeat is made of. "
+        "The final number of tuning executions is: "
+        "(1 + number * repeat)",
+    )
+    parser.add_argument(
+        "-o",
+        "--output",
+        required=True,
+        help="output file to store the tuning records for the tuning process",
+    )
+    parser.add_argument(
+        "--parallel",
+        default=4,
+        type=int,
+        help="the maximum number of parallel devices to use when tuning",
+    )
+    parser.add_argument(
+        "--repeat",
+        type=int,
+        default=1,
+        help="how many times to repeat each measurement",
+    )
+    parser.add_argument(
+        "--rpc-key",
+        help="the RPC tracker key of the target device. Required when --rpc-tracker is provided.",
+    )
+    parser.add_argument(
+        "--rpc-tracker",
+        help="hostname (required) and port (optional, defaults to 9090) of the RPC tracker, "
+        "e.g. '192.168.0.100:9999'",
+    )
+    parser.add_argument(
+        "--target",
+        help="compilation target as plain string, inline JSON or path to a JSON file",
+        required=True,
+    )
+    parser.add_argument(
+        "--target-host",
+        help="the host compilation target, defaults to 'llvm'",
+        default="llvm",
+    )
+    parser.add_argument("--timeout", default=10, help="compilation timeout, in seconds")
+    parser.add_argument(
+        "--trials",
+        type=int,
+        default=1000,
+        help="the maximum number of tuning trials to perform",
+    )
+    parser.add_argument(
+        "--tuning-records",
+        metavar="PATH",
+        help="path to an auto-tuning json file produced by the autoscheduler",
+    )
+    parser.add_argument(
+        "--desired-layout",
+        choices=["NCHW", "NHWC"],
+        default=None,
+        help="change the data layout of the whole graph",
+    )
+    # TODO (@leandron) This is a path to a physical file, but
+    #     can be improved in future to add integration with a modelzoo
+    #     or URL, for example.
+    parser.add_argument("FILE", help="path to the input model file")
+
+
 def convert_graph_layout(mod, desired_layout):
     """Alter the layout of the input graph.
 
diff --git a/python/tvm/driver/tvmc/compiler.py b/python/tvm/driver/tvmc/compiler.py
index 57071476b073..5bd9959fc539 100644
--- a/python/tvm/driver/tvmc/compiler.py
+++ b/python/tvm/driver/tvmc/compiler.py
@@ -23,7 +23,7 @@
 from pathlib import Path
 
 import tvm
-from tvm import autotvm
+from tvm import autotvm, auto_scheduler
 from tvm import relay
 from tvm.contrib import cc
 from tvm.contrib import utils
@@ -82,6 +82,11 @@ def add_compile_parser(subparsers):
         help="path to an auto-tuning log file by AutoTVM. If not presented, "
         "the fallback/tophub configs will be used",
     )
+    parser.add_argument(
+        "--use-autoscheduler",
+        action="store_true",
+        help="use the autoscheduler to generate the compute schedules",
+    )
     parser.add_argument("-v", "--verbose", action="count", default=0, help="increase verbosity")
     # TODO (@leandron) This is a path to a physical file, but
     #     can be improved in future to add integration with a modelzoo
@@ -111,6 +116,7 @@ def drive_compile(args):
         None,
         args.model_format,
         args.tuning_records,
+        args.use_autoscheduler,
         args.desired_layout,
     )
 
@@ -128,6 +134,7 @@ def compile_model(
     target_host=None,
     model_format=None,
     tuning_records=None,
+    use_autoscheduler=False,
     alter_layout=None,
 ):
     """Compile a model from a supported framework into a TVM module.
@@ -182,10 +189,23 @@ def compile_model(
 
     if tuning_records and os.path.exists(tuning_records):
         logger.debug("tuning records file provided: %s", tuning_records)
-        with autotvm.apply_history_best(tuning_records):
-            with tvm.transform.PassContext(opt_level=3):
-                logger.debug("building relay graph with tuning records")
-                graph_module = relay.build(mod, tvm_target, params=params, target_host=target_host)
+
+        if use_autoscheduler:
+            with auto_scheduler.ApplyHistoryBest(tuning_records):
+                with tvm.transform.PassContext(
+                    opt_level=3, config={"relay.backend.use_auto_scheduler": True}
+                ):
+                    logger.debug("building relay graph with autoscheduler")
+                    graph_module = relay.build(
+                        mod, target=target, params=params, target_host=target_host
+                    )
+        else:
+            with autotvm.apply_history_best(tuning_records):
+                with tvm.transform.PassContext(opt_level=3):
+                    logger.debug("building relay graph with tuning records")
+                    graph_module = relay.build(
+                        mod, tvm_target, params=params, target_host=target_host
+                    )
     else:
         with tvm.transform.PassContext(opt_level=3):
             logger.debug("building relay graph (no tuning records provided)")
diff --git a/tests/python/driver/tvmc/test_autoscheduler.py b/tests/python/driver/tvmc/test_autoscheduler.py
new file mode 100644
index 000000000000..4f95c07ee177
--- /dev/null
+++ b/tests/python/driver/tvmc/test_autoscheduler.py
@@ -0,0 +1,101 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+import json
+import pytest
+import os
+import tarfile
+
+from os import path
+
+from tvm import auto_scheduler
+from tvm.driver import tvmc
+
+
+def _get_tasks(model):
+    mod, params = tvmc.frontends.load_model(model)
+    tasks, weights = tvmc.autoscheduler.get_tuning_tasks(mod, params, "llvm")
+    return (tasks, weights)
+
+
+def _autoscheduler_test_helper(
+    model, tmpdir_name, tasks_weights=None, early_stopping=1, tuning_records=None
+):
+    tasks, weights = tasks_weights if tasks_weights else _get_tasks(model)
+    log_file = os.path.join(tmpdir_name, "autoscheduler.json")
+
+    tuning_options = auto_scheduler.TuningOptions(
+        num_measure_trials=1,
+        measure_callbacks=[auto_scheduler.RecordToFile(log_file)],
+        runner="local",
+        builder="local",
+        verbose=0,
+        early_stopping=early_stopping,
+    )
+
+    tvmc.autoscheduler.schedule_tasks(tasks[:1], weights[:1], tuning_options, tuning_records)
+
+    # testing whether the log file was produced
+    assert path.exists(log_file), "autoscheduler log file should exist"
+
+    with auto_scheduler.ApplyHistoryBest(log_file) as best:
+        assert isinstance(
+            best, auto_scheduler.dispatcher.ApplyHistoryBest
+        ), "unable to load the best results of tuning"
+
+    return log_file
+
+
+def test_get_tuning_tasks(onnx_resnet50):
+    pytest.importorskip("onnx")
+
+    tasks, weights = _get_tasks(onnx_resnet50)
+    expected_task_type = auto_scheduler.SearchTask
+
+    assert type(tasks) is list
+    assert len(tasks) > 0
+    assert all([type(x) is expected_task_type for x in tasks]) is True
+
+
+def test_tune_tasks(onnx_resnet50, tmpdir_factory):
+    pytest.importorskip("onnx")
+
+    tmpdir_name = tmpdir_factory.mktemp("data")
+    _autoscheduler_test_helper(onnx_resnet50, tmpdir_name)
+
+
+def test_tune_tasks__tuning_records(onnx_resnet50, tmpdir_factory):
+    pytest.importorskip("onnx")
+
+    tmpdir_name = tmpdir_factory.mktemp("data")
+    output_log_phase_1 = _autoscheduler_test_helper(onnx_resnet50, tmpdir_name)
+
+    # Exercises transfer learning by making sure a previous log exists
+    _autoscheduler_test_helper(onnx_resnet50, tmpdir_name, tuning_records=output_log_phase_1)
+
+
+def test_tune_tasks__no_early_stopping(onnx_resnet50, tmpdir_factory):
+    pytest.importorskip("onnx")
+
+    tmpdir_name = tmpdir_factory.mktemp("data")
+    _autoscheduler_test_helper(onnx_resnet50, tmpdir_name, tasks_weights=None, early_stopping=None)
+
+
+def test_tune_tasks__no_tuning_records(onnx_resnet50, tmpdir_factory):
+    pytest.importorskip("onnx")
+
+    tmpdir_name = tmpdir_factory.mktemp("data")
+    _autoscheduler_test_helper(onnx_resnet50, tmpdir_name, tasks_weights=None, tuning_records=None)

From 0be86efb0e6778b1eae15b4ca1bba8a1c8d800b3 Mon Sep 17 00:00:00 2001
From: Giuseppe Rossini <giuseppe.rossini@arm.com>
Date: Wed, 9 Dec 2020 16:38:29 +0000
Subject: [PATCH 2/8] Linting and small bug-fixing

---
 python/tvm/driver/tvmc/autoscheduler.py | 25 ++++++++++---------------
 1 file changed, 10 insertions(+), 15 deletions(-)

diff --git a/python/tvm/driver/tvmc/autoscheduler.py b/python/tvm/driver/tvmc/autoscheduler.py
index d8772f1ed54d..31cdeaa48797 100644
--- a/python/tvm/driver/tvmc/autoscheduler.py
+++ b/python/tvm/driver/tvmc/autoscheduler.py
@@ -14,14 +14,10 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
-from docutils.nodes import target
-
 """
-Provides support to auto-tuning networks using AutoTVM.
+Provides support to auto-tuning networks using AutoScheduler.
 """
-import os.path
 import logging
-import time
 
 from urllib.parse import urlparse
 
@@ -29,7 +25,7 @@
 from tvm.auto_scheduler.auto_schedule import HardwareParams
 
 from . import common, frontends
-from .common import TVMCException, add_tuning_options
+from .common import add_tuning_options
 from .main import register_parser
 
 
@@ -67,7 +63,7 @@ def add_autoscheduler_parser(subparsers):
 
 
 def drive_autoschedule(args):
-    """Invoke auto-tuning with command line arguments
+    """Invoke auto-scheduling with command line arguments
 
     Parameters
     ----------
@@ -132,7 +128,9 @@ def drive_autoschedule(args):
     )
 
     # Specify hardware parameters
-    hardware_params = HardwareParams(args.num_cores, args.vector_unit_bytes, args.cache_line_bytes)
+    hardware_params = HardwareParams(
+        args.num_cores, args.vector_unit_bytes, args.cache_line_bytes, None, None, None, None, None
+    )
 
     # Extract the tasks from the model
     tasks, weights = get_tuning_tasks(
@@ -141,13 +139,10 @@ def drive_autoschedule(args):
 
     # Schedule the tasks (i.e., produce a schedule for each task)
     schedule_tasks(
-        mod,
-        params,
-        target,
-        target_host=args.target_host,
-        tuning_records=args.tuning_records,
-        tuning_options=tuning_options,
-        alter_layout=args.desired_layout,
+        tasks,
+        weights,
+        tuning_options,
+        args.tuning_records,
     )
 
 

From 88b4a1d5f78294ae684eabec4b805b666c854888 Mon Sep 17 00:00:00 2001
From: Giuseppe Rossini <giuseppe.rossini@arm.com>
Date: Thu, 10 Dec 2020 18:05:46 +0000
Subject: [PATCH 3/8] Addressing comments and refactoring

---
 python/tvm/driver/tvmc/__init__.py            |   1 -
 python/tvm/driver/tvmc/autoscheduler.py       | 212 --------------
 python/tvm/driver/tvmc/autotuner.py           | 269 ++++++++++++++++--
 python/tvm/driver/tvmc/compiler.py            |  14 +-
 .../python/driver/tvmc/test_autoscheduler.py  | 101 -------
 5 files changed, 249 insertions(+), 348 deletions(-)
 delete mode 100644 python/tvm/driver/tvmc/autoscheduler.py
 delete mode 100644 tests/python/driver/tvmc/test_autoscheduler.py

diff --git a/python/tvm/driver/tvmc/__init__.py b/python/tvm/driver/tvmc/__init__.py
index 37344af42d42..d96a725877eb 100644
--- a/python/tvm/driver/tvmc/__init__.py
+++ b/python/tvm/driver/tvmc/__init__.py
@@ -21,4 +21,3 @@
 from . import autotuner
 from . import compiler
 from . import runner
-from . import autoscheduler
diff --git a/python/tvm/driver/tvmc/autoscheduler.py b/python/tvm/driver/tvmc/autoscheduler.py
deleted file mode 100644
index 31cdeaa48797..000000000000
--- a/python/tvm/driver/tvmc/autoscheduler.py
+++ /dev/null
@@ -1,212 +0,0 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements.  See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership.  The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License.  You may obtain a copy of the License at
-#
-#   http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied.  See the License for the
-# specific language governing permissions and limitations
-# under the License.
-"""
-Provides support to auto-tuning networks using AutoScheduler.
-"""
-import logging
-
-from urllib.parse import urlparse
-
-from tvm import auto_scheduler
-from tvm.auto_scheduler.auto_schedule import HardwareParams
-
-from . import common, frontends
-from .common import add_tuning_options
-from .main import register_parser
-
-
-# pylint: disable=invalid-name
-logger = logging.getLogger("TVMC")
-
-
-@register_parser
-def add_autoscheduler_parser(subparsers):
-    """ Include parser for 'autoschedule' subcommand """
-    parser = subparsers.add_parser("autoschedule", help="auto-schedule a model")
-    parser.set_defaults(func=drive_autoschedule)
-    add_tuning_options(parser)
-
-    parser.add_argument(
-        "--cache-line-bytes",
-        default=64,
-        help="the size of cache line in bytes",
-    )
-    parser.add_argument(
-        "--num-cores",
-        default=4,
-        help="the number of device cores",
-    )
-    parser.add_argument(
-        "--vector-unit-bytes",
-        default=16,
-        help="the width of vector units in bytes",
-    )
-    parser.add_argument(
-        "--model-format",
-        choices=frontends.get_frontend_names(),
-        help="specify input model format",
-    )
-
-
-def drive_autoschedule(args):
-    """Invoke auto-scheduling with command line arguments
-
-    Parameters
-    ----------
-    args: argparse.Namespace
-        Arguments from command line parser.
-    """
-
-    # extra arguments validation before importing the model, so that obvious errors
-    # are pointed in advance.
-    if args.rpc_tracker:
-        parsed_url = urlparse("//%s" % args.rpc_tracker)
-        rpc_hostname = parsed_url.hostname
-        rpc_port = parsed_url.port or 9090
-        logger.info("RPC tracker hostname: %s", rpc_hostname)
-        logger.info("RPC tracker port: %s", rpc_port)
-
-        if not args.rpc_key:
-            raise common.TVMCException(
-                "need to provide an RPC tracker key (--rpc-key) for remote tuning"
-            )
-
-    target = common.target_from_cli(args.target)
-    mod, params = frontends.load_model(args.FILE, args.model_format)
-
-    # min_repeat_ms should be:
-    # a. the value provided by the user, if any, or
-    # b. 0ms in case target is "cpu"; otherwise 1000ms
-    if args.min_repeat_ms is not None:
-        min_repeat_ms = args.min_repeat_ms
-    else:
-        min_repeat_ms = 0 if target.keys[0] == "cpu" else 1000
-        logger.debug("Default --min-repeat-ms for this target is %s", min_repeat_ms)
-
-    if args.rpc_tracker:
-
-        runner = auto_scheduler.RPCRunner(
-            key=args.rpc_key,
-            host=rpc_hostname,
-            port=rpc_port,
-            number=args.number,
-            repeat=args.repeat,
-            n_parallel=args.parallel,
-            timeout=args.timeout,
-            min_repeat_ms=min_repeat_ms,
-        )
-    else:
-        logger.info("starting localhost tuning")
-        runner = auto_scheduler.LocalRunner(
-            number=args.number,
-            repeat=args.repeat,
-            timeout=args.timeout,
-            min_repeat_ms=min_repeat_ms,
-        )
-
-    # Create the autoscheduler tuning options
-    tuning_options = auto_scheduler.TuningOptions(
-        num_measure_trials=args.trials,
-        measure_callbacks=[auto_scheduler.RecordToFile(args.output)],
-        runner=runner,
-        builder="local",
-        early_stopping=args.early_stopping,
-    )
-
-    # Specify hardware parameters
-    hardware_params = HardwareParams(
-        args.num_cores, args.vector_unit_bytes, args.cache_line_bytes, None, None, None, None, None
-    )
-
-    # Extract the tasks from the model
-    tasks, weights = get_tuning_tasks(
-        mod, params, target, target_host, args.desired_layout, hardware_params
-    )
-
-    # Schedule the tasks (i.e., produce a schedule for each task)
-    schedule_tasks(
-        tasks,
-        weights,
-        tuning_options,
-        args.tuning_records,
-    )
-
-
-def get_tuning_tasks(
-    mod, params, target, target_host=None, alter_layout=None, hardware_params=None
-):
-    """Get the tuning tasks for a given relay module.
-
-    Parameters
-    ----------
-    mod : tvm.relay.Module
-        The relay module from which to extract tuning tasks.
-    params : dict
-        The params for the relay module.
-    target : tvm.target.Target
-        The compilation target.
-    target_host : str, optional
-        The compilation target for the host.
-    alter_layout : str, optional
-        The layout to convert the graph to. Note, the convert layout
-        pass doesn't currently guarantee the whole of the graph will
-        be converted to the chosen layout.
-
-    Returns
-    -------
-    tasks : list of autotvm.Tasks
-        list of tasks to be tuned
-    weights : List[int]
-        the weight (i.e. the number of appearance) of extracted tasks
-    """
-    if alter_layout:
-        mod = common.convert_graph_layout(mod, alter_layout)
-
-    # Extract the tasks
-    tasks, task_weights = auto_scheduler.extract_tasks(
-        mod["main"], params, target=target, target_host=target_host, hardware_params=hardware_params
-    )
-
-    return tasks, task_weights
-
-
-def schedule_tasks(
-    tasks,
-    task_weights,
-    tuning_options,
-    tuning_records=None,
-):
-    """Generate the schedules for the different tasks (i.e., subgraphs) contained in the module.
-    Store the schedules in a json file that will be used later by the compiler.
-
-    Parameters
-    ----------
-    tasks : list
-        A list of autotvm.Tasks to tune.
-    task_weights : list
-        The weight (i.e. the number of appearance) of extracted tasks
-    tuning_records : str, optional
-        The json file used to preload the autoscheduler
-    tuning_options:
-        The options of tuning
-    """
-
-    # Create the scheduler
-    tuner = auto_scheduler.TaskScheduler(tasks, task_weights, load_log_file=tuning_records)
-
-    # Tune the tasks
-    tuner.tune(tuning_options)
diff --git a/python/tvm/driver/tvmc/autotuner.py b/python/tvm/driver/tvmc/autotuner.py
index dcbc63778973..1a63503549da 100644
--- a/python/tvm/driver/tvmc/autotuner.py
+++ b/python/tvm/driver/tvmc/autotuner.py
@@ -23,14 +23,14 @@
 
 from urllib.parse import urlparse
 
-from tvm import autotvm
+from tvm import autotvm, auto_scheduler
 from tvm.autotvm.tuner import GATuner
 from tvm.autotvm.tuner import GridSearchTuner
 from tvm.autotvm.tuner import RandomTuner
 from tvm.autotvm.tuner import XGBTuner
 
 from . import common, frontends
-from .common import TVMCException, add_tuning_options
+from .common import TVMCException
 from .main import register_parser
 
 
@@ -44,18 +44,132 @@ def add_tune_parser(subparsers):
 
     parser = subparsers.add_parser("tune", help="auto-tune a model")
     parser.set_defaults(func=drive_tune)
-    add_tuning_options(parser)
+    parser.add_argument(
+        "--early-stopping",
+        type=int,
+        help="minimum number of trials before early stopping",
+    )
+
+    # There is some extra processing required to define the actual default value
+    # for --min-repeat-ms. This is done in `drive_tune`.
+    parser.add_argument(
+        "--min-repeat-ms",
+        default=None,
+        type=int,
+        help="minimum time to run each trial, in milliseconds. "
+        "Defaults to 0 on x86 and 1000 on all other targets",
+    )
     parser.add_argument(
         "--model-format",
         choices=frontends.get_frontend_names(),
         help="specify input model format",
     )
     parser.add_argument(
+        "--number",
+        default=10,
+        type=int,
+        help="number of runs a single repeat is made of. "
+        "The final number of tuning executions is: "
+        "(1 + number * repeat)",
+    )
+    parser.add_argument(
+        "-o",
+        "--output",
+        required=True,
+        help="output file to store the tuning records for the tuning process",
+    )
+    parser.add_argument(
+        "--parallel",
+        default=4,
+        type=int,
+        help="the maximum number of parallel devices to use when tuning",
+    )
+    parser.add_argument(
+        "--repeat",
+        type=int,
+        default=1,
+        help="how many times to repeat each measurement",
+    )
+    parser.add_argument(
+        "--rpc-key",
+        help="the RPC tracker key of the target device. Required when --rpc-tracker is provided.",
+    )
+    parser.add_argument(
+        "--rpc-tracker",
+        help="hostname (required) and port (optional, defaults to 9090) of the RPC tracker, "
+        "e.g. '192.168.0.100:9999'",
+    )
+    parser.add_argument(
+        "--target",
+        help="compilation target as plain string, inline JSON or path to a JSON file",
+        required=True,
+    )
+    parser.add_argument(
+        "--target-host",
+        help="the host compilation target, defaults to 'llvm'",
+        default="llvm",
+    )
+    parser.add_argument("--timeout", default=10, help="compilation timeout, in seconds")
+    parser.add_argument(
+        "--trials",
+        type=int,
+        default=1000,
+        help="the maximum number of tuning trials to perform",
+    )
+    parser.add_argument(
+        "--tuning-records",
+        metavar="PATH",
+        help="path to an auto-tuning log file by AutoTVM.",
+    )
+    parser.add_argument(
+        "--desired-layout",
+        choices=["NCHW", "NHWC"],
+        default=None,
+        help="change the data layout of the whole graph",
+    )
+    parser.add_argument(
+        "--enable-autoscheduler",
+        help="enable tuning the graph through the autoscheduler",
+        action="store_true",
+    )
+
+    auto_scheduler_group = parser.add_argument_group(
+        "Autoscheduler options",
+        "Autoscheduler options, used when --enabled-auto-scheduler is provided",
+    )
+
+    auto_scheduler_group.add_argument(
+        "--cache-line-bytes",
+        type=int,
+        default=64,
+        help="the size of cache line in bytes",
+    )
+    auto_scheduler_group.add_argument(
+        "--num-cores",
+        type=int,
+        default=4,
+        help="the number of device cores",
+    )
+    auto_scheduler_group.add_argument(
+        "--vector-unit-bytes",
+        type=int,
+        default=16,
+        help="the width of vector units in bytes",
+    )
+    auto_tuning_group = parser.add_argument_group(
+        "Autotuning options",
+        "Autotuning options, used when the autoscheduler is not enabled",
+    )
+    auto_tuning_group.add_argument(
         "--tuner",
         choices=["ga", "gridsearch", "random", "xgb", "xgb_knob", "xgb-rank"],
         default="xgb",
-        help="type of tuner to use",
+        help="type of tuner to use when autotuning. If --enable-autoscheduler is set, this option won't have any effect",
     )
+    # TODO (@leandron) This is a path to a physical file, but
+    #     can be improved in future to add integration with a modelzoo
+    #     or URL, for example.
+    parser.add_argument("FILE", help="path to the input model file")
 
 
 def drive_tune(args):
@@ -66,7 +180,6 @@ def drive_tune(args):
     args: argparse.Namespace
         Arguments from command line parser.
     """
-
     # extra arguments validation before importing the model, so that obvious errors
     # are pointed in advance.
     if args.rpc_tracker:
@@ -93,17 +206,26 @@ def drive_tune(args):
         min_repeat_ms = 0 if target.keys[0] == "cpu" else 1000
         logger.debug("Default --min-repeat-ms for this target is %s", min_repeat_ms)
 
-    tasks = get_tuning_tasks(
-        mod=mod,
-        params=params,
-        target=target,
-        target_host=args.target_host,
-        alter_layout=args.desired_layout,
-    )
+    if args.enable_autoscheduler:
+        tasks, weights = autoscheduler_get_tuning_tasks(
+            mod=mod,
+            params=params,
+            target=target,
+            target_host=args.target_host,
+            alter_layout=args.desired_layout,
+        )
+    else:
+        tasks = autotuner_get_tuning_tasks(
+            mod=mod,
+            params=params,
+            target=target,
+            target_host=args.target_host,
+            alter_layout=args.desired_layout,
+        )
 
     if args.rpc_tracker:
-
-        runner = autotvm.RPCRunner(
+        runner_ctor = auto_scheduler.RPCRunner if args.enable_autoscheduler else autotvm.RPCRunner
+        runner = runner_ctor(
             key=args.rpc_key,
             host=rpc_hostname,
             port=rpc_port,
@@ -115,28 +237,55 @@ def drive_tune(args):
         )
     else:
         logger.info("starting localhost tuning")
-        runner = autotvm.LocalRunner(
+        runner_ctor = (
+            auto_scheduler.LocalRunner if args.enable_autoscheduler else autotvm.LocalRunner
+        )
+        runner = runner_ctor(
             number=args.number,
             repeat=args.repeat,
             timeout=args.timeout,
             min_repeat_ms=min_repeat_ms,
         )
 
-    tuning_option = {
-        "tuner": args.tuner,
-        "trials": args.trials,
-        "early_stopping": args.early_stopping,
-        "measure_option": autotvm.measure_option(
-            builder=autotvm.LocalBuilder(build_func="default"), runner=runner
-        ),
-        "tuning_records": args.tuning_records,
-    }
-    logger.debug(" tuning options: %s", tuning_option)
+    if args.enable_autoscheduler:
+        # Create the autoscheduler tuning options
+        tuning_options = auto_scheduler.TuningOptions(
+            num_measure_trials=args.trials,
+            measure_callbacks=[auto_scheduler.RecordToFile(args.output)],
+            runner=runner,
+            early_stopping=args.early_stopping,
+        )
 
-    tune_tasks(tasks, args.output, **tuning_option)
+        # Specify hardware parameters
+        print(type(args.vector_unit_bytes))
+        hardware_params = auto_scheduler.HardwareParams(
+            args.num_cores, args.vector_unit_bytes, args.cache_line_bytes, 0, 0, 0, 0, 0
+        )
 
+        # Schedule the tasks (i.e., produce a schedule for each task)
+        schedule_tasks(
+            tasks,
+            weights,
+            tuning_options,
+            args.tuning_records,
+        )
 
-def get_tuning_tasks(mod, params, target, target_host=None, alter_layout=None):
+    else:
+        tuning_option = {
+            "tuner": args.tuner,
+            "trials": args.trials,
+            "early_stopping": args.early_stopping,
+            "measure_option": autotvm.measure_option(
+                builder=autotvm.LocalBuilder(build_func="default"), runner=runner
+            ),
+            "tuning_records": args.tuning_records,
+        }
+        logger.debug(" tuning options: %s", tuning_option)
+
+        tune_tasks(tasks, args.output, **tuning_option)
+
+
+def autotuner_get_tuning_tasks(mod, params, target, target_host=None, alter_layout=None):
     """Get the tuning tasks for a given relay module.
 
     Parameters
@@ -172,6 +321,72 @@ def get_tuning_tasks(mod, params, target, target_host=None, alter_layout=None):
     return tasks
 
 
+def autoscheduler_get_tuning_tasks(
+    mod, params, target, target_host=None, alter_layout=None, hardware_params=None
+):
+    """Get the tuning tasks for a given relay module.
+
+    Parameters
+    ----------
+    mod : tvm.relay.Module
+        The relay module from which to extract tuning tasks.
+    params : dict
+        The params for the relay module.
+    target : tvm.target.Target
+        The compilation target.
+    target_host : str, optional
+        The compilation target for the host.
+    alter_layout : str, optional
+        The layout to convert the graph to. Note, the convert layout
+        pass doesn't currently guarantee the whole of the graph will
+        be converted to the chosen layout.
+
+    Returns
+    -------
+    tasks : list of autotvm.Tasks
+        list of tasks to be tuned
+    weights : List[int]
+        the weight (i.e. the number of appearance) of extracted tasks
+    """
+    if alter_layout:
+        mod = common.convert_graph_layout(mod, alter_layout)
+
+    # Extract the tasks
+    tasks, task_weights = auto_scheduler.extract_tasks(
+        mod["main"], params, target=target, target_host=target_host, hardware_params=hardware_params
+    )
+
+    return tasks, task_weights
+
+
+def schedule_tasks(
+    tasks,
+    task_weights,
+    tuning_options,
+    tuning_records=None,
+):
+    """Generate the schedules for the different tasks (i.e., subgraphs) contained in the module.
+    Store the schedules in a json file that will be used later by the compiler.
+
+    Parameters
+    ----------
+    tasks : list
+        A list of auto_scheduler.SearchTask to tune.
+    task_weights : list
+        The weight (i.e. the number of appearance) of extracted tasks
+    tuning_options:
+        The options of tuning
+    tuning_records : str, optional
+        The json file used to preload the autoscheduler
+    """
+
+    # Create the scheduler
+    tuner = auto_scheduler.TaskScheduler(tasks, task_weights, load_log_file=tuning_records)
+
+    # Tune the tasks
+    tuner.tune(tuning_options)
+
+
 def tune_tasks(
     tasks,
     log_file,
diff --git a/python/tvm/driver/tvmc/compiler.py b/python/tvm/driver/tvmc/compiler.py
index 5bd9959fc539..b1d3655b7738 100644
--- a/python/tvm/driver/tvmc/compiler.py
+++ b/python/tvm/driver/tvmc/compiler.py
@@ -82,11 +82,6 @@ def add_compile_parser(subparsers):
         help="path to an auto-tuning log file by AutoTVM. If not presented, "
         "the fallback/tophub configs will be used",
     )
-    parser.add_argument(
-        "--use-autoscheduler",
-        action="store_true",
-        help="use the autoscheduler to generate the compute schedules",
-    )
     parser.add_argument("-v", "--verbose", action="count", default=0, help="increase verbosity")
     # TODO (@leandron) This is a path to a physical file, but
     #     can be improved in future to add integration with a modelzoo
@@ -116,7 +111,6 @@ def drive_compile(args):
         None,
         args.model_format,
         args.tuning_records,
-        args.use_autoscheduler,
         args.desired_layout,
     )
 
@@ -134,7 +128,6 @@ def compile_model(
     target_host=None,
     model_format=None,
     tuning_records=None,
-    use_autoscheduler=False,
     alter_layout=None,
 ):
     """Compile a model from a supported framework into a TVM module.
@@ -190,6 +183,13 @@ def compile_model(
     if tuning_records and os.path.exists(tuning_records):
         logger.debug("tuning records file provided: %s", tuning_records)
 
+        use_autoscheduler = True
+
+        try:
+            auto_scheduler.load_records(tuning_records)
+        except Exception:
+            use_autoscheduler = False
+
         if use_autoscheduler:
             with auto_scheduler.ApplyHistoryBest(tuning_records):
                 with tvm.transform.PassContext(
diff --git a/tests/python/driver/tvmc/test_autoscheduler.py b/tests/python/driver/tvmc/test_autoscheduler.py
deleted file mode 100644
index 4f95c07ee177..000000000000
--- a/tests/python/driver/tvmc/test_autoscheduler.py
+++ /dev/null
@@ -1,101 +0,0 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements.  See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership.  The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License.  You may obtain a copy of the License at
-#
-#   http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied.  See the License for the
-# specific language governing permissions and limitations
-# under the License.
-import json
-import pytest
-import os
-import tarfile
-
-from os import path
-
-from tvm import auto_scheduler
-from tvm.driver import tvmc
-
-
-def _get_tasks(model):
-    mod, params = tvmc.frontends.load_model(model)
-    tasks, weights = tvmc.autoscheduler.get_tuning_tasks(mod, params, "llvm")
-    return (tasks, weights)
-
-
-def _autoscheduler_test_helper(
-    model, tmpdir_name, tasks_weights=None, early_stopping=1, tuning_records=None
-):
-    tasks, weights = tasks_weights if tasks_weights else _get_tasks(model)
-    log_file = os.path.join(tmpdir_name, "autoscheduler.json")
-
-    tuning_options = auto_scheduler.TuningOptions(
-        num_measure_trials=1,
-        measure_callbacks=[auto_scheduler.RecordToFile(log_file)],
-        runner="local",
-        builder="local",
-        verbose=0,
-        early_stopping=early_stopping,
-    )
-
-    tvmc.autoscheduler.schedule_tasks(tasks[:1], weights[:1], tuning_options, tuning_records)
-
-    # testing whether the log file was produced
-    assert path.exists(log_file), "autoscheduler log file should exist"
-
-    with auto_scheduler.ApplyHistoryBest(log_file) as best:
-        assert isinstance(
-            best, auto_scheduler.dispatcher.ApplyHistoryBest
-        ), "unable to load the best results of tuning"
-
-    return log_file
-
-
-def test_get_tuning_tasks(onnx_resnet50):
-    pytest.importorskip("onnx")
-
-    tasks, weights = _get_tasks(onnx_resnet50)
-    expected_task_type = auto_scheduler.SearchTask
-
-    assert type(tasks) is list
-    assert len(tasks) > 0
-    assert all([type(x) is expected_task_type for x in tasks]) is True
-
-
-def test_tune_tasks(onnx_resnet50, tmpdir_factory):
-    pytest.importorskip("onnx")
-
-    tmpdir_name = tmpdir_factory.mktemp("data")
-    _autoscheduler_test_helper(onnx_resnet50, tmpdir_name)
-
-
-def test_tune_tasks__tuning_records(onnx_resnet50, tmpdir_factory):
-    pytest.importorskip("onnx")
-
-    tmpdir_name = tmpdir_factory.mktemp("data")
-    output_log_phase_1 = _autoscheduler_test_helper(onnx_resnet50, tmpdir_name)
-
-    # Exercises transfer learning by making sure a previous log exists
-    _autoscheduler_test_helper(onnx_resnet50, tmpdir_name, tuning_records=output_log_phase_1)
-
-
-def test_tune_tasks__no_early_stopping(onnx_resnet50, tmpdir_factory):
-    pytest.importorskip("onnx")
-
-    tmpdir_name = tmpdir_factory.mktemp("data")
-    _autoscheduler_test_helper(onnx_resnet50, tmpdir_name, tasks_weights=None, early_stopping=None)
-
-
-def test_tune_tasks__no_tuning_records(onnx_resnet50, tmpdir_factory):
-    pytest.importorskip("onnx")
-
-    tmpdir_name = tmpdir_factory.mktemp("data")
-    _autoscheduler_test_helper(onnx_resnet50, tmpdir_name, tasks_weights=None, tuning_records=None)

From def986504af90feb0c99dafdf12edb9a74f689d5 Mon Sep 17 00:00:00 2001
From: Giuseppe Rossini <giuseppe.rossini@arm.com>
Date: Thu, 10 Dec 2020 18:58:55 +0000
Subject: [PATCH 4/8] Fix linting

---
 python/tvm/driver/tvmc/autotuner.py | 13 ++++++-------
 python/tvm/driver/tvmc/compiler.py  |  2 +-
 2 files changed, 7 insertions(+), 8 deletions(-)

diff --git a/python/tvm/driver/tvmc/autotuner.py b/python/tvm/driver/tvmc/autotuner.py
index 1a63503549da..122d12036af9 100644
--- a/python/tvm/driver/tvmc/autotuner.py
+++ b/python/tvm/driver/tvmc/autotuner.py
@@ -164,7 +164,7 @@ def add_tune_parser(subparsers):
         "--tuner",
         choices=["ga", "gridsearch", "random", "xgb", "xgb_knob", "xgb-rank"],
         default="xgb",
-        help="type of tuner to use when autotuning. If --enable-autoscheduler is set, this option won't have any effect",
+        help="type of tuner to use when autotuning.",
     )
     # TODO (@leandron) This is a path to a physical file, but
     #     can be improved in future to add integration with a modelzoo
@@ -207,12 +207,17 @@ def drive_tune(args):
         logger.debug("Default --min-repeat-ms for this target is %s", min_repeat_ms)
 
     if args.enable_autoscheduler:
+        # Specify hardware parameters
+        hardware_params = auto_scheduler.HardwareParams(
+            args.num_cores, args.vector_unit_bytes, args.cache_line_bytes, 0, 0, 0, 0, 0
+        )
         tasks, weights = autoscheduler_get_tuning_tasks(
             mod=mod,
             params=params,
             target=target,
             target_host=args.target_host,
             alter_layout=args.desired_layout,
+            hardware_params=hardware_params,
         )
     else:
         tasks = autotuner_get_tuning_tasks(
@@ -256,12 +261,6 @@ def drive_tune(args):
             early_stopping=args.early_stopping,
         )
 
-        # Specify hardware parameters
-        print(type(args.vector_unit_bytes))
-        hardware_params = auto_scheduler.HardwareParams(
-            args.num_cores, args.vector_unit_bytes, args.cache_line_bytes, 0, 0, 0, 0, 0
-        )
-
         # Schedule the tasks (i.e., produce a schedule for each task)
         schedule_tasks(
             tasks,
diff --git a/python/tvm/driver/tvmc/compiler.py b/python/tvm/driver/tvmc/compiler.py
index b1d3655b7738..3dc043ab521e 100644
--- a/python/tvm/driver/tvmc/compiler.py
+++ b/python/tvm/driver/tvmc/compiler.py
@@ -187,7 +187,7 @@ def compile_model(
 
         try:
             auto_scheduler.load_records(tuning_records)
-        except Exception:
+        except tvm._ffi.base.TVMError:
             use_autoscheduler = False
 
         if use_autoscheduler:

From 8c5df5e7f345e9a1149e284b6a93ff4c37a98714 Mon Sep 17 00:00:00 2001
From: Giuseppe Rossini <giuseppe.rossini@arm.com>
Date: Thu, 10 Dec 2020 19:08:27 +0000
Subject: [PATCH 5/8] rebasing

---
 python/tvm/driver/tvmc/common.py | 87 --------------------------------
 1 file changed, 87 deletions(-)

diff --git a/python/tvm/driver/tvmc/common.py b/python/tvm/driver/tvmc/common.py
index afcf3afe0ce0..9db22f3f3390 100644
--- a/python/tvm/driver/tvmc/common.py
+++ b/python/tvm/driver/tvmc/common.py
@@ -36,93 +36,6 @@ class TVMCException(Exception):
     """TVMC Exception"""
 
 
-def add_tuning_options(parser):
-    """ Add common tuning options for the autotuner/autoscheduler."""
-
-    parser.add_argument(
-        "--early-stopping",
-        type=int,
-        help="minimum number of trials before early stopping",
-    )
-
-    # There is some extra processing required to define the actual default value
-    # for --min-repeat-ms. This is done in `drive_tune`.
-    parser.add_argument(
-        "--min-repeat-ms",
-        default=None,
-        type=int,
-        help="minimum time to run each trial, in milliseconds. "
-        "Defaults to 0 on x86 and 1000 on all other targets",
-    )
-    parser.add_argument(
-        "--number",
-        default=10,
-        type=int,
-        help="number of runs a single repeat is made of. "
-        "The final number of tuning executions is: "
-        "(1 + number * repeat)",
-    )
-    parser.add_argument(
-        "-o",
-        "--output",
-        required=True,
-        help="output file to store the tuning records for the tuning process",
-    )
-    parser.add_argument(
-        "--parallel",
-        default=4,
-        type=int,
-        help="the maximum number of parallel devices to use when tuning",
-    )
-    parser.add_argument(
-        "--repeat",
-        type=int,
-        default=1,
-        help="how many times to repeat each measurement",
-    )
-    parser.add_argument(
-        "--rpc-key",
-        help="the RPC tracker key of the target device. Required when --rpc-tracker is provided.",
-    )
-    parser.add_argument(
-        "--rpc-tracker",
-        help="hostname (required) and port (optional, defaults to 9090) of the RPC tracker, "
-        "e.g. '192.168.0.100:9999'",
-    )
-    parser.add_argument(
-        "--target",
-        help="compilation target as plain string, inline JSON or path to a JSON file",
-        required=True,
-    )
-    parser.add_argument(
-        "--target-host",
-        help="the host compilation target, defaults to 'llvm'",
-        default="llvm",
-    )
-    parser.add_argument("--timeout", default=10, help="compilation timeout, in seconds")
-    parser.add_argument(
-        "--trials",
-        type=int,
-        default=1000,
-        help="the maximum number of tuning trials to perform",
-    )
-    parser.add_argument(
-        "--tuning-records",
-        metavar="PATH",
-        help="path to an auto-tuning json file produced by the autoscheduler",
-    )
-    parser.add_argument(
-        "--desired-layout",
-        choices=["NCHW", "NHWC"],
-        default=None,
-        help="change the data layout of the whole graph",
-    )
-    # TODO (@leandron) This is a path to a physical file, but
-    #     can be improved in future to add integration with a modelzoo
-    #     or URL, for example.
-    parser.add_argument("FILE", help="path to the input model file")
-
-
 def convert_graph_layout(mod, desired_layout):
     """Alter the layout of the input graph.
 

From c59d4b7bbffcd5707600b183d3328b0a11cc2468 Mon Sep 17 00:00:00 2001
From: Giuseppe Rossini <giuseppe.rossini@arm.com>
Date: Thu, 10 Dec 2020 19:40:28 +0000
Subject: [PATCH 6/8] Addressing comments - 2

---
 python/tvm/driver/tvmc/autotuner.py        | 91 +++++++++++++++-------
 python/tvm/driver/tvmc/compiler.py         |  1 -
 tests/python/driver/tvmc/test_autotuner.py |  2 +-
 3 files changed, 65 insertions(+), 29 deletions(-)

diff --git a/python/tvm/driver/tvmc/autotuner.py b/python/tvm/driver/tvmc/autotuner.py
index 122d12036af9..37188261080e 100644
--- a/python/tvm/driver/tvmc/autotuner.py
+++ b/python/tvm/driver/tvmc/autotuner.py
@@ -156,6 +156,36 @@ def add_tune_parser(subparsers):
         default=16,
         help="the width of vector units in bytes",
     )
+    auto_scheduler_group.add_argument(
+        "--max-shared-memory-per-block",
+        type=int,
+        default=0,
+        help="the max shared memory per block in bytes",
+    )
+    auto_scheduler_group.add_argument(
+        "--max-local-memory-per-block",
+        type=int,
+        default=0,
+        help="the max local memory per block in bytes",
+    )
+    auto_scheduler_group.add_argument(
+        "--max-threads-per-block",
+        type=int,
+        default=0,
+        help="the max number of threads per block",
+    )
+    auto_scheduler_group.add_argument(
+        "--max-vthread-extent",
+        type=int,
+        default=0,
+        help="the max vthread extent",
+    )
+    auto_scheduler_group.add_argument(
+        "--warp-size",
+        type=int,
+        default=0,
+        help="the thread numbers of a warp",
+    )
     auto_tuning_group = parser.add_argument_group(
         "Autotuning options",
         "Autotuning options, used when the autoscheduler is not enabled",
@@ -206,28 +236,6 @@ def drive_tune(args):
         min_repeat_ms = 0 if target.keys[0] == "cpu" else 1000
         logger.debug("Default --min-repeat-ms for this target is %s", min_repeat_ms)
 
-    if args.enable_autoscheduler:
-        # Specify hardware parameters
-        hardware_params = auto_scheduler.HardwareParams(
-            args.num_cores, args.vector_unit_bytes, args.cache_line_bytes, 0, 0, 0, 0, 0
-        )
-        tasks, weights = autoscheduler_get_tuning_tasks(
-            mod=mod,
-            params=params,
-            target=target,
-            target_host=args.target_host,
-            alter_layout=args.desired_layout,
-            hardware_params=hardware_params,
-        )
-    else:
-        tasks = autotuner_get_tuning_tasks(
-            mod=mod,
-            params=params,
-            target=target,
-            target_host=args.target_host,
-            alter_layout=args.desired_layout,
-        )
-
     if args.rpc_tracker:
         runner_ctor = auto_scheduler.RPCRunner if args.enable_autoscheduler else autotvm.RPCRunner
         runner = runner_ctor(
@@ -253,6 +261,26 @@ def drive_tune(args):
         )
 
     if args.enable_autoscheduler:
+        # Specify hardware parameters
+        hardware_params = auto_scheduler.HardwareParams(
+            args.num_cores,
+            args.vector_unit_bytes,
+            args.cache_line_bytes,
+            args.max_shared_memory_per_block,
+            args.max_local_memory_per_block,
+            args.max_threads_per_block,
+            args.max_vthread_extent,
+            args.warp_size,
+        )
+        tasks, weights = autoscheduler_get_tuning_tasks(
+            mod=mod,
+            params=params,
+            target=target,
+            target_host=args.target_host,
+            alter_layout=args.desired_layout,
+            hardware_params=hardware_params,
+        )
+
         # Create the autoscheduler tuning options
         tuning_options = auto_scheduler.TuningOptions(
             num_measure_trials=args.trials,
@@ -268,8 +296,15 @@ def drive_tune(args):
             tuning_options,
             args.tuning_records,
         )
-
     else:
+        tasks = autotvm_get_tuning_tasks(
+            mod=mod,
+            params=params,
+            target=target,
+            target_host=args.target_host,
+            alter_layout=args.desired_layout,
+        )
+
         tuning_option = {
             "tuner": args.tuner,
             "trials": args.trials,
@@ -284,8 +319,8 @@ def drive_tune(args):
         tune_tasks(tasks, args.output, **tuning_option)
 
 
-def autotuner_get_tuning_tasks(mod, params, target, target_host=None, alter_layout=None):
-    """Get the tuning tasks for a given relay module.
+def autotvm_get_tuning_tasks(mod, params, target, target_host=None, alter_layout=None):
+    """Get the autotvm tuning tasks for a given relay module.
 
     Parameters
     ----------
@@ -323,7 +358,7 @@ def autotuner_get_tuning_tasks(mod, params, target, target_host=None, alter_layo
 def autoscheduler_get_tuning_tasks(
     mod, params, target, target_host=None, alter_layout=None, hardware_params=None
 ):
-    """Get the tuning tasks for a given relay module.
+    """Get the autoscheduler tuning tasks for a given relay module.
 
     Parameters
     ----------
@@ -339,6 +374,8 @@ def autoscheduler_get_tuning_tasks(
         The layout to convert the graph to. Note, the convert layout
         pass doesn't currently guarantee the whole of the graph will
         be converted to the chosen layout.
+    hardware_params : Optional[HardwareParams]
+        Hardware parameters used for the search tasks
 
     Returns
     -------
@@ -373,7 +410,7 @@ def schedule_tasks(
         A list of auto_scheduler.SearchTask to tune.
     task_weights : list
         The weight (i.e. the number of appearance) of extracted tasks
-    tuning_options:
+    tuning_options: dict
         The options of tuning
     tuning_records : str, optional
         The json file used to preload the autoscheduler
diff --git a/python/tvm/driver/tvmc/compiler.py b/python/tvm/driver/tvmc/compiler.py
index 3dc043ab521e..90b0aceaa17a 100644
--- a/python/tvm/driver/tvmc/compiler.py
+++ b/python/tvm/driver/tvmc/compiler.py
@@ -184,7 +184,6 @@ def compile_model(
         logger.debug("tuning records file provided: %s", tuning_records)
 
         use_autoscheduler = True
-
         try:
             auto_scheduler.load_records(tuning_records)
         except tvm._ffi.base.TVMError:
diff --git a/tests/python/driver/tvmc/test_autotuner.py b/tests/python/driver/tvmc/test_autotuner.py
index bdad167cfe3a..5ce4ca95c810 100644
--- a/tests/python/driver/tvmc/test_autotuner.py
+++ b/tests/python/driver/tvmc/test_autotuner.py
@@ -27,7 +27,7 @@
 
 def _get_tasks(model):
     mod, params = tvmc.frontends.load_model(model)
-    return tvmc.autotuner.get_tuning_tasks(mod, params, "llvm")
+    return tvmc.autotuner.autotvm_get_tuning_tasks(mod, params, "llvm")
 
 
 def _get_measure_options():

From 881dde04b58dfffa64ff601654a786837d695f5a Mon Sep 17 00:00:00 2001
From: Giuseppe Rossini <giuseppe.rossini@arm.com>
Date: Tue, 15 Dec 2020 17:30:07 +0000
Subject: [PATCH 7/8] Addressing comments -3

Change-Id: I207872757473210681d9db04bfdcd2c5e6deaa05
---
 python/tvm/driver/tvmc/autotuner.py           |  47 ++++++--
 .../python/driver/tvmc/test_autoscheduler.py  | 101 ++++++++++++++++++
 2 files changed, 137 insertions(+), 11 deletions(-)
 create mode 100644 tests/python/driver/tvmc/test_autoscheduler.py

diff --git a/python/tvm/driver/tvmc/autotuner.py b/python/tvm/driver/tvmc/autotuner.py
index 37188261080e..04bc993b548c 100644
--- a/python/tvm/driver/tvmc/autotuner.py
+++ b/python/tvm/driver/tvmc/autotuner.py
@@ -186,6 +186,16 @@ def add_tune_parser(subparsers):
         default=0,
         help="the thread numbers of a warp",
     )
+    auto_scheduler_group.add_argument(
+        "--include-simple-tasks",
+        help="Whether to extract simple tasks that do not include complicated ops",
+        action="store_true",
+    )
+    auto_scheduler_group.add_argument(
+        "--log-estimated-latency",
+        help="Whether to log the estimated latency to the file after tuning a task",
+        action="store_true",
+    )
     auto_tuning_group = parser.add_argument_group(
         "Autotuning options",
         "Autotuning options, used when the autoscheduler is not enabled",
@@ -279,6 +289,7 @@ def drive_tune(args):
             target_host=args.target_host,
             alter_layout=args.desired_layout,
             hardware_params=hardware_params,
+            include_simple_tasks=args.include_simple_tasks,
         )
 
         # Create the autoscheduler tuning options
@@ -291,10 +302,7 @@ def drive_tune(args):
 
         # Schedule the tasks (i.e., produce a schedule for each task)
         schedule_tasks(
-            tasks,
-            weights,
-            tuning_options,
-            args.tuning_records,
+            tasks, weights, tuning_options, args.tuning_records, args.log_estimated_latency
         )
     else:
         tasks = autotvm_get_tuning_tasks(
@@ -356,7 +364,13 @@ def autotvm_get_tuning_tasks(mod, params, target, target_host=None, alter_layout
 
 
 def autoscheduler_get_tuning_tasks(
-    mod, params, target, target_host=None, alter_layout=None, hardware_params=None
+    mod,
+    params,
+    target,
+    target_host=None,
+    alter_layout=None,
+    hardware_params=None,
+    include_simple_tasks=False,
 ):
     """Get the autoscheduler tuning tasks for a given relay module.
 
@@ -389,17 +403,19 @@ def autoscheduler_get_tuning_tasks(
 
     # Extract the tasks
     tasks, task_weights = auto_scheduler.extract_tasks(
-        mod["main"], params, target=target, target_host=target_host, hardware_params=hardware_params
+        mod["main"],
+        params,
+        target=target,
+        target_host=target_host,
+        hardware_params=hardware_params,
+        include_simple_tasks=include_simple_tasks,
     )
 
     return tasks, task_weights
 
 
 def schedule_tasks(
-    tasks,
-    task_weights,
-    tuning_options,
-    tuning_records=None,
+    tasks, task_weights, tuning_options, tuning_records=None, log_estimated_latency=False
 ):
     """Generate the schedules for the different tasks (i.e., subgraphs) contained in the module.
     Store the schedules in a json file that will be used later by the compiler.
@@ -415,9 +431,18 @@ def schedule_tasks(
     tuning_records : str, optional
         The json file used to preload the autoscheduler
     """
+    if not log_estimated_latency:
+        callbacks = [auto_scheduler.task_scheduler.PrintTableInfo()]
+    else:
+        callbacks = [
+            auto_scheduler.task_scheduler.PrintTableInfo(),
+            auto_scheduler.task_scheduler.LogEstimatedLatency(("total_latency.tsv")),
+        ]
 
     # Create the scheduler
-    tuner = auto_scheduler.TaskScheduler(tasks, task_weights, load_log_file=tuning_records)
+    tuner = auto_scheduler.TaskScheduler(
+        tasks, task_weights, load_log_file=tuning_records, callbacks=callbacks
+    )
 
     # Tune the tasks
     tuner.tune(tuning_options)
diff --git a/tests/python/driver/tvmc/test_autoscheduler.py b/tests/python/driver/tvmc/test_autoscheduler.py
new file mode 100644
index 000000000000..25525eb9ce97
--- /dev/null
+++ b/tests/python/driver/tvmc/test_autoscheduler.py
@@ -0,0 +1,101 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+import json
+import pytest
+import os
+import tarfile
+
+from os import path
+
+from tvm import auto_scheduler
+from tvm.driver import tvmc
+
+
+def _get_tasks(model):
+    mod, params = tvmc.frontends.load_model(model)
+    tasks, weights = tvmc.autotuner.autoscheduler_get_tuning_tasks(mod, params, "llvm")
+    return (tasks, weights)
+
+
+def _autoscheduler_test_helper(
+    model, tmpdir_name, tasks_weights=None, early_stopping=1, tuning_records=None
+):
+    tasks, weights = tasks_weights if tasks_weights else _get_tasks(model)
+    log_file = os.path.join(tmpdir_name, "autoscheduler.json")
+
+    tuning_options = auto_scheduler.TuningOptions(
+        num_measure_trials=1,
+        measure_callbacks=[auto_scheduler.RecordToFile(log_file)],
+        runner="local",
+        builder="local",
+        verbose=0,
+        early_stopping=early_stopping,
+    )
+
+    tvmc.autotuner.schedule_tasks(tasks[:1], weights[:1], tuning_options, tuning_records)
+
+    # testing whether the log file was produced
+    assert path.exists(log_file), "autoscheduler log file should exist"
+
+    with auto_scheduler.ApplyHistoryBest(log_file) as best:
+        assert isinstance(
+            best, auto_scheduler.dispatcher.ApplyHistoryBest
+        ), "unable to load the best results of tuning"
+
+    return log_file
+
+
+def test_get_tuning_tasks(onnx_resnet50):
+    pytest.importorskip("onnx")
+
+    tasks, weights = _get_tasks(onnx_resnet50)
+    expected_task_type = auto_scheduler.SearchTask
+
+    assert type(tasks) is list
+    assert len(tasks) > 0
+    assert all([type(x) is expected_task_type for x in tasks]) is True
+
+
+def test_tune_tasks(onnx_resnet50, tmpdir_factory):
+    pytest.importorskip("onnx")
+
+    tmpdir_name = tmpdir_factory.mktemp("data")
+    _autoscheduler_test_helper(onnx_resnet50, tmpdir_name)
+
+
+def test_tune_tasks__tuning_records(onnx_resnet50, tmpdir_factory):
+    pytest.importorskip("onnx")
+
+    tmpdir_name = tmpdir_factory.mktemp("data")
+    output_log_phase_1 = _autoscheduler_test_helper(onnx_resnet50, tmpdir_name)
+
+    # Exercises transfer learning by making sure a previous log exists
+    _autoscheduler_test_helper(onnx_resnet50, tmpdir_name, tuning_records=output_log_phase_1)
+
+
+def test_tune_tasks__no_early_stopping(onnx_resnet50, tmpdir_factory):
+    pytest.importorskip("onnx")
+
+    tmpdir_name = tmpdir_factory.mktemp("data")
+    _autoscheduler_test_helper(onnx_resnet50, tmpdir_name, tasks_weights=None, early_stopping=None)
+
+
+def test_tune_tasks__no_tuning_records(onnx_resnet50, tmpdir_factory):
+    pytest.importorskip("onnx")
+
+    tmpdir_name = tmpdir_factory.mktemp("data")
+    _autoscheduler_test_helper(onnx_resnet50, tmpdir_name, tasks_weights=None, tuning_records=None)

From df6eaf604bb99a40e8e4e17d7c28a2d62558f6db Mon Sep 17 00:00:00 2001
From: Giuseppe Rossini <giuseppe.rossini@arm.com>
Date: Wed, 16 Dec 2020 12:11:14 +0000
Subject: [PATCH 8/8] Addressing comments - 4

Change-Id: I11f73c9b32e83c013cfb2224ccce06f60a128af7
---
 python/tvm/driver/tvmc/autotuner.py | 14 +++++++-------
 1 file changed, 7 insertions(+), 7 deletions(-)

diff --git a/python/tvm/driver/tvmc/autotuner.py b/python/tvm/driver/tvmc/autotuner.py
index 04bc993b548c..71ccc8546e8b 100644
--- a/python/tvm/driver/tvmc/autotuner.py
+++ b/python/tvm/driver/tvmc/autotuner.py
@@ -188,23 +188,23 @@ def add_tune_parser(subparsers):
     )
     auto_scheduler_group.add_argument(
         "--include-simple-tasks",
-        help="Whether to extract simple tasks that do not include complicated ops",
+        help="whether to extract simple tasks that do not include complicated ops",
         action="store_true",
     )
     auto_scheduler_group.add_argument(
         "--log-estimated-latency",
-        help="Whether to log the estimated latency to the file after tuning a task",
+        help="whether to log the estimated latency to the file after tuning a task",
         action="store_true",
     )
-    auto_tuning_group = parser.add_argument_group(
-        "Autotuning options",
-        "Autotuning options, used when the autoscheduler is not enabled",
+    autotvm_group = parser.add_argument_group(
+        "autotvm options",
+        "autotvm options, used when the autoscheduler is not enabled",
     )
-    auto_tuning_group.add_argument(
+    autotvm_group.add_argument(
         "--tuner",
         choices=["ga", "gridsearch", "random", "xgb", "xgb_knob", "xgb-rank"],
         default="xgb",
-        help="type of tuner to use when autotuning.",
+        help="type of tuner to use when tuning with autotvm.",
     )
     # TODO (@leandron) This is a path to a physical file, but
     #     can be improved in future to add integration with a modelzoo