# 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
# pylint: disable=redefined-builtin
"""The base Relax operators."""

from typing import Dict, Union, List, Tuple, Optional, Callable


import tvm
import tvm.runtime
from tvm.runtime.object import Object
from tvm.runtime import ObjectGeneric

from . import _ffi_api
from ..expr import Expr, StringImm, ShapeExpr, Call, ExternFunc, GlobalVar, Var
from ..struct_info import StructInfo, TensorStructInfo
from ...ir import PrimExpr
from ..utils import args_converter


py_print = print  # pylint: disable=invalid-name



[文档]
def register_gradient(
    op_name: str,
    fgradient: Callable[[Var, Call, Var, "BlockBuilder"], List[Expr]] = None,
    level: int = 10,
):
    """Register operator gradient function for a relax operator.

    Parameters
    ----------
    op_name: str
        The name of the op.

    fgradient: function (orig_var: Var, orig_call: Call, output_grad: Var, ctx: BlockBuilder)
         -> partials: List[Expr]
        The gradient function being used.

    level: int
        The priority level
    """
    return tvm.ir.register_op_attr(op_name, "FPrimalGradient", fgradient, level)




[文档]
def null_value() -> Call:
    """Create a call node that represents a null value object.

    Returns
    -------
    ret: Call
        The created call node.
    """
    return _ffi_api.null_value()  # type: ignore



def _wrap_inline_arg_tuple(args) -> Expr:
    """Helper function to wrap argument tuple

    Normalize the arguments provided the functions that accept a tuple
    of arguments, and require the tuple of arguments to be written
    in-line.  If the arguments provided are a single relax expression,
    and are not a reference to a relax tuple, then wrap them into an
    in-line relax Tuple.

    """
    if (
        isinstance(args, Expr)
        and not isinstance(args, tvm.relax.Tuple)
        and (
            args.struct_info_ is None
            or not isinstance(args.struct_info_, tvm.relax.TupleStructInfo)
        )
    ):
        return tvm.relax.Tuple([args])
    else:
        return args



[文档]
@args_converter.auto
def call_tir(
    gvar: GlobalVar,
    args: Expr,
    out_sinfo: Union[TensorStructInfo, List[TensorStructInfo]],
    tir_vars: Optional[Union[ShapeExpr, Tuple[PrimExpr], List[PrimExpr]]] = None,
) -> Call:
    """
    Call a tir.prim_func and return the output.

    Parameters
    ----------
    gvar : GlobalVar
        The GlobalVar referring to a tir PrimFunc.

    args : Expr
        The input arguments.

    out_sinfo : Union[TensorStructInfo, List[TensorStructInfo]]
        The structure info of the call_tir output.
        It should be a single or a list of TensorStructInfo. Each one denotes the
        structure info of a returned tensor.

    tir_vars : Optional[Union[ShapeExpr, Tuple[PrimExpr], List[PrimExpr]]]
        ShapeExpr representing a tuple of integers to unpack when calling func. Is null if not used

    Returns
    -------
    ret: Call
        A call node for the call_tir operator.
    """
    args = _wrap_inline_arg_tuple(args)

    if not isinstance(out_sinfo, list):
        out_sinfo = [out_sinfo]

    if isinstance(tir_vars, (list, tuple)):
        tir_vars = ShapeExpr(tir_vars)

    return _ffi_api.call_tir(gvar, args, out_sinfo, tir_vars)  # type: ignore




