[Ansor][AutoTVM v2.0] Phase 2: Layout Rewrite in AutoScheduler (apach…

…e#6297)

* enable layout rewrite for auto scheduler

* refine

* update

* fix CI

* fix CI

* fix CI

* resolve review comments

* add ut

* resolve comments

* resolve comments

* fix coding style

include/tvm/auto_scheduler/compute_dag.h

@@ -205,19 +205,29 @@ class ComputeDAG : public ObjectRef {
    */
   TVM_DLL explicit ComputeDAG(Array<te::Tensor> tensors);
 
+  /*!
+   * \brief Rewrite the layout of placeholder specified by attr `layout_free_placeholders`
+   * according to the loop nest derived with `transform_steps`.
+   * \param transform_steps Transform steps of a state.
+   */
+  void RewriteLayout(const Array<Step>& transform_steps);
+
   /*!
    * \brief Apply the history transform steps to get a TVM schedule.
    * \param transform_steps Transform steps of a state.
    * \param stages The list of stages after applying the steps.
    * Pass a valid pointer if this information needs to be used outside this function.
    * \param stage_to_axes The map that stores all axes for one stage.
    * Pass a valid pointer if this information needs to be used outside this function.
+   * \param layout_rewrite Rewrite the layout of placeholders specified by
+   * attr `layout_free_placeholders`
    * \return A `te.schedule` and the an Array of `te.Tensor` to be used in `tvm.lower`
    * or `tvm.build`.
    */
-  std::pair<te::Schedule, Array<te::Tensor>> ApplySteps(
-      const Array<Step>& transform_steps, Array<te::Stage>* stages = nullptr,
-      StageToAxesMap* stage_to_axes = nullptr) const;
+  std::pair<te::Schedule, Array<te::Tensor>> ApplySteps(const Array<Step>& transform_steps,
+                                                        Array<te::Stage>* stages = nullptr,
+                                                        StageToAxesMap* stage_to_axes = nullptr,
+                                                        bool layout_rewrite = false) const;
 
   /*!
    * \brief Print transform steps as equivalent python schedule API.
@@ -262,6 +272,8 @@ class ComputeDAG : public ObjectRef {
    */
   ComputeDAG ReplayAndGetDAG(const Array<Step>& steps) const;
 
+  static constexpr const char* layout_free_placeholders_key = "layout_free_placeholders";
+
   TVM_DEFINE_OBJECT_REF_METHODS(ComputeDAG, ObjectRef, ComputeDAGNode);
   TVM_DEFINE_OBJECT_REF_COW_METHOD(ComputeDAGNode);
 };

include/tvm/auto_scheduler/transform_step.h

@@ -50,6 +50,8 @@
 #include <tvm/node/node.h>
 #include <tvm/te/schedule.h>
 
+#include <vector>
+
 namespace tvm {
 namespace auto_scheduler {
 
@@ -104,6 +106,9 @@ enum class IteratorAnnotation : int {
 
 extern const char* IteratorAnnotationString[];
 
+// forward declaration
+class Iterator;
+
 /*!
  * \brief An iterator of a for-loop
  * Similar to tvm::IterVar in `include/tvm/tir/expr.h`
@@ -118,6 +123,8 @@ class IteratorNode : public Object {
   IteratorKind iter_kind;
   /*! \brief The annotation type of this iterator. */
   IteratorAnnotation annotation;
+  /*! The original iterators before fusion. */
+  std::vector<Iterator> orig_iters;
 
   void VisitAttrs(tvm::AttrVisitor* v) {
     v->Visit("name", &name);
@@ -142,8 +149,10 @@ class Iterator : public ObjectRef {
    * \param range The range of this iterator.
    * \param iter_kind The iterator type of this iterator.
    * \param annotation The annotation type of this iterator.
+   * \param orig_iters The original iterators before fusion
    */
-  Iterator(String name, Range range, IteratorKind iter_kind, IteratorAnnotation annotation);
+  Iterator(String name, Range range, IteratorKind iter_kind, IteratorAnnotation annotation,
+           const std::vector<Iterator>* orig_iters = nullptr);
 
   TVM_DEFINE_OBJECT_REF_METHODS(Iterator, ObjectRef, IteratorNode);
 };

python/tvm/auto_scheduler/compute_dag.py

@@ -72,7 +72,7 @@ def get_init_state(self):
         """
         return State(self.init_state, self)
 
-    def apply_steps_from_state(self, state):
+    def apply_steps_from_state(self, state, layout_rewrite=False):
         """
         Apply the history transform steps from a State to get a TVM schedule.
 
@@ -81,12 +81,16 @@ def apply_steps_from_state(self, state):
         state : Union[State, StateObject]
             The state from which we get transform steps.
 
+        layout_rewrite: Bool
+            Rewrite the layout of placeholders specified by "layout_free_placeholders" attr
+            to make it most friendly for the generated schedule to read from.
+
         Returns
         -------
             A `te.schedule` and the a list of `te.Tensor` to be used in `tvm.lower` or `tvm.build`.
         """
         state_obj = state if isinstance(state, StateObject) else state.state_object
-        return _ffi_api.ComputeDAGApplyStepsFromState(self, state_obj)
+        return _ffi_api.ComputeDAGApplyStepsFromState(self, state_obj, layout_rewrite)
 
     def print_python_code_from_state(self, state):
         """

python/tvm/auto_scheduler/measure.py

@@ -435,7 +435,7 @@ def timed_func():
 
         try:
             sch, args = task.compute_dag.apply_steps_from_state(
-                inp.state)
+                inp.state, layout_rewrite=True)
         # pylint: disable=broad-except
         except Exception:
             error_no = MeasureErrorNo.INSTANTIATION_ERROR