stmt.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_STMT_H_
#define TVM_S_TIR_STMT_H_
 
#include <tvm/ffi/container/tuple.h>
#include <tvm/s_tir/iter_var.h>
#include <tvm/tirx/stmt.h>
 
namespace tvm {
namespace s_tir {
 
using TensorAllocatedAddresses = ffi::Array<ffi::Tuple<Var, ffi::Array<PrimExpr>>>;
 
class MatchTensorRegionNode : public ffi::Object {
 public:
  explicit MatchTensorRegionNode(ffi::UnsafeInit tag) : tensor(tag), source(tag) {}
 
  MatchTensorRegionNode(tirx::TensorVar tensor, TensorRegion source)
      : tensor(std::move(tensor)), source(std::move(source)) {}
 
  tirx::TensorVar tensor;
  TensorRegion source;
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<MatchTensorRegionNode>()
        .def_ro("tensor", &MatchTensorRegionNode::tensor,
                refl::AttachFieldFlag::SEqHashDefPattern())
        .def_ro("source", &MatchTensorRegionNode::source);
  }
 
  static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode;
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("s_tir.MatchTensorRegion", MatchTensorRegionNode, ffi::Object);
};
 
class MatchTensorRegion : public ffi::ObjectRef {
 public:
  explicit MatchTensorRegion(ffi::UnsafeInit tag, ffi::ObjectPtr<MatchTensorRegionNode> data)
      : ffi::ObjectRef(tag) {
    data_ = std::move(data);
  }
 
  TVM_DLL explicit MatchTensorRegion(tirx::TensorVar tensor, TensorRegion source);
 
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(MatchTensorRegion, ffi::ObjectRef,
                                                MatchTensorRegionNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(MatchTensorRegionNode);
};
 
class SBlockNode : public StmtNode {
 public:
  explicit SBlockNode(ffi::UnsafeInit tag) : body(tag) {}
 
  explicit SBlockNode(SeqStmt body) : body(std::move(body)) {}
 
  ffi::Array<s_tir::IterVar> iter_vars;
  ffi::Array<TensorRegion> reads;
  ffi::Array<TensorRegion> writes;
  ffi::String name_hint;
  ffi::Array<tirx::TensorVar> alloc_tensors;
  ffi::Array<MatchTensorRegion> match_tensors;
  ffi::Map<ffi::String, ffi::Any> annotations;
  ffi::Optional<SeqStmt> init;
  SeqStmt body;
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<SBlockNode>()
        .def_ro("iter_vars", &SBlockNode::iter_vars)
        .def_ro("reads", &SBlockNode::reads)
        .def_ro("writes", &SBlockNode::writes)
        .def_ro("name_hint", &SBlockNode::name_hint, refl::AttachFieldFlag::SEqHashIgnore())
        .def_ro("alloc_tensors", &SBlockNode::alloc_tensors,
                refl::AttachFieldFlag::SEqHashDefSimple())
        .def_ro("match_tensors", &SBlockNode::match_tensors)
        .def_ro("annotations", &SBlockNode::annotations)
        .def_ro("init", &SBlockNode::init)
        .def_ro("body", &SBlockNode::body);
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("s_tir.SBlock", SBlockNode, StmtNode);
};
 
class SBlock : public Stmt {
 public:
  TVM_DLL explicit SBlock(
      ffi::Array<s_tir::IterVar> iter_vars, ffi::Array<TensorRegion> reads,
      ffi::Array<TensorRegion> writes, ffi::String name_hint, SeqStmt body,
      ffi::Optional<SeqStmt> init = std::nullopt,
      ffi::Array<tirx::TensorVar> alloc_tensors = ffi::Array<tirx::TensorVar>(),
      ffi::Array<MatchTensorRegion> match_tensors = ffi::Array<MatchTensorRegion>(),
      ffi::Map<ffi::String, ffi::Any> annotations = ffi::Map<ffi::String, ffi::Any>(),
      Location loc = UnknownLoc());
 
  TVM_DLL explicit SBlock(ffi::String name_hint, SeqStmt body,
                          ffi::Array<tirx::TensorVar> alloc_tensors = ffi::Array<tirx::TensorVar>(),
                          Location loc = UnknownLoc());
 
  explicit SBlock(ffi::ObjectPtr<SBlockNode> node) : Stmt(std::move(node)) {}
 
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(SBlock, Stmt, SBlockNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(SBlockNode);
};
 
class SBlockRealizeNode : public StmtNode {
 public:
  explicit SBlockRealizeNode(ffi::UnsafeInit tag) : predicate(tag), block(tag) {}
 
  SBlockRealizeNode(PrimExpr predicate, SBlock block)
      : predicate(std::move(predicate)), block(std::move(block)) {}
 
  ffi::Array<PrimExpr> iter_values;
  PrimExpr predicate;
  SBlock block;
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<SBlockRealizeNode>()
        .def_ro("iter_values", &SBlockRealizeNode::iter_values)
        .def_ro("predicate", &SBlockRealizeNode::predicate)
        .def_ro("block", &SBlockRealizeNode::block);
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("s_tir.SBlockRealize", SBlockRealizeNode, StmtNode);
};
 
class SBlockRealize : public Stmt {
 public:
  TVM_DLL explicit SBlockRealize(ffi::Array<PrimExpr> iter_values, PrimExpr predicate, SBlock block,
                                 Location loc = UnknownLoc());
 
