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.
 */
// Acknowledgement: Many low-level stmts originate from Halide.
#ifndef TVM_TIRX_STMT_H_
#define TVM_TIRX_STMT_H_
 
#include <tvm/ffi/reflection/registry.h>
#include <tvm/script/printer/config.h>
#include <tvm/tirx/exec_scope.h>
#include <tvm/tirx/expr.h>
#include <tvm/tirx/layout.h>
 
#include <optional>
#include <string>
#include <type_traits>
#include <utility>
 
namespace tvm {
namespace tirx {
 
class StmtNode : public ffi::Object {
 public:
  mutable Span span;
 
  StmtNode() = default;
  explicit StmtNode(Span span) : span(span) {}
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<StmtNode>().def_ro("span", &StmtNode::span);
  }
 
  TVM_OBJECT_ENABLE_SCRIPT_PRINTER();
 
  static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode;
 
  static constexpr const uint32_t _type_child_slots = 15;
  TVM_FFI_DECLARE_OBJECT_INFO("tirx.Stmt", StmtNode, ffi::Object);
};
 
class Stmt : public ffi::ObjectRef {
 public:
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Stmt, ffi::ObjectRef, StmtNode);
};
 
class BindNode : public StmtNode {
 public:
  Var var;
  PrimExpr value;
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<BindNode>()
        .def_ro("var", &BindNode::var, refl::AttachFieldFlag::SEqHashDef())
        .def_ro("value", &BindNode::value);
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.Bind", BindNode, StmtNode);
};
 
class Bind : public Stmt {
 public:
  TVM_DLL Bind(Var var, PrimExpr value, Span span = Span());
 
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Bind, Stmt, BindNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(BindNode);
};
 
class AttrStmtNode : public StmtNode {
 public:
  ffi::Any node;
  ffi::String attr_key;
  PrimExpr value;
  Stmt body;
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<AttrStmtNode>()
        .def_ro("node", &AttrStmtNode::node)
        .def_ro("attr_key", &AttrStmtNode::attr_key)
        .def_ro("value", &AttrStmtNode::value)
        .def_ro("body", &AttrStmtNode::body);
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.AttrStmt", AttrStmtNode, StmtNode);
};
 
class AttrStmt : public Stmt {
 public:
  TVM_DLL AttrStmt(ffi::Any node, ffi::String attr_key, PrimExpr value, Stmt body,
                   Span span = Span());
 
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(AttrStmt, Stmt, AttrStmtNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(AttrStmtNode);
};
 
class AssertStmtNode : public StmtNode {
 public:
  PrimExpr condition;
  StringImm error_kind;
  ffi::Array<StringImm> message_parts;
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<AssertStmtNode>()
        .def_ro("condition", &AssertStmtNode::condition)
        .def_ro("error_kind", &AssertStmtNode::error_kind)
        .def_ro("message_parts", &AssertStmtNode::message_parts);
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.AssertStmt", AssertStmtNode, StmtNode);
};
 
class AssertStmt : public Stmt {
 public:
  TVM_DLL AssertStmt(PrimExpr condition, StringImm error_kind, ffi::Array<StringImm> message_parts,
                     Span span = Span());
 
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(AssertStmt, Stmt, AssertStmtNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(AssertStmtNode);
};
 
class BufferStoreNode : public StmtNode {
 public:
  Buffer buffer;
  PrimExpr value;
  ffi::Array<PrimExpr> indices;
  ffi::Optional<PrimExpr> predicate;
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<BufferStoreNode>()
        .def_ro("buffer", &BufferStoreNode::buffer)
        .def_ro("value", &BufferStoreNode::value)
        .def_ro("indices", &BufferStoreNode::indices)
        .def_ro("predicate", &BufferStoreNode::predicate);
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.BufferStore", BufferStoreNode, StmtNode);
};
 
class BufferStore : public Stmt {
 public:
  TVM_DLL explicit BufferStore(Buffer buffer, PrimExpr value, ffi::Array<PrimExpr> indices,
                               ffi::Optional<PrimExpr> predicate = std::nullopt,
                               Span span = Span());
 
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(BufferStore, Stmt, BufferStoreNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(BufferStoreNode);
};
 
class DeclBufferNode : public StmtNode {
 public:
  Buffer buffer;
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<DeclBufferNode>().def_ro("buffer", &DeclBufferNode::buffer);
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.DeclBuffer", DeclBufferNode, StmtNode);
};
 
class DeclBuffer : public Stmt {
 public:
  TVM_DLL DeclBuffer(Buffer buffer, Span span = Span());
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(DeclBuffer, Stmt, DeclBufferNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(DeclBufferNode);
};
 
class AllocBufferNode : public StmtNode {
 public:
  Buffer buffer;
  ffi::Map<ffi::String, ffi::Any> annotations;
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<AllocBufferNode>()
        .def_ro("buffer", &AllocBufferNode::buffer, refl::AttachFieldFlag::SEqHashDef())
        .def_ro("annotations", &AllocBufferNode::annotations);
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.AllocBuffer", AllocBufferNode, StmtNode);
};
 
class AllocBuffer : public Stmt {
 public:
  TVM_DLL AllocBuffer(
      Buffer buffer,
      ffi::Map<ffi::String, ffi::Any> annotations = ffi::Map<ffi::String, ffi::Any>(),
      Span span = Span());
  std::optional<int64_t> ConstantAllocationSize() const {
    int64_t result = 1;
    for (const PrimExpr& extent : (*this)->buffer->shape) {
      if (const auto* int_size = extent.as<IntImmNode>()) {
        result *= int_size->value;
      } else {
        return std::nullopt;
      }
    }
    return result;
  }
 
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(AllocBuffer, Stmt, AllocBufferNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(AllocBufferNode);
};
 
class SeqStmtNode : public StmtNode {
 public:
  ffi::Array<Stmt> seq;
 
