Revert "[Frontend] Add Span filling for frontends to Relay (apache#9723)" (apache#10072) (#246)

Because of the failure of LSTM conversion from Pytorch

Co-authored-by: Chun-I Tsai <[email protected]>

Author: Jian Sheng

python/tvm/relay/expr.py

@@ -316,13 +316,10 @@ class TupleGetItem(ExprWithOp):
 
     index: int
         The index.
-
-    span: Optional[tvm.relay.Span]
-        Span that points to original source code
     """
 
-    def __init__(self, tuple_value, index, span=None):
-        self.__init_handle_by_constructor__(_ffi_api.TupleGetItem, tuple_value, index, span)
+    def __init__(self, tuple_value, index):
+        self.__init_handle_by_constructor__(_ffi_api.TupleGetItem, tuple_value, index)
 
 
 @tvm._ffi.register_object("relay.RefCreate")

python/tvm/relay/frontend/common.py

@@ -25,7 +25,6 @@
 from tvm.topi.utils import get_const_tuple
 
 from .. import expr as _expr
-from ..expr_functor import ExprMutator
 from .. import function as _function
 from .. import transform as _transform
 from .. import op as _op
@@ -996,6 +995,7 @@ def visit_tuple(self, tup):
         def visit_tuple_getitem(self, op):
             if op.span is None:
                 self.distance_from_leaf += 1
+                # pylint: disable=too-many-function-args
                 return _expr.TupleGetItem(self.visit(op.tuple_value), op.index, self._create_span())
             return op
 

python/tvm/relay/frontend/pytorch.py

@@ -45,7 +45,6 @@
 from .common import infer_value as _infer_value
 from .common import infer_value_simulated as _infer_value_simulated
 from .common import lstm_cell, try_infer_value, unbind
-from .common import set_span
 from .pytorch_utils import is_version_greater_than
 
 __all__ = ["from_pytorch"]
@@ -3257,9 +3256,6 @@ def body(*current_vals):
 
     def convert_operators(self, operators, outputs, ret_names):
         """Convert each Torch IR operators to Relay equivalent"""
-        # an op node might not belong to any of scope in trace info natively
-        # use a cunter to prevent from messing up its scope in span
-        empty_counter = 0
         for node_name, op_node in operators:
             operator = op_node.kind()
             inputs = _get_op_inputs(op_node, outputs)
@@ -3309,9 +3305,6 @@ def convert_operators(self, operators, outputs, ret_names):
                 relay_out = relay_op(
                     inputs, _get_input_types(op_node, outputs, default_dtype=self.default_dtype)
                 )
-                span_str, empty_counter = self._get_torch_span(op_node, empty_counter)
-                relay_out = set_span(relay_out, span_str)
-
                 self.record_output_type(relay_out)
 
                 if isinstance(relay_out, tuple):
@@ -3325,18 +3318,6 @@ def convert_operators(self, operators, outputs, ret_names):
 
         return [_wrap_const(outputs[ret_name]) for ret_name in ret_names]
 
-    def _get_torch_span(self, node, empty_counter):
-        # torch span looks like
-        # %input.5 : Float(...) = aten::relu_(%input.3), scope: __module.relu # ${torch}/nn file
-        # the scope part might not exist
-        if node.scopeName():
-            scope_name_str = "jit._trace.TopLevelTracedModule: " + node.scopeName()
-        else:
-            scope_name_str = "warning: no trace info " + str(empty_counter)
-            empty_counter += 1
-        span_str = "C.graph: {}, {}".format(node.kind(), scope_name_str)
-        return span_str, empty_counter
-
 
 def _pytorch_result_type(dtypes, non_tensor_inputs):
     """This promotes TVM dtypes like PyTorch would"""

python/tvm/relay/frontend/tensorflow.py

@@ -37,7 +37,6 @@
 from .common import infer_type as _infer_type
 from .common import infer_shape as _infer_shape
 from .common import infer_value as _infer_value
-from .common import set_span
 
 from .tensorflow_ops import _convert_map
 from .tensorflow_ops import _need_prelude_for_shape_inference
@@ -1029,10 +1028,24 @@ def _convert_operator(
         else:
             raise NotImplementedError("Operator {} not implemented.".format(op_name))
 
-        sym = set_span(sym, node_name)
+        sym = self._set_span(sym, node_name)
 
         return sym
 
+    @staticmethod
+    def _set_span(sym, node_name):
+        span = tvm.relay.Span(tvm.relay.SourceName(node_name), 0, 0, 0, 0)
+        if isinstance(sym, _expr.Call) and sym.span is None:
+            sym = _expr.Call(sym.op, sym.args, sym.attrs, sym.type_args, span)
+        elif isinstance(sym, _expr.TupleWrapper):
+            tuple_value = sym.tuple_value
+            if isinstance(tuple_value, _expr.Call) and tuple_value.span is None:
+                tuple_value = _expr.Call(
+                    tuple_value.op, tuple_value.args, tuple_value.attrs, tuple_value.type_args, span
+                )
+                sym = _expr.TupleWrapper(tuple_value, sym.size)
+        return sym
+
     def _licm_construct(self, loop_name, node_name):
         """Construct a node by considering whether it is
         loop invariant with the given while loop. If yes, we

python/tvm/relay/frontend/tensorflow2.py

@@ -36,7 +36,6 @@
 from .. import function as _function
 from ..loops import while_loop as _while_loop
 from .common import infer_type as _infer_type
-from .common import set_span
 
 from .tensorflow_ops import _convert_map as _convert_map_common
 from .tensorflow_ops import _get_more_static_shape_rank
@@ -59,6 +58,22 @@ def _infer_type_with_prelude(val, prelude):
     return body.checked_type
 
 
+def set_span(sym, node_name):
+    """set span of symbol"""
+
+    span = tvm.relay.Span(tvm.relay.SourceName(node_name), 0, 0, 0, 0)
+    if isinstance(sym, _expr.Call):
+        sym = _expr.Call(sym.op, sym.args, sym.attrs, sym.type_args, span)
+    elif isinstance(sym, _expr.TupleWrapper):
+        tuple_value = sym.tuple_value
+        if isinstance(tuple_value, _expr.Call):
+            tuple_value = _expr.Call(
+                tuple_value.op, tuple_value.args, tuple_value.attrs, tuple_value.type_args, span
+            )
+            sym = _expr.TupleWrapper(tuple_value, sym.size)
+    return sym
+
+
 def is_tensor_list_constuctor(tf_node):
     """Check whether is tensor list constructor node."""
     return tf_node.op == "TensorListReserve"

python/tvm/relay/frontend/tflite.py

@@ -32,7 +32,6 @@
 from .. import qnn as _qnn
 from .common import ExprTable
 from .common import infer_shape as _infer_shape
-from .common import set_span
 from .common import to_int_list
 from .tflite_flexbuffer import FlexBufferDecoder
 
@@ -240,17 +239,12 @@ def convert_op_to_relay(self):
 
             if len(output_tensors) == 1:
                 tensor_idx = output_tensors[0].tensor_idx
-                curr_output = get_tensor_name(self.subgraph, tensor_idx)
-                ret = set_span(ret, "location: {}, output_name: {}".format(op_idx, curr_output))
-                self.exp_tab.set_expr(curr_output, ret)
+                self.exp_tab.set_expr(get_tensor_name(self.subgraph, tensor_idx), ret)
             else:
-                out_names = []
-                for output_tensor in output_tensors:
-                    out_names.append(get_tensor_name(self.subgraph, output_tensor.tensor_idx))
-                curr_output = ", ".join(out_names)
-                ret = set_span(ret, "location: {}, output_name: {}".format(op_idx, curr_output))
-                for idx, out_name in enumerate(out_names):
-                    self.exp_tab.set_expr(out_name, ret[idx])
+                for idx, output_tensor in enumerate(output_tensors):
+                    self.exp_tab.set_expr(
+                        get_tensor_name(self.subgraph, output_tensor.tensor_idx), ret[idx]
+                    )
 
     def get_op_code_str(self, op):
         """Get TFLite ops string representation"""

src/printer/relay_text_printer.cc

@@ -389,21 +389,12 @@ Doc RelayTextPrinter::VisitExpr_(const TupleNode* op) {
   if (op->fields.size() == 1) {
     doc << ",";
   }
-  doc << ")";
-  if (op->span.defined()) {
-    doc << " /* " << PrintSpan(op->span) << " */";
-  }
-  return doc;
+  return doc << ")";
 }
 
 Doc RelayTextPrinter::VisitExpr_(const TupleGetItemNode* op) {
   Doc doc;
-  doc << Print(op->tuple) << "." << op->index;
-
-  if (op->span.defined()) {
-    doc << " /* " << PrintSpan(op->span) << " */";
-  }
-  return doc;
+  return doc << Print(op->tuple) << "." << op->index;
 }
 
 Doc RelayTextPrinter::VisitExpr_(const IfNode* op) {
@@ -977,13 +968,11 @@ Doc RelayTextPrinter::PrintMapAsAttributeValue(const Map<ObjectRef, ObjectRef>&
   return doc;
 }
 
-Doc RelayTextPrinter::PrintSpan(const Span& span, bool include_spans) {
+Doc RelayTextPrinter::PrintSpan(const Span& span) {
   Doc doc;
-  if (include_spans) {
-    const auto* span_node = span.as<SpanNode>();
-    ICHECK(span_node);
-    doc << span_node->source_name->name;
-  }
+  const auto* span_node = span.as<SpanNode>();
+  ICHECK(span_node);
+  doc << span_node->source_name->name;
   return doc;
 }
 

src/printer/text_printer.h

@@ -113,7 +113,7 @@ class RelayTextPrinter : public ExprFunctor<Doc(const Expr&)>,
    */
   Doc PrintMapAsAttributeValue(const Map<ObjectRef, ObjectRef>& map);
 
-  Doc PrintSpan(const Span& span, bool include_spans = true);
+  Doc PrintSpan(const Span& span);
 
   Doc Print(const ObjectRef& node, bool meta = false, bool try_inline = false);
 

src/relay/ir/expr.cc

@@ -362,8 +362,8 @@ TupleGetItem WithFields(TupleGetItem tuple_get_item, Optional<Expr> opt_tuple,
 
 TVM_REGISTER_NODE_TYPE(TupleGetItemNode);
 
-TVM_REGISTER_GLOBAL("relay.ir.TupleGetItem").set_body_typed([](Expr tuple, int index, Span span) {
-  return TupleGetItem(tuple, index, span);
+TVM_REGISTER_GLOBAL("relay.ir.TupleGetItem").set_body_typed([](Expr tuple, int index) {
+  return TupleGetItem(tuple, index);
 });
 
 TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable)

tests/python/frontend/pytorch/test_forward.py

@@ -247,53 +247,6 @@ def visit(op):
     torch.cuda.empty_cache()
 
 
-def verify_span(model_name, input_data=[], custom_convert_map={}):
-    if isinstance(model_name, str):
-        baseline_model, baseline_input = load_model(model_name)
-    elif isinstance(input_data, list):
-        baseline_model = model_name
-        baseline_input = input_data
-    elif isinstance(input_data, torch.Tensor) or len(input_data.shape) == 0:
-        baseline_model = model_name
-        baseline_input = [input_data]
-    else:
-        assert False, "Unexpected input format"
-
-    trace = torch.jit.trace(baseline_model, [input.clone() for input in baseline_input])
-    if isinstance(baseline_model, torch.nn.Module):
-        trace = trace.float().eval()
-
-        if torch.cuda.is_available():
-            trace = trace.cuda()
-        else:
-            trace = trace.cpu()
-
-    input_names = ["input{}".format(idx) for idx, inp in enumerate(baseline_input)]
-    input_shapes = list(zip(input_names, [inp.shape for inp in baseline_input]))
-    mod, params = relay.frontend.from_pytorch(trace, input_shapes, custom_convert_map)
-
-    # collect fail cases for the convenience of further improvement
-    fail_cases = []
-    mod_main_start = False
-    for line in str(mod.__str__).split("\n"):
-        if "@main" in line:
-            mod_main_start = True
-            continue
-
-        if mod_main_start == True:
-            if "}" == line:
-                break
-            elif not ("/*" in line and "*/" in line):
-                fail_cases.append(line)
-
-    print(fail_cases)
-    assert len(fail_cases) == 0
-
-
-def test_span():
-    verify_span("resnet18")
-
-
 # Single operator tests
 @tvm.testing.uses_gpu
 def test_forward_pixel_shuffle():