teeny_cuda/device/
program.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::ffi::CString;
18use std::marker::PhantomData;
19
20use teeny_core::device::program::{Kernel, Program};
21
22use crate::cuda;
23use crate::errors::{Error, Result};
24
25/// Kernel resource metadata parsed from `// meta:key=value` PTX comments
26/// appended by the Triton CUDA backend during compilation.
27#[derive(Debug, Default, Clone)]
28pub(crate) struct KernelMetadata {
29    pub(crate) name: String,
30    pub(crate) num_warps: u32,
31    pub(crate) num_ctas: u32,
32    pub(crate) shared: u32,
33    pub(crate) tmem_size: u32,
34    pub(crate) global_scratch_size: u64,
35    pub(crate) global_scratch_align: u64,
36    pub(crate) profile_scratch_size: u32,
37    pub(crate) profile_scratch_align: u32,
38}
39
40impl KernelMetadata {
41    fn parse(ptx: &str) -> Self {
42        let mut m = KernelMetadata {
43            num_ctas: 1,
44            global_scratch_align: 1,
45            profile_scratch_align: 1,
46            ..Default::default()
47        };
48        let mut reqntid: Option<u32> = None;
49        let mut visible_entry_name = String::new();
50
51        for line in ptx.lines() {
52            let trimmed = line.trim();
53
54            // Parse Triton metadata block: `// meta:key=value`
55            if let Some(rest) = trimmed.strip_prefix("// meta:") {
56                if let Some((key, val)) = rest.split_once('=') {
57                    match key {
58                        "name" => m.name = val.to_owned(),
59                        "num_warps" => m.num_warps = val.parse().unwrap_or(0),
60                        "num_ctas" => m.num_ctas = val.parse().unwrap_or(1),
61                        "shared" => m.shared = val.parse().unwrap_or(0),
62                        "tmem_size" => m.tmem_size = val.parse().unwrap_or(0),
63                        "global_scratch_size" => m.global_scratch_size = val.parse().unwrap_or(0),
64                        "global_scratch_align" => m.global_scratch_align = val.parse().unwrap_or(1),
65                        "profile_scratch_size" => m.profile_scratch_size = val.parse().unwrap_or(0),
66                        "profile_scratch_align" => {
67                            m.profile_scratch_align = val.parse().unwrap_or(1)
68                        }
69                        _ => {}
70                    }
71                }
72                continue;
73            }
74
75            // Parse `.reqntid X` — PTX thread-count directive emitted by NVPTX backend.
76            // Used as fallback when Triton metadata is absent (e.g. NVPTX-compiled kernels).
77            if let Some(rest) = trimmed.strip_prefix(".reqntid ") {
78                let x_str = rest.split(',').next().unwrap_or("").trim();
79                if let Ok(x) = x_str.parse::<u32>() {
80                    reqntid = Some(x);
81                }
82            }
83
84            // Parse `.visible .entry name(` — the PTX kernel symbol name emitted by NVPTX.
85            // Used as a fallback when no `// meta:name=` comment is present.
86            if visible_entry_name.is_empty()
87                && let Some(rest) = trimmed.strip_prefix(".visible .entry ")
88            {
89                let name_end = rest
90                    .find('(')
91                    .unwrap_or(rest.find(' ').unwrap_or(rest.len()));
92                let parsed = rest[..name_end].trim();
93                if !parsed.is_empty() {
94                    visible_entry_name = parsed.to_owned();
95                }
96            }
97
98            // Parse legacy Triton comment format emitted by older PTX cached before the
99            // `// meta:` block was introduced. These act as low-priority fallbacks: a
100            // later `// meta:` line for the same key will have already set the value, so
101            // we only write when the field still holds its zero/unit default.
102            if let Some(v) = trimmed.strip_prefix("// TRITON_SHARED_MEM_BYTES: ") {
103                if m.shared == 0 {
104                    m.shared = v.trim().parse().unwrap_or(0);
105                }
106            } else if let Some(v) = trimmed.strip_prefix("// TRITON_GLOBAL_SCRATCH_BYTES_PER_CTA: ")
107            {
108                if m.global_scratch_size == 0 {
109                    m.global_scratch_size = v.trim().parse().unwrap_or(0);
110                }
111            } else if let Some(v) = trimmed.strip_prefix("// TRITON_GLOBAL_SCRATCH_ALIGN: ")
112                && m.global_scratch_align == 1
113            {
114                m.global_scratch_align = v.trim().parse::<u64>().unwrap_or(1).max(1);
115            }
116        }
117
118        // Fill name: prefer `// meta:name=`, then `.visible .entry`, then fallback.
119        if m.name.is_empty() {
120            m.name = if !visible_entry_name.is_empty() {
121                visible_entry_name
122            } else {
123                "entry_point".to_owned()
124            };
125        }
126        if m.num_warps == 0 {
127            // Derive num_warps from .reqntid (round up to full warps).
128            m.num_warps = reqntid.unwrap_or(128).div_ceil(32);
129        }
130
131        m
132    }
133
134    /// Threads per block, derived from num_warps (CUDA warp size is always 32).
135    pub(crate) fn threads_per_block(&self) -> u32 {
136        self.num_warps * 32
137    }
138}
139
140/// Strip DWARF debug sections from PTX source.
141///
142/// The Rust NVPTX backend emits PTX object files that include DWARF debug
143/// sections with relocation references (e.g. `.b32 .debug_abbrev`). These
144/// relocations are resolved by a PTX linker but are invalid when passing PTX
145/// directly to `cuModuleLoadData` for driver JIT compilation.
146///
147/// This function truncates the PTX at the first `.file` or `.section .debug`
148/// directive, which always appears after the kernel body.
149fn strip_debug_sections(ptx: &[u8]) -> &[u8] {
150    let mut pos = 0;
151    while pos < ptx.len() {
152        let line_end = ptx[pos..]
153            .iter()
154            .position(|&b| b == b'\n')
155            .map(|i| pos + i)
156            .unwrap_or(ptx.len());
157
158        let line = &ptx[pos..line_end];
159        let trim_start = line
160            .iter()
161            .position(|&b| b != b' ' && b != b'\t')
162            .unwrap_or(line.len());
163        let trimmed = &line[trim_start..];
164
165        // .file directives and .section .debug_* mark the start of DWARF content
166        if trimmed.starts_with(b".file")
167            || (trimmed.starts_with(b".section") && line.windows(7).any(|w| w == b".debug_"))
168        {
169            return &ptx[..pos];
170        }
171
172        pos = line_end + 1;
173    }
174    ptx
175}
176
177/// A loaded CUDA program: the cubin is loaded into a `CUmodule` and the
178/// entry-point function is resolved to a `CUfunction` ready to launch.
179pub struct CudaProgram<'a, K: Kernel> {
180    pub(crate) module: cuda::CUmodule,
181    pub(crate) function: cuda::CUfunction,
182    /// Kernel resource metadata parsed from `// meta:key=value` PTX comments.
183    pub(crate) metadata: KernelMetadata,
184    _unused: PhantomData<&'a ()>,
185    _kernel: PhantomData<K>,
186}
187
188impl<'a, K: Kernel> CudaProgram<'a, K> {
189    /// The loaded module's raw `CUmodule` handle, as a `usize`.
190    pub fn module_ptr(&self) -> usize {
191        self.module as usize
192    }
193    /// The resolved kernel entry point's raw `CUfunction` handle, as a `usize`.
194    pub fn function_ptr(&self) -> usize {
195        self.function as usize
196    }
197    /// Threads per block, from the kernel's parsed metadata.
198    pub fn threads_per_block(&self) -> u32 {
199        self.metadata.threads_per_block()
200    }
201    /// Number of CTAs (thread blocks), from the kernel's parsed metadata.
202    pub fn num_ctas(&self) -> u32 {
203        self.metadata.num_ctas
204    }
205
206    /// Load a cubin image into the current CUDA context and resolve `entry_point`.
207    ///
208    /// Metadata is not available from a cubin; default values are used.
209    pub fn try_new(cubin: &[u8], entry_point: &str) -> Result<Self> {
210        let mut module = cuda::CUmodule::default();
211        let status = unsafe { cuda::cuModuleLoadData(&mut module, cubin.as_ptr().cast()) };
212        if status != cuda::cudaError_enum_CUDA_SUCCESS {
213            return Err(Error::from_cuda_error(status).into());
214        }
215
216        Self::resolve_function(module, entry_point, KernelMetadata::default())
217    }
218
219    /// JIT-compile PTX source via the CUDA driver.
220    ///
221    /// The entry-point function name and all resource metadata are read from
222    /// the `// meta:key=value` block that the Triton CUDA backend appends to
223    /// every PTX output. `ptx` must be ASCII PTX text; a null terminator is
224    /// appended automatically.
225    pub fn try_from_ptx(ptx: &[u8]) -> Result<Self> {
226        let ptx_str = std::str::from_utf8(ptx).unwrap_or("");
227        let metadata = KernelMetadata::parse(ptx_str);
228
229        // Strip DWARF debug sections — they contain relocations (.b32 .debug_abbrev)
230        // that are only valid in linked PTX object files, not for driver JIT.
231        let ptx = strip_debug_sections(ptx);
232
233        // cuModuleLoadData expects a null-terminated string for PTX input.
234        let mut ptx_ntstr = ptx.to_vec();
235        ptx_ntstr.push(0);
236
237        // Use cuModuleLoadDataEx with error logging only. Do NOT override CU_JIT_TARGET:
238        // our PTX targets sm_90a (Hopper with TMA), and Blackwell runs sm_90a code in
239        // forward-compatible mode. Forcing CU_JIT_TARGET=120 causes ptxas to attempt a
240        // cross-architecture compile that it refuses ("cannot be compiled to future
241        // architectures"). Without a target override, the CUDA driver uses the PTX's
242        // declared target and JITs it for the current device automatically.
243        const LOG_SIZE: usize = 65536;
244        let mut error_log = vec![0u8; LOG_SIZE];
245        let mut info_log = vec![0u8; LOG_SIZE];
246
247        // CUDA cuModuleLoadDataEx option values:
248        // - Buffer pointers: pass as pointer (void*)
249        // - Size/enum values: pass as the VALUE itself cast to void* (not a pointer to value)
250        let mut options = [
251            cuda::CUjit_option_enum_CU_JIT_ERROR_LOG_BUFFER,
252            cuda::CUjit_option_enum_CU_JIT_ERROR_LOG_BUFFER_SIZE_BYTES,
253            cuda::CUjit_option_enum_CU_JIT_INFO_LOG_BUFFER,
254            cuda::CUjit_option_enum_CU_JIT_INFO_LOG_BUFFER_SIZE_BYTES,
255        ];
256        #[allow(clippy::cast_ptr_alignment)]
257        let mut option_values: [*mut std::ffi::c_void; 4] = [
258            error_log.as_mut_ptr().cast(),
259            LOG_SIZE as *mut std::ffi::c_void, // size value, not a pointer
260            info_log.as_mut_ptr().cast(),
261            LOG_SIZE as *mut std::ffi::c_void, // size value, not a pointer
262        ];
263
264        let mut module = cuda::CUmodule::default();
265        let status = unsafe {
266            cuda::cuModuleLoadDataEx(
267                &mut module,
268                ptx_ntstr.as_ptr().cast(),
269                4,
270                options.as_mut_ptr(),
271                option_values.as_mut_ptr(),
272            )
273        };
274        if status != cuda::cudaError_enum_CUDA_SUCCESS {
275            let err_len = error_log.iter().position(|&b| b == 0).unwrap_or(LOG_SIZE);
276            let error_str = std::str::from_utf8(&error_log[..err_len]).unwrap_or("<invalid utf8>");
277            eprintln!("[CUDA-JIT] error log: {}", error_str);
278            return Err(Error::from_cuda_error(status).into());
279        }
280
281        let info_len = info_log.iter().position(|&b| b == 0).unwrap_or(0);
282        if info_len > 0 {
283            let info_str = std::str::from_utf8(&info_log[..info_len]).unwrap_or("<invalid utf8>");
284            eprintln!("[CUDA-JIT] info: {}", info_str);
285        }
286
287        Self::resolve_function(module, &metadata.name.clone(), metadata)
288    }
289
290    fn resolve_function(
291        module: cuda::CUmodule,
292        entry_point: &str,
293        metadata: KernelMetadata,
294    ) -> Result<Self> {
295        let name = CString::new(entry_point).map_err(Error::CStringError)?;
296        let mut function = cuda::CUfunction::default();
297        let status = unsafe { cuda::cuModuleGetFunction(&mut function, module, name.as_ptr()) };
298        if status != cuda::cudaError_enum_CUDA_SUCCESS {
299            unsafe { cuda::cuModuleUnload(module) };
300            return Err(Error::from_cuda_error(status).into());
301        }
302
303        Ok(Self {
304            module,
305            function,
306            metadata,
307            _unused: PhantomData,
308            _kernel: PhantomData,
309        })
310    }
311}
312
313impl<'a, K: Kernel> Drop for CudaProgram<'a, K> {
314    fn drop(&mut self) {
315        let status = unsafe { cuda::cuModuleUnload(self.module) };
316        if status != cuda::cudaError_enum_CUDA_SUCCESS {
317            eprintln!("Failed to unload CUDA module: {}", status);
318        }
319    }
320}
321
322impl<'a, K: Kernel> Program<'a, K> for CudaProgram<'a, K> {}
323
324/// A dummy `Kernel` marker used to load PTX without a concrete kernel type.
325/// Enables `CudaProgram<ErasedKernel>` for type-erased model execution.
326pub struct ErasedKernel;
327
328impl Kernel for ErasedKernel {
329    type Args<'a> = ();
330    fn name(&self) -> &str {
331        ""
332    }
333    fn source(&self) -> &str {
334        ""
335    }
336    fn kernel_source(&self) -> &str {
337        ""
338    }
339    fn entry_point_source(&self) -> &str {
340        ""
341    }
342}
teeny_cuda/device/program.rs

teeny_cuda/device/
program.rs
