expr.h

Go to the documentation of this file.

/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */
#ifndef TVM_RELAX_EXPR_H_
#define TVM_RELAX_EXPR_H_
 
#include <tvm/ffi/container/array.h>
#include <tvm/ffi/container/map.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/cow.h>
#include <tvm/ir/expr.h>
#include <tvm/ir/function.h>
#include <tvm/ir/location.h>
#include <tvm/ir/prim/expr.h>
#include <tvm/ir/prim/op.h>
#include <tvm/relax/type.h>
#include <tvm/runtime/tensor.h>
 
#include <functional>
 
namespace tvm {
namespace relax {
 
// Compatibility aliases. Tuple expressions are defined in the common IR.
using ::tvm::Tuple;
using ::tvm::TupleGetItem;
using ::tvm::TupleGetItemNode;
using ::tvm::TupleNode;
 
class ShapeExprNode : public ExprNode {
 public:
  ffi::Array<PrimExpr> values;
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<ShapeExprNode>().def_ro("values", &ShapeExprNode::values);
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.expr.ShapeExpr", ShapeExprNode, ExprNode);
};
 
class ShapeExpr : public Expr {
 public:
  TVM_DLL explicit ShapeExpr(ffi::Array<PrimExpr> values, Location loc = UnknownLoc());
  explicit ShapeExpr(ffi::ObjectPtr<ShapeExprNode> node) : Expr(std::move(node)) {}
 
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(ShapeExpr, Expr, ShapeExprNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(ShapeExprNode);
};
 
class DataflowVarNode : public VarNode {
 public:
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<DataflowVarNode>();
  }
 
  static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindFreeVar;
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.expr.DataflowVar", DataflowVarNode, VarNode);
};
 
class DataflowVar : public Var {
 public:
  TVM_DLL explicit DataflowVar(ffi::String name, ffi::Optional<Type> ty_annotation,
                               Location loc = UnknownLoc());
 
  explicit DataflowVar(ffi::ObjectPtr<DataflowVarNode> node) : Var(std::move(node)) {}
 
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(DataflowVar, Var, DataflowVarNode);
};
 
TVM_DLL GenericConst MakeTensorConst(ffi::Tensor data,
                                     ffi::Optional<Type> ty_annotation = std::nullopt,
                                     Location loc = UnknownLoc());
 
class BindingNode : public ffi::Object {
 public:
  explicit BindingNode(Var var) : var(std::move(var)) {}
  explicit BindingNode(ffi::UnsafeInit) : var(ffi::UnsafeInit{}) {}
 
  mutable Location loc = UnknownLoc();
  Var var;
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<BindingNode>()
        .def_ro("loc", &BindingNode::loc, refl::AttachFieldFlag::SEqHashIgnore())
        // TODO(tqchen): use SEqHashDefSimple after the next pypi tvm-ffi release
        .def_ro("var", &BindingNode::var, refl::AttachFieldFlag::SEqHashDefPattern());
  }
 
  static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode;
 
  TVM_FFI_DECLARE_OBJECT_INFO("relax.expr.Binding", BindingNode, ffi::Object);
};
 
class Binding : public ffi::ObjectRef {
 protected:
  Binding() = default;
 
 public:
  explicit Binding(ffi::UnsafeInit tag) : ffi::ObjectRef(tag) {}
  Binding(const Binding&) = default;
  Binding(Binding&&) = default;
  Binding& operator=(const Binding&) = default;
  Binding& operator=(Binding&&) = default;
  const BindingNode* operator->() const { return static_cast<const BindingNode*>(data_.get()); }
  const BindingNode* get() const { return operator->(); }
  static constexpr bool _type_is_nullable = false;
  using ContainerType = BindingNode;
};
 
class MatchCastNode : public BindingNode {
 public:
  explicit MatchCastNode(Var var, Expr value)
      : BindingNode(std::move(var)), value(std::move(value)) {}
  explicit MatchCastNode(ffi::UnsafeInit)
      : BindingNode(ffi::UnsafeInit{}), value(ffi::UnsafeInit{}) {}
 
