From 80265451107bcd549ddfe7d6f29542a2bb5eebd7 Mon Sep 17 00:00:00 2001 From: Emil M George Date: Mon, 23 Sep 2024 00:07:33 +0530 Subject: [PATCH 1/7] Add failing tests for TensorAllocator --- crates/kornia-core/src/storage.rs | 74 +++++++++++++++++++++++++++++++ 1 file changed, 74 insertions(+) diff --git a/crates/kornia-core/src/storage.rs b/crates/kornia-core/src/storage.rs index 0eaa127b..c38fccde 100644 --- a/crates/kornia-core/src/storage.rs +++ b/crates/kornia-core/src/storage.rs @@ -253,6 +253,8 @@ mod tests { use super::*; use crate::allocator::CpuAllocator; use std::alloc::Layout; + use std::cell::RefCell; + use std::rc::Rc; #[test] fn test_tensor_storage() -> Result<(), TensorAllocatorError> { @@ -365,4 +367,76 @@ mod tests { assert_eq!(result_vec.capacity(), original_vec_capacity); assert!(std::ptr::eq(result_vec.as_ptr(), original_vec_ptr)); } + + #[test] + fn test_tensor_storage_allocator() { + // A test TensorAllocator that keeps a count of the bytes that are allocated but not yet + // deallocated via the allocator. + #[derive(Clone)] + struct TestAllocator { + bytes_allocated: Rc>, + } + impl TensorAllocator for TestAllocator { + fn alloc(&self, layout: Layout) -> Result<*mut u8, TensorAllocatorError> { + *self.bytes_allocated.borrow_mut() += layout.size() as i32; + CpuAllocator.alloc(layout) + } + fn dealloc(&self, ptr: *mut u8, layout: Layout) { + *self.bytes_allocated.borrow_mut() -= layout.size() as i32; + CpuAllocator.dealloc(ptr, layout) + } + } + + let allocator = TestAllocator { + bytes_allocated: Rc::new(RefCell::new(0)), + }; + let len = 1024; + + // TensorStorage::new() + // Deallocation should happen when `storage` goes out of scope. + { + let _storage = TensorStorage::::new(len, allocator.clone()).unwrap(); + assert_eq!(*allocator.bytes_allocated.borrow(), len as i32); + } + assert_eq!(*allocator.bytes_allocated.borrow(), 0); + + // TensorStorage::new() -> TensorStorage::into_vec() + // TensorStorage::into_vec() consumes the storage and creates a copy (in this case). + // This should cause deallocation of the original memory. + { + let storage = TensorStorage::::new(len, allocator.clone()).unwrap(); + assert_eq!(*allocator.bytes_allocated.borrow(), len as i32); + + let _vec = storage.into_vec(); + assert_eq!(*allocator.bytes_allocated.borrow(), 0); + } + assert_eq!(*allocator.bytes_allocated.borrow(), 0); + + // TensorStorage::from_vec() -> TensorStorage::into_vec() + // TensorStorage::from_vec() currently does not use the custom allocator, so the + // bytes_allocated value should not change. + { + let vec = Vec::::with_capacity(len); + let storage = TensorStorage::::from_vec(vec, allocator.clone()); + assert_eq!(*allocator.bytes_allocated.borrow(), 0); + + let _vec = storage.into_vec(); + assert_eq!(*allocator.bytes_allocated.borrow(), 0); + } + assert_eq!(*allocator.bytes_allocated.borrow(), 0); + + // TensorStorage::from_ptr() + // TensorStorage::from_ptr() does not take ownership of buffer. So the memory should not be + // deallocated when the TensorStorage goes out of scope. + // In this case, the memory will be deallocated when the vector goes out of scope. + { + let mut vec = Vec::::with_capacity(len); + { + let _storage = + unsafe { TensorStorage::::from_ptr(vec.as_mut_ptr(), len, &allocator) }; + assert_eq!(*allocator.bytes_allocated.borrow(), 0); + } + assert_eq!(*allocator.bytes_allocated.borrow(), 0); + } + } } From 90e8c6e383eb8a5d5b8ce43de99f60e23e8ffaae Mon Sep 17 00:00:00 2001 From: Emil M George Date: Sun, 22 Sep 2024 11:36:17 +0530 Subject: [PATCH 2/7] Fix tensor memory leaks --- crates/kornia-core/src/serde.rs | 5 ++-- crates/kornia-core/src/storage.rs | 42 ++++++++++++++++++++++++++----- crates/kornia-core/src/tensor.rs | 4 +-- crates/kornia-core/src/view.rs | 2 +- 4 files changed, 42 insertions(+), 11 deletions(-) diff --git a/crates/kornia-core/src/serde.rs b/crates/kornia-core/src/serde.rs index f72f0069..db2b0625 100644 --- a/crates/kornia-core/src/serde.rs +++ b/crates/kornia-core/src/serde.rs @@ -7,9 +7,10 @@ use crate::{ use serde::ser::SerializeStruct; use serde::Deserialize; -impl serde::Serialize for Tensor +impl serde::Serialize for Tensor where T: serde::Serialize + SafeTensorType, + A: TensorAllocator + 'static, { fn serialize(&self, serializer: S) -> Result where @@ -23,7 +24,7 @@ where } } -impl<'de, T, const N: usize, A: TensorAllocator + Default> serde::Deserialize<'de> +impl<'de, T, const N: usize, A: TensorAllocator + Default + 'static> serde::Deserialize<'de> for Tensor where T: serde::Deserialize<'de> + SafeTensorType, diff --git a/crates/kornia-core/src/storage.rs b/crates/kornia-core/src/storage.rs index c38fccde..9d4c9ccf 100644 --- a/crates/kornia-core/src/storage.rs +++ b/crates/kornia-core/src/storage.rs @@ -18,6 +18,30 @@ impl SafeTensorType for i64 {} impl SafeTensorType for f32 {} impl SafeTensorType for f64 {} +/// Represents the owner of custom Arrow Buffer memory allocations. +/// +/// This struct is used to facilitate the automatic deallocation of the memory it owns, +/// using the `Drop` trait. +pub struct TensorCustomAllocationOwner { + /// The allocator used to allocate the tensor storage. + alloc: A, + /// The layout used for the allocation. + layout: Layout, + /// The pointer to the allocated memory + ptr: *const u8, +} + +// SAFETY: TensorCustomAllocationOwner is never modifed from multiple threads. +impl std::panic::RefUnwindSafe for TensorCustomAllocationOwner {} +unsafe impl Sync for TensorCustomAllocationOwner {} +unsafe impl Send for TensorCustomAllocationOwner {} + +impl Drop for TensorCustomAllocationOwner { + fn drop(&mut self) { + self.alloc.dealloc(self.ptr as *mut u8, self.layout); + } +} + /// Represents a contiguous memory region that can be shared with other buffers and across thread boundaries. /// /// This struct provides methods to create, access, and manage tensor storage using a custom allocator. @@ -35,9 +59,10 @@ where alloc: A, } -impl TensorStorage +impl TensorStorage where T: SafeTensorType + Clone, + A: TensorAllocator + 'static, { /// Creates a new tensor storage with the given length and allocator. /// @@ -51,8 +76,13 @@ where /// A new tensor storage if successful, otherwise an error. pub fn new(len: usize, alloc: A) -> Result { // allocate memory for tensor storage - let ptr = - alloc.alloc(Layout::array::(len).map_err(TensorAllocatorError::LayoutError)?)?; + let layout = Layout::array::(len).map_err(TensorAllocatorError::LayoutError)?; + let ptr = alloc.alloc(layout)?; + let owner = TensorCustomAllocationOwner { + alloc: alloc.clone(), + layout, + ptr, + }; // create the buffer let buffer = unsafe { @@ -60,7 +90,7 @@ where Buffer::from_custom_allocation( NonNull::new_unchecked(ptr), len * std::mem::size_of::(), - Arc::new(Vec::::with_capacity(len)), + Arc::new(owner), ) }; @@ -223,7 +253,7 @@ where let buffer = Buffer::from_custom_allocation( NonNull::new_unchecked(ptr as *mut u8), len * std::mem::size_of::(), - Arc::new(Vec::::with_capacity(len)), + Arc::new(()), ); // create tensor storage @@ -238,7 +268,7 @@ where impl Clone for TensorStorage where T: SafeTensorType + Clone, - A: TensorAllocator + Clone, + A: TensorAllocator + Clone + 'static, { fn clone(&self) -> Self { let mut new_storage = Self::new(self.len(), self.alloc.clone()) diff --git a/crates/kornia-core/src/tensor.rs b/crates/kornia-core/src/tensor.rs index f8a4d9d1..d17ce070 100644 --- a/crates/kornia-core/src/tensor.rs +++ b/crates/kornia-core/src/tensor.rs @@ -90,7 +90,7 @@ where impl Tensor where T: SafeTensorType, - A: TensorAllocator, + A: TensorAllocator + 'static, { /// Create a new `Tensor` with uninitialized data. /// @@ -875,7 +875,7 @@ where impl Clone for Tensor where T: SafeTensorType + Clone, - A: TensorAllocator + Clone, + A: TensorAllocator + Clone + 'static, { fn clone(&self) -> Self { Self { diff --git a/crates/kornia-core/src/view.rs b/crates/kornia-core/src/view.rs index e777d898..6cafd2df 100644 --- a/crates/kornia-core/src/view.rs +++ b/crates/kornia-core/src/view.rs @@ -14,7 +14,7 @@ pub struct TensorView<'a, T: SafeTensorType, const N: usize, A: TensorAllocator> pub strides: [usize; N], } -impl<'a, T: SafeTensorType, const N: usize, A: TensorAllocator> TensorView<'a, T, N, A> { +impl<'a, T: SafeTensorType, const N: usize, A: TensorAllocator + 'static> TensorView<'a, T, N, A> { /// Returns the data slice of the tensor. #[inline] pub fn as_slice(&self) -> &[T] { From ef8246a8f8d3d5f3fc4c474a907d8bf12e15647b Mon Sep 17 00:00:00 2001 From: Emil M George Date: Wed, 25 Sep 2024 19:36:36 +0530 Subject: [PATCH 3/7] Use NonNull instead of raw pointer --- crates/kornia-core/src/storage.rs | 12 ++++-------- 1 file changed, 4 insertions(+), 8 deletions(-) diff --git a/crates/kornia-core/src/storage.rs b/crates/kornia-core/src/storage.rs index 9d4c9ccf..a4fd1481 100644 --- a/crates/kornia-core/src/storage.rs +++ b/crates/kornia-core/src/storage.rs @@ -28,7 +28,7 @@ pub struct TensorCustomAllocationOwner { /// The layout used for the allocation. layout: Layout, /// The pointer to the allocated memory - ptr: *const u8, + ptr: NonNull, } // SAFETY: TensorCustomAllocationOwner is never modifed from multiple threads. @@ -38,7 +38,7 @@ unsafe impl Send for TensorCustomAllocationOwner {} impl Drop for TensorCustomAllocationOwner { fn drop(&mut self) { - self.alloc.dealloc(self.ptr as *mut u8, self.layout); + self.alloc.dealloc(self.ptr.as_ptr(), self.layout); } } @@ -77,7 +77,7 @@ where pub fn new(len: usize, alloc: A) -> Result { // allocate memory for tensor storage let layout = Layout::array::(len).map_err(TensorAllocatorError::LayoutError)?; - let ptr = alloc.alloc(layout)?; + let ptr = NonNull::new(alloc.alloc(layout)?).ok_or(TensorAllocatorError::NullPointer)?; let owner = TensorCustomAllocationOwner { alloc: alloc.clone(), layout, @@ -87,11 +87,7 @@ where // create the buffer let buffer = unsafe { // SAFETY: `ptr` is non-null and properly aligned, and `len` is the correct size. - Buffer::from_custom_allocation( - NonNull::new_unchecked(ptr), - len * std::mem::size_of::(), - Arc::new(owner), - ) + Buffer::from_custom_allocation(ptr, len * std::mem::size_of::(), Arc::new(owner)) }; Ok(Self { From 7ed404fa501cc8792550e21069cf43b160b0ee04 Mon Sep 17 00:00:00 2001 From: Emil M George Date: Wed, 25 Sep 2024 19:39:17 +0530 Subject: [PATCH 4/7] Return result from test and use ? operator --- crates/kornia-core/src/storage.rs | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/crates/kornia-core/src/storage.rs b/crates/kornia-core/src/storage.rs index a4fd1481..41e08b7f 100644 --- a/crates/kornia-core/src/storage.rs +++ b/crates/kornia-core/src/storage.rs @@ -395,7 +395,7 @@ mod tests { } #[test] - fn test_tensor_storage_allocator() { + fn test_tensor_storage_allocator() -> Result<(), TensorAllocatorError> { // A test TensorAllocator that keeps a count of the bytes that are allocated but not yet // deallocated via the allocator. #[derive(Clone)] @@ -421,7 +421,7 @@ mod tests { // TensorStorage::new() // Deallocation should happen when `storage` goes out of scope. { - let _storage = TensorStorage::::new(len, allocator.clone()).unwrap(); + let _storage = TensorStorage::::new(len, allocator.clone())?; assert_eq!(*allocator.bytes_allocated.borrow(), len as i32); } assert_eq!(*allocator.bytes_allocated.borrow(), 0); @@ -430,7 +430,7 @@ mod tests { // TensorStorage::into_vec() consumes the storage and creates a copy (in this case). // This should cause deallocation of the original memory. { - let storage = TensorStorage::::new(len, allocator.clone()).unwrap(); + let storage = TensorStorage::::new(len, allocator.clone())?; assert_eq!(*allocator.bytes_allocated.borrow(), len as i32); let _vec = storage.into_vec(); @@ -464,5 +464,7 @@ mod tests { } assert_eq!(*allocator.bytes_allocated.borrow(), 0); } + + Ok(()) } } From 9458c6055e32ef66b76280ea70dd896dc622d322 Mon Sep 17 00:00:00 2001 From: Emil M George Date: Thu, 26 Sep 2024 20:39:08 +0530 Subject: [PATCH 5/7] Add ptr sanity checks in tests --- crates/kornia-core/src/storage.rs | 22 ++++++++++++++++------ 1 file changed, 16 insertions(+), 6 deletions(-) diff --git a/crates/kornia-core/src/storage.rs b/crates/kornia-core/src/storage.rs index 41e08b7f..1ba6ffe2 100644 --- a/crates/kornia-core/src/storage.rs +++ b/crates/kornia-core/src/storage.rs @@ -442,12 +442,18 @@ mod tests { // TensorStorage::from_vec() currently does not use the custom allocator, so the // bytes_allocated value should not change. { - let vec = Vec::::with_capacity(len); - let storage = TensorStorage::::from_vec(vec, allocator.clone()); + let original_vec = Vec::::with_capacity(len); + let original_vec_ptr = original_vec.as_ptr(); + let original_vec_capacity = original_vec.capacity(); + + let storage = TensorStorage::::from_vec(original_vec, allocator.clone()); assert_eq!(*allocator.bytes_allocated.borrow(), 0); - let _vec = storage.into_vec(); + let result_vec = storage.into_vec(); assert_eq!(*allocator.bytes_allocated.borrow(), 0); + + assert_eq!(result_vec.capacity(), original_vec_capacity); + assert!(std::ptr::eq(result_vec.as_ptr(), original_vec_ptr)); } assert_eq!(*allocator.bytes_allocated.borrow(), 0); @@ -456,11 +462,15 @@ mod tests { // deallocated when the TensorStorage goes out of scope. // In this case, the memory will be deallocated when the vector goes out of scope. { - let mut vec = Vec::::with_capacity(len); + let mut original_vec = Vec::::with_capacity(len); + let original_ptr = original_vec.as_ptr(); { - let _storage = - unsafe { TensorStorage::::from_ptr(vec.as_mut_ptr(), len, &allocator) }; + let storage = unsafe { + TensorStorage::::from_ptr(original_vec.as_mut_ptr(), len, &allocator) + }; assert_eq!(*allocator.bytes_allocated.borrow(), 0); + + assert_eq!(storage.as_ptr(), original_ptr); } assert_eq!(*allocator.bytes_allocated.borrow(), 0); } From 27ddb676e1bd722f325514054aefa74f39c953ec Mon Sep 17 00:00:00 2001 From: Emil M George Date: Thu, 26 Sep 2024 23:21:46 +0530 Subject: [PATCH 6/7] Use Global allocator in CpuAllocator --- crates/kornia-core/src/allocator.rs | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/crates/kornia-core/src/allocator.rs b/crates/kornia-core/src/allocator.rs index 823b8098..14234c59 100644 --- a/crates/kornia-core/src/allocator.rs +++ b/crates/kornia-core/src/allocator.rs @@ -1,4 +1,5 @@ -use std::alloc::{GlobalAlloc, Layout, System}; +use std::alloc; +use std::alloc::Layout; use thiserror::Error; @@ -55,7 +56,7 @@ impl TensorAllocator for CpuAllocator { /// /// A non-null pointer to the allocated memory if successful, otherwise an error. fn alloc(&self, layout: Layout) -> Result<*mut u8, TensorAllocatorError> { - let ptr = unsafe { System.alloc(layout) }; + let ptr = unsafe { alloc::alloc(layout) }; if ptr.is_null() { Err(TensorAllocatorError::NullPointer)? } @@ -74,7 +75,7 @@ impl TensorAllocator for CpuAllocator { /// The pointer must be non-null and the layout must be correct. #[allow(clippy::not_unsafe_ptr_arg_deref)] fn dealloc(&self, ptr: *mut u8, layout: Layout) { - unsafe { System.dealloc(ptr, layout) } + unsafe { alloc::dealloc(ptr, layout) } } } From 46db49f5cf13e1ac466841404cc06ef2b3090855 Mon Sep 17 00:00:00 2001 From: Emil M George Date: Fri, 27 Sep 2024 16:00:08 +0530 Subject: [PATCH 7/7] Wrap alloc in an Arc<> --- crates/kornia-core/src/storage.rs | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/crates/kornia-core/src/storage.rs b/crates/kornia-core/src/storage.rs index 1ba6ffe2..c9b6ec32 100644 --- a/crates/kornia-core/src/storage.rs +++ b/crates/kornia-core/src/storage.rs @@ -24,7 +24,7 @@ impl SafeTensorType for f64 {} /// using the `Drop` trait. pub struct TensorCustomAllocationOwner { /// The allocator used to allocate the tensor storage. - alloc: A, + alloc: Arc, /// The layout used for the allocation. layout: Layout, /// The pointer to the allocated memory @@ -56,7 +56,7 @@ where /// The buffer containing the tensor storage. data: ScalarBuffer, /// The allocator used to allocate the tensor storage. - alloc: A, + alloc: Arc, } impl TensorStorage @@ -78,6 +78,7 @@ where // allocate memory for tensor storage let layout = Layout::array::(len).map_err(TensorAllocatorError::LayoutError)?; let ptr = NonNull::new(alloc.alloc(layout)?).ok_or(TensorAllocatorError::NullPointer)?; + let alloc = Arc::new(alloc); let owner = TensorCustomAllocationOwner { alloc: alloc.clone(), layout, @@ -125,7 +126,7 @@ where // create tensor storage Self { data: buffer.into(), - alloc, + alloc: Arc::new(alloc), } } @@ -255,7 +256,7 @@ where // create tensor storage Self { data: buffer.into(), - alloc: alloc.clone(), + alloc: Arc::new(alloc.clone()), } } } @@ -267,7 +268,7 @@ where A: TensorAllocator + Clone + 'static, { fn clone(&self) -> Self { - let mut new_storage = Self::new(self.len(), self.alloc.clone()) + let mut new_storage = Self::new(self.len(), (*self.alloc).clone()) .expect("Failed to allocate memory for cloned TensorStorage"); new_storage.as_mut_slice().clone_from_slice(self.as_slice()); new_storage