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/tir/expr.h>

#include <string>
#include <type_traits>
#include <utility>
#include <vector>

namespace tvm {
namespace tir {

class StmtNode : public Object {
 public:
  mutable Span span;

  StmtNode() = default;
  explicit StmtNode(Span span) : span(span) {}

  static constexpr const char* _type_key = "tir.Stmt";
  static constexpr const bool _type_has_method_sequal_reduce = true;
  static constexpr const bool _type_has_method_shash_reduce = true;
  static constexpr const uint32_t _type_child_slots = 15;
  TVM_DECLARE_BASE_OBJECT_INFO(StmtNode, Object);
};

class Stmt : public ObjectRef {
 public:
  TVM_DEFINE_OBJECT_REF_METHODS(Stmt, ObjectRef, StmtNode);
};

class LetStmtNode : public StmtNode {
 public:
  Var var;
  PrimExpr value;
  Stmt body;

  void VisitAttrs(AttrVisitor* v) {
    v->Visit("var", &var);
    v->Visit("value", &value);
    v->Visit("body", &body);
    v->Visit("span", &span);
  }

  bool SEqualReduce(const LetStmtNode* other, SEqualReducer equal) const {
    return equal.DefEqual(var, other->var) && equal(value, other->value) &&
           equal(body, other->body);
  }

  void SHashReduce(SHashReducer hash_reduce) const {
    hash_reduce.DefHash(var);
    hash_reduce(value);
    hash_reduce(body);
  }

  static constexpr const char* _type_key = "tir.LetStmt";
  TVM_DECLARE_FINAL_OBJECT_INFO(LetStmtNode, StmtNode);
};

class LetStmt : public Stmt {
 public:
  TVM_DLL LetStmt(Var var, PrimExpr value, Stmt body, Span span = Span());

  TVM_DEFINE_OBJECT_REF_METHODS(LetStmt, Stmt, LetStmtNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(LetStmtNode);
};

class AttrStmtNode : public StmtNode {
 public:
  ObjectRef node;
  String attr_key;
  PrimExpr value;
  Stmt body;

  void VisitAttrs(AttrVisitor* v) {
    v->Visit("node", &node);
    v->Visit("attr_key", &attr_key);
    v->Visit("value", &value);
    v->Visit("body", &body);
    v->Visit("span", &span);
  }

  bool SEqualReduce(const AttrStmtNode* other, SEqualReducer equal) const {
    return equal(node, other->node) && equal(attr_key, other->attr_key) &&
           equal(value, other->value) && equal(body, other->body);
  }

  void SHashReduce(SHashReducer hash_reduce) const {
    hash_reduce(node);
    hash_reduce(attr_key);
    hash_reduce(value);
    hash_reduce(body);
  }

  static constexpr const char* _type_key = "tir.AttrStmt";
  TVM_DECLARE_FINAL_OBJECT_INFO(AttrStmtNode, StmtNode);
};

class AttrStmt : public Stmt {
 public:
  TVM_DLL AttrStmt(ObjectRef node, String attr_key, PrimExpr value, Stmt body, Span span = Span());

  TVM_DEFINE_OBJECT_REF_METHODS(AttrStmt, Stmt, AttrStmtNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(AttrStmtNode);
};

class AssertStmtNode : public StmtNode {
 public:
  PrimExpr condition;
  PrimExpr message;
  Stmt body;

  void VisitAttrs(AttrVisitor* v) {
    v->Visit("condition", &condition);
    v->Visit("message", &message);
    v->Visit("body", &body);
    v->Visit("span", &span);
  }

  bool SEqualReduce(const AssertStmtNode* other, SEqualReducer equal) const {
    return equal(condition, other->condition) && equal(message, other->message) &&
           equal(body, other->body);
  }

  void SHashReduce(SHashReducer hash_reduce) const {
    hash_reduce(condition);
    hash_reduce(message);
    hash_reduce(body);
  }

  static constexpr const char* _type_key = "tir.AssertStmt";
  TVM_DECLARE_FINAL_OBJECT_INFO(AssertStmtNode, StmtNode);
};

class AssertStmt : public Stmt {
 public:
  TVM_DLL AssertStmt(PrimExpr condition, PrimExpr message, Stmt body, Span span = Span());