  Expr value;
  Type ty = Type::Missing();
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<MatchCastNode>()
        .def_ro("value", &MatchCastNode::value)
        // TODO(tqchen): use SEqHashDefSimple after the next pypi tvm-ffi release
        .def_ro("ty", &MatchCastNode::ty, refl::AttachFieldFlag::SEqHashDefPattern());
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.expr.MatchCast", MatchCastNode, BindingNode);
};
 
class MatchCast : public Binding {
 public:
  TVM_DLL explicit MatchCast(Var var, Expr value, Type ty, Location loc = UnknownLoc());
 
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(MatchCast, Binding, MatchCastNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(MatchCastNode);
};
 
class VarBindingNode : public BindingNode {
 public:
  explicit VarBindingNode(Var var, Expr value)
      : BindingNode(std::move(var)), value(std::move(value)) {}
  explicit VarBindingNode(ffi::UnsafeInit)
      : BindingNode(ffi::UnsafeInit{}), value(ffi::UnsafeInit{}) {}
 
  Expr value;
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<VarBindingNode>().def_ro("value", &VarBindingNode::value);
    // customize the SEqual and SHash methods for better error messages
    refl::TypeAttrDef<VarBindingNode>()
        .def(tvm::ffi::reflection::type_attr::kSEqual, &VarBindingNode::SEqual)
        .def(tvm::ffi::reflection::type_attr::kSHash, &VarBindingNode::SHash);
  }
 
  bool SEqual(const VarBindingNode* other,
              ffi::TypedFunction<bool(AnyView, AnyView, bool, AnyView)> equal) const;
  int64_t SHash(int64_t init_hash, ffi::TypedFunction<int64_t(AnyView, int64_t, bool)> hash) const;
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.expr.VarBinding", VarBindingNode, BindingNode);
};
 
class VarBinding : public Binding {
 public:
  TVM_DLL explicit VarBinding(Var var, Expr value, Location loc = UnknownLoc());
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(VarBinding, Binding, VarBindingNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(VarBindingNode);
};
 
class BindingBlockNode : public ffi::Object {
 public:
  ffi::Array<Binding> bindings;
  mutable Location loc = UnknownLoc();
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<BindingBlockNode>()
        .def_ro("bindings", &BindingBlockNode::bindings)
        .def_ro("loc", &BindingBlockNode::loc, refl::AttachFieldFlag::SEqHashIgnore(),
                refl::DefaultValue(UnknownLoc()));
  }
 
  static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode;
  TVM_FFI_DECLARE_OBJECT_INFO("relax.expr.BindingBlock", BindingBlockNode, ffi::Object);
};
 
class BindingBlock : public ffi::ObjectRef {
 public:
  TVM_DLL explicit BindingBlock(ffi::Array<Binding> bindings, Location loc = UnknownLoc());
  explicit BindingBlock(ffi::UnsafeInit tag, ffi::ObjectPtr<BindingBlockNode> n)
      : ffi::ObjectRef(tag) {
    data_ = std::move(n);
  }
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(BindingBlock, ffi::ObjectRef, BindingBlockNode);
 
  BindingBlockNode* CopyOnWrite();
};
 
class DataflowBlockNode : public BindingBlockNode {
 public:
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<DataflowBlockNode>();
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.expr.DataflowBlock", DataflowBlockNode,
                                    BindingBlockNode);
};
 
class DataflowBlock : public BindingBlock {
 public:
  TVM_DLL explicit DataflowBlock(ffi::Array<Binding> bindings, Location loc = UnknownLoc());
  explicit DataflowBlock(ffi::UnsafeInit tag, ffi::ObjectPtr<DataflowBlockNode> n)
      : BindingBlock(tag) {
    data_ = std::move(n);
  }
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(DataflowBlock, BindingBlock, DataflowBlockNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(DataflowBlockNode);
};
 
class SeqExprNode : public ExprNode {
 public:
  explicit SeqExprNode(Expr body) : body(std::move(body)) {}
  explicit SeqExprNode(ffi::UnsafeInit) : body(ffi::UnsafeInit{}) {}
 
