[TVMScript] A1: Relax Parser infra (#240)

This PR introduces the tvm script parser for relax dialect.

Note that it's only a PR for infrastructure, which does not support all features yet.

Author: Hzfengsy

python/tvm/script/__init__.py

@@ -22,6 +22,7 @@
 from .parser import ir_module as ir_module_v2
 from .parser import parse as from_source_v2
 from .parser import tir as tir_v2
+from .parser import relax as relax_v2
 
 #############
 from .parser_v1 import from_source as from_source_v1

python/tvm/script/parser/__init__.py

@@ -15,7 +15,8 @@
 # specific language governing permissions and limitations
 # under the Licens.
 """The parser"""
-from . import _core, ir, tir
+from . import _core, ir, tir, relax
 from ._core import parse
 from .ir import ir_module
 from .tir import prim_func
+from .relax import function

python/tvm/script/parser/ir/__init__.py

@@ -16,6 +16,6 @@
 # under the License.
 # pylint: disable=missing-docstring
 from . import parser as _parser
-from .entry import ir_module
+from .entry import ir_module, is_defined_in_class
 
-__all__ = ["ir_module"]
+__all__ = ["ir_module", "is_defined_in_class"]

python/tvm/script/parser/ir/entry.py

@@ -23,6 +23,21 @@
 from .._core import parse, utils
 
 
+def is_defined_in_class(frames):
+    if len(frames) > 2:
+        maybe_class_frame = frames[2]
+        statement_list = maybe_class_frame[4]
+        if statement_list is None:
+            return False
+        first_statement = statement_list[0]
+        line = first_statement.strip()
+        if line.startswith("class "):
+            return True
+        if line.startswith("@") and "ir_module" in line:
+            return True
+    return False
+
+
 def ir_module(f: Type) -> IRModule:
     if not inspect.isclass(f):
         raise TypeError(f"Expect a class, but got: {f}")

python/tvm/script/parser/relax/__init__.py

@@ -0,0 +1,24 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=missing-docstring
+from ...ir_builder.relax import *  # pylint: disable=redefined-builtin
+from ...ir_builder.relax import ir as _relax
+from . import parser as _parser
+from .entry import function, Tensor
+
+
+__all__ = _relax.__all__ + ["function", "Tensor"]

python/tvm/script/parser/relax/entry.py

@@ -0,0 +1,56 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=missing-docstring
+import inspect
+from typing import Callable, List, Optional, Union, TypeVar
+
+from tvm.relax import Function, Var
+from tvm.tir import PrimExpr
+
+from ...ir_builder.relax import tensor, TensorType
+from .._core import parse, utils
+from ..ir import is_defined_in_class
+
+
+FType = TypeVar("FType", bound=Callable)
+
+
+def function(f: FType) -> Union[Function, FType]:
+    if not inspect.isfunction(f):
+        raise TypeError(f"Expect a function, but got: {f}")
+    if is_defined_in_class(inspect.stack()):
+        return f
+    return parse(f, utils.inspect_function_capture(f))
+
+
+setattr(function, "dispatch_token", "relax")
+
+
+class TensorProxy:
+    def __call__(
+        self,
+        shape: Optional[List[Union[PrimExpr, str]]] = None,
+        dtype: str = None,
+        ndim: int = -1,
+    ) -> TensorType:
+        return tensor(shape, dtype, ndim)
+
+    def __getitem__(self, keys) -> Var:
+        return self(*keys)  # pylint: disable=no-member # type: ignore
+
+
+Tensor = TensorProxy()  # pylint: disable=invalid-name

python/tvm/script/parser/relax/parser.py

@@ -0,0 +1,91 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=missing-docstring,
+
+from typing import Any
+
+from tvm import relax
+
+from ...ir_builder import relax as R
+from ...ir_builder.base import IRBuilder
+from .._core import Parser, dispatch, doc
+
+
+def bind_assign_value(self: Parser, node: doc.expr, var_name: str, value: Any) -> Any:
+    # pylint: disable=unused-argument
+    if isinstance(value, relax.Expr):
+        var = R.emit(value)
+        IRBuilder.name(var_name, var)
+        return var
+    else:
+        raise TypeError(f"Unsupported type {type(value)} in assignment")
+
+
+@dispatch.register(token="relax", type_name="FunctionDef")
+def visit_function_def(self: Parser, node: doc.FunctionDef) -> None:
+    with self.var_table.with_frame():
+        with R.function():
+            R.func_name(node.name)
+            if node.returns is not None:
+                R.func_ret_type(self.eval_expr(node.returns).type)
+            with self.with_dispatch_token("relax"):
+                self.visit(node.args)
+                self.visit_body(node.body)
+
+
+@dispatch.register(token="relax", type_name="Expr")
+def visit_expr_stmt(self: Parser, node: doc.FunctionDef) -> None:
+    self.eval_expr(node.value)
+
+
+@dispatch.register(token="relax", type_name="arguments")
+def visit_arguments(self: Parser, node: doc.arguments) -> None:
+    arg: doc.arg
+    for arg in node.args:
+        if arg.annotation is None:
+            self.report_error(arg, "Type annotation is required for function parameters.")
+        param_type = self.visit_tvm_annotation(arg.annotation)
+        param = R.arg(arg.arg, param_type)
+
+        self.var_table.add(arg.arg, param)
+
+
+@dispatch.register(token="relax", type_name="tvm_annotation")
+def visit_tvm_annotation(self: Parser, node: doc.expr):
+    annotation = self.eval_expr(node)
+    if callable(annotation):
+        annotation = annotation()
+    return annotation
+
+
+@dispatch.register(token="relax", type_name="Assign")
+def visit_assign(self: Parser, node: doc.Assign) -> None:
+    if len(node.targets) != 1:
+        self.report_error(node, "Consequential assignments like 'a = b = c' are not supported.")
+    lhs = node.targets[0]
+    rhs = self.eval_expr(node.value)
+    self.eval_assign(target=lhs, source=rhs, bind_value=bind_assign_value)
+
+
+@dispatch.register(token="relax", type_name="Return")
+def visit_return(self: Parser, node: doc.Assign) -> None:
+    value = self.eval_expr(node.value)
+
+    if isinstance(value, relax.Expr):
+        R.func_ret_value(value)
+    else:
+        self.report_error(node, f"Unsupported return value type {type(value)}.")

python/tvm/script/parser/tir/entry.py

@@ -22,27 +22,13 @@
 
 from ...ir_builder.tir import buffer_decl, ptr
 from .._core import parse, utils
-
-
-def _is_defined_in_class(frames):
-    if len(frames) > 2:
-        maybe_class_frame = frames[2]
-        statement_list = maybe_class_frame[4]
-        if statement_list is None:
-            return False
-        first_statement = statement_list[0]
-        line = first_statement.strip()
-        if line.startswith("class "):
-            return True
-        if line.startswith("@") and "ir_module" in line:
-            return True
-    return False
+from ..ir import is_defined_in_class
 
 
 def prim_func(f: Callable) -> Union[PrimFunc, Callable]:
     if not inspect.isfunction(f):
         raise TypeError(f"Expect a function, but got: {f}")
-    if _is_defined_in_class(inspect.stack()):
+    if is_defined_in_class(inspect.stack()):
         return f
     return parse(f, utils.inspect_function_capture(f))
 

tests/python/relax/test_tvmscript_parser.py

@@ -0,0 +1,51 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from typing import Union
+import tvm
+import tvm.testing
+
+from tvm import relax
+from tvm import IRModule
+from tvm.script.parser import ir as I, tir as T, relax as R
+
+
+def _check(
+    parsed: Union[relax.Function, IRModule],
+    expect: Union[relax.Function, IRModule],
+):
+    # TODO(siyuan): add round-trip tests
+    tvm.ir.assert_structural_equal(parsed, expect)
+
+
+def test_simple_func():
+    @R.function
+    def foo(x: R.Tensor((128, 128), "float32")) -> R.Tensor(None, "float32", ndim=2):
+        R.func_attr({"Primitive": 1})
+        gv0 = R.call_tir("extern_func", x, (128, 128), dtype="float32")
+        return gv0
+
+    x = relax.Var("x", [128, 128], relax.DynTensorType(2, "float32"))
+    bb = relax.BlockBuilder()
+    with bb.function("foo", (x,), attrs={"Primitive": 1}):
+        out = bb.emit(relax.call_tir("extern_func", x, (128, 128), dtype="float32"))
+        bb.emit_func_output(out)
+
+    _check(foo, bb.get()["foo"])
+
+
+if __name__ == "__main__":
+    tvm.testing.main()