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_RELAY_EXPR_H_
#define TVM_RELAY_EXPR_H_

#include <tvm/ir/attrs.h>
#include <tvm/ir/expr.h>
#include <tvm/ir/module.h>
#include <tvm/ir/op.h>
#include <tvm/target/virtual_device.h>

#include <functional>
#include <stack>
#include <string>
#include <utility>

#include "./base.h"
#include "./type.h"

namespace tvm {

GlobalVar WithFields(GlobalVar global_var, Optional<String> opt_name_hint = {},
                     Optional<Type> opt_type = {}, Optional<VirtualDevice> opt_virtual_device = {},
                     Optional<Span> opt_span = {});

namespace relay {

using Expr = tvm::RelayExpr;
using ExprNode = tvm::RelayExprNode;
using BaseFunc = tvm::BaseFunc;
using BaseFuncNode = tvm::BaseFuncNode;
using GlobalVar = tvm::GlobalVar;
using GlobalVarNode = tvm::GlobalVarNode;
using tvm::PrettyPrint;

class Constant;
class ConstantNode : public ExprNode {
 public:
  runtime::NDArray data;

  TensorType tensor_type() const;

  bool is_scalar() const { return data->ndim == 0; }

  void VisitAttrs(tvm::AttrVisitor* v) {
    v->Visit("data", &data);
    v->Visit("virtual_device_", &virtual_device_);
    v->Visit("span", &span);
    v->Visit("_checked_type_", &checked_type_);
  }

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

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

  static constexpr const char* _type_key = "relay.Constant";
  TVM_DECLARE_FINAL_OBJECT_INFO(ConstantNode, ExprNode);
};

class Constant : public Expr {
 public:
  TVM_DLL explicit Constant(runtime::NDArray data, Span span = Span());

  TVM_DEFINE_OBJECT_REF_METHODS(Constant, RelayExpr, ConstantNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(ConstantNode);
};

Constant WithFields(Constant constant, Optional<runtime::NDArray> opt_data = {},
                    Optional<VirtualDevice> opt_virtual_device = {}, Optional<Span> opt_span = {});

class Tuple;
class TupleNode : public ExprNode {
 public:
  tvm::Array<relay::Expr> fields;

  void VisitAttrs(tvm::AttrVisitor* v) {
    v->Visit("fields", &fields);
    v->Visit("virtual_device_", &virtual_device_);
    v->Visit("span", &span);
    v->Visit("_checked_type_", &checked_type_);
  }

  bool SEqualReduce(const TupleNode* other, SEqualReducer equal) const {
    // specially handle empty tuple as a constant is not a graph node.
    if (fields.size() == other->fields.size() && fields.size() == 0) {
      return true;
    } else {
      equal->MarkGraphNode();
      return equal(fields, other->fields);
    }
  }

  void SHashReduce(SHashReducer hash_reduce) const {
    if (fields.size() != 0) {
      hash_reduce->MarkGraphNode();
      hash_reduce(fields);
    }
  }

  static constexpr const char* _type_key = "relay.Tuple";
  TVM_DECLARE_FINAL_OBJECT_INFO(TupleNode, ExprNode);
};

class Tuple : public Expr {
 public:
  TVM_DLL explicit Tuple(tvm::Array<relay::Expr> fields, Span span = Span());

  TVM_DEFINE_OBJECT_REF_METHODS(Tuple, RelayExpr, TupleNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(TupleNode);
};

Tuple WithFields(Tuple tuple, Optional<Array<Expr>> opt_fields = Optional<Array<Expr>>(),
                 Optional<VirtualDevice> opt_virtual_device = Optional<VirtualDevice>(),
                 Optional<Span> opt_span = Optional<Span>());

class Var;
class VarNode : public ExprNode {
 public:
  Id vid;
  Type type_annotation;

  const String& name_hint() const { return vid->name_hint; }

  void VisitAttrs(tvm::AttrVisitor* v) {
    v->Visit("vid", &vid);
    v->Visit("type_annotation", &type_annotation);
    v->Visit("virtual_device_", &virtual_device_);
    v->Visit("span", &span);
    v->Visit("_checked_type_", &checked_type_);
  }

  bool SEqualReduce(const VarNode* other, SEqualReducer equal) const {
    equal->MarkGraphNode();
    return equal(type_annotation, other->type_annotation) && equal(vid, other->vid) &&
           equal(virtual_device_, other->virtual_device_);
  }

  void SHashReduce(SHashReducer hash_reduce) const {
    hash_reduce->MarkGraphNode();
    hash_reduce(type_annotation);
    hash_reduce(vid);
  }

