[Meta Schedule] Feature Extractor & Cost Model

gallery/how_to/extend_tvm/bring_your_own_datatypes.py

@@ -313,7 +313,7 @@ def convert_ndarray(dst_dtype, array):
     print(str(e).split("\n")[-1])
 
 ######################################################################
-# When we attempt to run the model, we get a familiar error telling us that more funcions need to be registerd for myfloat.
+# When we attempt to run the model, we get a familiar error telling us that more functions need to be registerd for myfloat.
 #
 # Because this is a neural network, many more operations are required.
 # Here, we register all the needed functions:

include/tvm/auto_scheduler/cost_model.h

@@ -122,11 +122,11 @@ class RandomModel : public CostModel {
  *  This class will call functions defined in the python */
 class PythonBasedModelNode : public CostModelNode {
  public:
-  /*! \brief Pointer to the update funcion in python */
+  /*! \brief Pointer to the update function in python */
   PackedFunc update_func;
-  /*! \brief Pointer to the predict funcion in python */
+  /*! \brief Pointer to the predict function in python */
   PackedFunc predict_func;
-  /*! \brief Pointer to the predict funcion in python */
+  /*! \brief Pointer to the predict function in python */
   PackedFunc predict_stage_func;
 
   void Update(const Array<MeasureInput>& inputs, const Array<MeasureResult>& results) final;

include/tvm/auto_scheduler/measure.h

@@ -236,7 +236,7 @@ class MeasureCallback : public ObjectRef {
  *  This class will call functions defined in the python */
 class PythonBasedMeasureCallbackNode : public MeasureCallbackNode {
  public:
-  /*! \brief Pointer to the callback funcion in python */
+  /*! \brief Pointer to the callback function in python */
   PackedFunc callback_func;
 
   void Callback(const SearchPolicy& policy, const Array<MeasureInput>& inputs,

include/tvm/meta_schedule/cost_model.h

@@ -0,0 +1,185 @@
+/*
+ * 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_COST_MODEL_H_
+#define TVM_META_SCHEDULE_COST_MODEL_H_
+
+#include <tvm/meta_schedule/search_strategy.h>
+
+namespace tvm {
+namespace meta_schedule {
+
+class TuneContext;
+
+/*! \brief Cost model. */
+class CostModelNode : public runtime::Object {
+ public:
+  /*! \brief Virtual destructor. */
+  virtual ~CostModelNode() = default;
+
+  void VisitAttrs(tvm::AttrVisitor* v) {}
+
+  /*!
+   * \brief Load the cost model from given file location.
+   * \param file_location The file location.
+   * \return Whether cost model was loaded successfully.
+   */
+  virtual bool Load(const String& file_location) = 0;
+
+  /*!
+   * \brief Save the cost model to given file location.
+   * \param file_location The file location.
+   * \return Whether cost model was saved successfully.
+   */
+  virtual bool Save(const String& file_location) = 0;
+
+  /*!
+   * \brief Update the cost model given running results.
+   * \param tune_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,
+                      const Array<RunnerResult>& results) = 0;
+
+  /*!
+   * \brief Predict the running results of given measure candidates.
+   * \param tune_context The tuning context.
+   * \param candidates The measure candidates.
+   * \return The predicted running results.
+   */
+  virtual std::vector<double> Predict(const TuneContext& tune_context,
+                                      const Array<MeasureCandidate>& candidates) = 0;
+
+  static constexpr const char* _type_key = "meta_schedule.CostModel";
+  TVM_DECLARE_BASE_OBJECT_INFO(CostModelNode, Object);
+};
+
+/*! \brief The cost model with customized methods on the python-side. */
+class PyCostModelNode : public CostModelNode {
+ public:
+  /*!
+   * \brief Load the cost model from given file location.
+   * \param file_location The file location.
+   * \return Whether cost model was loaded successfully.
+   */
+  using FLoad = runtime::TypedPackedFunc<bool(String)>;
+  /*!
+   * \brief Save the cost model to given file location.
+   * \param file_location The file location.
+   * \return Whether cost model was saved successfully.
+   */
+  using FSave = runtime::TypedPackedFunc<bool(String)>;
+  /*!
+   * \brief Update the cost model given running results.
+   * \param tune_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.
+   */
+  using FUpdate = runtime::TypedPackedFunc<void(const TuneContext&, const Array<MeasureCandidate>&,
+                                                const Array<RunnerResult>&)>;
+  /*!
+   * \brief Predict the running results of given measure candidates.
+   * \param tune_context The tuning context.
+   * \param candidates The measure candidates.
+   * \param p_addr The address to save the the estimated running results.
+   */
+  using FPredict = runtime::TypedPackedFunc<void(const TuneContext&, const Array<MeasureCandidate>&,
+                                                 void* p_addr)>;
+  /*!
+   * \brief Get the cost model as string with name.
+   * \return The string representation of the cost model.
+   */
+  using FAsString = runtime::TypedPackedFunc<String()>;
+
+  /*! \brief The packed function to the `Load` function. */
+  FLoad f_load;
+  /*! \brief The packed function to the `Save` function. */
+  FSave f_save;
+  /*! \brief The packed function to the `Update` function. */
+  FUpdate f_update;
+  /*! \brief The packed function to the `Predict` function. */
+  FPredict f_predict;
+  /*! \brief The packed function to the `AsString` function. */
+  FAsString f_as_string;
+
+  void VisitAttrs(tvm::AttrVisitor* v) {
+    // `f_load` is not visited
+    // `f_save` is not visited
+    // `f_update` is not visited
+    // `f_predict` is not visited
+    // `f_as_string` is not visited
+  }
+
+  bool Load(const String& file_location) {
+    ICHECK(f_load != nullptr) << "PyCostModel's Load method not implemented!";
+    return f_load(file_location);
+  }
+
+  bool Save(const String& file_location) {
+    ICHECK(f_save != nullptr) << "PyCostModel's Save method not implemented!";
+    return f_save(file_location);
+  }
+
+  void Update(const TuneContext& tune_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);
+  }
+
+  std::vector<double> Predict(const TuneContext& tune_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());
+    return result;
+  }
+
+  static constexpr const char* _type_key = "meta_schedule.PyCostModel";
+  TVM_DECLARE_FINAL_OBJECT_INFO(PyCostModelNode, CostModelNode);
+};
+
+/*!
+ * \brief Managed reference to CostModelNode
+ * \sa CostModelNode
+ */
+class CostModel : public runtime::ObjectRef {
+ public:
+  /*!
+   * \brief Create a feature extractor with customized methods on the python-side.
+   * \param f_load The packed function of `Load`.
+   * \param f_save The packed function of `Save`.
+   * \param f_update The packed function of `Update`.
+   * \param f_predict The packed function of `Predict`.
+   * \param f_as_string The packed function of `AsString`.
+   * \return The feature extractor created.
+   */
+  TVM_DLL static CostModel PyCostModel(PyCostModelNode::FLoad f_load,        //
+                                       PyCostModelNode::FSave f_save,        //
+                                       PyCostModelNode::FUpdate f_update,    //
+                                       PyCostModelNode::FPredict f_predict,  //
+                                       PyCostModelNode::FAsString f_as_string);
+  TVM_DEFINE_MUTABLE_OBJECT_REF_METHODS(CostModel, ObjectRef, CostModelNode);
+};
+
+}  // namespace meta_schedule
+}  // namespace tvm
+
+#endif  // TVM_META_SCHEDULE_COST_MODEL_H_

include/tvm/meta_schedule/feature_extractor.h

@@ -0,0 +1,109 @@
+/*
+ * 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_FEATURE_EXTRACTOR_H_
+#define TVM_META_SCHEDULE_FEATURE_EXTRACTOR_H_
+
+#include <tvm/meta_schedule/search_strategy.h>
+
+namespace tvm {
+namespace meta_schedule {
+
+class TuneContext;
+
+/*! \brief Extractor for features from measure candidates for use in cost model. */
+class FeatureExtractorNode : public runtime::Object {
+ public:
+  /*! \brief Virtual destructor. */
+  virtual ~FeatureExtractorNode() = default;
+
+  void VisitAttrs(tvm::AttrVisitor* v) {}
+
+  /*!
+   * \brief Extract features from the given measure candidate.
+   * \param tune_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,
+                                                   const Array<MeasureCandidate>& candidates) = 0;
+
+  static constexpr const char* _type_key = "meta_schedule.FeatureExtractor";
+  TVM_DECLARE_BASE_OBJECT_INFO(FeatureExtractorNode, Object);
+};
+
+/*! \brief The feature extractor with customized methods on the python-side. */
+class PyFeatureExtractorNode : public FeatureExtractorNode {
+ public:
+  /*!
+   * \brief Extract features from the given measure candidate.
+   * \param tune_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)>;
+  /*!
+   * \brief Get the feature extractor as string with name.
+   * \return The string of the feature extractor.
+   */
+  using FAsString = runtime::TypedPackedFunc<String()>;
+
+  /*! \brief The packed function to the `ExtractFrom` function. */
+  FExtractFrom f_extract_from;
+  /*! \brief The packed function to the `AsString` function. */
+  FAsString f_as_string;
+
+  void VisitAttrs(tvm::AttrVisitor* v) {
+    // `f_extract_from` is not visited
+    // `f_as_string` is not visited
+  }
+
+  Array<tvm::runtime::NDArray> ExtractFrom(const TuneContext& tune_context,
+                                           const Array<MeasureCandidate>& candidates) {
+    ICHECK(f_extract_from != nullptr) << "PyFeatureExtractor's ExtractFrom method not implemented!";
+    return f_extract_from(tune_context, candidates);
+  }
+
+  static constexpr const char* _type_key = "meta_schedule.PyFeatureExtractor";
+  TVM_DECLARE_FINAL_OBJECT_INFO(PyFeatureExtractorNode, FeatureExtractorNode);
+};
+
+/*!
+ * \brief Managed reference to FeatureExtractorNode
+ * \sa FeatureExtractorNode
+ */
+class FeatureExtractor : public runtime::ObjectRef {
+ public:
+  /*!
+   * \brief Create a feature extractor with customized methods on the python-side.
+   * \param f_extract_from The packed function of `ExtractFrom`.
+   * \param f_as_string The packed function of `AsString`.
+   * \return The feature extractor created.
+   */
+  TVM_DLL static FeatureExtractor PyFeatureExtractor(
+      PyFeatureExtractorNode::FExtractFrom f_extract_from,  //
+      PyFeatureExtractorNode::FAsString f_as_string);
+  TVM_DEFINE_MUTABLE_OBJECT_REF_METHODS(FeatureExtractor, ObjectRef, FeatureExtractorNode);
+};
+
+}  // namespace meta_schedule
+}  // namespace tvm
+
+#endif  // TVM_META_SCHEDULE_FEATURE_EXTRACTOR_H_

