Program Listing for File stmt.h

↰ Return to documentation for file (tvm/ir/stmt.h)

/*
 * 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_IR_STMT_H_
#define TVM_IR_STMT_H_

#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/op.h>
#include <tvm/ir/prim/expr.h>

#include <initializer_list>
#include <optional>
#include <string>
#include <utility>

namespace tvm {

class StmtObj : public ffi::Object {
 public:
  mutable Location loc = UnknownLoc();

  StmtObj() = default;
  explicit StmtObj(Location loc) : loc(loc) {}

  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<StmtObj>().def_ro("loc", &StmtObj::loc, refl::AttachFieldFlag::SEqHashIgnore());
  }

  static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode;
  static constexpr bool _type_s_eq_hash_subclass_kind_fixed = true;

  static constexpr const uint32_t _type_child_slots = 15;
  TVM_FFI_DECLARE_OBJECT_INFO("ir.Stmt", StmtObj, ffi::Object);
};

class Stmt : public ffi::ObjectRef {
 public:
  explicit Stmt(ffi::ObjectPtr<StmtObj> node) : ffi::ObjectRef(std::move(node)) {
    TVM_FFI_CHECK(defined(), ValueError) << "Stmt expects a defined node";
  }

  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(Stmt, ffi::ObjectRef, StmtObj);
};

class SeqStmtObj : public StmtObj {
 public:
  ffi::Array<Stmt> seq;

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

  TVM_TIRX_DLL static void RegisterReflection();
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ir.SeqStmt", SeqStmtObj, StmtObj);
};

class SeqStmt : public Stmt {
 public:
  TVM_TIRX_DLL explicit SeqStmt(ffi::Array<Stmt> seq, Location loc = UnknownLoc());
  TVM_TIRX_DLL SeqStmt(Stmt stmt, ffi::Optional<Location> loc = std::nullopt);
  SeqStmt(std::initializer_list<Stmt> seq, Location loc = UnknownLoc())
      : SeqStmt(ffi::Array<Stmt>(seq), std::move(loc)) {}

  size_t size() const { return operator->()->size(); }
  Stmt operator[](size_t index) const { return (*(operator->()))[index]; }
  explicit SeqStmt(ffi::ObjectPtr<SeqStmtObj> node) : Stmt(std::move(node)) {}

  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(SeqStmt, Stmt, SeqStmtObj);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(SeqStmtObj);
};

namespace ffi {
template <>
inline constexpr bool use_default_type_traits_v<SeqStmt> = false;

template <>
struct TypeTraits<SeqStmt> : public ObjectRefWithFallbackTraitsBase<SeqStmt, Stmt, Array<Stmt>> {
  static SeqStmt ConvertFallbackValue(Stmt stmt) { return SeqStmt(std::move(stmt)); }
  static SeqStmt ConvertFallbackValue(Array<Stmt> seq) { return SeqStmt(std::move(seq)); }
};
}  // namespace ffi

class BindObj : public StmtObj {
 public:
  explicit BindObj(ffi::UnsafeInit tag) : var(tag), value(tag) {}

  BindObj(Var var, Expr value) : var(std::move(var)), value(std::move(value)) {}

  Var var;
  Expr value;

  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<BindObj>()
        .def_ro("var", &BindObj::var, refl::AttachFieldFlag::SEqHashDefSimple())
        .def_ro("value", &BindObj::value);
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ir.Bind", BindObj, StmtObj);
};

class Bind : public Stmt {
 public:
  TVM_TIRX_DLL Bind(Var var, Expr value, Location loc = UnknownLoc());

  explicit Bind(ffi::ObjectPtr<BindObj> node) : Stmt(std::move(node)) {}

  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(Bind, Stmt, BindObj);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(BindObj);
};

TVM_TIRX_DLL bool IsRegionOp(const Op& op);

TVM_TIRX_DLL ffi::Array<Var> GetRegionBodyParams(Op op, ffi::Array<Expr> args, DictAttrs attrs);

class RegionStmtObj : public StmtObj {
 public:
  Op op;
  ffi::Array<Expr> args;
  ffi::Array<Var> body_params;
  DictAttrs attrs;
  SeqStmt body;
  ffi::Array<Var> result_vars;

  explicit RegionStmtObj(ffi::UnsafeInit tag) : op(tag), body(tag) {}

  RegionStmtObj(Op op, SeqStmt body) : op(std::move(op)), body(std::move(body)) {}

  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<RegionStmtObj>()
        .def_ro("op", &RegionStmtObj::op)
        .def_ro("attrs", &RegionStmtObj::attrs)
        .def_ro("args", &RegionStmtObj::args)
        .def_ro("body_params", &RegionStmtObj::body_params,
                refl::AttachFieldFlag::SEqHashDefSimple())
        .def_ro("body", &RegionStmtObj::body)
        .def_ro("result_vars", &RegionStmtObj::result_vars,
                refl::AttachFieldFlag::SEqHashDefSimple());
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ir.RegionStmt", RegionStmtObj, StmtObj);
};

class RegionStmt : public Stmt {
 public:
  TVM_TIRX_DLL RegionStmt(Op op, ffi::Array<Expr> args, ffi::Array<Var> body_params,
                          DictAttrs attrs, SeqStmt body, ffi::Array<Var> result_vars = {},
                          Location loc = UnknownLoc());

  explicit RegionStmt(ffi::ObjectPtr<RegionStmtObj> node) : Stmt(std::move(node)) {}

  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(RegionStmt, Stmt, RegionStmtObj);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(RegionStmtObj);
};

namespace op_attr {
inline constexpr ffi::reflection::AttrKey<
    ffi::reflection::NativeFunctionView<ffi::Array<Var>(const CallObj* call)>>
    kRegionGetBodyParams{"FRegionGetBodyParams"};
inline constexpr ffi::reflection::AttrKey<
    ffi::reflection::NativeFunctionView<void(const RegionStmtObj*)>>
    kRegionValidate{"FRegionValidate"};
}  // namespace op_attr

// Compatibility names; new callers can deduce the value type from the key.
using FRegionGetBodyParams = decltype(op_attr::kRegionGetBodyParams)::ValueType;
using FRegionValidate = decltype(op_attr::kRegionValidate)::ValueType;

class AssertStmtObj : public StmtObj {
 public:
  explicit AssertStmtObj(ffi::UnsafeInit tag) : condition(tag), error_kind(tag) {}

  AssertStmtObj(PrimExpr condition, StringImm error_kind)
      : condition(std::move(condition)), error_kind(std::move(error_kind)) {}

  PrimExpr condition;
  StringImm error_kind;
  ffi::Array<StringImm> message_parts;

  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<AssertStmtObj>()
        .def_ro("condition", &AssertStmtObj::condition)
        .def_ro("error_kind", &AssertStmtObj::error_kind)
        .def_ro("message_parts", &AssertStmtObj::message_parts);
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ir.AssertStmt", AssertStmtObj, StmtObj);
};

class AssertStmt : public Stmt {
 public:
  TVM_TIRX_DLL AssertStmt(PrimExpr condition, StringImm error_kind,
                          ffi::Array<StringImm> message_parts, Location loc = UnknownLoc());

  explicit AssertStmt(ffi::ObjectPtr<AssertStmtObj> node) : Stmt(std::move(node)) {}

  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(AssertStmt, Stmt, AssertStmtObj);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(AssertStmtObj);
};

class TensorStoreObj : public StmtObj {
 public:
  explicit TensorStoreObj(ffi::UnsafeInit tag) : dest(tag), value(tag) {}

  TensorStoreObj(Expr dest, PrimExpr value) : dest(std::move(dest)), value(std::move(value)) {}

  Expr dest;
  ffi::Array<PrimExpr> indices;
  PrimExpr value;

  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<TensorStoreObj>()
        .def_ro("dest", &TensorStoreObj::dest)
        .def_ro("indices", &TensorStoreObj::indices)
        .def_ro("value", &TensorStoreObj::value);
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ir.TensorStore", TensorStoreObj, StmtObj);
};

class TensorStore : public Stmt {
 public:
  TVM_TIRX_DLL explicit TensorStore(Expr dest, ffi::Array<PrimExpr> indices, PrimExpr value,
                                    Location loc = UnknownLoc());

  explicit TensorStore(ffi::ObjectPtr<TensorStoreObj> node) : Stmt(std::move(node)) {}

  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(TensorStore, Stmt, TensorStoreObj);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(TensorStoreObj);
};

namespace type_attr {
inline constexpr ffi::reflection::AttrKey<ffi::reflection::NativeFunctionView<void(
    const Expr&, const ffi::Array<PrimExpr>&, const PrimExpr&)>>
    kTensorStoreValidateByTy{"__tvm_tensor_store_validate__"};
}  // namespace type_attr

class EvaluateObj : public StmtObj {
 public:
  explicit EvaluateObj(ffi::UnsafeInit tag) : value(tag) {}

  explicit EvaluateObj(Expr value) : value(std::move(value)) {}

  Expr value;

  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<EvaluateObj>().def_ro("value", &EvaluateObj::value);
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ir.Evaluate", EvaluateObj, StmtObj);
};

class Evaluate : public Stmt {
 public:
  TVM_TIRX_DLL explicit Evaluate(Expr value, Location loc = UnknownLoc());

  explicit Evaluate(int value, Location loc = UnknownLoc()) : Evaluate(PrimExpr(value), loc) {}

  explicit Evaluate(ffi::ObjectPtr<EvaluateObj> node) : Stmt(std::move(node)) {}

  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(Evaluate, Stmt, EvaluateObj);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(EvaluateObj);
};

class IfObj : public StmtObj {
 public:
  explicit IfObj(ffi::UnsafeInit tag) : condition(tag), then_case(tag) {}

  IfObj(PrimExpr condition, SeqStmt then_case)
      : condition(std::move(condition)), then_case(std::move(then_case)) {}

  PrimExpr condition;
  SeqStmt then_case;
  ffi::Optional<SeqStmt> else_case;

  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<IfObj>()
        .def_ro("condition", &IfObj::condition)
        .def_ro("then_case", &IfObj::then_case)
        .def_ro("else_case", &IfObj::else_case);
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ir.If", IfObj, StmtObj);
};

class If : public Stmt {
 public:
  TVM_TIRX_DLL If(PrimExpr condition, SeqStmt then_case,
                  ffi::Optional<SeqStmt> else_case = std::nullopt, Location loc = UnknownLoc());

  explicit If(ffi::ObjectPtr<IfObj> node) : Stmt(std::move(node)) {}

  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(If, Stmt, IfObj);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(IfObj);
};

enum class ForKind : int {
  kDefault = 0,
  kParallel = 1,
  kVectorized = 2,
  kUnrolled = 3,
};

class ForObj : public StmtObj {
 public:
  explicit ForObj(ffi::UnsafeInit tag) : var(tag), begin(tag), end(tag), step(tag), body(tag) {}

  ForObj(PrimVar var, PrimExpr begin, PrimExpr end, PrimExpr step, ForKind kind, SeqStmt body)
      : var(std::move(var)),
        begin(std::move(begin)),
        end(std::move(end)),
        step(std::move(step)),
        kind(kind),
        body(std::move(body)) {}

  PrimVar var;
  PrimExpr begin;
  PrimExpr end;
  PrimExpr step;
  ForKind kind;
  SeqStmt body;
  ffi::Map<ffi::String, ffi::Any> annotations;

  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<ForObj>()
        .def_ro("var", &ForObj::var, refl::AttachFieldFlag::SEqHashDefSimple())
        .def_ro("begin", &ForObj::begin)
        .def_ro("end", &ForObj::end)
        // Reflection stores one typed default. Wider unit steps remain distinct.
        .def_ro("step", &ForObj::step, refl::DefaultValue(IntImm::Int32(1)))
        .def_ro("kind", &ForObj::kind, refl::DefaultValue(ForKind::kDefault))
        .def_ro("body", &ForObj::body)
        .def_ro("annotations", &ForObj::annotations,
                refl::DefaultValue(ffi::Map<ffi::String, ffi::Any>()));
  }

  TVM_TIRX_DLL bool HasTrivialStep() const;

  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ir.For", ForObj, StmtObj);
};

class For : public Stmt {
 public:
  TVM_TIRX_DLL For(PrimVar var, PrimExpr begin, PrimExpr end, PrimExpr step, ForKind kind,
                   SeqStmt body, ffi::Map<ffi::String, ffi::Any> annotations = {},
                   Location loc = UnknownLoc());

  explicit For(ffi::ObjectPtr<ForObj> node) : Stmt(std::move(node)) {}

  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(For, Stmt, ForObj);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(ForObj);
};

class WhileObj : public StmtObj {
 public:
  explicit WhileObj(ffi::UnsafeInit tag) : condition(tag), body(tag) {}

  WhileObj(PrimExpr condition, SeqStmt body)
      : condition(std::move(condition)), body(std::move(body)) {}

  PrimExpr condition;
  SeqStmt body;

  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<WhileObj>()
        .def_ro("condition", &WhileObj::condition)
        .def_ro("body", &WhileObj::body);
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ir.While", WhileObj, StmtObj);
};

class While : public Stmt {
 public:
  TVM_TIRX_DLL While(PrimExpr condition, SeqStmt body, Location loc = UnknownLoc());

  explicit While(ffi::ObjectPtr<WhileObj> node) : Stmt(std::move(node)) {}

  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(While, Stmt, WhileObj);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(WhileObj);
};

class ReturnObj : public StmtObj {
 public:
  explicit ReturnObj(ffi::UnsafeInit tag) : value(tag) {}

  explicit ReturnObj(Expr value) : value(std::move(value)) {}

  Expr value;

  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<ReturnObj>().def_ro("value", &ReturnObj::value);
  }

  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ir.Return", ReturnObj, StmtObj);
};

class Return : public Stmt {
 public:
  TVM_TIRX_DLL explicit Return(Expr value, Location loc = UnknownLoc());

  explicit Return(ffi::ObjectPtr<ReturnObj> node) : Stmt(std::move(node)) {}

  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(Return, Stmt, ReturnObj);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(ReturnObj);
};

class BreakObj : public StmtObj {
 public:
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<BreakObj>();
  }

  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ir.Break", BreakObj, StmtObj);
};

class Break : public Stmt {
 public:
  TVM_TIRX_DLL explicit Break(Location loc);

  explicit Break(ffi::ObjectPtr<BreakObj> node) : Stmt(std::move(node)) {}

  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(Break, Stmt, BreakObj);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(BreakObj);
};

class ContinueObj : public StmtObj {
 public:
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<ContinueObj>();
  }

  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ir.Continue", ContinueObj, StmtObj);
};

class Continue : public Stmt {
 public:
  TVM_TIRX_DLL explicit Continue(Location loc);

  explicit Continue(ffi::ObjectPtr<ContinueObj> node) : Stmt(std::move(node)) {}

  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(Continue, Stmt, ContinueObj);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(ContinueObj);
};

// overload printing of for type.
TVM_TIRX_DLL std::ostream& operator<<(std::ostream& os, ForKind kind);

// inline implementations
inline const char* ForKind2String(ForKind t) {
  switch (t) {
    case ForKind::kDefault:
      return "serial";
    case ForKind::kParallel:
      return "parallel";
    case ForKind::kVectorized:
      return "vectorized";
    case ForKind::kUnrolled:
      return "unroll";
  }
  TVM_FFI_THROW(InternalError) << "Unknown ForKind" << t;
  TVM_FFI_UNREACHABLE();
}

}  // namespace tvm
#endif  // TVM_IR_STMT_H_