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.
 */
#ifndef TVM_TIR_SCHEDULE_SCHEDULE_H_
#define TVM_TIR_SCHEDULE_SCHEDULE_H_

#include <tvm/support/random_engine.h>
#include <tvm/tir/index_map.h>
#include <tvm/tir/schedule/state.h>
#include <tvm/tir/schedule/trace.h>

namespace tvm {
namespace tir {

enum class ScheduleErrorRenderLevel : int32_t {
  kDetail = 0,
  kFast = 1,
  kNone = 2,
};

enum class BufferIndexType : int32_t {
  kRead = 0,
  kWrite = 1,
};

/**************** Random variable: BlockRV ****************/

class BlockRVNode : public runtime::Object {
 public:
  void VisitAttrs(tvm::AttrVisitor* v) {}
  static constexpr const char* _type_key = "tir.BlockRV";
  TVM_DECLARE_FINAL_OBJECT_INFO(BlockRVNode, runtime::Object);
};

class BlockRV : public runtime::ObjectRef {
 public:
  TVM_DLL BlockRV();
  TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(BlockRV, runtime::ObjectRef, BlockRVNode);
};

/**************** Random variable: LoopRV ****************/

class LoopRVNode : public runtime::Object {
 public:
  void VisitAttrs(tvm::AttrVisitor* v) {}
  static constexpr const char* _type_key = "tir.LoopRV";
  TVM_DECLARE_FINAL_OBJECT_INFO(LoopRVNode, runtime::Object);
};

class LoopRV : public runtime::ObjectRef {
 public:
  TVM_DLL LoopRV();
  TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(LoopRV, runtime::ObjectRef, LoopRVNode);
};

/**************** Random variable: ExprRV ****************/

using ExprRV = PrimExpr;

using ExprRVNode = PrimExprNode;

/**************** The Schedule class ****************/

class Schedule;

class ScheduleNode : public runtime::Object {
  friend class Schedule;

 public:
  virtual ~ScheduleNode() = default;

  static constexpr const char* _type_key = "tir.Schedule";
  TVM_DECLARE_FINAL_OBJECT_INFO(ScheduleNode, runtime::Object);

 public:
  virtual IRModule mod() const { return state()->mod; }
  virtual ScheduleState state() const = 0;
  virtual Optional<Trace> trace() const = 0;
  virtual void WorkOn(const String& func_name) = 0;
  virtual Schedule Copy() = 0;
  virtual void Seed(support::LinearCongruentialEngine::TRandState seed) = 0;
  virtual support::LinearCongruentialEngine::TRandState ForkSeed() = 0;

 public:
  /******** Lookup/Remove random variables ********/
  virtual Block Get(const BlockRV& block_rv) const = 0;
  virtual For Get(const LoopRV& loop_rv) const = 0;
  virtual PrimExpr Get(const ExprRV& expr_rv) const = 0;
  virtual StmtSRef GetSRef(const BlockRV& block_rv) const = 0;
  virtual StmtSRef GetSRef(const LoopRV& loop_rv) const = 0;
  virtual bool HasBlock(const BlockRV& block_rv) const = 0;
  virtual StmtSRef GetSRef(const StmtNode* stmt) const;
  StmtSRef GetSRef(const Stmt& stmt) const { return this->GetSRef(stmt.get()); }
  virtual void RemoveRV(const BlockRV& block_rv) = 0;
  virtual void RemoveRV(const LoopRV& loop_rv) = 0;
  virtual void RemoveRV(const ExprRV& expr_rv) = 0;

 public:
  /******** Schedule: Sampling ********/
  virtual ExprRV SampleCategorical(const Array<Integer>& candidates, const Array<FloatImm>& probs,
                                   Optional<Integer> decision = NullOpt) = 0;
  virtual Array<ExprRV> SamplePerfectTile(const LoopRV& loop_rv, int n, int max_innermost_factor,
                                          Optional<Array<Integer>> decision = NullOpt) = 0;
  virtual LoopRV SampleComputeLocation(const BlockRV& block_rv,
                                       Optional<Integer> decision = NullOpt) = 0;

