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_TIR_STMT_H_
#define TVM_TIR_STMT_H_
 
#include <tvm/ffi/reflection/registry.h>
#include <tvm/tir/expr.h>
 
#include <string>
#include <type_traits>
#include <utility>
 
namespace tvm {
namespace tir {
 
class StmtNode : public 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("tir.Stmt", StmtNode, Object);
};
 
class Stmt : public ObjectRef {
 public:
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Stmt, ObjectRef, StmtNode);
};
 
class LetStmtNode : public StmtNode {
 public:
  Var var;
  PrimExpr value;
  Stmt body;
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<LetStmtNode>()
        .def_ro("var", &LetStmtNode::var, refl::AttachFieldFlag::SEqHashDef())
        .def_ro("value", &LetStmtNode::value)
        .def_ro("body", &LetStmtNode::body);
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tir.LetStmt", LetStmtNode, StmtNode);
};
 
class LetStmt : public Stmt {
 public:
  TVM_DLL LetStmt(Var var, PrimExpr value, Stmt body, Span span = Span());
 
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(LetStmt, Stmt, LetStmtNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(LetStmtNode);
};
 
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("tir.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;
  PrimExpr message;
  Stmt body;
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<AssertStmtNode>()
        .def_ro("condition", &AssertStmtNode::condition)
        .def_ro("message", &AssertStmtNode::message)
        .def_ro("body", &AssertStmtNode::body);
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tir.AssertStmt", AssertStmtNode, StmtNode);
};
 
class AssertStmt : public Stmt {
 public:
  TVM_DLL AssertStmt(PrimExpr condition, PrimExpr message, Stmt body, 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("tir.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 BufferRealizeNode : public StmtNode {
 public:
  Buffer buffer;
  ffi::Array<Range> bounds;
  PrimExpr condition;
  Stmt body;
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<BufferRealizeNode>()
        .def_ro("buffer", &BufferRealizeNode::buffer)
        .def_ro("bounds", &BufferRealizeNode::bounds)
        .def_ro("condition", &BufferRealizeNode::condition)
        .def_ro("body", &BufferRealizeNode::body);
  }
 
  BufferRealizeNode() = default;
  BufferRealizeNode(Buffer buffer, ffi::Array<Range> bounds, PrimExpr condition, Stmt body,
                    Span span = Span())
      : StmtNode(span), buffer(buffer), bounds(bounds), condition(condition), body(body) {}
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tir.BufferRealize", BufferRealizeNode, StmtNode);
};
 
class BufferRealize : public Stmt {
 public:
  TVM_DLL explicit BufferRealize(Buffer buffer, ffi::Array<Range> bounds, PrimExpr condition,
                                 Stmt body, Span span = Span());
 
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(BufferRealize, Stmt, BufferRealizeNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(BufferRealizeNode);
};
 
class AllocateNode : public StmtNode {
 public:
  Var buffer_var;
  DataType dtype;
  ffi::Array<PrimExpr> extents;
  PrimExpr condition;
  Stmt body;
  ffi::Map<ffi::String, ffi::Any> annotations;
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<AllocateNode>()
        .def_ro("buffer_var", &AllocateNode::buffer_var, refl::AttachFieldFlag::SEqHashDef())
        .def_ro("dtype", &AllocateNode::dtype)
        .def_ro("extents", &AllocateNode::extents)
        .def_ro("condition", &AllocateNode::condition)
        .def_ro("body", &AllocateNode::body)
        .def_ro("annotations", &AllocateNode::annotations);
  }
 
  int64_t ConstantAllocationSize() const { return ConstantAllocationSize(extents); }
  TVM_DLL static int64_t ConstantAllocationSize(const ffi::Array<PrimExpr>& extents);
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tir.Allocate", AllocateNode, StmtNode);
};
 
class Allocate : public Stmt {
 public:
  TVM_DLL Allocate(Var buffer_var, DataType dtype, ffi::Array<PrimExpr> extents, PrimExpr condition,
                   Stmt body,
                   ffi::Map<ffi::String, ffi::Any> annotations = ffi::Map<ffi::String, ffi::Any>(),
                   Span span = Span());
 
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Allocate, Stmt, AllocateNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(AllocateNode);
};
 
class AllocateConstNode : public StmtNode {
 public:
  Var buffer_var;
  ffi::Optional<runtime::Tensor> data;
  ffi::Optional<Integer> irmod_storage_idx;
  DataType dtype;
  ffi::Array<PrimExpr> extents;
  Stmt body;
  ffi::Map<ffi::String, ffi::Any> annotations;
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<AllocateConstNode>()
        .def_ro("buffer_var", &AllocateConstNode::buffer_var, refl::AttachFieldFlag::SEqHashDef())
        .def_ro("data", &AllocateConstNode::data)
        .def_ro("irmod_storage_idx", &AllocateConstNode::irmod_storage_idx)
        .def_ro("dtype", &AllocateConstNode::dtype)
        .def_ro("extents", &AllocateConstNode::extents)
        .def_ro("body", &AllocateConstNode::body)
        .def_ro("annotations", &AllocateConstNode::annotations);
  }
 
