[Relay][Prelude] Use the Relay parser to define the Relay prelude

python/tvm/relay/_parser.py

@@ -242,10 +242,12 @@ def visitProg(self, ctx):
             self.visit_list(ctx.defn())
             return self.module
 
-        return self.visit(ctx.expr())
+        if ctx.expr():
+            return self.visit(ctx.expr())
 
-    # Exprs
+        return self.module
 
+    # Exprs
     def visitOpIdent(self, ctx):
         # type: (RelayParser.OpIdentContext) -> op.Op
         return op.get(ctx.CNAME().getText())
@@ -368,14 +370,25 @@ def mk_func(self, ctx):
         self.enter_var_scope()
         # Capture type params in params.
         self.enter_type_param_scope()
+        type_params = ctx.typeParamSeq()
+
+        if type_params is not None:
+            type_params = type_params.ident()
+            assert type_params
+            for ty_param in type_params:
+                name = ty_param.getText()
+                self.mk_typ(name, ty.Kind.Type)
+
         var_list, attr_list = self.visit(ctx.argList())
         ret_type = self.getType_(ctx.type_())
 
+        body = self.visit(ctx.body())
+        # NB(@jroesch): you must stay in the type parameter scope until
+        # after you exit the body, you can reference the type parameters
+        # of your parent scopes.
         type_params = list(self.exit_type_param_scope())
         if type_params:
             _, type_params = zip(*type_params)
-
-        body = self.visit(ctx.body())
         self.exit_var_scope()
 
         attrs = tvm.make.node("DictAttrs", **attr_list) if attr_list is not None else None
@@ -453,16 +466,23 @@ def visitIncompleteType(self, ctx):
         # type (RelayParser.IncompleteTypeContext) -> None:
         return None
 
-    def visitIdentType(self, ctx):
-        # type: (RelayParser.IdentTypeContext) -> Union[ty.TensorType, str]
-        ident_type = ctx.CNAME().getText()
+    def visitTypeIdent(self, ctx):
+        # type: (RelayParser.TypeIdentContext) -> Union[ty.TensorType, str]
+        '''
+        Handle type identifier.
+        '''
+        type_ident = ctx.CNAME().getText()
 
-        # look through all type prefixes for a match
+        # Look through all type prefixes for a match
         for type_prefix in TYPE_PREFIXES:
-            if ident_type.startswith(type_prefix):
-                return ty.scalar_type(ident_type)
+            if type_ident.startswith(type_prefix):
+                return ty.scalar_type(type_ident)
+
+        type_param = lookup(self.type_param_scopes, type_ident)
+        if type_param is not None:
+            return type_param
 
-        raise ParseError("Unknown builtin type: {}".format(ident_type))
+        raise ParseError("Unknown builtin type: {}".format(type_ident))
 
     # def visitCallType(self, ctx):
     #     # type: (RelayParser.CallTypeContext) -> Union[expr.Expr, ty.TensorType]

python/tvm/relay/grammar/Relay.g4

@@ -19,7 +19,7 @@
 
 grammar Relay;
 
-SEMVER: 'v0.0.1' ;
+SEMVER: 'v0.0.2' ;
 
 // Lexing
 // comments
@@ -111,8 +111,8 @@ expr
   // | 'debug'                                # debug
   ;
 
-func: 'fn'        '(' argList ')' ('->' type_)? body ;
-defn: 'def' ident '(' argList ')' ('->' type_)? body ;
+func: 'fn'        typeParamSeq? '(' argList ')' ('->' type_)? body ;
+defn: 'def' ident typeParamSeq? '(' argList ')' ('->' type_)? body ;
 
 argList
   : varList
@@ -132,15 +132,20 @@ attr: CNAME '=' expr ;
 // relations: 'where' relation (',' relation)* ;
 // relation: ident '(' (type_ (',' type_)*)? ')' ;
 
+typeParamSeq
+  : '[' ']'
+  | '[' ident (',' ident)* ']'
+  ;
+
 type_
   : '(' ')'                                         # tupleType
   | '(' type_ ',' ')'                               # tupleType
   | '(' type_ (',' type_)+ ')'                      # tupleType
-  | identType                                       # identTypeType
+  | typeIdent                                       # typeIdentType
   | 'Tensor' '[' shapeSeq ',' type_ ']'             # tensorType
   // currently unused
-  // | identType '[' (type_ (',' type_)*)? ']'         # callType
-  | 'fn' '(' (type_ (',' type_)*)? ')' '->' type_   # funcType
+  // | typeIdent '[' (type_ (',' type_)*)? ']'         # callType
+  | 'fn' typeParamSeq? '(' (type_ (',' type_)*)? ')' '->' type_   # funcType
   | '_'                                             # incompleteType
   | NAT                                             # intType
   ;
@@ -158,7 +163,7 @@ shape
   | NAT                             # intShape
   ;
 
-identType: CNAME ;
+typeIdent : CNAME ;
 // int8, int16, int32, int64
 // uint8, uint16, uint32, uint64
 // float16, float32, float64

python/tvm/relay/grammar/py2/.gitattributes

@@ -0,0 +1,3 @@
+Relay* binary
+Relay* linguist-generated=true
+Relay* linguist-detectable=false