  size_t size() const { return seq.size(); }
  Stmt operator[](size_t index) const { return seq[index]; }
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<SeqStmtNode>().def_ro("seq", &SeqStmtNode::seq);
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.SeqStmt", SeqStmtNode, StmtNode);
};
 
class EvaluateNode : public StmtNode {
 public:
  PrimExpr value;
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<EvaluateNode>().def_ro("value", &EvaluateNode::value);
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.Evaluate", EvaluateNode, StmtNode);
};
 
class Evaluate : public Stmt {
 public:
  TVM_DLL explicit Evaluate(PrimExpr value, Span span = Span());
 
  explicit Evaluate(int value, Span span = Span()) : Evaluate(PrimExpr(value), span) {}
 
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Evaluate, Stmt, EvaluateNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(EvaluateNode);
};
 
class SeqStmt : public Stmt {
 public:
  TVM_DLL explicit SeqStmt(ffi::Array<Stmt> seq, Span span = Span());
 
  size_t size() const { return operator->()->size(); }
  Stmt operator[](size_t index) const { return (*(operator->()))[index]; }
  template <typename... Args>
  static Stmt Flatten(Args&&... seq_args) {
    ffi::Array<Stmt> seq;
 
    ffi::details::for_each(Flattener(&seq), std::forward<Args>(seq_args)...);
 
    if (seq.empty()) {
      return Evaluate(0);
    } else if (seq.size() == 1) {
      return seq[0];
    }
 
    // If the argument is a single SeqStmt argument with no
    // flattening or unwrapping required, then we may
    // return the SeqStmt as-is.
    if constexpr (sizeof...(seq_args) == 1) {
      if (auto opt = Flattener::AsSeqStmt(std::forward<Args>(seq_args)...)) {
        SeqStmt original = opt.value();
        bool all_same = [&]() {
          if (original->seq.size() != seq.size()) {
            return false;
          }
          for (size_t i = 0; i < seq.size(); i++) {
            if (!original->seq[i].same_as(seq[i])) {
              return false;
            }
          }
          return true;
        }();
        if (all_same) {
          return original;
        }
      }
    }
 
    return SeqStmt(seq);
  }
  class Flattener {
   public:
    explicit Flattener(ffi::Array<Stmt>* seq) : seq_(seq) {}
 
    template <typename T>
    static ffi::Optional<SeqStmt> AsSeqStmt(const T& t) {
      if constexpr (std::is_same_v<T, SeqStmt>) {
        return t;
      }
      if constexpr (!std::is_base_of_v<T, SeqStmt>) {
        return std::nullopt;
      }
      if constexpr (std::is_base_of_v<Stmt, T>) {
        if (const SeqStmtNode* ptr = t.template as<SeqStmtNode>()) {
          return ffi::GetRef<SeqStmt>(ptr);
        } else {
          return std::nullopt;
        }
      }
      return std::nullopt;
    }
 