[文档]
@args_converter.auto
def call_tir_with_grad(
    gvar: GlobalVar,
    args: Expr,
    out_sinfo: Union[TensorStructInfo, List[TensorStructInfo]],
    te_grad_name: str,
    te_grad_kwargs: Dict[str, Object] = None,
    tir_vars: Optional[Union[ShapeExpr, Tuple[PrimExpr], List[PrimExpr]]] = None,
) -> Call:
    """
    Call a tir.prim_func and return the output. This intrinsic will bind a te gradient function
    (refered by te_grad_name) to the call_tir_with_grad node. The te gradient function will be
    called by the Gradient pass.

    Parameters
    ----------
    gvar : GlobalVar
        The GlobalVar referring to a tir PrimFunc.

    args : Expr
        The input arguments.

    out_sinfo : Union[TensorStructInfo, List[TensorStructInfo]]
        The structure info of the call_tir_with_grad output.
        It should be a single or a list of TensorStructInfo. Each one denotes the
        structure info of a returned tensor.

    te_grad_name : str
        The registered name of the te gradient function associated with the call_tir_with_grad
        node. Must be provided as a keyword argument.

    te_grad_kwargs : Dict[str, Object], optional
        The keyword arguments passed to the te gradient function.
        Optionally provided as a keyword argument. Default: {}.

    tir_vars : Optional[Union[ShapeExpr, Tuple[PrimExpr], List[PrimExpr]]]
        ShapeExpr representing a tuple of integers to unpack when calling func. Is null if not used

    Returns
    -------
    ret: Call
        A call node for the call_tir_with_grad operator.
    """
    args = _wrap_inline_arg_tuple(args)

    if not isinstance(out_sinfo, list):
        out_sinfo = [out_sinfo]

    if isinstance(tir_vars, (list, tuple)):
        tir_vars = ShapeExpr(tir_vars)

    if te_grad_kwargs is None:
        te_grad_kwargs = {}

    return _ffi_api.call_tir_with_grad(  # type: ignore
        gvar, args, out_sinfo, te_grad_name, te_grad_kwargs, tir_vars
    )




[文档]
@args_converter.auto
def call_tir_inplace(
    gvar: GlobalVar,
    args: Expr,
    inplace_indices: Union[int, List[int]],
    out_sinfo: Union[TensorStructInfo, List[TensorStructInfo]],
    tir_vars: Optional[Union[ShapeExpr, Tuple[PrimExpr], List[PrimExpr]]] = None,
) -> Call:
    """
    Call a TIR PrimFunc and return the result, doing the specified computations in-place
    (based on the `inplace_indices` argument; outputs will alias the inputs
    selected by in-place indices).

    Warning: This operator is considered pure by the type system but actually mutates
    the arguments specified by `inplace_indices`. This operator should not be used directly,
    but rather should be inserted by passes that have checked whether it is safe to perform
    operations in-place (i.e., none of the arguments specified as an output is aliased or is
    live after calling call_tir_inplace).

    Direct calls to this operator should be done for testing purposes only.

    Parameters
    ----------
    gvar : GlobalVar
        The GlobalVar referring to a TIR PrimFunc.

    args : Expr
        The input arguments.

    inplace_indices : Union[int, List[int]]
        Specify which arguments should be used for in-place computations.
        If `inplace_indices` is a single integer, it will be made into a singleton list.
        Suppose `inplace_indices[i] = j`, where `j >= 0`. Then the `i`th output
        will be an alias of `args[j]`.
        If `inplace_indices[i] = -1`, then the `i`th output will be a freshly allocated tensor.
        At least one member of `inplace_indices` must not be -1.

    out_sinfo : Union[TensorStructInfo, List[TensorStructInfo]]
        The structure info of the call_tir_inplace output.
        It should be a single `TensorStructInfo` or a list of `TensorStructInfo`.
        Each one denotes the structure info of a returned tensor.
        If a list of `TensorStructInfo` is given, the result will be a tuple of `TensorStructInfo`.

    tir_vars : Optional[Union[ShapeExpr, Tuple[PrimExpr], List[PrimExpr]]]
        ShapeExpr representing a tuple of integers to unpack when calling func. Is null if not used

    Returns
    -------
    ret: Call
        A call node for the call_tir operator.
    """
    args = _wrap_inline_arg_tuple(args)

    if not isinstance(inplace_indices, list):
        inplace_indices = [inplace_indices]

    if not isinstance(out_sinfo, list):
        out_sinfo = [out_sinfo]

    if isinstance(tir_vars, (list, tuple)):
        tir_vars = ShapeExpr(tir_vars)

    return _ffi_api.call_tir_inplace(  # type: ignore
        gvar,
        args,
        inplace_indices,
        out_sinfo,
        tir_vars,
    )




