chore: import upstream snapshot with attribution

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

@@ -0,0 +1,21 @@
+# isort: skip_file
+# 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.
+"""Relax memory primitives."""
+
+from .memory import alloc_storage, alloc_tensor, kill_storage, kill_tensor
+from .view import view, ensure_zero_offset

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

@@ -0,0 +1,20 @@
+# 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
+"""FFI APIs for tvm.relax.op.memory"""
+
+import tvm_ffi
+
+tvm_ffi.init_ffi_api("relax.op.memory", __name__)

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

@@ -0,0 +1,135 @@
+# 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
+"""Relax memory primitives."""
+
+from tvm.ir import Call
+
+from ...expr import DataTypeImm, Expr, StringImm, prim_value
+from ...utils import convert_to_expr
+from . import _ffi_api
+
+
+def alloc_storage(
+    size: Expr,
+    virtual_device_index: int | Expr,
+    storage_scope: str | Expr,
+    dtype: str | Expr,
+) -> Call:
+    """Construct a Call to allocate a storage with specific size, virtual_device_index,
+    storage_scope and dtype.
+
+    Parameters
+    ----------
+    size : Expr
+        The size of the storage to be allocated.
+
+    virtual_device_index : Union[int, Expr]
+        The virtual device index indicating on which device the storage is to be allocated.
+        Index -1 is reserved for the host device.
+
+    storage_scope : Union[str, Expr]
+        The storage scope to allocate the storage to.
+
+    dtype : Union[str, Expr]
+        The datatype of the storage to be allocated.
+
+    Returns
+    -------
+    result : Call
+        A relax Call, which gets the allocated storage.
+    """
+    size = convert_to_expr(size)
+    if isinstance(dtype, str):
+        dtype = DataTypeImm(dtype)
+    if isinstance(storage_scope, str):
+        storage_scope = StringImm(storage_scope)
+    if isinstance(virtual_device_index, int):
+        virtual_device_index = prim_value(virtual_device_index)
+    return _ffi_api.alloc_storage(size, virtual_device_index, storage_scope, dtype)  # type: ignore
+
+
+def alloc_tensor(
+    storage: Expr,
+    offset: int | Expr,
+    shape: Expr,
+    dtype: str | Expr,
+    runtime_device_ind: int | Expr = prim_value(0),
+) -> Call:
+    """Construct a Call to allocate a tensor on a certain storage starting from the given offset.
+
+    Parameters
+    ----------
+    storage : Expr
+        The storage to allocate the tensor to.
+
+    offset : Union[int, Expr]
+        The storage offset to allocate the tensor.
+
+    shape : Expr
+        The shape of the tensor to be allocated.
+
+    dtype : Union[str, Expr]
+        The datatype of the tensor to be allocated.
+
+    runtime_device_ind: Union[int, Expr]
+        The device index indicating on which device the tensor is to be
+        allocated at runtime. Index -1 is reserved for the host device.
+
+    Returns
+    -------
+    result : Call
+        A relax Call, which gets the allocated tensor.
+    """
+    if isinstance(offset, int):
+        offset = prim_value(offset)
+    shape = convert_to_expr(shape)
+    if isinstance(dtype, str):
+        dtype = DataTypeImm(dtype)
+    if isinstance(runtime_device_ind, int):
+        runtime_device_ind = prim_value(runtime_device_ind)
+    return _ffi_api.alloc_tensor(storage, offset, shape, dtype, runtime_device_ind)  # type: ignore
+
+
+def kill_storage(storage: Expr) -> Call:
+    """Construct a Call to kill a storage.
+
+    Parameters
+    ----------
+    storage : Expr
+        The storage to be killed.
+
+    Returns
+    -------
+    result : Call
+        A relax Call to kill a storage.
+    """
+    return _ffi_api.kill_storage(storage)  # type: ignore
+
+
+def kill_tensor(tensor: Expr) -> Call:
+    """Construct a Call to kill a tensor.
+
+    Parameters
+    ----------
+    tensor : Expr
+        The tensor to be killed.
+
+    Returns
+    -------
+    result : Call
+        A relax Call to kill a tensor.
+    """
+    return _ffi_api.kill_tensor(tensor)  # type: ignore

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