    template <typename T>
    void operator()(size_t i, const T& stmt_or_seq) const {
      if constexpr (std::is_base_of_v<ObjectRef, T>) {
        // Early bail-out, applicable to any ObjectRef
        if (!stmt_or_seq.defined()) {
          return;
        }
      }
 
      if constexpr (std::is_same_v<T, SeqStmt>) {
        // Static type-checking for a SeqStmt that could be flattened.
        (*this)(0, stmt_or_seq->seq);
        return;
      }
 
      if constexpr (std::is_base_of_v<T, SeqStmt>) {
        // Dynamic type-checking for a SeqStmt that could be
        // flattened.
        if (auto* op = stmt_or_seq.template as<SeqStmtNode>()) {
          operator()(0, op->seq);
          return;
        }
      }
 
      if constexpr (std::is_base_of_v<T, Evaluate>) {
        // Evaluate(0) is used to represent a no-op, and may be
        // generated by previous calls to SeqStmt::Flatten().  These
        // should be removed to ensure that Flatten(a+b) is equivalent
        // to Flatten(Flatten(a), Flatten(b)).
        if (auto* op = stmt_or_seq.template as<EvaluateNode>()) {
          if (auto* as_int = op->value.template as<IntImmNode>(); as_int && as_int->value == 0) {
            return;
          }
        }
      }
 
      if constexpr (std::is_base_of_v<Stmt, T>) {
        // Any other Stmt type just gets appended.
        seq_->push_back(stmt_or_seq);
      } else {
        // Anything else is treated as an iterable of Stmt.
        for (auto v : stmt_or_seq) {
          this->operator()(0, v);
        }
      }
    }
 
   private:
    ffi::Array<Stmt>* seq_;
  };
 
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(SeqStmt, Stmt, SeqStmtNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(SeqStmtNode);
};
 
class IfThenElseNode : public StmtNode {
 public:
  PrimExpr condition;
  Stmt then_case;
  ffi::Optional<Stmt> else_case;
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<IfThenElseNode>()
        .def_ro("condition", &IfThenElseNode::condition)
        .def_ro("then_case", &IfThenElseNode::then_case)
        .def_ro("else_case", &IfThenElseNode::else_case);
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.IfThenElse", IfThenElseNode, StmtNode);
};
 
class IfThenElse : public Stmt {
 public:
  TVM_DLL IfThenElse(PrimExpr condition, Stmt then_case,
                     ffi::Optional<Stmt> else_case = std::nullopt, Span span = Span());
 
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(IfThenElse, Stmt, IfThenElseNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(IfThenElseNode);
};
 
enum class ForKind : int {
  kSerial = 0,
  kParallel = 1,
  kVectorized = 2,
  kUnrolled = 3,
  kThreadBinding = 4
};
 
class ForNode : public StmtNode {
 public:
  Var loop_var;
  PrimExpr min;
  PrimExpr extent;
  ForKind kind;
  Stmt body;
  ffi::Optional<IterVar> thread_binding;
  ffi::Map<ffi::String, ffi::Any> annotations;
  ffi::Optional<PrimExpr> step;
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<ForNode>()
        .def_ro("loop_var", &ForNode::loop_var, refl::AttachFieldFlag::SEqHashDef())
        .def_ro("min", &ForNode::min)
        .def_ro("extent", &ForNode::extent)
        .def_ro("kind", &ForNode::kind)
        .def_ro("body", &ForNode::body)
        .def_ro("thread_binding", &ForNode::thread_binding)
        .def_ro("annotations", &ForNode::annotations)
        .def_ro("step", &ForNode::step);
  }
 
