[TVMScript] Support T.buffer_decl using data pointer from Let/Allocate #10099

Lunderberg · 2022-01-28T18:03:57Z

Supporting this use case required additional support on both the parsing side and the tvmscript printing side.

Parse pointer objects (e.g. T.Ptr[T.int32]) in TVMScript. Previously, there was no handling of ast.TypeApply in TVMScriptParser.
Call T.buffer_decl in the function body if the data variable is defined by Let/Allocate. Previously, these buffer declarations were hoisted to the function header, prior to the declaration of the data variable they require.

These constructions come up in multiple cases in the updated unit tests for #9727, typically in round-trip unit tests for PrimFuncs that have been passed through MakePackedAPI. This change does not depend on the changes included in #9727, and is therefore kept separate.

These can be generated when exporting to TVMscript, but were not parsable after being generated.

LetStmt and AllocateNode can both be used to generate handles that are used in Buffer objects. In these cases, the Buffer declarations must go after the handle declaration, not in the function header.

Lunderberg · 2022-01-28T18:29:03Z

Updated to resolve lint errors.

junrushao · 2022-01-28T21:11:15Z

CC: @Hzfengsy @vinx13 @spectrometerHBH

tests/python/unittest/test_tvmscript_roundtrip.py

shingjan · 2022-02-03T00:54:12Z

python/tvm/script/parser.py

+        """
+        func = self.transform(node.func_name)
+
+        if not isinstance(func, ty.TypeGeneric):


Use of GenericPtrType here maybe more accurate

From my understanding, the TypeApply ast node would also be used for T.Tuple[type1,type2] type annotations. While I couldn't find any usage of T.Tuple type annotations, GenericTupleType exists in tvm.script.tir.ty and would also use TypeApply in the synr ast. I've added a check for hasattr(func, __getitem__), since that expresses the intent that this is a type that can accept type arguments, and would work for both T.Ptr and T.Tuple.

shingjan · 2022-02-03T00:55:02Z

python/tvm/script/parser.py

+        func = self.transform(node.func_name)
+
+        if not isinstance(func, ty.TypeGeneric):
+            self.report_error(f"Expected a type but found {type(func).__name__}", node.span)


Expect a type -> Expect a GenericPtrType maybe be better

Similar to above, I'd like to avoid breaking any usage of T.Tuple, if it exists outside of the implementation above. I've updated the error message to read "Use of type arguments requires a type that accepts type arguments (e.g. T.Ptr), but found {type(func).__name__} instead." to clarify that the error would be in applying a type argument, not in the use of a type annotation.

shingjan · 2022-02-03T00:56:41Z

python/tvm/script/parser.py

+        param_types = []
+        for param in node.params:
+            param_type = self.transform(param)
+            if not isinstance(param_type, ty.TypeGeneric):


Does this for-loop imply that all params of TypeApply call will be transformed into GenericPtrType? If so can we specify that here as well instead of ty.TypeGeneric

It doesn't, no. The parameters of TypeApply typically aren't GenericPtrType. For example, T.Ptr[T.int32], the T.int32 parameter is a ConcreteType.

At some point, I may see if there's support for renaming ty.TypeGeneric to ty.Type. As it is, ty.ConcreteType is a subclass of ty.TypeGeneric, which doesn't make very much sense to me as there since it require any generic parameters in the user-supplied tvmscript.

shingjan · 2022-02-03T00:58:40Z

python/tvm/script/parser.py

+        if not isinstance(func, ty.TypeGeneric):
+            self.report_error(f"Expected a type but found {type(func).__name__}", node.span)
+
+        param_types = []


there is a transform_TypeTuple impl that may be useful here

I had looked at the transform_TypeTuple implementation, and I don't think it's directly applicable. It assumes that there is a synr.ast.TypeTuple node, with types in node.values, and doesn't have a way to accept a list of types directly or to access the node.params of a TypeApply node. The TVMScriptParser.parse_arg_list is the closest I found to the desired functionality, but it requires the node to be an intrinsic, scope handler, or special statement, and doesn't have a case for type annotations.

shingjan · 2022-02-03T01:01:07Z

python/tvm/script/tir/ty.py

@@ -44,6 +44,9 @@ def __init__(self, vtype):
        self.type = vtype

    def evaluate(self):
+        if isinstance(self.type, tvm.ir.Type):


I thought this change only impact GenericPtrType instead of ConcreteType here. Is that the case?

This check gets hit when calling .evaluate() on the return type of GenericPtrType.__getitem__. The return type is a ConcreteType that has been passed a tvm.ir.PointerType, which should then be returned directly when the concrete type is evaluated, rather than the usual behavior of wrapping it in a tvm.ir.PrimType.

After thinking on it, I think that it would be cleaner if the type conversion happens during the call to __init__, rather than being delayed until the .evaluate() call. That way, the self.type object has a consistent type, regardless of whether it was initialized with a string to represent a primitive, or was initialized with a tvm.ir.Type to represent that type.

shingjan · 2022-02-03T01:01:35Z

python/tvm/script/tir/ty.py

@@ -65,6 +70,8 @@ class GenericTupleType(TypeGeneric):  # pylint: disable=abstract-method
    """

    def __getitem__(self, vtypes):
