[Relay] WithFields method for Call, Function, Var, TupleGetItem, If, Let, RefCreate, RefRead, RefWrite, Match, and Clause (#9569)

* Implement WithFields for Relay exprs

* lint

Author: electriclilies

include/tvm/relay/adt.h

@@ -260,8 +260,25 @@ class Clause : public ObjectRef {
   TVM_DLL explicit Clause(Pattern lhs, Expr rhs);
 
   TVM_DEFINE_OBJECT_REF_METHODS(Clause, ObjectRef, ClauseNode);
+  TVM_DEFINE_OBJECT_REF_COW_METHOD(ClauseNode);
 };
 
+/*!
+ * \brief Returns the clause with given properties. A null property denotes 'no change'.
+ * Returns clause if all properties are unchanged. Otherwise, returns a copy with the new fields.
+ * \param clause The clause to copy.
+ * \param opt_lhs The (optional) lhs for the copied clause. If none, ret_clause->lhs = clause->lhs.
+ * \param opt_rhs The (optional) rhs for the copied clause. If none,
+ * ret_clause->rhs = clause->rhs.
+ * \return If all
+ * properties are null or the same as the property in the input clause (i.e., opt_lhs is null or
+ * opt_lhs.value() == clause->lhs, etc.), then we return clause. Otherwise, we return a copy of
+ * clause with the different fields overwritten. (i.e., if opt_lhs.value() != clause->lhs, then
+ * ret_clause->lhs = opt_lhs.value()).
+ */
+Clause WithFields(Clause clause, Optional<Pattern> opt_lhs = Optional<Pattern>(),
+                  Optional<Expr> opt_rhs = Optional<Expr>());
+
 /*! \brief ADT pattern matching exression. */
 class Match;
 /*! \brief Match container node. */
@@ -315,8 +332,30 @@ class Match : public Expr {
   TVM_DLL Match(Expr data, tvm::Array<Clause> clauses, bool complete = true, Span span = Span());
 
   TVM_DEFINE_OBJECT_REF_METHODS(Match, RelayExpr, MatchNode);
+  TVM_DEFINE_OBJECT_REF_COW_METHOD(MatchNode);
 };
 
+/*!
+ * \brief Returns the match with given properties. A null property denotes 'no change'.
+ * Returns match if all properties are unchanged. Otherwise, returns a copy with the new fields.
+ * \param match The match to copy.
+ * \param opt_data The (optional) data for the copied match. If none, ret_match->data = match->data.
+ * \param opt_clauses The (optional) clauses for the copied match. If none, ret_match->clauses =
+ * match->clauses.
+ * \param opt_complete The (optional) complete for the copied match. If none, ret_match->complete =
+ * match->complete.
+ * \param opt_span The (optional) span for the copied match. If none, ret_match->span = match->span.
+ * \return If all properties are null or the same as the
+ * property in the input match (i.e., opt_clauses is null or opt_clauses.value() == match->clauses,
+ * etc.), then we return match. Otherwise, we return a copy of match with the different fields
+ * overwritten. (i.e., if opt_clauses.value() != match->clauses, then ret_match->clauses =
+ * opt_clauses.value()).
+ */
+Match WithFields(Match match, Optional<Expr> opt_data = Optional<Expr>(),
+                 Optional<Array<Clause>> opt_clauses = Optional<Array<Clause>>(),
+                 Optional<Bool> opt_complete = Optional<Bool>(),
+                 Optional<Span> opt_span = Optional<Span>());
+
 }  // namespace relay
 }  // namespace tvm
 

include/tvm/relay/expr.h

@@ -230,8 +230,26 @@ class Var : public Expr {
   TVM_DLL Var(Id vid, Type type_annotation, Span span = Span());
 
   TVM_DEFINE_OBJECT_REF_METHODS(Var, RelayExpr, VarNode);
+  TVM_DEFINE_OBJECT_REF_COW_METHOD(VarNode);
 };
 
