[MetaSchedule][M3c] Update TuneContext, TaskScheduler & Search Strategy Design

include/tvm/meta_schedule/cost_model.h

@@ -51,20 +51,20 @@ class CostModelNode : public runtime::Object {
 
   /*!
    * \brief Update the cost model given running results.
-   * \param tune_context The tuning context.
+   * \param context The tuning context.
    * \param candidates The measure candidates.
    * \param results The running results of the measure candidates.
    */
-  virtual void Update(const TuneContext& tune_context, const Array<MeasureCandidate>& candidates,
+  virtual void Update(const TuneContext& context, const Array<MeasureCandidate>& candidates,
                       const Array<RunnerResult>& results) = 0;
 
   /*!
    * \brief Predict the normalized score (the larger the better) of given measure candidates.
-   * \param tune_context The tuning context.
+   * \param context The tuning context.
    * \param candidates The measure candidates.
    * \return The predicted normalized score.
    */
-  virtual std::vector<double> Predict(const TuneContext& tune_context,
+  virtual std::vector<double> Predict(const TuneContext& context,
                                       const Array<MeasureCandidate>& candidates) = 0;
 
   static constexpr const char* _type_key = "meta_schedule.CostModel";
@@ -86,7 +86,7 @@ class PyCostModelNode : public CostModelNode {
   using FSave = runtime::TypedPackedFunc<void(String)>;
   /*!
    * \brief Update the cost model given running results.
-   * \param tune_context The tuning context.
+   * \param context The tuning context.
    * \param candidates The measure candidates.
    * \param results The running results of the measure candidates.
    * \return Whether cost model was updated successfully.
@@ -95,7 +95,7 @@ class PyCostModelNode : public CostModelNode {
                                                 const Array<RunnerResult>&)>;
   /*!
    * \brief Predict the running results of given measure candidates.
-   * \param tune_context The tuning context.
+   * \param context The tuning context.
    * \param candidates The measure candidates.
    * \param p_addr The address to save the the estimated running results.
    */
@@ -135,17 +135,17 @@ class PyCostModelNode : public CostModelNode {
     ICHECK(f_save != nullptr) << "PyCostModel's Save method not implemented!";
     f_save(path);
   }
-  void Update(const TuneContext& tune_context, const Array<MeasureCandidate>& candidates,
+  void Update(const TuneContext& context, const Array<MeasureCandidate>& candidates,
               const Array<RunnerResult>& results) {
     ICHECK(f_update != nullptr) << "PyCostModel's Update method not implemented!";
-    f_update(tune_context, candidates, results);
+    f_update(context, candidates, results);
   }
 
-  std::vector<double> Predict(const TuneContext& tune_context,
+  std::vector<double> Predict(const TuneContext& context,
                               const Array<MeasureCandidate>& candidates) {
     ICHECK(f_predict != nullptr) << "PyCostModel's Predict method not implemented!";
     std::vector<double> result(candidates.size(), 0.0);
-    f_predict(tune_context, candidates, result.data());
+    f_predict(context, candidates, result.data());
     return result;
   }
 

include/tvm/meta_schedule/feature_extractor.h

@@ -37,11 +37,11 @@ class FeatureExtractorNode : public runtime::Object {
 
   /*!
    * \brief Extract features from the given measure candidate.
-   * \param tune_context The tuning context for feature extraction.
+   * \param context The tuning context for feature extraction.
    * \param candidates The measure candidates to extract features from.
    * \return The feature ndarray extracted.
    */
-  virtual Array<tvm::runtime::NDArray> ExtractFrom(const TuneContext& tune_context,
+  virtual Array<tvm::runtime::NDArray> ExtractFrom(const TuneContext& context,
                                                    const Array<MeasureCandidate>& candidates) = 0;
 
   static constexpr const char* _type_key = "meta_schedule.FeatureExtractor";
@@ -53,12 +53,12 @@ class PyFeatureExtractorNode : public FeatureExtractorNode {
  public:
   /*!
    * \brief Extract features from the given measure candidate.
-   * \param tune_context The tuning context for feature extraction.
+   * \param context The tuning context for feature extraction.
    * \param candidates The measure candidates to extract features from.
    * \return The feature ndarray extracted.
    */
   using FExtractFrom = runtime::TypedPackedFunc<Array<tvm::runtime::NDArray>(
-      const TuneContext& tune_context, const Array<MeasureCandidate>& candidates)>;
+      const TuneContext& context, const Array<MeasureCandidate>& candidates)>;
   /*!
    * \brief Get the feature extractor as string with name.
    * \return The string of the feature extractor.
@@ -75,10 +75,10 @@ class PyFeatureExtractorNode : public FeatureExtractorNode {
     // `f_as_string` is not visited
   }
 
-  Array<tvm::runtime::NDArray> ExtractFrom(const TuneContext& tune_context,
+  Array<tvm::runtime::NDArray> ExtractFrom(const TuneContext& context,
                                            const Array<MeasureCandidate>& candidates) {
     ICHECK(f_extract_from != nullptr) << "PyFeatureExtractor's ExtractFrom method not implemented!";
-    return f_extract_from(tune_context, candidates);
+    return f_extract_from(context, candidates);
   }
 
   static constexpr const char* _type_key = "meta_schedule.PyFeatureExtractor";

include/tvm/meta_schedule/mutator.h

@@ -0,0 +1,146 @@
+/*
+ * 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.
+ */
+
+#ifndef TVM_META_SCHEDULE_MUTATOR_H_
+#define TVM_META_SCHEDULE_MUTATOR_H_
+
+#include <tvm/tir/schedule/schedule.h>
+
+namespace tvm {
+namespace meta_schedule {
+
+class TuneContext;
+
+/*! \brief Mutator is designed to mutate the trace to explore the design space. */
+class MutatorNode : public runtime::Object {
+ public:
+  /*! \brief Virtual destructor. */
+  virtual ~MutatorNode() = default;
+
+  void VisitAttrs(tvm::AttrVisitor* v) {}
+
+  /*!
+   * \brief Initialize the design space generator with tuning context.
+   * \param context The tuning context for initialization.
+   * \note This method is supposed to be called only once before every other method.
+   */
+  virtual void InitializeWithTuneContext(const TuneContext& context) = 0;
+
+  /*!
+   * \brief Apply the mutator function to the given trace.
+   * \param trace The given trace for mutation.
+   * \param rand_state The random state for mutation.
+   * \return None if mutator failed, otherwise return the mutated trace.
+   */
+  virtual Optional<tir::Trace> Apply(const tir::Trace& trace,
+                                     support::LinearCongruentialEngine::TRandState* rand_state) = 0;
+
+  static constexpr const char* _type_key = "meta_schedule.Mutator";
+  TVM_DECLARE_BASE_OBJECT_INFO(MutatorNode, Object);
+};
+
+/*! \brief The mutator with customized methods on the python-side. */
+class PyMutatorNode : public MutatorNode {
+ public:
+  /*!
+   * \brief The function type of `InitializeWithTuneContext` method.
+   * \param context The tuning context for initialization.
+   */
+  using FInitializeWithTuneContext = runtime::TypedPackedFunc<void(const TuneContext&)>;
+  /*!
+   * \brief Apply the mutator function to the given trace.
+   * \param trace The given trace for mutation.
+   * \return None if mutator failed, otherwise return the mutated trace.
+   */
+  using FApply = runtime::TypedPackedFunc<Optional<tir::Trace>(
+      const tir::Trace&, support::LinearCongruentialEngine::TRandState rand_state)>;
+  /*!
+   * \brief Get the mutator as string with name.
+   * \return The string of the mutator.
+   */
+  using FAsString = runtime::TypedPackedFunc<String()>;
+
+  /*! \brief The packed function to the `InitializeWithTuneContext` function. */
+  FInitializeWithTuneContext f_initialize_with_tune_context;
+  /*! \brief The packed function to the `Apply` function. */
+  FApply f_apply;
+  /*! \brief The packed function to the `AsString` function. */
+  FAsString f_as_string;
+
+  void VisitAttrs(tvm::AttrVisitor* v) {
+    // `f_initialize_with_tune_context` is not visited
+    // `f_apply` is not visited
+    // `f_as_string` is not visited
+  }
+
+  void InitializeWithTuneContext(const TuneContext& context) final {
+    ICHECK(f_initialize_with_tune_context != nullptr)
+        << "PyMutator's InitializeWithTuneContext method not implemented!";
+    this->f_initialize_with_tune_context(context);
+  }
+
+  Optional<tir::Trace> Apply(const tir::Trace& trace,
+                             support::LinearCongruentialEngine::TRandState* rand_state) final {
+    ICHECK(f_apply != nullptr) << "PyMutator's Apply method not implemented!";
+    return this->f_apply(trace, *rand_state);
+  }
+
+  static constexpr const char* _type_key = "meta_schedule.PyMutator";
+  TVM_DECLARE_FINAL_OBJECT_INFO(PyMutatorNode, MutatorNode);
+};
+
+/*!
+ * \brief Managed reference to MutatorNode
+ * \sa MutatorNode
+ */
+class Mutator : public runtime::ObjectRef {
+ public:
+  /*! \brief Create a Mutator that mutates the tile size. */
+  TVM_DLL static Mutator MutateTileSize();
+  /*!
+   * \brief Create a Mutator that mutates the parallel extent
+   * \param max_jobs_per_core The maximum number of parallel jobs per core.
+   * \return The created mutator.
+   */
+  TVM_DLL static Mutator MutateParallel(int64_t max_jobs_per_core);
+  /*! \brief Create a Mutator that mutates auto unroll step */
+  TVM_DLL static Mutator MutateUnroll();
+  /*!
+   * \brief Create a Mutator that mutates the outcome of SampleComputeLocation
+   * \return The mutator created
+   */
+  TVM_DLL static Mutator MutateComputeLocation();
+  /*!
+   * \brief Create a mutator with customized methods on the python-side.
+   * \param f_initialize_with_tune_context The packed function of `InitializeWithTuneContext`.
+   * \param f_apply The packed function of `Apply`.
+   * \param f_as_string The packed function of `AsString`.
+   * \return The mutator created.
+   */
+  TVM_DLL static Mutator PyMutator(
+      PyMutatorNode::FInitializeWithTuneContext f_initialize_with_tune_context,  //
+      PyMutatorNode::FApply f_apply,                                             //
+      PyMutatorNode::FAsString f_as_string);
+  TVM_DEFINE_MUTABLE_OBJECT_REF_METHODS(Mutator, ObjectRef, MutatorNode);
+};
+
+}  // namespace meta_schedule
+}  // namespace tvm
+
+#endif  // TVM_META_SCHEDULE_MUTATOR_H_