chore: import upstream snapshot with attribution

python/tvm/relax/op/memory/view.py

@@ -0,0 +1,116 @@
+# 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