  ffi::Array<BindingBlock> blocks;
  Expr body;
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<SeqExprNode>()
        .def_ro("blocks", &SeqExprNode::blocks)
        .def_ro("body", &SeqExprNode::body);
    refl::TypeAttrDef<SeqExprNode>()
        .def(tvm::ffi::reflection::type_attr::kSEqual, &SeqExprNode::SEqual)
        .def(tvm::ffi::reflection::type_attr::kSHash, &SeqExprNode::SHash);
  }
 
  bool SEqual(const SeqExprNode* other,
              ffi::TypedFunction<bool(AnyView, AnyView, bool, AnyView)> equal) const {
    // Establish mappings for symbolic variables defined by bindings before
    // comparing their uses in the SeqExpr result type and body.
    return equal(blocks, other->blocks, false, "blocks") && equal(ty, other->ty, false, "ty") &&
           equal(body, other->body, false, "body");
  }
 
  int64_t SHash(int64_t init_hash, ffi::TypedFunction<int64_t(AnyView, int64_t, bool)> hash) const {
    int64_t hash_value = init_hash;
    hash_value = hash(blocks, hash_value, false);
    hash_value = hash(ty, hash_value, false);
    hash_value = hash(body, hash_value, false);
    return hash_value;
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.expr.SeqExpr", SeqExprNode, ExprNode);
};
 
class SeqExpr : public Expr {
 public:
  /* \brief Implicit conversion constructor
   *
   * Relax nodes that introduce a new scope (e.g. `relax::Function`)
   * are required to be held as SeqExpr.  This implicit conversion
   * provides allows callsites to use these member variables when the
   * C++ compile-time type is a `relax::Expr`.  For example,
   * a transform may use `func.CopyOnWrite()->body = expr;`.
   *
   * If the expression is already a `relax::SeqExpr`, the same
   * underlying `relax::SeqExprNode` is used, and no copies are made.
   */
  TVM_DLL SeqExpr(Expr body);  // NOLINT(*)
 
  TVM_DLL explicit SeqExpr(ffi::Array<BindingBlock> blocks, Expr body, Location loc = UnknownLoc());
  explicit SeqExpr(ffi::ObjectPtr<SeqExprNode> node) : Expr(std::move(node)) {}
 
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(SeqExpr, Expr, SeqExprNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(SeqExprNode);
};
 
class IfExprNode : public ExprNode {
 public:
  explicit IfExprNode(Expr cond, SeqExpr true_branch, SeqExpr false_branch)
      : cond(std::move(cond)),
        true_branch(std::move(true_branch)),
        false_branch(std::move(false_branch)) {}
  explicit IfExprNode(ffi::UnsafeInit)
      : cond(ffi::UnsafeInit{}), true_branch(ffi::UnsafeInit{}), false_branch(ffi::UnsafeInit{}) {}
 
  Expr cond;
  SeqExpr true_branch;
  SeqExpr false_branch;
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<IfExprNode>()
        .def_ro("cond", &IfExprNode::cond)
        .def_ro("true_branch", &IfExprNode::true_branch)
        .def_ro("false_branch", &IfExprNode::false_branch);
  }
 
  static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindDAGNode;
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.expr.IfExpr", IfExprNode, ExprNode);
};
 
class IfExpr : public Expr {
 public:
  TVM_DLL IfExpr(Expr cond, Expr true_branch, Expr false_branch, Location loc = UnknownLoc());
 
  explicit IfExpr(ffi::ObjectPtr<IfExprNode> node) : Expr(std::move(node)) {}
 
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(IfExpr, Expr, IfExprNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(IfExprNode);
};
 
class FunctionNode : public BaseFuncNode {
 public:
  explicit FunctionNode(SeqExpr body) : body(std::move(body)) {}
  explicit FunctionNode(ffi::UnsafeInit) : body(ffi::UnsafeInit{}) {}
 
