teeny_cuda/device/
mem.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
17//! Thin safe wrappers around the CUDA driver memory API.
18
19use crate::cuda;
20use crate::errors::{Error, Result};
21
22/// Opaque CUDA device pointer (a virtual address in device memory).
23pub type DevicePtr = cuda::CUdeviceptr;
24
25/// Allocate `byte_size` bytes of device memory in the current CUDA context.
26pub fn alloc(byte_size: usize) -> Result<DevicePtr> {
27    let mut ptr: DevicePtr = 0;
28    let status = unsafe { cuda::cuMemAlloc_v2(&mut ptr, byte_size) };
29    if status != cuda::cudaError_enum_CUDA_SUCCESS {
30        return Err(Error::from_cuda_error(status).into());
31    }
32    Ok(ptr)
33}
34
35/// Free a device allocation previously returned by [`alloc`].
36pub fn free(ptr: DevicePtr) -> Result<()> {
37    let status = unsafe { cuda::cuMemFree_v2(ptr) };
38    if status != cuda::cudaError_enum_CUDA_SUCCESS {
39        return Err(Error::from_cuda_error(status).into());
40    }
41    Ok(())
42}
43
44/// Copy `count` elements of type `T` from host memory to device memory.
45///
46/// # Safety
47/// `src` must be valid for `count` reads. `dst` must be a valid device
48/// allocation of at least `count * size_of::<T>()` bytes.
49pub unsafe fn copy_h_to_d<T>(dst: DevicePtr, src: *const T, count: usize) -> Result<()> {
50    let status =
51        unsafe { cuda::cuMemcpyHtoD_v2(dst, src.cast(), count * std::mem::size_of::<T>()) };
52    if status != cuda::cudaError_enum_CUDA_SUCCESS {
53        return Err(Error::from_cuda_error(status).into());
54    }
55    Ok(())
56}
57
58/// Copy `count` elements of type `T` from device memory to host memory.
59///
60/// # Safety
61/// `dst` must be valid for `count` writes. `src` must be a valid device
62/// allocation of at least `count * size_of::<T>()` bytes.
63pub unsafe fn copy_d_to_h<T>(dst: *mut T, src: DevicePtr, count: usize) -> Result<()> {
64    let status =
65        unsafe { cuda::cuMemcpyDtoH_v2(dst.cast(), src, count * std::mem::size_of::<T>()) };
66    if status != cuda::cudaError_enum_CUDA_SUCCESS {
67        return Err(Error::from_cuda_error(status).into());
68    }
69    Ok(())
70}
71
72/// Allocate `n_elems` elements of `T` in page-locked (pinned) host memory.
73///
74/// Pinned memory enables direct DMA for `cuMemcpyHtoD`/`cuMemcpyDtoH`,
75/// bypassing the driver's internal staging buffer and achieving full PCIe
76/// bandwidth.
77pub fn alloc_host<T>(n_elems: usize) -> Result<*mut T> {
78    let mut ptr: *mut std::ffi::c_void = std::ptr::null_mut();
79    let status = unsafe { cuda::cuMemAllocHost_v2(&mut ptr, n_elems * std::mem::size_of::<T>()) };
80    if status != cuda::cudaError_enum_CUDA_SUCCESS {
81        return Err(Error::from_cuda_error(status).into());
82    }
83    Ok(ptr.cast())
84}
85
86/// Free a pinned host allocation returned by [`alloc_host`].
87///
88/// # Safety
89/// `ptr` must have been returned by [`alloc_host`] and not yet freed.
90pub unsafe fn free_host<T>(ptr: *mut T) -> Result<()> {
91    let status = unsafe { cuda::cuMemFreeHost(ptr.cast()) };
92    if status != cuda::cudaError_enum_CUDA_SUCCESS {
93        return Err(Error::from_cuda_error(status).into());
94    }
95    Ok(())
96}
97
98/// Copy `num_rows` rows of `row_bytes` bytes from a device buffer with
99/// `src_stride_bytes` row stride to a device buffer with `dst_stride_bytes`
100/// row stride.  Used to depad TMA-aligned output tensors back to tight NCHW.
101///
102/// Issues a single `cuMemcpy2D` rather than one `cuMemcpyDtoD` per row,
103/// letting the driver schedule the whole transfer as one unit.
104pub fn copy_rows_d_to_d(
105    dst: DevicePtr,
106    dst_stride_bytes: usize,
107    src: DevicePtr,
108    src_stride_bytes: usize,
109    row_bytes: usize,
110    num_rows: usize,
111) -> Result<()> {
112    let params = cuda::CUDA_MEMCPY2D {
113        srcMemoryType: cuda::CUmemorytype_enum_CU_MEMORYTYPE_DEVICE,
114        srcDevice: src,
115        srcPitch: src_stride_bytes,
116        dstMemoryType: cuda::CUmemorytype_enum_CU_MEMORYTYPE_DEVICE,
117        dstDevice: dst,
118        dstPitch: dst_stride_bytes,
119        WidthInBytes: row_bytes,
120        Height: num_rows,
121        ..Default::default()
122    };
123    let status = unsafe { cuda::cuMemcpy2D_v2(&params) };
124    if status != cuda::cudaError_enum_CUDA_SUCCESS {
125        return Err(Error::from_cuda_error(status).into());
126    }
127    Ok(())
128}
teeny_cuda/device/mem.rs

teeny_cuda/device/
mem.rs