+        if isinstance(vtypes, TypeGeneric):


This seems to be a bug fix right?

That's correct. If transform_TypeApply was to be able to handle both T.Ptr and T.Tuple cases, then I wanted to make sure that they both accepted the same types of arguments. The two options I considered were (a) always passing a tuple of parameter types even if there is only 1, or (b) passing a bare type when there is only 1 and otherwise passing a tuple of parameter types. Previously, T.Ptr implicitly followed convention (b), while T.Tuple followed convention (a). Option (b) matches python's subscripting syntax, so that's the one that I chose.

Hzfengsy

LGTM

shingjan

LGTM

tmoreau89 · 2022-02-03T20:05:05Z

Thank you @Lunderberg @shingjan @Hzfengsy - the PR has been merged

@shingjan

apache#10099) * [TVMScript] Added unit tests demonstrating desired functionality * [TVMScript] Implemented parsing of T.Ptr[...] These can be generated when exporting to TVMscript, but were not parsable after being generated. * [TVMScript] Updated buffer_var printing LetStmt and AllocateNode can both be used to generate handles that are used in Buffer objects. In these cases, the Buffer declarations must go after the handle declaration, not in the function header. * Moved printing of var and buffer_decl into separate statements. * Updated following @shingjan's review comments.

@shingjan

apache#10099) * [TVMScript] Added unit tests demonstrating desired functionality * [TVMScript] Implemented parsing of T.Ptr[...] These can be generated when exporting to TVMscript, but were not parsable after being generated. * [TVMScript] Updated buffer_var printing LetStmt and AllocateNode can both be used to generate handles that are used in Buffer objects. In these cases, the Buffer declarations must go after the handle declaration, not in the function header. * Moved printing of var and buffer_decl into separate statements. * Updated following @shingjan's review comments.

[TVMScript] Added unit tests demonstrating desired functionality

fcab55e

Lunderberg requested review from areusch, comaniac, icemelon, jroesch, junrushao, merrymercy, tqchen, yzhliu and zhiics as code owners January 28, 2022 18:03

Lunderberg added 2 commits January 28, 2022 12:27

[TVMScript] Implemented parsing of T.Ptr[...]

0c589aa

These can be generated when exporting to TVMscript, but were not parsable after being generated.

[TVMScript] Updated buffer_var printing

1b70fa2

LetStmt and AllocateNode can both be used to generate handles that are used in Buffer objects. In these cases, the Buffer declarations must go after the handle declaration, not in the function header.

Lunderberg force-pushed the tvmscript_t_ptr branch from f9041dc to 1b70fa2 Compare January 28, 2022 18:28

Moved printing of var and buffer_decl into separate statements.

a18667c

Lunderberg added the status: need review label Feb 2, 2022

shingjan suggested changes Feb 3, 2022

View reviewed changes

Hzfengsy approved these changes Feb 3, 2022

View reviewed changes

Updated following @shingjan's review comments.

2454682

shingjan approved these changes Feb 3, 2022

View reviewed changes

tmoreau89 merged commit 455c02a into apache:main Feb 3, 2022

Lunderberg mentioned this pull request Feb 3, 2022

[TE][TIR] Implement layout transformations, non-flat memory buffers #9727

Merged

Lunderberg deleted the tvmscript_t_ptr branch February 5, 2022 03:23

driazati mentioned this pull request Jul 14, 2022

TVM v0.9.0.rc0 Release Candidate Notes #12102

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[TVMScript] Support T.buffer_decl using data pointer from Let/Allocate #10099

[TVMScript] Support T.buffer_decl using data pointer from Let/Allocate #10099

Lunderberg commented Jan 28, 2022 •

edited

Loading

Lunderberg commented Jan 28, 2022

junrushao commented Jan 28, 2022

shingjan Feb 3, 2022

Lunderberg Feb 3, 2022 •

edited

Loading

shingjan Feb 3, 2022

Lunderberg Feb 3, 2022 •

edited

Loading

shingjan Feb 3, 2022

Lunderberg Feb 3, 2022

shingjan Feb 3, 2022

Lunderberg Feb 3, 2022

shingjan Feb 3, 2022

Lunderberg Feb 3, 2022

Lunderberg Feb 3, 2022

shingjan Feb 3, 2022

Lunderberg Feb 3, 2022

Hzfengsy left a comment

shingjan left a comment

tmoreau89 commented Feb 3, 2022

[TVMScript] Support T.buffer_decl using data pointer from Let/Allocate #10099

[TVMScript] Support T.buffer_decl using data pointer from Let/Allocate #10099

Conversation

Lunderberg commented Jan 28, 2022 • edited Loading

Lunderberg commented Jan 28, 2022

junrushao commented Jan 28, 2022

Choose a reason for hiding this comment

Lunderberg Feb 3, 2022 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Lunderberg Feb 3, 2022 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Hzfengsy left a comment

Choose a reason for hiding this comment

shingjan left a comment

Choose a reason for hiding this comment

tmoreau89 commented Feb 3, 2022

Lunderberg commented Jan 28, 2022 •

edited

Loading

Lunderberg Feb 3, 2022 •

edited

Loading

Lunderberg Feb 3, 2022 •

edited

Loading