[TVMScript] Syntax sugar for reads & writes

python/tvm/script/tir/__init__.pyi

@@ -210,9 +210,14 @@ def alloc_buffer(
 """
 special_stmt - Reads/Writes
 """
-
-def reads(read_regions: Union[BufferSlice, List[BufferSlice]]) -> None: ...
-def writes(write_region: Union[BufferSlice, List[BufferSlice]]) -> None: ...
+@overload
+def reads(read_regions: List[BufferSlice]) -> None: ...
+@overload
+def reads(*read_regions: BufferSlice) -> None: ...
+@overload
+def writes(write_region: List[BufferSlice]) -> None: ...
+@overload
+def writes(*write_region: BufferSlice) -> None: ...
 def block_attr(attrs: Mapping[str, Object]) -> None: ...
 
 """

python/tvm/script/tir/special_stmt.py

@@ -300,7 +300,12 @@ def alloc_buffer(
 
 @register
 class BlockReads(SpecialStmt):
-    """Special function reads([read_buffer_regions])
+    """Special function reads([read_regions], *other_regions)
+
+    Note
+    ----
+    *other_region is an unpackable list of BufferSlice to support
+    reads syntax sugar like reads(BufferRegion1, BufferRegion2, ...)
 
     Example
     -------
@@ -310,7 +315,11 @@ class BlockReads(SpecialStmt):
     """
 
     def __init__(self):
-        def reads(read_regions: Union[BufferSlice, List[BufferSlice]], span: Span = None):
+        def reads(
+            read_regions: Union[BufferSlice, List[BufferSlice]],
+            *other_regions: BufferSlice,
+            span: Span = None,
+        ):
             assert self.context, "call 'exit_scope' before 'enter_scope'"
             block_scope = self.context.current_block_scope()
             if block_scope is None:
@@ -327,6 +336,8 @@ def reads(read_regions: Union[BufferSlice, List[BufferSlice]], span: Span = None
                 )
             if isinstance(read_regions, BufferSlice):
                 read_regions = [read_regions]
+                for region in other_regions:
+                    read_regions.append(region)
             if not isinstance(read_regions, list):
                 self.context.report_error(
                     "Incorrect input type. "
@@ -340,7 +351,12 @@ def reads(read_regions: Union[BufferSlice, List[BufferSlice]], span: Span = None
 
 @register
 class BlockWrites(SpecialStmt):
-    """Special function writes([write_buffer_regions])
+    """Special function writes([write_regions], *other_regions)
+
+    Note
+    ----
+    *other_region is an unpackable list of BufferSlice to support
+    writes syntax sugar like writes(BufferRegion1, BufferRegion2, ...)
 
     Example
     -------
@@ -350,7 +366,11 @@ class BlockWrites(SpecialStmt):
     """
 
     def __init__(self):
-        def writes(write_region: Union[BufferSlice, List[BufferSlice]], span: Span = None):
+        def writes(
+            write_region: Union[BufferSlice, List[BufferSlice]],
+            *other_region: BufferSlice,
+            span: Span = None,
+        ):
             assert self.context, "call 'exit_scope' before 'enter_scope'"
             block_scope = self.context.current_block_scope()
             if block_scope is None:
@@ -369,6 +389,8 @@ def writes(write_region: Union[BufferSlice, List[BufferSlice]], span: Span = Non
                 pass
             elif isinstance(write_region, BufferSlice):
                 write_region = [write_region]
+                for region in other_region:
+                    write_region.append(region)
             else:
                 self.context.report_error(
                     "Incorrect input type. "

tests/python/unittest/test_tvmscript_error_report.py

@@ -411,17 +411,6 @@ def test_error_index_with_stop_slice():
     check_error(error_bufferslice_index_with_stop, 8)
 
 
-def mismatch_args() -> None:
-    A = T.alloc_buffer((128, 128), "float32")
-    with T.block():
-        T.reads(A[0, 0], A[1, 1])  # error
-        T.evaluate(1.0)
-
-
-def test_mismatch_args():
-    check_error(mismatch_args, 4)
-
-
 def special_stmt_except() -> None:
     A = T.alloc_buffer("(128, 128)", "float32")  # error
     T.evaluate(1.0)

tests/python/unittest/test_tvmscript_syntax_sugar.py

@@ -0,0 +1,66 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=missing-function-docstring,missing-module-docstring,invalid-name,pointless-string-statement
+import sys
+
+import pytest
+from tvm.ir import assert_structural_equal
+from tvm.script import tir as T
+
+
+@T.prim_func
+def transformed_matmul_no_syntax_sugar(a: T.handle, b: T.handle, c: T.handle) -> None:
+    A = T.match_buffer(a, [128, 128])
+    B = T.match_buffer(b, [128, 128])
+    C = T.match_buffer(c, [128, 128])
+
+    for i0, i1, i2_outer, i2_inner_outer, i2_inner_inner in T.grid(128, 128, 4, 8, 4):
+        with T.block("update"):
+            vi, vj = T.axis.remap("SS", [i0, i1])
+            vk = T.axis.R(128, i2_outer * 32 + i2_inner_outer * 4 + i2_inner_inner)
+            T.reads([C[vi, vj], A[vi, vk], B[vj, vk]])
+            T.writes([C[vi, vj], A[vi, vk]])
+            with T.init():
+                C[vi, vj] = 0.0
+            A[vi, vk] = A[vi, vk] + B[vj, vk]
+            C[vi, vj] = C[vi, vj] + (A[vi, vk] * B[vj, vk])
+
+
+@T.prim_func
+def transformed_matmul_syntax_sugar(a: T.handle, b: T.handle, c: T.handle) -> None:
+    A = T.match_buffer(a, [128, 128])
+    B = T.match_buffer(b, [128, 128])
+    C = T.match_buffer(c, [128, 128])
+
+    for i0, i1, i2_outer, i2_inner_outer, i2_inner_inner in T.grid(128, 128, 4, 8, 4):
+        with T.block("update"):
+            vi, vj = T.axis.remap("SS", [i0, i1])
+            vk = T.axis.R(128, i2_outer * 32 + i2_inner_outer * 4 + i2_inner_inner)
+            T.reads(C[vi, vj], A[vi, vk], B[vj, vk])
+            T.writes(C[vi, vj], A[vi, vk])
+            with T.init():
+                C[vi, vj] = 0.0
+            A[vi, vk] = A[vi, vk] + B[vj, vk]
+            C[vi, vj] = C[vi, vj] + (A[vi, vk] * B[vj, vk])
+
+
+def test_reads_writes_syntax_sugar():
+    assert_structural_equal(transformed_matmul_no_syntax_sugar, transformed_matmul_syntax_sugar)
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__] + sys.argv[1:]))