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

cubecl_macros/
lib.rs
