From e94f053605ad5bc39aee97ebdbcc4de9a20f0b57 Mon Sep 17 00:00:00 2001 From: Eric Lunderberg Date: Fri, 26 Apr 2024 12:07:57 -0500 Subject: [PATCH 1/2] [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. --- 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(-) create mode 100644 tests/python/runtime/test_runtime_nd_array.py diff --git a/include/tvm/runtime/ndarray.h b/include/tvm/runtime/ndarray.h index d643355d2660..a9995d21d710 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, size_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 aadd5206bccc..082a28c7e204 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 @@ def copyto(self, target, mem_scope=None): 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 @@ def _create_view(self, shape): 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 6d03e2e01b51..d5d623832615 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, size_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 000000000000..8b30b7bba05c --- /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() From 6ead8054083071863f1afe8b5ac5e84321ea4dde Mon Sep 17 00:00:00 2001 From: Eric Lunderberg Date: Sat, 27 Apr 2024 08:10:13 -0500 Subject: [PATCH 2/2] 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 | 2 +- src/runtime/ndarray.cc | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/include/tvm/runtime/ndarray.h b/include/tvm/runtime/ndarray.h index a9995d21d710..5bdc883649c9 100644 --- a/include/tvm/runtime/ndarray.h +++ b/include/tvm/runtime/ndarray.h @@ -147,7 +147,7 @@ class NDArray : public ObjectRef { * outside the bounds of the current array, this function will * raise an exception. */ - TVM_DLL NDArray CreateView(ShapeTuple shape, DLDataType dtype, size_t relative_byte_offset = 0); + TVM_DLL NDArray CreateView(ShapeTuple shape, DLDataType dtype, uint64_t relative_byte_offset = 0); /*! * \brief Create a reference view of NDArray that diff --git a/src/runtime/ndarray.cc b/src/runtime/ndarray.cc index d5d623832615..c2efa79c0c83 100644 --- a/src/runtime/ndarray.cc +++ b/src/runtime/ndarray.cc @@ -179,7 +179,7 @@ struct NDArray::Internal { } }; -NDArray NDArray::CreateView(ShapeTuple shape, DLDataType dtype, size_t relative_byte_offset) { +NDArray NDArray::CreateView(ShapeTuple shape, DLDataType dtype, uint64_t relative_byte_offset) { ICHECK(data_ != nullptr); const DLTensor& orig = get_mutable()->dl_tensor;