This is an automated email from the ASF dual-hosted git repository.
junrushao 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 d9a5f9e [NDArray] Expose NDArray::CreateView to python (#10712)
d9a5f9e is described below
commit d9a5f9ec7ec35db8beb60dd06c813866992bc259
Author: Eric Lunderberg <[email protected]>
AuthorDate: Thu Mar 24 15:06:11 2022 -0500
[NDArray] Expose NDArray::CreateView to python (#10712)
Modifying the array view is needed for Hexagon targets, in order to
first call `tvm.nd.array` with the physical dimensions, then update the
shape to contain the logical dimensions.
---
python/tvm/runtime/ndarray.py | 27 +++++++++++++++++++++++++++
src/runtime/ndarray.cc | 5 +++++
2 files changed, 32 insertions(+)
diff --git a/python/tvm/runtime/ndarray.py b/python/tvm/runtime/ndarray.py
index 8400a59..22d0874 100644
--- a/python/tvm/runtime/ndarray.py
+++ b/python/tvm/runtime/ndarray.py
@@ -249,6 +249,33 @@ class NDArray(NDArrayBase):
return self._copyto(res)
raise ValueError("Unsupported target type %s" % str(type(target)))
+ def _create_view(self, shape):
+ """Create a view into an existing array.
+
+ The view shares the same allocation and datatype as the
+ existing array, but can have a different array shape. This is
+ useful for runtimes that support non-flat memory, where both
+ the physical shape of an allocation and the logical shape of
+ the tensor it represents may need to be independently
+ specified.
+
+ Warning: This function should not be used outside of low-level
+ manipulations, as it breaks non-aliasing assumptions made by
+ TVM. This function may also be removed/replaced in the
+ future.
+
+ Parameters
+ ----------
+ shape: Union[tvm.runtime.ShapeTuple, Sequence[typing.SupportsInt]]
+
+ The shape of the view.
+ """
+
+ if not isinstance(shape, tvm.runtime.ShapeTuple):
+ shape = tvm.runtime.ShapeTuple([int(dim) for dim in shape])
+
+ return _ffi_api.TVMArrayCreateView(self, shape)
+
def device(dev_type, dev_id=0):
"""Construct a TVM device with given device type and id.
diff --git a/src/runtime/ndarray.cc b/src/runtime/ndarray.cc
index edb0aed..bcf78af 100644
--- a/src/runtime/ndarray.cc
+++ b/src/runtime/ndarray.cc
@@ -293,6 +293,11 @@
TVM_REGISTER_GLOBAL("runtime.TVMArrayAllocWithScope").set_body([](TVMArgs args,
*ret = ndarray;
});
+TVM_REGISTER_GLOBAL("runtime.TVMArrayCreateView").set_body_typed([](NDArray
arr, ShapeTuple shape) {
+ NDArray view = arr.CreateView(shape, arr->dtype);
+ return view;
+});
+
int TVMArrayFree(TVMArrayHandle handle) {
API_BEGIN();
NDArray::Internal::FFIDecRef(handle);
