[tvm] branch main updated: [IRPrinter] Prevent multiple printing of optional info (#8279)

[ma...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

masahi pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git


The following commit(s) were added to refs/heads/main by this push:
     new 1f0f8f1  [IRPrinter] Prevent multiple printing of optional info (#8279)
1f0f8f1 is described below

commit 1f0f8f19b98093332ba3253d96551b1e499bf984
Author: Huang, Guangtai <gu...@amazon.com>
AuthorDate: Sat Jun 19 05:59:28 2021 +0800

    [IRPrinter] Prevent multiple printing of optional info (#8279)
    
    * fix
    
    * test
---
 src/printer/relay_text_printer.cc          | 4 +++-
 src/printer/text_printer.h                 | 2 ++
 tests/python/relay/test_ir_text_printer.py | 9 +++++++++
 3 files changed, 14 insertions(+), 1 deletion(-)

diff --git a/src/printer/relay_text_printer.cc b/src/printer/relay_text_printer.cc
index 2de331b..aad42fc 100644
--- 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
index 52ab701..7a529cc 100644
--- 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
index 4968660..b4d02e4 100644
--- 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__])