wehub-resource-sync
117 lines
3.6 KiB
Python
117 lines
3.6 KiB
Python
# 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 collections.abc import Sequence
|
|
|
|
from tvm.relax import DataTypeImm, Expr, ShapeExpr
|
|
from tvm.relax.expr import prim_value
|
|
|
|
from ..base import null_value
|
|
from . import _ffi_api
|
|
|
|
PrimExprLike = int | Expr
|
|
|
|
|
|
def view(
|
|
data: Expr,
|
|
shape: Sequence[PrimExprLike] | Expr | None = None,
|
|
dtype: Expr | None = None,
|
|
relative_byte_offset: Expr | None = 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 Tensor, 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 = null_value() if dtype is None else _normalize(dtype, DataTypeImm)
|
|
relative_byte_offset = (
|
|
relative_byte_offset
|
|
if relative_byte_offset is None or isinstance(relative_byte_offset, Expr)
|
|
else prim_value(relative_byte_offset)
|
|
)
|
|
|
|
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
|