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)

[MetaSchedule] Parallel-Vectorize-Unroll & Random-Compute-Location (#516)

[Meta Schedule] Per-Store-Feature (#521)

Add traced schedule for blockize & tensorize (#526)

[Meta Schedule] Add XGBoost Model & Random Model (#519)

User-Interface: Tune-TIR (#525)

User-Interface: Tune-TE (#527)

[Minor] More logging on python (#528)

Get CUDA tuning working (#529)

[MetaSchedule] TensorRT BYOC (#518)

[BugFix] LocalBuilder API (#531)

[Meta Schedule] Add Cost Model Update Measure Callback (#530)

[Bugfix] BuilderInput with default params (#532)

[MetaSchedule] Mutator-Tile-Size, Mutate-Parallel, Mutate-Unroll (#534)

[Meta Schedule] Evolutionary Search (#522)

[BugFix] Remove duplicated definition of MakeMultinomialSampler (#535)

[Meta Schedule] Fix some bugs (#537)

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/builder.h

@@ -32,10 +32,13 @@ class BuilderInputNode : public runtime::Object {
   IRModule mod;
   /*! \brief The target to be built for. */
   Target target;
+  /*! \brief Parameters for Relay build module. */
+  Optional<Map<String, runtime::NDArray>> params;
 
   void VisitAttrs(tvm::AttrVisitor* v) {
     v->Visit("mod", &mod);
     v->Visit("target", &target);
+    v->Visit("params", &params);
   }
 
   static constexpr const char* _type_key = "meta_schedule.BuilderInput";
@@ -52,8 +55,10 @@ class BuilderInput : public runtime::ObjectRef {
    * \brief Constructor of BuilderInput.
    * \param mod The IRModule to be built.
    * \param target The target to be built for.
+   * \param params Parameters for Relay build module.
    */
-  TVM_DLL explicit BuilderInput(IRModule mod, Target target);
+  TVM_DLL explicit BuilderInput(IRModule mod, Target target,
+                                Optional<Map<String, runtime::NDArray>> params = NullOpt);
   TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(BuilderInput, runtime::ObjectRef, BuilderInputNode);
 };
 

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_

include/tvm/meta_schedule/database.h

@@ -155,6 +155,12 @@ class DatabaseNode : public runtime::Object {
  public:
   /*! \brief Default destructor */
   virtual ~DatabaseNode() = default;
+  /*!
+   * \brief Check if the database has the given workload.
+   * \param mod The IRModule to be searched for.
+   * \return Whether the database has the given workload.
+   */
+  virtual bool HasWorkload(const IRModule& mod) = 0;
   /*!
    * \brief Look up or add workload to the database if missing.
    * \param mod The IRModule to be searched for or added.
@@ -186,6 +192,12 @@ class DatabaseNode : public runtime::Object {
 /*! \brief The database with customized methods on the python-side. */
 class PyDatabaseNode : public DatabaseNode {
  public:
+  /*!
+   * \brief The function type of `HasWorkload` method.
+   * \param mod The IRModule to be searched for.
+   * \return Whether the database has the given workload.
+   */
+  using FHasWorkload = runtime::TypedPackedFunc<bool(const IRModule&)>;
   /*!
    * \brief The function type of `CommitWorkload` method.
    * \param mod The IRModule to be searched for or added.
@@ -210,6 +222,8 @@ class PyDatabaseNode : public DatabaseNode {
    */
   using FSize = runtime::TypedPackedFunc<int64_t()>;
 
+  /*! \brief The packed function to the `HasWorkload` function. */
+  FHasWorkload f_has_workload;
   /*! \brief The packed function to the `CommitWorkload` function. */
   FCommitWorkload f_commit_workload;
   /*! \brief The packed function to the `CommitTuningRecord` function. */
@@ -224,12 +238,18 @@ class PyDatabaseNode : public DatabaseNode {
     // so it cannot be accessible on the python side. If there is such need from the future,
     // we can then add corresponding accessor methods to help access on python.
     //
+    // `f_has_workload` is not visited
     // `f_commit_workload` is not visited
     // `f_commit_tuning_record` is not visited
     // `f_get_top_k` is not visited
     // `f_size` is not visited
   }
 
+  bool HasWorkload(const IRModule& mod) final {
+    ICHECK(f_has_workload != nullptr) << "PyDatabase's HasWorkload method not implemented!";
+    return f_has_workload(mod);
+  }
+
   Workload CommitWorkload(const IRModule& mod) final {
     ICHECK(f_commit_workload != nullptr) << "PyDatabase's CommitWorkload method not implemented!";
     return f_commit_workload(mod);
@@ -271,13 +291,15 @@ class Database : public runtime::ObjectRef {
                                        bool allow_missing);
   /*!
    * \brief Create a database with customized methods on the python-side.
+   * \param f_has_workload The packed function of `HasWorkload`.
    * \param f_commit_workload The packed function of `CommitWorkload`.
    * \param f_commit_tuning_record The packed function of `CommitTuningRecord`.
    * \param f_get_top_k The packed function of `GetTopK`.
    * \param f_size The packed function of `Size`.
    * \return The created database.
    */
-  TVM_DLL static Database PyDatabase(PyDatabaseNode::FCommitWorkload f_commit_workload,
+  TVM_DLL static Database PyDatabase(PyDatabaseNode::FHasWorkload f_has_workload,
+                                     PyDatabaseNode::FCommitWorkload f_commit_workload,
                                      PyDatabaseNode::FCommitTuningRecord f_commit_tuning_record,
                                      PyDatabaseNode::FGetTopK f_get_top_k,
                                      PyDatabaseNode::FSize f_size);