  bool HasTrivialStep() const;
 
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.For", ForNode, StmtNode);
};
 
class For : public Stmt {
 public:
  TVM_DLL For(Var loop_var, PrimExpr min, PrimExpr extent, ForKind kind, Stmt body,
              ffi::Optional<IterVar> thread_binding = std::nullopt,
              ffi::Map<ffi::String, ffi::Any> annotations = {},
              ffi::Optional<PrimExpr> step = std::nullopt, Span span = Span());
 
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(For, Stmt, ForNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(ForNode);
};
 
class WhileNode : public StmtNode {
 public:
  PrimExpr condition;
  Stmt body;
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<WhileNode>()
        .def_ro("condition", &WhileNode::condition)
        .def_ro("body", &WhileNode::body);
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.While", WhileNode, StmtNode);
};
 
class While : public Stmt {
 public:
  TVM_DLL While(PrimExpr condition, Stmt body, Span span = Span());
 
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(While, Stmt, WhileNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(WhileNode);
};
 
class BreakNode : public StmtNode {
 public:
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<BreakNode>();
  }
 
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.Break", BreakNode, StmtNode);
};
 
class Break : public Stmt {
 public:
  TVM_DLL explicit Break(Span span);
 
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Break, Stmt, BreakNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(BreakNode);
};
 
class ContinueNode : public StmtNode {
 public:
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<ContinueNode>();
  }
 
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.Continue", ContinueNode, StmtNode);
};
 
class Continue : public Stmt {
 public:
  TVM_DLL explicit Continue(Span span);
 
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Continue, Stmt, ContinueNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(ContinueNode);
};
 
class BufferRegionNode : public PrimExprConvertibleNode {
 public:
  Buffer buffer;
  ffi::Array<Range> region;
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<BufferRegionNode>()
        .def_ro("buffer", &BufferRegionNode::buffer)
        .def_ro("region", &BufferRegionNode::region);
  }
 
  TVM_DLL PrimExpr ToPrimExpr() const final;
 
  static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode;
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.BufferRegion", BufferRegionNode, PrimExprConvertibleNode);
};
 
class BufferRegion : public PrimExprConvertible {
 public:
  TVM_DLL explicit BufferRegion(Buffer buffer, ffi::Array<Range> region);
 
  TVM_DLL static BufferRegion FullRegion(Buffer buffer);
 
  TVM_DLL static BufferRegion FromPoint(Buffer buffer, ffi::Array<PrimExpr> indices);
 
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(BufferRegion, PrimExprConvertible, BufferRegionNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(BufferRegionNode);
};
 
class MatchBufferRegionNode : public ffi::Object {
 public:
  Buffer buffer;
  BufferRegion source;
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<MatchBufferRegionNode>()
        .def_ro("buffer", &MatchBufferRegionNode::buffer)
        .def_ro("source", &MatchBufferRegionNode::source);
  }
 
  static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode;
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.MatchBufferRegion", MatchBufferRegionNode, ffi::Object);
};
 
class MatchBufferRegion : public ffi::ObjectRef {
 public:
  TVM_DLL explicit MatchBufferRegion(Buffer buffer, BufferRegion source);
 
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(MatchBufferRegion, ffi::ObjectRef,
                                             MatchBufferRegionNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(MatchBufferRegionNode);
};
 
class SBlockNode : public StmtNode {
 public:
  ffi::Array<IterVar> iter_vars;
  ffi::Array<BufferRegion> reads;
  ffi::Array<BufferRegion> writes;
  ffi::String name_hint;
  ffi::Array<Buffer> alloc_buffers;
  ffi::Array<MatchBufferRegion> match_buffers;
  ffi::Map<ffi::String, ffi::Any> annotations;
  ffi::Optional<Stmt> init;
  Stmt body;
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<SBlockNode>()
        .def_ro("iter_vars", &SBlockNode::iter_vars, refl::AttachFieldFlag::SEqHashDef())
        .def_ro("reads", &SBlockNode::reads)
        .def_ro("writes", &SBlockNode::writes)
        .def_ro("name_hint", &SBlockNode::name_hint, refl::AttachFieldFlag::SEqHashIgnore())
        .def_ro("alloc_buffers", &SBlockNode::alloc_buffers)
        .def_ro("match_buffers", &SBlockNode::match_buffers)
        .def_ro("annotations", &SBlockNode::annotations)
        .def_ro("init", &SBlockNode::init)
        .def_ro("body", &SBlockNode::body);
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.SBlock", SBlockNode, StmtNode);
};
 
class SBlock : public Stmt {
 public:
  TVM_DLL explicit SBlock(
      ffi::Array<IterVar> iter_vars, ffi::Array<BufferRegion> reads,
      ffi::Array<BufferRegion> writes, ffi::String name_hint, Stmt body,
      ffi::Optional<Stmt> init = std::nullopt,
      ffi::Array<Buffer> alloc_buffers = ffi::Array<Buffer>(),
      ffi::Array<MatchBufferRegion> match_buffers = ffi::Array<MatchBufferRegion>(),
      ffi::Map<ffi::String, ffi::Any> annotations = ffi::Map<ffi::String, ffi::Any>(),
      Span span = Span());
 
