add test case and concrete type for buffer

Author: shingjan

python/tvm/script/tir/__init__.py

@@ -18,6 +18,6 @@
 
 # Type system
 from .ty import int8, int16, int32, int64, float16, float32, float64
-from .ty import boolean, handle, Ptr, Tuple
+from .ty import boolean, handle, Ptr, Tuple, Buffer
 
 from .prim_func import prim_func

python/tvm/script/tir/ty.py

@@ -46,6 +46,19 @@ def evaluate(self):
         return tvm.ir.PrimType(self.type)
 
 
+class ConcreteBufferType(TypeGeneric):  # pylint: disable=too-few-public-methods, abstract-method
+    """TVM script typing class for uniform Type objects"""
+
+    def __init__(self, vtype):
+        self.type = vtype
+
+    def evaluate(self):
+        return tvm.ir.PrimType(self.type)
+
+    def __call__(self, shape, dtype, elem_offset):
+        pass
+
+
 class GenericPtrType(TypeGeneric):  # pylint: disable=abstract-method
     """TVM script typing class generator for PtrType
 
@@ -78,3 +91,4 @@ def __getitem__(self, vtypes):
 handle = ConcreteType("handle")
 Ptr = GenericPtrType()
 Tuple = GenericTupleType()
+Buffer = ConcreteBufferType("Buffer")

tests/python/unittest/test_tvmscript_syntax_sugar.py

@@ -0,0 +1,52 @@
+# 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.script import tir as T
+
+# match buffer - use kwargs
+@T.prim_func
+def elementwise(
+    a: T.Buffer(shape=(128, 128, 128, 128), dtype="float32", elem_offset=1),
+    b: T.Buffer(shape=(128, 128, 128, 128), dtype="float32", elem_offset=2),
+) -> None:
+    # A = T.match_buffer(a, (128, 128, 128, 128))
+    # B = T.match_buffer(b, (128, 128, 128, 128))
+    for i, j, k, l in T.grid(128, 128, 128, 128):
+        with T.block("B"):
+            vi, vj, vk, vl = T.axis.remap("SSSS", [i, j, k, l])
+            b[vi, vj, vk, vl] = a[vi, vj, vk, vl] * 2.0
+
+
+# match buffer - no kwargs
+@T.prim_func
+def elementwise(
+    a: T.Buffer[(128, 128, 128, 128), "float32"],
+    b: T.Buffer[(128, 128, 128, 128), "float32"],
+) -> None:
+    # A = T.match_buffer(a, (128, 128, 128, 128))
+    # B = T.match_buffer(b, (128, 128, 128, 128))
+    for i, j, k, l in T.grid(128, 128, 128, 128):
+        with T.block("B"):
+            vi, vj, vk, vl = T.axis.remap("SSSS", [i, j, k, l])
+            b[vi, vj, vk, vl] = a[vi, vj, vk, vl] * 2.0
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__] + sys.argv[1:]))