[文档]
@args_converter.auto
def call_dps_packed(
    func: Union[str, Expr],
    args: Expr,
    out_sinfo: Union[TensorStructInfo, List[TensorStructInfo]],
) -> Call:
    """
    Call a destination-passing-style packed function and return the output.

    Note: The called function is assumed to be _pure_ (other than modifying the designated
    output arguments). If the function _does_ result in other side effects, then the compiler
    may end up removing, reordering, or repeating those effects--no guarantees can be made.

    Parameters
    ----------
    func : Union[str, Expr]
        The destination-passing-style function, can be ExternFunc.

    args : Expr
        The input arguments.

    out_sinfo : Union[TensorStructInfo, List[TensorStructInfo]]
        The structure info of the call_dps_packed output.
        It should be a single or a list of TensorStructInfo. Each one denotes the
        structure info of a returned tensor.

    Returns
    -------
    ret: Call
        A call node for the call_dps_packed operator.
    """
    if isinstance(func, str):
        func = ExternFunc(func)

    args = _wrap_inline_arg_tuple(args)

    if not isinstance(out_sinfo, list):
        out_sinfo = [out_sinfo]

    return _ffi_api.call_dps_packed(func, args, out_sinfo)  # type: ignore




[文档]
@args_converter.auto
def call_builtin_with_ctx(
    func: Union[str, Expr],
    args: Expr,
    *,
    sinfo_args: Optional[Union[StructInfo, List[StructInfo]]] = None,
) -> Call:
    """Call a builtin function func.

    Parameters
    ----------
    func : Expr
        The builtin function to be called.

    args : Expr
        The input arguments.

    sinfo_args: Optional[Union[StructInfo, List[StructInfo]]]
        The struct info arguments to the call node.

    Returns
    -------
    ret: Call
        The created call node.
    """
    if isinstance(func, str):
        func = ExternFunc(func)

    if sinfo_args is not None and not isinstance(sinfo_args, (list, tuple)):
        sinfo_args = [sinfo_args]

    return _ffi_api.call_builtin_with_ctx(  # type: ignore
        func,
        args,
        sinfo_args,  # type: ignore
    )




[文档]
@args_converter.auto
def make_closure(
    func: Expr,
    args: Expr,
) -> Object:
    """
    Create a closure with free variables and return the closure.

    Parameters
    ----------
    func : Expr
        The closure, can be ExternFunc or PrimFunc.

    args : Expr
        The input arguments.


    Returns
    -------
    ret: Object
        The VMClosure.
    """

    return _ffi_api.make_closure(func, args)  # type: ignore




[文档]
@args_converter.auto
def invoke_closure(
    closure: Expr,
    args: Expr,
    sinfo_args: Union[List[StructInfo], StructInfo],
) -> Call:
    """
    Invoke a closure.

    Parameters
    ----------
    closure : Expr
        The VMClosure object.

    args : Expr
        The input arguments.

    type_args: Union[List[StructInfo], StructInfo]
        The structure info arguments of the CallNode

    Returns
    -------
    ret: Call
        A call to `invoke_closure`.
    """

    if not isinstance(sinfo_args, (list, tuple)):
        sinfo_args = [sinfo_args]

    return _ffi_api.invoke_closure(closure, args, sinfo_args)  # type: ignore



def render_object(val: tvm.Object) -> str:
    """
    Given a TVM Object, renders it in string form. Used for Relax printing and assertions.

    Parameters
    ----------
    val: tvm.Object
        An object to render

    Returns
    -------
    ret: str
        A string representing the value, ideally human-readable
    """
    if isinstance(val, tvm.nd.NDArray):
        return str(val)
    if isinstance(val, tvm.ir.Array):
        fields = ", ".join([render_object(val[i]) for i in range(len(val))])
        return f"({fields})"
    return str(val)


