transform_step.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.
 */
 
#ifndef TVM_AUTO_SCHEDULER_TRANSFORM_STEP_H_
#define TVM_AUTO_SCHEDULER_TRANSFORM_STEP_H_
 
#include <dmlc/common.h>
#include <dmlc/json.h>
#include <tvm/node/node.h>
#include <tvm/te/schedule.h>
 
#include <vector>
 
namespace tvm {
namespace auto_scheduler {
 
typedef Map<tvm::te::Stage, Array<tir::IterVar>, ObjectHash, ObjectEqual> StageToAxesMap;
 
void UpdateStageToAxesMap(const te::Stage& stage, StageToAxesMap* stage_to_axes);
 
enum class IteratorKind : int {
  kSpatial = 0,
  kReduction = 1,
  kMixed = 2,
  kSpecial = 3
};
 
enum class IteratorAnnotation : int {
  kNone = 0,
  kUnroll = 1,
  kVectorize = 2,
  kParallel = 3,
  kVThread = 4,
  kBlockX = 5,
  kThreadX = 6,
  kBlockY = 7,
  kThreadY = 8,
  kBlockZ = 9,
  kThreadZ = 10,
  kTensorize = 11
};
 
extern const char* IteratorAnnotationString[];
 
// forward declaration
class Iterator;
 
class IteratorNode : public Object {
 public:
  String name;
  Range range;
  IteratorKind iter_kind;
  IteratorAnnotation annotation;
  std::vector<Iterator> orig_iters;
 
  void VisitAttrs(tvm::AttrVisitor* v) {
    v->Visit("name", &name);
    v->Visit("range", &range);
    v->Visit("iter_kind", &iter_kind);
    v->Visit("annotation", &annotation);
  }
 
  static constexpr const char* _type_key = "auto_scheduler.Iterator";
  TVM_DECLARE_FINAL_OBJECT_INFO(IteratorNode, Object);
};
 
class Iterator : public ObjectRef {
 public:
  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);
};
 
class StepNode : public Object {
 public:
  int stage_id;
 
  virtual void WriteToRecord(dmlc::JSONWriter* writer) const = 0;
 
  static constexpr const char* _type_key = "auto_scheduler.Step";
  TVM_DECLARE_BASE_OBJECT_INFO(StepNode, Object);
};
 
class Step : public ObjectRef {
 public:
  StepNode* CopyOnWrite();
 
  TVM_DEFINE_OBJECT_REF_METHODS(Step, ObjectRef, StepNode);
};
 
// Forward declaration
class State;
class ComputeDAG;
 
Step StepReadFromRecord(dmlc::JSONReader* reader);
 
void StepApplyToState(const Step& step, State* state, const ComputeDAG& dag);
 
void StepApplyToSchedule(const Step& step, Array<te::Stage>* stages, StageToAxesMap* stage_to_axes,
                         te::Schedule* schedule, const Array<Step>& transform_steps);
 
String StepPrintAsPythonAPI(const Step& step, Array<te::Stage>* stages,
                            StageToAxesMap* stage_to_axes, te::Schedule* schedule,
                            const Array<Step>& transform_steps);
 
/********** Steps working on single stage **********/
 
class AnnotationStepNode : public StepNode {
 public:
  int iter_id;
  IteratorAnnotation annotation;
 
  void WriteToRecord(dmlc::JSONWriter* writer) const final;
 
  Iterator ApplyToState(State* state) const;
 
  void ApplyToSchedule(Array<te::Stage>* stages, StageToAxesMap* stage_to_axes) const;
 
  String PrintAsPythonAPI(Array<te::Stage>* stages, StageToAxesMap* stage_to_axes) const;
 
  static constexpr const char* record_prefix_str = "AN";
 
