[IRPrinter] Prevent multiple printing of optional info (#8279)

* fix * test
apache · Jun 18, 2021 · 1f0f8f1 · 1f0f8f1
1 parent edb7e77
commit 1f0f8f1
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 1 deletion.
diff --git a/src/printer/relay_text_printer.cc b/src/printer/relay_text_printer.cc
@@ -54,6 +54,9 @@ namespace relay {
  */
 Doc RelayTextPrinter::PrintOptionalInfo(const Expr& expr) {
   Doc doc;
+  if (!opt_info_memo_.insert(expr).second) {
+    return doc;
+  }
   // default annotations
   if (annotate_ == nullptr) {
     if ((expr.as<ConstantNode>() || expr.as<CallNode>()) && expr->checked_type_.defined()) {
@@ -65,7 +68,6 @@ Doc RelayTextPrinter::PrintOptionalInfo(const Expr& expr) {
       doc << annotated_expr;
     }
   }
-
   return doc;
 }
 

diff --git a/src/printer/text_printer.h b/src/printer/text_printer.h
@@ -176,6 +176,8 @@ class RelayTextPrinter : public ExprFunctor<Doc(const Expr&)>,
   std::vector<Doc> doc_stack_{};
   /*! \brief Set for introduced vars */
   std::unordered_set<Expr, ObjectPtrHash, ObjectPtrEqual> var_memo_;
+  /*! \brief Set for exprs have been printed optional information */
+  std::unordered_set<Expr, ObjectPtrHash, ObjectPtrEqual> opt_info_memo_;
   /*! \brief Map for result and memo_ diffs for visited expression */
   std::unordered_map<Expr, Doc, ObjectPtrHash, ObjectPtrEqual> result_memo_;
   /*! \brief Map from Expr to Doc */

diff --git a/tests/python/relay/test_ir_text_printer.py b/tests/python/relay/test_ir_text_printer.py
@@ -276,5 +276,14 @@ def test_span():
     assert "Add1" in txt
 
 
+def test_optional_info():
+    c = relay.const(1)
+    call = relay.add(c, c)
+    m = tvm.IRModule.from_expr(call)
+    m = relay.transform.InferType()(m)
+    txt = astext(m)
+    assert txt.count("/* ty=int32 */") == 3
+
+
 if __name__ == "__main__":
     pytest.main([__file__])