[Meta Schedule][M3b] Runner (apache#9111)

This PR is part of the meta schedule project (apache#8473) that adds the
asynchronous program runner interface, as well as a reference
implementation of RPCRunner. LocalRunner will be implemented with
PopenPool executor in a follow-up PR.

Co-authored-by: Xiyou Zhou <xiyou@octoml.ai>
Co-authored-by: Bohan Hou <32121147+spectrometerHBH@users.noreply.github.com>
Co-authored-by: Ruihang Lai <lairuihangdongdong@qq.com>
Co-authored-by: Hongyi Jin <3231950289@qq.com>
Co-authored-by: Wuwei Lin <wuwei@apache.org>
Co-authored-by: Siyuan Feng <Hzfengsy@sjtu.edu.cn>

Address comments

Co-authored-by: Cody Yu <comaniac0422@gmail.com>

fix lint

include/tvm/meta_schedule/runner.h

@@ -20,16 +20,53 @@
 #define TVM_META_SCHEDULE_RUNNER_H_
 
 #include <tvm/ir/expr.h>
+#include <tvm/meta_schedule/arg_info.h>
 
 namespace tvm {
 namespace meta_schedule {
 
-/*! \brief Runner's output containing measurement result of MeasureCandidate or error msg if any. */
+/*! \brief The runner's input. */
+class RunnerInputNode : public runtime::Object {
+ public:
+  /*! \brief The path to the built artifact. */
+  String artifact_path;
+  /*! \brief The type of device. */
+  String device_type;
+  /*! \brief The argument information. */
+  Array<ArgInfo> args_info;
+
+  void VisitAttrs(tvm::AttrVisitor* v) {
+    v->Visit("artifact_path", &artifact_path);
+    v->Visit("device_type", &device_type);
+    v->Visit("args_info", &args_info);
+  }
+
+  static constexpr const char* _type_key = "meta_schedule.RunnerInput";
+  TVM_DECLARE_FINAL_OBJECT_INFO(RunnerInputNode, runtime::Object);
+};
+
+/*!
+ * \brief Managed reference to RunnerInputNode
+ * \sa RunnerInputNode
+ */
+class RunnerInput : public runtime::ObjectRef {
+ public:
+  /*!
+   * \brief Constructor of RunnerInput
+   * \param artifact_path The path to the built artifact.
+   * \param device_type The type of device.
+   * \param args_info The argument information.
+   */
+  TVM_DLL explicit RunnerInput(String artifact_path, String device_type, Array<ArgInfo> args_info);
+  TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(RunnerInput, runtime::ObjectRef, RunnerInputNode);
+};
+
+/*! \brief The runner's output. */
 class RunnerResultNode : public runtime::Object {
  public:
-  /*! \brief The run time in seconds. If not None, error_msg should be None. */
+  /*! \brief The run time in seconds.*/
   Optional<Array<FloatImm>> run_secs;
-  /*! \brief The error message, if any. If not None, run_secs should be None. */
+  /*! \brief The error message, if any. */
   Optional<String> error_msg;
 
   void VisitAttrs(tvm::AttrVisitor* v) {
@@ -48,14 +85,134 @@ class RunnerResultNode : public runtime::Object {
 class RunnerResult : public runtime::ObjectRef {
  public:
   /*!
-   * \brief Constructor for RunnerResult.
-   * \param run_secs The run time in seconds.
-   * \param error_msg The error message, if any.
+   * \brief Constructor
+   * \brief The run time in seconds.
+   * \brief The error message, if any.
    */
   TVM_DLL explicit RunnerResult(Optional<Array<FloatImm>> run_secs, Optional<String> error_msg);
   TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(RunnerResult, runtime::ObjectRef, RunnerResultNode);
 };
 
+/*!
+ * \brief A class to asynchronously fetch runner's output.
+ * \note The API design is consistent with python's concurrent.futures.Future:
+ * https://docs.python.org/3/library/concurrent.futures.html#concurrent.futures.Future
+ */
+class RunnerFutureNode : public runtime::Object {
+ public:
+  /*!
+   * \brief The function type to check whether the runner has finished.
+   * \return Whether the runner's output is ready.
+   */
+  using FDone = runtime::TypedPackedFunc<bool()>;
+  /*!
+   * \brief The function type to fetch runner output if it is ready.
+   * \return The runner's output.
+   */
+  using FResult = runtime::TypedPackedFunc<RunnerResult()>;
+
+  /*! \brief The packed function to check whether the runner has finished. */
+  FDone f_done;
+  /*! \brief The packed function to fetch runner output if it is ready. */
+  FResult f_result;
+
+  void VisitAttrs(tvm::AttrVisitor* v) {
+    // `f_done` is not visited
+    // `f_result` is not visited
+  }
+
+  /*!
+   * \brief Check whether the runner has finished.
+   * \return A boolean indicating whether the runner has finished.
+   */
+  bool Done() const { return f_done(); }
+  /*!
+   * \brief Fetch the runner's output if it is ready.
+   * \return The runner's output.
+   */
+  RunnerResult Result() const { return f_result(); }
+
+  static constexpr const char* _type_key = "meta_schedule.RunnerFuture";
+  TVM_DECLARE_FINAL_OBJECT_INFO(RunnerFutureNode, runtime::Object);
+};
+
+/*!
+ * \brief Managed reference to RunnerFutureNode
+ * \sa RunnerFutureNode
+ */
+class RunnerFuture : public runtime::ObjectRef {
+ public:
+  using FDone = RunnerFutureNode::FDone;
+  using FResult = RunnerFutureNode::FResult;
+
+  /*!
+   * \brief Constructor of RunnerFuture
+   * \param f_done The packed function to check whether the runner has finished.
+   * \param f_result The packed function to fetch runner output if it is ready.
+   */
+  TVM_DLL explicit RunnerFuture(FDone f_done, FResult f_result);
+  TVM_DEFINE_MUTABLE_NOTNULLABLE_OBJECT_REF_METHODS(RunnerFuture, runtime::ObjectRef,
+                                                    RunnerFutureNode);
+};
+
+/*! \brief The abstract runner interface. */
+class RunnerNode : public runtime::Object {
+ public:
+  /*!
+   * \brief The function type to run the built artifacts and get runner futures.
+   * \param input The runner's inputs.
+   * \return The runner futures.
+   * \sa RunnerFuture
+   */
+  using FRun = runtime::TypedPackedFunc<Array<RunnerFuture>(Array<RunnerInput>)>;
+
+  /*! \brief Default destructor */
+  virtual ~RunnerNode() = default;
+
+  /*!
+   * \brief Run the built artifact and get runner futures.
+   * \param runner_inputs The runner's inputs.
+   * \return The runner futures.
+   */
+  virtual Array<RunnerFuture> Run(Array<RunnerInput> runner_inputs) = 0;
+
+  static constexpr const char* _type_key = "meta_schedule.Runner";
+  TVM_DECLARE_BASE_OBJECT_INFO(RunnerNode, runtime::Object);
+};
+
+/*!
+ * \brief Managed reference to RunnerNode
+ * \sa RunnerNode
+ */
+class Runner : public runtime::ObjectRef {
+ public:
+  using FRun = RunnerNode::FRun;
+
+  /*!
+   * \brief Create a runner with customized build method on the python-side.
+   * \param f_run The packed function to run the built artifacts and get runner futures.
+   * \return The runner created.
+   */
+  TVM_DLL static Runner PyRunner(FRun f_run);
+  TVM_DEFINE_MUTABLE_NOTNULLABLE_OBJECT_REF_METHODS(Runner, runtime::ObjectRef, RunnerNode);
+};
+
+/*! \brief An abstract runner with customized build method on the python-side. */
+class PyRunnerNode : public RunnerNode {
+ public:
+  /*! \brief The packed function to run the built artifacts and get runner futures. */
+  FRun f_run;
+
+  void VisitAttrs(tvm::AttrVisitor* v) {
+    // `f_run` is not visited
+  }
+
+  Array<RunnerFuture> Run(Array<RunnerInput> runner_inputs) final { return f_run(runner_inputs); }
+
+  static constexpr const char* _type_key = "meta_schedule.PyRunner";
+  TVM_DECLARE_FINAL_OBJECT_INFO(PyRunnerNode, RunnerNode);
+};
+
 }  // namespace meta_schedule
 }  // namespace tvm
 

python/tvm/meta_schedule/__init__.py

@@ -16,10 +16,9 @@
 # under the License.
 """Package `tvm.meta_schedule`. The meta schedule infrastructure."""
 from . import arg_info
-from . import builder
 from . import database
+from . import builder
+from . import runner
 from . import space_generator
 from . import search_strategy
-from . import runner
-from .database import TuningRecord
 from .tune_context import TuneContext

python/tvm/meta_schedule/builder/local_builder.py

@@ -48,11 +48,20 @@ class LocalBuilder(PyBuilder):
     Attributes
     ----------
     T_BUILD : typing._GenericAlias
-        The signature of the build function `f_build`, which is
-        `Callable[[IRModule, Target], Module]`
+        The signature of the function `f_build`, which is
+
+        .. code-block:: python
+
+        def default_build(mod: IRModule, target: Target) -> Module:
+            ...
+
     T_EXPORT : typing._GenericAlias
-        The signature of the build function `f_export`, which is
-        `Callable[[Module], str]`
+        The signature of the function `f_export`, which is
+
+        .. code-block:: python
+
+        def default_export(mod: Module) -> str:
+            ...
 
     Note
     ----

python/tvm/meta_schedule/runner/__init__.py

@@ -14,5 +14,10 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
-"""meta_schedule.runner"""
-from .runner import RunnerResult
+"""
+The tvm.meta_schedule.runner package.
+Meta Schedule runners that runs an artifact either locally or through the RPC interface
+"""
+from .config import EvaluatorConfig, RPCConfig
+from .rpc_runner import RPCRunner
+from .runner import PyRunner, Runner, RunnerFuture, RunnerInput, RunnerResult