oxirs_vec/gpu/
buffer.rs

1//! GPU memory buffer management
2//!
3//! The published `oxirs-vec` build is 100% Pure Rust: this buffer is backed by
4//! host memory. Real CUDA device memory (`cudaMalloc`/`cudaMemcpy`/`cudaFree`)
5//! lives in the quarantined `oxirs-vec-adapter-cuda` crate (publish = false) per
6//! the COOLJAPAN Pure Rust Policy v2.
7
8use anyhow::{anyhow, Result};
9
10/// GPU memory buffer for vector data
11#[derive(Debug)]
12pub struct GpuBuffer {
13    ptr: *mut f32,
14    size: usize,
15    device_id: i32,
16}
17
18unsafe impl Send for GpuBuffer {}
19unsafe impl Sync for GpuBuffer {}
20
21impl GpuBuffer {
22    pub fn new(size: usize, device_id: i32) -> Result<Self> {
23        let ptr = Self::allocate_gpu_memory(size * std::mem::size_of::<f32>(), device_id)?;
24        Ok(Self {
25            ptr: ptr as *mut f32,
26            size,
27            device_id,
28        })
29    }
30
31    pub fn copy_from_host(&mut self, data: &[f32]) -> Result<()> {
32        if data.len() > self.size {
33            return Err(anyhow!("Data size exceeds buffer capacity"));
34        }
35        self.copy_host_to_device(data.as_ptr(), self.ptr, std::mem::size_of_val(data))
36    }
37
38    pub fn copy_to_host(&self, data: &mut [f32]) -> Result<()> {
39        if data.len() > self.size {
40            return Err(anyhow!("Host buffer too small"));
41        }
42        self.copy_device_to_host(self.ptr, data.as_mut_ptr(), std::mem::size_of_val(data))
43    }
44
45    #[allow(unused_variables)]
46    fn allocate_gpu_memory(size: usize, device_id: i32) -> Result<*mut u8> {
47        // Host-memory allocation (Pure Rust). CUDA-backed allocation is provided
48        // by oxirs-vec-adapter-cuda.
49        let layout = std::alloc::Layout::from_size_align(size, std::mem::align_of::<f32>())
50            .map_err(|e| anyhow!("Invalid memory layout: {}", e))?;
51        unsafe {
52            let ptr = std::alloc::alloc(layout);
53            if ptr.is_null() {
54                return Err(anyhow!("Failed to allocate memory"));
55            }
56            Ok(ptr)
57        }
58    }
59
60    fn copy_host_to_device(&self, src: *const f32, dst: *mut f32, size: usize) -> Result<()> {
61        // Pure Rust build: host-to-host copy.
62        unsafe {
63            std::ptr::copy_nonoverlapping(src, dst, size / std::mem::size_of::<f32>());
64        }
65        Ok(())
66    }
67
68    fn copy_device_to_host(&self, src: *const f32, dst: *mut f32, size: usize) -> Result<()> {
69        // Pure Rust build: host-to-host copy.
70        unsafe {
71            std::ptr::copy_nonoverlapping(src, dst, size / std::mem::size_of::<f32>());
72        }
73        Ok(())
74    }
75
76    pub fn ptr(&self) -> *mut f32 {
77        self.ptr
78    }
79
80    pub fn size(&self) -> usize {
81        self.size
82    }
83
84    pub fn device_id(&self) -> i32 {
85        self.device_id
86    }
87
88    pub fn is_valid(&self) -> bool {
89        !self.ptr.is_null()
90    }
91
92    /// Zero out the buffer
93    pub fn zero(&mut self) -> Result<()> {
94        unsafe {
95            std::ptr::write_bytes(self.ptr, 0, self.size);
96        }
97        Ok(())
98    }
99}
100
101impl Drop for GpuBuffer {
102    fn drop(&mut self) {
103        if !self.ptr.is_null() {
104            let layout = std::alloc::Layout::from_size_align(
105                self.size * std::mem::size_of::<f32>(),
106                std::mem::align_of::<f32>(),
107            );
108            if let Ok(layout) = layout {
109                unsafe {
110                    std::alloc::dealloc(self.ptr as *mut u8, layout);
111                }
112            }
113        }
114    }
115}
oxirs_vec/gpu/buffer.rs

oxirs_vec/gpu/
buffer.rs