  explicit SBlockRealize(ffi::ObjectPtr<SBlockRealizeNode> node) : Stmt(std::move(node)) {}
 
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(SBlockRealize, Stmt, SBlockRealizeNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(SBlockRealizeNode);
};
 
TVM_DLL const Op& async_copy_scope();
TVM_DLL const Op& async_commit();
TVM_DLL const Op& async_wait();
TVM_DLL const Op& manual_sync();
 
namespace attr {
constexpr const char* kThreadBinding = "thread_binding";
 
constexpr const char* kDoubleBufferScope = "double_buffer_scope";
 
constexpr const char* kFragmentShape = "fragment_shape";
 
constexpr const char* kFragmentLayout = "fragment_layout";
 
constexpr const char* kLoopPartitionHint = "loop_partition_hint";
 
// -----------------------------------------------------------------------
// meta_schedule annotations
// -----------------------------------------------------------------------
 
constexpr const char* kMetaScheduleTilingStructure = "meta_schedule.tiling_structure";
 
constexpr const char* kMetaScheduleCooperativeFetch = "meta_schedule.cooperative_fetch";
 
constexpr const char* kMetaScheduleThreadExtentLowInclusive =
    "meta_schedule.thread_extent_low_inclusive";
 
constexpr const char* kMetaScheduleThreadExtentHighInclusive =
    "meta_schedule.thread_extent_high_inclusive";
 
constexpr const char* kMetaScheduleRandomComputeProducer = "meta_schedule.random_compute_producer";
 
constexpr const char* kMetaScheduleParallel = "meta_schedule.parallel";
 
constexpr const char* kMetaScheduleVectorize = "meta_schedule.vectorize";
 
constexpr const char* kMetaScheduleUnrollExplicit = "meta_schedule.unroll_explicit";
 
constexpr const char* kMetaScheduleUnrollImplicit = "meta_schedule.unroll_implicit";
 
constexpr const char* kMetaScheduleAutoTensorize = "meta_schedule.auto_tensorize";
 
constexpr const char* kMetaScheduleLayoutRewritePreproc = "meta_schedule.layout_rewrite_preproc";
 
constexpr const char* kMetaScheduleAutoTensorizeInit = "meta_schedule.auto_tensorize_init";
 
constexpr const char* kMetaScheduleInlineRule = "meta_schedule.inline_rule";
 
// -----------------------------------------------------------------------
// Schedule primitive / SBlock annotations
// -----------------------------------------------------------------------
 
constexpr const char* kTensorAllocatedAddr = "s_tir.tensor_allocated_addr";
 
constexpr const char* kScriptParsingDetectAccess = "tirx.script_parsing_detect_access";
 
constexpr const char* kRequireBlockVarBoundPredicate = "require_bound_predicate";
 
constexpr const char* kSoftwarePipelineStage = "software_pipeline_stage";
 
constexpr const char* kSoftwarePipelineOrder = "software_pipeline_order";
 
constexpr const char* kSoftwarePipelineAsyncStages = "software_pipeline_async_stages";
 
constexpr const char* kManifestSharedMemoryLocalStage = "tirx.manifest_shared_memory_local_stage";
 
constexpr const char* kTensorDimAlign = "tensor_dim_align";
 
constexpr const char* kExplicitReadRegion = "explicit_read_region";
 
constexpr const char* kExplicitWriteRegion = "explicit_write_region";
 
constexpr const char* kIrregularLoopMark = "irregular_loop_mark";
 
constexpr const char* kAutoCopy = "auto_copy";
 
constexpr const char* kLocalStage = "local_stage";
 
constexpr const char* kVectorBytes = "vector_bytes";
 
constexpr const char* kWarpExecution = "warp_execution";
 
constexpr const char* kPermutedLayout = "permuted_layout";
 
constexpr const char* kScheduleRule = "schedule_rule";
 
constexpr const char* kMetaScheduleWriteCacheLevel = "s_tir.meta_schedule.write_cache_level";
 
}  // namespace attr
 
inline ffi::Optional<ffi::String> GetThreadBinding(const ForNode* loop) {
  if (auto tag = loop->annotations.Get(tvm::s_tir::attr::kThreadBinding)) {
    TVM_FFI_CHECK(loop->kind == ForKind::kParallel, ValueError)
        << "thread_binding requires a parallel loop";
    TVM_FFI_CHECK(tag->as<ffi::String>().has_value(), TypeError)
        << "thread_binding annotation must be a string";
    TVM_FFI_CHECK(loop->HasTrivialStep(), ValueError) << "Thread binding loops require a unit step";
    return tag->cast<ffi::String>();
  }
  return std::nullopt;
}
inline ffi::Optional<ffi::String> GetThreadBinding(const For& loop) {
  return GetThreadBinding(loop.get());
}
 
}  // namespace s_tir
}  // namespace tvm
#endif  // TVM_S_TIR_STMT_H_