  TVM_DEFINE_OBJECT_REF_METHODS(AssertStmt, Stmt, AssertStmtNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(AssertStmtNode);
};

class StoreNode : public StmtNode {
 public:
  Var buffer_var;
  PrimExpr value;
  PrimExpr index;
  PrimExpr predicate;

  void VisitAttrs(AttrVisitor* v) {
    v->Visit("buffer_var", &buffer_var);
    v->Visit("value", &value);
    v->Visit("index", &index);
    v->Visit("predicate", &predicate);
    v->Visit("span", &span);
  }

  bool SEqualReduce(const StoreNode* other, SEqualReducer equal) const {
    return equal(buffer_var, other->buffer_var) && equal(value, other->value) &&
           equal(index, other->index) && equal(predicate, other->predicate);
  }

  void SHashReduce(SHashReducer hash_reduce) const {
    hash_reduce(buffer_var);
    hash_reduce(value);
    hash_reduce(index);
    hash_reduce(predicate);
  }

  static constexpr const char* _type_key = "tir.Store";
  TVM_DECLARE_FINAL_OBJECT_INFO(StoreNode, StmtNode);
};

class Store : public Stmt {
 public:
  TVM_DLL Store(Var buffer_var, PrimExpr value, PrimExpr index, PrimExpr predicate,
                Span span = Span());

  TVM_DEFINE_OBJECT_REF_METHODS(Store, Stmt, StoreNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(StoreNode);
};

class BufferStoreNode : public StmtNode {
 public:
  Buffer buffer;
  PrimExpr value;
  Array<PrimExpr> indices;

  void VisitAttrs(AttrVisitor* v) {
    v->Visit("buffer", &buffer);
    v->Visit("value", &value);
    v->Visit("indices", &indices);
    v->Visit("span", &span);
  }

  bool SEqualReduce(const BufferStoreNode* other, SEqualReducer equal) const {
    return equal(buffer, other->buffer) && equal(value, other->value) &&
           equal(indices, other->indices);
  }

  void SHashReduce(SHashReducer hash_reduce) const {
    hash_reduce(buffer);
    hash_reduce(value);
    hash_reduce(indices);
  }

  static constexpr const char* _type_key = "tir.BufferStore";
  TVM_DECLARE_FINAL_OBJECT_INFO(BufferStoreNode, StmtNode);
};

class BufferStore : public Stmt {
 public:
  TVM_DLL explicit BufferStore(Buffer buffer, PrimExpr value, Array<PrimExpr> indices,
                               Span span = Span());

  TVM_DEFINE_OBJECT_REF_METHODS(BufferStore, Stmt, BufferStoreNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(BufferStoreNode);
};

class BufferRealizeNode : public StmtNode {
 public:
  Buffer buffer;
  Array<Range> bounds;
  PrimExpr condition;
  Stmt body;

  void VisitAttrs(AttrVisitor* v) {
    v->Visit("buffer", &buffer);
    v->Visit("bounds", &bounds);
    v->Visit("condition", &condition);
    v->Visit("body", &body);
    v->Visit("span", &span);
  }

  bool SEqualReduce(const BufferRealizeNode* other, SEqualReducer equal) const {
    return equal(buffer, other->buffer) && equal(bounds, other->bounds) &&
           equal(condition, other->condition) && equal(body, other->body);
  }

  void SHashReduce(SHashReducer hash_reduce) const {
    hash_reduce(buffer);
    hash_reduce(bounds);
    hash_reduce(condition);
    hash_reduce(body);
  }

  BufferRealizeNode() = default;
  BufferRealizeNode(Buffer buffer, Array<Range> bounds, PrimExpr condition, Stmt body,
                    Span span = Span())
      : StmtNode(span), buffer(buffer), bounds(bounds), condition(condition), body(body) {}

  static constexpr const char* _type_key = "tir.BufferRealize";
  TVM_DECLARE_FINAL_OBJECT_INFO(BufferRealizeNode, StmtNode);
};

class BufferRealize : public Stmt {
 public:
  TVM_DLL explicit BufferRealize(Buffer buffer, Array<Range> bounds, PrimExpr condition, Stmt body,
                                 Span span = Span());

  TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(BufferRealize, Stmt, BufferRealizeNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(BufferRealizeNode);
};

class ProducerStoreNode : public StmtNode {
 public:
  DataProducer producer;
  PrimExpr value;
  Array<PrimExpr> indices;

