Squashed commit

[Meta Schedule][M3c] Schedule Rules, Mutator & Postprocs (#485)

[Meta Schedule][M3c] PostOrderApply (#486)

Fix Post Order Apply (#490)

[MetaSchedule] Relay Integration (#489)

[M3c][Meta Schedule] Add Trace Correctness Test for PostOrderApply (#492)

Fix replay trace. (#493)

[M3c][Meta Schedule] Implement the Replay Func class. (#495)

[PR] Test script for meta-schedule task extraction. Interface to load… (#494)

[Meta Schedule Refactor] Get child blocks (#500)

Read-at && Write-at (#497)

[M3c][Meta Schedule] Measure Callbacks (#498)

[Bug] Fix Infinite Loop Caused When Calling Methods Not Overrided In PyClass (#496)

[MetaSchedule] Sample-Perfect-Tile (#501)

[MetaSchedule] TE Workloads (#502)

[TensorIR] GetProducer, GetConsumer (#506)

[MetaScheduleRefactor] Annotate&Unannotate (#505)

[MetaSchedule] Multi-Level-Tiling & Auto-Inline (#503)

[Tests] Add unittests for auto-inline and multi-level-tiling (#508)

[Meta Schedule] Minor Fixes (#507)

[MetaSchedule] Rewrite Cooperative-Fetching / Unbound-Block / Reduction-Block (#509)

[MetaSchedule] Rewrite Parallel-Vectorize-Unroll / Verify-GPU / Disallow-Dynamic-Loops (#499)

[Meta Schedule] Add Helper Function & Minor Modification (#512)

[MetaSchedule] Test for Rewrite Parallel-Vectorize-Unroll  (#513)

[Meta Schedule] Feature Extractor & Cost Model (#510)

Blockize & Tensorize (#514)

Layout Rewriting: Suggest-Index-Map (#520)

Co-authored-by: Siyuan Feng <Hzfengsy@sjtu.edu.cn>
Co-authored-by: Bohan Hou <32121147+spectrometerHBH@users.noreply.github.com>
Co-authored-by: Hongyi Jin <3231950289@qq.com>
Co-authored-by: Ruihang Lai <lairuihangdongdong@qq.com>
Co-authored-by: Junru Shao <junrushao1994@gmail.com>
Co-authored-by: Wuwei Lin <wuwei@apache.org>
Co-authored-by: Sunghyun Park <49998730+sunggg@users.noreply.github.com>

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/arith/iter_affine_map.h

@@ -346,6 +346,8 @@ Array<Array<IterMark>> SubspaceDivide(const Array<PrimExpr>& bindings,
                                       const Array<Var>& sub_iters, const PrimExpr& predicate,
                                       bool require_bijective, arith::Analyzer* analyzer);
 
+PrimExpr NormalizeIterMapToExpr(const IterMapExpr& expr);
+
 }  // namespace arith
 }  // namespace tvm
 #endif  // TVM_ARITH_ITER_AFFINE_MAP_H_

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,187 @@
+/*
+ * 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 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_