Skip to content

Commit

Permalink
[TVMScript] More concise T.allocate syntax printing (apache#13830)
Browse files Browse the repository at this point in the history
This PR is the follow up of apache#13813. We simplify the printing output of `T.allocate` with `T.decl_buffer`. For example, we have a code snippet as 

```python
buffer_data = T.allocate(...)
buffer = T.decl_buffer(..., data=buffer_data)
T.evaluate(buffer_data)
```

Originally, we skip the `T.allocate` only if the var `buffer_data` defined by `T.allocate` is used only once by the following `T.decl_buffer`. This was due to the limitation of the old printer design. 

But in the new printer, we may automatically replace the `buffer_data` with `buffer.data` if skipping the definition of `buffer_data`. We are able to link all `buffer_data` usages together. So the new output result will be like

```python
buffer = T.decl_buffer(...)
T.evaluate(buffer.data)
```
  • Loading branch information
cyx-6 authored and fzi-peccia committed Mar 27, 2023
1 parent cf48327 commit f7b5c10
Show file tree
Hide file tree
Showing 2 changed files with 4 additions and 7 deletions.
4 changes: 1 addition & 3 deletions src/script/printer/tir/stmt.cc
Original file line number Diff line number Diff line change
Expand Up @@ -175,9 +175,7 @@ TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable)
.set_dispatch<tir::Allocate>( //
"", [](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<ExprDoc> args;
Expand Down
7 changes: 3 additions & 4 deletions tests/python/unittest/test_tvmscript_printer_tir.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand All @@ -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)
""",
)

Expand Down

0 comments on commit f7b5c10

Please sign in to comment.