  static constexpr const char* _type_key = "auto_scheduler.AnnotationStep";
  TVM_DECLARE_FINAL_OBJECT_INFO(AnnotationStepNode, StepNode);
};
 
class AnnotationStep : public Step {
 public:
  AnnotationStep(int stage_id, int iter_id, IteratorAnnotation ann);
 
  explicit AnnotationStep(dmlc::JSONReader* reader);
 
  TVM_DEFINE_OBJECT_REF_METHODS(AnnotationStep, Step, AnnotationStepNode);
};
 
class FuseStepNode : public StepNode {
 public:
  Array<Integer> fused_ids;
 
  void WriteToRecord(dmlc::JSONWriter* writer) const final;
 
  Iterator ApplyToState(State* state) const;
 
  tir::IterVar ApplyToSchedule(Array<te::Stage>* stages, StageToAxesMap* stage_to_axes) const;
 
  String PrintAsPythonAPI(Array<te::Stage>* stages, StageToAxesMap* stage_to_axes) const;
 
  static constexpr const char* record_prefix_str = "FU";
 
  static constexpr const char* _type_key = "auto_scheduler.FuseStep";
  TVM_DECLARE_FINAL_OBJECT_INFO(FuseStepNode, StepNode);
};
 
class FuseStep : public Step {
 public:
  FuseStep(int stage_id, const Array<Integer>& fused_ids);
 
  explicit FuseStep(dmlc::JSONReader* reader);
 
  TVM_DEFINE_OBJECT_REF_METHODS(FuseStep, Step, FuseStepNode);
};
 
class PragmaStepNode : public StepNode {
 public:
  int iter_id;
  String pragma_type;
 
  void WriteToRecord(dmlc::JSONWriter* writer) const final;
 
  void ApplyToState(State* state) const;
 
  void ApplyToSchedule(Array<te::Stage>* stages, StageToAxesMap* stage_to_axes) const;
 
  String PrintAsPythonAPI(Array<te::Stage>* stages, StageToAxesMap* stage_to_axes) const;
 
  static constexpr const char* record_prefix_str = "PR";
 
  static constexpr const char* _type_key = "auto_scheduler.PragmaStep";
  TVM_DECLARE_FINAL_OBJECT_INFO(PragmaStepNode, StepNode);
};
 
class PragmaStep : public Step {
 public:
  PragmaStep(int stage_id, int iter_id, String pragma_type);
 
  explicit PragmaStep(dmlc::JSONReader* reader);
 
  TVM_DEFINE_OBJECT_REF_METHODS(PragmaStep, Step, PragmaStepNode);
};
 
class ReorderStepNode : public StepNode {
 public:
  Array<Integer> after_ids;
 
  void WriteToRecord(dmlc::JSONWriter* writer) const final;
 
  void ApplyToState(State* state) const;
 
  void ApplyToSchedule(Array<te::Stage>* stages, StageToAxesMap* stage_to_axes) const;
 
  String PrintAsPythonAPI(Array<te::Stage>* stages, StageToAxesMap* stage_to_axes) const;
 
  static constexpr const char* record_prefix_str = "RE";
 
  static constexpr const char* _type_key = "auto_scheduler.ReorderStep";
  TVM_DECLARE_FINAL_OBJECT_INFO(ReorderStepNode, StepNode);
};
 
class ReorderStep : public Step {
 public:
  ReorderStep(int stage_id, const Array<Integer>& after_ids);
 
  explicit ReorderStep(dmlc::JSONReader* reader);
 
  TVM_DEFINE_OBJECT_REF_METHODS(ReorderStep, Step, ReorderStepNode);
};
 
class SplitStepNode : public StepNode {
 public:
  int iter_id;
  Optional<PrimExpr> extent;
  Array<Optional<Integer>> lengths;
  bool inner_to_outer;
 
  void WriteToRecord(dmlc::JSONWriter* writer) const final;
 
  Array<Iterator> ApplyToState(State* state) const;
 
  Array<tir::IterVar> ApplyToSchedule(Array<te::Stage>* stages,
                                      StageToAxesMap* stage_to_axes) const;
 
