[tvm] branch main updated: [TVMScript] Avoid visiting repetition tensor in SetCommonPrefix Visitor (#15083)

[syfeng]

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

syfeng 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 77c75d1ef3 [TVMScript] Avoid visiting repetition tensor in 
SetCommonPrefix Visitor (#15083)
77c75d1ef3 is described below

commit 77c75d1ef3d4ad096ee79cbe07adbfc514721b63
Author: chenxinli <39092231+cccxi...@users.noreply.github.com>
AuthorDate: Sat Jun 17 15:41:50 2023 +0800

    [TVMScript] Avoid visiting repetition tensor in SetCommonPrefix Visitor 
(#15083)
    
    * [TVMScript] Avoid visiting repetition tensor in SetCommonPrefix Visitor
    
    * Switch vector to unordered_set
---
 src/script/printer/ir_docsifier.cc | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/src/script/printer/ir_docsifier.cc 
b/src/script/printer/ir_docsifier.cc
index 62084d17be..7cd27057f4 100644
--- a/src/script/printer/ir_docsifier.cc
+++ b/src/script/printer/ir_docsifier.cc
@@ -95,6 +95,13 @@ void IRDocsifierNode::SetCommonPrefix(const ObjectRef& root,
       if (obj == nullptr) {
         return;
       }
+      if (visited_.count(obj)) {
+        if (is_var(GetRef<ObjectRef>(obj))) {
+          HandleVar(obj);
+        }
+        return;
+      }
+      visited_.insert(obj);
       stack_.push_back(obj);
       if (obj->IsInstance<ArrayNode>()) {
         const ArrayNode* array = static_cast<const ArrayNode*>(obj);
@@ -134,6 +141,7 @@ void IRDocsifierNode::SetCommonPrefix(const ObjectRef& root,
 
     ReflectionVTable* vtable_ = ReflectionVTable::Global();
     std::vector<const Object*> stack_;
+    std::unordered_set<const Object*> visited_;
 
    public:
     runtime::TypedPackedFunc<bool(ObjectRef)> is_var;