# 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, unused-import
"""The struct info nodes of the Relax language."""
from typing import List, Optional, Union

import tvm._ffi
import tvm

from tvm.ir import Span, EnvFunc, Array, VDevice
from tvm.tir import PrimExpr
from tvm.runtime import DataType
from .expr import StructInfo, Expr, ShapeExpr

from . import _ffi_api, ty, expr



[文档]
@tvm._ffi.register_object("relax.ObjectStructInfo")
class ObjectStructInfo(StructInfo):
    """StructInfo of an Object."""

    def __init__(self, span: Span = None) -> None:
        self.__init_handle_by_constructor__(_ffi_api.ObjectStructInfo, span)  # type: ignore




[文档]
@tvm._ffi.register_object("relax.PrimStructInfo")
class PrimStructInfo(StructInfo):
    """StructInfo of a primitive POD value.

    Parameters
    ----------
    dtype_or_expr : Union[str, DataType, PrimExpr]

       The data type of the prim value, or a known expression for the prim
       value.
    """

    value: Optional[PrimExpr]
    dtype: str

    def __init__(
        self,
        dtype: Optional[Union[str, DataType]] = None,
        value: Optional[Union[int, float, PrimExpr]] = None,
        span: Span = None,
    ) -> None:
        # Guard against incorrect usage.  For backwards compatibility,
        # the dtype and value are in the opposite order from most
        # usages.  While PrimStructInfo could take a single positional
        # argument and check the type, this would require an API
        # difference from TVMScript's PrimProxy, which cannot.
        # (PrimProxy uses string arguments for datatype, and also for
        # inline variable definitions when used in a function
        # signature, and requires separate arguments to distinguish
        # the two cases.)
        if isinstance(dtype, (PrimExpr, int, float)):
            raise TypeError(
                f"The first positional argument of PrimStructInfo must be the datatype, "
                f", but received {type(dtype)}.  "
                f"The value can be specified as a keyword argument "
                f"without needing specifying the dtype: "
                f"PrimStructInfo(value=arg)."
            )

        if dtype is None and value is None:
            raise TypeError(
                "PrimStructInfo.__init__ missing required argument.  "
                "Must provide either 'dtype' or 'value'"
            )

        if dtype is not None:
            if isinstance(value, PrimExpr):
                assert value.dtype == dtype, (
                    "When providing both 'value' and 'dtype' to PrimStructInfo.__init__, "
                    "they must be consistent with each other.  "
                    "However, the value {value} has dtype {value.dtype}, "
                    "but the specified dtype was {dtype}."
                )
            elif isinstance(value, (int, float)):
                value = tvm.tir.const(value, dtype)

        # Use relax's default integer type if not otherwise specified.
        if isinstance(value, int):
            value = tvm.tir.IntImm("int64", value)

        if value is None:
            self.__init_handle_by_constructor__(
                _ffi_api.PrimStructInfoFromDtype, dtype, span
            )  # type: ignore
        else:
            self.__init_handle_by_constructor__(
                _ffi_api.PrimStructInfoFromValue, value, span
            )  # type: ignore




[文档]
@tvm._ffi.register_object("relax.ShapeStructInfo")
class ShapeStructInfo(StructInfo):
    """StructInfo of a shape value.

    Parameters
    ----------
    values : Optional[List[PrimExpr]]
       The symbolic shape values if known.

    ndim : Optional[int]
       The size of the shape.

    Note
    ----
    Do not specify values and ndim at the same time.
    """

    values: Optional[List[PrimExpr]]
    ndim: int
    span: Span

    def __init__(
        self, values: Optional[List[PrimExpr]] = None, ndim: int = -1, span: Span = None
    ) -> None:
        self.__init_handle_by_constructor__(
            _ffi_api.ShapeStructInfo, values, ndim, span  # type: ignore
        )




[文档]
@tvm._ffi.register_object("relax.TensorStructInfo")
class TensorStructInfo(StructInfo):
    """StructInfo of a Tensor value.

    Parameters
    ----------
    shape : Optional[Expr]
       The shape expression.

    dtype : Optional[str]
        The content data type.

    vdevice : Optional[Vdevice]
        The virtual device.

    ndim : Optional[int]
       The number of dimensions of the tensor.

    Note
    ----
    Do not specify shape and ndim at the same time.
    """

    shape: Optional[Expr]
    dtype: str
    vdevice: Optional[VDevice]
    ndim: int
    span: Span

    def __init__(
        self,
        shape: Union[Optional[Expr], List[PrimExpr]] = None,
        dtype: str = "float32",
        vdevice: Union[Optional[VDevice], str] = None,
        ndim: int = -1,
        span: Span = None,
    ) -> None:
        if isinstance(shape, (list, tuple, Array)):
            shape = ShapeExpr(shape)
        self.__init_handle_by_constructor__(
            _ffi_api.TensorStructInfo, shape, dtype, ndim, vdevice, span  # type: ignore
        )




[文档]
@tvm._ffi.register_object("relax.TupleStructInfo")
class TupleStructInfo(StructInfo):
    """StructInfo of a Tuple value.

    Parameters
    ----------
    fields: List[StructInfo]
        The struct info of the fields.
    """

    fields: List[StructInfo]
    span: Span

    def __init__(self, fields: List[StructInfo], span: Span = None) -> None:
        self.__init_handle_by_constructor__(_ffi_api.TupleStructInfo, fields, span)  # type: ignore




[文档]
@tvm._ffi.register_object("relax.FuncStructInfo")
class FuncStructInfo(StructInfo):
    """StructInfo of a function value.

    Parameters
    ----------
    params: List[StructInfo]
        The struct info of the fields.

    ret: StructInfo
        The struct info of return value

    purity: bool
        Whether the function is pure (has no visible side effects).
        Note: We consider a function to be pure only if it is pure on all inputs.
        If a function can have visible side effects only in some cases,
        we still consider it impure.
    """

    params: Optional[List[StructInfo]]
    ret: StructInfo
    derive_func: Optional[EnvFunc]
    purity: bool
    span: Span

    def __init__(
        self, params: List[StructInfo], ret: StructInfo, purity: bool = True, span: Span = None
    ) -> None:
        self.__init_handle_by_constructor__(
            _ffi_api.FuncStructInfo, params, ret, purity, span  # type: ignore
        )


[文档]
    @staticmethod
    def opaque_func(
        *,
        ret: Optional[StructInfo] = None,
        derive_func: Optional[Union[str, EnvFunc]] = None,
        purity: bool = False,
        span: Span = None,
    ) -> "FuncStructInfo":
        """
        Create an opaque FuncStructInfo.

        The opaque function takes either a ret
        that specificies the struct info of the return value
        or a derive_func that provides a customized derivation rule.

        Parameters
        ----------
        ret: Optional[StructInfo]
           The struct info of the function return value.

        derive_func: Optional[Union[str,EnvFunc]]
           The environment function used for derivation

        purity: bool
           Whether the function is pure (false by default, as most opaque functions are not pure)

        span: Optional[Span]
           Optional span information of the ast.

        Returns
        -------
        info: FuncStructInfo

        Note
        ----
        We cannot specify ret and derive_func simultaneously.
        """

        if isinstance(derive_func, str):
            derive_func = tvm.ir.EnvFunc.get("tvm.relax.struct_info.infer_view_sinfo")
        return _ffi_api.FuncStructInfoOpaqueFunc(ret, derive_func, purity, span)  # type: ignore