  String PrintAsPythonAPI(Array<te::Stage>* stages, StageToAxesMap* stage_to_axes) const;
 
  static constexpr const char* record_prefix_str = "SP";
 
  static constexpr const char* _type_key = "auto_scheduler.SplitStep";
  TVM_DECLARE_FINAL_OBJECT_INFO(SplitStepNode, StepNode);
};
 
class SplitStep : public Step {
 public:
  SplitStep(int stage_id, int iter_id, Optional<PrimExpr> extent,
            const Array<Optional<Integer>>& lengths, bool inner_to_outer);
 
  explicit SplitStep(dmlc::JSONReader* reader);
 
  TVM_DEFINE_OBJECT_REF_METHODS(SplitStep, Step, SplitStepNode);
};
 
class FollowSplitStepNode : public StepNode {
 public:
  int iter_id;
  int src_step_id;
  int n_split;
 
  void WriteToRecord(dmlc::JSONWriter* writer) const final;
 
  Array<Optional<Integer>> ExtractSplitLengths(const Array<Step>& transform_steps) const;
 
  Array<Iterator> ApplyToState(State* state) const;
 
  Array<tir::IterVar> ApplyToSchedule(Array<te::Stage>* stages, StageToAxesMap* stage_to_axes,
                                      const Array<Step>& transform_steps) const;
 
  String PrintAsPythonAPI(Array<te::Stage>* stages, StageToAxesMap* stage_to_axes,
                          const Array<Step>& transform_steps) const;
 
  static constexpr const char* record_prefix_str = "FSP";
 
  static constexpr const char* _type_key = "auto_scheduler.FollowSplitStep";
  TVM_DECLARE_FINAL_OBJECT_INFO(FollowSplitStepNode, StepNode);
};
 
class FollowSplitStep : public Step {
 public:
  FollowSplitStep(int stage_id, int iter_id, int src_step_id, int n_split);
 
  explicit FollowSplitStep(dmlc::JSONReader* reader);
 
  TVM_DEFINE_OBJECT_REF_METHODS(FollowSplitStep, Step, FollowSplitStepNode);
};
 
class FollowFusedSplitStepNode : public StepNode {
 public:
  int iter_id;
  Array<Integer> src_step_ids;
  int level;
  bool factor_or_nparts;
 
  void WriteToRecord(dmlc::JSONWriter* writer) const final;
 
  Optional<Integer> ExtractSplitLength(const Array<Step>& transform_steps) const;
 
  Array<Iterator> ApplyToState(State* state) const;
 
  Array<tir::IterVar> ApplyToSchedule(Array<te::Stage>* stages, StageToAxesMap* stage_to_axes,
                                      const Array<Step>& transform_steps) const;
 
  String PrintAsPythonAPI(Array<te::Stage>* stages, StageToAxesMap* stage_to_axes,
                          const Array<Step>& transform_steps) const;
 
  static constexpr const char* record_prefix_str = "FFSP";
 
  static constexpr const char* _type_key = "auto_scheduler.FollowFusedSplitStep";
  TVM_DECLARE_FINAL_OBJECT_INFO(FollowFusedSplitStepNode, StepNode);
};
 
class FollowFusedSplitStep : public Step {
 public:
  FollowFusedSplitStep(int stage_id, int iter_id, const Array<Integer>& src_step_ids, int level,
                       bool factor_or_nparts);
 
  explicit FollowFusedSplitStep(dmlc::JSONReader* reader);
 
  TVM_DEFINE_OBJECT_REF_METHODS(FollowFusedSplitStep, Step, FollowFusedSplitStepNode);
};
 
class StorageAlignStepNode : public StepNode {
 public:
  int iter_id;
  int factor;
  int offset;
 
  void WriteToRecord(dmlc::JSONWriter* writer) const final;
 
