From 6bf3bd98c369956b2f94070055a17efbac788c3a Mon Sep 17 00:00:00 2001 From: Eric Lunderberg Date: Wed, 25 May 2022 09:50:06 -0500 Subject: [PATCH 1/2] [TVMScript] Allow T.Buffer[] arg annotation to use int as shape Both the function `tvm.tir.decl_buffer` and the TVMScript `T.match_buffer` expression allow a `PrimExpr` to be passed as the buffer shape, which is interpreted as a 1-d buffer of that size. This allows the same behavior to be used in the `T.Buffer` syntactic sugar. (e.g. `A: T.Buffer[16, "float32"]` instead of `A: T.Buffer[(16,), "float32"`) --- python/tvm/script/tir/ty.py | 9 ++++++++- src/printer/tvmscript_printer.cc | 7 ++++++- .../unittest/test_tvmscript_syntax_sugar.py | 15 +++++++++++++++ 3 files changed, 29 insertions(+), 2 deletions(-) diff --git a/python/tvm/script/tir/ty.py b/python/tvm/script/tir/ty.py index 7d90dec64617..878f029e55dd 100644 --- a/python/tvm/script/tir/ty.py +++ b/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" diff --git a/src/printer/tvmscript_printer.cc b/src/printer/tvmscript_printer.cc index 99d1a7845d3f..94f8aeaca00f 100644 --- a/src/printer/tvmscript_printer.cc +++ b/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()); + doc << tir_prefix_ << ".Buffer["; + if (buffer->shape.size() == 1) { + doc << Print(buffer->shape[0]); + } else { + doc << PrintTuple(buffer->shape.as()); + } doc << ", " << PrintDType(buffer->dtype) << "]"; return doc; } diff --git a/tests/python/unittest/test_tvmscript_syntax_sugar.py b/tests/python/unittest/test_tvmscript_syntax_sugar.py index 0da80d80cf21..aebc606528ba 100644 --- a/tests/python/unittest/test_tvmscript_syntax_sugar.py +++ b/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: From 92e138eb0314a32aeefecb189769c8ab5fd2e988 Mon Sep 17 00:00:00 2001 From: Eric Lunderberg Date: Thu, 26 May 2022 09:48:33 -0500 Subject: [PATCH 2/2] Fixed round-trip when buffer size contains an expression --- python/tvm/script/parser.py | 27 ++++++++++++++++++++++++++- 1 file changed, 26 insertions(+), 1 deletion(-) diff --git a/python/tvm/script/parser.py b/python/tvm/script/parser.py index a376cb7eb08d..e4bdd1206506 100644 --- a/python/tvm/script/parser.py +++ b/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: