[SpanFillingCommonAPI]

- Expose Relay Expr WithFields APIs to python side
- Change the APIs in _SpanFiller from creating new instance to
  WithFields.
- Change the control of frontend span filler from env var to the
  passcontext config

Author: Joey Tsai

python/tvm/relay/expr.py

@@ -180,6 +180,21 @@ def __init__(self, data, span=None):
         self.__init_handle_by_constructor__(_ffi_api.Constant, data, span)
 
 
+@tvm._ffi.register_func("relay.ConstantWithFields")
+def ConstantWithFields(
+    constant,
+    data=None,
+    virtual_device=None,
+    span=None,
+):
+    """
+    Returns constant with the given properties. A None property denotes 'no change'.
+    Returns constant if all properties are unchanged. Otherwise, returns a copy with the new
+    fields.
+    """
+    return _ffi_api.ConstantWithFields(constant, data, virtual_device, span)
+
+
 @tvm._ffi.register_object("relay.Tuple")
 class Tuple(ExprWithOp):
     """Tuple expression that groups several fields together.
@@ -208,6 +223,16 @@ def astype(self, _):
         raise TypeError("astype cannot be used on tuple")
 
 
+@tvm._ffi.register_func("relay.TupleWithFields")
+def TupleWithFields(tup, fields=None, virtual_device=None, span=None):
+    """
+    Returns tuple with the given properties. A None property denotes 'no change'.
+    Returns tuple if all properties are unchanged. Otherwise, returns a copy with the new
+    fields.
+    """
+    return _ffi_api.TupleWithFields(tup, fields, virtual_device, span)
+
+
 @tvm._ffi.register_object("relay.Var")
 class Var(ExprWithOp):
     """A local variable in Relay.
@@ -239,6 +264,16 @@ def name_hint(self):
         return name
 
 
+@tvm._ffi.register_func("relay.VarWithFields")
+def VarWithFields(variable, vid=None, type_annotation=None, virtual_device=None, span=None):
+    """
+    Returns var with the given properties. A None property denotes 'no change'.
+    Returns var if all properties are unchanged. Otherwise, returns a copy with the new
+    fields.
+    """
+    return _ffi_api.VarWithFields(variable, vid, type_annotation, virtual_device, span)
+
+
 @tvm._ffi.register_object("relay.Call")
 class Call(ExprWithOp):
     """Function call node in Relay.
@@ -271,6 +306,18 @@ def __init__(self, op, args, attrs=None, type_args=None, span=None):
         self.__init_handle_by_constructor__(_ffi_api.Call, op, args, attrs, type_args, span)
 
 
+@tvm._ffi.register_func("relay.CallWithFields")
+def CallWithFields(
+    call, op=None, args=None, attrs=None, type_args=None, virtual_device=None, span=None
+):
+    """
+    Returns call with the given properties. A None property denotes 'no change'.
+    Returns call if all properties are unchanged. Otherwise, returns a copy with the new
+    fields.
+    """
+    return _ffi_api.CallWithFields(call, op, args, attrs, type_args, virtual_device, span)
+
+
 @tvm._ffi.register_object("relay.Let")
 class Let(ExprWithOp):
     """Let variable binding expression.
@@ -294,6 +341,16 @@ def __init__(self, variable, value, body, span=None):
         self.__init_handle_by_constructor__(_ffi_api.Let, variable, value, body, span)
 
 
+@tvm._ffi.register_func("relay.LetWithFields")
+def LetWithFields(let, variable=None, value=None, body=None, virtual_device=None, span=None):
+    """
+    Returns let with the given properties. A None property denotes 'no change'.
+    Returns let if all properties are unchanged. Otherwise, returns a copy with the new
+    fields.
+    """
+    return _ffi_api.LetWithFields(let, variable, value, body, virtual_device, span)
+
+
 @tvm._ffi.register_object("relay.If")
 class If(ExprWithOp):
     """A conditional expression in Relay.
@@ -317,6 +374,18 @@ def __init__(self, cond, true_branch, false_branch, span=None):
         self.__init_handle_by_constructor__(_ffi_api.If, cond, true_branch, false_branch, span)
 
 
+@tvm._ffi.register_func("relay.IfWithFields")
+def IfWithFields(
+    if_expr, cond=None, true_branch=None, false_branch=None, virtual_device=None, span=None
+):
+    """
+    Returns if with the given properties. A None property denotes 'no change'.
+    Returns if if all properties are unchanged. Otherwise, returns a copy with the new
+    fields.
+    """
+    return _ffi_api.IfWithFields(if_expr, cond, true_branch, false_branch, virtual_device, span)
+
+
 @tvm._ffi.register_object("relay.TupleGetItem")
 class TupleGetItem(ExprWithOp):
     """Get index-th item from a tuple.
