[GitHub] [tvm] tkonolige commented on a change in pull request #8215: [TVMSCRIPT] Fix printing of rank 0 buffer access

[GitBox <gi...@apache.org>]

tkonolige commented on a change in pull request #8215:
URL: https://github.com/apache/tvm/pull/8215#discussion_r648490597



##########
File path: tests/python/unittest/test_tvmscript_roundtrip.py
##########
@@ -2888,6 +2888,67 @@ def test_opaque_block():
     assert len(root_block.body.body[1].block.iter_vars) == 0
 
 
+@tvm.script.tir
+def rank0(a: ty.handle) -> None:
+    A = tir.match_buffer(a, (), "float32")
+    B = tir.alloc_buffer((), "float32")
+    A[()] = 2
+    B[()] = A[()]
+
+
+def test_rank0_buffers():
+    func = rank0
+    rt_func = tvm.script.from_source(tvm.script.asscript(func, True))
+    tvm.ir.assert_structural_equal(func, rt_func)
+
+
+@tvm.script.tir
+def rank0_block(a: ty.handle) -> None:
+    A = tir.match_buffer(a, (), "float32")
+    B = tir.alloc_buffer((), "float32")
+    tir.store(B.data, 0, tir.load("float32", A.data, 0))
+
+    with tir.block([], "update") as []:
+        tir.reads([A[()]])
+        tir.writes([B[()]])
+        for i in range(0, 1):
+            B[()] = A[()]
+
+
+def test_rank0_blocks():
+    func = rank0_block
+    print(tvm.script.asscript(func, True))

Review comment:
       Totally forgot to remove the print statements. Done.




-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
users@infra.apache.org