From 41dfc578fb6699e031924095311e478adb2e4fe0 Mon Sep 17 00:00:00 2001 From: Edgar Riba Date: Sun, 13 Mar 2022 18:51:42 +0100 Subject: [PATCH 1/5] implement __dlpack__, __dlpack_device__, for numpy conversions --- src/dlpack_py.rs | 24 +++--------------------- src/tensor.rs | 43 +++++++++++++++++++++++++++++++++++++++++++ test/test_io.py | 11 +++++++---- test/test_tensor.py | 14 +++++++++++++- 4 files changed, 66 insertions(+), 26 deletions(-) diff --git a/src/dlpack_py.rs b/src/dlpack_py.rs index db9bc5cd..03457147 100644 --- a/src/dlpack_py.rs +++ b/src/dlpack_py.rs @@ -24,17 +24,7 @@ unsafe extern "C" fn destructor(o: *mut pyo3::ffi::PyObject) { // println!("Delete by Python"); } -unsafe extern "C" fn deleter(x: *mut dlpack::DLManagedTensor) { - // println!("DLManagedTensor deleter"); - - let ctx = (*x).manager_ctx as *mut cv::Tensor; - ctx.drop_in_place(); - (*x).dl_tensor.shape.drop_in_place(); - (*x).dl_tensor.strides.drop_in_place(); - x.drop_in_place(); -} - -fn cvtensor_to_dltensor(x: &cv::Tensor) -> dlpack::DLTensor { +pub fn cvtensor_to_dltensor(x: &cv::Tensor) -> dlpack::DLTensor { dlpack::DLTensor { data: x.data.as_ptr() as *mut c_void, device: dlpack::DLDevice { @@ -54,17 +44,9 @@ fn cvtensor_to_dltensor(x: &cv::Tensor) -> dlpack::DLTensor { } #[pyfunction] -pub fn cvtensor_to_dlpack(x: cv::Tensor) -> PyResult<*mut pyo3::ffi::PyObject> { - let tensor_bx = Box::new(x); - let dl_tensor = cvtensor_to_dltensor(&tensor_bx); - - // create dlpack managed tensor - let dlm_tensor = dlpack::DLManagedTensor { - dl_tensor, - manager_ctx: Box::into_raw(tensor_bx) as *mut c_void, - deleter: Some(deleter), - }; +pub fn cvtensor_to_dlpack(x: &cv::Tensor) -> PyResult<*mut pyo3::ffi::PyObject> { + let dlm_tensor: dlpack::DLManagedTensor = x.to_dlpack(); let dlm_tensor_bx = Box::new(dlm_tensor); let name = CString::new("dltensor").unwrap(); diff --git a/src/tensor.rs b/src/tensor.rs index 001c3195..50ae6e7b 100644 --- a/src/tensor.rs +++ b/src/tensor.rs @@ -1,6 +1,19 @@ pub mod cv { use pyo3::prelude::*; + use crate::dlpack; + use crate::dlpack_py::{cvtensor_to_dltensor, cvtensor_to_dlpack}; + use std::ffi::{c_void}; + + unsafe extern "C" fn deleter(x: *mut dlpack::DLManagedTensor) { + // println!("DLManagedTensor deleter"); + + let ctx = (*x).manager_ctx as *mut Tensor; + ctx.drop_in_place(); + (*x).dl_tensor.shape.drop_in_place(); + (*x).dl_tensor.strides.drop_in_place(); + x.drop_in_place(); + } fn get_strides_from_shape(shape: &[i64]) -> Vec { let mut strides = vec![0i64; shape.len()]; @@ -37,7 +50,37 @@ pub mod cv { strides, } } + + #[pyo3(name = "__dlpack__")] + pub fn to_dlpack_py(&self) -> PyResult<*mut pyo3::ffi::PyObject> { + return cvtensor_to_dlpack(self); + } + + #[pyo3(name = "__dlpack_device__")] + pub fn to_dlpack_device_py(&self) -> (u32, i32) { + let tensor_bx = Box::new(self); + let dl_tensor = cvtensor_to_dltensor(&tensor_bx); + (dl_tensor.device.device_type, dl_tensor.device.device_id) + } } + + impl Tensor { + pub fn to_dlpack(&self) -> dlpack::DLManagedTensor { + // we need to clone to avoid race conditions + // TODO: check how to avoid that + let tensor_bx = Box::new(self.clone()); + let dl_tensor = cvtensor_to_dltensor(&tensor_bx); + + // create dlpack managed tensor + let dlm_tensor = dlpack::DLManagedTensor { + dl_tensor, + manager_ctx: Box::into_raw(tensor_bx) as *mut c_void, + deleter: Some(deleter), + }; + dlm_tensor + } + } + } // namespace cv // TODO(carlos): enable tests later diff --git a/test/test_io.py b/test/test_io.py index 8c66f91d..7f7d7d4a 100644 --- a/test/test_io.py +++ b/test/test_io.py @@ -4,6 +4,7 @@ from kornia_rs import Tensor as cvTensor import torch +import jax import numpy as np DATA_DIR = Path(__file__).parent / "data" @@ -20,11 +21,13 @@ def test_read_image_jpeg(): th_tensor = torch.utils.dlpack.from_dlpack(dlpack) assert th_tensor.shape == (195, 258, 3) - # TODO: needs to be fixed + # test __dlpack__() + torch.testing.assert_close( + th_tensor, torch.utils.dlpack.from_dlpack(cv_tensor)) + # convert to dlpack to import to numpy - #dlpack = K.cvtensor_to_dlpack(cv_tensor) - #np_array = np._from_dlpack(dlpack) - #assert np_array.shape == (195, 258, 3) + np_array = np._from_dlpack(cv_tensor) + assert np_array.shape == (195, 258, 3) def test_read_image_rs(): # load an image with image-rs diff --git a/test/test_tensor.py b/test/test_tensor.py index 18e81be8..e4ee0d87 100644 --- a/test/test_tensor.py +++ b/test/test_tensor.py @@ -2,6 +2,7 @@ from kornia_rs import Tensor as cvTensor import torch +import numpy as np def test_smoke(): # dumy test @@ -12,7 +13,18 @@ def test_smoke(): assert len(data) == len(cv_tensor.data) assert cv_tensor.strides == [6, 3, 1] - # to dlpack / torch +def test_conversions(): + H, W, C = 2, 2, 3 + data = [i for i in range(H * W * C)] + cv_tensor = cvTensor([H, W, C], data) + + # to dlpack / torch / numpy dlpack = K.cvtensor_to_dlpack(cv_tensor) th_tensor = torch.utils.dlpack.from_dlpack(dlpack) assert [x for x in th_tensor.shape] == cv_tensor.shape + + th_tensor_hat = torch.utils.dlpack.from_dlpack(cv_tensor) + torch.testing.assert_close(th_tensor, th_tensor_hat) + + np_array = np._from_dlpack(cv_tensor) + np.testing.assert_array_equal(np_array, th_tensor_hat.numpy()) \ No newline at end of file From eb92751068285c04b2de659674fc799cffc2f8b2 Mon Sep 17 00:00:00 2001 From: Edgar Riba Date: Sun, 13 Mar 2022 18:54:51 +0100 Subject: [PATCH 2/5] update docs --- README.md | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/README.md b/README.md index 95c063e3..efb20a42 100644 --- a/README.md +++ b/README.md @@ -24,15 +24,15 @@ The visualisation API is based on `vviz`: https://github.com/strasdat/vviz import kornia_rs as K from kornia_rs import Tensor as cvTensor - img_path: Path = DATA_DIR / "dog.jpeg" - cv_tensor: cvTensor = K.read_image_jpeg(str(img_path.absolute())) + cv_tensor: cvTensor = K.read_image_jpeg("dog.jpeg") assert cv_tensor.shape == [195, 258, 3] - # convert to dlpack to import to torch - # NOTE: later we will support to numpy, jax and mxnet. - dlpack = K.cvtensor_to_dlpack(cv_tensor) - th_tensor = torch.utils.dlpack.from_dlpack(dlpack) + # convert to dlpack to import to torch and numpy + # NOTE: later we will support to jax and mxnet. + th_tensor = torch.utils.dlpack.from_dlpack(cv_tensor) + np_tensor = np._from_dlpack(cv_tensor) assert th_tensor.shape == (195, 258, 3) + assert np_tensor.shape == (195, 258, 3) ``` ## TODO From 29c93961ba5baa09c992bcba731a6c9ffc1f28eb Mon Sep 17 00:00:00 2001 From: Edgar Riba Date: Sun, 13 Mar 2022 19:52:16 +0100 Subject: [PATCH 3/5] empty the desctructor --- src/tensor.rs | 20 +++++++++++++------- test/test_tensor.py | 10 +++++++--- 2 files changed, 20 insertions(+), 10 deletions(-) diff --git a/src/tensor.rs b/src/tensor.rs index 50ae6e7b..85c5194a 100644 --- a/src/tensor.rs +++ b/src/tensor.rs @@ -5,14 +5,15 @@ pub mod cv { use crate::dlpack_py::{cvtensor_to_dltensor, cvtensor_to_dlpack}; use std::ffi::{c_void}; - unsafe extern "C" fn deleter(x: *mut dlpack::DLManagedTensor) { + // in our case we don not want to delete the data + unsafe extern "C" fn deleter(_: *mut dlpack::DLManagedTensor) { // println!("DLManagedTensor deleter"); - let ctx = (*x).manager_ctx as *mut Tensor; - ctx.drop_in_place(); - (*x).dl_tensor.shape.drop_in_place(); - (*x).dl_tensor.strides.drop_in_place(); - x.drop_in_place(); + //let ctx = (*x).manager_ctx as *mut Tensor; + //ctx.drop_in_place(); + //(*x).dl_tensor.shape.drop_in_place(); + //(*x).dl_tensor.strides.drop_in_place(); + //x.drop_in_place(); } fn get_strides_from_shape(shape: &[i64]) -> Vec { @@ -28,6 +29,11 @@ pub mod cv { strides } + struct CvDlmTensor<'a> { + pub handle: &'a Tensor, + pub tensor: dlpack::DLManagedTensor, + } + #[pyclass] #[derive(Debug, Clone, PartialEq)] pub struct Tensor { @@ -68,7 +74,7 @@ pub mod cv { pub fn to_dlpack(&self) -> dlpack::DLManagedTensor { // we need to clone to avoid race conditions // TODO: check how to avoid that - let tensor_bx = Box::new(self.clone()); + let tensor_bx = Box::new(self); let dl_tensor = cvtensor_to_dltensor(&tensor_bx); // create dlpack managed tensor diff --git a/test/test_tensor.py b/test/test_tensor.py index e4ee0d87..908461c3 100644 --- a/test/test_tensor.py +++ b/test/test_tensor.py @@ -23,8 +23,12 @@ def test_conversions(): th_tensor = torch.utils.dlpack.from_dlpack(dlpack) assert [x for x in th_tensor.shape] == cv_tensor.shape - th_tensor_hat = torch.utils.dlpack.from_dlpack(cv_tensor) - torch.testing.assert_close(th_tensor, th_tensor_hat) +def test_conversions2(): + H, W, C = 2, 2, 3 + data = [i for i in range(H * W * C)] + cv_tensor = cvTensor([H, W, C], data) + # to dlpack / torch / numpy + th_tensor = torch.utils.dlpack.from_dlpack(cv_tensor) np_array = np._from_dlpack(cv_tensor) - np.testing.assert_array_equal(np_array, th_tensor_hat.numpy()) \ No newline at end of file + np.testing.assert_array_equal(np_array, th_tensor.numpy()) \ No newline at end of file From d799a3d027508503550767962a562aaa006e2dbf Mon Sep 17 00:00:00 2001 From: Edgar Riba Date: Sun, 13 Mar 2022 20:07:40 +0100 Subject: [PATCH 4/5] apply clippy! --- src/dlpack_py.rs | 1 - src/tensor.rs | 22 +++++++--------------- 2 files changed, 7 insertions(+), 16 deletions(-) diff --git a/src/dlpack_py.rs b/src/dlpack_py.rs index 03457147..73e9dba5 100644 --- a/src/dlpack_py.rs +++ b/src/dlpack_py.rs @@ -45,7 +45,6 @@ pub fn cvtensor_to_dltensor(x: &cv::Tensor) -> dlpack::DLTensor { #[pyfunction] pub fn cvtensor_to_dlpack(x: &cv::Tensor) -> PyResult<*mut pyo3::ffi::PyObject> { - let dlm_tensor: dlpack::DLManagedTensor = x.to_dlpack(); let dlm_tensor_bx = Box::new(dlm_tensor); diff --git a/src/tensor.rs b/src/tensor.rs index 85c5194a..bf2aad58 100644 --- a/src/tensor.rs +++ b/src/tensor.rs @@ -1,9 +1,9 @@ pub mod cv { - use pyo3::prelude::*; use crate::dlpack; - use crate::dlpack_py::{cvtensor_to_dltensor, cvtensor_to_dlpack}; - use std::ffi::{c_void}; + use crate::dlpack_py::{cvtensor_to_dlpack, cvtensor_to_dltensor}; + use pyo3::prelude::*; + use std::ffi::c_void; // in our case we don not want to delete the data unsafe extern "C" fn deleter(_: *mut dlpack::DLManagedTensor) { @@ -29,11 +29,6 @@ pub mod cv { strides } - struct CvDlmTensor<'a> { - pub handle: &'a Tensor, - pub tensor: dlpack::DLManagedTensor, - } - #[pyclass] #[derive(Debug, Clone, PartialEq)] pub struct Tensor { @@ -59,7 +54,7 @@ pub mod cv { #[pyo3(name = "__dlpack__")] pub fn to_dlpack_py(&self) -> PyResult<*mut pyo3::ffi::PyObject> { - return cvtensor_to_dlpack(self); + cvtensor_to_dlpack(self) } #[pyo3(name = "__dlpack_device__")] @@ -72,21 +67,18 @@ pub mod cv { impl Tensor { pub fn to_dlpack(&self) -> dlpack::DLManagedTensor { - // we need to clone to avoid race conditions - // TODO: check how to avoid that let tensor_bx = Box::new(self); let dl_tensor = cvtensor_to_dltensor(&tensor_bx); // create dlpack managed tensor - let dlm_tensor = dlpack::DLManagedTensor { + + dlpack::DLManagedTensor { dl_tensor, manager_ctx: Box::into_raw(tensor_bx) as *mut c_void, deleter: Some(deleter), - }; - dlm_tensor + } } } - } // namespace cv // TODO(carlos): enable tests later From c301d30d913f622d897ec185d9443a5fea504665 Mon Sep 17 00:00:00 2001 From: Edgar Riba Date: Sun, 13 Mar 2022 20:39:13 +0100 Subject: [PATCH 5/5] Update test_io.py --- test/test_io.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/test/test_io.py b/test/test_io.py index 7f7d7d4a..1121918d 100644 --- a/test/test_io.py +++ b/test/test_io.py @@ -4,7 +4,6 @@ from kornia_rs import Tensor as cvTensor import torch -import jax import numpy as np DATA_DIR = Path(__file__).parent / "data" @@ -38,4 +37,4 @@ def test_read_image_rs(): # convert to dlpack to import to torch dlpack = K.cvtensor_to_dlpack(cv_tensor) th_tensor = torch.utils.dlpack.from_dlpack(dlpack) - assert th_tensor.shape == (195, 258, 3) \ No newline at end of file + assert th_tensor.shape == (195, 258, 3)