[AUTOTVM] Decouple build and run in measurement (#1661)

Author: merrymercy

docs/api/python/autotvm.rst

@@ -16,6 +16,11 @@ tvm.autotvm.measure
 
 .. autofunction:: tvm.autotvm.measure.create_measure_batch
 
+.. autoclass:: tvm.autotvm.measure.measure_methods.LocalBuilder
+
+.. autoclass:: tvm.autotvm.measure.measure_methods.RPCRunner
+
+.. autoclass:: tvm.autotvm.measure.measure_methods.LocalRunner
 
 tvm.autotvm.tuner
 ~~~~~~~~~~~~~~~~~

python/tvm/autotvm/__init__.py

@@ -22,7 +22,8 @@
 from . import tophub
 
 # some shortcuts
-from .measure import measure_option, MeasureInput, MeasureResult, MeasureErrorNo
+from .measure import measure_option, MeasureInput, MeasureResult, MeasureErrorNo, \
+    LocalBuilder, LocalRunner, RPCRunner
 from .tuner import callback
 from .task import template, get_config, create, ConfigSpace, ConfigEntity, \
     register_topi_compute, register_topi_schedule, \

python/tvm/autotvm/measure/__init__.py

@@ -1,7 +1,7 @@
 """Distributed executor infrastructure to scale up the tuning"""
 
-from .measure import MeasureInput, MeasureResult, MeasureErrorNo, measure_option
-from .measure_methods import request_remote, check_remote, create_measure_batch, rpc
-
+from .measure import MeasureInput, MeasureResult, MeasureErrorNo, measure_option, \
+    create_measure_batch
+from .measure_methods import LocalBuilder, LocalRunner, RPCRunner, request_remote
+from .executor import Executor
 from .local_executor import LocalExecutor
-from .executor import Future, Executor

python/tvm/autotvm/measure/local_executor.py

@@ -37,25 +37,21 @@ def _execute_func(func, queue, args, kwargs):
         res = exc
     queue.put(res)
 
-def timeout_monitor(queue, timeout, func, args, kwargs):
+
+def call_with_timeout(queue, timeout, func, args, kwargs):
     """A wrapper to support timeout of a function call"""
 
     # start a new process for timeout (cannot use thread because we have c function)
     p = Process(target=_execute_func, args=(func, queue, args, kwargs))
     p.start()
     p.join(timeout=timeout)
 
-    alive = p.is_alive()
+    queue.put(executor.TimeoutError())
+
     kill_child_processes(p.pid)
     p.terminate()
     p.join()
 
-    if alive:
-        queue.put(executor.TimeoutError())
-    else:
-        if queue.empty():
-            queue.put(executor.ExecutionError("Fatal error in local executor"))
-
 
 class LocalFuture(executor.Future):
     """Local wrapper for the future
@@ -134,7 +130,7 @@ def submit(self, func, *args, **kwargs):
             return LocalFutureNoFork(func(*args, **kwargs))
 
         queue = Queue(2)
-        process = Process(target=timeout_monitor,
+        process = Process(target=call_with_timeout,
                           args=(queue, self.timeout, func, args, kwargs))
         process.start()
         return LocalFuture(process, queue)

python/tvm/autotvm/measure/measure.py

@@ -1,5 +1,6 @@
 # pylint: disable=pointless-string-statement,consider-using-enumerate,invalid-name
 """User facing API for specifying how to measure the generated code"""
+import multiprocessing
 from collections import namedtuple
 
 class MeasureInput(namedtuple("MeasureInput", ["target", "task", "config"])):
@@ -16,15 +17,16 @@ class MeasureInput(namedtuple("MeasureInput", ["target", "task", "config"])):
         Specific configuration.
     """
 
+
 class MeasureResult(namedtuple("MeasureResult", ["costs", "error_no", "all_cost", "timestamp"])):
     """
     Stores all the results of a measurement
 
     Parameters
     ----------
     costs: Array of float or Array of Exception
-        If no error occurs for this measurement, it is an array of measured running times.
-        If some error occurs during the measurement, it is an array of the exception objections.
+        If no error occurs during measurement, it is an array of measured running times.
+        If an error occurs during measurement, it is an array of the exception objections.
     error_no: int
         Denote error type, defined by MeasureErrorNo
     all_cost: float
@@ -37,92 +39,185 @@ class MeasureResult(namedtuple("MeasureResult", ["costs", "error_no", "all_cost"
 class MeasureErrorNo(object):
     """Error type for MeasureResult"""
     NO_ERROR = 0              # no error
-    INSTANTIATION_ERROR = 1   # error when calling template function
+    INSTANTIATION_ERROR = 1   # actively detected error in instantiating a template with a config
     COMPILE_HOST = 2          # error when compiling code on host (e.g. tvm.build)
-    COMPILE_DEVICE = 3        # error when compiling code on device (e.g. opencl JIT on device)
+    COMPILE_DEVICE = 3        # error when compiling code on device (e.g. OpenCL JIT on the device)
     RUNTIME_DEVICE = 4        # error when run program on device
     WRONG_ANSWER = 5          # answer is wrong when compared to a golden output
-    FLEET_ERROR = 6           # error of measure infrastructure
+    BUILD_TIMEOUT = 6         # timeout during compilation
+    RUN_TIMEOUT = 7           # timeout during run
+    UNKNOWN_ERROR = 8         # unknown error
+
 
+class Builder(object):
+    """Builder that builds programs in tuning
 
-def measure_option(measure_func,
-                   number=1,
-                   repeat=1,
-                   timeout=60,
-                   n_parallel=1,
-                   do_fork=True,
-                   build_func='default',
-                   check_correctness=False,
-                   replay_db=None):
-    """Configure how to do measurement
+    Parameters
+    ----------
+    timeout: float, optional
+        The timeout of a build task
+    n_parallel: int, optional
+        The number of tasks submitted in parallel
+        By default it will use all cpu cores
+    """
+    def __init__(self, timeout=10, n_parallel=None):
+        self.timeout = timeout
+        self.n_parallel = n_parallel or multiprocessing.cpu_count()
+        self.build_kwargs = {}
+        self.task = None
+
+    def set_task(self, task, build_kwargs=None):
+        """
+        Initialize for a new tuning task
+
+        Parameters
+        ----------
+        task: Task
+            The tuning task
+        build_kwargs: dict, optional
+            The additional kwargs for build function
+        """
+        self.task = task
+        self.build_kwargs = build_kwargs
+
+    def build(self, measure_inputs):
+        """Build programs
+
+        Parameters
+        ----------
+        measure_inputs: List of MeasureInput
+            The measure input
+
+        Returns
+        -------
+        build_results: List of BuildResult
+            The build result.
+        """
+        raise NotImplementedError()
+
+
+class Runner(object):
+    """Runner that runs and measures the time cost of a generated program in tuning
 
     Parameters
     ----------
-    measure_func: str or callable
-        'local': use the local device for measurement. The tuner will start a tracker
-        and a RPC server silently for the user.
-
-        callable: It is a callable function for measurement.
-                  See the return value of measure/measure_methods.py::rpc for example.
-    number : int, optional
-        Number of times to do the measurement for average
-    repeat : int, optional
-        Number of times to repeat the measurement.
-        In total, the generated code will be run (1 + number x repeat) times,
-        where the first one is warm up. The returned result contains `repeat` costs,
-        each of which is the average of `number` test run.
-    timeout: int, optional
-        Timeout for a whole batch. TimeoutError will be returned as the result if a
-        task timeouts.
+    timeout: float, optional
+        The timeout of a build task
     n_parallel: int, optional
-        The number of measurement task that can run in parallel.
-        Set this according to the number of cpu cores (for compilation) and
-        the number of devices you have (for measuring generate code).
-    do_fork: bool, optional
-        Whether use multiprocessing (based on fork) for running measure jobs in parallel.
-        Set this to False if you want to debug (see trackback) or using fork is not suitable.
-        NOTE: If this is False, parallel and timeout do not work.
-    build_func: str or callable, optional
-        'default': call default builder. This works for normal target (llvm, cuda)
-
-        'ndk': use Android NDK to create shared library. Use this for android target.
-
-        callable: customized build function for other backends (e.g. VTA).
-                  See measure/measure_methods.py::default_build_func for example.
-    check_correctness: bool, optional
-        Whether check correctness after measurement. This will use llvm cpu target to generate
-        reference output.
-    replay_db : Database, optional
-        The database that we retrieve saved MeasureResult from.
+        The number of tasks submitted in parallel
+        By default it will use all cpu cores
+    """
+    def __init__(self, timeout=5, n_parallel=None):
+        self.timeout = timeout
+        self.n_parallel = n_parallel or multiprocessing.cpu_count()
+        self.task = None
+
+    def set_task(self, task):
+        """
+        Initialize for a new tuning task
+
+        Parameters
+        ----------
+        task: Task
+            The tuning task
+        """
+        self.task = task
+
+    def get_build_kwargs(self):
+        """
+        Get device specific build arguments (e.g. maximum shared memory size)
+
+        Returns
+        ----------
+        kwargs: dict
+            The additional keyword arguments
+        """
+        raise NotImplementedError()
+
+    def run(self, measure_inputs, build_results):
+        """Run amd measure built programs
+
+        Parameters
+        ----------
+        measure_inputs: List of MeasureInput
+            The raw measure input
+        build_results: List of BuildResults
+            The build results
+
+        Returns
+        -------
+        measure_results: List of MeasureResult
+            The final results of measurement
+        """
+        raise NotImplementedError()
+
+
+def measure_option(builder, runner):
+    """
+    Set options for measure. To measure a config, we will build it and run it.
+    So we have to set options for these two steps.
+    They have their own options on timeout, parallel, etc.
+
+    Parameters
+    ----------
+    builder: Builder
+        Specify how to build programs
+    runner: Runner
+        Specify how to run programs
+    """
+    from .measure_methods import LocalBuilder, LocalRunner
+
+    if isinstance(builder, str):
+        if builder == 'local':
+            builder = LocalBuilder()
+        else:
+            raise ValueError("Invalid builder: " + builder)
+
+    if isinstance(runner, str):
+        if runner == 'local':
+            runner = LocalRunner()
+        else:
+            raise ValueError("Invalid runner: " + runner)
+
+    opt = {
+        'builder': builder,
+        'runner': runner,
+    }
+
+    return opt
+
+
+def create_measure_batch(task, option):
+    """Get a standard measure_batch function.
+
+    Parameters
+    ----------
+    task: tvm.autotvm.task.Task
+        The tuning task
+    option: dict
+        The option for measuring generated code.
+        You should use the return value of function :any:`measure_option` for this argument.
 
     Returns
     -------
-    options: dict
-        A dict to store all options
-
-    Note
-    ----
-    To support customized measure, you can pass callable `measure_func` or
-    `build_func` in. The `measure_func` will call `build_func` to build binary library
-    and handle the logic of measurement.
-
-    Signature:
-    * measure_func (see the return value of measure/measure_methods.py::rpc for example)
-    def measure_func(input_pack, build_func, build_kwargs, number, repeat, ref_input, ref_output):
-        return measure_results
-
-    * build_func (see measure/measure_methods.py::default_build_func for example)
-    def build_func(inp, tmp_dir, **kwargs):
-        return func, args, filename
+    measure_batch: callable
+        a callback function to measure a batch of configs
     """
-    return {
-        'measure_func': measure_func,
-        'number': number,
-        'repeat': repeat,
-        'timeout': timeout,
-        'n_parallel': n_parallel,
-        'do_fork': do_fork,
-        'build_func': build_func,
-        'check_correctness': check_correctness,
-        'replay_db': replay_db,
-    }
+    builder = option['builder']
+    runner = option['runner']
+
+    attach_objects = runner.set_task(task)
+
+    # feed device related information from runner to builder
+    # (e.g. max shared memory for validity checking)
+    build_kwargs = runner.get_build_kwargs()
+    builder.set_task(task, build_kwargs)
+
+    def measure_batch(measure_inputs):
+        build_results = builder.build(measure_inputs)
+        results = runner.run(measure_inputs, build_results)
+        return results
+
+    measure_batch.n_parallel = builder.n_parallel
+    measure_batch.attach_objects = attach_objects
+    return measure_batch