[TVMScript] Allow T.Buffer[] arg annotation to use int as shape

python/tvm/script/parser.py

@@ -170,6 +170,7 @@ def __init__(
         self.tir_namespace = tir_namespace
         self.closure_vars = closure_vars
         self.meta = None
+        self._inside_buffer_sugar = False
 
     def init_function_parsing_env(self):
         """Initialize function parsing environment"""
@@ -1216,6 +1217,9 @@ def transform_TypeConstant(self, node):
 
         See `transform_Constant`.
         """
+        if self._inside_buffer_sugar:
+            return self.transform_Constant(node)
+
         return node.value
 
     def transform_TypeTuple(self, node):
@@ -1225,6 +1229,22 @@ def transform_TypeTuple(self, node):
         """
         return [self.transform(value) for value in node.values]
 
+    def transform_TypeCall(self, node):
+        """TypeCall visitor
+
+        This occurs when an expression is used inside a T.Buffer
+        parameter annotation.
+        """
+
+        # ast.Call has the BuiltinOp as node.func_name.name, where
+        # ast.TypeCall has the BuiltinOp as node.func_name.  So we can
+        # delegate to self.transform_Call, but the error messages for
+        # unsupported operations will highlight the entire expression
+        # and not just the function itself.
+        op = ast.Op(node.span, node.func_name)
+        call = ast.Call(node.span, op, node.params, node.keyword_params)
+        return self.transform_Call(call)
+
     def transform_TypeApply(self, node):
         """Visitor for Type[Type] expressions.
 
@@ -1265,7 +1285,12 @@ def handle_match_buffer_type(self, node, buffer_name):
         assert isinstance(func, SpecialStmt)
 
         # parse args and kwargs for TypeCall and TypeApply
-        arg_list = self.parse_arg_list(func, node)
+        self._inside_buffer_sugar = True
+        try:
+            arg_list = self.parse_arg_list(func, node)
+        finally:
+            self._inside_buffer_sugar = False
+
         # Note that the third element in arg_list would always be the 'name'
         # TODO: This index is hardcoded as a workaround. Better to make it programmatic
         if arg_list[2] is None:

python/tvm/script/tir/ty.py

@@ -20,6 +20,8 @@
 a wrapper for uniform Type system in IR
 """
 # pylint: disable=invalid-name
+from numbers import Integral
+
 import tvm
 from .special_stmt import SpecialStmt, convert_to_int
 
@@ -177,8 +179,13 @@ def __getitem__(self, args):
         """
         if len(args) < 2:
             raise ValueError("T.Buffer[...] needs at least two arguments: shape and dtype.")
+
         shape = args[0]
-        if not isinstance(shape, tuple):
+        dtype = args[1]
+
+        valid_shape = isinstance(shape, (tvm.ir.PrimExpr, Integral, tuple, list))
+        valid_dtype = isinstance(dtype, str)
+        if not (valid_shape and valid_dtype):
             raise ValueError(
                 "The first argument of T.Buffer[...] needs to be a tuple, "
                 "followed by the second argument dtype as a string"

src/printer/tvmscript_printer.cc

@@ -618,7 +618,12 @@ bool TVMScriptPrinter::IsSimpleBuffer(const Buffer& buf) {
 
 Doc TVMScriptPrinter::PrintInlineBufferBind(const Buffer& buffer) {
   Doc doc;
-  doc << tir_prefix_ << ".Buffer[" << PrintTuple(buffer->shape.as<ArrayNode>());
+  doc << tir_prefix_ << ".Buffer[";
+  if (buffer->shape.size() == 1) {
+    doc << Print(buffer->shape[0]);
+  } else {
+    doc << PrintTuple(buffer->shape.as<ArrayNode>());
+  }
   doc << ", " << PrintDType(buffer->dtype) << "]";
   return doc;
 }

tests/python/unittest/test_tvmscript_syntax_sugar.py

@@ -148,6 +148,21 @@ def test_match_buffer_syntax_sugar():
     assert_structural_equal(elementwise_handle, elementwise_buffer_no_kwargs)
 
 
+def test_match_buffer_1d():
+    @T.prim_func
+    def func_no_sugar(a: T.handle):
+        A = T.match_buffer(a, shape=(16,))
+        for i in T.serial(16):
+            A[i] = 0.0
+
+    @T.prim_func
+    def func_with_sugar(A: T.Buffer[16, "float32"]):
+        for i in T.serial(16):
+            A[i] = 0.0
+
+    assert_structural_equal(func_no_sugar, func_with_sugar)
+
+
 # match buffer failed case
 def test_match_buffer_no_kwargs_failed():
     with pytest.raises(ValueError) as e: