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_S_TIR_SCHEDULE_SCHEDULE_H_
#define TVM_S_TIR_SCHEDULE_SCHEDULE_H_
 
#include <tvm/ir/prim/expr.h>
#include <tvm/s_tir/random_engine.h>
#include <tvm/s_tir/schedule/state.h>
#include <tvm/s_tir/schedule/trace.h>
#include <tvm/tirx/index_map.h>
 
namespace tvm {
namespace s_tir {
using namespace tvm::prim;
using namespace tvm::tirx;
 
enum class ScheduleErrorRenderLevel : int32_t {
  kDetail = 0,
  kFast = 1,
  kNone = 2,
};
 
enum class BufferIndexType : int32_t {
  kRead = 0,
  kWrite = 1,
};
 
/**************** Random variable: SBlockRV ****************/
 
class SBlockRVNode : public ffi::Object {
 public:
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<SBlockRVNode>();
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("s_tir.SBlockRV", SBlockRVNode, ffi::Object);
};
 
class SBlockRV : public ffi::ObjectRef {
 public:
  TVM_DLL SBlockRV();
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(SBlockRV, ffi::ObjectRef, SBlockRVNode);
};
 
/**************** Random variable: LoopRV ****************/
 
class LoopRVNode : public ffi::Object {
 public:
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<LoopRVNode>();
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("s_tir.LoopRV", LoopRVNode, ffi::Object);
};
 
class LoopRV : public ffi::ObjectRef {
 public:
  TVM_DLL LoopRV();
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(LoopRV, ffi::ObjectRef, LoopRVNode);
};
 
/**************** Random variable: ExprRV ****************/
 
using ExprRV = PrimExpr;
 
using ExprRVNode = ExprNode;
 
/**************** The Schedule class ****************/
 
class Schedule;
 
class ScheduleNode : public ffi::Object {
  friend class Schedule;
 
 public:
  virtual ~ScheduleNode() = default;
 
  static constexpr const bool _type_mutable = true;
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("s_tir.Schedule", ScheduleNode, ffi::Object);
 
 public:
  virtual IRModule mod() const { return state()->mod; }
  virtual ScheduleState state() const = 0;
  virtual ffi::Optional<Trace> trace() const = 0;
  virtual ffi::Optional<GlobalVar> func_working_on() const = 0;
  virtual void WorkOn(const ffi::String& func_name) = 0;
  virtual Schedule Copy() = 0;
  virtual void Seed(LinearCongruentialEngine::TRandState seed) = 0;
  virtual LinearCongruentialEngine::TRandState ForkSeed() = 0;
 
 public:
  /******** Lookup/Remove random variables ********/
  virtual SBlock Get(const SBlockRV& 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 SBlockRV& block_rv) const = 0;
  virtual StmtSRef GetSRef(const LoopRV& loop_rv) const = 0;
  virtual bool HasBlock(const SBlockRV& 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 SBlockRV& 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 ffi::Array<int64_t>& candidates,
                                   const ffi::Array<FloatImm>& probs,
                                   ffi::Optional<int64_t> decision = std::nullopt) = 0;
  virtual ffi::Array<ExprRV> SamplePerfectTile(
      const LoopRV& loop_rv, int n, int max_innermost_factor,
      ffi::Optional<ffi::Array<int64_t>> decision = std::nullopt) = 0;
  virtual ffi::Array<ExprRV> SamplePartitionedTile(
      const LoopRV& loop_rv, int n, int partition_pos, int innerpart_factor,
      ffi::Optional<ffi::Array<int64_t>> decision = std::nullopt) = 0;
  virtual LoopRV SampleComputeLocation(const SBlockRV& block_rv,
                                       ffi::Optional<int64_t> decision = std::nullopt) = 0;
 
