schedule.h

Go to the documentation of this file.

/*
 * 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.
 */

// Acknowledgement: Many schedule primitives originate from Halide and Loopy.
#ifndef TVM_TE_SCHEDULE_H_
#define TVM_TE_SCHEDULE_H_

#include <tvm/support/with.h>
#include <tvm/te/tensor.h>
#include <tvm/te/tensor_intrin.h>
#include <tvm/tir/expr.h>
#include <tvm/tir/index_map.h>

#include <string>
#include <unordered_map>

namespace tvm {
namespace te {
// Node container for Stage
class StageNode;
// Node container for Schedule
class ScheduleNode;
// Node container for IterVarRelation
class IterVarRelationNode;
// Attribute of itervar.
class IterVarAttrNode;

enum AttachType : int {
  kGroupRoot = 1,
  kInline = 2,
  kInlinedAlready = 3,
  kScope = 4,
  kScanUpdate = 5
};

class Stage : public ObjectRef {
 public:
  Stage() {}
  explicit Stage(ObjectPtr<Object> n) : ObjectRef(n) {}
  explicit Stage(Operation op, const ScheduleNode* sch);
  inline const StageNode* operator->() const;
  inline StageNode* operator->();
  TVM_DLL Stage& set_scope(std::string scope);  // NOLINT(*)
  TVM_DLL Stage& compute_at(Stage parent, IterVar scope);  // NOLINT(*)
  TVM_DLL Stage& compute_inline();  // NOLINT(*)
  TVM_DLL Stage& compute_root();  // NOLINT(*)
  TVM_DLL Stage& bind(IterVar ivar, IterVar thread_ivar);
  TVM_DLL Stage& set_store_predicate(PrimExpr predicate);
  TVM_DLL Stage& env_threads(Array<IterVar> threads);
  TVM_DLL Stage& split(IterVar parent, PrimExpr factor, IterVar* p_outer,
                       IterVar* p_inner);  // NOLINT(*)
  TVM_DLL Stage& split_by_nparts(IterVar parent, PrimExpr nparts, IterVar* p_outer,
                                 IterVar* p_inner);  // NOLINT(*)
  TVM_DLL Stage& fuse(IterVar outer, IterVar inner, IterVar* p_target);  // NOLINT(*)
  TVM_DLL Stage& fuse(const Array<IterVar>& axes, IterVar* p_target);  // NOLINT(*)
  TVM_DLL Stage& reorder(const Array<IterVar>& order);  // NOLINT(*)
  TVM_DLL Stage& tile(IterVar x_parent, IterVar y_parent,  // NOLINT(*)
                      PrimExpr x_factor, PrimExpr y_factor, IterVar* p_x_outer, IterVar* p_y_outer,
                      IterVar* p_x_inner, IterVar* p_y_inner);
  TVM_DLL Stage& vectorize(IterVar var);  // NOLINT(*)
  TVM_DLL Stage& tensorize(IterVar var, TensorIntrin f);  // NOLINT(*)
  TVM_DLL Stage& unroll(IterVar var);  // NOLINT(*)
  TVM_DLL Stage& parallel(IterVar var);  // NOLINT(*)
  TVM_DLL Stage& pragma(IterVar var, const std::string& pragma_type,
                        const PrimExpr& pragma_value = PrimExpr());  // NOLINT(*)
  TVM_DLL Stage& prefetch(const Tensor& domain, IterVar var, PrimExpr offset);  // NOLINT(*)
  TVM_DLL Stage& storage_align(IterVar axis, int factor, int offset);  // NOLINT(*)
  TVM_DLL Stage& double_buffer();  // NOLINT(*)
  TVM_DLL Stage& rolling_buffer();  // NOLINT(*)
  TVM_DLL Stage& transform_layout(const Array<Var>& initial_indices,
                                  const Array<PrimExpr>& final_indices,
                                  Array<IterVar>* out_iter_vars = nullptr);
  TVM_DLL Stage& set_axis_separators(const Array<IntImm>& axis_separators);
  bool is_scheduled() const;
  Stage GetAttachSpec() const;
  // declare container type
  using ContainerType = StageNode;
};

class Schedule : public ObjectRef {
 public:
  Schedule() {}
  explicit Schedule(ObjectPtr<Object> n) : ObjectRef(n) {}
  TVM_DLL explicit Schedule(Array<Operation> ops);
  Schedule copy() const;
  TVM_DLL Stage operator[](const Operation& op);
  TVM_DLL Stage operator[](const Tensor& tensor) { return this->operator[](tensor->op); }
  TVM_DLL Stage create_group(const Array<Tensor>& outputs, const Array<Tensor>& inputs,
                             bool include_inputs = false);
  TVM_DLL Tensor cache_read(const Tensor& tensor, const std::string& scope,
                            const Array<Operation>& readers);
  TVM_DLL Array<Tensor> cache_write(const Array<Tensor>& tensor, const std::string& scope);
  TVM_DLL Tensor cache_write(const Tensor& tensor, const std::string& scope);
  TVM_DLL Array<Tensor> rfactor(const Tensor& tensor, const IterVar& axis, int factor_axis = 0);
  Schedule normalize();