  void VisitAttrs(AttrVisitor* v) {
    v->Visit("producer", &producer);
    v->Visit("value", &value);
    v->Visit("indices", &indices);
    v->Visit("span", &span);
  }

  bool SEqualReduce(const ProducerStoreNode* other, SEqualReducer equal) const {
    return equal(producer, other->producer) && equal(value, other->value) &&
           equal(indices, other->indices);
  }

  void SHashReduce(SHashReducer hash_reduce) const {
    hash_reduce(producer);
    hash_reduce(value);
    hash_reduce(indices);
  }

  static constexpr const char* _type_key = "tir.ProducerStore";
  TVM_DECLARE_FINAL_OBJECT_INFO(ProducerStoreNode, StmtNode);
};

class ProducerStore : public Stmt {
 public:
  TVM_DLL ProducerStore(DataProducer producer, PrimExpr value, Array<PrimExpr> indices,
                        Span span = Span());

  TVM_DEFINE_OBJECT_REF_METHODS(ProducerStore, Stmt, ProducerStoreNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(ProducerStoreNode);
};

class ProducerRealizeNode : public StmtNode {
 public:
  DataProducer producer;
  Region bounds;
  PrimExpr condition;
  Stmt body;
  String storage_scope;

  void VisitAttrs(AttrVisitor* v) {
    v->Visit("producer", &producer);
    v->Visit("bounds", &bounds);
    v->Visit("condition", &condition);
    v->Visit("body", &body);
    v->Visit("storage_scope", &storage_scope);
    v->Visit("span", &span);
  }

  bool SEqualReduce(const ProducerRealizeNode* other, SEqualReducer equal) const {
    return equal(producer, other->producer) && equal(bounds, other->bounds) &&
           equal(condition, other->condition) && equal(body, other->body) &&
           equal(storage_scope, other->storage_scope);
  }

  void SHashReduce(SHashReducer hash_reduce) const {
    hash_reduce(producer);
    hash_reduce(bounds);
    hash_reduce(condition);
    hash_reduce(body);
    hash_reduce(storage_scope);
  }

  static constexpr const char* _type_key = "tir.ProducerRealize";
  TVM_DECLARE_FINAL_OBJECT_INFO(ProducerRealizeNode, StmtNode);
};

class ProducerRealize : public Stmt {
 public:
  TVM_DLL ProducerRealize(DataProducer producer, Region bounds, PrimExpr condition, Stmt body,
                          String storage_scope = "", Span span = Span());

  TVM_DEFINE_OBJECT_REF_METHODS(ProducerRealize, Stmt, ProducerRealizeNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(ProducerRealizeNode);
};

class AllocateNode : public StmtNode {
 public:
  Var buffer_var;
  DataType dtype;
  Array<PrimExpr> extents;
  PrimExpr condition;
  Stmt body;
  Map<String, ObjectRef> annotations;

  void VisitAttrs(AttrVisitor* v) {
    v->Visit("buffer_var", &buffer_var);
    v->Visit("dtype", &dtype);
    v->Visit("extents", &extents);
    v->Visit("condition", &condition);
    v->Visit("body", &body);
    v->Visit("annotations", &annotations);
    v->Visit("span", &span);
  }

  bool SEqualReduce(const AllocateNode* other, SEqualReducer equal) const {
    return equal.DefEqual(buffer_var, other->buffer_var) && equal(dtype, other->dtype) &&
           equal(extents, other->extents) && equal(condition, other->condition) &&
           equal(body, other->body) && equal(annotations, other->annotations);
  }

  void SHashReduce(SHashReducer hash_reduce) const {
    hash_reduce.DefHash(buffer_var);
    hash_reduce(dtype);
    hash_reduce(extents);
    hash_reduce(condition);
    hash_reduce(body);
    hash_reduce(annotations);
  }

  int64_t ConstantAllocationSize() const { return ConstantAllocationSize(extents); }
  TVM_DLL static int64_t ConstantAllocationSize(const Array<PrimExpr>& extents);

  static constexpr const char* _type_key = "tir.Allocate";
  static constexpr const bool _type_has_method_sequal_reduce = true;
  static constexpr const bool _type_has_method_shash_reduce = true;
  TVM_DECLARE_FINAL_OBJECT_INFO(AllocateNode, StmtNode);
};

class Allocate : public Stmt {
 public:
  TVM_DLL Allocate(Var buffer_var, DataType dtype, Array<PrimExpr> extents, PrimExpr condition,
                   Stmt body, Map<String, ObjectRef> annotations = Map<String, ObjectRef>(),
                   Span span = Span());

