[Relay][Text Format] Text Printer Refactor and Debug Printing (#2605)

include/tvm/relay/expr.h

@@ -551,10 +551,9 @@ inline const TTypeNode* ExprNode::type_as() const {
  *        additional comment block to an expr.
  * \return The text representation.
  */
-std::string RelayPrint(
-    const NodeRef& node,
-    bool show_meta_data = true,
-    runtime::TypedPackedFunc<std::string(Expr)> annotate = nullptr);
+std::string RelayPrint(const NodeRef& node,
+                       bool show_meta_data = true,
+                       runtime::TypedPackedFunc<std::string(Expr)> annotate = nullptr);
 }  // namespace relay
 }  // namespace tvm
 #endif  // TVM_RELAY_EXPR_H_

python/tvm/relay/_parser.py

@@ -512,6 +512,9 @@ __source_name_counter__ = 0
 def fromtext(data, source_name=None):
     # type: (str, str) -> Union[expr.Expr, module.Module]
     """Parse a Relay program."""
+    if data == "":
+        raise ParseError("Cannot parse the empty string.")
+
     global __source_name_counter__
 
     if source_name is None:

python/tvm/relay/base.py

@@ -53,7 +53,7 @@ class RelayNode(NodeBase):
 
         Note
         ----
-        The metadata section is necessary to fully parse the text format.
+        The meta data section is necessary to fully parse the text format.
         However, it can contain dumps that are big (e.g constant weights),
         so it can be helpful to skip printing the meta data section.
 

python/tvm/relay/ir_pass.py