  Schedule normalize_for_feature_extraction();

  inline const ScheduleNode* operator->() const;
  inline ScheduleNode* operator->();
  // declare container type
  using ContainerType = ScheduleNode;
};

class ScheduleContext {
 private:
  friend class With<ScheduleContext>;
  ScheduleContext(const ScheduleNode* sch_node, String current_primitive_name);
  void EnterWithScope();
  void ExitWithScope();

  Schedule sch_;
  String current_primitive_name_;
};

class IterVarRelation : public ObjectRef {
 public:
  IterVarRelation() {}
  explicit IterVarRelation(ObjectPtr<Object> n) : ObjectRef(n) {}
  inline const IterVarRelationNode* operator->() const;
};

class IterVarAttr : public ObjectRef {
 public:
  IterVarAttr() {}
  explicit IterVarAttr(ObjectPtr<Object> n) : ObjectRef(n) {}
  inline const IterVarAttrNode* operator->() const;
};

class StageNode : public Object {
 public:
  Operation op;
  Operation origin_op;
  Array<IterVar> all_iter_vars;
  Array<IterVar> leaf_iter_vars;
  Array<IterVar> env_threads;
  PrimExpr store_predicate;
  Array<IterVarRelation> relations;
  Map<IterVar, IterVarAttr> iter_var_attrs;
  AttachType attach_type{kGroupRoot};
  IterVar attach_ivar;
  Stage attach_stage;
  const ScheduleNode* attach_sch;
  std::string scope;
  bool is_output{false};
  bool double_buffer{false};
  bool rolling_buffer{false};
  Array<IndexMap> layout_transforms;
  Array<IntImm> axis_separators;
  Stage group;
  int num_child_stages{0};

  void VisitAttrs(AttrVisitor* v) {
    v->Visit("op", &op);
    v->Visit("origin_op", &origin_op);
    v->Visit("all_iter_vars", &all_iter_vars);
    v->Visit("leaf_iter_vars", &leaf_iter_vars);
    v->Visit("env_threads", &env_threads);
    v->Visit("relations", &relations);
    v->Visit("iter_var_attrs", &iter_var_attrs);
    v->Visit("attach_type", &attach_type);
    v->Visit("attach_ivar", &attach_ivar);
    v->Visit("attach_stage", &attach_stage);
    v->Visit("scope", &scope);
    v->Visit("is_output", &is_output);
    v->Visit("double_buffer", &double_buffer);
    v->Visit("layout_transforms", &layout_transforms);
    v->Visit("axis_separators", &axis_separators);
    v->Visit("group", &group);
    v->Visit("num_child_stages", &num_child_stages);
  }

  static constexpr const char* _type_key = "Stage";
  TVM_DECLARE_FINAL_OBJECT_INFO(StageNode, Object);
};

class ScheduleNode : public Object {
 public:
  Array<Operation> outputs;
  Array<Stage> stages;
  Array<Stage> groups;
  Map<Operation, Stage> stage_map;
  std::unordered_map<const Object*, Stage> op2stage_cache_;
  Array<Schedule> schedule_record;
  Array<String> primitive_record;
  Optional<Bool> keep_schedule_record;