  TVM_DEFINE_OBJECT_REF_METHODS(Allocate, Stmt, AllocateNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(AllocateNode);
};

class AllocateConstNode : public StmtNode {
 public:
  Var buffer_var;
  Optional<runtime::NDArray> data;
  Optional<Integer> irmod_storage_idx;
  DataType dtype;
  Array<PrimExpr> extents;
  Stmt body;
  Map<String, ObjectRef> annotations;

  void VisitAttrs(AttrVisitor* v) {
    v->Visit("buffer_var", &buffer_var);
    v->Visit("data", &data);
    v->Visit("irmod_storage_idx", &irmod_storage_idx);
    v->Visit("dtype", &dtype);
    v->Visit("extents", &extents);
    v->Visit("body", &body);
    v->Visit("annotations", &annotations);
    v->Visit("span", &span);
  }

  bool SEqualReduce(const AllocateConstNode* other, SEqualReducer equal) const {
    return equal.DefEqual(buffer_var, other->buffer_var) && equal(dtype, other->dtype) &&
           equal(extents, other->extents) && equal(data, other->data) && equal(body, other->body) &&
           equal(annotations, other->annotations);
  }

  void SHashReduce(SHashReducer hash_reduce) const {
    hash_reduce.DefHash(buffer_var);
    hash_reduce(dtype);
    hash_reduce(extents);
    hash_reduce(body);
    hash_reduce(annotations);
    hash_reduce(data);
  }

  int64_t ConstantAllocationSize() const { return ConstantAllocationSize(extents); }
  TVM_DLL static int64_t ConstantAllocationSize(const Array<PrimExpr>& extents);

  static constexpr const char* _type_key = "tir.AllocateConst";
  static constexpr const bool _type_has_method_sequal_reduce = true;
  static constexpr const bool _type_has_method_shash_reduce = true;
  TVM_DECLARE_FINAL_OBJECT_INFO(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, Array<PrimExpr> extents,
                        ObjectRef data_or_idx, Stmt body,
                        Map<String, ObjectRef> annotations = Map<String, ObjectRef>(),
                        Span span = Span());
  TVM_DEFINE_OBJECT_REF_METHODS(AllocateConst, Stmt, AllocateConstNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(AllocateConstNode);
};

class DeclBufferNode : public StmtNode {
 public:
  Buffer buffer;
  Stmt body;

  void VisitAttrs(AttrVisitor* v) {
    v->Visit("buffer", &buffer);
    v->Visit("body", &body);
    v->Visit("span", &span);
  }

  bool SEqualReduce(const DeclBufferNode* other, SEqualReducer equal) const {
    return equal(buffer, other->buffer) && equal(body, other->body);
  }

  void SHashReduce(SHashReducer hash_reduce) const {
    hash_reduce(buffer);
    hash_reduce(body);
  }

  static constexpr const char* _type_key = "tir.DeclBuffer";
  TVM_DECLARE_FINAL_OBJECT_INFO(DeclBufferNode, StmtNode);
};

class DeclBuffer : public Stmt {
 public:
  TVM_DLL DeclBuffer(Buffer buffer, Stmt body, Span span = Span());
  TVM_DEFINE_OBJECT_REF_METHODS(DeclBuffer, Stmt, DeclBufferNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(DeclBufferNode);
};

class SeqStmtNode : public StmtNode {
 public:
  Array<Stmt> seq;

  size_t size() const { return seq.size(); }
  Stmt operator[](size_t index) const { return seq[index]; }

  void VisitAttrs(AttrVisitor* v) {
    v->Visit("seq", &seq);
    v->Visit("span", &span);
  }

  bool SEqualReduce(const SeqStmtNode* other, SEqualReducer equal) const {
    return equal(seq, other->seq);
  }

  void SHashReduce(SHashReducer hash_reduce) const { hash_reduce(seq); }

