From e5985ff9eb18d0bc90a8468f69828579a6b8902f Mon Sep 17 00:00:00 2001 From: Yaxing Cai Date: Mon, 23 Jan 2023 19:34:34 -0800 Subject: [PATCH] allocate syntax sugar improvement --- src/script/printer/tir/stmt.cc | 4 +--- tests/python/unittest/test_tvmscript_printer_tir.py | 7 +++---- 2 files changed, 4 insertions(+), 7 deletions(-) diff --git a/src/script/printer/tir/stmt.cc b/src/script/printer/tir/stmt.cc index 7c8d44c10e72..acdfd7da472b 100644 --- a/src/script/printer/tir/stmt.cc +++ b/src/script/printer/tir/stmt.cc @@ -175,9 +175,7 @@ TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable) .set_dispatch( // "", [](tir::Allocate stmt, ObjectPath stmt_p, IRDocsifier d) -> Doc { bool concise = AllowConciseScoping(d); - OccurrenceCounter counter(stmt->buffer_var.get()); - counter(stmt->body); - if (counter.count == 1 && IsAllocateDeclBufferPattern(stmt.get())) { + if (IsAllocateDeclBufferPattern(stmt.get())) { return d->AsDoc(stmt->body, stmt_p->Attr("body")); } Array args; diff --git a/tests/python/unittest/test_tvmscript_printer_tir.py b/tests/python/unittest/test_tvmscript_printer_tir.py index d57d10467077..c73ae291930c 100644 --- a/tests/python/unittest/test_tvmscript_printer_tir.py +++ b/tests/python/unittest/test_tvmscript_printer_tir.py @@ -343,7 +343,7 @@ def test_allocate_with_decl_buffer_sugar(): ) -def test_allocate_with_decl_buffer_no_sugar_multi_usage(): +def test_allocate_with_decl_buffer_sugar_multi_usage(): with IRBuilder() as ib: with T.allocate([128, 128], "float32") as buffer_data: with T.decl_buffer([128, 128], "float32", data=buffer_data) as buffer: @@ -352,9 +352,8 @@ def test_allocate_with_decl_buffer_no_sugar_multi_usage(): _assert_print( obj, """ -with T.allocate([128, 128], "float32", "global") as v: - buffer = T.decl_buffer((128, 128), data=v) - T.evaluate(v) +with T.decl_buffer((128, 128)) as buffer: + T.evaluate(buffer.data) """, )