  /******** Schedule: Get blocks & loops ********/
  virtual SBlockRV GetSBlock(const ffi::String& name,
                             const ffi::Optional<ffi::String>& func_name = std::nullopt) = 0;
  virtual ffi::Array<LoopRV> GetLoops(const SBlockRV& block_rv) = 0;
  virtual ffi::Array<SBlockRV> GetChildBlocks(const SBlockRV& block_rv) = 0;
  virtual ffi::Array<SBlockRV> GetChildBlocks(const LoopRV& loop_rv) = 0;
  virtual ffi::Array<SBlockRV> GetProducers(const SBlockRV& block_rv) = 0;
  virtual ffi::Array<SBlockRV> GetConsumers(const SBlockRV& block_rv) = 0;
  virtual ffi::Array<SBlockRV> GetOutputBlocks(const SBlockRV& scope_block_rv) = 0;
  /******** Schedule: Transform loops ********/
  virtual LoopRV Merge(const ffi::Array<LoopRV>& loop_rvs) = 0;
  virtual LoopRV Fuse(const ffi::Array<LoopRV>& loop_rvs, bool preserve_unit_iters = true) = 0;
  virtual ffi::Array<LoopRV> Split(const LoopRV& loop_rv,
                                   const ffi::Array<ffi::Optional<ExprRV>>& factors,
                                   bool preserve_unit_iters = true,
                                   bool disable_predication = false) = 0;
  virtual ffi::Array<LoopRV> LoopPartition(const LoopRV& loop_rv,
                                           const ffi::Array<ffi::Optional<ExprRV>>& factors,
                                           bool preserve_unit_iters = true) = 0;
  virtual void Reorder(const ffi::Array<LoopRV>& ordered_loop_rvs) = 0;
  virtual void ReorderBlockIterVar(const SBlockRV& block_rv,
                                   const ffi::Array<int64_t> new_order) = 0;
  virtual LoopRV AddUnitLoop(const SBlockRV& 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 ffi::String& thread_axis) = 0;
  virtual void Unroll(const LoopRV& loop_rv) = 0;
  /******** Schedule: Insert cache stages ********/
  virtual SBlockRV CacheRead(const SBlockRV& block_rv, int read_buffer_index,
                             const ffi::String& storage_scope,
                             const ffi::Array<SBlockRV> consumer_blocks = {}) = 0;
  virtual SBlockRV CacheWrite(const SBlockRV& block_rv, int write_buffer_index,
                              const ffi::String& storage_scope,
                              const ffi::Array<SBlockRV> consumer_blocks = {}) = 0;
  virtual SBlockRV ReindexCacheRead(const SBlockRV& block_rv, int read_buffer_index,
                                    const ffi::String& storage_scope,
                                    const IndexMap& index_map) = 0;
  virtual SBlockRV ReindexCacheWrite(const SBlockRV& block_rv, int write_buffer_index,
                                     const ffi::String& storage_scope,
                                     const IndexMap& index_map) = 0;
  virtual ffi::Array<SBlockRV> CacheInplace(const SBlockRV& block_rv, int read_buffer_index,
                                            const ffi::String& storage_scope) = 0;
  virtual ffi::Array<SBlockRV> CacheIndex(const SBlockRV& block_rv,
                                          const ffi::String& storage_scope, int cse_thresh) = 0;
  virtual SBlockRV ReIndex(const SBlockRV& block_rv, int buffer_index,
                           BufferIndexType buffer_index_type) = 0;
  /******** Schedule: Data movement ********/
  virtual SBlockRV ReadAt(const LoopRV& loop_rv, const SBlockRV& block_rv, int read_buffer_index,
                          const ffi::String& storage_scope) = 0;
  virtual SBlockRV WriteAt(const LoopRV& loop_rv, const SBlockRV& block_rv, int write_buffer_index,
                           const ffi::String& storage_scope) = 0;
  /******** Schedule: Compute location ********/
  virtual void ComputeAt(const SBlockRV& block_rv, const LoopRV& loop_rv, bool preserve_unit_loops,
                         int index = -1) = 0;
  virtual void ReverseComputeAt(const SBlockRV& block_rv, const LoopRV& loop_rv,
                                bool preserve_unit_loops, int index = -1) = 0;
  virtual void ComputeInline(const SBlockRV& block) = 0;
  virtual void ReverseComputeInline(const SBlockRV& block) = 0;
  virtual void FuseReductionEpilogue(const SBlockRV& reduction_block,
                                     const SBlockRV& epilogue_block) = 0;
  /******** Schedule: Reduction ********/
  virtual SBlockRV DecomposeReduction(const SBlockRV& block_rv, const LoopRV& loop_rv) = 0;
  virtual SBlockRV RFactor(const LoopRV& loop_rv, int factor_axis) = 0;
  /******** Schedule: SBlock annotation ********/
  virtual void StorageAlign(const SBlockRV& block_rv, int buffer_index, int axis, int factor,
                            int offset) = 0;
  virtual void SetScope(const SBlockRV& block_rv, int buffer_index,
                        const ffi::String& storage_scope) = 0;
  virtual void UnsafeSetDType(const SBlockRV& block_rv, int buffer_index,
                              const ffi::String& dtype) = 0;
  /******** Schedule: Blockize & Tensorize ********/
  virtual SBlockRV Blockize(const LoopRV& loop_rv, bool preserve_unit_iters = true) = 0;
  virtual SBlockRV Blockize(const ffi::Array<SBlockRV>& blocks,
                            bool preserve_unit_iters = true) = 0;
  virtual void Tensorize(const LoopRV& loop_rv, const ffi::String& intrin,
                         bool preserve_unit_iters = true) = 0;
  virtual void Tensorize(const SBlockRV& block_rv, const ffi::String& intrin,
                         bool preserve_unit_iters = true) = 0;
 
