MarisaKirisame
diff --git a/‎include/tvm/relay/expr.h‎
Lines changed: 40 additions & 1 deletion b/‎include/tvm/relay/expr.h‎
Lines changed: 40 additions & 1 deletion
diff --git a/‎include/tvm/relay/expr_functor.h‎
Lines changed: 4 additions & 0 deletions b/‎include/tvm/relay/expr_functor.h‎
Lines changed: 4 additions & 0 deletions
@@ -184,6 +184,26 @@ class VarNode : public ExprNode {
 
 RELAY_DEFINE_NODE_REF(Var, VarNode, Expr);
 
+/*! \brief Hash Var by it's id.
+ * Different VarNode might has same vid, and they are considered to be the same var in such case.
+ * Use VarHash to hash Var by id.
+ */
+struct VarHash {
+  size_t operator()(const Var& v) const {
+    return v->vid.hash();
+  }
+};
+
+/*! \brief Compare Var by it's id.
+ * Different VarNode might has same vid, and they are considered to be the same var in such case.
+ * Use VarEqual to compare Var by id.
+ */
+struct VarEqual {
+  bool operator()(const Var& l, const Var& r) const {
+    return l->vid.get() == r->vid.get();
+  }
+};
+
 /*!
  * \brief Global variable that leaves in the top-level module.
  * This is used to enable recursive calls between function.
@@ -521,7 +541,7 @@ RELAY_DEFINE_NODE_REF(RefWrite, RefWriteNode, Expr);
  * rewriting pass such as layout or type transformation.
  *
  * Subclass TempExprNode allows us to pattern match on
- * specific kind TempExpr and use them for expression rewriting.
+ * specific kind of TempExpr and use them for expression rewriting.
  *
  * TempExpr should only be used within a pass,
  */
@@ -539,6 +559,25 @@ class TempExprNode : public ExprNode {
 
 RELAY_DEFINE_NODE_REF(TempExpr, TempExprNode, Expr);
 
+class Annotate;
+class AnnotateNode : public ExprNode {
+ public:
+  Expr expr;
+  NodeRef annotation;
+  void VisitAttrs(tvm::AttrVisitor* v) final {
+    v->Visit("expr", &expr);
+    v->Visit("annotation", &annotation);
+    v->Visit("_checked_type_", &checked_type_);
+  }
+
+  TVM_DLL static Annotate make(Expr expr, NodeRef annotation);
+
+  static constexpr const char* _type_key = "relay.AnnotateNode";
+  TVM_DECLARE_NODE_TYPE_INFO(AnnotateNode, ExprNode);
+};
+
+RELAY_DEFINE_NODE_REF(Annotate, AnnotateNode, Expr);
+
 // implementataions
 inline const Type& ExprNode::checked_type() const {
   CHECK(checked_type_.defined()) << "internal error: the type checker has "
 
@@ -116,6 +116,7 @@ class ExprFunctor<R(const Expr& n, Args...)> {
   virtual R VisitExpr_(const RefWriteNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
   virtual R VisitExpr_(const ConstructorNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
   virtual R VisitExpr_(const MatchNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
+  virtual R VisitExpr_(const AnnotateNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
   virtual R VisitExprDefault_(const Node* op, Args...) {
     throw Error(std::string("Do not have a default for ") + op->type_key());
   }
@@ -140,6 +141,7 @@ class ExprFunctor<R(const Expr& n, Args...)> {
     RELAY_EXPR_FUNCTOR_DISPATCH(RefWriteNode);
     RELAY_EXPR_FUNCTOR_DISPATCH(ConstructorNode);
     RELAY_EXPR_FUNCTOR_DISPATCH(MatchNode);
+    RELAY_EXPR_FUNCTOR_DISPATCH(AnnotateNode);
     return vtable;
   }
 };
@@ -170,6 +172,7 @@ class ExprVisitor
   void VisitExpr_(const RefWriteNode* op) override;
   void VisitExpr_(const ConstructorNode* op) override;
   void VisitExpr_(const MatchNode* op) override;
+  void VisitExpr_(const AnnotateNode* op) override;
   virtual void VisitType(const Type& t);
   virtual void VisitClause(const Clause& c);
   virtual void VisitPattern(const Pattern& c);
@@ -212,6 +215,7 @@ class ExprMutator
   Expr VisitExpr_(const RefWriteNode* op) override;
   Expr VisitExpr_(const ConstructorNode* op) override;
   Expr VisitExpr_(const MatchNode* op) override;
+  Expr VisitExpr_(const AnnotateNode* op) override;
 
   /*!
    * \brief Used to visit the types inside of expressions.