# 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=invalid-name,too-many-locals

"""Utility functions for Relax"""

import itertools
import string

from typing import Tuple as typing_Tuple
from typing import Any, Callable, List, Dict, Optional

import tvm
from .. import tir
from ..tir import PrimExpr
from ..runtime import String, convert_to_object
from . import _ffi_api
from .expr import Tuple as rx_Tuple
from .expr import Expr, ShapeExpr, Function, PrimValue, StringImm, te_tensor
from ..te import Tensor as te_Tensor, create_prim_func
from ..ir import Array, Attrs, Type, Map, VDevice
from .struct_info import PrimStructInfo, ShapeStructInfo, TensorStructInfo

# Re-export `args_converter` here for backwards compatibility
from .type_converter import args_converter  # pylint: disable=unused-import


def metadata_partitioner(rx_txt: str) -> List[str]:
    """Extract Relax program and metadata section.

    Parameters
    ----------
    rx_txt : str
        The input relax text.

    Returns
    -------
    output : List[str]
        The result list of partitioned text, the first element
        is the relax program, and the second is metadata section.
    """
    partitions = []
    left_curly = 0
    meta_start = 0
    meta_end = 0
    for i, char in enumerate(rx_txt):
        if i < 0:
            raise ValueError("The program is invalid.")
        if char == "{":
            if meta_start == 0:
                meta_start = i
            left_curly += 1
        elif char == "}":
            left_curly -= 1
            if left_curly == 0:
                meta_end = i + 1
                break

    if meta_end == 0:
        raise ValueError("The metadata section was not found.")
    metadata = rx_txt[meta_start:meta_end]
    rx_program = rx_txt[meta_end:-1]

    partitions.append(rx_program)
    partitions.append(metadata)

    return partitions



[文档]
def convert_to_expr(value: Any) -> Expr:
    """Helper function to convert the input to Expr, which follows the rules:
    1. Return the input itself if it's already a `relax.Expr`;
    2. Return `relax.PrimValue` if the input is a `PrimExpr`;
    3. Return `relax.StringImm` if the input is `tvm.String` or `str`;
    4. Return `relax.Tuple` if the input is a tuple/list of `Expr`.

    Notes
    -----
    1. `tvm.tir.StringImm` is not allowed because of ambiguity,
       which can be either `relax.StringImm` or `relax.PrimValue`.
    """
    if isinstance(value, int):
        return PrimValue(tir.IntImm("int64", value))

    if isinstance(value, float):
        return PrimValue(tir.FloatImm("float64", value))

    tvm_value = convert_to_object(value)
    # Case 1
    if isinstance(tvm_value, Expr):  # type: ignore
        return tvm_value
    # Note`` 1
    if isinstance(tvm_value, tir.StringImm):
        raise TypeError(
            "Cannot convert `tir.StringImm` to `relax.Expr` because of ambiguity,"
            "which can be either `relax.StringImm` or `relax.PrimValue` "
        )
    # Case 2
    if isinstance(tvm_value, PrimExpr):
        return PrimValue(value)
    # Case 3
    if isinstance(tvm_value, String):
        return StringImm(value)
    # Case 4
    if isinstance(value, (tuple, list)):
        # `convert_to_expr` ensures that all elements are `Expr` if no exception raises
        return rx_Tuple([convert_to_expr(v) for v in value])
    raise TypeError(f"Cannot convert {value} with type {type(value)} to `relax.Expr`")



def copy_with_new_vars(func: Function) -> Function:
    """Copy the given function. All variables that are bound inside the original function
    would be copied to satisfy the restriction in the well-formed check: Variables in
    Relax must be bound exactly once. This also ensures that both the function and its copy
    can be inserted into the same IRModule, and be asserted on the structural equality
    agaisnt IRModule created by TVMScript.

    Parameters
    ----------
    func : Function
        The relax function to copy.

    Returns
    -------
    ret : Function
        The copied function.
    """
    return _ffi_api.CopyWithNewVars(func)  # type: ignore