  ffi::Array<Var> params;
  SeqExpr body;
  Type ret_ty = Type::Missing();
  bool is_pure;
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<FunctionNode>()
        .def_ro("params", &FunctionNode::params, refl::AttachFieldFlag::SEqHashDefPattern())
        .def_ro("body", &FunctionNode::body)
        .def_ro("ret_ty", &FunctionNode::ret_ty)
        .def_ro("is_pure", &FunctionNode::is_pure);
  }
 
  static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindDAGNode;
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.expr.Function", FunctionNode, BaseFuncNode);
};
 
class Function : public BaseFunc {
 public:
  TVM_DLL explicit Function(ffi::Array<Var> params, Expr body, ffi::Optional<Type> ret_ty,
                            bool is_pure = true, DictAttrs attrs = DictAttrs(),
                            Location loc = UnknownLoc());
 
  TVM_DLL static Function CreateEmpty(ffi::Array<Var> params, Type ret_ty, bool is_pure = true,
                                      DictAttrs attrs = DictAttrs(), Location loc = UnknownLoc());
 
  explicit Function(ffi::ObjectPtr<FunctionNode> node) : BaseFunc(std::move(node)) {}
 
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(Function, BaseFunc, FunctionNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(FunctionNode);
};
 
class ExternFuncNode : public BaseFuncNode {
 public:
  ffi::String global_symbol;
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<ExternFuncNode>().def_ro("global_symbol", &ExternFuncNode::global_symbol);
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.expr.ExternFunc", ExternFuncNode, BaseFuncNode);
};
 
class ExternFunc : public BaseFunc {
 public:
  TVM_DLL ExternFunc(ffi::String global_symbol, Location loc = UnknownLoc());
  TVM_DLL ExternFunc(ffi::String global_symbol, Type ty, Location loc = UnknownLoc());
 
  explicit ExternFunc(ffi::ObjectPtr<ExternFuncNode> node) : BaseFunc(std::move(node)) {}
 
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(ExternFunc, BaseFunc, ExternFuncNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(ExternFuncNode);
};
 
TVM_DLL Expr GetShapeOf(const Expr& expr);
 
// TODO(@sunggg): Investigate the exact usage of kComposite, kPartitionedFromPattern, and
// kPrimitive.
namespace attr {
constexpr const char* kPrimitive = "Primitive";
constexpr const char* kCodegen = "Codegen";
constexpr const char* kComposite = "Composite";
constexpr const char* kPartitionedFromPattern = "PartitionedFromPattern";
constexpr const char* kWorkspaceSize = "WorkspaceSize";
 
// Note: in the future, we prefer snake_case instead of CamelCase for attributes.
// Past ones will be kept for backwards compatibility.
constexpr const char* kForcePure = "relax.force_pure";
 
constexpr const char* kNumInput = "num_input";
constexpr const char* kOpPattern = "op_pattern";
 
constexpr const char* kOpAttrs = "op_attrs";
 
constexpr const char* kOperatorName = "operator_name";
 
constexpr const char* kLibraryKernel = "library_kernel";
 
constexpr const char* kCSource = "c_source";
 
constexpr const char* kCSourceFmt = "c_source_fmt";
 
constexpr const char* kTirVarUpperBound = "tir_var_upper_bound";
 
constexpr const char* kTirVarLowerBound = "tir_var_lower_bound";
 
constexpr const char* kTirNonNegativeVar = "tir_non_negative_var";
 
constexpr const char* kMemoryPlanDynamicFuncOutput = "relax.memory_plan_dynamic_func_output";
 
constexpr const char* kCaptureSymbolicVars = "relax.rewrite_cuda_graph.capture_symbolic_vars";
 
constexpr const char* kScoped = "scoped";
 
}  // namespace attr
 
}  // namespace relax
}  // namespace tvm
 
#endif  // TVM_RELAX_EXPR_H_