  /******** Schedule: Get blocks & loops ********/
  virtual BlockRV GetBlock(const String& name, const Optional<String>& func_name = NullOpt) = 0;
  virtual Array<LoopRV> GetLoops(const BlockRV& block_rv) = 0;
  virtual Array<BlockRV> GetChildBlocks(const BlockRV& block_rv) = 0;
  virtual Array<BlockRV> GetChildBlocks(const LoopRV& loop_rv) = 0;
  virtual Array<BlockRV> GetProducers(const BlockRV& block_rv) = 0;
  virtual Array<BlockRV> GetConsumers(const BlockRV& block_rv) = 0;
  virtual Array<BlockRV> GetOutputBlocks(const BlockRV& scope_block_rv) = 0;
  /******** Schedule: Transform loops ********/
  virtual LoopRV Merge(const Array<LoopRV>& loop_rvs) = 0;
  virtual LoopRV Fuse(const Array<LoopRV>& loop_rvs, bool preserve_unit_iters = true) = 0;
  virtual Array<LoopRV> Split(const LoopRV& loop_rv, const Array<Optional<ExprRV>>& factors,
                              bool preserve_unit_iters = true) = 0;
  virtual void Reorder(const Array<LoopRV>& ordered_loop_rvs) = 0;
  virtual void ReorderBlockIterVar(const BlockRV& block_rv, const Array<Integer> new_order) = 0;
  virtual LoopRV AddUnitLoop(const BlockRV& block_rv) = 0;
  virtual LoopRV AddUnitLoop(const LoopRV& loop_rv) = 0;
  /******** Schedule: Manipulate ForKind ********/
  virtual void Parallel(const LoopRV& loop_rv) = 0;
  virtual void Vectorize(const LoopRV& loop_rv) = 0;
  virtual void Bind(const LoopRV& loop_rv, const String& thread_axis) = 0;
  virtual void Unroll(const LoopRV& loop_rv) = 0;
  /******** Schedule: Insert cache stages ********/
  virtual BlockRV CacheRead(const BlockRV& block_rv, int read_buffer_index,
                            const String& storage_scope,
                            const Array<BlockRV> consumer_blocks = {}) = 0;
  virtual BlockRV CacheWrite(const BlockRV& block_rv, int write_buffer_index,
                             const String& storage_scope,
                             const Array<BlockRV> consumer_blocks = {}) = 0;
  virtual BlockRV ReindexCacheRead(const BlockRV& block_rv, int read_buffer_index,
                                   const String& storage_scope, const IndexMap& index_map) = 0;
  virtual BlockRV ReindexCacheWrite(const BlockRV& block_rv, int write_buffer_index,
                                    const String& storage_scope, const IndexMap& index_map) = 0;
  virtual Array<BlockRV> CacheInplace(const BlockRV& block_rv, int read_buffer_index,
                                      const String& storage_scope) = 0;
  virtual Array<BlockRV> CacheIndex(const BlockRV& block_rv, const String& storage_scope,
                                    int cse_thresh) = 0;
  virtual BlockRV ReIndex(const BlockRV& block_rv, int buffer_index,
                          BufferIndexType buffer_index_type) = 0;
  /******** Schedule: Data movement ********/
  virtual BlockRV ReadAt(const LoopRV& loop_rv, const BlockRV& block_rv, int read_buffer_index,
                         const String& storage_scope) = 0;
  virtual BlockRV WriteAt(const LoopRV& loop_rv, const BlockRV& block_rv, int write_buffer_index,
                          const String& storage_scope) = 0;
  /******** Schedule: Compute location ********/
  virtual void ComputeAt(const BlockRV& block_rv, const LoopRV& loop_rv, bool preserve_unit_loops,
                         int index = -1) = 0;
  virtual void ReverseComputeAt(const BlockRV& block_rv, const LoopRV& loop_rv,
                                bool preserve_unit_loops, int index = -1) = 0;
  virtual void ComputeInline(const BlockRV& block) = 0;
  virtual void ReverseComputeInline(const BlockRV& block) = 0;
  /******** Schedule: Reduction ********/
  virtual BlockRV DecomposeReduction(const BlockRV& block_rv, const LoopRV& loop_rv) = 0;
  virtual BlockRV RFactor(const LoopRV& loop_rv, int factor_axis) = 0;
  /******** Schedule: Block annotation ********/
  virtual void StorageAlign(const BlockRV& block_rv, int buffer_index, int axis, int factor,
                            int offset) = 0;
  virtual void SetScope(const BlockRV& block_rv, int buffer_index, const String& storage_scope) = 0;
  virtual void UnsafeSetDType(const BlockRV& block_rv, int buffer_index, const String& dtype) = 0;
  /******** Schedule: Blockize & Tensorize ********/
  virtual BlockRV Blockize(const LoopRV& loop_rv, bool preserve_unit_iters = true) = 0;
  virtual void Tensorize(const LoopRV& loop_rv, const String& intrin,
                         bool preserve_unit_iters = true) = 0;
  virtual void Tensorize(const BlockRV& block_rv, const String& intrin,
                         bool preserve_unit_iters = true) = 0;