  static constexpr const char* _type_key = "tir.SeqStmt";
  TVM_DECLARE_FINAL_OBJECT_INFO(SeqStmtNode, StmtNode);
};

class SeqStmt : public Stmt {
 public:
  TVM_DLL explicit SeqStmt(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) {
    Array<Stmt> seq;
    runtime::detail::for_each(Flattener(&seq), std::forward<Args>(seq_args)...);
    if (seq.size() == 1) return seq[0];
    return SeqStmt(seq);
  }
  class Flattener {
   public:
    explicit Flattener(Array<Stmt>* seq) : seq_(seq) {}

    void operator()(size_t i, const Stmt& stmt) const {
      if (!stmt.defined()) return;
      if (auto* op = stmt.as<SeqStmtNode>()) {
        operator()(0, op->seq);
      } else {
        seq_->push_back(stmt);
      }
    }

    template <typename T>
    void operator()(size_t i, const T& seq) const {
      for (auto v : seq) {
        this->operator()(0, v);
      }
    }

   private:
    Array<Stmt>* seq_;
  };

  TVM_DEFINE_OBJECT_REF_METHODS(SeqStmt, Stmt, SeqStmtNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(SeqStmtNode);
};

class IfThenElseNode : public StmtNode {
 public:
  PrimExpr condition;
  Stmt then_case;
  Optional<Stmt> else_case;

  void VisitAttrs(AttrVisitor* v) {
    v->Visit("condition", &condition);
    v->Visit("then_case", &then_case);
    v->Visit("else_case", &else_case);
    v->Visit("span", &span);
  }

  bool SEqualReduce(const IfThenElseNode* other, SEqualReducer equal) const {
    return equal(condition, other->condition) && equal(then_case, other->then_case) &&
           equal(else_case, other->else_case);
  }

  void SHashReduce(SHashReducer hash_reduce) const {
    hash_reduce(condition);
    hash_reduce(then_case);
    hash_reduce(else_case);
  }

  static constexpr const char* _type_key = "tir.IfThenElse";
  TVM_DECLARE_FINAL_OBJECT_INFO(IfThenElseNode, StmtNode);
};

class IfThenElse : public Stmt {
 public:
  TVM_DLL IfThenElse(PrimExpr condition, Stmt then_case, Optional<Stmt> else_case = NullOpt,
                     Span span = Span());

  TVM_DEFINE_OBJECT_REF_METHODS(IfThenElse, Stmt, IfThenElseNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(IfThenElseNode);
};

class EvaluateNode : public StmtNode {
 public:
  PrimExpr value;

  void VisitAttrs(AttrVisitor* v) {
    v->Visit("value", &value);
    v->Visit("span", &span);
  }

  bool SEqualReduce(const EvaluateNode* other, SEqualReducer equal) const {
    return equal(value, other->value);
  }

  void SHashReduce(SHashReducer hash_reduce) const { hash_reduce(value); }

  static constexpr const char* _type_key = "tir.Evaluate";
  TVM_DECLARE_FINAL_OBJECT_INFO(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_DEFINE_OBJECT_REF_METHODS(Evaluate, Stmt, EvaluateNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(EvaluateNode);
};

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;
  Optional<IterVar> thread_binding;
  Map<String, ObjectRef> annotations;

  void VisitAttrs(AttrVisitor* v) {
    v->Visit("loop_var", &loop_var);
    v->Visit("min", &min);
    v->Visit("extent", &extent);
    v->Visit("kind", &kind);
    v->Visit("body", &body);
    v->Visit("thread_binding", &thread_binding);
    v->Visit("annotations", &annotations);
    v->Visit("span", &span);
  }

  bool SEqualReduce(const ForNode* other, SEqualReducer equal) const {
    return equal.DefEqual(loop_var, other->loop_var) && equal(min, other->min) &&
           equal(extent, other->extent) && equal(kind, other->kind) && equal(body, other->body) &&
           equal(thread_binding, other->thread_binding) && equal(annotations, other->annotations);
  }

  void SHashReduce(SHashReducer hash_reduce) const {
    hash_reduce.DefHash(loop_var);
    hash_reduce(min);
    hash_reduce(extent);
    hash_reduce(kind);
    hash_reduce(body);
    hash_reduce(thread_binding);
    hash_reduce(annotations);
  }

  static constexpr const char* _type_key = "tir.For";
  TVM_DECLARE_FINAL_OBJECT_INFO(ForNode, StmtNode);
};

class For : public Stmt {
 public:
  TVM_DLL For(Var loop_var, PrimExpr min, PrimExpr extent, ForKind kind, Stmt body,
              Optional<IterVar> thread_binding = NullOpt,
              Map<String, ObjectRef> annotations = Map<String, ObjectRef>(), Span span = Span());