  int64_t ConstantAllocationSize() const { return ConstantAllocationSize(extents); }
  TVM_DLL static int64_t ConstantAllocationSize(const ffi::Array<PrimExpr>& extents);
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tir.AllocateConst", AllocateConstNode, StmtNode);
};
 
class AllocateConst : public Stmt {
 public:
  /* The constructor to create a IRNode with constant data
   * depending on the type of ObjectRef, it will either
   * create AllocateConstNode with irmod_storage_idx or data
   */
  TVM_DLL AllocateConst(
      Var buffer_var, DataType dtype, ffi::Array<PrimExpr> extents, ObjectRef data_or_idx,
      Stmt body, ffi::Map<ffi::String, ffi::Any> annotations = ffi::Map<ffi::String, ffi::Any>(),
      Span span = Span());
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(AllocateConst, Stmt, AllocateConstNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(AllocateConstNode);
};
 
class DeclBufferNode : public StmtNode {
 public:
  Buffer buffer;
  Stmt body;
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<DeclBufferNode>()
        .def_ro("buffer", &DeclBufferNode::buffer)
        .def_ro("body", &DeclBufferNode::body);
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tir.DeclBuffer", DeclBufferNode, StmtNode);
};
 
class DeclBuffer : public Stmt {
 public:
  TVM_DLL DeclBuffer(Buffer buffer, Stmt body, Span span = Span());
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(DeclBuffer, Stmt, DeclBufferNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(DeclBufferNode);
};
 
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("tir.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("tir.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("tir.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("tir.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("tir.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 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("tir.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 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("tir.MatchBufferRegion", MatchBufferRegionNode, Object);
};
 
class MatchBufferRegion : public ObjectRef {
 public:
  TVM_DLL explicit MatchBufferRegion(Buffer buffer, BufferRegion source);
 
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(MatchBufferRegion, ObjectRef, MatchBufferRegionNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(MatchBufferRegionNode);
};
 
class BlockNode : 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<BlockNode>()
        .def_ro("iter_vars", &BlockNode::iter_vars, refl::AttachFieldFlag::SEqHashDef())
        .def_ro("reads", &BlockNode::reads)
        .def_ro("writes", &BlockNode::writes)
        .def_ro("name_hint", &BlockNode::name_hint, refl::AttachFieldFlag::SEqHashIgnore())
        .def_ro("alloc_buffers", &BlockNode::alloc_buffers)
        .def_ro("match_buffers", &BlockNode::match_buffers)
        .def_ro("annotations", &BlockNode::annotations)
        .def_ro("init", &BlockNode::init)
        .def_ro("body", &BlockNode::body);
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tir.Block", BlockNode, StmtNode);
};
 
class Block : public Stmt {
 public:
  TVM_DLL explicit Block(
      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_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Block, Stmt, BlockNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(BlockNode);
};
 
class BlockRealizeNode : public StmtNode {
 public:
  ffi::Array<PrimExpr> iter_values;
  PrimExpr predicate;
  Block block;
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<BlockRealizeNode>()
        .def_ro("iter_values", &BlockRealizeNode::iter_values)
        .def_ro("predicate", &BlockRealizeNode::predicate)
        .def_ro("block", &BlockRealizeNode::block);
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tir.BlockRealize", BlockRealizeNode, StmtNode);
};
 
class BlockRealize : public Stmt {
 public:
  TVM_DLL explicit BlockRealize(ffi::Array<PrimExpr> iter_values, PrimExpr predicate, Block block,
                                Span span = Span());
 
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(BlockRealize, Stmt, BlockRealizeNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(BlockRealizeNode);
};
 
namespace attr {
// The above attr does not pass to ir stage.
constexpr const char* thread_extent = "thread_extent";
constexpr const char* virtual_thread = "virtual_thread";
constexpr const char* coproc_scope = "coproc_scope";
constexpr const char* coproc_uop_scope = "coproc_uop_scope";
constexpr const char* volatile_scope = "volatile_scope";
constexpr const char* extern_scope = "extern_scope";
constexpr const char* compute_scope = "compute_scope";
constexpr const char* storage_alignment = "storage_alignment";
constexpr const char* realize_scope = "realize_scope";
constexpr const char* device_id = "device_id";
constexpr const char* device_type = "device_type";
constexpr const char* loop_scope = "loop_scope";
constexpr const char* reduce_scope = "reduce_scope";
constexpr const char* pragma_auto_unroll_max_step = "pragma_auto_unroll_max_step";
constexpr const char* pragma_unroll_explicit = "pragma_unroll_explicit";
constexpr const char* pragma_scope_prefix = "pragma_";
constexpr const char* pragma_import_c = "pragma_import_c";
constexpr const char* pragma_import_llvm = "pragma_import_llvm";
constexpr const char* pragma_tensor_core = "pragma_tensor_core";
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* rolling_buffer_scope = "rolling_buffer_scope";
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_bound = "buffer_bound";
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* device_scope = "device_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 = "tir.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 = "tir.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* irregular_loop_mark = "irregular_loop_mark";
 
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";
  }
  LOG(FATAL) << "Unknown ForKind" << t;
}
 
}  // namespace tir
}  // namespace tvm
#endif  // TVM_TIR_STMT_H_