  /******** Schedule: Annotation ********/
  virtual void Annotate(const LoopRV& loop_rv, const String& ann_key, const ObjectRef& ann_val) = 0;
  virtual void Annotate(const BlockRV& block_rv, const String& ann_key,
                        const ObjectRef& ann_val) = 0;
  virtual void Unannotate(const LoopRV& loop_rv, const String& ann_key) = 0;
  virtual void Unannotate(const BlockRV& block_rv, const String& ann_key) = 0;

  /******** Schedule: Layout transformation ********/
  virtual void TransformLayout(const BlockRV& block_rv, int buffer_index,
                               BufferIndexType buffer_index_type, const IndexMap& index_map,
                               const Optional<IndexMap>& pad_value = NullOpt,
                               bool assume_injective_transform = false) = 0;

  virtual void TransformBlockLayout(const BlockRV& block_rv, const IndexMap& index_map) = 0;

  virtual void SetAxisSeparator(const BlockRV& block_rv, int buffer_index,
                                BufferIndexType buffer_index_type,
                                const Array<IntImm>& axis_separators) = 0;

  /******** Schedule: Padding ********/
  virtual BlockRV DecomposePadding(const BlockRV& block_rv, const LoopRV& loop_rv) = 0;

  virtual void PadEinsum(const BlockRV& block_rv, const Array<Integer>& padding) = 0;

  /******** Schedule: Buffer transformation ********/
  virtual void RollingBuffer(const BlockRV& block_rv, int write_buffer_index) = 0;

  /******** Schedule: Misc ********/
  virtual void EnterPostproc() = 0;
};

class Schedule : public runtime::ObjectRef {
 public:
  TVM_DLL static Schedule Concrete(IRModule mod, support::LinearCongruentialEngine::TRandState seed,
                                   int debug_mask, ScheduleErrorRenderLevel error_render_level,
                                   bool enable_check = true);
  TVM_DLL static Schedule Traced(IRModule mod, support::LinearCongruentialEngine::TRandState seed,
                                 int debug_mask, ScheduleErrorRenderLevel error_render_level,
                                 bool enable_check = true);
  TVM_DEFINE_MUTABLE_OBJECT_REF_METHODS(Schedule, runtime::ObjectRef, ScheduleNode);
};

}  // namespace tir
}  // namespace tvm

#endif  // TVM_TIR_SCHEDULE_SCHEDULE_H_