  void VisitAttrs(AttrVisitor* v) {
    v->Visit("outputs", &outputs);
    v->Visit("stages", &stages);
    v->Visit("groups", &groups);
    v->Visit("stage_map", &stage_map);
    v->Visit("schedule_record", &schedule_record);
    v->Visit("primitive_record", &primitive_record);
    v->Visit("keep_schedule_record", &keep_schedule_record);
  }

  void InitCache();
  void InvalidateCache();

  TVM_DLL bool Contain(const Operation& op) const;

  TVM_DLL bool Contain(const Tensor& tensor) const { return Contain(tensor->op); }

  static constexpr const char* _type_key = "Schedule";
  TVM_DECLARE_FINAL_OBJECT_INFO(ScheduleNode, Object);
};

inline Schedule create_schedule(Array<Operation> ops) { return Schedule(ops); }

class IterVarAttrNode : public Object {
 public:
  IterVarType iter_type{kDataPar};
  IterVar bind_thread;
  Array<Tensor> prefetch_data;
  Array<PrimExpr> prefetch_offset;
  TensorIntrin tensor_intrin;
  int dim_align_factor{0};
  int dim_align_offset{0};
  Array<PrimExpr> pragma_keys;
  Array<PrimExpr> pragma_values;

  void VisitAttrs(AttrVisitor* v) {
    v->Visit("iter_type", &iter_type);
    v->Visit("bind_thread", &bind_thread);
    v->Visit("prefetch_data", &prefetch_data);
    v->Visit("prefetch_offset", &prefetch_offset);
    v->Visit("tensor_intrin", &tensor_intrin);
    v->Visit("dim_align_factor", &dim_align_factor);
    v->Visit("dim_align_offset", &dim_align_offset);
    v->Visit("pragma_keys", &pragma_keys);
    v->Visit("pragma_values", &pragma_values);
  }

  static constexpr const char* _type_key = "IterVarAttr";
  TVM_DECLARE_FINAL_OBJECT_INFO(IterVarAttrNode, Object);
};

class IterVarRelationNode : public Object {
 public:
  static constexpr const char* _type_key = "IterVarRelation";
  TVM_DECLARE_BASE_OBJECT_INFO(IterVarRelationNode, Object);
};

class SplitNode : public IterVarRelationNode {
 public:
  IterVar parent;
  IterVar outer;
  IterVar inner;
  PrimExpr factor;
  PrimExpr nparts;

  void VisitAttrs(AttrVisitor* v) {
    v->Visit("parent", &parent);
    v->Visit("outer", &outer);
    v->Visit("inner", &inner);
    v->Visit("factor", &factor);
    v->Visit("nparts", &nparts);
  }

  static constexpr const char* _type_key = "Split";
  TVM_DECLARE_FINAL_OBJECT_INFO(SplitNode, IterVarRelationNode);
};

class Split : public IterVarRelation {
 public:
  TVM_DLL Split(IterVar parent, IterVar outer, IterVar inner, PrimExpr factor, PrimExpr nparts);

  TVM_DEFINE_OBJECT_REF_METHODS(Split, IterVarRelation, SplitNode);
};

class FuseNode : public IterVarRelationNode {
 public:
  IterVar outer;
  IterVar inner;
  IterVar fused;

  void VisitAttrs(AttrVisitor* v) {
    v->Visit("outer", &outer);
    v->Visit("inner", &inner);
    v->Visit("fused", &fused);
  }

  static constexpr const char* _type_key = "Fuse";
  TVM_DECLARE_FINAL_OBJECT_INFO(FuseNode, IterVarRelationNode);
};

class Fuse : public IterVarRelation {
 public:
  TVM_DLL Fuse(IterVar outer, IterVar inner, IterVar fused);

