teeny_cuda/device/
mod.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::{
20    device::{
21        context::DeviceInfo,
22        program::{ArgVisitor, Kernel, KernelArgs},
23        {Device, LaunchConfig},
24    },
25    dtype::Num,
26};
27
28use crate::{
29    cuda,
30    device::buffer::CudaBuffer,
31    device::program::CudaProgram,
32    errors::{Error, Result},
33};
34
35/// Device memory buffers.
36pub mod buffer;
37/// Device/context management.
38pub mod context;
39/// Memory-related helpers.
40pub mod mem;
41/// Compiled kernel programs.
42pub mod program;
43
44/// Packs kernel arguments into the `void**` array expected by `cuLaunchKernel`.
45///
46/// Each argument's value is stored as raw bytes in `values`. After visiting all
47/// args, `as_ptrs()` returns a mutable slice of `*mut c_void` pointing into
48/// those buffers — the slice lifetime is tied to `self`.
49pub struct CudaArgPacker {
50    values: Vec<Vec<u8>>,
51}
52
53impl Default for CudaArgPacker {
54    fn default() -> Self {
55        Self::new()
56    }
57}
58
59impl CudaArgPacker {
60    /// Creates an empty argument packer.
61    pub fn new() -> Self {
62        Self { values: Vec::new() }
63    }
64
65    fn push_bytes(&mut self, bytes: &[u8]) {
66        self.values.push(bytes.to_vec());
67    }
68
69    /// Returns a `Vec` of `*mut c_void` pointers, one per argument, each
70    /// pointing at the argument's value buffer. The caller must not outlive
71    /// `self`.
72    fn as_ptrs(&mut self) -> Vec<*mut core::ffi::c_void> {
73        self.values
74            .iter_mut()
75            .map(|v| v.as_mut_ptr().cast())
76            .collect()
77    }
78}
79
80impl ArgVisitor for CudaArgPacker {
81    fn visit_ptr(&mut self, ptr: *mut core::ffi::c_void) {
82        self.push_bytes(&(ptr as usize).to_ne_bytes());
83    }
84    fn visit_bool(&mut self, val: bool) {
85        self.push_bytes(&[val as u8]);
86    }
87    fn visit_i8(&mut self, val: i8) {
88        self.push_bytes(&val.to_ne_bytes());
89    }
90    fn visit_i16(&mut self, val: i16) {
91        self.push_bytes(&val.to_ne_bytes());
92    }
93    fn visit_i32(&mut self, val: i32) {
94        self.push_bytes(&val.to_ne_bytes());
95    }
96    fn visit_i64(&mut self, val: i64) {
97        self.push_bytes(&val.to_ne_bytes());
98    }
99    fn visit_u8(&mut self, val: u8) {
100        self.push_bytes(&[val]);
101    }
102    fn visit_u16(&mut self, val: u16) {
103        self.push_bytes(&val.to_ne_bytes());
104    }
105    fn visit_u32(&mut self, val: u32) {
106        self.push_bytes(&val.to_ne_bytes());
107    }
108    fn visit_u64(&mut self, val: u64) {
109        self.push_bytes(&val.to_ne_bytes());
110    }
111    fn visit_f32(&mut self, val: f32) {
112        self.push_bytes(&val.to_ne_bytes());
113    }
114    fn visit_f64(&mut self, val: f64) {
115        self.push_bytes(&val.to_ne_bytes());
116    }
117}
118
119/// A kernel launch's grid/block/cluster dimensions (in `(x, y, z)` order).
120pub struct CudaLaunchConfig {
121    /// Grid dimensions (number of blocks per dimension).
122    pub grid: [u32; 3],
123    /// Block dimensions (number of threads per block, per dimension).
124    pub block: [u32; 3],
125    /// Thread-block-cluster dimensions.
126    pub cluster: [u32; 3],
127}
128
129impl LaunchConfig for CudaLaunchConfig {}
130
131/// A CUDA device's static properties, read from `cudaGetDeviceProperties` at
132/// [`CudaDevice::try_new`] time.
133#[derive(Debug, Clone)]
134pub struct CudaDeviceInfo {
135    /// The device's ordinal ID.
136    pub id: i32,
137    /// The device's name (e.g. `"NVIDIA GeForce RTX 5070"`).
138    pub name: String,
139    /// Compute capability major version.
140    pub major: i32,
141    /// Compute capability minor version.
142    pub minor: i32,
143    /// Number of streaming multiprocessors.
144    pub multi_processor_count: i32,
145    /// Total global memory, in bytes.
146    pub total_global_mem: usize,
147    /// Shared memory available per block, in bytes.
148    pub shared_mem_per_block: usize,
149    /// Number of 32-bit registers available per block.
150    pub regs_per_block: i32,
151    /// Warp size in threads.
152    pub warp_size: i32,
153    /// Maximum threads per block.
154    pub max_threads_per_block: i32,
155    /// Maximum resident threads per multiprocessor.
156    pub max_threads_per_multi_processor: i32,
157    /// Maximum resident blocks per multiprocessor.
158    pub max_blocks_per_multi_processor: i32,
159    /// Maximum block size, per dimension.
160    pub max_threads_dim: [i32; 3],
161    /// Maximum grid size, per dimension.
162    pub max_grid_size: [i32; 3],
163    /// Global memory bus width, in bits.
164    pub memory_bus_width: i32,
165    /// L2 cache size, in bytes.
166    pub l2_cache_size: i32,
167    /// Whether the device supports executing multiple kernels concurrently.
168    pub concurrent_kernels: i32,
169}
170
171impl CudaDeviceInfo {
172    /// Builds a `CudaDeviceInfo` from a raw `cudaDeviceProp`.
173    pub fn new(id: i32, props: cuda::cudaDeviceProp) -> Self {
174        let name = unsafe { std::ffi::CStr::from_ptr(props.name.as_ptr()) };
175        let name = name.to_string_lossy().to_string();
176
177        CudaDeviceInfo {
178            id,
179            name,
180            major: props.major,
181            minor: props.minor,
182            multi_processor_count: props.multiProcessorCount,
183            total_global_mem: props.totalGlobalMem,
184            shared_mem_per_block: props.sharedMemPerBlock,
185            regs_per_block: props.regsPerBlock,
186            warp_size: props.warpSize,
187            max_threads_per_block: props.maxThreadsPerBlock,
188            max_threads_per_multi_processor: props.maxThreadsPerMultiProcessor,
189            max_blocks_per_multi_processor: props.maxBlocksPerMultiProcessor,
190            max_threads_dim: props.maxThreadsDim,
191            max_grid_size: props.maxGridSize,
192            memory_bus_width: props.memoryBusWidth,
193            l2_cache_size: props.l2CacheSize,
194            concurrent_kernels: props.concurrentKernels,
195        }
196    }
197}
198impl DeviceInfo for CudaDeviceInfo {
199    type Id = i32;
200
201    fn id(&self) -> Self::Id {
202        self.id
203    }
204
205    fn name(&self) -> &str {
206        &self.name
207    }
208}
209
210/// An open CUDA device and its context. Destroys the context on drop.
211#[derive(Debug, Clone)]
212pub struct CudaDevice<'a> {
213    /// The device's static properties.
214    pub info: CudaDeviceInfo,
215    // Retained for future device-property queries; only `context` is currently used for CUDA
216    // API calls.
217    #[allow(dead_code)]
218    device: cuda::CUdevice,
219    context: cuda::CUcontext,
220    _unused: PhantomData<&'a ()>,
221}
222
223impl<'a> CudaDevice<'a> {
224    /// Opens device `id`, creating a new CUDA context for it.
225    pub fn try_new(id: i32) -> Result<Self> {
226        let device_id = id;
227        let mut device = cuda::CUdevice::default();
228        let status = unsafe { cuda::cuDeviceGet(&mut device, device_id) };
229        if status != cuda::cudaError_enum_CUDA_SUCCESS {
230            return Err(Error::from_cuda_error(status).into());
231        }
232
233        let mut props = cuda::cudaDeviceProp::default();
234        #[cfg(cuda_props_v2)]
235        let status = unsafe { cuda::cudaGetDeviceProperties_v2(&mut props, device_id) };
236        #[cfg(not(cuda_props_v2))]
237        let status = unsafe { cuda::cudaGetDeviceProperties(&mut props, device_id) };
238        if status != cuda::cudaError_enum_CUDA_SUCCESS {
239            return Err(Error::from_cuda_error(status).into());
240        }
241
242        let info = CudaDeviceInfo::new(device_id, props);
243
244        let mut context = cuda::CUcontext::default();
245        let mut params = cuda::CUctxCreateParams::default();
246        let flags = 0;
247        let status = unsafe { cuda::cuCtxCreate_v4(&mut context, &mut params, flags, device) };
248        if status != cuda::cudaError_enum_CUDA_SUCCESS {
249            return Err(Error::from_cuda_error(status).into());
250        }
251
252        Ok(Self {
253            device,
254            context,
255            info,
256            _unused: PhantomData,
257        })
258    }
259
260    /// This device's static properties.
261    pub fn info(&self) -> &CudaDeviceInfo {
262        &self.info
263    }
264}
265
266impl<'a> Drop for CudaDevice<'a> {
267    fn drop(&mut self) {
268        let result = unsafe { cuda::cuCtxDestroy_v2(self.context) };
269        if result != cuda::cudaError_enum_CUDA_SUCCESS {
270            // just log, we can't do anything about it
271            eprintln!("Failed to destroy CUDA context: {}", result);
272        }
273    }
274}
275
276impl<'a> Device<'a> for CudaDevice<'a> {
277    type Buffer<N: Num> = CudaBuffer<'a, N>;
278    type Program<K: teeny_core::device::program::Kernel> = CudaProgram<'a, K>;
279    type LaunchConfig = CudaLaunchConfig;
280
281    fn buffer<N: Num>(&self, count: usize) -> teeny_core::errors::Result<Self::Buffer<N>> {
282        CudaBuffer::try_new(count)
283    }
284
285    fn launch<K: Kernel>(
286        &self,
287        program: &Self::Program<K>,
288        cfg: &Self::LaunchConfig,
289        args: K::Args<'a>,
290    ) -> teeny_core::errors::Result<()> {
291        // Allocate global scratch memory for TMA descriptors if the kernel requires it.
292        // Total scratch = per-CTA scratch * number of CTAs in the launch grid.
293        let num_ctas = (cfg.grid[0] * cfg.grid[1] * cfg.grid[2]) as u64;
294        let scratch_total = program.metadata.global_scratch_size * num_ctas;
295        let mut scratch_ptr: cuda::CUdeviceptr = 0;
296        if scratch_total > 0 {
297            // cuMemAlloc_v2 guarantees 256-byte alignment, which satisfies Triton's
298            // scratch alignment requirement (typically 128 bytes).
299            let alloc_status =
300                unsafe { cuda::cuMemAlloc_v2(&mut scratch_ptr, scratch_total as usize) };
301            if alloc_status != cuda::cudaError_enum_CUDA_SUCCESS {
302                return Err(Error::from_cuda_error(alloc_status).into());
303            }
304            // Zero-initialize the scratch pad so TMA descriptors start in a clean state.
305            unsafe { cuda::cuMemsetD8_v2(scratch_ptr, 0, scratch_total as usize) };
306        }
307
308        let mut packer = CudaArgPacker::new();
309        args.visit_args(&mut packer);
310        // Trailing Triton kernel parameters: global scratch pad + profile scratch pad.
311        packer.visit_ptr(scratch_ptr as *mut std::ffi::c_void); // global scratch pad
312        packer.visit_ptr(std::ptr::null_mut()); // profile scratch pad (unused)
313        // Build the pointer array while `packer` is still alive — both must
314        // remain live for the entire duration of `cuLaunchKernel`.
315        let mut ptrs = packer.as_ptrs();
316
317        let status = unsafe {
318            cuda::cuLaunchKernel(
319                program.function,
320                cfg.grid[0],
321                cfg.grid[1],
322                cfg.grid[2],
323                cfg.block[0],
324                cfg.block[1],
325                cfg.block[2],
326                program.metadata.shared, // dynamic shared memory required by Triton kernel
327                std::ptr::null_mut(),    // hStream (default/null stream)
328                ptrs.as_mut_ptr(),
329                std::ptr::null_mut(), // extra
330            )
331        };
332
333        if status != cuda::cudaError_enum_CUDA_SUCCESS {
334            return Err(Error::from_cuda_error(status).into());
335        }
336
337        // `cuLaunchKernel` returns immediately; the kernel runs asynchronously.
338        // Synchronize here so that any GPU-side fault (bad pointer, out-of-bounds
339        // access) surfaces as a CUDA error code rather than a later SIGSEGV.
340        let sync_status = unsafe { cuda::cuCtxSynchronize() };
341
342        if sync_status != cuda::cudaError_enum_CUDA_SUCCESS {
343            if scratch_ptr != 0 {
344                unsafe { cuda::cuMemFree_v2(scratch_ptr) };
345            }
346            return Err(Error::from_cuda_error(sync_status).into());
347        }
348
349        if scratch_ptr != 0 {
350            unsafe { cuda::cuMemFree_v2(scratch_ptr) };
351        }
352
353        Ok(())
354    }
355}
356
357impl<'a> CudaDevice<'a> {
358    /// Launch a kernel on the given stream without allocating scratch or
359    /// synchronising. Used during CUDA graph capture: the caller must have
360    /// already appended the global-scratch-pad and profile-scratch-pad pointers
361    /// to `packer` before calling this.
362    pub(crate) fn launch_on_stream<K: Kernel>(
363        &self,
364        program: &CudaProgram<'_, K>,
365        cfg: &CudaLaunchConfig,
366        packer: &mut CudaArgPacker,
367        stream: cuda::CUstream,
368    ) -> Result<()> {
369        let mut ptrs = packer.as_ptrs();
370        let status = unsafe {
371            cuda::cuLaunchKernel(
372                program.function,
373                cfg.grid[0],
374                cfg.grid[1],
375                cfg.grid[2],
376                cfg.block[0],
377                cfg.block[1],
378                cfg.block[2],
379                program.metadata.shared,
380                stream,
381                ptrs.as_mut_ptr(),
382                std::ptr::null_mut(),
383            )
384        };
385        if status != cuda::cudaError_enum_CUDA_SUCCESS {
386            Err(Error::from_cuda_error(status).into())
387        } else {
388            Ok(())
389        }
390    }
391
392    /// Launch a pre-loaded kernel using a dynamically-built arg list.
393    ///
394    /// Use this instead of `launch` when the kernel type is erased (e.g. in
395    /// `LoadedModel::forward`) and arguments are packed by `CudaArgPacker`
396    /// via a `RuntimeOp` rather than a static `Kernel::Args` tuple.
397    pub fn launch_with_packer<K: Kernel>(
398        &self,
399        program: &CudaProgram<'_, K>,
400        cfg: &CudaLaunchConfig,
401        packer: &mut CudaArgPacker,
402    ) -> Result<()> {
403        let num_ctas = (cfg.grid[0] * cfg.grid[1] * cfg.grid[2]) as u64;
404        let scratch_total = program.metadata.global_scratch_size * num_ctas;
405        let mut scratch_ptr: cuda::CUdeviceptr = 0;
406        if scratch_total > 0 {
407            let alloc_status =
408                unsafe { cuda::cuMemAlloc_v2(&mut scratch_ptr, scratch_total as usize) };
409            if alloc_status != cuda::cudaError_enum_CUDA_SUCCESS {
410                return Err(Error::from_cuda_error(alloc_status).into());
411            }
412            unsafe { cuda::cuMemsetD8_v2(scratch_ptr, 0, scratch_total as usize) };
413        }
414
415        packer.visit_ptr(scratch_ptr as *mut std::ffi::c_void);
416        packer.visit_ptr(std::ptr::null_mut());
417        let mut ptrs = packer.as_ptrs();
418
419        let status = unsafe {
420            cuda::cuLaunchKernel(
421                program.function,
422                cfg.grid[0],
423                cfg.grid[1],
424                cfg.grid[2],
425                cfg.block[0],
426                cfg.block[1],
427                cfg.block[2],
428                program.metadata.shared,
429                std::ptr::null_mut(),
430                ptrs.as_mut_ptr(),
431                std::ptr::null_mut(),
432            )
433        };
434
435        if status != cuda::cudaError_enum_CUDA_SUCCESS {
436            if scratch_ptr != 0 {
437                unsafe { cuda::cuMemFree_v2(scratch_ptr) };
438            }
439            return Err(Error::from_cuda_error(status).into());
440        }
441
442        let sync_status = unsafe { cuda::cuCtxSynchronize() };
443        if scratch_ptr != 0 {
444            unsafe { cuda::cuMemFree_v2(scratch_ptr) };
445        }
446
447        if sync_status != cuda::cudaError_enum_CUDA_SUCCESS {
448            return Err(Error::from_cuda_error(sync_status).into());
449        }
450        Ok(())
451    }
452}
teeny_cuda/device/mod.rs

teeny_cuda/device/
mod.rs