include/tvm/meta_schedule/measure_callback.h

@@ -81,9 +81,9 @@ class PyMeasureCallbackNode : public MeasureCallbackNode {
    */
   using FAsString = runtime::TypedPackedFunc<String()>;
 
-  /*! \brief The packed function to the `Apply` funcion. */
+  /*! \brief The packed function to the `Apply` function. */
   FApply f_apply;
-  /*! \brief The packed function to the `AsString` funcion. */
+  /*! \brief The packed function to the `AsString` function. */
   FAsString f_as_string;
 
   void VisitAttrs(tvm::AttrVisitor* v) {

include/tvm/meta_schedule/mutator.h

@@ -36,8 +36,9 @@ class MutatorNode : public runtime::Object {
   void VisitAttrs(tvm::AttrVisitor* v) {}
 
   /*!
-   * \brief The function type of `InitializeWithTuneContext` method.
+   * \brief Initialize the design space generator with tuning context.
    * \param tune_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;
 
@@ -72,11 +73,11 @@ class PyMutatorNode : public MutatorNode {
    */
   using FAsString = runtime::TypedPackedFunc<String()>;
 
-  /*! \brief The packed function to the `InitializeWithTuneContext` funcion. */
+  /*! \brief The packed function to the `InitializeWithTuneContext` function. */
   FInitializeWithTuneContext f_initialize_with_tune_context;
-  /*! \brief The packed function to the `Apply` funcion. */
+  /*! \brief The packed function to the `Apply` function. */
   FApply f_apply;
-  /*! \brief The packed function to the `AsString` funcion. */
+  /*! \brief The packed function to the `AsString` function. */
   FAsString f_as_string;
 
   void VisitAttrs(tvm::AttrVisitor* v) {
@@ -110,6 +111,7 @@ class Mutator : public runtime::ObjectRef {
    * \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(

include/tvm/meta_schedule/postproc.h

@@ -38,8 +38,9 @@ class PostprocNode : public runtime::Object {
   void VisitAttrs(tvm::AttrVisitor* v) {}
 
   /*!
-   * \brief The function type of `InitializeWithTuneContext` method.
+   * \brief Initialize the design space generator with tuning context.
    * \param tune_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;
 
@@ -112,6 +113,7 @@ class Postproc : public runtime::ObjectRef {
    * \brief Create a postprocessor 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 postprocessor created.
    */
   TVM_DLL static Postproc PyPostproc(