  static constexpr const char* _type_key = "relay.Var";
  TVM_DECLARE_FINAL_OBJECT_INFO(VarNode, ExprNode);
};

class Var : public Expr {
 public:
  TVM_DLL Var(String name_hint, Type type_annotation, Span span = Span())
      : Var(Id(name_hint), type_annotation, span) {}

  TVM_DLL Var(Id vid, Type type_annotation, Span span = Span());

  static Var GenSym(Type type_annotation = {}, Span span = {});

  TVM_DEFINE_OBJECT_REF_METHODS(Var, RelayExpr, VarNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(VarNode);
};

Var WithFields(Var var, Optional<Id> opt_vid = Optional<Id>(),
               Optional<Type> opt_type_annotation = Optional<Type>(),
               Optional<VirtualDevice> opt_virtual_device = Optional<VirtualDevice>(),
               Optional<Span> opt_span = Optional<Span>());

class Call;
class CallNode : public ExprNode {
 protected:
  // CallNode uses own deleter to indirectly call non-recursive destructor
  Object::FDeleter saved_deleter_;
  static void Deleter_(Object* ptr);

 public:
  Expr op;

  tvm::Array<relay::Expr> args;

  Attrs attrs;

  tvm::Array<Type> type_args;

  void VisitAttrs(tvm::AttrVisitor* v) {
    v->Visit("op", &op);
    v->Visit("args", &args);
    v->Visit("attrs", &attrs);
    v->Visit("type_args", &type_args);
    v->Visit("virtual_device_", &virtual_device_);
    v->Visit("span", &span);
    v->Visit("_checked_type_", &checked_type_);
  }

  bool SEqualReduce(const CallNode* other, SEqualReducer equal) const {
    // skip type_args check for primitive ops.
    equal->MarkGraphNode();
    return equal(op, other->op) && equal(args, other->args) && equal(attrs, other->attrs) &&
           (IsPrimitiveOp(op) || equal(type_args, other->type_args));
  }

  void SHashReduce(SHashReducer hash_reduce) const {
    hash_reduce->MarkGraphNode();
    hash_reduce(op);
    hash_reduce(args);
    hash_reduce(attrs);
    if (!IsPrimitiveOp(op)) {
      hash_reduce(type_args);
    }
  }

  static constexpr const char* _type_key = "relay.Call";
  TVM_DECLARE_FINAL_OBJECT_INFO(CallNode, ExprNode);
  template <typename>
  friend class runtime::ObjAllocatorBase;
  friend class Call;
};

class Call : public Expr {
 public:
  ~Call();

  TVM_DLL Call(Expr op, Array<Expr> args, Attrs attrs = Attrs(),
               Array<Type> type_args = Array<Type>(), Span span = Span());

  TVM_DEFINE_OBJECT_REF_METHODS(Call, RelayExpr, CallNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(CallNode);
};

Call WithFields(Call call, Optional<Expr> opt_op = Optional<Expr>(),
                Optional<Array<Expr>> opt_args = Optional<Array<Expr>>(),
                Optional<Attrs> opt_attrs = Optional<Attrs>(),
                Optional<Array<Type>> opt_type_args = Optional<Array<Type>>(),
                Optional<VirtualDevice> opt_virtual_device = Optional<VirtualDevice>(),
                Optional<Span> opt_span = Optional<Span>());

class Let;
class LetNode : public ExprNode {
 protected:
  // LetNode uses own deleter to indirectly call non-recursive destructor
  Object::FDeleter saved_deleter_;
  static void Deleter_(Object* ptr);

 public:
  Var var;
  Expr value;
  Expr body;

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

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

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

  static constexpr const char* _type_key = "relay.Let";
  TVM_DECLARE_FINAL_OBJECT_INFO(LetNode, ExprNode);
  template <typename>
  friend class runtime::ObjAllocatorBase;
  friend class Let;
};

class Let : public Expr {
 public:
  ~Let();

  TVM_DLL Let(Var var, Expr value, Expr body, Span span = Span());

  TVM_DEFINE_OBJECT_REF_METHODS(Let, RelayExpr, LetNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(LetNode);
};

Let WithFields(Let let, Optional<Var> opt_var = Optional<Var>(),
               Optional<Expr> opt_value = Optional<Expr>(),
               Optional<Expr> opt_body = Optional<Expr>(),
               Optional<VirtualDevice> opt_virtual_device = Optional<VirtualDevice>(),
               Optional<Span> opt_span = Optional<Span>());

class If;
class IfNode : public ExprNode {
 public:
  Expr cond;
  Expr true_branch;
  Expr false_branch;