  TVM_DEFINE_OBJECT_REF_METHODS(Fuse, IterVarRelation, FuseNode);
};

class RebaseNode : public IterVarRelationNode {
 public:
  IterVar parent;
  IterVar rebased;

  void VisitAttrs(AttrVisitor* v) {
    v->Visit("parent", &parent);
    v->Visit("rebased", &rebased);
  }

  static constexpr const char* _type_key = "Rebase";
  TVM_DECLARE_FINAL_OBJECT_INFO(RebaseNode, IterVarRelationNode);
};

class Rebase : public IterVarRelation {
 public:
  TVM_DLL Rebase(IterVar parent, IterVar rebased);

  TVM_DEFINE_OBJECT_REF_METHODS(Rebase, IterVarRelation, RebaseNode);
};

class SingletonNode : public IterVarRelationNode {
 public:
  IterVar iter;

  void VisitAttrs(AttrVisitor* v) { v->Visit("iter", &iter); }

  static constexpr const char* _type_key = "Singleton";
  TVM_DECLARE_FINAL_OBJECT_INFO(SingletonNode, IterVarRelationNode);
};

class Singleton : public IterVarRelation {
 public:
  TVM_DLL explicit Singleton(IterVar iter);

  TVM_DEFINE_OBJECT_REF_METHODS(Singleton, IterVarRelation, SingletonNode);
};

class TransformNode : public IterVarRelationNode {
 public:
  Array<IterVar> original_variables;

  Array<IterVar> transformed_variables;

  IndexMap forward_transformation;

  IndexMap inverse_transformation;

  void VisitAttrs(AttrVisitor* v) {
    v->Visit("original_variables", &original_variables);
    v->Visit("transformed_variables", &transformed_variables);
    v->Visit("forward_transformation", &forward_transformation);
    v->Visit("inverse_transformation", &inverse_transformation);
  }

  static constexpr const char* _type_key = "Transform";
  TVM_DECLARE_FINAL_OBJECT_INFO(TransformNode, IterVarRelationNode);
};

class Transform : public IterVarRelation {
 public:
  TVM_DLL explicit Transform(Array<IterVar> original_variables,
                             Array<IterVar> transformed_variables, IndexMap forward_transformation,
                             IndexMap inverse_transformation);

  TVM_DEFINE_OBJECT_REF_METHODS(Transform, IterVarRelation, TransformNode);
};

class SpecializedConditionNode : public Object {
 public:
  Array<PrimExpr> clauses;

  void VisitAttrs(AttrVisitor* v) { v->Visit("clauses", &clauses); }

  static constexpr const char* _type_key = "SpecializedCondition";
  TVM_DECLARE_FINAL_OBJECT_INFO(SpecializedConditionNode, Object);
};

class SpecializedCondition : public ObjectRef {
 public:
  TVM_DLL SpecializedCondition(Array<PrimExpr> conditions);  // NOLINT(*)

  TVM_DLL static SpecializedCondition Current();

  TVM_DEFINE_OBJECT_REF_METHODS(SpecializedCondition, ObjectRef, SpecializedConditionNode);
  class Internal;

 private:
  // enable with syntax.
  friend class Internal;
  friend class With<SpecializedCondition>;
  TVM_DLL void EnterWithScope();
  TVM_DLL void ExitWithScope();
};

// implementations
inline const StageNode* Stage::operator->() const { return static_cast<const StageNode*>(get()); }
inline StageNode* Stage::operator->() { return static_cast<StageNode*>(get_mutable()); }

inline const ScheduleNode* Schedule::operator->() const {
  return static_cast<const ScheduleNode*>(get());
}
inline ScheduleNode* Schedule::operator->() { return static_cast<ScheduleNode*>(get_mutable()); }

inline const IterVarRelationNode* IterVarRelation::operator->() const {
  return static_cast<const IterVarRelationNode*>(get());
}

inline const IterVarAttrNode* IterVarAttr::operator->() const {
  return static_cast<const IterVarAttrNode*>(get());
}

}  // namespace te
}  // namespace tvm
#endif  // TVM_TE_SCHEDULE_H_