openvm_cuda_common/
copy.rs

1use std::ffi::c_void;
2
3use crate::{
4    d_buffer::DeviceBuffer,
5    error::{check, MemCopyError},
6    stream::{cudaStream_t, GpuDeviceCtx},
7};
8
9#[repr(i32)]
10#[non_exhaustive]
11#[allow(non_camel_case_types)]
12#[derive(Debug, Copy, Clone, Hash, PartialEq, Eq)]
13pub enum cudaMemcpyKind {
14    cudaMemcpyHostToHost = 0,
15    cudaMemcpyHostToDevice = 1,
16    cudaMemcpyDeviceToHost = 2,
17    cudaMemcpyDeviceToDevice = 3,
18    cudaMemcpyDefault = 4,
19}
20
21#[link(name = "cudart")]
22extern "C" {
23    fn cudaMemcpyAsync(
24        dst: *mut c_void,
25        src: *const c_void,
26        count: usize,
27        kind: cudaMemcpyKind,
28        stream: cudaStream_t,
29    ) -> i32;
30}
31
32/// FFI binding for `cudaMemcpyAsync` on an explicit stream.
33///
34/// # Safety
35/// Must follow the rules of the `cudaMemcpyAsync` function from the CUDA runtime API.
36pub unsafe fn cuda_memcpy_on<const SRC_DEVICE: bool, const DST_DEVICE: bool>(
37    dst: *mut c_void,
38    src: *const c_void,
39    size_bytes: usize,
40    device_ctx: &GpuDeviceCtx,
41) -> Result<(), MemCopyError> {
42    check(unsafe {
43        cudaMemcpyAsync(
44            dst,
45            src,
46            size_bytes,
47            std::mem::transmute::<i32, cudaMemcpyKind>(
48                if DST_DEVICE { 1 } else { 0 } + if SRC_DEVICE { 2 } else { 0 },
49            ),
50            device_ctx.stream.as_raw(),
51        )
52    })
53    .map_err(MemCopyError::from)
54}
55
56// ---- Host -> Device ----
57
58pub trait MemCopyH2D<T> {
59    fn copy_to_on(
60        &self,
61        dst: &mut DeviceBuffer<T>,
62        device_ctx: &GpuDeviceCtx,
63    ) -> Result<(), MemCopyError>;
64    fn to_device_on(&self, device_ctx: &GpuDeviceCtx) -> Result<DeviceBuffer<T>, MemCopyError>;
65}
66
67impl<T> MemCopyH2D<T> for [T] {
68    fn copy_to_on(
69        &self,
70        dst: &mut DeviceBuffer<T>,
71        device_ctx: &GpuDeviceCtx,
72    ) -> Result<(), MemCopyError> {
73        if self.len() > dst.len() {
74            return Err(MemCopyError::SizeMismatch {
75                operation: "copy_to_on",
76                src_len: self.len(),
77                dst_len: dst.len(),
78            });
79        }
80        let size_bytes = std::mem::size_of_val(self);
81        check(unsafe {
82            cudaMemcpyAsync(
83                dst.as_mut_raw_ptr(),
84                self.as_ptr() as *const c_void,
85                size_bytes,
86                cudaMemcpyKind::cudaMemcpyHostToDevice,
87                device_ctx.stream.as_raw(),
88            )
89        })
90        .map_err(MemCopyError::from)
91    }
92
93    fn to_device_on(&self, device_ctx: &GpuDeviceCtx) -> Result<DeviceBuffer<T>, MemCopyError> {
94        let mut dst = DeviceBuffer::with_capacity_on(self.len(), device_ctx);
95        self.copy_to_on(&mut dst, device_ctx)?;
96        Ok(dst)
97    }
98}
99
100// ---- Device -> Host ----
101
102pub trait MemCopyD2H<T> {
103    fn to_host_on(&self, device_ctx: &GpuDeviceCtx) -> Result<Vec<T>, MemCopyError>;
104}
105
106impl<T> MemCopyD2H<T> for DeviceBuffer<T> {
107    fn to_host_on(&self, device_ctx: &GpuDeviceCtx) -> Result<Vec<T>, MemCopyError> {
108        let mut host_vec = Vec::with_capacity(self.len());
109        let size_bytes = std::mem::size_of::<T>() * self.len();
110
111        check(unsafe {
112            cudaMemcpyAsync(
113                host_vec.as_mut_ptr() as *mut c_void,
114                self.as_raw_ptr(),
115                size_bytes,
116                cudaMemcpyKind::cudaMemcpyDeviceToHost,
117                device_ctx.stream.as_raw(),
118            )
119        })?;
120        device_ctx.stream.to_host_sync()?;
121        unsafe {
122            host_vec.set_len(self.len());
123        }
124
125        Ok(host_vec)
126    }
127}
128
129// ---- Device -> Device ----
130
131pub trait MemCopyD2D<T> {
132    fn device_copy_on(&self, device_ctx: &GpuDeviceCtx) -> Result<DeviceBuffer<T>, MemCopyError>;
133    fn device_copy_to_on(
134        &self,
135        dst: &mut DeviceBuffer<T>,
136        device_ctx: &GpuDeviceCtx,
137    ) -> Result<(), MemCopyError>;
138}
139
140impl<T> MemCopyD2D<T> for DeviceBuffer<T> {
141    fn device_copy_on(&self, device_ctx: &GpuDeviceCtx) -> Result<DeviceBuffer<T>, MemCopyError> {
142        let mut dst = DeviceBuffer::<T>::with_capacity_on(self.len(), device_ctx);
143        self.device_copy_to_on(&mut dst, device_ctx)?;
144        Ok(dst)
145    }
146
147    fn device_copy_to_on(
148        &self,
149        dst: &mut DeviceBuffer<T>,
150        device_ctx: &GpuDeviceCtx,
151    ) -> Result<(), MemCopyError> {
152        if self.len() > dst.len() {
153            return Err(MemCopyError::SizeMismatch {
154                operation: "device_copy_to_on",
155                src_len: self.len(),
156                dst_len: dst.len(),
157            });
158        }
159        let size_bytes = std::mem::size_of::<T>() * self.len();
160
161        check(unsafe {
162            cudaMemcpyAsync(
163                dst.as_mut_raw_ptr(),
164                self.as_raw_ptr(),
165                size_bytes,
166                cudaMemcpyKind::cudaMemcpyDeviceToDevice,
167                device_ctx.stream.as_raw(),
168            )
169        })
170        .map_err(MemCopyError::from)
171    }
172}
173
174#[cfg(test)]
175mod tests {
176    use super::*;
177    use crate::d_buffer::DeviceBuffer;
178
179    fn test_ctx() -> GpuDeviceCtx {
180        GpuDeviceCtx::for_current_device().unwrap()
181    }
182
183    #[test]
184    fn test_mem_copy() {
185        let device_ctx = test_ctx();
186        let h = vec![1, 2, 3, 4, 5];
187
188        let d1 = h.to_device_on(&device_ctx).unwrap();
189        let mut d2 = DeviceBuffer::<i32>::with_capacity_on(h.len(), &device_ctx);
190        h.copy_to_on(&mut d2, &device_ctx).unwrap();
191
192        let h1 = d1.to_host_on(&device_ctx).unwrap();
193        let h2 = d2.to_host_on(&device_ctx).unwrap();
194
195        assert_eq!(h, h1, "First device buffer mismatch");
196        assert_eq!(h, h2, "Second device buffer mismatch");
197    }
198}
openvm_cuda_common/copy.rs

openvm_cuda_common/
copy.rs
