How does Buffer::from_custom_allocation work ?

[edgarriba]

hi, I have the kornia-rs crate where I implemented a custom TensorStorage struct out of ScalarBuffer and I have some issues while extending a method TensorStorage::from_vec. This struct has the particularity that holds an Allocator which the intention is that the user can specify, e.g for CpuTensor or CudaTensor (or more advanced devices/backends).

Initally, I had implemented this using the Buffer::from_custom_allocation method aiming to create the storage with a zero-copy cost. However, I recently discovered after adding some tests that something fishy is going on with the memory pointers of the vector.

I can quickly illustrate what I have and what I'm trying to figure out. From kornia/kornia-rs#125

pub struct TensorStorage<T, A: TensorAllocator>
where
    T: SafeTensorType,
{
    /// The buffer containing the tensor storage.
    data: ScalarBuffer<T>,
    /// The allocator used to allocate the tensor storage.
    alloc: A,
}

impl<T, A: TensorAllocator> TensorStorage<T, A>
where
    T: SafeTensorType + Clone,
{
pub fn from_vec(vec: Vec<T>, alloc: A) -> Self {
        // NOTE: this is a temporary solution until we have a custom allocator for the buffer
        // create immutable buffer from vec
        //let buffer = unsafe {
        //    // SAFETY: `vec` is properly aligned and has the correct length.
        //    Buffer::from_custom_allocation(
        //        NonNull::new_unchecked(vec.as_ptr() as *mut u8),
        //        vec.len() * std::mem::size_of::<T>(),
        //        Arc::new(vec),
        //    )
        //};

        // create immutable buffer from vec
        // NOTE: this is a temporary solution until we have a custom allocator for the buffer
        let buffer = Buffer::from_vec(vec);

        // create tensor storage
        Self {
            data: buffer.into(),
            alloc,
        }
    }
}

Then the crashing test

#[test]
    fn test_tensor_storage_into_vec() {
        let allocator = CpuAllocator;
        let original_vec = vec![1, 2, 3, 4, 5];
        let original_vec_ptr = original_vec.as_ptr();
        let original_vec_capacity = original_vec.capacity();

        let storage = TensorStorage::<i32, _>::from_vec(original_vec, allocator);

        // Convert the storage back to a vector
        let result_vec = storage.into_vec();

        // check NO copy
        assert_eq!(result_vec.capacity(), original_vec_capacity);

        // THIS TEST IS FAILING
        assert!(std::ptr::eq(result_vec.as_ptr(), original_vec_ptr));
    }

Thanks in advance!

[jhorstmann]

At a glance, this looks correct. I use it like the following, where v is an Arc<Vec<T, Allocator>>:

let len = v.len() * std::mem::size_of::<T>();
let buffer = unsafe {
    let ptr = NonNull::new_unchecked(v.as_ptr() as *mut u8);
    Buffer::from_custom_allocation(ptr, len, v)
};

I'm being careful here to not use any references to the vector while a mutable pointer exists, as that could lead to undefined behavior. But that does not seem to be the problem here. Maybe the issue is with the into_vec method instead?

[edgarriba]

this is the implementation of into_vec: https://github.com/kornia/kornia-rs/blob/65ab143f1b14cc991b9b01cd01f406aaaab57060/crates/kornia-core/src/storage.rs#L123-L134

    pub fn into_vec(self) -> Vec<T> {
        match self.data.into_inner().into_vec() {
            Ok(vec) => vec,
            Err(buf) => unsafe {
                std::slice::from_raw_parts(
                    buf.as_ptr() as *const T,
                    buf.len() / std::mem::size_of::<T>(),
                )
                .to_vec()
            },
        }
    }

[jhorstmann]

Buffer::into_vec cannot really support custom allocations without copying:

arrow-rs/arrow-buffer/src/buffer/immutable.rs

Line 318 in 704f90b

_ => return Err(self), // Custom allocation

This usecase might get supported by #6336 or a followup of that.