Skip to content

Commit

Permalink
Add new Tensor ctor
Browse files Browse the repository at this point in the history
  • Loading branch information
p-wysocki committed Nov 25, 2024
1 parent f6e0ba0 commit 0a1a533
Show file tree
Hide file tree
Showing 4 changed files with 46 additions and 0 deletions.
4 changes: 4 additions & 0 deletions src/bindings/python/src/pyopenvino/core/common.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -490,6 +490,10 @@ ov::Tensor create_shared(py::array& array) {
OPENVINO_THROW("SHARED MEMORY MODE FOR THIS TENSOR IS NOT APPLICABLE! Passed numpy array must be C contiguous.");
}

ov::Tensor tensor_from_pointer(int64_t data_ptr, const ov::Shape& shape, const ov::element::Type& type) {
return ov::Tensor(type, shape, (void*)(data_ptr), {});
}

ov::Tensor tensor_from_pointer(py::array& array, const ov::Shape& shape, const ov::element::Type& type) {
if (type_helpers::get_ov_type(array) == ov::element::string) {
OPENVINO_THROW("SHARED MEMORY MODE FOR THIS TENSOR IS NOT APPLICABLE! String types can be only copied.");
Expand Down
2 changes: 2 additions & 0 deletions src/bindings/python/src/pyopenvino/core/common.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -157,6 +157,8 @@ T object_from_data(D& data, bool shared_memory) {
return create_copied<T>(data);
}

ov::Tensor tensor_from_pointer(int64_t data_ptr, const ov::Shape& shape, const ov::element::Type& ov_type);

ov::Tensor tensor_from_pointer(py::array& array, const ov::Shape& shape, const ov::element::Type& ov_type);

ov::Tensor tensor_from_pointer(py::array& array, const ov::Output<const ov::Node>& port);
Expand Down
32 changes: 32 additions & 0 deletions src/bindings/python/src/pyopenvino/core/tensor.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,38 @@ void regclass_Tensor(py::module m) {
:type shared_memory: bool
)");

cls.def(py::init([](int64_t data_ptr, const ov::Shape& shape, const ov::element::Type& ov_type) {
return Common::tensor_from_pointer(data_ptr, shape, ov_type);
}),
py::arg("data_ptr"),
py::arg("shape"),
py::arg("type") = ov::element::undefined,
R"(
Tensor's special constructor .
Represents an array in the memory with given shape and element type.
It's recommended to use this constructor when initializing Tensor
with data stored in non-RAM memory (for example GPU memory). It creates
an OpenVINO Tensor without the need for additional data copies.
:param data_ptr: Pointer to a C_CONTIGUOUS array which will be wrapped in
openvino.runtime.Tensor with given parameters (shape
and element_type). Array's memory is being shared with
a host, that means the responsibility of keeping host memory is
on the side of a user. Any action performed on the host
memory will be reflected on this Tensor's memory!
:type data_ptr: int
:param shape: Shape of the new tensor.
:type shape: openvino.runtime.Shape
:param type: Element type
:type type: openvino.runtime.Type
:Example:
.. code-block:: python
from openvino import Tensor, Shape, Type
import torch
arr = torch.rand(128).to(torch.device("xpu"))
data_ptr = arr.detach().data_ptr()
t = Tensor(data_ptr, Shape(arr.shape), Type.f32)
)");

cls.def(py::init([](py::array& array, const ov::Shape& shape, const ov::element::Type& ov_type) {
return Common::tensor_from_pointer(array, shape, ov_type);
}),
Expand Down
8 changes: 8 additions & 0 deletions src/bindings/python/tests/test_runtime/test_tensor.py
Original file line number Diff line number Diff line change
Expand Up @@ -617,3 +617,11 @@ def get_tensor():

tensor = get_tensor()
assert np.allclose(tensor.data[0][0][0:3], [0, 0, 1])


def test_init_from_data_ptr():
arr = np.ones((8, 16, 300), dtype=np.float32)
pointer, _ = arr.__array_interface__['data']
tensor = ov.Tensor(pointer, ov.Shape(arr.shape), ov.Type.f32)
arr[0][0][0:2] = 0
assert np.allclose(tensor.data[0][0][0:3], [0, 0, 1])

0 comments on commit 0a1a533

Please sign in to comment.