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}