teeny_cuda/compiler/
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::ffi::{c_char, c_void};
18use std::ptr;
19
20use teeny_core::device::program::Kernel;
21
22use crate::compiler::options::Options;
23use crate::cuda;
24use crate::device::program::CudaProgram;
25use crate::errors::{Error, Result};
26
27/// Ahead-of-time kernel compilation.
28pub mod aot;
29/// Compiling a `teeny-core` graph's kernels.
30pub mod graph;
31/// `nvptxcompiler` compile options.
32pub mod options;
33/// CUDA compilation target descriptions.
34pub mod target;
35
36/// Wraps `nvptxcompiler` (NVIDIA's standalone PTX-to-cubin compiler).
37pub struct PtxCompiler {
38    compiler: cuda::nvPTXCompilerHandle,
39}
40
41impl PtxCompiler {
42    /// Creates a compiler for the given PTX source.
43    ///
44    /// `ptx` is the raw PTX source bytes (ASCII; the C API takes a byte pointer
45    /// and length, so no UTF-8 validation or null-termination is required).
46    pub fn try_new(ptx: &[u8]) -> Result<Self> {
47        let mut compiler = cuda::nvPTXCompilerHandle::default();
48        let result = unsafe {
49            cuda::nvPTXCompilerCreate(&mut compiler, ptx.len(), ptx.as_ptr().cast::<c_char>())
50        };
51
52        if result != cuda::nvPTXCompileResult_NVPTXCOMPILE_SUCCESS {
53            return Err(Error::NvptxCompileError {
54                code: result,
55                log: String::new(),
56            }
57            .into());
58        }
59        Ok(PtxCompiler { compiler })
60    }
61
62    /// Compiles the PTX to a cubin binary using `options`.
63    pub fn compile(&mut self, options: &Options) -> Result<Vec<u8>> {
64        let compile_options = options.to_compile_options();
65        let num_options = compile_options.len() as i32;
66        eprintln!("[nvPTX] compile options: {compile_options:?}");
67
68        // Convert Vec<String> to Vec<*const c_char>
69        let cstrs: Vec<std::ffi::CString> = compile_options
70            .iter()
71            .map(|s| std::ffi::CString::new(s.as_str()).map_err(|e| Error::CStringError(e).into()))
72            .collect::<Result<Vec<std::ffi::CString>>>()?;
73        let cptrs: Vec<*const c_char> = cstrs.iter().map(|cs| cs.as_ptr()).collect();
74
75        let result =
76            unsafe { cuda::nvPTXCompilerCompile(self.compiler, num_options, cptrs.as_ptr()) };
77
78        if result != cuda::nvPTXCompileResult_NVPTXCOMPILE_SUCCESS {
79            // Retrieve the compiler error log for a human-readable message.
80            let log = self.error_log();
81            return Err(Error::NvptxCompileError { code: result, log }.into());
82        }
83
84        let mut binary_size = 0usize;
85        let result =
86            unsafe { cuda::nvPTXCompilerGetCompiledProgramSize(self.compiler, &mut binary_size) };
87        if result != cuda::nvPTXCompileResult_NVPTXCOMPILE_SUCCESS {
88            return Err(Error::NvptxCompileError {
89                code: result,
90                log: String::new(),
91            }
92            .into());
93        }
94
95        let mut binary = vec![0u8; binary_size];
96        let result = unsafe {
97            cuda::nvPTXCompilerGetCompiledProgram(
98                self.compiler,
99                binary.as_mut_ptr().cast::<c_void>(),
100            )
101        };
102        if result != cuda::nvPTXCompileResult_NVPTXCOMPILE_SUCCESS {
103            return Err(Error::NvptxCompileError {
104                code: result,
105                log: String::new(),
106            }
107            .into());
108        }
109
110        Ok(binary)
111    }
112
113    fn error_log(&self) -> String {
114        let mut log_size = 0usize;
115        let result = unsafe { cuda::nvPTXCompilerGetErrorLogSize(self.compiler, &mut log_size) };
116        if result != cuda::nvPTXCompileResult_NVPTXCOMPILE_SUCCESS || log_size == 0 {
117            return String::new();
118        }
119        let mut buf = vec![0u8; log_size];
120        let result = unsafe {
121            cuda::nvPTXCompilerGetErrorLog(self.compiler, buf.as_mut_ptr().cast::<c_char>())
122        };
123        if result != cuda::nvPTXCompileResult_NVPTXCOMPILE_SUCCESS {
124            return String::new();
125        }
126        String::from_utf8_lossy(&buf)
127            .trim_end_matches('\0')
128            .to_string()
129    }
130
131    /// Compile the PTX to a cubin and load it into the current CUDA context,
132    /// returning a ready-to-launch `CudaProgram`. The entry-point symbol is
133    /// always `"entry_point"` (the name emitted by the `#[kernel]` proc macro).
134    pub fn compile_program<K: Kernel>(
135        &mut self,
136        options: &Options,
137    ) -> Result<CudaProgram<'static, K>> {
138        let cubin = self.compile(options)?;
139        CudaProgram::try_new(&cubin, options.entry.as_str())
140    }
141}
142
143impl Drop for PtxCompiler {
144    fn drop(&mut self) {
145        let result = unsafe { cuda::nvPTXCompilerDestroy(ptr::addr_of_mut!(self.compiler)) };
146        if result != cuda::nvPTXCompileResult_NVPTXCOMPILE_SUCCESS {
147            eprintln!("Failed to destroy NVPTX compiler: {}", result);
148        }
149    }
150}
teeny_cuda/compiler/mod.rs

teeny_cuda/compiler/
mod.rs
