[MetaSchedule] Random Feature Extractor (#9760)

Co-authored-by: Junru Shao <junrushao1994@gmail.com> 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> 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>
apache · Dec 17, 2021 · e1255c9 · e1255c9
1 parent a374cdd
commit e1255c9
Show file tree

Hide file tree

Showing 16 changed files with 1,137 additions and 3 deletions.
diff --git a/include/tvm/meta_schedule/cost_model.h b/include/tvm/meta_schedule/cost_model.h
@@ -0,0 +1,182 @@
+/*
+ * 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>
+
+#include <vector>
+
+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 path The file path.
+   */
+  virtual void Load(const String& path) = 0;
+
+  /*!
+   * \brief Save the cost model to given file location.
+   * \param path The file path.
+   */
+  virtual void Save(const String& path) = 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 normalized score (the larger the better) of given measure candidates.
+   * \param tune_context The tuning context.
+   * \param candidates The measure candidates.
+   * \return The predicted normalized score.
+   */
+  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 path The file path.
+   */
+  using FLoad = runtime::TypedPackedFunc<void(String)>;
+  /*!
+   * \brief Save the cost model to given file location.
+   * \param path The file path.
+   */
+  using FSave = runtime::TypedPackedFunc<void(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
+  }
+
+  void Load(const String& path) {
+    ICHECK(f_load != nullptr) << "PyCostModel's Load method not implemented!";
+    f_load(path);
+  }
+
+  void Save(const String& path) {
+    ICHECK(f_save != nullptr) << "PyCostModel's Save method not implemented!";
+    f_save(path);
+  }
+  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_
diff --git a/include/tvm/meta_schedule/feature_extractor.h b/include/tvm/meta_schedule/feature_extractor.h
@@ -0,0 +1,121 @@
+/*
+ * 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 that extracts features from each BufferStore
+   * \param buffers_per_store The number of buffers in each BufferStore; Pad or truncate if
+   * necessary.
+   * \param arith_intensity_curve_num_samples The number of samples used in the arithmetic intensity
+   * curve.
+   * \param cache_line_bytes The number of bytes in a cache line.
+   * \return The feature extractor created.
+   */
+  TVM_DLL static FeatureExtractor PerStoreFeature(int buffers_per_store = 5,
+                                                  int arith_intensity_curve_num_samples = 10,
+                                                  int cache_line_bytes = 64);
+  /*!
+   * \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_
diff --git a/python/tvm/meta_schedule/cost_model/__init__.py b/python/tvm/meta_schedule/cost_model/__init__.py
@@ -0,0 +1,21 @@
+# 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.
+"""
+The tvm.meta_schedule.cost_model package.
+"""
+from .cost_model import CostModel, PyCostModel
+from .random_model import RandomModel