@tvm.register_func("relax.run.shape_to_tensor")
def relax_shape_to_tensor(shape_tuple: tvm.runtime.ShapeTuple) -> tvm.nd.NDArray:
    """
    Takes a ShapeTuple and convert it to NDArray.

    Parameters
    ----------
    shape_tuple: tvm.runtime.ShapeTuple
        Shape tuple that we want to convert to NDArray at runtime
    """
    return tvm.nd.array([int(v) for v in shape_tuple])


@tvm.register_func("relax.run.print")
def relax_print(format_str: str, *format_args: tvm.Object) -> None:
    """
    Takes a list of values to print, formats with the given format string.
    If the format string is empty, simply prints.

    Call from TVM script like this:
    `relax.print(value1, value2, ..., valueN, format=format_str)`
    or
    `relax.print(value1, value2, ..., valueN) # format_str defaults to ""`

    Parameters
    ----------
    format_str: str
        The last argument is a Python-style format string for printing the value

    format_args: List[Object]
        The values to print.
    """
    val_strs = map(render_object, format_args)
    if format_str == "":
        py_print(*val_strs)
    else:
        py_print(format_str.format(*val_strs))



[文档]
def print(*values: List[Expr], format: Union[str, Expr] = "") -> Expr:
    """Print op to print the values

    Parameters
    ----------
    values : List[Expr]
        The values to print.

    format: Union[str, Expr]
        The format string or StringImm.

    Returns
    -------
    result : Expr
        A relax Call, which will print the value during runtime.
    """
    if isinstance(format, str):
        format = StringImm(format)

    return _ffi_api.print(values, format)  # type: ignore # pylint: disable=no-member



@tvm.register_func("relax.run.assert_op")
def relax_assert_op(condition: tvm.Object, format_str: str, *format_args: tvm.Object) -> None:
    """
    A variadic function. The first value serves as the assertion condition:
    If the condition is true, then the operator does nothing.
    If the condition is false, then the operator raises an assertion error.

    Arguments after the first value serve as format arguments for the error message;
    the last argument must be a format string for the error message (empty by default).
    If the format string is the empty string, then the error message will simply include
    a comma-separated list of the format arguments.
    The condition argument is not included in the format string.

    Parameters
    ----------
    condition: tvm.Object
        The assertion condition. Must be a boolean scalar.

    format_str: str
        The last argument is a Python-style format string for printing the value

    format_args: List[tvm.Object]
        Values used for formatting the string.
    """
    if not isinstance(format_str, str):
        raise ValueError(
            f"The format string argument to assert must be a string, given {type(format_str)})"
        )

    if isinstance(condition, (bool, int)):
        val = condition
    elif isinstance(condition, tvm.nd.NDArray):
        # may happen if the original program had unknown shape or dtype for the tensor's type
        dtype = condition.dtype
        if dtype != "bool":
            raise ValueError(f"The condition must be a bool scalar, but given a {dtype} tensor")
        shape = condition.shape
        if len(shape) != 0:
            raise ValueError(f"The condition must be a scalar, but it has a shape of {shape}")

        val = condition.numpy()

    else:
        # should be guaranteed by the type system
        raise ValueError(
            f"The condition for relax assert must be a bool, int, or NDArray, "
            f"but received a {type(condition)}."
        )

    if not val:
        error_message = "Assertion Failed"
        if format_args or format_str != "":
            rendered = map(render_object, format_args)
            if format_str != "":
                error_message = format_str.format(*rendered)
            else:
                error_message = ", ".join(rendered)
        raise AssertionError(error_message)



[文档]
def assert_op(
    condition: Union[Expr, PrimExpr],
    format_args: Optional[Union[Expr, List[Expr]]] = None,
    format: Union[str, Expr] = "",
) -> Expr:
    """
    Create a call to Relax's assert_op operation (`assert` is reserved in Python,
    so the name must be distinct).

    Parameters
    ----------
    condition: Union[Expr, PrimExpr]
        The assertion condition.

    format_args: Optional[Union[Expr, List[Expr]]]
        Format arguments for the error message if the condition fails.

    format: Union[str, Expr]
        The format string or StringImm for the error message.

    Returns
    -------
    result : Expr
        A Call to the Relax assert operation.
    """
    if not isinstance(condition, Expr):
        condition = tvm.relax.PrimValue(condition)

    if format_args is None:
        format_args = []
    elif isinstance(format_args, Expr):
        format_args = [format_args]

    if isinstance(format, str):
        format = StringImm(format)

    return _ffi_api.assert_op(condition, format_args, format)  # type: ignore




