chore: import upstream snapshot with attribution

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

@@ -0,0 +1,456 @@
+# 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, unused-argument
+
+import functools
+import numbers
+from typing import Any
+
+import tvm_ffi
+
+import tvm
+from tvm import relax, tirx
+from tvm.ir import GlobalVar
+from tvm.relax import Expr, Type
+from tvm.relax.script import builder as R
+from tvm.relax.script.builder.frame import BindingBlockFrame
+from tvm.relax.utils import convert_to_expr
+from tvm.script.ir_builder import ir as I
+from tvm.script.ir_builder.base import IRBuilder
+from tvm.script.parser._core import Parser, dispatch, doc
+
+from .entry import (
+    MatchCastPair,
+    TypeProxy,
+    _normalize_ty,
+    _normalize_ty_proxy,
+)
+
+relax.Expr._dispatch_type = relax.Expr  # pylint: disable=protected-access
+dispatch.register_op(relax.Expr, doc.GtE, 0)(lambda lhs, rhs: lhs >= rhs)
+dispatch.register_op(relax.Expr, doc.Gt, 0)(lambda lhs, rhs: lhs > rhs)
+dispatch.register_op(relax.Expr, doc.LtE, 0)(lambda lhs, rhs: lhs <= rhs)
+dispatch.register_op(relax.Expr, doc.Lt, 0)(lambda lhs, rhs: lhs < rhs)
+
+
+def bind_assign_value(
+    self: Parser,
+    node: doc.expr,
+    var_name: str,
+    value: Any,
+    anno_ty: Type | None = None,
+    emit_prim_expr: bool = False,
+) -> Any:
+    var_table = self.var_table.get()
+
+    if isinstance(value, tirx.Var):
+        if value.name and var_name != value.name:
+            self.report_error(
+                node,
+                "Cannot define TIR variables with different names. The LHS of binding should "
+                "has the same name provided in RHS.",
+            )
+        if var_name in var_table:
+            prev_value = var_table[var_name]
+            if not isinstance(prev_value, tirx.Var):
+                self.report_error(
+                    node,
+                    "Cannot redefine a non-TIR-variable object to a TIR variable. Please "
+                    "define the TIR variable with another name.",
+                )
+            if prev_value.ty != value.ty:
+                self.report_error(
+                    node,
+                    f"Expected the same dtype for TIR vars but got {value.ty} vs {prev_value.ty}",
+                )
+            if not isinstance(value, type(prev_value)):
+                self.report_error(
+                    node,
+                    f"Expected the same IR type for TIR vars "
+                    f"but existing value {type(value)} is mismatched "
+                    f"to previous {type(prev_value)}",
+                )
+            value = prev_value
+        IRBuilder.name(var_name, value)
+        return value
+
+    if tvm.ir.is_prim_expr(value):
+        if not emit_prim_expr:
+            return value
+
+    if isinstance(value, tuple):
+        value = convert_to_expr(value)
+    if isinstance(value, numbers.Number):
+        value = R.const(value)
+
+    if isinstance(value, relax.Expr):
+        var = R.emit(value, anno_ty)
+    elif isinstance(value, MatchCastPair):
+        if anno_ty is not None and not tvm_ffi.structural_equal(anno_ty, value.ty):
+            self.report_error(
+                node, "Cannot specify inconsistent annotation for a match cast pair. "
+            )
+        var = R.emit_match_cast(value.value, value.ty)
+    else:
+        return value
+
+    IRBuilder.name(var_name, var)
+    return var
+
+
+def is_prim_value_call(node: doc.expr) -> bool:
+    return isinstance(node, doc.Call) and getattr(node.func, "attr", None) == "prim_value"
+
+
+def eval_ty_proxy(self: Parser, node: doc.expr) -> TypeProxy:
+    try:
+        annotation = self.eval_expr(node)
+        return _normalize_ty_proxy(annotation)
+    except Exception as err:  # pylint: disable=broad-except
+        self.report_error(node, err)
+        raise
+
+
+def eval_ty(self: Parser, node: doc.expr, eval_str: bool = False) -> Type:
+    var_table = self.var_table.get() if eval_str else None
+    try:
+        ty = self.eval_expr(node)
+        return _normalize_ty(ty, var_table)
+    except Exception as err:  # pylint: disable=broad-except
+        self.report_error(node, err)
+        raise
+
+
+def is_called(node: Any, func_name: str) -> bool:
+    # Check if it calls into a func
+    if isinstance(node, doc.Call):
+        # Recursive call was found
+        if isinstance(node.func, doc.Name) and node.func.id == func_name:
+            return True
+    elif isinstance(node, list | tuple):
+        for stmt in node:
+            if is_called(stmt, func_name):
+                return True
+    elif isinstance(node, doc.AnnAssign | doc.Assign | doc.Return | doc.Expr):
+        return is_called(node.value, func_name)
+    elif isinstance(node, doc.With):
+        return is_called(node.body, func_name)
+    elif isinstance(node, doc.If):
+        smts = []
+        if node.body is not None:
+            smts = smts + list(node.body)
+        if node.orelse is not None:
+            smts = smts + list(node.orelse)
+        return is_called(smts, func_name)
+    return False
+
+
+def is_recursive(node: doc.FunctionDef) -> bool:
+    # Check if it is a recursive function
+    for stmt in node.body:
+        if is_called(stmt, node.name):
+            return True
+    return False
+
+
+def collect_symbolic_var_from_prelude(
+    self: Parser, node: doc.FunctionDef, symbolic_vars: dict[str, tirx.Var]
+) -> dict[str, tirx.Var]:
+    prelude_vars = {}
+    for stmt in node.body:
+        if isinstance(stmt, doc.Assign) and all(
+            isinstance(target, doc.Name) and target.id in symbolic_vars for target in stmt.targets
+        ):
+            values = self.eval_expr(stmt.value)
+
+            try:
+                iter(values)
+            except TypeError:
+                values = [values]
+
+            assert len(stmt.targets) == len(values)
+            for target, value in zip(stmt.targets, values):
+                name = target.id
+                prelude_vars[name] = value
+
+    return {**symbolic_vars, **prelude_vars}
+
+
+def collect_symbolic_var_from_params(self: Parser, node: doc.FunctionDef) -> None:
+    # Collect symbolic vars from parameters
+    symbolic_vars = {}
+    for arg in node.args.args:
+        if arg.annotation is None:
+            self.report_error(arg, "Type annotation is required for function parameters.")
+        param_ty_proxy = eval_ty_proxy(self, arg.annotation)
+
+        for var_name in param_ty_proxy.get_symbolic_vars():
+            if var_name not in symbolic_vars:
+                symbolic_vars[var_name] = tirx.Var(var_name, "int64")
+
+    # Update symbolic vars based on
+    symbolic_vars = collect_symbolic_var_from_prelude(self, node, symbolic_vars)
+
+    # Define symbolic vars to the current var_table frame
+    for var_name, var in symbolic_vars.items():
+        self.var_table.add(var_name, var, allow_shadowing=False)
+
+
+@dispatch.register(token="relax", type_name="FunctionDef")
+def visit_function_def(self: Parser, node: doc.FunctionDef) -> None:
+    is_inner_function = self.inside_function
+    self.inside_function = True
+
+    # reserve a var for local function
+    func_val = self.var_table.get().get(node.name)
+    if not func_val and is_recursive(node):
+        collect_symbolic_var_from_params(self, node)
+        if node.returns is None:
+            ret_ty = relax.TupleType([])
+        else:
+            ret_ty = eval_ty(self, node.returns, eval_str=True)
+        params_ty = []
+        for arg in node.args.args:
+            if arg.annotation is None:
+                self.report_error(arg, "Type annotation is required for function parameters.")
+            param_ty = eval_ty(self, arg.annotation, eval_str=True)
+            params_ty.append(param_ty)
+        # created a var for the local function, the same var could be used for recursive call
+        local_func_var = relax.Var(node.name, relax.FuncType(params_ty, ret_ty))
+        self.var_table.add(node.name, local_func_var)
+
+    purity = find_decorator_annotation(node, "pure")
+    # treat the function as private if we are inside another function
+    # or if it has a privacy annotation
+    privacy = is_inner_function or find_decorator_annotation(node, "private", default=False)
+
+    with self.var_table.with_frame():
+        with self.with_dispatch_token("relax"):
+            with R.function(is_pure=purity, is_private=privacy):
+                R.func_name(node.name)
+                collect_symbolic_var_from_params(self, node)
+
+                if node.returns is not None:
+                    ann_ty = eval_ty(self, node.returns, eval_str=True)
+                    R.func_ret_ty(ann_ty)
+
+                self.visit(node.args)
+
+                for stmt in node.body:
+                    if isinstance(stmt, doc.FunctionDef):
+                        if not stmt.decorator_list:
+                            self.report_error(stmt, "Function must be decorated")
+                        dec = self.eval_expr(stmt.decorator_list[-1])
+                        # inline prim_func was found
+                        if dec.dispatch_token == "tirx":
+                            self.report_error(stmt, "inline prim_func is disallowed in Relax IR")
+
+                self.visit_body(node.body)
+    self.inside_function = is_inner_function
+
+
+def find_decorator_annotation(node: doc.FunctionDef, annotation: str, default: bool = True) -> bool:
+    """
+    Check the value of given annotation (argument name) in the function decorator.
+    Returns the value of the annotation if present, otherwise giving the default value.
+    """
+    # look for the named argument in the function decorator
+    for dec in node.decorator_list:
+        if not isinstance(dec, doc.Call) or dec.func.attr != "function":
+            continue
+        for keyword in dec.keywords:
+            if keyword.arg == annotation:
+                return keyword.value.value
+    return default
+
+
+@dispatch.register(token="relax", type_name="tvm_declare_function")
+def visit_tvm_declare_function(self: Parser, node: doc.FunctionDef) -> GlobalVar:
+    with self.var_table.with_frame():
+        collect_symbolic_var_from_params(self, node)
+
+        if node.returns is None:
+            # Use AnyType as unknown return type
+            # NOTE: Cannot use VoidType here because the return type can be refined later.
+            ret_ty = relax.AnyType()
+        else:
+            ret_ty = eval_ty(self, node.returns, eval_str=True)
+        params = []
+        for arg in node.args.args:
+            if arg.annotation is None:
+                self.report_error(arg, "Type annotation is required for function parameters.")
+            param_ty = eval_ty(self, arg.annotation, eval_str=True)
+            params.append(relax.Var(arg.arg, param_ty))
+
+    is_pure = find_decorator_annotation(node, "pure")
+
+    func_signature = relax.Function.create_empty(params, ret_ty, is_pure=is_pure)
+    return I.decl_function(node.name, func_signature)
+
+
+@dispatch.register(token="relax", type_name="pre_visit_local_function")
+def pre_visit_local_function(self: Parser, node: doc.Expr) -> None:
+    ir_builder = IRBuilder()
+    ir_builder.__enter__()
+
+
+@dispatch.register(token="relax", type_name="post_visit_local_function")
+def post_visit_local_function(self: Parser, node: doc.Expr) -> None:
+    ir_builder = IRBuilder.current()
+    result = ir_builder.get()
+    ir_builder.__exit__(None, None, None)
+    # reuse var if it is reserved
+    reserved_var = self.var_table.get().get(node.name)
+    if reserved_var:
+        var = R.emit_var_binding(relax.VarBinding(reserved_var, result))
+    else:
+        var = R.emit(result)
+    IRBuilder.name(node.name, var)
+    self.var_table.add(node.name, var, allow_shadowing=False)
+
+
+@dispatch.register(token="relax", type_name="Expr")
+def visit_expr_stmt(self: Parser, node: doc.Expr) -> None:
+    value = self.eval_expr(node.value)
+    if isinstance(value, relax.Expr):
+        var = R.emit(value)
+        IRBuilder.name("_", var)
+        is_void_value = isinstance(var.ty, relax.TupleType) and len(var.ty.fields) == 0
+
+        if not is_void_value:
+            self.report_error(
+                node,
+                f"Non-void relax expressions must be bound to a variable, "
+                f"but expression of type {var.ty} was used as a statement.",
+            )
+
+    elif value is not None:
+        self.report_error(node, f"Unsupported Expr stmt type {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_ty = eval_ty(self, arg.annotation, eval_str=True)
+        param = R.arg(arg.arg, param_ty)
+
+        self.var_table.add(arg.arg, param)
+
+
+@dispatch.register(token="relax", type_name="tvm_annotation")
+def visit_tvm_annotation(self: Parser, node: doc.expr) -> Type:
+    return eval_ty(self, node, eval_str=False)
+
+
+@dispatch.register(token="relax", type_name="With")
+def visit_with(self: Parser, node: doc.With) -> None:
+    # Currently only `with R.dataflow()` is supported
+    if len(node.items) != 1:
+        self.report_error(node, "Only one item is allowed.")
+    item = node.items[0]
+    if item.optional_vars is not None:
+        self.report_error(
+            item.context_expr,
+            "Relax syntax doesn't allow binding expressions in `with` to variables",
+        )
+    frame = self.eval_expr(item.context_expr)
+    with self.var_table.with_frame():
+        with frame:
+            self.visit(node.body)
+    if isinstance(frame, BindingBlockFrame) and frame.is_dataflow:
+        output_vars = frame.output_vars
+        for var in output_vars:
+            self.var_table.add(var.name_hint, var, allow_shadowing=True)
+
+
+@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=functools.partial(
+            bind_assign_value,
+            emit_prim_expr=is_prim_value_call(node.value),
+        ),
+        allow_shadowing=True,
+    )
+
+
+@dispatch.register(token="relax", type_name="AnnAssign")
+def visit_ann_assign(self: Parser, node: doc.AnnAssign) -> None:
+    lhs = node.target
+    rhs = self.eval_expr(node.value)
+    anno_ty = self.visit_tvm_annotation(node.annotation)
+    self.eval_assign(
+        target=lhs,
+        source=rhs,
+        bind_value=functools.partial(
+            bind_assign_value,
+            anno_ty=anno_ty,
+            emit_prim_expr=is_prim_value_call(node.value),
+        ),
+        allow_shadowing=True,
+    )
+
+
+@dispatch.register(token="relax", type_name="Return")
+def visit_return(self: Parser, node: doc.Assign) -> None:
+    value = self.eval_expr(node.value)
+    value = convert_to_expr(value)
+    R.func_ret_value(value)
+
+
+@dispatch.register(token="relax", type_name="If")
+def visit_if(self: Parser, node: doc.If) -> None:
+    if node.orelse is None:
+        raise ValueError("Else statements are required for relax dialect.")
+    with R.If(self.eval_expr(node.test)) as if_frame:
+        with self.var_table.with_frame():
+            with R.Then():
+                self.visit_body(node.body)
+        with self.var_table.with_frame():
+            with R.Else():
+                self.visit_body(node.orelse)
+    self.var_table.add(if_frame.var_name, if_frame.var, allow_shadowing=True)
+
+
+@dispatch.register(token="relax", type_name="enter_token")
+def enter_token(self: Parser) -> dict[str, Any]:
+    def relax_call(self, *args) -> Expr:
+        args = [convert_to_expr(arg) if isinstance(arg, tuple) else arg for arg in args]
+
+        if all(isinstance(x, Expr) for x in args):
+            return relax.Call(self, args)
+        arg_types = [type(x) for x in args]
+        raise RuntimeError(f"Do not know how to handle GlobalVar.__call__ for types {arg_types}")
+
+    context = {"GlobalVar.__call__": GlobalVar.__call__}
+    GlobalVar.__call__ = relax_call
+    return context
+
+
+@dispatch.register(token="relax", type_name="exit_token")
+def exit_token(self: Parser, context: dict[str, Any]) -> None:
+    assert "GlobalVar.__call__" in context
+    GlobalVar.__call__ = context.get("GlobalVar.__call__")