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 CudaDevice<'_> {
277    /// Raise the per-function dynamic shared-memory ceiling when Triton asks for
278    /// more than the default 48 KiB carveout.
279    ///
280    /// Without this, `cuLaunchKernel` returns `CUDA_ERROR_INVALID_ARGUMENT` for
281    /// kernels whose `meta:shared` is in (48 KiB, opt-in max] — e.g.
282    /// `conv2d_bn_silu_gemm` at 128×128 tiles (~64 KiB) on sm_120. Tiles that
283    /// still exceed the device opt-in max (~99 KiB on RTX 5070) must be capped
284    /// at compile/dispatch time instead.
285    pub(crate) fn ensure_dynamic_shared(
286        function: cuda::CUfunction,
287        shared_bytes: u32,
288    ) -> Result<()> {
289        if shared_bytes == 0 {
290            return Ok(());
291        }
292        let status = unsafe {
293            cuda::cuFuncSetAttribute(
294                function,
295                cuda::CUfunction_attribute_enum_CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
296                shared_bytes as i32,
297            )
298        };
299        if status != cuda::cudaError_enum_CUDA_SUCCESS {
300            return Err(Error::from_cuda_error(status).into());
301        }
302        Ok(())
303    }
304}
305
306impl<'a> Device<'a> for CudaDevice<'a> {
307    type Buffer<N: Num> = CudaBuffer<'a, N>;
308    type Program<K: teeny_core::device::program::Kernel> = CudaProgram<'a, K>;
309    type LaunchConfig = CudaLaunchConfig;
310
311    fn buffer<N: Num>(&self, count: usize) -> teeny_core::errors::Result<Self::Buffer<N>> {
312        CudaBuffer::try_new(count)
313    }
314
315    fn launch<K: Kernel>(
316        &self,
317        program: &Self::Program<K>,
318        cfg: &Self::LaunchConfig,
319        args: K::Args<'a>,
320    ) -> teeny_core::errors::Result<()> {
321        // Allocate global scratch memory for TMA descriptors if the kernel requires it.
322        // Total scratch = per-CTA scratch * number of CTAs in the launch grid.
323        let num_ctas = (cfg.grid[0] * cfg.grid[1] * cfg.grid[2]) as u64;
324        let scratch_total = program.metadata.global_scratch_size * num_ctas;
325        let mut scratch_ptr: cuda::CUdeviceptr = 0;
326        if scratch_total > 0 {
327            // cuMemAlloc_v2 guarantees 256-byte alignment, which satisfies Triton's
328            // scratch alignment requirement (typically 128 bytes).
329            let alloc_status =
330                unsafe { cuda::cuMemAlloc_v2(&mut scratch_ptr, scratch_total as usize) };
331            if alloc_status != cuda::cudaError_enum_CUDA_SUCCESS {
332                return Err(Error::from_cuda_error(alloc_status).into());
333            }
334            // Zero-initialize the scratch pad so TMA descriptors start in a clean state.
335            unsafe { cuda::cuMemsetD8_v2(scratch_ptr, 0, scratch_total as usize) };
336        }
337
338        let mut packer = CudaArgPacker::new();
339        args.visit_args(&mut packer);
340        // Trailing Triton kernel parameters: global scratch pad + profile scratch pad.
341        packer.visit_ptr(scratch_ptr as *mut std::ffi::c_void); // global scratch pad
342        packer.visit_ptr(std::ptr::null_mut()); // profile scratch pad (unused)
343        // Build the pointer array while `packer` is still alive — both must
344        // remain live for the entire duration of `cuLaunchKernel`.
345        let mut ptrs = packer.as_ptrs();
346
347        Self::ensure_dynamic_shared(program.function, program.metadata.shared)?;
348
349        let status = unsafe {
350            cuda::cuLaunchKernel(
351                program.function,
352                cfg.grid[0],
353                cfg.grid[1],
354                cfg.grid[2],
355                cfg.block[0],
356                cfg.block[1],
357                cfg.block[2],
358                program.metadata.shared, // dynamic shared memory required by Triton kernel
359                std::ptr::null_mut(),    // hStream (default/null stream)
360                ptrs.as_mut_ptr(),
361                std::ptr::null_mut(), // extra
362            )
363        };
364
365        if status != cuda::cudaError_enum_CUDA_SUCCESS {
366            return Err(Error::from_cuda_error(status).into());
367        }
368
369        // `cuLaunchKernel` returns immediately; the kernel runs asynchronously.
370        // Synchronize here so that any GPU-side fault (bad pointer, out-of-bounds
371        // access) surfaces as a CUDA error code rather than a later SIGSEGV.
372        let sync_status = unsafe { cuda::cuCtxSynchronize() };
373
374        if sync_status != cuda::cudaError_enum_CUDA_SUCCESS {
375            if scratch_ptr != 0 {
376                unsafe { cuda::cuMemFree_v2(scratch_ptr) };
377            }
378            return Err(Error::from_cuda_error(sync_status).into());
379        }
380
381        if scratch_ptr != 0 {
382            unsafe { cuda::cuMemFree_v2(scratch_ptr) };
383        }
384
385        Ok(())
386    }
387}
388
389impl<'a> CudaDevice<'a> {
390    /// Launch a kernel on the given stream without allocating scratch or
391    /// synchronising. Used during CUDA graph capture: the caller must have
392    /// already appended the global-scratch-pad and profile-scratch-pad pointers
393    /// to `packer` before calling this.
394    pub(crate) fn launch_on_stream<K: Kernel>(
395        &self,
396        program: &CudaProgram<'_, K>,
397        cfg: &CudaLaunchConfig,
398        packer: &mut CudaArgPacker,
399        stream: cuda::CUstream,
400    ) -> Result<()> {
401        let mut ptrs = packer.as_ptrs();
402        Self::ensure_dynamic_shared(program.function, program.metadata.shared)?;
403        let status = unsafe {
404            cuda::cuLaunchKernel(
405                program.function,
406                cfg.grid[0],
407                cfg.grid[1],
408                cfg.grid[2],
409                cfg.block[0],
410                cfg.block[1],
411                cfg.block[2],
412                program.metadata.shared,
413                stream,
414                ptrs.as_mut_ptr(),
415                std::ptr::null_mut(),
416            )
417        };
418        if status != cuda::cudaError_enum_CUDA_SUCCESS {
419            Err(Error::from_cuda_error(status).into())
420        } else {
421            Ok(())
422        }
423    }
424
425    /// Launch a pre-loaded kernel using a dynamically-built arg list.
426    ///
427    /// Use this instead of `launch` when the kernel type is erased (e.g. in
428    /// `LoadedModel::forward`) and arguments are packed by `CudaArgPacker`
429    /// via a `RuntimeOp` rather than a static `Kernel::Args` tuple.
430    pub fn launch_with_packer<K: Kernel>(
431        &self,
432        program: &CudaProgram<'_, K>,
433        cfg: &CudaLaunchConfig,
434        packer: &mut CudaArgPacker,
435    ) -> Result<()> {
436        let num_ctas = (cfg.grid[0] * cfg.grid[1] * cfg.grid[2]) as u64;
437        let scratch_total = program.metadata.global_scratch_size * num_ctas;
438        let mut scratch_ptr: cuda::CUdeviceptr = 0;
439        if scratch_total > 0 {
440            let alloc_status =
441                unsafe { cuda::cuMemAlloc_v2(&mut scratch_ptr, scratch_total as usize) };
442            if alloc_status != cuda::cudaError_enum_CUDA_SUCCESS {
443                return Err(Error::from_cuda_error(alloc_status).into());
444            }
445            unsafe { cuda::cuMemsetD8_v2(scratch_ptr, 0, scratch_total as usize) };
446        }
447
448        packer.visit_ptr(scratch_ptr as *mut std::ffi::c_void);
449        packer.visit_ptr(std::ptr::null_mut());
450        let mut ptrs = packer.as_ptrs();
451
452        Self::ensure_dynamic_shared(program.function, program.metadata.shared)?;
453
454        let status = unsafe {
455            cuda::cuLaunchKernel(
456                program.function,
457                cfg.grid[0],
458                cfg.grid[1],
459                cfg.grid[2],
460                cfg.block[0],
461                cfg.block[1],
462                cfg.block[2],
463                program.metadata.shared,
464                std::ptr::null_mut(),
465                ptrs.as_mut_ptr(),
466                std::ptr::null_mut(),
467            )
468        };
469
470        if status != cuda::cudaError_enum_CUDA_SUCCESS {
471            if scratch_ptr != 0 {
472                unsafe { cuda::cuMemFree_v2(scratch_ptr) };
473            }
474            return Err(Error::from_cuda_error(status).into());
475        }
476
477        let sync_status = unsafe { cuda::cuCtxSynchronize() };
478        if scratch_ptr != 0 {
479            unsafe { cuda::cuMemFree_v2(scratch_ptr) };
480        }
481
482        if sync_status != cuda::cudaError_enum_CUDA_SUCCESS {
483            return Err(Error::from_cuda_error(sync_status).into());
484        }
485        Ok(())
486    }
487}
teeny_cuda/device/mod.rs

teeny_cuda/device/
mod.rs