def gen_call_tir_inputs(
    func: Callable, *args: Any, **kwargs: Any
) -> typing_Tuple[tir.PrimFunc, Expr, List[TensorStructInfo], Optional[ShapeExpr]]:
    """Generate the inputs for call_tir according to the te function.
    This function converts arguments from relax expression to te tensor,
    The callback func should return a te tensor or a list of te tensors.

    Parameters
    ----------
    func : Callable
        A function that returns a te tensor or a list of te tensors.

    args : Any, optional
        arguments passed to the function.

    kwargs : Any, optional
        The keyword arguments passed to the function.
        Note that the keyword args 'primfunc_attrs' is reserved for passing func
        attributes to be added to the PrimFunc that gets created.

    Returns
    -------
    ret : Tuple[tir.PrimFunc, Expr, List[TensorStructInfo], Optional[ShapeExpr]]
        ret contains the inputs for call_tir, including a tir prim_func, args,
        out_sinfo, and tir_vars.
    """

    tir_var_map: Dict[tir.Var, tir.PrimExpr] = {}

    call_tir_args = []
    create_primfunc_args = []
    # extra list of tir expression arguments
    # that are not covered by Tensor
    extra_tir_args_list = []

    def _copy_undefined_var(expr: tir.PrimExpr):
        def _visit_expr(e: tir.PrimExpr):
            if isinstance(e, tir.Var) and e not in tir_var_map:
                new_var = tir.Var(e.name, e.dtype)
                tir_var_map[e] = new_var

        tir.stmt_functor.post_order_visit(expr, _visit_expr)

    def _convert_te_arg(te_args: Any) -> Any:
        """Helper function used to convert Relax expressions to TE tensor.

        In the common case, the type of te_args is a Relax expression and is converted
        into a TE tensor.
        If te_args is a nested or recursive datatype (i.e list, dict, tvm.ir.Map, tvm.ir.Array),
        we recursive and convert any value of type Relax expression into a TE tensor.
        Common values of type int, float, and str are preserved.

        In dynamic shape cases, the passed in arguments may contain TIR variable.
        For example, the argument can be a Relax Var with TensorStructInfo, which
        has symbolic shape, or the argument can be a ShapeExpr with symbolic variables.
        To make the PrimFunc generated has independent variables with
        the caller Relax function, we will substitute the TIR variables in the input
        arguments with fresh ones, which is done by maintaining a TIR variable mapping.

        Parameters
        ----------
        te_args : Any
            Argument to convert to TE

        tir_var_map : Dict[tir.Var, tir.PrimExpr]
            The TIR variable mapping, which maps TIR variables on the Relax function
            side to the new set of variables used on the PrimFunc side.

        Returns
        -------
        ret : (Any, [tvm.te.Tensor])
            A tuple of the converted te_args, and a list of te tensors for each converted
            Relax expression
        """

        def _convert_te_arg_helper(arg):
            if isinstance(arg, Expr):  # type: ignore
                if isinstance(arg.struct_info, TensorStructInfo):
                    assert isinstance(
                        arg.struct_info.shape, ShapeExpr
                    ), "emit_te now only supports Tensor that has ShapeExpr shape"
                    for shape_value in arg.struct_info.shape.values:
                        _copy_undefined_var(shape_value)

                    n_args = len(create_primfunc_args)
                    if isinstance(arg, tvm.relax.Var):
                        name = arg.name_hint
                    elif n_args < len(string.ascii_uppercase):
                        name = string.ascii_uppercase[n_args]
                    else:
                        name = f"tensor_input_{n_args}"

                    te_arg = te_tensor(arg, tir_var_map, name)

                    call_tir_args.append(arg)
                    create_primfunc_args.append(te_arg)

                    return te_arg

                if isinstance(arg.struct_info, ShapeStructInfo):
                    assert isinstance(
                        arg, ShapeExpr
                    ), "For Expr having ShapeStructInfo, emit_te now only supports ShapeExpr"
                    return [_convert_te_arg_helper(val) for val in arg.values]

                if isinstance(arg.struct_info, PrimStructInfo):
                    if arg.struct_info.value is None:
                        n_args = len(create_primfunc_args)
                        if isinstance(arg, tvm.relax.Var):
                            name = arg.name_hint
                        elif n_args < len(string.ascii_lowercase):
                            name = string.ascii_lowercase[n_args]
                        else:
                            name = f"scalar_input_{n_args}"

                        tir_param = tir.Var(name, arg.struct_info.dtype)

                        call_tir_args.append(arg)
                        create_primfunc_args.append(tir_param)

                        return tir_param
                    else:
                        return _convert_te_arg_helper(arg.struct_info.value)

            elif isinstance(arg, (list, Array)):
                return [_convert_te_arg_helper(x) for x in arg]
            elif isinstance(arg, tuple):
                return tuple(_convert_te_arg_helper(x) for x in arg)
            elif isinstance(arg, (dict, Map)):
                for key in arg:
                    assert isinstance(
                        key, str
                    ), "emit_te only supports dict with string as the key currently"
                return {k: _convert_te_arg_helper(arg[k]) for k in arg}
            elif isinstance(arg, tir.PrimExpr):
                _copy_undefined_var(arg)
                new_arg = tir.stmt_functor.substitute(arg, tir_var_map)
                extra_tir_args_list.append(new_arg)
                return new_arg
            elif isinstance(arg, (int, float, str, Type, Attrs)) or arg is None:
                return arg
            raise TypeError("not supported type in emit_te: {}".format(type(arg)))

        new_arg = _convert_te_arg_helper(te_args)
        return new_arg

    def _get_unbound_tir_vars(
        args: List[te_Tensor], extra_tir_args: List[PrimExpr]
    ) -> List[tir.Var]:
        """get unbound TIR vars (i.e TIR vars used in the shape but is not
        itself a dimension of a shape)"""

        bound_vars = set()
        used_vars = set()

        def _populate_bound_vars(expr):
            if isinstance(expr, te_Tensor):
                for dim in expr.shape:
                    _populate_bound_vars(dim)
            elif isinstance(expr, tir.Var):
                bound_vars.add(expr)

        def _populate_used_vars(expr):
            if isinstance(expr, te_Tensor):
                for dim in expr.shape:
                    _populate_used_vars(dim)
            elif isinstance(expr, tir.PrimExpr):
                used_vars.update(tir.analysis.undefined_vars(expr))

        for arg in itertools.chain(args, extra_tir_args):
            _populate_used_vars(arg)

        for arg in args:
            _populate_bound_vars(arg)

        diff = used_vars - bound_vars
        return list(diff)

    def _get_vdevice(arg: Any) -> Optional[VDevice]:
        """get the virtual device from arguments."""
        vdevice = None
        if isinstance(arg, Expr):  # type: ignore
            if isinstance(arg.struct_info, TensorStructInfo):
                vdevice = arg.struct_info.vdevice
        elif isinstance(arg, (list, Array, tuple)):
            for x in arg:
                vdevice = _get_vdevice(x)
                if vdevice is not None:
                    return vdevice
        elif isinstance(arg, (dict, Map)):
            for k in arg:
                vdevice = _get_vdevice(arg[k])
                if vdevice is not None:
                    return vdevice
        return vdevice

    def _shape_with_old_tir_var(
        shape_values: List[tir.PrimExpr], tir_var_inverse_map: Dict[tir.Var, tir.PrimExpr]
    ):
        return ShapeExpr(
            [tir.stmt_functor.substitute(value, tir_var_inverse_map) for value in shape_values]
        )

    primfunc_attrs = kwargs.pop("primfunc_attrs", None)

    te_args = _convert_te_arg(args)
    te_kwargs = _convert_te_arg(kwargs)

    te_out = func(*te_args, **te_kwargs)
    assert isinstance(te_out, te_Tensor) or (
        isinstance(te_out, (tuple, list, Array)) and all(isinstance(t, te_Tensor) for t in te_out)
    ), "only support te.tensor or tuple/list/Array of te.tensor as function output"

    outs = [te_out] if isinstance(te_out, te_Tensor) else list(te_out)
    unbound_tir_vars = _get_unbound_tir_vars([*create_primfunc_args, *outs], extra_tir_args_list)

    inputs = [*create_primfunc_args] + outs + unbound_tir_vars
    tir_func = create_prim_func(inputs, "int64")

    if primfunc_attrs:
        tir_func = tir_func.with_attrs(primfunc_attrs)

    tir_func = tir_func.without_attr("global_symbol")

    # Invert the TIR variable mapping, to convert the output shape back
    # with old set of variables.
    tir_var_inverse_map = {v: k for k, v in tir_var_map.items()}

    output_sinfo = [
        TensorStructInfo(
            _shape_with_old_tir_var(out.shape, tir_var_inverse_map),
            out.dtype,
            _get_vdevice(args),
        )
        for out in outs
    ]

    tir_vars = None
    if len(unbound_tir_vars) > 0:
        tir_vars = _shape_with_old_tir_var(unbound_tir_vars, tir_var_inverse_map)

    return (tir_func, call_tir_args, output_sinfo, tir_vars)