[文档]
def shape_of(expr: Expr) -> Expr:
    """Get shape of a tensor.

    Parameters
    ----------
    expr : Expr
        The input Expr.

    Returns
    -------
    result : Expr
        A relax Call, which gets the shape of the input
    """
    return _ffi_api.shape_of(expr)  # type: ignore # pylint: disable=no-member




[文档]
def tensor_to_shape(expr: Expr) -> Expr:
    """Convert tensor to shape expr.
    Parameters
    ----------
    expr : Expr
        The input Expr
    Returns
    -------
    result : Expr
        A relax Call, which transforms the tensor values to the shape
    """
    return _ffi_api.tensor_to_shape(expr)  # type: ignore # pylint: disable=no-member




[文档]
def shape_to_tensor(expr: Expr) -> Expr:
    """Convert shape to tensor expr.
    Parameters
    ----------
    expr : Expr
        The input Expr
    Returns
    -------
    result : Expr
        A relax Call, which transforms the shape values to the tensor
    """
    return _ffi_api.shape_to_tensor(expr)  # type: ignore # pylint: disable=no-member




[文档]
@args_converter.auto
def call_inplace_packed(
    func: Union[str, ExternFunc, GlobalVar],
    *args: Expr,
    inplace_indices: Union[int, List[int]],
    sinfo_args: Union[StructInfo, List[StructInfo]],
) -> Expr:
    """
    Construct a call to a packed function that consumes some of its arguments "in-place"
    and returns the mutated arguments (aliased), but should be considered to be otherwise pure.
    The `inplace_indices` argument indicates which of the outputs are mutated arguments.

    The resulting call will have the same semantics as calling the packed function directly.

    Note: This should be used for cases when the user knows that calling the packed function
    with these arguments will **in reality** not cause any other side effects.
    If it is used for a call that **does** result in other side effects, then the compiler
    may end up removing, reordering, or repeating that call, with no guarantees
    made about any side effects from the callee.

    Warning: This operator as treated as pure by the type system even though it *is* performing
    side effects (mutating some arguments). It is therefore incumbent upon the user to ensure
    that it is being used safely (viz., that mutated arguments are not live after the mutation,
    that they do not alias values live after the mutation).

    Parameters
    ----------
    func : Union[str, ExternFunc]
      The name (global symbol) for a PackedFunc or an ExternFunc node.

    args: Expr
      The arguments for the PackedFunc.

    inplace_indices : Union[int, List[int]]
      Specify which arguments should be used for in-place computations.
      If `inplace_indices` is a single integer, it will be made into a singleton list.
      Suppose `inplace_indices[i] = j`, where `j >= 0`. Then the `i`th output
      will be an alias of `args[j]`.
      If `inplace_indices[i] = -1`, then the `i`th output will be a freshly allocated tensor.
      At least one member of `inplace_indices` must not be -1.

    sinfo_args: Union[StructInfo, List[StructInfo]]
        The list of structure info arguments (giving the structural info for the returned value).

    Returns
    -------
    result : Expr
      A Relax call, corresponding to
      `call_pure_packed(ExternFunc(func), args, DictAttrs(kwargs), sinfo_args)`
    """
    if isinstance(func, ExternFunc):
        func = func.global_symbol

    op = ExternFunc(func)
    if sinfo_args is None:
        raise ValueError("R.call_pure_packed is required to have type_args")
    if isinstance(sinfo_args, tuple):  # type: ignore
        sinfo_args = list(sinfo_args)
    elif not isinstance(sinfo_args, list):
        sinfo_args = [sinfo_args]
    if not isinstance(inplace_indices, list):
        inplace_indices = [inplace_indices]

    return _ffi_api.call_inplace_packed(op, args, inplace_indices, sinfo_args)  # type: ignore # pylint: disable=no-member