  TVM_DEFINE_OBJECT_REF_METHODS(For, Stmt, ForNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(ForNode);
};

class WhileNode : public StmtNode {
 public:
  PrimExpr condition;
  Stmt body;

  void VisitAttrs(AttrVisitor* v) {
    v->Visit("condition", &condition);
    v->Visit("body", &body);
    v->Visit("span", &span);
  }

  bool SEqualReduce(const WhileNode* other, SEqualReducer equal) const {
    return equal(condition, other->condition) && equal(body, other->body);
  }

  void SHashReduce(SHashReducer hash_reduce) const {
    hash_reduce(condition);
    hash_reduce(body);
  }

  static constexpr const char* _type_key = "tir.While";
  TVM_DECLARE_FINAL_OBJECT_INFO(WhileNode, StmtNode);
};

class While : public Stmt {
 public:
  TVM_DLL While(PrimExpr condition, Stmt body, Span span = Span());

  TVM_DEFINE_OBJECT_REF_METHODS(While, Stmt, WhileNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(WhileNode);
};

class PrefetchNode : public StmtNode {
 public:
  Buffer buffer;
  Array<Range> bounds;

  void VisitAttrs(AttrVisitor* v) {
    v->Visit("buffer", &buffer);
    v->Visit("bounds", &bounds);
    v->Visit("span", &span);
  }

  bool SEqualReduce(const PrefetchNode* other, SEqualReducer equal) const {
    return equal(buffer, other->buffer) && equal(bounds, other->bounds);
  }

  void SHashReduce(SHashReducer hash_reduce) const {
    hash_reduce(buffer);
    hash_reduce(bounds);
  }

  PrefetchNode() = default;
  PrefetchNode(Buffer buffer, Array<Range> bounds, Span span = Span())
      : StmtNode(span), buffer(buffer), bounds(bounds) {}

  static constexpr const char* _type_key = "tir.Prefetch";
  TVM_DECLARE_FINAL_OBJECT_INFO(PrefetchNode, StmtNode);
};

class Prefetch : public Stmt {
 public:
  TVM_DLL explicit Prefetch(Buffer buffer, Array<Range> bounds, Span span = Span());

  TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(Prefetch, Stmt, PrefetchNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(PrefetchNode);
};

class BufferRegionNode : public Object {
 public:
  Buffer buffer;
  Array<Range> region;

  void VisitAttrs(AttrVisitor* v) {
    v->Visit("buffer", &buffer);
    v->Visit("region", &region);
  }

  bool SEqualReduce(const BufferRegionNode* other, SEqualReducer equal) const {
    return equal(buffer, other->buffer) && equal(region, other->region);
  }

  void SHashReduce(SHashReducer hash_reduce) const {
    hash_reduce(buffer);
    hash_reduce(region);
  }

  static constexpr const char* _type_key = "tir.BufferRegion";
  static constexpr const bool _type_has_method_sequal_reduce = true;
  static constexpr const bool _type_has_method_shash_reduce = true;
  TVM_DECLARE_FINAL_OBJECT_INFO(BufferRegionNode, Object);
};

class BufferRegion : public ObjectRef {
 public:
  TVM_DLL explicit BufferRegion(Buffer buffer, Array<Range> region);

  TVM_DLL static BufferRegion FullRegion(Buffer buffer);

  TVM_DLL static BufferRegion FromPoint(Buffer buffer, Array<PrimExpr> indices);

  TVM_DEFINE_OBJECT_REF_METHODS(BufferRegion, ObjectRef, BufferRegionNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(BufferRegionNode);
};

class MatchBufferRegionNode : public Object {
 public:
  Buffer buffer;
  BufferRegion source;

  void VisitAttrs(AttrVisitor* v) {
    v->Visit("buffer", &buffer);
    v->Visit("source", &source);
  }

  bool SEqualReduce(const MatchBufferRegionNode* other, SEqualReducer equal) const {
    return equal(buffer, other->buffer) && equal(source, other->source);
  }

  void SHashReduce(SHashReducer hash_reduce) const {
    hash_reduce(buffer);
    hash_reduce(source);
  }