  void VisitAttrs(tvm::AttrVisitor* v) {
    v->Visit("cond", &cond);
    v->Visit("true_branch", &true_branch);
    v->Visit("false_branch", &false_branch);
    v->Visit("virtual_device_", &virtual_device_);
    v->Visit("span", &span);
    v->Visit("_checked_type_", &checked_type_);
  }

  bool SEqualReduce(const IfNode* other, SEqualReducer equal) const {
    equal->MarkGraphNode();
    return equal(cond, other->cond) && equal(true_branch, other->true_branch) &&
           equal(false_branch, other->false_branch);
  }

  void SHashReduce(SHashReducer hash_reduce) const {
    hash_reduce->MarkGraphNode();
    hash_reduce(cond);
    hash_reduce(true_branch);
    hash_reduce(false_branch);
  }

  static constexpr const char* _type_key = "relay.If";
  TVM_DECLARE_FINAL_OBJECT_INFO(IfNode, ExprNode);
};

class If : public Expr {
 public:
  TVM_DLL If(Expr cond, Expr true_branch, Expr false_branch, Span span = Span());

  TVM_DEFINE_OBJECT_REF_METHODS(If, RelayExpr, IfNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(IfNode);
};

If WithFields(If if_expr, Optional<Expr> opt_cond = Optional<Expr>(),
              Optional<Expr> opt_true_branch = Optional<Expr>(),
              Optional<Expr> opt_false_branch = Optional<Expr>(),
              Optional<VirtualDevice> opt_virtual_device = Optional<VirtualDevice>(),
              Optional<Span> opt_span = Optional<Span>());

class TupleGetItem;
class TupleGetItemNode : public ExprNode {
 public:
  Expr tuple;
  int index;

  void VisitAttrs(tvm::AttrVisitor* v) {
    v->Visit("tuple_value", &tuple);
    v->Visit("index", &index);
    v->Visit("virtual_device_", &virtual_device_);
    v->Visit("span", &span);
    v->Visit("_checked_type_", &checked_type_);
  }

  bool SEqualReduce(const TupleGetItemNode* other, SEqualReducer equal) const {
    return equal(tuple, other->tuple) && equal(index, other->index);
  }

  void SHashReduce(SHashReducer hash_reduce) const {
    hash_reduce(tuple);
    hash_reduce(index);
  }

  static constexpr const char* _type_key = "relay.TupleGetItem";
  TVM_DECLARE_FINAL_OBJECT_INFO(TupleGetItemNode, ExprNode);
};

class TupleGetItem : public Expr {
 public:
  TVM_DLL TupleGetItem(Expr tuple, int index, Span span = Span());

  TVM_DEFINE_OBJECT_REF_METHODS(TupleGetItem, RelayExpr, TupleGetItemNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(TupleGetItemNode);
};

TupleGetItem WithFields(TupleGetItem tuple_get_item, Optional<Expr> opt_tuple = Optional<Expr>(),
                        Optional<Integer> opt_index = Optional<Integer>(),
                        Optional<VirtualDevice> opt_virtual_device = Optional<VirtualDevice>(),
                        Optional<Span> opt_span = Optional<Span>());

class RefCreate;
class RefCreateNode : public ExprNode {
 public:
  Expr value;

  void VisitAttrs(tvm::AttrVisitor* v) {
    v->Visit("value", &value);
    v->Visit("virtual_device_", &virtual_device_);
    v->Visit("span", &span);
    v->Visit("_checked_type_", &checked_type_);
  }

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

  void SHashReduce(SHashReducer hash_reduce) const {
    hash_reduce->MarkGraphNode();
    hash_reduce(value);
  }

  static constexpr const char* _type_key = "relay.RefCreate";
  TVM_DECLARE_FINAL_OBJECT_INFO(RefCreateNode, ExprNode);
};

class RefCreate : public Expr {
 public:
  TVM_DLL explicit RefCreate(Expr value, Span span = Span());

  TVM_DEFINE_OBJECT_REF_METHODS(RefCreate, RelayExpr, RefCreateNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(RefCreateNode);
};

RefCreate WithFields(RefCreate ref_create, Optional<Expr> opt_value = Optional<Expr>(),
                     Optional<VirtualDevice> opt_virtual_device = Optional<VirtualDevice>(),
                     Optional<Span> opt_span = Optional<Span>());

class RefRead;
class RefReadNode : public ExprNode {
 public:
  Expr ref;

  void VisitAttrs(tvm::AttrVisitor* v) {
    v->Visit("ref", &ref);
    v->Visit("virtual_device_", &virtual_device_);
    v->Visit("span", &span);
    v->Visit("_checked_type_", &checked_type_);
  }