@@ -337,6 +406,18 @@ def __init__(self, tuple_value, index, span=None):
         self.__init_handle_by_constructor__(_ffi_api.TupleGetItem, tuple_value, index, span)
 
 
+@tvm._ffi.register_func("relay.TupleGetItemWithFields")
+def TupleGetItemWithFields(
+    tuple_get_item, tuple_value=None, index=None, virtual_device=None, span=None
+):
+    """
+    Returns tuple_get_item with the given properties. A None property denotes 'no change'.
+    Returns tuple_get_item if all properties are unchanged. Otherwise, returns a copy with the new
+    fields.
+    """
+    return _ffi_api.TupleGetItemWithFields(tuple_get_item, tuple_value, index, virtual_device, span)
+
+
 @tvm._ffi.register_object("relay.RefCreate")
 class RefCreate(ExprWithOp):
     """Create a new reference from initial value.
@@ -353,6 +434,21 @@ def __init__(self, value, span=None):
         self.__init_handle_by_constructor__(_ffi_api.RefCreate, value, span)
 
 
+@tvm._ffi.register_func("relay.RefCreateWithFields")
+def RefCreateWithFields(
+    ref_create,
+    value=None,
+    virtual_device=None,
+    span=None,
+):
+    """
+    Returns ref_create with the given properties. A None property denotes 'no change'.
+    Returns ref_create if all properties are unchanged. Otherwise, returns a copy with the new
+    fields.
+    """
+    return _ffi_api.RefCreateWithFields(ref_create, value, virtual_device, span)
+
+
 @tvm._ffi.register_object("relay.RefRead")
 class RefRead(ExprWithOp):
     """Get the value inside the reference.
@@ -369,6 +465,21 @@ def __init__(self, ref, span=None):
         self.__init_handle_by_constructor__(_ffi_api.RefRead, ref, span)
 
 
+@tvm._ffi.register_func("relay.RefReadWithFields")
+def RefReadWithFields(
+    ref_read,
+    ref=None,
+    virtual_device=None,
+    span=None,
+):
+    """
+    Returns ref_read with the given properties. A None property denotes 'no change'.
+    Returns ref_read if all properties are unchanged. Otherwise, returns a copy with the new
+    fields.
+    """
+    return _ffi_api.RefReadWithFields(ref_read, ref, virtual_device, span)
+
+
 @tvm._ffi.register_object("relay.RefWrite")
 class RefWrite(ExprWithOp):
     """
@@ -390,6 +501,22 @@ def __init__(self, ref, value, span=None):
         self.__init_handle_by_constructor__(_ffi_api.RefWrite, ref, value, span)
 
 
+@tvm._ffi.register_func("relay.RefWriteWithFields")
+def RefWriteWithFields(
+    ref_write,
+    ref=None,
+    value=None,
+    virtual_device=None,
+    span=None,
+):
+    """
+    Returns ref_write with the given properties. A None property denotes 'no change'.
+    Returns ref_write if all properties are unchanged. Otherwise, returns a copy with the new
+    fields.
+    """
+    return _ffi_api.RefWriteWithFields(ref_write, ref, value, virtual_device, span)
+
+
 class TempExpr(ExprWithOp):
     """Baseclass of all TempExpr.
 

python/tvm/relay/frontend/common.py

@@ -18,7 +18,6 @@
 """Common utilities"""
 from __future__ import absolute_import as _abs
 import logging
-import os
 import numpy as np
 
 import tvm
@@ -1027,42 +1026,50 @@ def visit(self, expr):
     def visit_function(self, fn):
         new_params = [self.visit(x) for x in fn.params]
         new_body = self.visit(fn.body)