+/*!
+ * \brief Returns the var with given properties. A null property denotes 'no change'.
+ * Returns var if all properties are unchanged. Otherwise, returns a copy with the new fields.
+ * \param var The var to copy.
+ * \param opt_vid The (optional) vid for the copied var. If none, ret_var->vid = var->vid.
+ * \param opt_type_annotation The (optional) type_annotation for the copied var. If none,
+ * ret_var->type_annotation = var->type_annotation.
+ * \param opt_span The (optional) span for the copied var. If none, ret_var->span = var->span.
+ * \return If all properties are null or the same as the property in the input var
+ * (i.e., opt_vid is null or opt_vid.value() == var->vid, etc.), then we return var. Otherwise,
+ * we return a copy of call with the different fields overwritten. (i.e., if
+ * opt_vid.value() != var->vid, then ret_var->vid = opt_.value()).
+ */
+Var WithFields(Var var, Optional<Id> opt_vid = Optional<Id>(),
+               Optional<Type> opt_type_annotation = Optional<Type>(),
+               Optional<Span> opt_span = Optional<Span>());
+
 /*!
  * \brief Call corresponds to operator invocation.
  *  Corresponds to the operator in computational graph terminology.
@@ -331,8 +349,31 @@ class Call : public Expr {
                Array<Type> type_args = Array<Type>(), Span span = Span());
 
   TVM_DEFINE_OBJECT_REF_METHODS(Call, RelayExpr, CallNode);
+  TVM_DEFINE_OBJECT_REF_COW_METHOD(CallNode);
 };
 
+/*!
+ * \brief Returns the call with given properties. A null property denotes 'no change'.
+ * Returns call if all properties are unchanged. Otherwise, returns a copy with the new fields.
+ * \param call The call to copy.
+ * \param opt_op The (optional) op for the copied call. If none, ret_call->op = call->op.
+ * \param opt_args The (optional) args for the copied call. If none, ret_call->args = call->args.
+ * \param opt_attrs The (optional) attrs for the copied call. If none, ret_call->attrs =
+ * call->attrs.
+ * \param opt_type_args The (optional) type args for the copied call. If none,
+ * ret_call->type_args = call->type_args.
+ * \param opt_span The (optional) span for the copied call. If none, ret_call->span = call->span.
+ * \return If all properties are null or the same as the property in the input call
+ * (i.e., opt_op is null or opt_op.value() == call->op, etc.), then we return call. Otherwise, we
+ * return a copy of call with the different fields overwritten. (i.e., if opt_op.value() !=
+ * call->op, then ret_call->op = opt_op.value()).
+ */
+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<Span> opt_span = Optional<Span>());
+
 /*!
  * \brief Let binding that binds a local var and optionally a type annotation.
  *
@@ -405,8 +446,27 @@ class Let : public Expr {
   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);
 };
 
+/*!
+ * \brief Returns the let with given properties. A null property denotes 'no change'.
+ * Returns let if all properties are unchanged. Otherwise, returns a copy with the new fields.
+ * \param let The let to copy.
+ * \param opt_var The (optional) var for the copied let. If none, ret_let->op = let->op.
+ * \param opt_value The (optional) value for the copied let. If none, ret_let->args = let->args.
+ * \param opt_body The (optional) body for the copied let. If none, ret_let->attrs = let->attrs.
+ * \param opt_span The (optional) span for the copied let. If none, ret_let->span = let->span.
+ * \return If all properties are null or the same as the property in the input let (i.e., opt_var is
+ * null or opt_var.value() == let->var, etc.), then we return let. Otherwise, we return a copy of
+ * let with the different fields overwritten. (i.e., if opt_var.value() != let->var, then
+ * ret_let->var = opt_var.value()).
+ */
+Let WithFields(Let let, Optional<Var> opt_var = Optional<Var>(),
+               Optional<Expr> opt_value = Optional<Expr>(),
+               Optional<Expr> opt_body = Optional<Expr>(),
+               Optional<Span> opt_span = Optional<Span>());
+
 /*!
  * \brief Condition expression
  *
@@ -466,8 +526,32 @@ class If : public Expr {
   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);
 };
 
+/*!
+ * \brief Returns the if_expr with given properties. A null property denotes 'no change'.
+ * Returns if_expr if all properties are unchanged. Otherwise, returns a copy with the new fields.
+ * \param if_expr The if expression to copy.
+ * \param opt_cond The (optional) cond for the copied if_expr. If none, ret_if->cond =
+ * if_expr->cond.
+ * \param opt_true_branch The (optional) true_branch for the copied if_expr. If none,
+ * ret_if->true_branch = ret_if->false_branch.
+ * \param opt_false_branch The (optional) false_branch
+ * for the copied if_expr. If none, ret_if->false_branch = if_expr->false_branch.
+ * \param opt_span
+ * The (optional) span for the copied if_expr. If none, ret_if->span = if_expr->span.
+ * \return If all
+ * properties are null or the same as the property in the input if_expr (i.e., opt_cond is null or
+ * opt_cond.value() == if_expr->cond, etc.), then we return if_expr. Otherwise, we return a copy of
+ * if_expr with the different fields overwritten. (i.e., if opt_cond.value() != if_expr->cond, then
+ * ret_if->cond = opt_cond.value()).
+ */
+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<Span> opt_span = Optional<Span>());
+
 /*! \brief Get index-th field out of a tuple. */
 class TupleGetItem;
 class TupleGetItemNode : public ExprNode {
@@ -508,8 +592,30 @@ class TupleGetItem : public Expr {
   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);
 };
 