  static constexpr const char* _type_key = "tir.MatchBufferRegion";
  static constexpr const bool _type_has_method_sequal_reduce = true;
  static constexpr const bool _type_has_method_shash_reduce = true;
  TVM_DECLARE_FINAL_OBJECT_INFO(MatchBufferRegionNode, Object);
};

class MatchBufferRegion : public ObjectRef {
 public:
  TVM_DLL explicit MatchBufferRegion(Buffer buffer, BufferRegion source);

  TVM_DEFINE_OBJECT_REF_METHODS(MatchBufferRegion, ObjectRef, MatchBufferRegionNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(MatchBufferRegionNode);
};

class BlockNode : public StmtNode {
 public:
  Array<IterVar> iter_vars;
  Array<BufferRegion> reads;
  Array<BufferRegion> writes;
  String name_hint;
  Stmt body;
  Optional<Stmt> init;
  Array<Buffer> alloc_buffers;
  Array<MatchBufferRegion> match_buffers;
  Map<String, ObjectRef> annotations;

  void VisitAttrs(AttrVisitor* v) {
    v->Visit("iter_vars", &iter_vars);
    v->Visit("reads", &reads);
    v->Visit("writes", &writes);
    v->Visit("name_hint", &name_hint);
    v->Visit("body", &body);
    v->Visit("init", &init);
    v->Visit("alloc_buffers", &alloc_buffers);
    v->Visit("match_buffers", &match_buffers);
    v->Visit("annotations", &annotations);
  }

  bool SEqualReduce(const BlockNode* other, SEqualReducer equal) const {
    // Need first reduce iter_vars, alloc_buffers and match_buffers to define new vars
    return equal.DefEqual(iter_vars, other->iter_vars) &&
           equal(alloc_buffers, other->alloc_buffers) &&
           equal(match_buffers, other->match_buffers) && equal(reads, other->reads) &&
           equal(writes, other->writes) && equal(body, other->body) && equal(init, other->init) &&
           equal(annotations, other->annotations);
  }

  void SHashReduce(SHashReducer hash_reduce) const {
    hash_reduce.DefHash(iter_vars);
    hash_reduce(alloc_buffers);
    hash_reduce(match_buffers);
    hash_reduce(reads);
    hash_reduce(writes);
    hash_reduce(body);
    hash_reduce(init);
    hash_reduce(annotations);
  }

  static constexpr const char* _type_key = "tir.Block";
  TVM_DECLARE_FINAL_OBJECT_INFO(BlockNode, StmtNode);
};

class Block : public Stmt {
 public:
  TVM_DLL explicit Block(Array<IterVar> iter_vars, Array<BufferRegion> reads,
                         Array<BufferRegion> writes, String name_hint, Stmt body,
                         Optional<Stmt> init = NullOpt,
                         Array<Buffer> alloc_buffers = Array<Buffer>(),
                         Array<MatchBufferRegion> match_buffers = Array<MatchBufferRegion>(),
                         Map<String, ObjectRef> annotations = Map<String, ObjectRef>(),
                         Span span = Span());

  TVM_DEFINE_OBJECT_REF_METHODS(Block, Stmt, BlockNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(BlockNode);
};

class BlockRealizeNode : public StmtNode {
 public:
  Array<PrimExpr> iter_values;
  PrimExpr predicate;
  Block block;

  void VisitAttrs(AttrVisitor* v) {
    v->Visit("iter_values", &iter_values);
    v->Visit("predicate", &predicate);
    v->Visit("block", &block);
  }

  bool SEqualReduce(const BlockRealizeNode* other, SEqualReducer equal) const {
    return equal(iter_values, other->iter_values) && equal(predicate, other->predicate) &&
           equal(block, other->block);
  }

  void SHashReduce(SHashReducer hash_reduce) const {
    hash_reduce(iter_values);
    hash_reduce(predicate);
    hash_reduce(block);
  }

  static constexpr const char* _type_key = "tir.BlockRealize";
  TVM_DECLARE_FINAL_OBJECT_INFO(BlockRealizeNode, StmtNode);
};

class BlockRealize : public Stmt {
 public:
  TVM_DLL explicit BlockRealize(Array<PrimExpr> iter_values, PrimExpr predicate, Block block,
                                Span span = Span());

  TVM_DEFINE_OBJECT_REF_METHODS(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* prefetch_scope = "prefetch_scope";
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* warp_execution = "warp_execution";

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_