tvm.relax.op.memory.view 源代码

# 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.

"""Operations that act on the DLTensor container

While most operations require inspecting the values stored within the
allocated buffers, some operations only require updating the fields in
a `DLTensor`, without touching the values that are stored within it.
For example, given an array of shape `[16,16]`, the slice at
`[0:8,0:16]` can be generated by changing the `DLTensor::shape` field,
while keeping the same underlying data.

"""
from typing import Optional, Sequence, Union

from tvm.tir import PrimExpr
from tvm.relax import Expr, ShapeExpr, DataTypeImm, PrimValue

from . import _ffi_api


PrimExprLike = Union[int, PrimExpr]




def view(
    data: Expr,
    shape: Optional[Union[Sequence[PrimExprLike], Expr]] = None,
    dtype: Optional[Expr] = None,
    relative_byte_offset: Optional[Expr] = None,
) -> Expr:
    """Provide a view into an existing tensor

    The view may have a different shape, may be a different datatype,
    and may start at an offset relative to the source array.

    Regardless of which combination of these options are used, the
    view may never access memory that was not accessible through the
    input `data` array.  This restriction applies even if the `data`
    array is itself a view into a shared backing array.

    Parameters
    ----------
    data : relax.Expr

        The input data to the operator.

    shape : Optional[Union[Sequence[PrimExprLike], Expr]]

        The target shape.  Should be a `relax.ShapeExpr`, or a
        collection that can be converted to a `relax.ShapeExpr`.

    dtype : Optional[Expr]

        The target datatype.  Should be a `relax.ShapeExpr`, or a
        collection that can be converted to a `relax.ShapeExpr`.

    relative_byte_offset: Optional[Expr]

        The offset of the output NDArray, relative to the byte offset
        of `data`.  If `None`, the offset of the view is the same as
        the offset of `data`.

    Returns
    -------
    result : relax.Expr
        The tensor view

    """

    def _normalize(expr, relax_cls):
        if expr is None or isinstance(expr, Expr):
            return expr
        else:
            return relax_cls(expr)

    shape = _normalize(shape, ShapeExpr)
    dtype = _normalize(dtype, DataTypeImm)
    relative_byte_offset = _normalize(relative_byte_offset, PrimValue)

    return _ffi_api.view(data, shape, dtype, relative_byte_offset)  # type: ignore





def ensure_zero_offset(data: Expr) -> Expr:
    """
    Ensure the tensor has elem_offset == 0. A copy will be made if necessary.

    Parameters
    ----------
    data : relax.Expr
        The input tensor

    Results
    -------
    result : relax.Expr
        The tensor with elem_offset == 0
    """
    return _ffi_api.ensure_zero_offset(data)  # type: ignore