+/*!
+ * \brief Returns the tuple_get_item with given properties. A null property denotes 'no change'.
+ * Returns if_expr if all properties are unchanged. Otherwise, returns a copy with the new fields.
+ * \param tuple_get_item The tuple_get_item to copy.
+ * \param opt_tuple The (optional) tuple for the copied tuple_get_item. If none,
+ * ret_tuple_get_item->tuple = tuple_get_item->tuple.
+ * \param opt_index The (optional) index for the copied tuple_get_item. If none,
+ * ret_tuple_get_item->index = tuple_get_item->index.
+ * \param
+ * opt_span The (optional) span for the copied tuple_get_item. If none,
+ * ret_tuple_get_item->span = tuple_get_item->span.
+ * \return If all properties are null or the same as the property in the input tuple_get_item
+ * (i.e., opt_tuple is null or opt_tuple.value() == tuple_get_item->tuple, etc.), then we return
+ * tuple_get_item. Otherwise, we return a copy of tuple_get_item with the different fields
+ * overwritten. (i.e., if opt_tuple.value() != tuple_get_item->tuple, then
+ * ret_tuple_get_item->tuple = opt_tuple.value()).
+ */
+TupleGetItem WithFields(TupleGetItem tuple_get_item, Optional<Expr> opt_tuple = Optional<Expr>(),
+                        Optional<Integer> opt_index = Optional<Integer>(),
+                        Optional<Span> opt_span = Optional<Span>());
+
 /*! \brief Create a new Reference out of initial value. */
 class RefCreate;
 class RefCreateNode : public ExprNode {
@@ -547,8 +653,27 @@ class RefCreate : public Expr {
   TVM_DLL explicit RefCreate(Expr value, Span span = Span());
 
   TVM_DEFINE_OBJECT_REF_METHODS(RefCreate, RelayExpr, RefCreateNode);
+  TVM_DEFINE_OBJECT_REF_COW_METHOD(RefCreateNode);
 };
 
+/*!
+ * \brief Returns the ref create with given properties. A null property denotes 'no change'.
+ * Returns ref_create if all properties are unchanged. Otherwise, returns a copy with the new
+ * fields.
+ * \param ref_create The ref_create to copy.
+ * \param opt_value The (optional) value for the copied ref_create. If none,
+ * ret_ref_create->value = ref_create->value.
+ * \param opt_span The (optional) span for the copied ref_create. If none,
+ * ret_ref_create->span = ref_create->span.
+ * \return If all properties are null or the same as the property in the input ref_create
+ * (i.e., opt_value is null or opt_value.value() == ref_create->value, etc.), then we return
+ * ref_create. Otherwise, we return a copy of ref_create with the different fields overwritten.
+ * (i.e., if opt_value.value() != ref_create->value, then
+ * ret_ref_create->value = opt_value.value()).
+ */
+RefCreate WithFields(RefCreate ref_create, Optional<Expr> opt_value = Optional<Expr>(),
+                     Optional<Span> opt_span = Optional<Span>());
+
 /*! \brief Get value out of Reference. */
 class RefRead;
 class RefReadNode : public ExprNode {
@@ -586,7 +711,26 @@ class RefRead : public Expr {
   TVM_DLL explicit RefRead(Expr ref, Span span = Span());
 
   TVM_DEFINE_OBJECT_REF_METHODS(RefRead, RelayExpr, RefReadNode);
+  TVM_DEFINE_OBJECT_REF_COW_METHOD(RefReadNode);
 };
+
+/*!
+ * \brief Returns the ref read with given properties. A null property denotes 'no change'.
+ * Returns ref_read if all properties are unchanged. Otherwise, returns a copy with the new fields.
+ * \param ref_read The ref_read to copy.
+ * \param opt_ref The (optional) ref for the copied ref_read. If none, ret_ref_read->ref =
+ * ref_read->ref.
+ * \param opt_span
+ * The (optional) span for the copied ref_read. If none, ret_ref_read->span = ref_read->span.
+ * \return If all properties are null or the same as the property in the input ref_read
+ * (i.e., opt_ref is null or opt_ref.value() == ref_read->ref, etc.), then we return ref_read.
+ * Otherwise, we return a copy of ref_read with the different fields overwritten.
+ * (i.e., if opt_ref.value() != ref_read->ref, then
+ * ret_ref_read->ref = opt_ref.value()).
+ */
+RefRead WithFields(RefRead ref_read, Optional<Expr> opt_ref = Optional<Expr>(),
+                   Optional<Span> opt_span = Optional<Span>());
+
 /*! \brief Set value of Reference. The whole expression evaluates to an Empty Tuple. */
 class RefWrite;
 class RefWriteNode : public ExprNode {
@@ -629,8 +773,29 @@ class RefWrite : public Expr {
   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);
 };
 
+/*!
+ * \brief Returns the ref write with given properties. A null property denotes 'no change'.
+ * Returns ref_write if all properties are unchanged. Otherwise, returns a copy with the new fields.
+ * \param ref_write The ref_write to copy.
+ * \param opt_ref The (optional) ref for the copied ref_write. If none,
+ * ret_ref_write->ref = ref_write->ref.
+ * \param opt_value The (optional) value for the copied ref_write. If none,
+ * ret_ref_write->value = ref_write->value.
+ * \param opt_span
+ * The (optional) span for the copied ref_write. If none, ret_ref_write->span = ref_write->span.
+ * \return If all properties are null or the same as the property in the input ref_write
+ * (i.e., opt_ref is null or opt_ref.value() == ref_write->ref, etc.), then we return ref_write.
+ * Otherwise, we return a copy of ref_write with the different fields overwritten.
+ * (i.e., if ref_write.value() != ref_write->ref, then
+ * ret_ref_write->ref = opt_ref.value()).
+ */
+RefWrite WithFields(RefWrite ref_write, Optional<Expr> opt_ref = Optional<Expr>(),
+                    Optional<Expr> opt_value = Optional<Expr>(),
+                    Optional<Span> opt_span = Optional<Span>());
+
 /*!
  * \brief Base class of the temporary expression.
  *

include/tvm/relay/function.h

@@ -119,6 +119,35 @@ class Function : public BaseFunc {
   TVM_DEFINE_OBJECT_REF_COW_METHOD(FunctionNode);
 };
 
+/*!
+ * \brief Returns the function with given properties. A null property denotes 'no change'.
+ * Returns function if all properties are unchanged. Otherwise, returns a copy with the new fields.
+ * \param function The function to copy.
+ * \param opt_params The (optional) params for the copied function. If none,
+ * ret_function->params = function->params.
+ * \param opt_body The (optional) body for the copied function. If none,
+ * ret_function->body = function->body.
+ * \param opt_ret_type The (optional) return type for the copied function. If none,
+ * ret_function->ret_type = function->ret_type.
+ * \param opt_ty_params The (optional) type params for the copied function. If none,
+ * ret_function->type_params = function->type_params.
+ * \param opt_attrs
+ * The (optional) attributes for the copied function. If none,
+ * ret_function->attrs = function->attrs.
+ * \param opt_span The (optional) span for the copied function. If none,
+ * ret_function->span = function->span.
+ * \return If all properties are null or the same as the property in the input function
+ * (i.e., opt_params is null or opt_params.value() == function->params, etc.), then we return
+ * function. Otherwise, we return a copy of function with the different fields overwritten. (i.e.,
+ * if opt_params.value() != function->params, then ret_function->params = opt_params.value()).
+ */
+Function WithFields(Function function, Optional<Array<Var>> opt_params = Optional<Array<Var>>(),
+                    Optional<Expr> opt_body = Optional<Expr>(),
+                    Optional<Type> opt_ret_type = Optional<Type>(),
+                    Optional<Array<TypeVar>> opt_ty_params = Optional<Array<TypeVar>>(),
+                    Optional<DictAttrs> opt_attrs = Optional<DictAttrs>(),
+                    Optional<Span> opt_span = Optional<Span>());
+
 /*!
  * \brief namespace of the attributes that can be attached to a relay::Function.
  */

src/relay/ir/adt.cc

@@ -104,6 +104,20 @@ Clause::Clause(Pattern lhs, Expr rhs) {
   data_ = std::move(n);
 }
 
+Clause WithFields(Clause clause, Optional<Pattern> opt_lhs, Optional<Expr> opt_rhs) {
+  Pattern lhs = opt_lhs.value_or(clause->lhs);
+  Expr rhs = opt_rhs.value_or(clause->rhs);
+
+  bool unchanged = lhs.same_as(clause->lhs) && rhs.same_as(clause->rhs);
+
+  if (!unchanged) {
+    ClauseNode* cow_clause_node = clause.CopyOnWrite();
+    cow_clause_node->lhs = lhs;
+    cow_clause_node->rhs = rhs;
+  }
+  return std::move(clause);
+}
+
 TVM_REGISTER_NODE_TYPE(ClauseNode);
 
 TVM_REGISTER_GLOBAL("relay.ir.Clause").set_body_typed([](Pattern lhs, Expr rhs) {
@@ -125,6 +139,38 @@ Match::Match(Expr data, tvm::Array<Clause> clauses, bool complete, Span span) {
   data_ = std::move(n);
 }
 
+Match WithFields(Match match, Optional<Expr> opt_data, Optional<Array<Clause>> opt_clauses,
+                 Optional<Bool> opt_complete, Optional<Span> opt_span) {
+  Expr data = opt_data.value_or(match->data);
+  Array<Clause> clauses = opt_clauses.value_or(match->clauses);
+  Bool complete = opt_complete.value_or(Bool(match->complete));
+  Span span = opt_span.value_or(match->span);
+
+  bool unchanged =
+      data.same_as(match->data) && (complete == match->complete) && span.same_as(match->span);
+
+  // Check that all clauses are unchanged
+  if (unchanged) {
+    bool all_clauses_unchanged = true;
+    if (clauses.size() == match->clauses.size()) {
+      for (size_t i = 0; i < clauses.size(); i++) {
+        all_clauses_unchanged &= clauses[i].same_as(match->clauses[i]);
+      }
+    } else {
+      all_clauses_unchanged = false;
+    }
+    unchanged &= all_clauses_unchanged;
+  }
+  if (!unchanged) {
+    MatchNode* cow_match_node = match.CopyOnWrite();
+    cow_match_node->data = data;
+    cow_match_node->clauses = clauses;
+    cow_match_node->complete = complete;
+    cow_match_node->span = span;
+  }
+  return std::move(match);
+}
+
 TVM_REGISTER_NODE_TYPE(MatchNode);
 
 TVM_REGISTER_GLOBAL("relay.ir.Match")