teeny_cuda/device/
buffer.rs

1/*
2 * Copyright (c) 2026 Teenygrad.
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 *   http://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 */
16
17use std::marker::PhantomData;
18
19use teeny_core::{device::buffer::Buffer, dtype::Num};
20
21use crate::device::mem::{self, DevicePtr};
22use crate::errors::{Error, Result};
23
24/// A device-side buffer holding `count` elements of type `N`.
25///
26/// The allocation size is derived from `N::BITS`: `count * N::BITS / 8` bytes.
27/// Memory is freed automatically on drop.
28pub struct CudaBuffer<'a, N: Num> {
29    ptr: DevicePtr,
30    count: usize,
31    _unused: PhantomData<&'a ()>,
32    _num: PhantomData<N>,
33}
34
35impl<'a, N: Num> CudaBuffer<'a, N> {
36    /// Allocates a new device buffer for `count` elements of `N`.
37    pub fn try_new(count: usize) -> Result<Self> {
38        let byte_size = count * N::BITS as usize / 8;
39        let ptr = mem::alloc(byte_size)?;
40        Ok(Self {
41            ptr,
42            count,
43            _unused: PhantomData,
44            _num: PhantomData,
45        })
46    }
47
48    /// The underlying device pointer.
49    pub fn as_device_ptr(&self) -> DevicePtr {
50        self.ptr
51    }
52
53    /// The number of elements this buffer holds.
54    pub fn count(&self) -> usize {
55        self.count
56    }
57}
58
59impl<'a, N: Num> Drop for CudaBuffer<'a, N> {
60    fn drop(&mut self) {
61        if let Err(e) = mem::free(self.ptr) {
62            eprintln!("Failed to free CUDA buffer: {e}");
63        }
64    }
65}
66
67impl<'a, N: Num> Buffer<'a, N> for CudaBuffer<'a, N> {
68    fn to_device(&mut self, src: &[N]) -> teeny_core::errors::Result<()> {
69        if src.len() > self.count {
70            return Err(Error::BufferOverflow {
71                src: src.len(),
72                buf: self.count,
73            }
74            .into());
75        }
76        // SAFETY: src slice is valid for src.len() reads by construction.
77        unsafe { mem::copy_h_to_d(self.ptr, src.as_ptr(), src.len()) }
78    }
79
80    fn to_host(&self, dst: &mut [N]) -> teeny_core::errors::Result<()> {
81        // SAFETY: dst slice is valid for dst.len() writes by construction.
82        unsafe { mem::copy_d_to_h(dst.as_mut_ptr(), self.ptr, dst.len()) }
83    }
84}
teeny_cuda/device/buffer.rs

teeny_cuda/device/
buffer.rs
