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

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

Author: Lunderberg

python/tvm/script/parser.py

@@ -47,6 +47,7 @@
 from .tir.node import Slice, BufferSlice
 from .tir.scope_handler import ScopeHandler, WithScopeHandler, ForScopeHandler
 from .tir.special_stmt import SpecialStmt
+from .tir import ty
 
 
 class CallArgumentReader(object):
@@ -447,7 +448,9 @@ def check_decorator(decorators: List[ast.Expr]) -> bool:
         # add parameters of function
         for arg in node.params:
             # Note that this case is for T.match_buffer syntax sugar
-            if isinstance(arg.ty, (ast.TypeCall, ast.TypeApply)):
+            if isinstance(arg.ty, (ast.TypeCall, ast.TypeApply)) and isinstance(
+                self.transform(arg.ty.func_name), ty.GenericBufferType
+            ):
                 result = self.handle_match_buffer_type(arg.ty, arg.name)
                 if not isinstance(result, buffer.Buffer):
                     self.report_error(
@@ -1138,6 +1141,25 @@ def transform_TypeTuple(self, node):
         """
         return [self.transform(value) for value in node.values]
 
+    def transform_TypeApply(self, node):
+        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)
+
+        param_types = []
+        for param in node.params:
+            param_type = self.transform(param)
+            if not isinstance(param_type, ty.TypeGeneric):
+                self.report_error(f"Expected a type but found {type(param).__name__}", param.span)
+
+            param_types.append(param_type)
+
+        if len(param_types) == 1:
+            return func[param_types[0]]
+        else:
+            return func[param_types]
+
     def handle_match_buffer_type(self, node, buffer_name):
         """special function to handle syntax sugar for match buffer.
 

python/tvm/script/tir/ty.py

@@ -44,7 +44,10 @@ def __init__(self, vtype):
         self.type = vtype
 
     def evaluate(self):
-        return tvm.ir.PrimType(self.type)
+        if isinstance(self.type, tvm.ir.Type):
+            return self.type
+        else:
+            return tvm.ir.PrimType(self.type)
 
 
 class GenericPtrType(TypeGeneric):  # pylint: disable=abstract-method
@@ -54,6 +57,8 @@ class GenericPtrType(TypeGeneric):  # pylint: disable=abstract-method
     """
 
     def __getitem__(self, vtype):
+        if not isinstance(vtype, TypeGeneric):
+            raise TypeError(f"Ptr expects a type argument, but received {type(vtype).__name__}")
         return ConcreteType(tvm.ir.PointerType(vtype.evaluate()))
 
 
@@ -65,6 +70,8 @@ class GenericTupleType(TypeGeneric):  # pylint: disable=abstract-method
     """
 
     def __getitem__(self, vtypes):
+        if isinstance(vtypes, TypeGeneric):
+            vtypes = [vtypes]
         return ConcreteType(tvm.ir.TupleType([vtype.evaluate() for vtype in vtypes]))