This is an automated email from the ASF dual-hosted git repository.
lunderberg pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new c0385c7523 [Runtime] Allow offset to be specified in
NDArray::CreateView (#16938)
c0385c7523 is described below
commit c0385c75230c3a352eeb1b19daf9e0638b962de0
Author: Eric Lunderberg <[email protected]>
AuthorDate: Mon Apr 29 09:57:31 2024 -0500
[Runtime] Allow offset to be specified in NDArray::CreateView (#16938)
* [Runtime] Allow offset to be specified in NDArray::CreateView
Prior to this commit, the `NDArray::CreateView` method could produce
an aliasing view of an existing array with a different shape or
datatype, but the view was required to have the same
`DLTensor::byte_offset` as the existing array. This commit updates
the `NDArray::CreateView` method with an additional parameter,
specifying the offset of the view relative to the existing array.
* Change type of `relative_byte_offset` from `size_t` to `uint64_t`
Both to match the type used in `DLTensor::byte_offset`, and to resolve
compilation errors on 32-bit platforms, which fail to compile due to a
missing `Type2Str` specialization.
---
include/tvm/runtime/ndarray.h | 20 +-
python/tvm/runtime/ndarray.py | 25 ++-
src/runtime/ndarray.cc | 70 +++----
tests/python/runtime/test_runtime_nd_array.py | 253 ++++++++++++++++++++++++++
4 files changed, 333 insertions(+), 35 deletions(-)
diff --git a/include/tvm/runtime/ndarray.h b/include/tvm/runtime/ndarray.h
index d643355d26..5bdc883649 100644
--- a/include/tvm/runtime/ndarray.h
+++ b/include/tvm/runtime/ndarray.h
@@ -126,13 +126,29 @@ class NDArray : public ObjectRef {
* \param stream The output data stream
*/
inline void Save(dmlc::Stream* stream) const;
+
/*!
* \brief Create a NDArray that shares the data memory with the current one.
+ *
* \param shape The shape of the new array.
+ *
* \param dtype The data type of the new array.
- * \note The memory size of new array must be smaller than the current one.
+ *
+ * \param relative_byte_offset The offset of the output NDArray,
+ * relative to the current byte offset.
+ *
+ * By default, the offset of the view is the same as the offset
+ * of the current array.
+ *
+ * \note The new array must not allow access of addresses which
+ * would be out of bounds in the current array. If the new
+ * array is larger than the current array, or if the
+ * `relative_byte_offset` would place the end of the new array
+ * outside the bounds of the current array, this function will
+ * raise an exception.
*/
- TVM_DLL NDArray CreateView(ShapeTuple shape, DLDataType dtype);
+ TVM_DLL NDArray CreateView(ShapeTuple shape, DLDataType dtype, uint64_t
relative_byte_offset = 0);
+
/*!
* \brief Create a reference view of NDArray that
* represents as DLManagedTensor.
diff --git a/python/tvm/runtime/ndarray.py b/python/tvm/runtime/ndarray.py
index aadd5206bc..082a28c7e2 100644
--- a/python/tvm/runtime/ndarray.py
+++ b/python/tvm/runtime/ndarray.py
@@ -18,6 +18,7 @@
"""Runtime NDArray API"""
import ctypes
import warnings
+from typing import Optional
import numpy as np
@@ -287,7 +288,7 @@ class NDArray(NDArrayBase):
return self._copyto(res)
raise ValueError(f"Unsupported target type {type(target)}")
- def _create_view(self, shape):
+ def _create_view(self, shape, dtype: Optional[str] = None,
relative_byte_offset: int = 0):
"""Create a view into an existing array.
The view shares the same allocation and datatype as the
@@ -307,12 +308,32 @@ class NDArray(NDArrayBase):
shape: Union[tvm.runtime.ShapeTuple, Sequence[typing.SupportsInt]]
The shape of the view.
+
+ dtype: Optional[str]
+
+ The datatype of the view. If None (default), the view
+ will be the same data type as the current array.
+
+ relative_byte_offset: int
+
+ The location of the view, relative to the location of the current
+ array.
+
+ Note: While the `DLTensor.byte_offset` field of the returned view
+ is usually the same as `relative_byte_offset`, this is not
+ guaranteed. The `DLTensor.byte_offset` field is relative to the
+ start of the backing allocation, while the `relative_byte_offset`
+ is relative to the start of `self`.
+
"""
if not isinstance(shape, tvm.runtime.ShapeTuple):
shape = tvm.runtime.ShapeTuple([int(dim) for dim in shape])
- return _ffi_api.TVMArrayCreateView(self, shape)
+ if dtype is None:
+ dtype = self.dtype
+
+ return _ffi_api.TVMArrayCreateView(self, shape, dtype,
relative_byte_offset)
def device(dev_type, dev_id=0):
diff --git a/src/runtime/ndarray.cc b/src/runtime/ndarray.cc
index 6d03e2e01b..c2efa79c0c 100644
--- a/src/runtime/ndarray.cc
+++ b/src/runtime/ndarray.cc
@@ -179,42 +179,53 @@ struct NDArray::Internal {
}
};
-NDArray NDArray::CreateView(ShapeTuple shape, DLDataType dtype) {
+NDArray NDArray::CreateView(ShapeTuple shape, DLDataType dtype, uint64_t
relative_byte_offset) {
ICHECK(data_ != nullptr);
const DLTensor& orig = get_mutable()->dl_tensor;
- ICHECK(IsContiguous()) << "Can only create view for compact tensor, but
found strides " <<
- [&orig]() {
- std::stringstream ss;
- ss << "[";
- for (int i = 0; i < orig.ndim; i++) {
- if (i) ss << ", ";
- ss << orig.strides[i];
- }
- ss << "]";
- return ss.str();
- }() << ", for shape "
- << [&]() {
- std::stringstream ss;
- ss << "[";
- for (int i = 0; i < orig.ndim; i++) {
- if (i) ss << ", ";
- ss << orig.shape[i];
- }
- ss << "]";
- return ss.str();
- }();
-
- NDArray ret = Internal::Create(shape, dtype,
get_mutable()->dl_tensor.device);
- ret.get_mutable()->dl_tensor.byte_offset =
this->get_mutable()->dl_tensor.byte_offset;
+ CHECK(IsContiguous()) << [&orig]() {
+ std::stringstream ss;
+ ss << "Can only create view for compact tensor, but found strides ";
+
+ ss << "[";
+ for (int i = 0; i < orig.ndim; i++) {
+ if (i) ss << ", ";
+ ss << orig.strides[i];
+ }
+ ss << "]";
+
+ ss << ", for shape ";
+ ss << "[";
+ for (int i = 0; i < orig.ndim; i++) {
+ if (i) ss << ", ";
+ ss << orig.shape[i];
+ }
+ ss << "]";
+ return ss.str();
+ }();
+
+ const auto& curr_dl_tensor = get_mutable()->dl_tensor;
+
+ NDArray ret = Internal::Create(shape, dtype, curr_dl_tensor.device);
+
size_t curr_size = GetDataSize(this->get_mutable()->dl_tensor);
size_t view_size = GetDataSize(ret.get_mutable()->dl_tensor);
- ICHECK_LE(view_size, curr_size)
- << "Tries to create a view that has bigger memory than current one";
+ CHECK_LE(relative_byte_offset + view_size, curr_size)
+ << "ValueError: "
+ << "View with shape " << shape << " and datatype " << dtype << " would
have a size of "
+ << view_size << " bytes. "
+ << "This would occupy bytes " << relative_byte_offset << " <= i_byte < "
+ << (relative_byte_offset + view_size) << " within the backing array. "
+ << "However, the NDArray being viewed only contains " << curr_size << "
bytes (shape = "
+ << ShapeTuple(curr_dl_tensor.shape, curr_dl_tensor.shape +
curr_dl_tensor.ndim)
+ << ", dtype= " << curr_dl_tensor.dtype << ").";
+
// increase ref count
get_mutable()->IncRef();
ret.get_mutable()->manager_ctx = get_mutable();
ret.get_mutable()->dl_tensor.data = get_mutable()->dl_tensor.data;
+ ret.get_mutable()->dl_tensor.byte_offset =
+ get_mutable()->dl_tensor.byte_offset + relative_byte_offset;
return ret;
}
@@ -372,10 +383,7 @@ int TVMArrayAlloc(const tvm_index_t* shape, int ndim, int
dtype_code, int dtype_
TVM_REGISTER_GLOBAL("runtime.TVMArrayAllocWithScope").set_body_typed(NDArray::Empty);
-TVM_REGISTER_GLOBAL("runtime.TVMArrayCreateView").set_body_typed([](NDArray
arr, ShapeTuple shape) {
- NDArray view = arr.CreateView(shape, arr->dtype);
- return view;
-});
+TVM_REGISTER_GLOBAL("runtime.TVMArrayCreateView").set_body_method(&NDArray::CreateView);
int TVMArrayFree(TVMArrayHandle handle) {
API_BEGIN();
diff --git a/tests/python/runtime/test_runtime_nd_array.py
b/tests/python/runtime/test_runtime_nd_array.py
new file mode 100644
index 0000000000..8b30b7bba0
--- /dev/null
+++ b/tests/python/runtime/test_runtime_nd_array.py
@@ -0,0 +1,253 @@
+# 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.
+
+import tvm
+import tvm.testing
+
+import numpy as np
+import pytest
+
+
+def test_1d_full_view_of_1d_arr():
+ """NDArray::CreateView may return the same array"""
+ np_input = np.arange(1024, dtype="int32")
+ tvm_input = tvm.nd.array(np_input)
+
+ tvm_output = tvm_input._create_view([1024])
+ np_expected = np_input
+
+ np.testing.assert_equal(tvm_output.numpy(), np_expected)
+
+
+def test_1d_view_of_first_half_of_1d_arr():
+ """NDArray::CreateView may return a subset of an array"""
+ np_input = np.arange(1024, dtype="int32")
+ tvm_input = tvm.nd.array(np_input)
+
+ tvm_output = tvm_input._create_view([512])
+ np_expected = np_input[0:512]
+
+ np.testing.assert_equal(tvm_output.numpy(), np_expected)
+
+
+def test_1d_view_of_first_half_of_1d_arr():
+ """Subset returned by NDArray::CreateView may have a byte offset"""
+ np_input = np.arange(1024, dtype="int32")
+ tvm_input = tvm.nd.array(np_input)
+
+ tvm_output = tvm_input._create_view([512], relative_byte_offset=512 * 4)
+ np_expected = np_input[512:1024]
+
+ np.testing.assert_equal(tvm_output.numpy(), np_expected)
+
+
+def test_view_larger_than_original_is_invalid():
+ """Subset may not be larger than the original array"""
+ np_input = np.arange(1024, dtype="int32")
+ tvm_input = tvm.nd.array(np_input)
+
+ with pytest.raises(ValueError, match="the NDArray being viewed only
contains 4096 bytes"):
+ tvm_input._create_view([2048])
+
+
+def test_view_entirely_outside_bounds_of_original_is_invalid():
+ """The byte_offset may not place a view outside the original array"""
+ np_input = np.arange(1024, dtype="int32")
+ tvm_input = tvm.nd.array(np_input)
+
+ with pytest.raises(ValueError, match="would occupy bytes 8192 <= i_byte <
12288"):
+ tvm_input._create_view([1024], relative_byte_offset=2048 * 4)
+
+
+def test_view_partially_outside_bounds_of_original_is_invalid():
+ """The byte_offset may not place any elements of a view outside the
original array"""
+ np_input = np.arange(1024, dtype="int32")
+ tvm_input = tvm.nd.array(np_input)
+
+ with pytest.raises(ValueError, match="would occupy bytes 2048 <= i_byte <
6144"):
+ tvm_input._create_view([1024], relative_byte_offset=512 * 4)
+
+
+def test_subview_first_half_of_first_half():
+ """NDArray::CreateView be applied to a view
+
+ The first view is at element offset 0 (byte offset 0). The second
+ view is at element offset 0 (byte offset 0) relative to the first
+ view, or element offset 0 (byte offset 0) relative to the original
+ array.
+
+ """
+ np_input = np.arange(1024, dtype="int32")
+ tvm_input = tvm.nd.array(np_input)
+
+ tvm_view = tvm_input._create_view(
+ [512],
+ relative_byte_offset=0,
+ )
+ tvm_subview = tvm_view._create_view(
+ [256],
+ relative_byte_offset=0,
+ )
+ np_expected = np_input[0:512][0:256]
+
+ np.testing.assert_equal(tvm_subview.numpy(), np_expected)
+
+
+def test_subview_first_half_of_second_half():
+ """NDArray::CreateView be applied to a view
+
+ The first view is at element offset 512 (byte offset 2048). The
+ second view is at element offset 0 (byte offset 0) relative to the
+ first view, or element offset 512 (byte offset 2048) relative to
+ the original array.
+
+ """
+ np_input = np.arange(1024, dtype="int32")
+ tvm_input = tvm.nd.array(np_input)
+
+ tvm_view = tvm_input._create_view(
+ [512],
+ relative_byte_offset=512 * 4,
+ )
+ tvm_subview = tvm_view._create_view(
+ [256],
+ relative_byte_offset=0,
+ )
+ np_expected = np_input[512:1024][0:256]
+
+ np.testing.assert_equal(tvm_subview.numpy(), np_expected)
+
+
+def test_subview_second_half_of_first_half():
+ """NDArray::CreateView be applied to a view
+
+ The first view is at element offset 0 (byte offset 0). The second
+ view is at element offset 256 (byte offset 1024) relative to the
+ first view, or element offset 256 (byte offset 1024) relative to
+ the original array.
+
+ """
+ np_input = np.arange(1024, dtype="int32")
+ tvm_input = tvm.nd.array(np_input)
+
+ tvm_view = tvm_input._create_view(
+ [512],
+ relative_byte_offset=0,
+ )
+ tvm_subview = tvm_view._create_view(
+ [256],
+ relative_byte_offset=256 * 4,
+ )
+ np_expected = np_input[0:512][256:512]
+
+ np.testing.assert_equal(tvm_subview.numpy(), np_expected)
+
+
+def test_subview_second_half_of_second_half():
+ """NDArray::CreateView be applied to a view
+
+ The first view is at element offset 512 (byte offset 2048). The
+ second view is at element offset 256 (byte offset 1024) relative
+ to the first view, or element offset 768 (byte offset 3072)
+ relative to the original array.
+
+ """
+ np_input = np.arange(1024, dtype="int32")
+ tvm_input = tvm.nd.array(np_input)
+
+ tvm_view = tvm_input._create_view(
+ [512],
+ relative_byte_offset=512 * 4,
+ )
+ tvm_subview = tvm_view._create_view(
+ [256],
+ relative_byte_offset=256 * 4,
+ )
+ np_expected = np_input[512:1024][256:512]
+
+ np.testing.assert_equal(tvm_subview.numpy(), np_expected)
+
+
+def test_subview_must_be_in_range_of_immediate_parent():
+ """Bounds-checking is applied relative to the NDArray
+
+ The first view is at location and covers bytes [0,2048). The
+ subview would occupy bytes [2048, 4096), and raises an error as
+ this is outside the range of the view.
+
+ """
+ np_input = np.arange(1024, dtype="int32")
+ tvm_input = tvm.nd.array(np_input)
+
+ tvm_view = tvm_input._create_view(
+ [512],
+ relative_byte_offset=0,
+ )
+
+ with pytest.raises(ValueError, match="would occupy bytes 2048 <= i_byte <
4096"):
+ tvm_view._create_view(
+ [512],
+ relative_byte_offset=512 * 4,
+ )
+
+
+def test_2d_view_into_1d_arr():
+ """NDArray::CreateView may change the dimensionality of an array"""
+ np_input = np.arange(1024, dtype="int32")
+ tvm_input = tvm.nd.array(np_input)
+
+ tvm_output = tvm_input._create_view([32, 32])
+ np_expected = np_input.reshape(32, 32)
+
+ np.testing.assert_equal(tvm_output.numpy(), np_expected)
+
+
+def test_2d_full_view_into_2d_arr():
+ """NDArray::CreateView may change the shape of an array"""
+ np_input = np.arange(1024, dtype="int32").reshape(32, 32)
+ tvm_input = tvm.nd.array(np_input)
+
+ tvm_output = tvm_input._create_view([16, 64])
+ np_expected = np_input.reshape(16, 64)
+
+ np.testing.assert_equal(tvm_output.numpy(), np_expected)
+
+
+def test_2d_view_of_first_half_of_2d_arr():
+ """NDArray::CreateView may return a multi-dimensional view"""
+ np_input = np.arange(1024, dtype="int32").reshape(32, 32)
+ tvm_input = tvm.nd.array(np_input)
+
+ tvm_output = tvm_input._create_view([16, 32])
+ np_expected = np_input[0:16, :]
+
+ np.testing.assert_equal(tvm_output.numpy(), np_expected)
+
+
+def test_2d_view_of_second_half_of_2d_arr():
+ """NDArray::CreateView may return a multi-dimensional view with byte
offset"""
+ np_input = np.arange(1024, dtype="int32").reshape(32, 32)
+ tvm_input = tvm.nd.array(np_input)
+
+ tvm_output = tvm_input._create_view([16, 32], relative_byte_offset=32 * 16
* 4)
+ np_expected = np_input[16:32, :]
+
+ np.testing.assert_equal(tvm_output.numpy(), np_expected)
+
+
+if __name__ == "__main__":
+ tvm.testing.main()
