[Relay][AST] Add WithFields for Constant and GlobalVar nodes

Co-authored-by: Mark Shields <[email protected]>

Author: Michalis Papapdimitriou

include/tvm/ir/expr.h

@@ -260,9 +260,10 @@ class GlobalVarNode : public RelayExprNode {
  */
 class GlobalVar : public RelayExpr {
  public:
-  TVM_DLL explicit GlobalVar(String name_hint, Type type = {});
+  TVM_DLL explicit GlobalVar(String name_hint, Type type = {}, Span span = {});
 
   TVM_DEFINE_OBJECT_REF_METHODS(GlobalVar, RelayExpr, GlobalVarNode);
+  + TVM_DEFINE_OBJECT_REF_COW_METHOD(GlobalVarNode);
 };
 
 // PrimExprs that are useful as runtime containers.

include/tvm/relay/expr.h

@@ -39,6 +39,22 @@
 #include "./type.h"
 
 namespace tvm {
+
+/*!
+ * \brief Returns the global_var with given properties. A null property denotes 'no change'.
+ * Returns this if all properties are unchanged. Otherwise, returns a copy with the new fields.
+ * \param global_var The global var to copy
+ * \param opt_name_hint The (optional) op name of the global var
+ * \param opt_type The (optional) type for the global var
+ * \param opt_virtual_device The (optional) virtual_device for the copied constant. If none,
+ * ret_constant->virtual_device = constant->virtual_device.
+ * \param opt_span The (optional) span for the copied global var. If none,
+ * ret_constant->span = constant->span.
+ */
+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;
@@ -97,8 +113,23 @@ class Constant : public Expr {
   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);
 };
 
+/*!
+ * \brief Returns the constant with given properties. A null property denotes 'no change'.
+ * Returns this if all properties are unchanged. Otherwise, returns a copy with the new fields.
+ * \param constant The constant to copy
+ * \param op_data The (optional) data for the copied constant. If none, ret_constant->data =
+ * constant->data.
+ * \param opt_virtual_device The (optional) virtual_device for the copied constant. If none,
+ * ret_constant->virtual_device = constant->virtual_device.
+ * \param opt_span The (optional) span for the copied constant. If none,
+ * ret_constant->span = constant->span.
+ */
+Constant WithFields(Constant constant, Optional<runtime::NDArray> opt_data = {},
+                    Optional<VirtualDevice> opt_virtual_device = {}, Optional<Span> opt_span = {});
+
 /*! \brief Tuple of multiple Exprs */
 class Tuple;
 /*! \brief Tuple container */

src/ir/expr.cc

@@ -141,10 +141,11 @@ TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable)
       p->stream << "range(min=" << op->min << ", ext=" << op->extent << ')';
     });
 
-GlobalVar::GlobalVar(String name_hint, Type type) {
+GlobalVar::GlobalVar(String name_hint, Type type, Span span) {
   ObjectPtr<GlobalVarNode> n = make_object<GlobalVarNode>();
   n->name_hint = std::move(name_hint);
   n->checked_type_ = std::move(type);
+  n->span = std::move(span);
   data_ = std::move(n);
 }
 

src/relay/ir/expr.cc

@@ -27,6 +27,26 @@
 
 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) {
+  String name_hint = opt_name_hint.value_or(global_var->name_hint);
+  Type type = opt_type.value_or(global_var->checked_type());
+  VirtualDevice virtual_device = opt_virtual_device.value_or(global_var->virtual_device());
+  Span span = opt_span.value_or(global_var->span);
+  bool all_fields_unchanged =
+      name_hint.same_as(global_var->name_hint) && type.same_as(global_var->checked_type()) &&
+      virtual_device.same_as(global_var->virtual_device()) && span.same_as(global_var->span);
+  if (!all_fields_unchanged) {
+    GlobalVarNode* cow_global_var_node = global_var.CopyOnWrite();
+    cow_global_var_node->name_hint = name_hint;
+    cow_global_var_node->checked_type_ = type;
+    cow_global_var_node->virtual_device_ = virtual_device;
+    cow_global_var_node->span = span;
+  }
+
+  return global_var;
+}
+
 VirtualDevice RelayExprNode::virtual_device() const {
   if (!this->virtual_device_.defined()) {
     // virtual_device_ should always be defined, unless we imported this node from JSON using an old
@@ -77,6 +97,25 @@ TensorType ConstantNode::tensor_type() const {
   return TensorType(shape, dtype);
 }
 
+Constant WithFields(Constant constant, Optional<runtime::NDArray> opt_data,
+                    Optional<VirtualDevice> opt_virtual_device, Optional<Span> opt_span) {
+  runtime::NDArray data = opt_data.value_or(constant->data);
+  VirtualDevice virtual_device = opt_virtual_device.value_or(constant->virtual_device());
+  Span span = opt_span.value_or(constant->span);
+
+  bool all_fields_unchanged = data.same_as(constant->data) &&
+                              virtual_device.same_as(constant->virtual_device()) &&
+                              span.same_as(constant->span);
+
+  if (!all_fields_unchanged) {
+    ConstantNode* cow_constant_node = constant.CopyOnWrite();
+    cow_constant_node->data = data;
+    cow_constant_node->virtual_device_ = virtual_device;
+    cow_constant_node->span = span;
+  }
+  return constant;
+}
+
 Tuple::Tuple(tvm::Array<relay::Expr> fields, Span span) {
   ObjectPtr<TupleNode> n = make_object<TupleNode>();
   n->fields = std::move(fields);