[AutoScheduler] Add task scheduler

include/tvm/auto_scheduler/measure.h

@@ -423,7 +423,7 @@ class RPCRunner : public ProgramRunner {
 
 /*!
  * \brief Measurer that measures the time costs of tvm programs
- * This class combines ProgramBuilder and ProgramRunner and provides a simpler API */
+ * This class combines ProgramBuilder and ProgramRunner, and provides a simpler API */
 class ProgramMeasurerNode : public Object {
  public:
   /*! \brief Measured programs counter. */
@@ -444,7 +444,7 @@ class ProgramMeasurerNode : public Object {
   Optional<Array<MeasureCallback>> callbacks;
   /*! \brief Verbosity level. 0 for silent, 1 to output information during program measuring. */
   int verbose;
-  /*! \brief The number of max continuous error. */
+  /*! \brief The number of allowed maximum continuous error before forcely stopping the tuning */
   int max_continuous_error;
 
   /*! \brief Reset book keeping variables */
@@ -454,13 +454,12 @@ class ProgramMeasurerNode : public Object {
    * \brief Do measurement.
    * \param task The current SearchTask.
    * \param policy The current SearchPolicy.
-   * \param inputs The MeasureInputs.
-   * \param results A pointer to a MeasureResult Array, this is used as output.
+   * \param inputs The inputs of measurement.
    * \param batch_size Number of programs to be measured in one batch.
+   * \return results The results of measurement.
    */
-  void Measure(const SearchTask& task, const SearchPolicy& policy,
-               const Array<MeasureInput>& inputs, Array<MeasureResult>* results,
-               int batch_size = -1);
+  Array<MeasureResult> Measure(const SearchTask& task, const SearchPolicy& policy,
+                               const Array<MeasureInput>& inputs, int batch_size = -1);
   /*!
    * \brief Do measurement silently.
    * This API will not print the measure results to screen.
@@ -486,12 +485,13 @@ class ProgramMeasurer : public ObjectRef {
  public:
   /*!
    * \brief The constructor.
-   * \param builder The ProgramBuilder to build each program.
-   * \param runner The ProgramRunner to measure each program.
-   * \param callbacks MeasureCallback to be called after each measure batch.
+   * \param builder The ProgramBuilder to build programs.
+   * \param runner The ProgramRunner to measure programs.
+   * \param callbacks MeasureCallback to be called after each measurement batch.
    * \param verbose Verbosity level. 0 for silent, 1 to output information during program
    * measuring.
-   * \param max_continuous_error The number of allowed maximum continuous error.
+   * \param max_continuous_error The number of allowed maximum continuous error before
+   * forcely stopping the tuning.
    */
   ProgramMeasurer(ProgramBuilder builder, ProgramRunner runner,
                   Optional<Array<MeasureCallback>> callbacks, int verbose,

include/tvm/auto_scheduler/search_policy.h

@@ -22,26 +22,6 @@
  * \brief The base class of search policies, including the abstract definition of search policy and
  * other supporting data structures.
  *
- * The basic schedule search process for the auto-scheduler is design to be:
- * `Program sampling` -> `Performance Tuning`.
- *
- * In `Program sampling`, we use some predefined precise or heuristic rules to generate several
- * initial schedules. Based on these initial starting points, we perform `Performance Tuning` which
- * uses cost model based evolutionary search to select schedules with the best performance.
- *
- * Candidate schedules are measured against the specific hardware target.
- *
- * We intend to introduce different level of automation on the schedule generation process:
- * - Level 0(the default level): For all kinds of ops/subgraphs, the search policy should be able
- *   to generate schedule automatically.
- * - Level 1: For some complicated ops/subgraphs(e.g. conv2d windograd), the default search space
- *   of level 0 may be too large to find a high performance schedule efficiently. We provide some
- *   op attributes to help reduce the total search space, see `SearchPolicyKey` below for more
- *   information.
- * - Level 2: For some further special ops/subgraphs, users may more likely to write their own
- *   template(just like AutoTVM). Search policy should be able to provide a flexible approach as
- *   well.
- *
  * \note How to add a new search policy.
  * In design, there's no need for users to implement their own search policy, our formal search
  * policy(will be brought later) should be enough to cover most use cases. Meanwhile, a custom rule
@@ -62,11 +42,13 @@
 #ifndef TVM_AUTO_SCHEDULER_SEARCH_POLICY_H_
 #define TVM_AUTO_SCHEDULER_SEARCH_POLICY_H_
 
+#include <tvm/auto_scheduler/measure.h>
 #include <tvm/auto_scheduler/search_task.h>
 #include <tvm/node/node.h>
 
 #include <string>
 #include <unordered_set>
+#include <utility>
 #include <vector>
 
 namespace tvm {
@@ -171,6 +153,15 @@ class SearchPolicyNode : public Object {
   virtual State Search(int num_measure_trials, int early_stopping, int num_measures_per_round,
                        ProgramMeasurer measurer) = 0;
 
+  /*!
+   * \brief Continue the search by doing an additional search round.
+   * \param num_measure The number of measurements
+   * \param measurer The measurer to measure programs
+   * \return The measurement records for measurements in this search round
+   */
+  virtual std::pair<Array<MeasureInput>, Array<MeasureResult>> ContinueSearchOneRound(
+      int num_measure, ProgramMeasurer measurer) = 0;
+
   /*!
    * \brief Preload measured states from a log file to resume the state of the search policy.
    * \param log_file The name of the record log file.

python/tvm/auto_scheduler/__init__.py

@@ -24,6 +24,7 @@
 from . import measure_record
 from . import search_policy
 from . import search_task
+from . import task_scheduler
 from . import utils
 from . import workload_registry
 
@@ -42,4 +43,5 @@
 from .measure_record import RecordToFile, RecordReader, load_best, load_records, save_records
 from .search_task import SearchTask
 from .search_policy import EmptyPolicy, SketchPolicy, PreloadMeasuredStates
+from .task_scheduler import TaskScheduler
 from .workload_registry import register_workload, make_workload_key

python/tvm/auto_scheduler/auto_schedule.py

@@ -16,16 +16,7 @@
 # under the License.
 
 """
-User interface for TVM Auto-scheduler.
-
-The basic schedule search process for TVM Auto-scheduler is designed to be:
-`Program sampling` -> `Performance Tuning`.
-
-In `Program sampling`, we use some predefined precise or heuristic rules to generate several
-initial schedules. Based on these initial starting points, we perform `Performance Tuning` which
-uses cost model based evolutionary search to select schedules with the best performance.
-
-Candidate schedules are measured against the specific hardware target.
+The user interface and tuning options of the TVM auto-scheduler.
 """
 
 import tvm._ffi

python/tvm/auto_scheduler/cost_model/xgb_model.py

@@ -20,6 +20,7 @@
 import multiprocessing
 import logging
 from collections import defaultdict
+import time
 
 import numpy as np
 import xgboost as xgb
@@ -76,7 +77,7 @@ def set(self, key, matrix, value):
 class XGBModel(PythonBasedModel):
     """Train a XGBoost model to predict the normalized throughputs of programs.
     Let the normalized throughput be the score of a program (higher is better). We predict
-    the (approximiate) score of a program = the sum of the scores of all stages in this program.
+    the (approximate) score of a program = the sum of the scores of all stages in this program.
     i.e. score(P) = score_s0 + score_s1 + ... + score_sn,
     where score_si is the score of Stage i in Program P.
     We extract feature for each stage and let the xgboost predict the score for each stage.
@@ -128,6 +129,7 @@ def update(self, inputs, results):
         if len(inputs) <= 0:
             return
         assert len(inputs) == len(results)
+        tic = time.time()
 
         self.inputs.extend(inputs)
         self.results.extend(results)
@@ -167,6 +169,8 @@ def update(self, inputs, results):
             ],
         )
 
+        logger.info("XGBModel Training time: %.2f s", time.time() - tic)
+
     def predict(self, task, states):
         """Predict the scores of states
         Parameters

python/tvm/auto_scheduler/measure.py

@@ -185,6 +185,33 @@ def run(self, measure_inputs, build_results, verbose=1):
         return _ffi_api.ProgramRunnerRun(self, measure_inputs, build_results, verbose)
 
 
+@tvm._ffi.register_object("auto_scheduler.ProgramMeasurer")
+class ProgramMeasurer(Object):
+    """
+    Measurer that measures the time costs of tvm programs
+    This class combines ProgramBuilder and ProgramRunner, and provides a simpler API.
+
+    Parameters
+    ----------
+    builder : ProgramBuilder
+        The ProgramBuilder to build programs
+    runner : ProgramRunner
+        The ProgramRunner to measure programs.
+    callbacks : List[MeasureCallback]
+        Callbacks to be called after each measurement batch
+    verbose : int
+        The Verbosity level: 0 for silent, 1 to output information during program
+    max_continuous_error : Optional[int]
+        The number of allowed maximum continuous error before stop the tuning
+    """
+
+    def __init__(self, builder, runner, callbacks, verbose, max_continuous_error=None):
+        max_continuous_error = max_continuous_error or -1  # -1 means using the default value
+        self.__init_handle_by_constructor__(
+            _ffi_api.ProgramMeasurer, builder, runner, callbacks, verbose, max_continuous_error
+        )
+
+
 @tvm._ffi.register_object("auto_scheduler.LocalBuilder")
 class LocalBuilder(ProgramBuilder):
     """LocalBuilder use local CPU cores to build programs in parallel.

python/tvm/auto_scheduler/search_policy.py

@@ -16,15 +16,17 @@
 # under the License.
 
 """
-The search policies for TVM Auto-scheduler.
+The search policies of TVM auto-scheduler.
 
-This contains the strategies to generate a schedule automatically. We provide an EmptyPolicy
-which always returns an unchanged initial state, and a more advanced SketchPolicy which can
-deal with various ops/subgraphs on different target devices.
+The auto-scheduler constructs a search space according to the compute declaration.
+It then randomly samples programs from the search space and uses evolutionary search with a
+learned cost model to fine tune the sampled programs.
+The final optimized programs are sent to actual hardware for measurement.
+The above process is repeated until the auto-scheduler runs out of time budget.
 
 Reference:
 L. Zheng, C. Jia, M. Sun, Z. Wu, C. Yu, et al. "Ansor : Generating High-Performance Tensor
-Programs for Deep Learning." arXiv preprint arXiv:2006.06762 (2020).
+Programs for Deep Learning." (OSDI 2020).
 """
 
 import random
@@ -63,11 +65,42 @@ def __init__(self, filename="auto_scheduler_tuning.json"):
 class SearchPolicy(Object):
     """ The base class of search policies. """
 
+    def continue_search_one_round(self, num_measure, measurer):
+        """
+        Continue the search by doing an additional search round.
+
+        Parameters
+        ----------
+        num_measure: int
+            The number of programs to measure in this round
+        measurer: ProgramMeasurer
+            The program measurer to measure programs
+
+        Returns
+        -------
+        inputs: List[MeasureInput]
+            The inputs of measurments in this search round
+        results: List[MeasureResult]
+            The results of measurments in this search round
+        """
+        return _ffi_api.SearchPolicyContinueSearchOneRound(self, num_measure, measurer)
+
+    def set_verbose(self, verbose):
+        """
+        Set the verbosity level of the search policy.
+
+        Parameters
+        ----------
+        verbose: int
+            The verbosity level
+        """
+        return _ffi_api.SearchPolicySetVerbose(self, verbose)
+
 
 @tvm._ffi.register_object("auto_scheduler.EmptyPolicy")
 class EmptyPolicy(SearchPolicy):
-    """This is an example empty search policy which will always generate
-    the init state of ComputeDAG.
+    """A simple example of the search policy which always returns
+    the initial naive schedule (state).
 
     Parameters
     ----------
@@ -195,15 +228,17 @@ def sample_initial_population(self, pop_size):
         return states
 
     def evolutionary_search(self, init_populations, out_size):
-        """Evolutionary search.
+        """Perform evolutionary search.
         This python interface is mainly used for debugging and testing.
         The actual search is all done in c++.
+
         Parameters
         ----------
         init_populations: List[State]
             The initial population states
         out_size : int
             The size of generated states
+
         Returns
         -------
         states: List[State]