[SpanFillingCommonAPI]

- Change based on comment
- Discard the change of pretty-print
- Add document to set_span

Author: Joey Tsai

python/tvm/relay/frontend/common.py

@@ -313,10 +313,9 @@ def new_const(self, value, shape=None, dtype="float32", source_name=None):
             shape = value.shape
         self.const_ctr += 1
         self.params[name] = value
-        tmp_var = _expr.var(name_hint=name, shape=shape, dtype=dtype)
+        self.exprs[name] = _expr.var(name_hint=name, shape=shape, dtype=dtype)
         if source_name:
-            tmp_var = set_span(tmp_var, source_name)
-        self.exprs[name] = tmp_var
+            self.exprs[name] = set_span(self.exprs[name], source_name)
         return self.exprs[name]
 
     def get_expr(self, name):
@@ -1130,7 +1129,39 @@ def _should_fill_span():
 
 
 def set_span(sym, span):
-    """Set up the sapn of relay expression(s) while converting OP"""
+    """
+    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.
+
+    Parameters
+    ----------
+    sym :
+        A symbol is generated from the conversion of a frontend operator. Raise an error when the
+        type of the symbol is not supported.
+
+    span : String, Span, or bytes
+        The source information of the corresponding symbol.
+
+    Returns
+    -------
+    result :
+        The symbol tagged with span.
+
+    Examples
+    --------
+    .. code-block:: python
+
+      x = set_span(relay.var("x", shape=(1, 64, 56, 56)), "x_var")
+      w = relay.const(np.ones([64, 64, 3, 3]), dtype="int64")
+      y = set_span(
+          relay.nn.conv2d(x, w, channels=64, kernel_size=(3, 3), padding=(1, 1)), "conv2d"
+      )
+      print(relay.Function([x], y))
+
+      #fn (%x: Tensor[(1, 64, 56, 56), float32] /* span=x_var:0:0 */) {
+      #  nn.conv2d(%x, meta[relay.Constant][0] /* span=conv2d:0:0 */, ...) /* span=conv2d:0:0 */
+      #}
+    """
 
     if _should_fill_span():
         return _SpanFiller(span).fill(sym)

src/printer/relay_text_printer.cc

@@ -65,20 +65,22 @@ Doc RelayTextPrinter::PrintOptionalInfo(const Expr& expr) {
   if (annotate_ == nullptr) {
     if ((expr.as<ConstantNode>() || expr.as<CallNode>() || expr.as<VarNode>() ||
          expr.as<FunctionNode>() || expr.as<TupleNode>() || expr.as<TupleGetItemNode>()) &&
-        expr->checked_type_.defined()) {
-      doc << " /* ty=" << Print(expr->checked_type()) << " */";
+        (expr->checked_type_.defined() || expr->span.defined())) {
+      doc << " /*";
+      if (expr->checked_type_.defined()) {
+        doc << " ty=" << Print(expr->checked_type());
+      }
+      if (expr->span.defined()) {
+        doc << " span=" << PrintSpan(expr->span);
+      }
+      doc << " */";
     }
   } else {
     std::string annotated_expr = annotate_(expr);
     if (annotated_expr != "") {
       doc << annotated_expr;
     }
   }
-
-  if (expr->span.defined()) {
-    doc << " /* si=" << Print(expr->span) << " */";
-  }
-
   return doc;
 }
 
@@ -130,10 +132,6 @@ Doc RelayTextPrinter::Print(const ObjectRef& node, bool meta, bool try_inline) {
     return PrintPattern(Downcast<Pattern>(node), meta);
   } else if (node.as<IRModuleNode>()) {
     return PrintMod(Downcast<IRModule>(node));
-  } else if (node.as<SpanNode>()) {
-    std::ostringstream os;
-    os << Downcast<Span>(node);
-    return Doc::RawText(os.str());
   } else {
     // default module.
     std::ostringstream os;
@@ -964,6 +962,14 @@ Doc RelayTextPrinter::PrintMapAsAttributeValue(const Map<ObjectRef, ObjectRef>&
   return doc;
 }
 
+Doc RelayTextPrinter::PrintSpan(const Span& span) {
+  Doc doc;
+  const auto* span_node = span.as<SpanNode>();
+  ICHECK(span_node);
+  doc << span_node->source_name->name << ":" << span_node->line << ":" << span_node->column;
+  return doc;
+}
+
 TVM_REGISTER_GLOBAL("ir.TextPrinter").set_body_typed([](ObjectRef node) {
   auto text = AsText(node, false, nullptr);
   return text;

src/printer/text_printer.h

@@ -113,6 +113,8 @@ class RelayTextPrinter : public ExprFunctor<Doc(const Expr&)>,
    */
   Doc PrintMapAsAttributeValue(const Map<ObjectRef, ObjectRef>& map);
 
+  Doc PrintSpan(const Span& span);
+
   Doc Print(const ObjectRef& node, bool meta = false, bool try_inline = false);
 
   Doc TempVar(int n);
@@ -470,25 +472,4 @@ class TextPrinter {
 };
 }  // namespace tvm
 
-namespace tvm {
-namespace runtime {
-
-inline std::ostream& operator<<(std::ostream& os, const SourceName& source_name) {  // NOLINT(*)
-  ICHECK(source_name->name.defined());
-  os << source_name->name;
-  return os;
-}
-
-inline std::ostream& operator<<(std::ostream& os, const Span& span) {  // NOLINT(*)
-  if (span.defined()) {
-    os << span->source_name;
-  } else {
-    os << "nullptr";
-  }
-  return os;
-}
-
-}  // namespace runtime
-}  // namespace tvm
-
 #endif  // TVM_PRINTER_TEXT_PRINTER_H_