  /******** Schedule: Annotation ********/
  virtual void Annotate(const LoopRV& loop_rv, const ffi::String& ann_key, const Any& ann_val) = 0;
  virtual void Annotate(const SBlockRV& block_rv, const ffi::String& ann_key,
                        const Any& ann_val) = 0;
  virtual void Unannotate(const LoopRV& loop_rv, const ffi::String& ann_key) = 0;
  virtual void Unannotate(const SBlockRV& block_rv, const ffi::String& ann_key) = 0;
 
  /******** Schedule: Layout transformation ********/
  virtual void TransformLayout(const SBlockRV& block_rv, int buffer_index,
                               BufferIndexType buffer_index_type, const IndexMap& index_map,
                               const ffi::Optional<IndexMap>& pad_value = std::nullopt,
                               bool assume_injective_transform = false) = 0;
 
  virtual void TransformBlockLayout(const SBlockRV& block_rv, const IndexMap& index_map) = 0;
 
  /******** Schedule: Padding ********/
  virtual SBlockRV DecomposePadding(const SBlockRV& block_rv, const LoopRV& loop_rv) = 0;
 
  virtual void PadEinsum(const SBlockRV& block_rv, const ffi::Array<int64_t>& padding) = 0;
 
  /******** Schedule: Buffer transformation ********/
  virtual void RollingBuffer(const SBlockRV& block_rv, int write_buffer_index) = 0;
 
  virtual void AnnotateBufferAccess(const SBlockRV& block_rv, int buffer_index,
                                    BufferIndexType buffer_index_type,
                                    const IndexMap& index_map) = 0;
 
  /******** Schedule: Misc ********/
  virtual void EnterPostproc() = 0;
 
  virtual void UnsafeHideBufferAccess(const SBlockRV& block_rv, const ffi::String& buf_type,
                                      const ffi::Array<IntImm>& buf_index_array) = 0;
};
 
class Schedule : public ffi::ObjectRef {
 public:
  TVM_DLL static Schedule Concrete(IRModule mod, LinearCongruentialEngine::TRandState seed,
                                   int debug_mask, ScheduleErrorRenderLevel error_render_level,
                                   bool enable_check = true);
  TVM_DLL static Schedule Traced(IRModule mod, LinearCongruentialEngine::TRandState seed,
                                 int debug_mask, ScheduleErrorRenderLevel error_render_level,
                                 bool enable_check = true);
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Schedule, ffi::ObjectRef, ScheduleNode);
};
 
}  // namespace s_tir
}  // namespace tvm
 
#endif  // TVM_S_TIR_SCHEDULE_SCHEDULE_H_