singe_cuda/
nvrtc.rs

1use std::{
2    cell::UnsafeCell,
3    ffi::{CStr, CString},
4    fmt::{self, Display},
5    ptr, result,
6    sync::atomic::{AtomicBool, Ordering},
7};
8
9use singe_core::impl_enum_display;
10use singe_cuda_sys::nvrtc as sys;
11
12use crate::{
13    architecture::GpuArchitecture,
14    error::{Error, Result},
15    module::ModuleImage,
16    try_nvrtc,
17};
18
19#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
20pub struct Version {
21    pub major: i32,
22    pub minor: i32,
23}
24
25/// NVRTC result code returned by compiler operations.
26#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
27#[non_exhaustive]
28pub enum Status {
29    Success,
30    OutOfMemory,
31    ProgramCreationFailure,
32    InvalidInput,
33    InvalidProgram,
34    InvalidOption,
35    Compilation,
36    BuiltinOperationFailure,
37    NoNameExpressionsAfterCompilation,
38    NoLoweredNamesBeforeCompilation,
39    NameExpressionNotValid,
40    InternalError,
41    TimeFileWriteFailed,
42    NoPchCreateAttempted,
43    PchCreateHeapExhausted,
44    PchCreate,
45    Cancelled,
46    TimeTraceFileWriteFailed,
47    Unknown(u32),
48}
49
50impl Status {
51    pub fn description(self) -> String {
52        match sys::nvrtcResult::try_from(self.raw()) {
53            Ok(status) => unsafe {
54                let ptr = sys::nvrtcGetErrorString(status);
55                if ptr.is_null() {
56                    String::from("unknown nvrtc error")
57                } else {
58                    CStr::from_ptr(ptr).to_string_lossy().into_owned()
59                }
60            },
61            Err(_) => String::from("unknown nvrtc error"),
62        }
63    }
64
65    pub const fn raw(self) -> u32 {
66        match self {
67            Self::Success => sys::nvrtcResult::NVRTC_SUCCESS as _,
68            Self::OutOfMemory => sys::nvrtcResult::NVRTC_ERROR_OUT_OF_MEMORY as _,
69            Self::ProgramCreationFailure => {
70                sys::nvrtcResult::NVRTC_ERROR_PROGRAM_CREATION_FAILURE as _
71            }
72            Self::InvalidInput => sys::nvrtcResult::NVRTC_ERROR_INVALID_INPUT as _,
73            Self::InvalidProgram => sys::nvrtcResult::NVRTC_ERROR_INVALID_PROGRAM as _,
74            Self::InvalidOption => sys::nvrtcResult::NVRTC_ERROR_INVALID_OPTION as _,
75            Self::Compilation => sys::nvrtcResult::NVRTC_ERROR_COMPILATION as _,
76            Self::BuiltinOperationFailure => {
77                sys::nvrtcResult::NVRTC_ERROR_BUILTIN_OPERATION_FAILURE as _
78            }
79            Self::NoNameExpressionsAfterCompilation => {
80                sys::nvrtcResult::NVRTC_ERROR_NO_NAME_EXPRESSIONS_AFTER_COMPILATION as _
81            }
82            Self::NoLoweredNamesBeforeCompilation => {
83                sys::nvrtcResult::NVRTC_ERROR_NO_LOWERED_NAMES_BEFORE_COMPILATION as _
84            }
85            Self::NameExpressionNotValid => {
86                sys::nvrtcResult::NVRTC_ERROR_NAME_EXPRESSION_NOT_VALID as _
87            }
88            Self::InternalError => sys::nvrtcResult::NVRTC_ERROR_INTERNAL_ERROR as _,
89            Self::TimeFileWriteFailed => sys::nvrtcResult::NVRTC_ERROR_TIME_FILE_WRITE_FAILED as _,
90            Self::NoPchCreateAttempted => {
91                sys::nvrtcResult::NVRTC_ERROR_NO_PCH_CREATE_ATTEMPTED as _
92            }
93            Self::PchCreateHeapExhausted => {
94                sys::nvrtcResult::NVRTC_ERROR_PCH_CREATE_HEAP_EXHAUSTED as _
95            }
96            Self::PchCreate => sys::nvrtcResult::NVRTC_ERROR_PCH_CREATE as _,
97            Self::Cancelled => sys::nvrtcResult::NVRTC_ERROR_CANCELLED as _,
98            Self::TimeTraceFileWriteFailed => {
99                sys::nvrtcResult::NVRTC_ERROR_TIME_TRACE_FILE_WRITE_FAILED as _
100            }
101            Self::Unknown(code) => code,
102        }
103    }
104}
105
106impl TryFrom<u32> for Status {
107    type Error = u32;
108
109    fn try_from(code: u32) -> result::Result<Self, u32> {
110        match code {
111            code if code == sys::nvrtcResult::NVRTC_SUCCESS as u32 => Ok(Self::Success),
112            code if code == sys::nvrtcResult::NVRTC_ERROR_OUT_OF_MEMORY as u32 => {
113                Ok(Self::OutOfMemory)
114            }
115            code if code == sys::nvrtcResult::NVRTC_ERROR_PROGRAM_CREATION_FAILURE as u32 => {
116                Ok(Self::ProgramCreationFailure)
117            }
118            code if code == sys::nvrtcResult::NVRTC_ERROR_INVALID_INPUT as u32 => {
119                Ok(Self::InvalidInput)
120            }
121            code if code == sys::nvrtcResult::NVRTC_ERROR_INVALID_PROGRAM as u32 => {
122                Ok(Self::InvalidProgram)
123            }
124            code if code == sys::nvrtcResult::NVRTC_ERROR_INVALID_OPTION as u32 => {
125                Ok(Self::InvalidOption)
126            }
127            code if code == sys::nvrtcResult::NVRTC_ERROR_COMPILATION as u32 => {
128                Ok(Self::Compilation)
129            }
130            code if code == sys::nvrtcResult::NVRTC_ERROR_BUILTIN_OPERATION_FAILURE as u32 => {
131                Ok(Self::BuiltinOperationFailure)
132            }
133            code if code
134                == sys::nvrtcResult::NVRTC_ERROR_NO_NAME_EXPRESSIONS_AFTER_COMPILATION as u32 =>
135            {
136                Ok(Self::NoNameExpressionsAfterCompilation)
137            }
138            code if code
139                == sys::nvrtcResult::NVRTC_ERROR_NO_LOWERED_NAMES_BEFORE_COMPILATION as u32 =>
140            {
141                Ok(Self::NoLoweredNamesBeforeCompilation)
142            }
143            code if code == sys::nvrtcResult::NVRTC_ERROR_NAME_EXPRESSION_NOT_VALID as u32 => {
144                Ok(Self::NameExpressionNotValid)
145            }
146            code if code == sys::nvrtcResult::NVRTC_ERROR_INTERNAL_ERROR as u32 => {
147                Ok(Self::InternalError)
148            }
149            code if code == sys::nvrtcResult::NVRTC_ERROR_TIME_FILE_WRITE_FAILED as u32 => {
150                Ok(Self::TimeFileWriteFailed)
151            }
152            code if code == sys::nvrtcResult::NVRTC_ERROR_NO_PCH_CREATE_ATTEMPTED as u32 => {
153                Ok(Self::NoPchCreateAttempted)
154            }
155            code if code == sys::nvrtcResult::NVRTC_ERROR_PCH_CREATE_HEAP_EXHAUSTED as u32 => {
156                Ok(Self::PchCreateHeapExhausted)
157            }
158            code if code == sys::nvrtcResult::NVRTC_ERROR_PCH_CREATE as u32 => Ok(Self::PchCreate),
159            code if code == sys::nvrtcResult::NVRTC_ERROR_CANCELLED as u32 => Ok(Self::Cancelled),
160            code if code == sys::nvrtcResult::NVRTC_ERROR_TIME_TRACE_FILE_WRITE_FAILED as u32 => {
161                Ok(Self::TimeTraceFileWriteFailed)
162            }
163            code => Err(code),
164        }
165    }
166}
167
168impl From<sys::nvrtcResult> for Status {
169    fn from(status: sys::nvrtcResult) -> Self {
170        Self::try_from(status as u32).unwrap_or_else(Self::Unknown)
171    }
172}
173
174impl TryFrom<Status> for sys::nvrtcResult {
175    type Error = Status;
176
177    fn try_from(status: Status) -> result::Result<Self, Status> {
178        match status {
179            Status::Success => Ok(Self::NVRTC_SUCCESS),
180            Status::OutOfMemory => Ok(Self::NVRTC_ERROR_OUT_OF_MEMORY),
181            Status::ProgramCreationFailure => Ok(Self::NVRTC_ERROR_PROGRAM_CREATION_FAILURE),
182            Status::InvalidInput => Ok(Self::NVRTC_ERROR_INVALID_INPUT),
183            Status::InvalidProgram => Ok(Self::NVRTC_ERROR_INVALID_PROGRAM),
184            Status::InvalidOption => Ok(Self::NVRTC_ERROR_INVALID_OPTION),
185            Status::Compilation => Ok(Self::NVRTC_ERROR_COMPILATION),
186            Status::BuiltinOperationFailure => Ok(Self::NVRTC_ERROR_BUILTIN_OPERATION_FAILURE),
187            Status::NoNameExpressionsAfterCompilation => {
188                Ok(Self::NVRTC_ERROR_NO_NAME_EXPRESSIONS_AFTER_COMPILATION)
189            }
190            Status::NoLoweredNamesBeforeCompilation => {
191                Ok(Self::NVRTC_ERROR_NO_LOWERED_NAMES_BEFORE_COMPILATION)
192            }
193            Status::NameExpressionNotValid => Ok(Self::NVRTC_ERROR_NAME_EXPRESSION_NOT_VALID),
194            Status::InternalError => Ok(Self::NVRTC_ERROR_INTERNAL_ERROR),
195            Status::TimeFileWriteFailed => Ok(Self::NVRTC_ERROR_TIME_FILE_WRITE_FAILED),
196            Status::NoPchCreateAttempted => Ok(Self::NVRTC_ERROR_NO_PCH_CREATE_ATTEMPTED),
197            Status::PchCreateHeapExhausted => Ok(Self::NVRTC_ERROR_PCH_CREATE_HEAP_EXHAUSTED),
198            Status::PchCreate => Ok(Self::NVRTC_ERROR_PCH_CREATE),
199            Status::Cancelled => Ok(Self::NVRTC_ERROR_CANCELLED),
200            Status::TimeTraceFileWriteFailed => Ok(Self::NVRTC_ERROR_TIME_TRACE_FILE_WRITE_FAILED),
201            Status::Unknown(_) => Err(status),
202        }
203    }
204}
205
206impl Display for Status {
207    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
208        match self {
209            Self::Success => f.write_str("NVRTC_SUCCESS"),
210            Self::OutOfMemory => f.write_str("NVRTC_ERROR_OUT_OF_MEMORY"),
211            Self::ProgramCreationFailure => f.write_str("NVRTC_ERROR_PROGRAM_CREATION_FAILURE"),
212            Self::InvalidInput => f.write_str("NVRTC_ERROR_INVALID_INPUT"),
213            Self::InvalidProgram => f.write_str("NVRTC_ERROR_INVALID_PROGRAM"),
214            Self::InvalidOption => f.write_str("NVRTC_ERROR_INVALID_OPTION"),
215            Self::Compilation => f.write_str("NVRTC_ERROR_COMPILATION"),
216            Self::BuiltinOperationFailure => f.write_str("NVRTC_ERROR_BUILTIN_OPERATION_FAILURE"),
217            Self::NoNameExpressionsAfterCompilation => {
218                f.write_str("NVRTC_ERROR_NO_NAME_EXPRESSIONS_AFTER_COMPILATION")
219            }
220            Self::NoLoweredNamesBeforeCompilation => {
221                f.write_str("NVRTC_ERROR_NO_LOWERED_NAMES_BEFORE_COMPILATION")
222            }
223            Self::NameExpressionNotValid => f.write_str("NVRTC_ERROR_NAME_EXPRESSION_NOT_VALID"),
224            Self::InternalError => f.write_str("NVRTC_ERROR_INTERNAL_ERROR"),
225            Self::TimeFileWriteFailed => f.write_str("NVRTC_ERROR_TIME_FILE_WRITE_FAILED"),
226            Self::NoPchCreateAttempted => f.write_str("NVRTC_ERROR_NO_PCH_CREATE_ATTEMPTED"),
227            Self::PchCreateHeapExhausted => f.write_str("NVRTC_ERROR_PCH_CREATE_HEAP_EXHAUSTED"),
228            Self::PchCreate => f.write_str("NVRTC_ERROR_PCH_CREATE"),
229            Self::Cancelled => f.write_str("NVRTC_ERROR_CANCELLED"),
230            Self::TimeTraceFileWriteFailed => {
231                f.write_str("NVRTC_ERROR_TIME_TRACE_FILE_WRITE_FAILED")
232            }
233            Self::Unknown(code) => write!(f, "UNKNOWN_NVRTC_STATUS({code})"),
234        }
235    }
236}
237
238#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
239pub struct Header<'a> {
240    pub source: &'a str,
241    pub include_name: &'a str,
242}
243
244#[derive(Debug, Clone)]
245struct OwnedHeader {
246    source: String,
247    include_name: String,
248}
249
250#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
251#[non_exhaustive]
252pub enum MacroDefinition<'a> {
253    Name(&'a str),
254    WithValue { name: &'a str, value: &'a str },
255}
256
257impl MacroDefinition<'_> {
258    fn format(self) -> String {
259        match self {
260            Self::Name(name) => name.to_string(),
261            Self::WithValue { name, value } => format!("{name}={value}"),
262        }
263    }
264}
265
266#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
267#[non_exhaustive]
268pub enum CppDialect {
269    Cpp03,
270    Cpp11,
271    Cpp14,
272    Cpp17,
273    Cpp20,
274}
275
276impl_enum_display!(CppDialect, {
277    Self::Cpp03 => "c++03",
278    Self::Cpp11 => "c++11",
279    Self::Cpp14 => "c++14",
280    Self::Cpp17 => "c++17",
281    Self::Cpp20 => "c++20",
282});
283
284#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
285#[non_exhaustive]
286pub enum FastCompileLevel {
287    Zero,
288    Min,
289    Mid,
290    Max,
291}
292
293impl_enum_display!(FastCompileLevel, {
294    Self::Zero => "0",
295    Self::Min => "min",
296    Self::Mid => "mid",
297    Self::Max => "max",
298});
299
300#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
301#[non_exhaustive]
302pub enum WarningAsErrorKind {
303    AllWarnings,
304    Reorder,
305    DeprecatedDeclarations,
306}
307
308impl_enum_display!(WarningAsErrorKind, {
309    Self::AllWarnings => "all-warnings",
310    Self::Reorder => "reorder",
311    Self::DeprecatedDeclarations => "deprecated-declarations",
312});
313
314#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
315#[non_exhaustive]
316pub enum OptimizationInfoKind {
317    Inline,
318}
319
320impl_enum_display!(OptimizationInfoKind, {
321    Self::Inline => "inline",
322});
323
324#[derive(Debug, Clone, Default)]
325pub struct CompileOptions<'a> {
326    pub gpu_architecture: Option<GpuArchitecture>,
327    pub relocatable_device_code: Option<bool>,
328    pub extensible_whole_program: bool,
329    pub device_debug: bool,
330    pub generate_line_info: bool,
331    pub device_optimization: Option<bool>,
332    pub fast_compile: Option<FastCompileLevel>,
333    pub ptxas_options: Vec<&'a str>,
334    pub max_register_count: Option<i32>,
335    pub flush_to_zero: Option<bool>,
336    pub precise_square_root: Option<bool>,
337    pub precise_division: Option<bool>,
338    pub fmad: Option<bool>,
339    pub use_fast_math: bool,
340    pub extra_device_vectorization: bool,
341    pub modify_stack_limit: Option<bool>,
342    pub dlink_time_optimization: bool,
343    pub generate_optimized_lto: bool,
344    pub optix_ir: bool,
345    pub jump_table_density: Option<u8>,
346    pub no_cache: bool,
347    pub random_seed: Option<&'a str>,
348    pub define_macros: Vec<MacroDefinition<'a>>,
349    pub undefine_macros: Vec<&'a str>,
350    pub include_paths: Vec<&'a str>,
351    pub pre_include_headers: Vec<&'a str>,
352    pub no_source_include: bool,
353    pub cpp_dialect: Option<CppDialect>,
354    pub builtin_move_forward: Option<bool>,
355    pub builtin_initializer_list: Option<bool>,
356    pub pch: bool,
357    pub create_pch: Option<&'a str>,
358    pub use_pch: Option<&'a str>,
359    pub pch_dir: Option<&'a str>,
360    pub pch_verbose: Option<bool>,
361    pub pch_messages: Option<bool>,
362    pub instantiate_templates_in_pch: Option<bool>,
363    pub disable_warnings: bool,
364    pub warning_as_error: Vec<WarningAsErrorKind>,
365    pub restrict_pointers: bool,
366    pub device_as_default_execution_space: bool,
367    pub device_int128: bool,
368    pub device_float128: bool,
369    pub optimization_info: Vec<OptimizationInfoKind>,
370    pub display_error_number: Option<bool>,
371    pub diag_error: Vec<i32>,
372    pub diag_suppress: Vec<i32>,
373    pub diag_warn: Vec<i32>,
374    pub brief_diagnostics: Option<bool>,
375    pub time: Option<&'a str>,
376    pub split_compile: Option<i32>,
377    pub device_syntax_only: bool,
378    pub minimal: bool,
379    pub device_stack_protector: Option<bool>,
380    pub device_time_trace: Option<&'a str>,
381    pub raw_options: Vec<&'a str>,
382}
383
384impl<'a> CompileOptions<'a> {
385    pub const fn new() -> Self {
386        Self {
387            gpu_architecture: None,
388            relocatable_device_code: None,
389            extensible_whole_program: false,
390            device_debug: false,
391            generate_line_info: false,
392            device_optimization: None,
393            fast_compile: None,
394            ptxas_options: Vec::new(),
395            max_register_count: None,
396            flush_to_zero: None,
397            precise_square_root: None,
398            precise_division: None,
399            fmad: None,
400            use_fast_math: false,
401            extra_device_vectorization: false,
402            modify_stack_limit: None,
403            dlink_time_optimization: false,
404            generate_optimized_lto: false,
405            optix_ir: false,
406            jump_table_density: None,
407            no_cache: false,
408            random_seed: None,
409            define_macros: Vec::new(),
410            undefine_macros: Vec::new(),
411            include_paths: Vec::new(),
412            pre_include_headers: Vec::new(),
413            no_source_include: false,
414            cpp_dialect: None,
415            builtin_move_forward: None,
416            builtin_initializer_list: None,
417            pch: false,
418            create_pch: None,
419            use_pch: None,
420            pch_dir: None,
421            pch_verbose: None,
422            pch_messages: None,
423            instantiate_templates_in_pch: None,
424            disable_warnings: false,
425            warning_as_error: Vec::new(),
426            restrict_pointers: false,
427            device_as_default_execution_space: false,
428            device_int128: false,
429            device_float128: false,
430            optimization_info: Vec::new(),
431            display_error_number: None,
432            diag_error: Vec::new(),
433            diag_suppress: Vec::new(),
434            diag_warn: Vec::new(),
435            brief_diagnostics: None,
436            time: None,
437            split_compile: None,
438            device_syntax_only: false,
439            minimal: false,
440            device_stack_protector: None,
441            device_time_trace: None,
442            raw_options: Vec::new(),
443        }
444    }
445
446    pub fn gpu_architecture(mut self, value: GpuArchitecture) -> Self {
447        self.gpu_architecture = Some(value);
448        self
449    }
450
451    pub fn relocatable_device_code(mut self, value: bool) -> Self {
452        self.relocatable_device_code = Some(value);
453        self
454    }
455
456    pub fn extensible_whole_program(mut self, value: bool) -> Self {
457        self.extensible_whole_program = value;
458        self
459    }
460
461    pub fn device_debug(mut self, value: bool) -> Self {
462        self.device_debug = value;
463        self
464    }
465
466    pub fn generate_line_info(mut self, value: bool) -> Self {
467        self.generate_line_info = value;
468        self
469    }
470
471    pub fn device_optimization(mut self, value: bool) -> Self {
472        self.device_optimization = Some(value);
473        self
474    }
475
476    pub fn fast_compile(mut self, value: FastCompileLevel) -> Self {
477        self.fast_compile = Some(value);
478        self
479    }
480
481    pub fn ptxas_option(mut self, value: &'a str) -> Self {
482        self.ptxas_options.push(value);
483        self
484    }
485
486    pub fn max_register_count(mut self, value: i32) -> Self {
487        self.max_register_count = Some(value);
488        self
489    }
490
491    pub fn flush_to_zero(mut self, value: bool) -> Self {
492        self.flush_to_zero = Some(value);
493        self
494    }
495
496    pub fn precise_square_root(mut self, value: bool) -> Self {
497        self.precise_square_root = Some(value);
498        self
499    }
500
501    pub fn precise_division(mut self, value: bool) -> Self {
502        self.precise_division = Some(value);
503        self
504    }
505
506    pub fn fmad(mut self, value: bool) -> Self {
507        self.fmad = Some(value);
508        self
509    }
510
511    pub fn use_fast_math(mut self, value: bool) -> Self {
512        self.use_fast_math = value;
513        self
514    }
515
516    pub fn extra_device_vectorization(mut self, value: bool) -> Self {
517        self.extra_device_vectorization = value;
518        self
519    }
520
521    pub fn modify_stack_limit(mut self, value: bool) -> Self {
522        self.modify_stack_limit = Some(value);
523        self
524    }
525
526    pub fn dlink_time_optimization(mut self, value: bool) -> Self {
527        self.dlink_time_optimization = value;
528        self
529    }
530
531    pub fn generate_optimized_lto(mut self, value: bool) -> Self {
532        self.generate_optimized_lto = value;
533        self
534    }
535
536    pub fn optix_ir(mut self, value: bool) -> Self {
537        self.optix_ir = value;
538        self
539    }
540
541    pub fn jump_table_density(mut self, value: u8) -> Self {
542        self.jump_table_density = Some(value.min(101));
543        self
544    }
545
546    pub fn no_cache(mut self, value: bool) -> Self {
547        self.no_cache = value;
548        self
549    }
550
551    pub fn random_seed(mut self, value: &'a str) -> Self {
552        self.random_seed = Some(value);
553        self
554    }
555
556    pub fn define_macro(mut self, value: MacroDefinition<'a>) -> Self {
557        self.define_macros.push(value);
558        self
559    }
560
561    pub fn undefine_macro(mut self, value: &'a str) -> Self {
562        self.undefine_macros.push(value);
563        self
564    }
565
566    pub fn include_path(mut self, value: &'a str) -> Self {
567        self.include_paths.push(value);
568        self
569    }
570
571    pub fn pre_include_header(mut self, value: &'a str) -> Self {
572        self.pre_include_headers.push(value);
573        self
574    }
575
576    pub fn no_source_include(mut self, value: bool) -> Self {
577        self.no_source_include = value;
578        self
579    }
580
581    pub fn cpp_dialect(mut self, value: CppDialect) -> Self {
582        self.cpp_dialect = Some(value);
583        self
584    }
585
586    pub fn builtin_move_forward(mut self, value: bool) -> Self {
587        self.builtin_move_forward = Some(value);
588        self
589    }
590
591    pub fn builtin_initializer_list(mut self, value: bool) -> Self {
592        self.builtin_initializer_list = Some(value);
593        self
594    }
595
596    pub fn pch(mut self, value: bool) -> Self {
597        self.pch = value;
598        self
599    }
600
601    pub fn create_pch(mut self, value: &'a str) -> Self {
602        self.create_pch = Some(value);
603        self
604    }
605
606    pub fn use_pch(mut self, value: &'a str) -> Self {
607        self.use_pch = Some(value);
608        self
609    }
610
611    pub fn pch_dir(mut self, value: &'a str) -> Self {
612        self.pch_dir = Some(value);
613        self
614    }
615
616    pub fn pch_verbose(mut self, value: bool) -> Self {
617        self.pch_verbose = Some(value);
618        self
619    }
620
621    pub fn pch_messages(mut self, value: bool) -> Self {
622        self.pch_messages = Some(value);
623        self
624    }
625
626    pub fn instantiate_templates_in_pch(mut self, value: bool) -> Self {
627        self.instantiate_templates_in_pch = Some(value);
628        self
629    }
630
631    pub fn disable_warnings(mut self, value: bool) -> Self {
632        self.disable_warnings = value;
633        self
634    }
635
636    pub fn warning_as_error(mut self, value: WarningAsErrorKind) -> Self {
637        self.warning_as_error.push(value);
638        self
639    }
640
641    pub fn restrict_pointers(mut self, value: bool) -> Self {
642        self.restrict_pointers = value;
643        self
644    }
645
646    pub fn device_as_default_execution_space(mut self, value: bool) -> Self {
647        self.device_as_default_execution_space = value;
648        self
649    }
650
651    pub fn device_int128(mut self, value: bool) -> Self {
652        self.device_int128 = value;
653        self
654    }
655
656    pub fn device_float128(mut self, value: bool) -> Self {
657        self.device_float128 = value;
658        self
659    }
660
661    pub fn optimization_info(mut self, value: OptimizationInfoKind) -> Self {
662        self.optimization_info.push(value);
663        self
664    }
665
666    pub fn display_error_number(mut self, value: bool) -> Self {
667        self.display_error_number = Some(value);
668        self
669    }
670
671    pub fn diag_error(mut self, value: i32) -> Self {
672        self.diag_error.push(value);
673        self
674    }
675
676    pub fn diag_suppress(mut self, value: i32) -> Self {
677        self.diag_suppress.push(value);
678        self
679    }
680
681    pub fn diag_warn(mut self, value: i32) -> Self {
682        self.diag_warn.push(value);
683        self
684    }
685
686    pub fn brief_diagnostics(mut self, value: bool) -> Self {
687        self.brief_diagnostics = Some(value);
688        self
689    }
690
691    pub fn time(mut self, value: &'a str) -> Self {
692        self.time = Some(value);
693        self
694    }
695
696    pub fn split_compile(mut self, value: i32) -> Self {
697        self.split_compile = Some(value);
698        self
699    }
700
701    pub fn device_syntax_only(mut self, value: bool) -> Self {
702        self.device_syntax_only = value;
703        self
704    }
705
706    pub fn minimal(mut self, value: bool) -> Self {
707        self.minimal = value;
708        self
709    }
710
711    pub fn device_stack_protector(mut self, value: bool) -> Self {
712        self.device_stack_protector = Some(value);
713        self
714    }
715
716    pub fn device_time_trace(mut self, value: &'a str) -> Self {
717        self.device_time_trace = Some(value);
718        self
719    }
720
721    pub fn raw_option(mut self, value: &'a str) -> Self {
722        self.raw_options.push(value);
723        self
724    }
725
726    pub fn as_arguments(&self) -> Vec<String> {
727        let mut arguments = Vec::new();
728
729        if let Some(value) = self.gpu_architecture {
730            arguments.push(format!("--gpu-architecture={value}"));
731        }
732        if let Some(value) = self.relocatable_device_code {
733            arguments.push(format!("--relocatable-device-code={}", bool_flag(value)));
734        }
735        if self.extensible_whole_program {
736            arguments.push(String::from("--extensible-whole-program"));
737        }
738        if self.device_debug {
739            arguments.push(String::from("--device-debug"));
740        }
741        if self.generate_line_info {
742            arguments.push(String::from("--generate-line-info"));
743        }
744        if let Some(value) = self.device_optimization
745            && value
746        {
747            arguments.push(String::from("--dopt=on"));
748        }
749        if let Some(value) = self.fast_compile {
750            arguments.push(format!("--Ofast-compile={value}"));
751        }
752        arguments.extend(
753            self.ptxas_options
754                .iter()
755                .map(|value| format!("--ptxas-options={value}")),
756        );
757        if let Some(value) = self.max_register_count {
758            arguments.push(format!("--maxrregcount={value}"));
759        }
760        if let Some(value) = self.flush_to_zero {
761            arguments.push(format!("--ftz={}", bool_flag(value)));
762        }
763        if let Some(value) = self.precise_square_root {
764            arguments.push(format!("--prec-sqrt={}", bool_flag(value)));
765        }
766        if let Some(value) = self.precise_division {
767            arguments.push(format!("--prec-div={}", bool_flag(value)));
768        }
769        if let Some(value) = self.fmad {
770            arguments.push(format!("--fmad={}", bool_flag(value)));
771        }
772        if self.use_fast_math {
773            arguments.push(String::from("--use_fast_math"));
774        }
775        if self.extra_device_vectorization {
776            arguments.push(String::from("--extra-device-vectorization"));
777        }
778        if let Some(value) = self.modify_stack_limit {
779            arguments.push(format!("--modify-stack-limit={}", bool_flag(value)));
780        }
781        if self.dlink_time_optimization {
782            arguments.push(String::from("--dlink-time-opt"));
783        }
784        if self.generate_optimized_lto {
785            arguments.push(String::from("--gen-opt-lto"));
786        }
787        if self.optix_ir {
788            arguments.push(String::from("--optix-ir"));
789        }
790        if let Some(value) = self.jump_table_density {
791            arguments.push(format!("--jump-table-density={value}"));
792        }
793        if self.no_cache {
794            arguments.push(String::from("--no-cache"));
795        }
796        if let Some(value) = self.random_seed {
797            arguments.push(format!("--frandom-seed={value}"));
798        }
799        arguments.extend(
800            self.define_macros
801                .iter()
802                .copied()
803                .map(|value| format!("--define-macro={}", value.format())),
804        );
805        arguments.extend(
806            self.undefine_macros
807                .iter()
808                .map(|value| format!("--undefine-macro={value}")),
809        );
810        arguments.extend(
811            self.include_paths
812                .iter()
813                .map(|value| format!("--include-path={value}")),
814        );
815        arguments.extend(
816            self.pre_include_headers
817                .iter()
818                .map(|value| format!("--pre-include={value}")),
819        );
820        if self.no_source_include {
821            arguments.push(String::from("--no-source-include"));
822        }
823        if let Some(value) = self.cpp_dialect {
824            arguments.push(format!("--std={value}"));
825        }
826        if let Some(value) = self.builtin_move_forward {
827            arguments.push(format!("--builtin-move-forward={}", bool_flag(value)));
828        }
829        if let Some(value) = self.builtin_initializer_list {
830            arguments.push(format!("--builtin-initializer-list={}", bool_flag(value)));
831        }
832        if self.pch {
833            arguments.push(String::from("--pch"));
834        }
835        if let Some(value) = self.create_pch {
836            arguments.push(format!("--create-pch={value}"));
837        }
838        if let Some(value) = self.use_pch {
839            arguments.push(format!("--use-pch={value}"));
840        }
841        if let Some(value) = self.pch_dir {
842            arguments.push(format!("--pch-dir={value}"));
843        }
844        if let Some(value) = self.pch_verbose {
845            arguments.push(format!("--pch-verbose={}", bool_flag(value)));
846        }
847        if let Some(value) = self.pch_messages {
848            arguments.push(format!("--pch-messages={}", bool_flag(value)));
849        }
850        if let Some(value) = self.instantiate_templates_in_pch {
851            arguments.push(format!(
852                "--instantiate-templates-in-pch={}",
853                bool_flag(value)
854            ));
855        }
856        if self.disable_warnings {
857            arguments.push(String::from("--disable-warnings"));
858        }
859        if !self.warning_as_error.is_empty() {
860            arguments.push(format!(
861                "--warning-as-error={}",
862                join_display(&self.warning_as_error)
863            ));
864        }
865        if self.restrict_pointers {
866            arguments.push(String::from("--restrict"));
867        }
868        if self.device_as_default_execution_space {
869            arguments.push(String::from("--device-as-default-execution-space"));
870        }
871        if self.device_int128 {
872            arguments.push(String::from("--device-int128"));
873        }
874        if self.device_float128 {
875            arguments.push(String::from("--device-float128"));
876        }
877        arguments.extend(
878            self.optimization_info
879                .iter()
880                .map(|value| format!("--optimization-info={value}")),
881        );
882        if let Some(value) = self.display_error_number {
883            arguments.push(if value {
884                String::from("--display-error-number")
885            } else {
886                String::from("--no-display-error-number")
887            });
888        }
889        if !self.diag_error.is_empty() {
890            arguments.push(format!("--diag-error={}", join_numbers(&self.diag_error)));
891        }
892        if !self.diag_suppress.is_empty() {
893            arguments.push(format!(
894                "--diag-suppress={}",
895                join_numbers(&self.diag_suppress)
896            ));
897        }
898        if !self.diag_warn.is_empty() {
899            arguments.push(format!("--diag-warn={}", join_numbers(&self.diag_warn)));
900        }
901        if let Some(value) = self.brief_diagnostics {
902            arguments.push(format!("--brief-diagnostics={}", bool_flag(value)));
903        }
904        if let Some(value) = self.time {
905            arguments.push(format!("--time={value}"));
906        }
907        if let Some(value) = self.split_compile {
908            arguments.push(format!("--split-compile={value}"));
909        }
910        if self.device_syntax_only {
911            arguments.push(String::from("--fdevice-syntax-only"));
912        }
913        if self.minimal {
914            arguments.push(String::from("--minimal"));
915        }
916        if let Some(value) = self.device_stack_protector {
917            arguments.push(format!("--device-stack-protector={}", bool_flag(value)));
918        }
919        if let Some(value) = self.device_time_trace {
920            arguments.push(format!("--fdevice-time-trace={value}"));
921        }
922
923        arguments.extend(self.raw_options.iter().map(|value| (*value).to_string()));
924        arguments
925    }
926}
927
928#[derive(Debug)]
929pub struct Program {
930    source: String,
931    name: Option<String>,
932    headers: Vec<OwnedHeader>,
933    handle: UnsafeCell<sys::nvrtcProgram>,
934}
935
936impl Program {
937    /// Creates a lazy NVRTC program from CUDA source text.
938    ///
939    /// The raw NVRTC handle is created when the program is compiled or
940    /// otherwise materialized.
941    pub fn from_source(source: &str) -> Self {
942        Self {
943            source: source.to_string(),
944            name: None,
945            headers: Vec::new(),
946            handle: UnsafeCell::new(ptr::null_mut()),
947        }
948    }
949
950    pub fn with_name(mut self, name: &str) -> Self {
951        self.name = Some(name.to_string());
952        self
953    }
954
955    pub fn with_header(mut self, header: Header<'_>) -> Self {
956        self.headers.push(OwnedHeader {
957            source: header.source.to_string(),
958            include_name: header.include_name.to_string(),
959        });
960        self
961    }
962
963    pub fn with_headers(mut self, headers: &[Header<'_>]) -> Self {
964        self.headers
965            .extend(headers.iter().map(|header| OwnedHeader {
966                source: header.source.to_string(),
967                include_name: header.include_name.to_string(),
968            }));
969        self
970    }
971
972    /// Takes ownership of a raw NVRTC program handle.
973    ///
974    /// # Safety
975    ///
976    /// `handle` must be a valid `nvrtcProgram` that is not owned by any other
977    /// wrapper. The returned wrapper destroys it with `nvrtcDestroyProgram`.
978    pub unsafe fn from_raw(handle: sys::nvrtcProgram) -> Result<Self> {
979        if handle.is_null() {
980            return Err(Error::NullHandle);
981        }
982
983        Ok(Self {
984            source: String::new(),
985            name: None,
986            headers: Vec::new(),
987            handle: UnsafeCell::new(handle),
988        })
989    }
990
991    pub fn compile(&self, options: &[&str]) -> Result<()> {
992        self.compile_raw(options)
993    }
994
995    pub fn compile_with_options(&self, options: &CompileOptions<'_>) -> Result<()> {
996        let arguments = options.as_arguments();
997        let argument_refs = arguments.iter().map(String::as_str).collect::<Vec<_>>();
998        self.compile_raw(&argument_refs)
999    }
1000
1001    /// Registers a flow callback that can cancel compilation.
1002    ///
1003    /// The callback function must satisfy the following constraints:
1004    ///
1005    /// (1) It must return 1 to cancel compilation or 0 to continue.
1006    /// Other return values are reserved for future use.
1007    ///
1008    /// (2) It must return consistent values.
1009    /// Once it returns 1 at one point, it must return 1 in all following invocations during the current nvrtcCompileProgram call in progress.
1010    ///
1011    /// (3) It must be thread-safe.
1012    ///
1013    /// (4) It must not invoke any NVRTC, libNVVM, or PTX APIs.
1014    pub fn compile_with_options_and_cancel_flag(
1015        &self,
1016        options: &CompileOptions<'_>,
1017        cancel: &AtomicBool,
1018    ) -> Result<()> {
1019        unsafe {
1020            try_nvrtc!(sys::nvrtcSetFlowCallback(
1021                self.handle()?,
1022                Some(cancel_if_requested_callback),
1023                ptr::from_ref(cancel).cast_mut().cast(),
1024            ))?;
1025        }
1026
1027        let compile_result = self.compile_with_options(options);
1028        let clear_result = clear_flow_callback(self);
1029
1030        match (compile_result, clear_result) {
1031            (Err(error), _) => Err(error),
1032            (Ok(()), Err(error)) => Err(error),
1033            (Ok(()), Ok(())) => Ok(()),
1034        }
1035    }
1036
1037    /// Registers a name expression for a `__global__`, `__device__`, or `__constant__` symbol.
1038    ///
1039    /// The identical name expression string must be provided to [`Program::lowered_name`] after compilation.
1040    ///
1041    /// # Errors
1042    ///
1043    /// Returns an error if `name_expression` contains an interior NUL byte, if
1044    /// the program handle is invalid, or if NVRTC rejects the expression.
1045    pub fn add_name_expression(&self, name_expression: &str) -> Result<()> {
1046        let name_expression = CString::new(name_expression)?;
1047        unsafe {
1048            try_nvrtc!(sys::nvrtcAddNameExpression(
1049                self.handle()?,
1050                name_expression.as_ptr(),
1051            ))
1052        }
1053    }
1054
1055    /// Returns the lowered (mangled) name for a registered name expression.
1056    ///
1057    /// The identical name expression must have been previously provided to [`Program::add_name_expression`].
1058    ///
1059    /// # Errors
1060    ///
1061    /// Returns an error if `name_expression` contains an interior NUL byte, if
1062    /// the program handle is invalid, or if NVRTC cannot return the lowered name.
1063    pub fn lowered_name(&self, name_expression: &str) -> Result<String> {
1064        let name_expression = CString::new(name_expression)?;
1065        let mut lowered_name = ptr::null();
1066        unsafe {
1067            try_nvrtc!(sys::nvrtcGetLoweredName(
1068                self.handle()?,
1069                name_expression.as_ptr(),
1070                &raw mut lowered_name,
1071            ))?;
1072            Ok(CStr::from_ptr(lowered_name).to_string_lossy().into_owned())
1073        }
1074    }
1075
1076    pub fn ptx(&self) -> Result<Vec<u8>> {
1077        self.bytes(sys::nvrtcGetPTXSize, sys::nvrtcGetPTX)
1078    }
1079
1080    pub fn ptx_image(&self) -> Result<ModuleImage<'static>> {
1081        Ok(ModuleImage::from_vec(self.ptx()?))
1082    }
1083
1084    pub fn ptx_string(&self) -> Result<String> {
1085        Ok(bytes_to_string(self.ptx()?))
1086    }
1087
1088    pub fn cubin(&self) -> Result<Vec<u8>> {
1089        self.bytes(sys::nvrtcGetCUBINSize, sys::nvrtcGetCUBIN)
1090    }
1091
1092    pub fn cubin_image(&self) -> Result<ModuleImage<'static>> {
1093        Ok(ModuleImage::from_vec(self.cubin()?))
1094    }
1095
1096    pub fn lto_ir(&self) -> Result<Vec<u8>> {
1097        self.bytes(sys::nvrtcGetLTOIRSize, sys::nvrtcGetLTOIR)
1098    }
1099
1100    pub fn lto_ir_image(&self) -> Result<ModuleImage<'static>> {
1101        Ok(ModuleImage::from_vec(self.lto_ir()?))
1102    }
1103
1104    pub fn optix_ir(&self) -> Result<Vec<u8>> {
1105        self.bytes(sys::nvrtcGetOptiXIRSize, sys::nvrtcGetOptiXIR)
1106    }
1107
1108    pub fn optix_ir_image(&self) -> Result<ModuleImage<'static>> {
1109        Ok(ModuleImage::from_vec(self.optix_ir()?))
1110    }
1111
1112    pub fn log(&self) -> Result<String> {
1113        Ok(bytes_to_string(self.bytes(
1114            sys::nvrtcGetProgramLogSize,
1115            sys::nvrtcGetProgramLog,
1116        )?))
1117    }
1118
1119    /// Returns the PCH creation status.
1120    ///
1121    /// [`Status::Success`] indicates that the PCH was successfully created.
1122    /// [`Status::NoPchCreateAttempted`] indicates that no PCH creation was attempted, either because PCH was not requested during the preceding compile call, or automatic PCH processing was requested, and the compiler chose not to create a PCH file.
1123    /// [`Status::PchCreateHeapExhausted`] indicates that a PCH file could potentially have been created, but the compiler ran out space in the PCH heap.
1124    /// In this scenario, use [`Program::pch_heap_size_required`] to query the required heap size, reallocate the heap with [`set_pch_heap_size`], and retry PCH creation with [`sys::nvrtcCompileProgram`](singe_cuda_sys::nvrtc::nvrtcCompileProgram) on a new NVRTC program instance.
1125    /// [`Status::PchCreate`] indicates that an error condition prevented the PCH file from being created.
1126    ///
1127    /// # Errors
1128    ///
1129    /// Returns an error if the program handle is invalid.
1130    pub fn pch_create_status(&self) -> Result<Status> {
1131        unsafe { Ok(sys::nvrtcGetPCHCreateStatus(self.handle()?).into()) }
1132    }
1133
1134    /// Returns the PCH heap size required to compile the given program.
1135    ///
1136    /// # Errors
1137    ///
1138    /// Returns an error if the program handle is invalid or if the returned size
1139    /// is not valid because [`Program::pch_create_status`] did not return
1140    /// [`Status::Success`] or [`Status::PchCreateHeapExhausted`].
1141    pub fn pch_heap_size_required(&self) -> Result<usize> {
1142        let mut size = 0;
1143        unsafe {
1144            try_nvrtc!(sys::nvrtcGetPCHHeapSizeRequired(
1145                self.handle()?,
1146                &raw mut size
1147            ))?;
1148        }
1149        Ok(size as usize)
1150    }
1151
1152    /// Returns the raw NVRTC program handle, creating it first if needed.
1153    ///
1154    /// If the program has not been materialized yet, this creates the raw
1155    /// handle first.
1156    ///
1157    /// # Errors
1158    ///
1159    /// Returns an error if NVRTC cannot create the program handle.
1160    pub fn materialize_raw(&self) -> Result<sys::nvrtcProgram> {
1161        self.handle()
1162    }
1163
1164    /// Returns the raw NVRTC program handle only if it has already been
1165    /// materialized.
1166    pub const fn as_raw_if_materialized(&self) -> Option<sys::nvrtcProgram> {
1167        let handle = unsafe { *self.handle.get() };
1168        if handle.is_null() { None } else { Some(handle) }
1169    }
1170
1171    /// Transfers ownership of the raw NVRTC program handle to the caller.
1172    ///
1173    /// If the program has not been materialized yet, this creates the raw
1174    /// handle first. The caller becomes responsible for destroying the returned
1175    /// handle with `nvrtcDestroyProgram`.
1176    pub fn into_raw(self) -> Result<sys::nvrtcProgram> {
1177        let handle = self.handle()?;
1178        std::mem::forget(self);
1179        Ok(handle)
1180    }
1181
1182    fn compile_raw(&self, options: &[&str]) -> Result<()> {
1183        let options = options
1184            .iter()
1185            .map(|option| CString::new(*option))
1186            .collect::<result::Result<Vec<_>, _>>()?;
1187        let option_ptrs = options
1188            .iter()
1189            .map(|value| value.as_ptr())
1190            .collect::<Vec<_>>();
1191
1192        unsafe {
1193            try_nvrtc!(sys::nvrtcCompileProgram(
1194                self.handle()?,
1195                option_ptrs.len() as _,
1196                if option_ptrs.is_empty() {
1197                    ptr::null()
1198                } else {
1199                    option_ptrs.as_ptr()
1200                },
1201            ))
1202        }
1203    }
1204
1205    fn bytes(
1206        &self,
1207        get_size: unsafe extern "C" fn(sys::nvrtcProgram, *mut sys::size_t) -> sys::nvrtcResult,
1208        get_data: unsafe extern "C" fn(sys::nvrtcProgram, *mut i8) -> sys::nvrtcResult,
1209    ) -> Result<Vec<u8>> {
1210        let mut size = 0;
1211        unsafe {
1212            try_nvrtc!(get_size(self.handle()?, &raw mut size))?;
1213        }
1214
1215        let mut bytes = vec![0u8; size as usize];
1216        if bytes.is_empty() {
1217            return Ok(bytes);
1218        }
1219
1220        unsafe {
1221            try_nvrtc!(get_data(self.handle()?, bytes.as_mut_ptr().cast()))?;
1222        }
1223        Ok(bytes)
1224    }
1225
1226    fn handle(&self) -> Result<sys::nvrtcProgram> {
1227        unsafe {
1228            let handle = self.handle.get();
1229            if (*handle).is_null() {
1230                *handle = self.create_handle()?;
1231            }
1232            Ok(*handle)
1233        }
1234    }
1235
1236    fn create_handle(&self) -> Result<sys::nvrtcProgram> {
1237        let source = CString::new(self.source.as_str())?;
1238        let name = self.name.as_deref().map(CString::new).transpose()?;
1239        let header_sources = self
1240            .headers
1241            .iter()
1242            .map(|header| CString::new(header.source.as_str()))
1243            .collect::<result::Result<Vec<_>, _>>()?;
1244        let include_names = self
1245            .headers
1246            .iter()
1247            .map(|header| CString::new(header.include_name.as_str()))
1248            .collect::<result::Result<Vec<_>, _>>()?;
1249        let header_ptrs = header_sources
1250            .iter()
1251            .map(|value| value.as_ptr())
1252            .collect::<Vec<_>>();
1253        let include_name_ptrs = include_names
1254            .iter()
1255            .map(|value| value.as_ptr())
1256            .collect::<Vec<_>>();
1257
1258        let mut handle = ptr::null_mut();
1259        unsafe {
1260            try_nvrtc!(sys::nvrtcCreateProgram(
1261                &raw mut handle,
1262                source.as_ptr(),
1263                name.as_ref().map_or(ptr::null(), |value| value.as_ptr()),
1264                self.headers.len() as _,
1265                if header_ptrs.is_empty() {
1266                    ptr::null()
1267                } else {
1268                    header_ptrs.as_ptr()
1269                },
1270                if include_name_ptrs.is_empty() {
1271                    ptr::null()
1272                } else {
1273                    include_name_ptrs.as_ptr()
1274                },
1275            ))?;
1276        }
1277
1278        Ok(handle)
1279    }
1280}
1281
1282impl Drop for Program {
1283    fn drop(&mut self) {
1284        unsafe {
1285            let handle = self.handle.get();
1286            if !(*handle).is_null() {
1287                let _ = sys::nvrtcDestroyProgram(handle);
1288            }
1289        }
1290    }
1291}
1292
1293#[non_exhaustive]
1294pub enum CompilationArtifact {
1295    Ptx(ModuleImage<'static>),
1296    Cubin(ModuleImage<'static>),
1297    LtoIr(ModuleImage<'static>),
1298    OptixIr(ModuleImage<'static>),
1299}
1300
1301impl CompilationArtifact {
1302    pub fn image(&self) -> &ModuleImage<'static> {
1303        match self {
1304            Self::Ptx(image) | Self::Cubin(image) | Self::LtoIr(image) | Self::OptixIr(image) => {
1305                image
1306            }
1307        }
1308    }
1309
1310    pub fn into_image(self) -> ModuleImage<'static> {
1311        match self {
1312            Self::Ptx(image) | Self::Cubin(image) | Self::LtoIr(image) | Self::OptixIr(image) => {
1313                image
1314            }
1315        }
1316    }
1317}
1318
1319#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
1320#[non_exhaustive]
1321pub enum OutputKind {
1322    Ptx,
1323    Cubin,
1324    LtoIr,
1325    OptixIr,
1326}
1327
1328impl Program {
1329    pub fn artifact(&self, kind: OutputKind) -> Result<CompilationArtifact> {
1330        match kind {
1331            OutputKind::Ptx => {
1332                let image = self.ptx_image()?;
1333                if image.as_bytes().is_empty() {
1334                    return Err(Error::InvalidValue);
1335                }
1336                Ok(CompilationArtifact::Ptx(image))
1337            }
1338            OutputKind::Cubin => {
1339                let image = self.cubin_image()?;
1340                if image.as_bytes().is_empty() {
1341                    return Err(Error::InvalidValue);
1342                }
1343                Ok(CompilationArtifact::Cubin(image))
1344            }
1345            OutputKind::LtoIr => {
1346                let image = self.lto_ir_image()?;
1347                if image.as_bytes().is_empty() {
1348                    return Err(Error::InvalidValue);
1349                }
1350                Ok(CompilationArtifact::LtoIr(image))
1351            }
1352            OutputKind::OptixIr => {
1353                let image = self.optix_ir_image()?;
1354                if image.as_bytes().is_empty() {
1355                    return Err(Error::InvalidValue);
1356                }
1357                Ok(CompilationArtifact::OptixIr(image))
1358            }
1359        }
1360    }
1361}
1362
1363/// Returns the CUDA Runtime Compilation version.
1364///
1365/// # Errors
1366///
1367/// Returns an error if NVRTC cannot report its version.
1368pub fn version() -> Result<Version> {
1369    let mut major = 0;
1370    let mut minor = 0;
1371    unsafe {
1372        try_nvrtc!(sys::nvrtcVersion(&raw mut major, &raw mut minor))?;
1373    }
1374    Ok(Version { major, minor })
1375}
1376
1377pub fn supported_architectures() -> Result<Vec<i32>> {
1378    let mut count = 0;
1379    unsafe {
1380        try_nvrtc!(sys::nvrtcGetNumSupportedArchs(&raw mut count))?;
1381    }
1382
1383    let mut architectures = vec![0; count as usize];
1384    if architectures.is_empty() {
1385        return Ok(Vec::new());
1386    }
1387
1388    unsafe {
1389        try_nvrtc!(sys::nvrtcGetSupportedArchs(architectures.as_mut_ptr()))?;
1390    }
1391    Ok(architectures)
1392}
1393
1394/// Returns the current size of the PCH heap.
1395///
1396/// # Errors
1397///
1398/// Returns an error if NVRTC cannot report the PCH heap size.
1399pub fn pch_heap_size() -> Result<usize> {
1400    let mut size = 0;
1401    unsafe {
1402        try_nvrtc!(sys::nvrtcGetPCHHeapSize(&raw mut size))?;
1403    }
1404    Ok(size as usize)
1405}
1406
1407/// Sets the size of the PCH heap.
1408///
1409/// The requested size may be rounded up to a platform-dependent alignment, such as the page size.
1410/// If the PCH heap has already been allocated, NVRTC frees it and allocates a new heap.
1411///
1412/// # Errors
1413///
1414/// Returns an error if NVRTC rejects the requested PCH heap size.
1415pub fn set_pch_heap_size(size: usize) -> Result<()> {
1416    unsafe { try_nvrtc!(sys::nvrtcSetPCHHeapSize(size as _)) }
1417}
1418
1419unsafe extern "C" fn noop_compile_callback(
1420    payload: *mut std::ffi::c_void,
1421    reserved: *mut std::ffi::c_void,
1422) -> i32 {
1423    let _ = payload;
1424    let _ = reserved;
1425    0
1426}
1427
1428unsafe extern "C" fn cancel_if_requested_callback(
1429    payload: *mut std::ffi::c_void,
1430    reserved: *mut std::ffi::c_void,
1431) -> i32 {
1432    let _ = reserved;
1433    if payload.is_null() {
1434        return 0;
1435    }
1436
1437    let cancel = unsafe { &*payload.cast::<AtomicBool>() };
1438    i32::from(cancel.load(Ordering::Relaxed))
1439}
1440
1441/// Clears the NVRTC flow callback by installing a no-op callback.
1442///
1443/// Use this after a compile attempt that installed a cancellation callback.
1444///
1445/// # Errors
1446///
1447/// Returns an error if the program handle is invalid or if NVRTC rejects the
1448/// callback update.
1449pub fn clear_flow_callback(program: &Program) -> Result<()> {
1450    unsafe {
1451        try_nvrtc!(sys::nvrtcSetFlowCallback(
1452            program.handle()?,
1453            Some(noop_compile_callback),
1454            ptr::null_mut(),
1455        ))
1456    }
1457}
1458
1459fn bool_flag(value: bool) -> &'static str {
1460    if value { "true" } else { "false" }
1461}
1462
1463fn join_display(values: &[impl Display]) -> String {
1464    values
1465        .iter()
1466        .map(ToString::to_string)
1467        .collect::<Vec<_>>()
1468        .join(",")
1469}
1470
1471fn join_numbers(values: &[i32]) -> String {
1472    values
1473        .iter()
1474        .map(ToString::to_string)
1475        .collect::<Vec<_>>()
1476        .join(",")
1477}
1478
1479fn bytes_to_string(mut bytes: Vec<u8>) -> String {
1480    while bytes.last() == Some(&0) {
1481        bytes.pop();
1482    }
1483    String::from_utf8_lossy(&bytes).into_owned()
1484}
1485
1486#[cfg(all(test, feature = "testing"))]
1487mod tests {
1488    use super::*;
1489    use crate::{
1490        device::Device,
1491        error::{Error, Result},
1492        memory::DeviceMemory,
1493        module::LaunchConfig,
1494        testing,
1495    };
1496
1497    fn current_device_sm_architecture() -> Result<GpuArchitecture> {
1498        let properties = Device::current()?.properties()?;
1499        Ok(match (properties.major, properties.minor) {
1500            (7, 5) => GpuArchitecture::Sm75,
1501            (8, 0) => GpuArchitecture::Sm80,
1502            (8, 6) => GpuArchitecture::Sm86,
1503            (8, 7) => GpuArchitecture::Sm87,
1504            (8, 9) => GpuArchitecture::Sm89,
1505            (9, 0) => GpuArchitecture::Sm90,
1506            (10, 0) => GpuArchitecture::Sm100,
1507            (10, 1) => GpuArchitecture::Sm101,
1508            (10, 3) => GpuArchitecture::Sm103,
1509            (12, 0) => GpuArchitecture::Sm120,
1510            (12, 1) => GpuArchitecture::Sm121,
1511            (major, minor) => panic!("unsupported device architecture sm_{major}{minor}"),
1512        })
1513    }
1514
1515    fn compile_or_skip(result: Result<()>) -> bool {
1516        match result {
1517            Ok(()) => true,
1518            Err(Error::Nvrtc {
1519                code: Status::BuiltinOperationFailure,
1520                ..
1521            }) => {
1522                eprintln!(
1523                    "warning: skipping NVRTC compile-dependent test because NVRTC builtins are unavailable"
1524                );
1525                false
1526            }
1527            Err(error) => panic!("NVRTC compilation failed: {error:?}"),
1528        }
1529    }
1530
1531    #[test]
1532    fn version_is_available() {
1533        let version = version().unwrap();
1534        assert_ne!(version.major, 0);
1535    }
1536
1537    #[test]
1538    fn supported_architectures_are_sorted() {
1539        let architectures = supported_architectures().unwrap();
1540        assert!(!architectures.is_empty());
1541        assert!(
1542            architectures
1543                .windows(2)
1544                .all(|window| window[0] <= window[1])
1545        );
1546    }
1547
1548    #[test]
1549    fn compile_options_build_expected_arguments() {
1550        let arguments = CompileOptions::new()
1551            .gpu_architecture(GpuArchitecture::Compute80)
1552            .device_debug(true)
1553            .generate_line_info(true)
1554            .define_macro(MacroDefinition::WithValue {
1555                name: "FOO",
1556                value: "42",
1557            })
1558            .include_path("include")
1559            .cpp_dialect(CppDialect::Cpp20)
1560            .warning_as_error(WarningAsErrorKind::Reorder)
1561            .diag_suppress(177)
1562            .raw_option("--custom-flag")
1563            .as_arguments();
1564
1565        assert_eq!(
1566            arguments,
1567            vec![
1568                "--gpu-architecture=compute_80",
1569                "--device-debug",
1570                "--generate-line-info",
1571                "--define-macro=FOO=42",
1572                "--include-path=include",
1573                "--std=c++20",
1574                "--warning-as-error=reorder",
1575                "--diag-suppress=177",
1576                "--custom-flag",
1577            ]
1578        );
1579    }
1580
1581    #[test]
1582    fn compiles_to_ptx() {
1583        let program = Program::from_source(
1584            r#"
1585            extern "C" __global__ void saxpy(float a, const float* x, const float* y, float* out, size_t n) {
1586                size_t tid = blockIdx.x * blockDim.x + threadIdx.x;
1587                if (tid < n) {
1588                    out[tid] = a * x[tid] + y[tid];
1589                }
1590            }
1591            "#,
1592        )
1593        .with_name("saxpy.cu");
1594        let options = CompileOptions::new()
1595            .gpu_architecture(GpuArchitecture::Compute80)
1596            .generate_line_info(true);
1597        if !compile_or_skip(program.compile_with_options(&options)) {
1598            return;
1599        }
1600
1601        let ptx = program.ptx_string().unwrap();
1602        assert!(ptx.contains(".visible .entry saxpy"));
1603    }
1604
1605    #[test]
1606    fn compile_with_cancel_flag_succeeds_when_not_cancelled() {
1607        let (_lock, _ctx) = testing::bootstrap().unwrap();
1608        let cancel = AtomicBool::new(false);
1609        let program = Program::from_source(
1610            r#"
1611            extern "C" __global__ void noop() {}
1612            "#,
1613        )
1614        .with_name("noop.cu");
1615        let options = CompileOptions::new().gpu_architecture(GpuArchitecture::Compute80);
1616
1617        if !compile_or_skip(program.compile_with_options_and_cancel_flag(&options, &cancel)) {
1618            return;
1619        }
1620        assert!(program.ptx_string().unwrap().contains("noop"));
1621    }
1622
1623    #[test]
1624    fn clear_flow_callback_is_allowed_before_compilation() {
1625        let program =
1626            Program::from_source("extern \"C\" __global__ void noop() {}").with_name("noop.cu");
1627        clear_flow_callback(&program).unwrap();
1628    }
1629
1630    #[test]
1631    fn cubin_artifact_loads_as_module() {
1632        let (_lock, ctx) = testing::bootstrap().unwrap();
1633        let program = Program::from_source(
1634            r#"
1635            extern "C" __global__ void saxpy(float a, const float* x, const float* y, float* out, size_t n) {
1636                size_t tid = blockIdx.x * blockDim.x + threadIdx.x;
1637                if (tid < n) {
1638                    out[tid] = a * x[tid] + y[tid];
1639                }
1640            }
1641            "#,
1642        )
1643        .with_name("saxpy_module.cu");
1644        let architecture = current_device_sm_architecture().unwrap();
1645        let options = CompileOptions::new().gpu_architecture(architecture);
1646        if !compile_or_skip(program.compile_with_options(&options)) {
1647            return;
1648        }
1649
1650        let module = ctx.load_nvrtc_module(&program, OutputKind::Cubin).unwrap();
1651        assert!(module.function("saxpy").is_ok());
1652    }
1653
1654    #[test]
1655    fn cubin_artifact_loads_as_module_with_jit_options() {
1656        let (_lock, ctx) = testing::bootstrap().unwrap();
1657        let program = Program::from_source(
1658            r#"
1659            extern "C" __global__ void noop() {}
1660            "#,
1661        )
1662        .with_name("noop_module_jit.cu");
1663        let architecture = current_device_sm_architecture().unwrap();
1664        let options = CompileOptions::new().gpu_architecture(architecture);
1665        if !compile_or_skip(program.compile_with_options(&options)) {
1666            return;
1667        }
1668
1669        let mut info_log = [0u8; 1024];
1670        let mut error_log = [0u8; 1024];
1671        let jit_options = crate::jit::JitOptions::default()
1672            .with_generate_line_info(true)
1673            .with_info_log(&mut info_log)
1674            .with_error_log(&mut error_log);
1675
1676        let module = ctx
1677            .load_nvrtc_module_with_options(&program, OutputKind::Cubin, jit_options)
1678            .unwrap();
1679        assert!(module.function("noop").is_ok());
1680    }
1681
1682    #[test]
1683    fn compiles_loads_launches_and_reads_back_results() {
1684        let (_lock, ctx) = testing::bootstrap().unwrap();
1685
1686        let input = vec![1.0f32, 2.0, 3.5, -4.0, 8.25];
1687        let mut output = vec![0.0f32; input.len()];
1688        let scalar = 2.5f32;
1689        let input_device = DeviceMemory::from_slice(&input).unwrap();
1690        let output_device = DeviceMemory::<f32>::zeroes(output.len()).unwrap();
1691        let length = input.len();
1692
1693        let program = Program::from_source(
1694            r#"
1695            extern "C" __global__ void scale_add(const float* input, float* output, float alpha, size_t len) {
1696                size_t i = blockIdx.x * blockDim.x + threadIdx.x;
1697                if (i < len) {
1698                    output[i] = input[i] * alpha + 1.0f;
1699                }
1700            }
1701            "#,
1702        )
1703        .with_name("scale_add.cu");
1704        let architecture = current_device_sm_architecture().unwrap();
1705        let compile_options = CompileOptions::new().gpu_architecture(architecture);
1706        if !compile_or_skip(program.compile_with_options(&compile_options)) {
1707            return;
1708        }
1709
1710        let module = ctx.load_nvrtc_module(&program, OutputKind::Cubin).unwrap();
1711        let function = module.function("scale_add").unwrap();
1712
1713        let config = LaunchConfig::for_1d_grid(input.len(), 128);
1714        let input_ptr = input_device.as_ptr();
1715        let mut output_ptr = output_device.as_mut_ptr();
1716
1717        function
1718            .launch(&config, (&input_ptr, &mut output_ptr, &scalar, &length))
1719            .unwrap();
1720        output_device.copy_to_host(&mut output).unwrap();
1721
1722        let expected = input
1723            .iter()
1724            .map(|value| value * scalar + 1.0)
1725            .collect::<Vec<_>>();
1726        assert_eq!(output, expected);
1727    }
1728
1729    #[test]
1730    fn cubin_artifact_loads_as_library() {
1731        let (_lock, ctx) = testing::bootstrap().unwrap();
1732        let program = Program::from_source(
1733            r#"
1734            extern "C" __global__ void noop() {}
1735            "#,
1736        )
1737        .with_name("noop_library.cu");
1738        let architecture = current_device_sm_architecture().unwrap();
1739        let options = CompileOptions::new().gpu_architecture(architecture);
1740        if !compile_or_skip(program.compile_with_options(&options)) {
1741            return;
1742        }
1743
1744        let library = ctx.load_nvrtc_library(&program, OutputKind::Cubin).unwrap();
1745        assert!(library.kernel_count().unwrap() >= 1);
1746    }
1747
1748    #[test]
1749    fn lto_ir_artifact_is_available_when_requested() {
1750        let program = Program::from_source(
1751            r#"
1752            extern "C" __global__ void noop() {}
1753            "#,
1754        )
1755        .with_name("noop_lto.cu");
1756        let options = CompileOptions::new().dlink_time_optimization(true);
1757        if !compile_or_skip(program.compile_with_options(&options)) {
1758            return;
1759        }
1760
1761        let artifact = program.artifact(OutputKind::LtoIr).unwrap();
1762        assert!(!artifact.image().as_bytes().is_empty());
1763        let ptx = program.artifact(OutputKind::Ptx).unwrap();
1764        assert!(!ptx.image().as_bytes().is_empty());
1765    }
1766
1767    #[test]
1768    fn optix_ir_artifact_is_available_when_requested() {
1769        let program = Program::from_source(
1770            r#"
1771            extern "C" __global__ void noop() {}
1772            "#,
1773        )
1774        .with_name("noop_optix.cu");
1775        let options = CompileOptions::new().optix_ir(true);
1776        if !compile_or_skip(program.compile_with_options(&options)) {
1777            return;
1778        }
1779
1780        let artifact = program.artifact(OutputKind::OptixIr).unwrap();
1781        assert!(!artifact.image().as_bytes().is_empty());
1782        let ptx = program.artifact(OutputKind::Ptx).unwrap();
1783        assert!(!ptx.image().as_bytes().is_empty());
1784    }
1785
1786    #[test]
1787    fn cubin_artifact_requires_real_architecture() {
1788        let program = Program::from_source(
1789            r#"
1790            extern "C" __global__ void noop() {}
1791            "#,
1792        )
1793        .with_name("noop_cubin.cu");
1794        let options = CompileOptions::new().gpu_architecture(GpuArchitecture::Compute80);
1795        if !compile_or_skip(program.compile_with_options(&options)) {
1796            return;
1797        }
1798
1799        assert!(program.artifact(OutputKind::Cubin).is_err());
1800    }
1801}
singe_cuda/nvrtc.rs

singe_cuda/
nvrtc.rs
