Skip to main content

oxicuda_driver/
module.rs

1//! PTX module loading and kernel function management.
2//!
3//! Modules are created from PTX source code and contain one or more
4//! kernel functions that can be launched on the GPU.
5//!
6//! # Example
7//!
8//! ```rust,no_run
9//! # use oxicuda_driver::module::{Module, JitOptions};
10//! # fn main() -> Result<(), oxicuda_driver::error::CudaError> {
11//! let ptx = r#"
12//! .version 7.0
13//! .target sm_70
14//! .address_size 64
15//! .visible .entry my_kernel() { ret; }
16//! "#;
17//!
18//! let module = Module::from_ptx(ptx)?;
19//! let func = module.get_function("my_kernel")?;
20//!
21//! // Or with JIT options and compilation logs:
22//! let opts = JitOptions { optimization_level: 4, ..Default::default() };
23//! let (module2, log) = Module::from_ptx_with_options(ptx, &opts)?;
24//! if !log.info.is_empty() {
25//!     println!("JIT info: {}", log.info);
26//! }
27//! # Ok(())
28//! # }
29//! ```
30
31use std::ffi::{CString, c_void};
32
33use crate::error::{CudaError, CudaResult};
34use crate::ffi::{CUfunction, CUjit_option, CUmodule};
35use crate::loader::try_driver;
36
37// ---------------------------------------------------------------------------
38// JitOptions
39// ---------------------------------------------------------------------------
40
41/// Options for JIT compilation of PTX to GPU binary.
42///
43/// These options control the behaviour of the CUDA JIT compiler when
44/// loading PTX source via [`Module::from_ptx_with_options`].
45#[derive(Debug, Clone)]
46pub struct JitOptions {
47    /// Maximum number of registers per thread (0 = no limit).
48    ///
49    /// Limiting register usage can increase occupancy at the cost of
50    /// potential register spilling to local memory.
51    pub max_registers: u32,
52    /// Optimisation level (0--4, default 4).
53    ///
54    /// Higher levels produce faster code but take longer to compile.
55    pub optimization_level: u32,
56    /// Whether to generate debug information in the compiled binary.
57    pub generate_debug_info: bool,
58    /// If `true`, the JIT compiler determines the target compute
59    /// capability from the current CUDA context.
60    pub target_from_context: bool,
61}
62
63impl Default for JitOptions {
64    /// Returns sensible defaults: no register limit, maximum
65    /// optimisation, no debug info, target derived from context.
66    fn default() -> Self {
67        Self {
68            max_registers: 0,
69            optimization_level: 4,
70            generate_debug_info: false,
71            target_from_context: true,
72        }
73    }
74}
75
76// ---------------------------------------------------------------------------
77// JitLog
78// ---------------------------------------------------------------------------
79
80/// Severity of a JIT compiler diagnostic message.
81#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
82pub enum JitSeverity {
83    /// A fatal error that prevents PTX compilation.
84    Fatal,
85    /// A non-fatal error.
86    Error,
87    /// A compiler warning (compilation may still succeed).
88    Warning,
89    /// An informational message (e.g. register usage).
90    Info,
91}
92
93impl std::fmt::Display for JitSeverity {
94    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
95        match self {
96            Self::Fatal => f.write_str("fatal"),
97            Self::Error => f.write_str("error"),
98            Self::Warning => f.write_str("warning"),
99            Self::Info => f.write_str("info"),
100        }
101    }
102}
103
104/// A single structured diagnostic emitted by the JIT compiler.
105///
106/// Parsed from the raw `ptxas` log lines that look like:
107///
108/// ```text
109/// ptxas error   : 'kernel', line 10; error   : Unknown instruction 'xyz'
110/// ptxas warning : 'kernel', line 15; warning : double-precision is slow
111/// ptxas info    : 'kernel' used 16 registers, 0 bytes smem
112/// ```
113#[derive(Debug, Clone, PartialEq, Eq, Hash)]
114pub struct JitDiagnostic {
115    /// Severity level.
116    pub severity: JitSeverity,
117    /// Kernel function name, if the message is function-scoped.
118    pub kernel: Option<String>,
119    /// Source line number, if present.
120    pub line: Option<u32>,
121    /// Human-readable message text.
122    pub message: String,
123}
124
125/// Log output from JIT compilation.
126///
127/// After calling [`Module::from_ptx_with_options`], this struct
128/// contains any informational or error messages emitted by the
129/// JIT compiler.
130///
131/// Use [`JitLog::parse_diagnostics`] to obtain structured
132/// [`JitDiagnostic`] entries instead of parsing the raw strings.
133#[derive(Debug, Clone, Default, PartialEq, Eq, Hash)]
134pub struct JitLog {
135    /// Informational messages from the JIT compiler.
136    pub info: String,
137    /// Error messages from the JIT compiler.
138    pub error: String,
139}
140
141impl JitLog {
142    /// Returns `true` if there are no messages in either buffer.
143    #[must_use]
144    pub fn is_empty(&self) -> bool {
145        self.info.is_empty() && self.error.is_empty()
146    }
147
148    /// Returns `true` if the error buffer is non-empty.
149    #[must_use]
150    pub fn has_errors(&self) -> bool {
151        !self.error.is_empty()
152    }
153
154    /// Parse both log buffers into a `Vec` of structured [`JitDiagnostic`]
155    /// entries.
156    ///
157    /// Lines that do not match the `ptxas` diagnostic format are included as
158    /// [`JitSeverity::Info`] diagnostics with no kernel or line information,
159    /// unless they are entirely blank.
160    ///
161    /// # Message format
162    ///
163    /// The CUDA JIT compiler emits lines in one of these formats:
164    ///
165    /// ```text
166    /// ptxas {severity}   : '{kernel}', line {n}; {type}   : {message}
167    /// ptxas {severity}   : '{kernel}' {message}
168    /// ptxas {severity}   : {message}
169    /// ```
170    ///
171    /// This method normalises all of those into [`JitDiagnostic`] values.
172    #[must_use]
173    pub fn parse_diagnostics(&self) -> Vec<JitDiagnostic> {
174        let mut out = Vec::new();
175        for line in self.error.lines().chain(self.info.lines()) {
176            if let Some(d) = parse_ptxas_line(line) {
177                out.push(d);
178            }
179        }
180        out
181    }
182
183    /// Return only the [`JitDiagnostic`] entries whose severity is
184    /// [`JitSeverity::Error`] or [`JitSeverity::Fatal`].
185    #[must_use]
186    pub fn errors(&self) -> Vec<JitDiagnostic> {
187        self.parse_diagnostics()
188            .into_iter()
189            .filter(|d| matches!(d.severity, JitSeverity::Error | JitSeverity::Fatal))
190            .collect()
191    }
192
193    /// Return only the [`JitDiagnostic`] entries whose severity is
194    /// [`JitSeverity::Warning`].
195    #[must_use]
196    pub fn warnings(&self) -> Vec<JitDiagnostic> {
197        self.parse_diagnostics()
198            .into_iter()
199            .filter(|d| matches!(d.severity, JitSeverity::Warning))
200            .collect()
201    }
202}
203
204// ── ptxas log line parser ─────────────────────────────────────────────────────
205
206/// Parse a single `ptxas` log line into a [`JitDiagnostic`], returning
207/// `None` for blank lines.
208///
209/// Handles these representative patterns:
210///
211/// ```text
212/// ptxas error   : 'vec_add', line 10; error   : Unknown instruction 'xyz'
213/// ptxas warning : 'vec_add', line 15; warning : slow double-precision
214/// ptxas info    : 'vec_add' used 16 registers, 0 bytes smem, 0 bytes cmem[0]
215/// ptxas fatal   : Unresolved extern function 'foo'
216/// ```
217fn parse_ptxas_line(line: &str) -> Option<JitDiagnostic> {
218    let line = line.trim();
219    if line.is_empty() {
220        return None;
221    }
222
223    // Must start with "ptxas " (case-insensitive is not needed — ptxas always
224    // uses lower-case).
225    let rest = line.strip_prefix("ptxas ")?;
226
227    // Extract severity word (first whitespace-delimited token).
228    let (sev_str, after_sev) = split_first_word(rest.trim_start());
229    let severity = match sev_str.to_ascii_lowercase().trim_end_matches(':') {
230        "fatal" => JitSeverity::Fatal,
231        "error" => JitSeverity::Error,
232        "warning" => JitSeverity::Warning,
233        "info" => JitSeverity::Info,
234        _ => JitSeverity::Info,
235    };
236
237    // Skip past `: ` after the severity keyword.
238    let body = skip_colon(after_sev.trim_start());
239
240    // Try to extract kernel name from `'kernel_name'` at the start.
241    let (kernel, after_kernel) = extract_kernel_name(body);
242
243    // Try to extract line number: `, line N;` or `, line N,`.
244    let (line_no, after_line) = extract_line_number(after_kernel);
245
246    // The remaining text — skip a leading type word if present (e.g. `error   : `).
247    let message = extract_message(after_line.trim());
248
249    Some(JitDiagnostic {
250        severity,
251        kernel,
252        line: line_no,
253        message: message.to_string(),
254    })
255}
256
257/// Split a `&str` at the first whitespace boundary; returns `("", s)` if
258/// there is no whitespace.
259fn split_first_word(s: &str) -> (&str, &str) {
260    match s.find(|c: char| c.is_whitespace()) {
261        Some(pos) => (&s[..pos], &s[pos..]),
262        None => (s, ""),
263    }
264}
265
266/// Skip past the first `: ` (colon + optional spaces) in `s`.
267fn skip_colon(s: &str) -> &str {
268    if let Some(pos) = s.find(':') {
269        s[pos + 1..].trim_start()
270    } else {
271        s
272    }
273}
274
275/// Attempt to extract `'kernel_name'` from the beginning of `s`.
276/// Returns `(Some(name), rest_after_name)` or `(None, s)`.
277fn extract_kernel_name(s: &str) -> (Option<String>, &str) {
278    let s = s.trim_start();
279    if !s.starts_with('\'') {
280        return (None, s);
281    }
282    let inner = &s[1..];
283    if let Some(end) = inner.find('\'') {
284        let name = inner[..end].to_string();
285        let after = &inner[end + 1..];
286        (Some(name), after)
287    } else {
288        (None, s)
289    }
290}
291
292/// Attempt to extract `, line N;` or `, line N,` from the start of `s`.
293/// Returns `(Some(n), rest)` or `(None, s)`.
294fn extract_line_number(s: &str) -> (Option<u32>, &str) {
295    // Accept `, line N` (with optional trailing `;` or `,`)
296    let s_trim = s.trim_start_matches([',', ' ', ';']);
297    let lower = s_trim.to_ascii_lowercase();
298    if !lower.starts_with("line ") {
299        return (None, s);
300    }
301    let after_line = &s_trim[5..]; // skip "line "
302    let (num_str, rest) = split_first_word(after_line.trim_start());
303    let num_clean: String = num_str.chars().filter(|c| c.is_ascii_digit()).collect();
304    if let Ok(n) = num_clean.parse::<u32>() {
305        (Some(n), rest)
306    } else {
307        (None, s)
308    }
309}
310
311/// Strip a leading `type   : ` prefix (e.g. `error   : ` or `warning : `)
312/// from a message if present; return the remaining text.
313fn extract_message(s: &str) -> &str {
314    // Pattern: word followed by optional spaces and `:`.
315    let (word, rest) = split_first_word(s);
316    let word_clean = word.trim_end_matches(':');
317    if matches!(
318        word_clean.to_ascii_lowercase().as_str(),
319        "error" | "warning" | "info" | "fatal"
320    ) {
321        skip_colon(rest.trim_start())
322    } else {
323        s
324    }
325}
326
327// ---------------------------------------------------------------------------
328// jit_failure — build a JitFailed error from raw log buffers
329// ---------------------------------------------------------------------------
330
331/// Build a [`CudaError::JitFailed`] by combining the raw JIT log buffers
332/// with the underlying CUDA error.
333///
334/// Both `info_buf` and `error_buf` are interpreted as UTF-8 (with lossy
335/// conversion), parsed for structured diagnostics, and wrapped together
336/// with `source` into a [`CudaError::JitFailed`] variant.
337///
338/// This is `pub(crate)` so that both `module.rs` and `link.rs` can call
339/// it without exposing it as part of the public API.
340///
341/// Only compiled on non-macOS platforms because `link.rs`'s GPU path
342/// is the sole caller.
343#[cfg(not(target_os = "macos"))]
344pub(crate) fn jit_failure(source: CudaError, info_buf: &[u8], error_buf: &[u8]) -> CudaError {
345    let info = String::from_utf8_lossy(info_buf).into_owned();
346    let error = String::from_utf8_lossy(error_buf).into_owned();
347
348    let log = JitLog { info, error };
349    let diagnostic_count = log.parse_diagnostics().len();
350
351    CudaError::JitFailed {
352        log: Box::new(log),
353        diagnostic_count,
354        source: Box::new(source),
355    }
356}
357
358/// Variant of [`jit_failure`] that accepts a pre-built [`JitLog`].
359///
360/// Used when the log has already been assembled (e.g. in
361/// [`Module::from_ptx_with_options`] where the log was extracted
362/// before checking for a compilation error).
363pub(crate) fn jit_failure_from_log(source: CudaError, log: JitLog) -> CudaError {
364    let diagnostic_count = log.parse_diagnostics().len();
365    CudaError::JitFailed {
366        log: Box::new(log),
367        diagnostic_count,
368        source: Box::new(source),
369    }
370}
371
372// ---------------------------------------------------------------------------
373// Module
374// ---------------------------------------------------------------------------
375
376/// A loaded CUDA module containing one or more kernel functions.
377///
378/// Modules are typically created from PTX source via [`Module::from_ptx`]
379/// or [`Module::from_ptx_with_options`]. Individual kernel functions
380/// are retrieved by name with [`Module::get_function`].
381///
382/// The module is unloaded when this struct is dropped.
383pub struct Module {
384    /// Raw CUDA module handle.
385    raw: CUmodule,
386    /// The context that owned this module at load time, used to skip the driver
387    /// unload if that context was torn down first (avoids a use-after-free).
388    /// `None` when no tracked context was current — see
389    /// [`crate::context::current_ctx_owner`].
390    owner: crate::context::CtxOwner,
391}
392
393// `Module` is `Send + Sync` by auto-derivation: its only field is a `CUmodule`
394// handle (a plain driver-side identifier). The CUDA Driver API is thread-safe,
395// so no manual `unsafe impl` is required.
396
397/// Size of the JIT log buffers in bytes.
398const JIT_LOG_BUFFER_SIZE: usize = 4096;
399
400/// Replaces every non-ASCII byte in `ptx` with `?`, returning `None` when the
401/// input is already pure ASCII (the overwhelmingly common case, which then costs
402/// no allocation).
403///
404/// `ptxas` and the driver's JIT reject **any** non-ASCII byte in a PTX module —
405/// including inside `//` comments, which carry no semantics at all. CUDA 11.x
406/// accepted such modules silently, so a generator that emitted a typographic
407/// character in a comment (`x`, `->`, box-drawing rules) worked for years and
408/// then began failing with `CUDA_ERROR_INVALID_PTX` on CUDA 12.9+ toolchains,
409/// with no indication of which character or which kernel was at fault.
410///
411/// Individual generators in this workspace have been corrected to emit ASCII,
412/// but every PTX string in every dependent crate funnels through `Module`, so
413/// scrubbing here means a stray non-ASCII comment can never again turn into an
414/// opaque load failure. Each non-ASCII character is replaced with one `?` per
415/// UTF-8 byte it occupies, so total byte length — and therefore all offsets —
416/// stays stable and JIT diagnostics still point at the right line and column.
417fn ascii_only(ptx: &str) -> Option<String> {
418    if ptx.is_ascii() {
419        return None;
420    }
421    let mut scrubbed = String::with_capacity(ptx.len());
422    for c in ptx.chars() {
423        if c.is_ascii() {
424            scrubbed.push(c);
425        } else {
426            for _ in 0..c.len_utf8() {
427                scrubbed.push('?');
428            }
429        }
430    }
431    Some(scrubbed)
432}
433
434impl Module {
435    /// Loads a module from PTX source with default JIT options.
436    ///
437    /// The PTX string is automatically null-terminated before being
438    /// passed to the driver.
439    ///
440    /// # Errors
441    ///
442    /// Returns [`CudaError::InvalidImage`] if
443    /// the PTX is malformed, or another [`CudaError`] if the driver
444    /// call fails (e.g. no current context).
445    pub fn from_ptx(ptx: &str) -> CudaResult<Self> {
446        let api = try_driver()?;
447        // See `ascii_only`: the JIT rejects non-ASCII bytes anywhere in the
448        // module, comments included.
449        let scrubbed = ascii_only(ptx);
450        let ptx = scrubbed.as_deref().unwrap_or(ptx);
451        let c_ptx = CString::new(ptx).map_err(|_| CudaError::InvalidValue)?;
452        let mut raw = CUmodule::default();
453        crate::cuda_call!((api.cu_module_load_data)(
454            &mut raw,
455            c_ptx.as_ptr().cast::<c_void>()
456        ))?;
457        Ok(Self {
458            raw,
459            owner: crate::context::current_ctx_owner(),
460        })
461    }
462
463    /// Loads a module from PTX source with explicit JIT compiler options.
464    ///
465    /// Returns the loaded module together with a [`JitLog`] containing
466    /// any informational or error messages from the JIT compiler.
467    ///
468    /// # Errors
469    ///
470    /// Returns a [`CudaError`] if JIT compilation fails or the driver
471    /// call otherwise errors.
472    pub fn from_ptx_with_options(ptx: &str, options: &JitOptions) -> CudaResult<(Self, JitLog)> {
473        let api = try_driver()?;
474        // See `ascii_only`.
475        let scrubbed = ascii_only(ptx);
476        let ptx = scrubbed.as_deref().unwrap_or(ptx);
477        let c_ptx = CString::new(ptx).map_err(|_| CudaError::InvalidValue)?;
478
479        // Allocate log buffers on the heap.
480        let mut info_buf: Vec<u8> = vec![0u8; JIT_LOG_BUFFER_SIZE];
481        let mut error_buf: Vec<u8> = vec![0u8; JIT_LOG_BUFFER_SIZE];
482
483        // Build the parallel option-key and option-value arrays.
484        //
485        // Each option is a (CUjit_option, *mut c_void) pair. The value
486        // pointer is reinterpreted according to the option key — scalar
487        // values are cast directly to pointer-width integers.
488        let mut opt_keys: Vec<CUjit_option> = Vec::with_capacity(8);
489        let mut opt_vals: Vec<*mut c_void> = Vec::with_capacity(8);
490
491        // Info log buffer.
492        opt_keys.push(CUjit_option::InfoLogBuffer);
493        opt_vals.push(info_buf.as_mut_ptr().cast::<c_void>());
494
495        opt_keys.push(CUjit_option::InfoLogBufferSizeBytes);
496        opt_vals.push(JIT_LOG_BUFFER_SIZE as *mut c_void);
497
498        // Error log buffer.
499        opt_keys.push(CUjit_option::ErrorLogBuffer);
500        opt_vals.push(error_buf.as_mut_ptr().cast::<c_void>());
501
502        opt_keys.push(CUjit_option::ErrorLogBufferSizeBytes);
503        opt_vals.push(JIT_LOG_BUFFER_SIZE as *mut c_void);
504
505        // Optimisation level.
506        opt_keys.push(CUjit_option::OptimizationLevel);
507        opt_vals.push(options.optimization_level as *mut c_void);
508
509        // Max registers (only if non-zero to avoid overriding defaults).
510        if options.max_registers > 0 {
511            opt_keys.push(CUjit_option::MaxRegisters);
512            opt_vals.push(options.max_registers as *mut c_void);
513        }
514
515        // Generate debug info.
516        if options.generate_debug_info {
517            opt_keys.push(CUjit_option::GenerateDebugInfo);
518            opt_vals.push(core::ptr::without_provenance_mut::<c_void>(1));
519        }
520
521        // Target from context.
522        if options.target_from_context {
523            opt_keys.push(CUjit_option::TargetFromCuContext);
524            opt_vals.push(core::ptr::without_provenance_mut::<c_void>(1));
525        }
526
527        let num_options = opt_keys.len() as u32;
528
529        let mut raw = CUmodule::default();
530        let result = crate::cuda_call!((api.cu_module_load_data_ex)(
531            &mut raw,
532            c_ptx.as_ptr().cast::<c_void>(),
533            num_options,
534            opt_keys.as_mut_ptr(),
535            opt_vals.as_mut_ptr(),
536        ));
537
538        // Extract log strings regardless of success or failure.
539        let log = JitLog {
540            info: buf_to_string(&info_buf),
541            error: buf_to_string(&error_buf),
542        };
543
544        // On failure, surface the full JIT diagnostic log in the error so
545        // callers can inspect exactly what went wrong.
546        if let Err(e) = result {
547            return Err(jit_failure_from_log(e, log));
548        }
549        Ok((
550            Self {
551                raw,
552                owner: crate::context::current_ctx_owner(),
553            },
554            log,
555        ))
556    }
557
558    /// Retrieves a kernel function by name from this module.
559    ///
560    /// The returned [`Function`] is a lightweight handle. The caller
561    /// must ensure that this `Module` outlives any `Function` handles
562    /// obtained from it.
563    ///
564    /// # Errors
565    ///
566    /// Returns [`CudaError::NotFound`] if no
567    /// function with the given name exists in the module, or another
568    /// [`CudaError`] on driver failure.
569    pub fn get_function(&self, name: &str) -> CudaResult<Function> {
570        let api = try_driver()?;
571        let c_name = CString::new(name).map_err(|_| CudaError::InvalidValue)?;
572        let mut raw = CUfunction::default();
573        crate::cuda_call!((api.cu_module_get_function)(
574            &mut raw,
575            self.raw,
576            c_name.as_ptr()
577        ))?;
578        Ok(Function { raw })
579    }
580
581    /// Returns the raw [`CUmodule`] handle.
582    ///
583    /// # Safety (caller)
584    ///
585    /// The caller must not unload or otherwise invalidate the handle
586    /// while this `Module` is still alive.
587    #[inline]
588    pub fn raw(&self) -> CUmodule {
589        self.raw
590    }
591}
592
593impl Drop for Module {
594    fn drop(&mut self) {
595        // Hold the registry lock across the unload, and skip it entirely if the
596        // owning context was already torn down (its `cuCtxDestroy` already
597        // unloaded this module — calling `cuModuleUnload` again would be a
598        // use-after-free).
599        let map = crate::context::lock_live_ctxs();
600        if !crate::context::owner_is_live(&map, self.owner) {
601            return;
602        }
603        if let Ok(api) = try_driver() {
604            let rc = unsafe { (api.cu_module_unload)(self.raw) };
605            if rc != 0 {
606                tracing::warn!(
607                    cuda_error = rc,
608                    module = ?self.raw,
609                    "cuModuleUnload failed during drop"
610                );
611            }
612        }
613    }
614}
615
616// ---------------------------------------------------------------------------
617// Function
618// ---------------------------------------------------------------------------
619
620/// A kernel function handle within a loaded module.
621///
622/// Functions are lightweight handles (a single pointer) — the lifetime
623/// is tied to the parent [`Module`]. The caller is responsible for
624/// ensuring the `Module` outlives any `Function` handles obtained
625/// from it.
626///
627/// Occupancy query methods are provided in the [`crate::occupancy`]
628/// module via an `impl Function` block.
629#[derive(Debug, Clone, Copy)]
630pub struct Function {
631    /// Raw CUDA function handle.
632    raw: CUfunction,
633}
634
635impl Function {
636    /// Returns the raw [`CUfunction`] handle.
637    ///
638    /// This is needed for kernel launches and occupancy queries
639    /// at the FFI level.
640    #[inline]
641    pub fn raw(&self) -> CUfunction {
642        self.raw
643    }
644}
645
646// ---------------------------------------------------------------------------
647// Helpers
648// ---------------------------------------------------------------------------
649
650/// Converts a null-terminated C buffer to a Rust [`String`], trimming
651/// trailing null bytes and whitespace.
652fn buf_to_string(buf: &[u8]) -> String {
653    // Find the first null byte (or use the whole buffer).
654    let len = buf.iter().position(|&b| b == 0).unwrap_or(buf.len());
655    String::from_utf8_lossy(&buf[..len]).trim().to_string()
656}
657
658// ---------------------------------------------------------------------------
659// Tests
660// ---------------------------------------------------------------------------
661
662#[cfg(test)]
663mod tests {
664    use super::*;
665
666    // ── ascii_only ────────────────────────────────────────────────────────────
667
668    #[test]
669    fn ascii_only_leaves_pure_ascii_untouched() {
670        // Pure-ASCII PTX must not be copied at all.
671        assert!(ascii_only(".version 7.1\n.target sm_86\n").is_none());
672    }
673
674    #[test]
675    fn ascii_only_scrubs_non_ascii_comment() {
676        // The exact shape that made CUDA 12.9+ reject the generated f64 GEMM:
677        // typographic characters inside a PTX comment.
678        let ptx = "    // total_elems = M*N (64-bit: 32×32→64 widening)\n";
679        let scrubbed = ascii_only(ptx).expect("non-ASCII input must be scrubbed");
680        assert!(scrubbed.is_ascii(), "output must be pure ASCII");
681        assert!(scrubbed.starts_with("    // total_elems = M*N (64-bit: 32"));
682        assert!(scrubbed.ends_with(" widening)\n"));
683    }
684
685    #[test]
686    fn ascii_only_preserves_byte_offsets() {
687        // Byte-for-byte substitution keeps JIT line/column diagnostics accurate.
688        let ptx = "// ── x\nmov.u32 %r0, 1;\n";
689        let scrubbed = ascii_only(ptx).expect("non-ASCII input must be scrubbed");
690        assert_eq!(scrubbed.len(), ptx.len());
691        assert_eq!(scrubbed.lines().count(), ptx.lines().count());
692        // Instruction lines are untouched.
693        assert!(scrubbed.contains("mov.u32 %r0, 1;"));
694    }
695
696    // ── parse_ptxas_line ──────────────────────────────────────────────────────
697
698    #[test]
699    fn parse_blank_line_returns_none() {
700        assert!(parse_ptxas_line("").is_none());
701        assert!(parse_ptxas_line("   ").is_none());
702    }
703
704    #[test]
705    fn parse_non_ptxas_line_returns_none() {
706        // Lines not starting with "ptxas " are ignored.
707        assert!(parse_ptxas_line("nvcc error: something").is_none());
708        assert!(parse_ptxas_line("  error: foo").is_none());
709    }
710
711    #[test]
712    fn parse_standard_error_with_kernel_and_line() {
713        let line = "ptxas error   : 'vec_add', line 42; error   : Unknown instruction 'xyz.f32'";
714        let d = parse_ptxas_line(line).expect("should parse");
715        assert_eq!(d.severity, JitSeverity::Error);
716        assert_eq!(d.kernel.as_deref(), Some("vec_add"));
717        assert_eq!(d.line, Some(42));
718        assert!(
719            d.message.contains("Unknown instruction"),
720            "msg: {}",
721            d.message
722        );
723    }
724
725    #[test]
726    fn parse_warning_with_kernel_and_line() {
727        let line = "ptxas warning : 'my_kernel', line 7; warning : Double-precision instructions will be slow";
728        let d = parse_ptxas_line(line).expect("should parse");
729        assert_eq!(d.severity, JitSeverity::Warning);
730        assert_eq!(d.kernel.as_deref(), Some("my_kernel"));
731        assert_eq!(d.line, Some(7));
732        assert!(d.message.contains("Double-precision"), "msg: {}", d.message);
733    }
734
735    #[test]
736    fn parse_info_register_usage() {
737        let line =
738            "ptxas info    : 'reduce_kernel' used 32 registers, 0 bytes smem, 0 bytes cmem[0]";
739        let d = parse_ptxas_line(line).expect("should parse");
740        assert_eq!(d.severity, JitSeverity::Info);
741        assert_eq!(d.kernel.as_deref(), Some("reduce_kernel"));
742        assert!(d.message.contains("32 registers"), "msg: {}", d.message);
743        assert!(d.line.is_none());
744    }
745
746    #[test]
747    fn parse_fatal_no_kernel() {
748        let line = "ptxas fatal   : Unresolved extern function 'missing_func'";
749        let d = parse_ptxas_line(line).expect("should parse");
750        assert_eq!(d.severity, JitSeverity::Fatal);
751        assert!(d.kernel.is_none());
752        assert!(d.message.contains("Unresolved"), "msg: {}", d.message);
753    }
754
755    #[test]
756    fn parse_error_no_kernel_no_line() {
757        let line = "ptxas error : syntax error near token ';'";
758        let d = parse_ptxas_line(line).expect("should parse");
759        assert_eq!(d.severity, JitSeverity::Error);
760        assert!(d.kernel.is_none());
761        assert!(d.line.is_none());
762        assert!(d.message.contains("syntax error"), "msg: {}", d.message);
763    }
764
765    // ── JitLog helpers ────────────────────────────────────────────────────────
766
767    #[test]
768    fn jitlog_is_empty_for_default() {
769        let log = JitLog::default();
770        assert!(log.is_empty());
771        assert!(!log.has_errors());
772    }
773
774    #[test]
775    fn jitlog_has_errors_when_error_buf_nonempty() {
776        let log = JitLog {
777            info: String::new(),
778            error: "ptxas error : something went wrong".to_string(),
779        };
780        assert!(log.has_errors());
781        assert!(!log.is_empty());
782    }
783
784    #[test]
785    fn jitlog_parse_diagnostics_multiline() {
786        let log = JitLog {
787            error: concat!(
788                "ptxas error   : 'k1', line 5; error   : bad opcode\n",
789                "ptxas warning : 'k1', line 8; warning : slow path\n",
790            )
791            .to_string(),
792            info: "ptxas info    : 'k1' used 8 registers, 0 bytes smem\n".to_string(),
793        };
794        let diags = log.parse_diagnostics();
795        assert_eq!(diags.len(), 3);
796        assert_eq!(diags[0].severity, JitSeverity::Error);
797        assert_eq!(diags[1].severity, JitSeverity::Warning);
798        assert_eq!(diags[2].severity, JitSeverity::Info);
799    }
800
801    #[test]
802    fn jitlog_errors_filter() {
803        let log = JitLog {
804            error: concat!(
805                "ptxas error   : 'k', line 1; error : bad\n",
806                "ptxas warning : 'k', line 2; warning : slow\n",
807            )
808            .to_string(),
809            info: "ptxas info    : 'k' used 4 registers\n".to_string(),
810        };
811        let errs = log.errors();
812        assert_eq!(errs.len(), 1);
813        assert_eq!(errs[0].severity, JitSeverity::Error);
814    }
815
816    #[test]
817    fn jitlog_warnings_filter() {
818        let log = JitLog {
819            error: "ptxas warning : 'k', line 3; warning : something slow\n".to_string(),
820            info: String::new(),
821        };
822        let warns = log.warnings();
823        assert_eq!(warns.len(), 1);
824        assert_eq!(warns[0].severity, JitSeverity::Warning);
825        assert_eq!(warns[0].line, Some(3));
826    }
827
828    // ── buf_to_string ─────────────────────────────────────────────────────────
829
830    #[test]
831    fn buf_to_string_null_terminated() {
832        let mut buf = b"hello\0\0\0".to_vec();
833        buf.extend_from_slice(&[0u8; 100]);
834        assert_eq!(buf_to_string(&buf), "hello");
835    }
836
837    #[test]
838    fn buf_to_string_empty() {
839        assert_eq!(buf_to_string(&[0u8; 10]), "");
840    }
841
842    #[test]
843    fn buf_to_string_no_null() {
844        let buf = b"abc".to_vec();
845        assert_eq!(buf_to_string(&buf), "abc");
846    }
847
848    // ── JitSeverity Display ───────────────────────────────────────────────────
849
850    #[test]
851    fn jit_severity_display() {
852        assert_eq!(JitSeverity::Fatal.to_string(), "fatal");
853        assert_eq!(JitSeverity::Error.to_string(), "error");
854        assert_eq!(JitSeverity::Warning.to_string(), "warning");
855        assert_eq!(JitSeverity::Info.to_string(), "info");
856    }
857}