-        return _function.Function(
-            list(new_params), new_body, fn.ret_type, fn.type_params, fn.attrs, self._span
+        return _function.FunctionWithFields(
+            fn, list(new_params), new_body, fn.ret_type, fn.type_params, fn.attrs, None, self._span
         )
 
     def visit_let(self, let):
         new_variable = self.visit(let.var)
         new_value = self.visit(let.value)
         new_body = self.visit(let.body)
-        return _expr.Let(new_variable, new_value, new_body, self._span)
+        return _expr.LetWithFields(let, new_variable, new_value, new_body, None, self._span)
 
     def visit_call(self, call):
         new_args = [self.visit(arg) for arg in call.args]
         # call.op might be RelayExpr or Op type
         # ExprMutator will return directly if subject belongs to Op type
         new_op = self.visit(call.op)
-        return _expr.Call(new_op, new_args, call.attrs, call.type_args, self._span)
+        return _expr.CallWithFields(
+            call, new_op, new_args, call.attrs, call.type_args, None, self._span
+        )
 
     def visit_var(self, var):
-        return _expr.Var(var.name_hint, var.type_annotation, self._span)
+        return _expr.VarWithFields(var, var.vid, var.type_annotation, None, self._span)
 
     def visit_if(self, ite):
-        return _expr.If(
+        return _expr.IfWithFields(
+            ite,
             self.visit(ite.cond),
             self.visit(ite.true_branch),
             self.visit(ite.false_branch),
+            None,
             self._span,
         )
 
     def visit_tuple(self, tup):
-        return _expr.Tuple([self.visit(field) for field in tup.fields], self._span)
+        return _expr.TupleWithFields(
+            tup, [self.visit(field) for field in tup.fields], None, self._span
+        )
 
     def visit_tuple_getitem(self, op):
-        return _expr.TupleGetItem(self.visit(op.tuple_value), op.index, self._span)
+        return _expr.TupleGetItemWithFields(
+            op, self.visit(op.tuple_value), op.index, None, self._span
+        )
 
     def visit_constant(self, const):
-        return _expr.Constant(const.data, self._span)
+        return _expr.ConstantWithFields(const, const.data, None, self._span)
 
     # TODO: Frontend model translation could not use following relay expressions so far,
     #       enable them when new models/impls leverage these kinds of relay expressions.
@@ -1115,23 +1122,10 @@ def fill(self, sym):
         raise RuntimeError(f"unsupported type {type(sym)}")
 
 
-def _should_fill_span():
-    should_fill_span = os.environ.get("TVM_SPANFILLING", "1")
-
-    try:
-        should_fill_span = bool(int(should_fill_span))
-    except ValueError:
-        raise ValueError(
-            f"invalid value for TVM_SPANFILLING {should_fill_span}, please set to 0 or 1."
-        )
-
-    return should_fill_span
-
-
 def set_span(sym, span):
     """
     Recursively tag the span to the symbol. Stop when it encounters a span-tagged expr. Disabled
-    when setting the environment variable "TVM_SPANFILLING" as 0.
+    when setting the "relay.frontend.fill_span" as False to the config of PassContext
 
     Parameters
     ----------
@@ -1163,6 +1157,6 @@ def set_span(sym, span):
       #}
     """
 
-    if _should_fill_span():
+    if tvm.transform.PassContext.current().config.get("relay.frontend.fill_span", True):
         return _SpanFiller(span).fill(sym)
     return sym

python/tvm/testing/utils.py

@@ -2085,17 +2085,14 @@ def pprint(name, obj):
 
 class _control_span_filling:
     def __init__(self, on=True):
-        self._old_state = os.environ["TVM_SPANFILLING"] if "TVM_SPANFILLING" in os.environ else None
         self._on = on
+        self._pass_ctx = tvm.transform.PassContext(config={"relay.frontend.fill_span": self._on})
 
     def __enter__(self):
-        os.environ["TVM_SPANFILLING"] = str(int(self._on))
+        self._pass_ctx.__enter__()
 
     def __exit__(self, exc_type, exc_val, exc_tb):
-        if self._old_state:
-            os.environ["TVM_SPANFILLING"] = self._old_state
-        else:
-            del os.environ["TVM_SPANFILLING"]
+        self._pass_ctx.__exit__(exc_type, exc_val, exc_tb)
 
 
 class enable_span_filling(_control_span_filling):

src/ir/span.cc

@@ -20,13 +20,17 @@
  * \file span.cc
  * \brief The span data structure.
  */
+#include <tvm/ir/expr.h>
 #include <tvm/ir/span.h>
+#include <tvm/ir/transform.h>
 #include <tvm/runtime/registry.h>
 
 #include <algorithm>
 
 namespace tvm {
 
+TVM_REGISTER_PASS_CONFIG_OPTION("relay.frontend.fill_span", Bool);
+
 ObjectPtr<Object> GetSourceNameNode(const String& name) {
   // always return pointer as the reference can change as map re-allocate.
   // or use another level of indirection by creating a unique_ptr