  void ApplyToState(State* state) const;
 
  void ApplyToSchedule(Array<te::Stage>* stages, StageToAxesMap* stage_to_axes) const;
 
  String PrintAsPythonAPI(Array<te::Stage>* stages, StageToAxesMap* stage_to_axes) const;
 
  static constexpr const char* record_prefix_str = "SA";
 
  static constexpr const char* _type_key = "auto_scheduler.StorageAlignStep";
  TVM_DECLARE_FINAL_OBJECT_INFO(StorageAlignStepNode, StepNode);
};
 
class StorageAlignStep : public Step {
 public:
  StorageAlignStep(int stage_id, int iter_id, int factor, int offset);
 
  explicit StorageAlignStep(dmlc::JSONReader* reader);
 
  TVM_DEFINE_OBJECT_REF_METHODS(StorageAlignStep, Step, StorageAlignStepNode);
};
 
/********** Steps working on multiple stages **********/
 
class ComputeAtStepNode : public StepNode {
 public:
  int target_stage_id;
  int target_iter_id;
 
  void WriteToRecord(dmlc::JSONWriter* writer) const final;
 
  void ApplyToState(State* state) const;
 
  void ApplyToSchedule(Array<te::Stage>* stages, StageToAxesMap* stage_to_axes) const;
 
  String PrintAsPythonAPI(Array<te::Stage>* stages, StageToAxesMap* stage_to_axes) const;
 
  static constexpr const char* record_prefix_str = "CA";
 
  static constexpr const char* _type_key = "auto_scheduler.ComputeAtStep";
  TVM_DECLARE_FINAL_OBJECT_INFO(ComputeAtStepNode, StepNode);
};
 
class ComputeAtStep : public Step {
 public:
  ComputeAtStep(int stage_id, int target_stage_id, int target_iter_id);
 
  explicit ComputeAtStep(dmlc::JSONReader* reader);
 
  TVM_DEFINE_OBJECT_REF_METHODS(ComputeAtStep, Step, ComputeAtStepNode);
};
 
class ComputeInlineStepNode : public StepNode {
 public:
  void WriteToRecord(dmlc::JSONWriter* writer) const final;
 
  void ApplyToState(State* state) const;
 
  void ApplyToSchedule(Array<te::Stage>* stages, StageToAxesMap* stage_to_axes) const;
 
  String PrintAsPythonAPI(Array<te::Stage>* stages, StageToAxesMap* stage_to_axes) const;
 
  static constexpr const char* record_prefix_str = "CI";
 
  static constexpr const char* _type_key = "auto_scheduler.ComputeInlineStep";
  TVM_DECLARE_FINAL_OBJECT_INFO(ComputeInlineStepNode, StepNode);
};
 
class ComputeInlineStep : public Step {
 public:
  explicit ComputeInlineStep(int stage_id);
 
  explicit ComputeInlineStep(dmlc::JSONReader* reader);
 
  TVM_DEFINE_OBJECT_REF_METHODS(ComputeInlineStep, Step, ComputeInlineStepNode);
};
 
class ComputeRootStepNode : public StepNode {
 public:
  void WriteToRecord(dmlc::JSONWriter* writer) const final;
 
  void ApplyToState(State* state) const;
 
  void ApplyToSchedule(Array<te::Stage>* stages, StageToAxesMap* stage_to_axes) const;
 
  String PrintAsPythonAPI(Array<te::Stage>* stages, StageToAxesMap* stage_to_axes) const;
 
  static constexpr const char* record_prefix_str = "CR";
 
  static constexpr const char* _type_key = "auto_scheduler.ComputeRootStep";
  TVM_DECLARE_FINAL_OBJECT_INFO(ComputeRootStepNode, StepNode);
};
 
class ComputeRootStep : public Step {
 public:
  explicit ComputeRootStep(int stage_id);
 
  explicit ComputeRootStep(dmlc::JSONReader* reader);
 