  bool SEqualReduce(const RefReadNode* other, SEqualReducer equal) const {
    equal->MarkGraphNode();
    return equal(ref, other->ref);
  }

  void SHashReduce(SHashReducer hash_reduce) const {
    hash_reduce->MarkGraphNode();
    hash_reduce(ref);
  }

  static constexpr const char* _type_key = "relay.RefRead";
  TVM_DECLARE_FINAL_OBJECT_INFO(RefReadNode, ExprNode);
};

class RefRead : public Expr {
 public:
  TVM_DLL explicit RefRead(Expr ref, Span span = Span());

  TVM_DEFINE_OBJECT_REF_METHODS(RefRead, RelayExpr, RefReadNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(RefReadNode);
};

RefRead WithFields(RefRead ref_read, Optional<Expr> opt_ref = Optional<Expr>(),
                   Optional<VirtualDevice> opt_virtual_device = Optional<VirtualDevice>(),
                   Optional<Span> opt_span = Optional<Span>());

class RefWrite;
class RefWriteNode : public ExprNode {
 public:
  Expr ref;
  Expr value;

  void VisitAttrs(tvm::AttrVisitor* v) {
    v->Visit("ref", &ref);
    v->Visit("value", &value);
    v->Visit("virtual_device_", &virtual_device_);
    v->Visit("span", &span);
    v->Visit("_checked_type_", &checked_type_);
  }

  bool SEqualReduce(const RefWriteNode* other, SEqualReducer equal) const {
    equal->MarkGraphNode();
    return equal(ref, other->ref) && equal(value, other->value);
  }

  void SHashReduce(SHashReducer hash_reduce) const {
    hash_reduce->MarkGraphNode();
    hash_reduce(ref);
    hash_reduce(value);
  }

  static constexpr const char* _type_key = "relay.RefWrite";
  TVM_DECLARE_FINAL_OBJECT_INFO(RefWriteNode, ExprNode);
};

class RefWrite : public Expr {
 public:
  TVM_DLL RefWrite(Expr ref, Expr value, Span span = Span());

  TVM_DEFINE_OBJECT_REF_METHODS(RefWrite, RelayExpr, RefWriteNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(RefWriteNode);
};

RefWrite WithFields(RefWrite ref_write, Optional<Expr> opt_ref = Optional<Expr>(),
                    Optional<Expr> opt_value = Optional<Expr>(),
                    Optional<VirtualDevice> opt_virtual_device = Optional<VirtualDevice>(),
                    Optional<Span> opt_span = Optional<Span>());

class TempExprNode : public ExprNode {
 public:
  virtual ~TempExprNode() {}
  virtual Expr Realize() const = 0;

  static constexpr const char* _type_key = "relay.TempExpr";
  static constexpr const bool _type_has_method_sequal_reduce = false;
  static constexpr const bool _type_has_method_shash_reduce = false;
  static constexpr const uint32_t _type_child_slots = 0;
  TVM_DECLARE_BASE_OBJECT_INFO(TempExprNode, ExprNode);
};

class TempExpr : public Expr {
 public:
  TVM_DEFINE_OBJECT_REF_METHODS(TempExpr, RelayExpr, TempExprNode);
};

}  // namespace relay

namespace runtime {

template <>
template <>
inline ObjectPtr<relay::LetNode>
ObjAllocatorBase<SimpleObjAllocator>::make_object<relay::LetNode>() {
  using Derived = SimpleObjAllocator;
  using T = relay::LetNode;
  using Handler = typename Derived::template Handler<T>;
  static_assert(std::is_base_of<Object, T>::value, "make can only be used to create Object");
  T* ptr = Handler::New(static_cast<Derived*>(this));
  ptr->type_index_ = T::RuntimeTypeIndex();
  ptr->saved_deleter_ = Handler::Deleter();
  ptr->deleter_ = relay::LetNode::Deleter_;
  return ObjectPtr<T>(ptr);
}

template <>
template <>
inline ObjectPtr<relay::CallNode>
ObjAllocatorBase<SimpleObjAllocator>::make_object<relay::CallNode>() {
  using Derived = SimpleObjAllocator;
  using T = relay::CallNode;
  using Handler = typename Derived::template Handler<T>;
  static_assert(std::is_base_of<Object, T>::value, "make can only be used to create Object");
  T* ptr = Handler::New(static_cast<Derived*>(this));
  ptr->type_index_ = T::RuntimeTypeIndex();
  ptr->saved_deleter_ = Handler::Deleter();
  ptr->deleter_ = relay::CallNode::Deleter_;
  return ObjectPtr<T>(ptr);
}

}  // namespace runtime

}  // namespace tvm
#endif  // TVM_RELAY_EXPR_H_