  TVM_DLL explicit SBlock(ffi::String name_hint, Stmt body,
                          ffi::Array<Buffer> alloc_buffers = ffi::Array<Buffer>(),
                          Span span = Span());
 
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(SBlock, Stmt, SBlockNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(SBlockNode);
};
 
class SBlockRealizeNode : public StmtNode {
 public:
  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("tirx.SBlockRealize", SBlockRealizeNode, StmtNode);
};
 
class SBlockRealize : public Stmt {
 public:
  TVM_DLL explicit SBlockRealize(ffi::Array<PrimExpr> iter_values, PrimExpr predicate, SBlock block,
                                 Span span = Span());
 
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(SBlockRealize, Stmt, SBlockRealizeNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(SBlockRealizeNode);
};
 
class ExecScopeStmtNode : public StmtNode {
 public:
  ExecScope exec_scope;
  Stmt body;
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<ExecScopeStmtNode>()
        .def_ro("exec_scope", &ExecScopeStmtNode::exec_scope)
        .def_ro("body", &ExecScopeStmtNode::body);
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.ExecScopeStmt", ExecScopeStmtNode, StmtNode);
};
 
class ExecScopeStmt : public Stmt {
 public:
  TVM_DLL ExecScopeStmt(ExecScope exec_scope, Stmt body, Span span = Span());
 
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(ExecScopeStmt, Stmt, ExecScopeStmtNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(ExecScopeStmtNode);
};
 
namespace attr {
constexpr const char* buffer_bound = "buffer_bound";
constexpr const char* compute_scope = "compute_scope";
constexpr const char* device_id = "device_id";
constexpr const char* device_scope = "device_scope";
constexpr const char* device_type = "device_type";
constexpr const char* extern_scope = "extern_scope";
constexpr const char* pragma_auto_unroll_max_step = "pragma_auto_unroll_max_step";
constexpr const char* pragma_import_c = "pragma_import_c";
constexpr const char* pragma_import_llvm = "pragma_import_llvm";
constexpr const char* pragma_scope_prefix = "pragma_";
constexpr const char* pragma_tensor_core = "pragma_tensor_core";
constexpr const char* pragma_unroll_explicit = "pragma_unroll_explicit";
constexpr const char* storage_alignment = "storage_alignment";
constexpr const char* thread_extent = "thread_extent";
constexpr const char* kVolatile = "tirx.volatile";
constexpr const char* layout_transforms = "layout_transforms";
constexpr const char* axis_separators = "axis_separators";
constexpr const char* double_buffer_scope = "double_buffer_scope";
constexpr const char* double_buffer_write = "double_buffer_write";
constexpr const char* scan_update_scope = "scan_update_scope";
constexpr const char* scan_init_scope = "scan_init_scope";
constexpr const char* buffer_dim_align = "buffer_dim_align";
constexpr const char* buffer_data_alignment = "buffer_data_alignment";
constexpr const char* buffer_allocated_addr = "buffer_allocated_addr";
constexpr const char* buffer_bind_scope = "buffer_bind_scope";
// Pipeline related attributes
constexpr const char* channel_read_scope = "channel_read_scope";
constexpr const char* channel_read_advance = "channel_read_advance";
constexpr const char* channel_write_scope = "channel_write_scope";
constexpr const char* channel_write_advance = "channel_write_advance";
constexpr const char* pipeline_stage_scope = "pipeline_stage_scope";
constexpr const char* pipeline_exec_scope = "pipeline_exec_scope";
 
constexpr const char* async_scope = "async_scope";
 
constexpr const char* async_commit_queue_scope = "async_commit_queue_scope";
constexpr const char* async_wait_queue_scope = "async_wait_queue_scope";
constexpr const char* async_wait_inflight_count = "async_wait_inflight_count";
 
constexpr const char* fragment_shape = "fragment_shape";
 
constexpr const char* fragment_layout = "fragment_layout";
 
constexpr const char* hand_threaded = "hand_threaded";
 
constexpr const char* script_parsing_detect_access = "tirx.script_parsing_detect_access";
 
constexpr const char* pragma_loop_partition_hint = "pragma_loop_partition_hint";
 
constexpr const char* software_pipeline_stage = "software_pipeline_stage";
 
constexpr const char* software_pipeline_order = "software_pipeline_order";
 
constexpr const char* software_pipeline_async_stages = "software_pipeline_async_stages";
 
constexpr const char* layout_free_buffers = "layout_free_buffers";
 
constexpr const char* manifest_shared_memory_local_stage =
    "tirx.manifest_shared_memory_local_stage";
 
constexpr const char* meta_schedule_tiling_structure = "meta_schedule.tiling_structure";
 
constexpr const char* meta_schedule_cooperative_fetch = "meta_schedule.cooperative_fetch";
 
constexpr const char* meta_schedule_thread_extent_low_inclusive =
    "meta_schedule.thread_extent_low_inclusive";
 
constexpr const char* meta_schedule_thread_extent_high_inclusive =
    "meta_schedule.thread_extent_high_inclusive";
 
constexpr const char* meta_schedule_random_compute_producer =
    "meta_schedule.random_compute_producer";
 
constexpr const char* meta_schedule_parallel = "meta_schedule.parallel";
 
constexpr const char* meta_schedule_vectorize = "meta_schedule.vectorize";
 
constexpr const char* meta_schedule_unroll_explicit = "meta_schedule.unroll_explicit";
 
constexpr const char* meta_schedule_unroll_implicit = "meta_schedule.unroll_implicit";
 
constexpr const char* meta_schedule_auto_tensorize = "meta_schedule.auto_tensorize";
 
constexpr const char* meta_schedule_layout_rewrite_preproc = "meta_schedule.layout_rewrite_preproc";
constexpr const char* meta_schedule_auto_tensorize_init = "meta_schedule.auto_tensorize_init";
 
constexpr const char* require_block_var_bound_predicate = "require_bound_predicate";
 
constexpr const char* meta_schedule_tensor_core_enabled = "meta_schedule.tensor_core_enabled";
 
constexpr const char* meta_schedule_cache_type = "meta_schedule.cache_type";
 
constexpr const int meta_schedule_cache_type_read = 0;
 
constexpr const int meta_schedule_cache_type_write = 1;
 
constexpr const char* auto_copy = "auto_copy";
 
constexpr const char* local_stage = "local_stage";
 
constexpr const char* vector_bytes = "vector_bytes";
 
constexpr const char* warp_execution = "warp_execution";
 
constexpr const char* meta_schedule_inline_rule = "meta_schedule.inline_rule";
 
constexpr const char* explicit_read_region = "explicit_read_region";
 
constexpr const char* explicit_write_region = "explicit_write_region";
constexpr const char* tensorized_nki_instruction = "tensorized_nki_instruction";
 
constexpr const char* irregular_loop_mark = "irregular_loop_mark";
 
constexpr const char* kPersistentKernel = "tirx.persistent_kernel";
 
inline bool IsPragmaKey(const std::string& attr_key) {
  return attr_key.compare(0, 7, "pragma_") == 0;
}
 
}  // namespace attr
TVM_DLL PrimExpr TypeAnnotation(DataType dtype, Span span = Span());
 
// overload printing of for type.
TVM_DLL std::ostream& operator<<(std::ostream& os, ForKind kind);
 
// inline implementations
inline const char* ForKind2String(ForKind t) {
  switch (t) {
    case ForKind::kSerial:
      return "serial";
    case ForKind::kParallel:
      return "parallel";
    case ForKind::kVectorized:
      return "vectorized";
    case ForKind::kUnrolled:
      return "unroll";
    case ForKind::kThreadBinding:
      return "thread_binding";
  }
  TVM_FFI_THROW(InternalError) << "Unknown ForKind" << t;
  TVM_FFI_UNREACHABLE();
}
 
}  // namespace tirx
}  // namespace tvm
#endif  // TVM_TIR_STMT_H_