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_