@@ -905,3 +905,35 @@ def eliminate_common_subexpr(expr, fskip=None):
       The output expression.
     """
     return _ir_pass.eliminate_common_subexpr(expr, fskip)
+
+
+def pass_debug_print(ast, show_meta_data=True, annotate=None, gnf=True):
+    """
+    THIS SHOULD BE USED ONLY FOR DEBUGGING, NOT AS AN INTERCHANGE FORMAT!
+    USE `.astext()` INSTEAD!
+
+    A version of the pretty printer intended for debugging passes. Contains
+    advanced printing options.
+
+    Parameters
+    ----------
+    ast : Union[relay.Expr, relay.Module, relay.Type]
+        The relay fragment to be turned into text.
+
+    show_meta_data : bool
+        Whether to include meta data section in the text
+        if there is meta data.
+
+    annotate: Optional[relay.Expr->str]
+        Optional annotate function to provide additional
+        information in the comment block.
+
+    gnf : bool
+        Whether to print in GNF. If it is disabled, pointers are left implicit.
+
+    Returns
+    -------
+    text : str
+        A text representation of `ast`.
+    """
+    return _ir_pass.pass_debug_print(ast, show_meta_data, annotate, gnf)

src/relay/ir/doc.cc

+/*!
+ *  Copyright (c) 2019 by Contributors
+ * \file src/tvm/relay/doc.cc
+ * \brief Doc ADT used for pretty printing.
+ * Based on Section 1 of https://homepages.inf.ed.ac.uk/wadler/papers/prettier/prettier.pdf.
+ */
+#include <memory>
+#include <vector>
+#include "doc.h"
+
+namespace tvm {
+namespace relay {
+
+// Text constructor
+DocAtom Text(const std::string& str) {
+  return std::make_shared<TextNode>(str);
+}
+
+// Line constructor
+DocAtom Line(int indent = 0) {
+  return std::make_shared<LineNode>(indent);
+}
+
+Doc::Doc(const std::string& str) {
+  if (str == "\n") {
+    this->stream_ = {Line()};
+  } else {
+    this->stream_ = {Text(str)};
+  }
+}
+
+// DSL function implementations
+
+Doc& Doc::operator<<(const Doc& right) {
+  assert(this != &right);
+  this->stream_.insert(this->stream_.end(), right.stream_.begin(), right.stream_.end());
+  return *this;
+}
+
+Doc& Doc::operator<<(const std::string& right) {
+  return *this << Doc(right);
+}
+
+Doc Indent(int indent, const Doc& doc) {
+  Doc ret;
+  for (auto atom : doc.stream_) {
+    if (auto text = std::dynamic_pointer_cast<TextNode>(atom)) {
+      ret.stream_.push_back(text);
+    } else if (auto line = std::dynamic_pointer_cast<LineNode>(atom)) {
+      ret.stream_.push_back(Line(indent + line->indent));
+    } else {assert(false);}
+  }
+  return ret;
+}
+
+std::string Doc::str() {
+  std::ostringstream os;
+  for (auto atom : this->stream_) {
+    if (auto text = std::dynamic_pointer_cast<TextNode>(atom)) {
+      os << text->str;
+    } else if (auto line = std::dynamic_pointer_cast<LineNode>(atom)) {
+      os << "\n" << std::string(line->indent, ' ');
+    } else {assert(false);}
+  }
+  return os.str();
+}
+
+Doc PrintVec(const std::vector<Doc>& vec, const Doc& sep) {
+  Doc seq;
+  if (vec.size() != 0) {
+    seq = vec[0];
+    for (size_t i = 1; i < vec.size(); i++) {
+      seq << sep << vec[i];
+    }
+  }
+  return seq;
+}
+
+Doc PrintBool(bool value) {
+  if (value) {
+    return Doc("True");
+  } else {
+    return Doc("False");
+  }
+}
+
+Doc PrintDType(DataType dtype) {
+  return Doc(runtime::TVMType2String(Type2TVMType(dtype)));
+}
+
+Doc PrintString(const std::string& value) {
+  // TODO(M.K.): add escape.
+  Doc doc;
+  return doc << "\"" << value << "\"";
+}
+
+}  // namespace relay
+}  // namespace tvm

src/relay/ir/doc.h

+/*!
+ *  Copyright (c) 2019 by Contributors
+ * \file tvm/relay/doc.h
+ * \brief Doc ADT used for pretty printing.
+ *  Based on Section 1 of
+ *  https://homepages.inf.ed.ac.uk/wadler/papers/prettier/prettier.pdf, but with
+ *  a vector instead of an implicitly linked list.
+ */
+#ifndef TVM_RELAY_IR_DOC_H_
+#define TVM_RELAY_IR_DOC_H_
+
+#include <tvm/relay/expr.h>
+#include <memory>
+#include <string>
+#include <vector>
+
+namespace tvm {
+namespace relay {
+
+// Doc Atom ADT
+struct DocAtomNode {
+  virtual ~DocAtomNode() = default;
+};
+
+using DocAtom = std::shared_ptr<DocAtomNode>;
+
+struct TextNode : DocAtomNode {
+  std::string str;
+
+  explicit TextNode(const std::string& str) : str(str) {}
+};
+
+struct LineNode : DocAtomNode {
+  int indent;
+
+  explicit LineNode(int indent) : indent(indent) {}
+};
+
+// Doc is a stream-like interface
+class Doc {
+ public:
+  Doc() {}
+  explicit Doc(const std::string& str);
+
+  // Append right to this.
+  Doc& operator<<(const Doc& right);
+  // Like above, but automatically lifts string to a Doc.
+  Doc& operator<<(const std::string& right);
+  // Like above, but converts right to a string first.
+  template<typename T>
+  Doc& operator<<(const T& right) {
+    std::ostringstream os;
+    os << right;
+    return *this << os.str();
+  }
+
+  // Indent a doc stream.
+  friend Doc Indent(int indent, const Doc& doc);
+
+  // Wadler's `layout`
+  std::string str();
+
+ private:
+  std::vector<DocAtom> stream_;
+};
+
+// DSL functions
+
+// Render vectors of docs with a separator. e.g. PrintVec([1, 2, 3], f) -> 1f2f3
+Doc PrintVec(const std::vector<Doc>& vec, const Doc& sep = Doc(", "));
+// Print a constant bool value.
+Doc PrintBool(bool value);
+// Print a data type.
+Doc PrintDType(DataType dtype);
+// Print a string.
+Doc PrintString(const std::string& value);
+/*!
+ * \brief special method to print out const scalar
+ * \param dtype The data type
+ * \param data The pointer to hold the data.
+ */
+template<typename T>
+Doc PrintConstScalar(DataType dtype, const T* data) {
+  std::ostringstream os;
+  if (dtype == Int(32)) {
+    os << data[0];
+  } else if (dtype == Float(32)) {
+    os << data[0] << 'f';
+  } else if (dtype == Bool()) {
+      return PrintBool(data[0] != 0);
+  } else {
+    os << dtype << "(" << data[0] << ")";
+  }
+  return Doc(os.str());
+}
+
+}  // namespace relay
+}  // namespace tvm
+#endif  // TVM_RELAY_IR_DOC_H_

tests/python/relay/test_ir_text_printer.py

@@ -33,8 +33,8 @@ def test_env():
     text = env.astext()
     assert "def @myf" in text
     assert "def @myf" in str(env)
-    assert "%1 = add(%0, %0) # ty=float32" in text
-    assert "%1 = add(%0, %0) # ty=float32" in str(env)
+    assert "%1 = add(%0, %0) // ty=float32" in text
+    assert "%1 = add(%0, %0) // ty=float32" in str(env)
     show(env.astext(annotate=lambda x: str(x.checked_type.dtype)))
     show(text)
 
@@ -95,7 +95,7 @@ def test_let_if_scope():
 
     f = relay.Function([x, y, cond], result)
     text = f.astext()
-    assert text.count("{") == 4
+    assert text.count("{") == 6
     assert "%cond: bool" in text
     show(f.astext())
 

tests/python/relay/test_type_infer.py

@@ -28,7 +28,7 @@ def initialize_box_adt(mod):
 
 
 def test_monomorphic_let():
-    "Program: let x = 1; return x"
+    "Program: let x = 1; x"
     sb = relay.ScopeBuilder()
     x = sb.let('x', relay.const(1.0, "float64"))
     sb.ret(x)
@@ -48,7 +48,7 @@ def test_add_broadcast_op():
     """
     Program:
         fn (x: Tensor[(10, 4), f32], y: Tensor[(5, 10, 1), f32]) -> Tensor[(5, 10, 4), f32] {
-            return x + y;
+            x + y
         }
     """
     x = relay.var('x', shape=(10, 4))
@@ -67,7 +67,7 @@ def test_dual_op():
        fn (x : Tensor[f32, (10, 10)]) {
          let t1 = log(x);
          let t2 = add(t1, x);
-         return t1;
+         t1
        }
     """
     tp = relay.TensorType((10, 10), "float32")
@@ -84,7 +84,7 @@ def test_dual_op():
 def test_decl():
     """Program:
        def f(x : Tensor[(10, 10), f32]) {
-           return log(x);
+           log(x)
        }
     """
     tp = relay.TensorType((10, 10))
@@ -99,9 +99,9 @@ def test_recursion():
     Program:
        def f(n: i32, data: f32) -> f32 {
           if (n == 0) {
-              return data;
+              data
           } else {
-              return f(n - 1, log(data));
+              f(n - 1, log(data))
           }
        }
     """