  TVM_DEFINE_OBJECT_REF_METHODS(ComputeRootStep, Step, ComputeRootStepNode);
};
 
/********** Steps adding new stages **********/
 
class CacheReadStepNode : public StepNode {
 public:
  String scope_name;
  Array<Integer> reader_stage_ids;
 
  void WriteToRecord(dmlc::JSONWriter* writer) const final;
 
  int ApplyToState(State* state, const ComputeDAG& dag) const;
 
  te::Tensor ApplyToSchedule(Array<te::Stage>* stages, StageToAxesMap* stage_to_axes,
                             te::Schedule* schedule) const;
 
  String PrintAsPythonAPI(Array<te::Stage>* stages, StageToAxesMap* stage_to_axes,
                          te::Schedule* schedule) const;
 
  static constexpr const char* record_prefix_str = "CHR";
 
  static constexpr const char* _type_key = "auto_scheduler.CacheReadStep";
  TVM_DECLARE_FINAL_OBJECT_INFO(CacheReadStepNode, StepNode);
};
 
class CacheReadStep : public Step {
 public:
  CacheReadStep(int stage_id, String scope_name, const Array<Integer>& reader_stage_ids);
 
  explicit CacheReadStep(dmlc::JSONReader* reader);
 
  TVM_DEFINE_OBJECT_REF_METHODS(CacheReadStep, Step, CacheReadStepNode);
};
 
class CacheWriteStepNode : public StepNode {
 public:
  String scope_name;
 
  void WriteToRecord(dmlc::JSONWriter* writer) const final;
 
  int ApplyToState(State* state, const ComputeDAG& dag) const;
 
  Array<te::Tensor> ApplyToSchedule(Array<te::Stage>* stages, StageToAxesMap* stage_to_axes,
                                    te::Schedule* schedule) const;
 
  String PrintAsPythonAPI(Array<te::Stage>* stages, StageToAxesMap* stage_to_axes,
                          te::Schedule* schedule) const;
 
  static constexpr const char* record_prefix_str = "CHW";
 
  static constexpr const char* _type_key = "auto_scheduler.CacheWriteStep";
  TVM_DECLARE_FINAL_OBJECT_INFO(CacheWriteStepNode, StepNode);
};
 
class CacheWriteStep : public Step {
 public:
  CacheWriteStep(int stage_id, String scope_name);
 
  explicit CacheWriteStep(dmlc::JSONReader* reader);
 
  TVM_DEFINE_OBJECT_REF_METHODS(CacheWriteStep, Step, CacheWriteStepNode);
};
 
class RfactorStepNode : public StepNode {
 public:
  int iter_id;
  int factor_iter_id;
 
  void WriteToRecord(dmlc::JSONWriter* writer) const final;
 
  int ApplyToState(State* state, const ComputeDAG& dag) const;
 
  Array<te::Tensor> ApplyToSchedule(Array<te::Stage>* stages, StageToAxesMap* stage_to_axes,
                                    te::Schedule* schedule) const;
 
  String PrintAsPythonAPI(Array<te::Stage>* stages, StageToAxesMap* stage_to_axes,
                          te::Schedule* schedule) const;
 
  static constexpr const char* record_prefix_str = "RF";
 
  static constexpr const char* _type_key = "auto_scheduler.RfactorStep";
  TVM_DECLARE_FINAL_OBJECT_INFO(RfactorStepNode, StepNode);
};
 
class RfactorStep : public Step {
 public:
  RfactorStep(int stage_id, int iter_id, int factor_iter_id);
 
  explicit RfactorStep(dmlc::JSONReader* reader);
 
  TVM_DEFINE_OBJECT_REF_METHODS(RfactorStep, Step, RfactorStepNode);
};
 
}  // namespace auto_scheduler
}  // namespace tvm
 
#endif  // TVM_AUTO_SCHEDULER_TRANSFORM_STEP_H_