vortex_cuda_macros/
lib.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4//! Proc macros for CUDA conditional compilation.
5//!
6//! These macros simplify gating code based on CUDA availability.
7//!
8//! # Usage
9//!
10//! ```ignore
11//! use vortex_cuda_macros::{cuda_available, cuda_not_available, cuda_test};
12//!
13//! // Only compiled when CUDA is available
14//! #[cuda_available]
15//! fn cuda_only_function() { /* ... */ }
16//!
17//! // Only compiled when CUDA is NOT available
18//! #[cuda_not_available]
19//! fn fallback_function() { /* ... */ }
20//!
21//! // Ignore tests when CUDA is not available
22//! #[crate::test]
23//! async fn my_test() {
24//! ...
25//! }
26//! ```
27
28use std::process::Command;
29use std::sync::LazyLock;
30
31use proc_macro::TokenStream;
32use quote::quote;
33use syn::parse_macro_input;
34
35/// Cached result of nvcc availability check.
36static NVCC_AVAILABLE: LazyLock<bool> = LazyLock::new(|| {
37    Command::new("nvcc")
38        .arg("--version")
39        .output()
40        .is_ok_and(|o| o.status.success())
41});
42
43/// Conditionally compiles the annotated item only when CUDA is available.
44#[proc_macro_attribute]
45pub fn cuda_available(_attr: TokenStream, item: TokenStream) -> TokenStream {
46    if *NVCC_AVAILABLE {
47        item
48    } else {
49        TokenStream::new()
50    }
51}
52
53/// Conditionally compiles the annotated item only when CUDA is not available.
54#[proc_macro_attribute]
55pub fn cuda_not_available(_attr: TokenStream, item: TokenStream) -> TokenStream {
56    if *NVCC_AVAILABLE {
57        TokenStream::new()
58    } else {
59        item
60    }
61}
62
63/// Test attribute to ignore tests if CUDA isn't available. Supports both sync and async tests (using tokio).
64///
65/// Must be named `test` to work with frameworks like `rstest`.
66#[proc_macro_attribute]
67pub fn test(_attr: TokenStream, item: TokenStream) -> TokenStream {
68    let item = parse_macro_input!(item as syn::ItemFn);
69    if *NVCC_AVAILABLE {
70        if item.sig.asyncness.is_some() {
71            quote! {
72                #[tokio::test]
73                #item
74            }
75        } else {
76            quote! {
77                #[test]
78                #item
79            }
80        }
81        .into()
82    } else {
83        if item.sig.asyncness.is_some() {
84            quote! {
85                #[tokio::test]
86                #[ignore]
87                #item
88            }
89        } else {
90            quote! {
91                #[test]
92                #[ignore]
93                #item
94            }
95        }
96        .into()
97    }
98}
vortex_cuda_macros/lib.rs

vortex_cuda_macros/
lib.rs
