cubecl_macros/
lib.rs

1#![allow(clippy::large_enum_variant)]
2
3use core::panic;
4
5use error::error_into_token_stream;
6use generate::autotune::generate_autotune_key;
7use parse::{
8    cube_impl::CubeImpl,
9    cube_trait::{CubeTrait, CubeTraitImpl},
10    helpers::RemoveHelpers,
11    kernel::{Launch, from_tokens},
12};
13use proc_macro::TokenStream;
14use quote::quote;
15use syn::{Item, visit_mut::VisitMut};
16
17use crate::{
18    generate::{
19        asm::generate_asm_unexpanded, assign::generate_cube_type_mut,
20        into_runtime::generate_into_runtime,
21    },
22    parse::{
23        cube_type::generate_cube_type, derive_expand::generate_derive_expand,
24        helpers::ReplaceDefines,
25    },
26};
27
28mod error;
29mod expression;
30mod generate;
31mod operator;
32mod parse;
33mod paths;
34mod scope;
35mod statement;
36
37/// Mark a cube function, trait or implementation for expansion.
38///
39/// # Arguments
40/// * `launch` - generates a function to launch the kernel
41/// * `launch_unchecked` - generates a launch function without checks
42/// * `debug` - panics after generation to print the output to console
43/// * `create_dummy_kernel` - Generates a function to create a kernel without launching it. Used for
44///   testing.
45///
46/// # Trait arguments
47/// * `expand_base_traits` - base traits for the expanded "second half" of a trait with methods.
48///
49/// # Example
50///
51/// ```ignored
52/// # use cubecl_macros::cube;
53/// #[cube]
54/// fn my_addition(a: u32, b: u32) -> u32 {
55///     a + b
56/// }
57/// ```
58#[proc_macro_attribute]
59pub fn cube(args: TokenStream, input: TokenStream) -> TokenStream {
60    match cube_impl(args, input.clone()) {
61        Ok(tokens) => tokens,
62        Err(e) => error_into_token_stream(e, input.into()).into(),
63    }
64}
65
66fn cube_impl(args: TokenStream, input: TokenStream) -> syn::Result<TokenStream> {
67    let mut item: Item = syn::parse(input)?;
68    let args = from_tokens(args.into())?;
69
70    let tokens = match item.clone() {
71        Item::Fn(kernel) => {
72            let kernel = Launch::from_item_fn(kernel, args)?;
73            RemoveHelpers.visit_item_mut(&mut item);
74            ReplaceDefines.visit_item_mut(&mut item);
75
76            let extra_allow = match kernel.func.context.is_intrinsic {
77                true => quote![#[allow(unused_variables)]],
78                false => quote![],
79            };
80
81            return Ok(TokenStream::from(quote! {
82                #[allow(dead_code, clippy::too_many_arguments)]
83                #extra_allow
84                #item
85                #kernel
86            }));
87        }
88        Item::Trait(kernel_trait) => {
89            let is_debug = args.debug.is_present();
90            let expand_trait = CubeTrait::from_item_trait(kernel_trait, args)?;
91
92            let tokens = TokenStream::from(quote! {
93                #expand_trait
94            });
95            if is_debug {
96                panic!("{tokens}");
97            }
98            return Ok(tokens);
99        }
100        Item::Impl(item_impl) => {
101            if item_impl.trait_.is_some() {
102                let mut expand_impl = CubeTraitImpl::from_item_impl(item_impl, &args)?;
103                let expand_impl = expand_impl.to_tokens_mut();
104
105                Ok(TokenStream::from(quote! {
106                    #expand_impl
107                }))
108            } else {
109                let mut expand_impl = CubeImpl::from_item_impl(item_impl, &args)?;
110                let expand_impl = expand_impl.to_tokens_mut();
111
112                Ok(TokenStream::from(quote! {
113                    #expand_impl
114                }))
115            }
116        }
117        item => Err(syn::Error::new_spanned(
118            item,
119            "`#[cube]` is only supported on traits and functions",
120        ))?,
121    };
122
123    if args.debug.is_present() {
124        match tokens {
125            Ok(tokens) => panic!("{tokens}"),
126            Err(err) => panic!("{err}"),
127        };
128    }
129
130    tokens
131}
132
133/// Derive macro to define a cube type that is launched with a kernel
134#[proc_macro_derive(CubeLaunch, attributes(cube, launch))]
135pub fn module_derive_cube_launch(input: TokenStream) -> TokenStream {
136    gen_cube_type(input, true)
137}
138
139/// Derive macro to define a cube type that is not launched
140#[proc_macro_derive(CubeType, attributes(cube, expand))]
141pub fn module_derive_cube_type(input: TokenStream) -> TokenStream {
142    gen_cube_type(input, false)
143}
144
145fn gen_cube_type(input: TokenStream, with_launch: bool) -> TokenStream {
146    let parsed = syn::parse(input);
147
148    let input = match &parsed {
149        Ok(val) => val,
150        Err(err) => return err.to_compile_error().into(),
151    };
152
153    match generate_cube_type(input, with_launch) {
154        Ok(val) => val.into(),
155        Err(err) => err.to_compile_error().into(),
156    }
157}
158
159/// Attribute macro to define a type that can be used as a kernel comptime
160/// argument This derive Debug, Hash, `PartialEq`, Eq, Clone, Copy
161#[proc_macro_attribute]
162pub fn derive_cube_comptime(_metadata: TokenStream, input: TokenStream) -> TokenStream {
163    let input: proc_macro2::TokenStream = input.into();
164    quote! {
165        #[derive(Debug, Hash, PartialEq, Eq, Clone, Copy)]
166        #input
167    }
168    .into()
169}
170
171/// Attribute macro to derive cube traits for existing structs, without redefining that struct.
172#[proc_macro_attribute]
173pub fn derive_expand(metadata: TokenStream, input: TokenStream) -> TokenStream {
174    match generate_derive_expand(input.into(), metadata.into()) {
175        Ok(val) => val.into(),
176        Err(err) => err.to_compile_error().into(),
177    }
178}
179
180/// Mark the contents of this macro as compile time values, turning off all
181/// expansion for this code and using it verbatim
182///
183/// # Example
184/// ```ignored
185/// #use cubecl_macros::cube;
186/// #fn some_rust_function(a: u32) -> u32 {}
187/// #[cube]
188/// fn do_stuff(input: u32) -> u32 {
189///     let comptime_value = comptime! { some_rust_function(3) };
190///     input + comptime_value
191/// }
192/// ```
193#[proc_macro]
194pub fn comptime(input: TokenStream) -> TokenStream {
195    let tokens: proc_macro2::TokenStream = input.into();
196    quote![{ #tokens }].into()
197}
198
199/// Mark the contents of this macro as an intrinsic, turning off all expansion
200/// for this code and calling it with the scope
201///
202/// # Example
203/// ```ignored
204/// #use cubecl_macros::cube;
205/// #[cube]
206/// fn do_stuff(input: u32) -> u32 {
207///     let comptime_value = intrinsic! { |scope| u32::elem_size(scope) };
208///     input + comptime_value
209/// }
210/// ```
211#[proc_macro]
212pub fn intrinsic(_input: TokenStream) -> TokenStream {
213    quote![{ cubecl::unexpanded!() }].into()
214}
215
216/// GPU version of [`asm`](std::arch::asm). Currently parses all the same options, but most are not
217/// applicable to GPU assembly architectures. Should validate and give proper errors at some point.
218/// Also adds a new register spec: the inferred register specifier (`_`). This is because the
219/// specifier isn't actually meaningful in PTX and is currently ignored, with constraints being
220/// inferred from the value type. The reason it's still present is because we may want to add an
221/// explicit `mem` specifier to allow fine-grained memory clobbering, or support for other assembly
222/// formats that do use different register types.
223///
224/// # Example
225/// ```ignored
226/// #use cubecl_macros::cube;
227/// #[cube]
228/// fn do_stuff(input: u32) -> u32 {
229///     let mut out: u32;
230///     gpu_asm!("some.custom.ptx {} {}", out(_) out, in(_) input)
231/// }
232/// ```
233#[proc_macro]
234pub fn gpu_asm(input: TokenStream) -> TokenStream {
235    match generate_asm_unexpanded(input.into()) {
236        Ok(val) => val.into(),
237        Err(err) => err.to_compile_error().into(),
238    }
239}
240
241/// Makes the function return a compile time value
242/// Useful in a cube trait to have a part of the trait return comptime values
243///
244/// # Example
245/// ```ignored
246/// #use cubecl_macros::cube;
247/// #[cube]
248/// fn do_stuff(#[comptime] input: u32) -> comptime_type!(u32) {
249///     input + 5   
250/// }
251/// ```
252///
253/// TODO: calling a trait method returning `comptime_type` from
254/// within another trait method does not work
255#[proc_macro]
256pub fn comptime_type(input: TokenStream) -> TokenStream {
257    let tokens: proc_macro2::TokenStream = input.into();
258    quote![ #tokens ].into()
259}
260
261/// Insert a literal comment into the kernel source code.
262///
263/// # Example
264/// ```ignored
265/// #use cubecl_macros::cube;
266/// #[cube]
267/// fn do_stuff(input: u32) -> u32 {
268///     comment!("Add five to the input");
269///     input + 5
270/// }
271/// ```
272#[proc_macro]
273pub fn comment(input: TokenStream) -> TokenStream {
274    let tokens: proc_macro2::TokenStream = input.into();
275    quote![{ #tokens }].into()
276}
277
278/// Terminate the execution of the kernel for the current unit.
279///
280/// This terminates the execution of the unit even if nested inside many
281/// functions.
282///
283/// # Example
284/// ```ignored
285/// #use cubecl_macros::cube;
286/// #[cube]
287/// fn stop_if_more_than_ten(input: u32)  {
288///     if input > 10 {
289///         terminate!();
290///     }
291/// }
292/// ```
293#[proc_macro]
294pub fn terminate(input: TokenStream) -> TokenStream {
295    let tokens: proc_macro2::TokenStream = input.into();
296    quote![{ #tokens }].into()
297}
298
299/// Implements display and initialization for autotune keys.
300///
301/// # Helper
302///
303/// Use the `#[autotune(anchor)]` helper attribute to anchor a numerical value.
304/// This groups multiple numerical values into the same bucket.
305///
306/// For now, only an exponential function is supported, and it can be modified
307/// with `exp`. By default, the base is '2' and there are no `min` or `max`
308/// provided.
309///
310/// # Example
311/// ```ignore
312/// #[derive(AutotuneKey, Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
313/// pub struct OperationKey {
314///     #[autotune(name = "Batch Size")]
315///     batch_size: usize,
316///     channels: usize,
317///     #[autotune(anchor(exp(min = 16, max = 1024, base = 2)))]
318///     height: usize,
319///     #[autotune(anchor)]
320///     width: usize,
321/// }
322/// ```
323#[proc_macro_derive(AutotuneKey, attributes(autotune))]
324pub fn derive_autotune_key(input: TokenStream) -> TokenStream {
325    let input = syn::parse(input).unwrap();
326    match generate_autotune_key(input) {
327        Ok(tokens) => tokens.into(),
328        Err(e) => e.into_compile_error().into(),
329    }
330}
331
332/// Implements `IntoRuntime` for a `CubeType`
333#[proc_macro_derive(IntoRuntime, attributes(cube))]
334pub fn derive_into_runtime(input: TokenStream) -> TokenStream {
335    let input = syn::parse(input).unwrap();
336    match generate_into_runtime(&input) {
337        Ok(tokens) => tokens.into(),
338        Err(e) => e.into_compile_error().into(),
339    }
340}
341
342/// Implements mutability for a `CubeType`
343#[proc_macro_derive(CubeTypeMut, attributes(cube))]
344pub fn derive_assign(input: TokenStream) -> TokenStream {
345    let input = syn::parse(input).unwrap();
346    match generate_cube_type_mut(&input) {
347        Ok(tokens) => tokens.into(),
348        Err(e) => e.into_compile_error().into(),
349    }
350}
cubecl_macros/lib.rs

cubecl_macros/
lib.rs