[文档]
@args_converter.auto
def call_pure_packed(
    func: Union[str, ExternFunc, GlobalVar],
    *args: Expr,
    sinfo_args: Union[StructInfo, List[StructInfo]],
) -> Expr:
    """
    Construct a call to a packed function that should be treated as pure,
    even though packed calls are normally not treated as pure.

    The resulting call will have the same semantics as calling the packed function directly.

    Note: This should be used for cases when the user knows that calling the packed function
    with these arguments will **in reality** not cause any side effects.
    If it is used for a call that **does** result in side effects, then the compiler
    may end up removing, reordering, or repeating that call, with no guarantees
    made about any side effects from the callee.

    Parameters
    ----------
    func : Union[str, ExternFunc]
      The name (global symbol) for a PackedFunc or an ExternFunc node.

    args: Expr
      The arguments for the PackedFunc.

    sinfo_args: Union[StructInfo, List[StructInfo]]
        The list of structure info arguments (giving the structural info for the returned value).

    Returns
    -------
    result : Expr
      A Relax call, corresponding to
      `call_pure_packed(ExternFunc(func), args, DictAttrs(kwargs), sinfo_args)`
    """
    if isinstance(func, ExternFunc):
        func = func.global_symbol

    op = ExternFunc(func)

    if sinfo_args is None:
        raise ValueError("R.call_pure_packed is required to have type_args")

    if isinstance(sinfo_args, tuple):  # type: ignore
        sinfo_args = list(sinfo_args)
    elif not isinstance(sinfo_args, list):
        sinfo_args = [sinfo_args]

    sinfo_args = [
        sinfo()
        if callable(sinfo)
        else sinfo.asobject()
        if isinstance(sinfo, ObjectGeneric)
        else sinfo
        for sinfo in sinfo_args
    ]

    # note: if we need attributes, we can also take them here

    return _ffi_api.call_pure_packed(op, args, None, sinfo_args)  # type: ignore # pylint: disable=no-member




[文档]
@args_converter.auto
def invoke_pure_closure(
    closure: Expr,
    args: Expr,
    sinfo_args: Union[List[StructInfo], StructInfo],
) -> Call:
    """
    Invoke a closure and indicate to the compiler that it is pure.

    Note: This should be used for cases when the user knows that calling the closure
    with these arguments will **in reality** not cause any side effects.
    If it is used for a call that _does_ result in side effects, then the compiler
    may end up removing, reordering, or repeating that call, with no guarantees
    made about any side effects from the callee.

    Parameters
    ----------
    closure : Expr
        The VMClosure object.

    args : Expr
        The input arguments.

    type_args: Union[List[StructInfo], StructInfo]
        The structure info arguments of the CallNode

    Returns
    -------
    ret: Call
        A call to `invoke_pure_closure`.
    """

    if not isinstance(sinfo_args, (list, tuple)):
        sinfo_args = [sinfo_args]

    return _ffi_api.invoke_pure_closure(closure, args, sinfo_args)  # type: ignore




[文档]
def to_vdevice(data, dst_vdevice) -> Expr:
    """Copy data to the destination device. This
    operator helps data transferring between difference devices for
    heterogeneous execution.

    Parameters
    ----------
    data : Expr
        The tensor to be copied.

    dst_device : VDevice
        The destination device where the data is copied to.

    Returns
    -------
    result : Expr
        The copied result.
    """
    return _ffi_api.to_vdevice(data, dst_vdevice)  # type: ignore




[文档]
def hint_on_device(data, dst_vdevice) -> Expr:
    """It provides a hint specifying the device on which the input data should be executed.
    This hint is utilized by RealizeVDevice to propagate the virtual device."

    Parameters
    ----------
    data : Expr
        The tensor to be copied.

    dst_device : VDevice
        The destination device where the data is supposed to be executed.

    Returns
    -------
    result : Expr
        The result.
    """
    return _ffi_api.hint_on_device(data, dst_vdevice)  # type: ignore