cubecl-macros 0.12.0-pre.1

#![allow(clippy::large_enum_variant)]

use core::panic;

use error::error_into_token_stream;
use generate::autotune::generate_autotune_key;
use parse::{
    cube_impl::CubeImpl,
    cube_trait::{CubeTrait, CubeTraitImpl},
    helpers::RemoveHelpers,
    kernel::{KernelArgs, Launch, from_tokens},
};
use proc_macro::TokenStream;
use quote::quote;
use syn::{Item, visit_mut::VisitMut};

use crate::{
    generate::{
        asm::generate_asm_unexpanded, assign::generate_cube_type_mut,
        into_runtime::generate_into_runtime,
    },
    parse::{
        cube_type::{generate_cube_harness, generate_cube_type},
        derive_expand::generate_derive_expand,
        helpers::ReplaceDefines,
    },
};

mod error;
mod expression;
mod generate;
mod operator;
mod parse;
mod paths;
mod scope;
mod statement;

/// Mark a cube function, trait or implementation for expansion.
///
/// # Arguments
/// * `launch` - generates a function to launch the kernel
/// * `launch_unchecked` - generates a launch function without checks
/// * `debug` - panics after generation to print the output to console
/// * `create_dummy_kernel` - Generates a function to create a kernel without launching it. Used for
///   testing.
/// * `harness` - generates a `harness` function that runs the function on a single unit of a
///   runtime and returns its result to the host, so a test can call `f::harness::<R>(..)` instead
///   of writing a kernel, buffers and a read back by hand. Every parameter must be a scalar, a
///   slice or array of scalars, a comptime scalar or a cube type that derives `CubeHarness`, and
///   the function must return them or a tuple of them, or nothing when it writes into a mutable
///   slice. The harness is gated:
///   * `harness` - same as `harness(cfg(test))`
///   * `harness(cfg(<predicate>))` - compiled under `#[cfg(<predicate>)]`. The predicate is
///     evaluated in the crate using the macro, so `feature = "..."` names one of its features.
///   * `harness(always)` - ungated
///
/// # Trait arguments
/// * `expand_base_traits` - base traits for the expanded "second half" of a trait with methods.
///
/// # Example
///
/// ```ignored
/// # use cubecl_macros::cube;
/// #[cube]
/// fn my_addition(a: u32, b: u32) -> u32 {
///     a + b
/// }
/// ```
///
/// # Harness example
///
/// ```ignored
/// #[cube(harness)]
/// fn add<F: Float>(x: F, y: F) -> F {
///     x + y
/// }
///
/// #[test]
/// fn adds() {
///     assert_eq!(add::harness::<CpuRuntime, f32>(1.0, 2.0), 3.0);
/// }
/// ```
#[proc_macro_attribute]
pub fn cube(args: TokenStream, input: TokenStream) -> TokenStream {
    match cube_impl(args, input.clone()) {
        Ok(tokens) => tokens,
        Err(e) => error_into_token_stream(e, input.into()).into(),
    }
}

fn cube_impl(args: TokenStream, input: TokenStream) -> syn::Result<TokenStream> {
    let mut item: Item = syn::parse(input)?;
    let args: KernelArgs = from_tokens(args.into())?;

    if args.harness.is_some() && !matches!(item, Item::Fn(_)) {
        return Err(syn::Error::new_spanned(
            item,
            "`harness` is only supported on functions",
        ));
    }

    let tokens = match item.clone() {
        Item::Fn(kernel) => {
            let kernel = Launch::from_item_fn(kernel, args)?;
            RemoveHelpers.visit_item_mut(&mut item);
            ReplaceDefines.visit_item_mut(&mut item);

            let extra_allow = match kernel.func.context.is_intrinsic {
                true => quote![#[allow(unused_variables)]],
                false => quote![],
            };

            return Ok(TokenStream::from(quote! {
                #[allow(dead_code, clippy::too_many_arguments)]
                #extra_allow
                #item
                #kernel
            }));
        }
        Item::Trait(kernel_trait) => {
            let is_debug = args.debug.is_present();
            let expand_trait = CubeTrait::from_item_trait(kernel_trait, args)?;

            let tokens = TokenStream::from(quote! {
                #expand_trait
            });
            if is_debug {
                panic!("{tokens}");
            }
            return Ok(tokens);
        }
        Item::Impl(item_impl) => {
            if item_impl.trait_.is_some() {
                let mut expand_impl = CubeTraitImpl::from_item_impl(item_impl, &args)?;
                let expand_impl = expand_impl.to_tokens_mut();

                Ok(TokenStream::from(quote! {
                    #expand_impl
                }))
            } else {
                let mut expand_impl = CubeImpl::from_item_impl(item_impl, &args)?;
                let expand_impl = expand_impl.to_tokens_mut();

                Ok(TokenStream::from(quote! {
                    #expand_impl
                }))
            }
        }
        item => Err(syn::Error::new_spanned(
            item,
            "`#[cube]` is only supported on traits and functions",
        ))?,
    };

    if args.debug.is_present() {
        match tokens {
            Ok(tokens) => panic!("{tokens}"),
            Err(err) => panic!("{err}"),
        };
    }

    tokens
}

/// Derive macro to define a cube type that is launched with a kernel
#[proc_macro_derive(CubeLaunch, attributes(cube, launch))]
pub fn module_derive_cube_launch(input: TokenStream) -> TokenStream {
    gen_cube_type(input, true)
}

/// Derive macro to define a cube type that is not launched
#[proc_macro_derive(CubeType, attributes(cube, expand))]
pub fn module_derive_cube_type(input: TokenStream) -> TokenStream {
    gen_cube_type(input, false)
}

/// Derive macro to let a cube type cross between the host and a cube
/// function's `harness`, as an input or an output. The type must also derive
/// `CubeType` and `CubeLaunch`.
///
/// Its impls are gated like a harness, by `cfg(test)` unless another gate is
/// given with `#[cube(harness(cfg(<predicate>)))]` or `#[cube(harness(always))]`.
/// It must be at least as permissive as the gate of every harness using the
/// type.
#[proc_macro_derive(CubeHarness, attributes(cube, launch))]
pub fn module_derive_cube_harness(input: TokenStream) -> TokenStream {
    let parsed = syn::parse(input);

    let input = match &parsed {
        Ok(val) => val,
        Err(err) => return err.to_compile_error().into(),
    };

    match generate_cube_harness(input) {
        Ok(val) => val.into(),
        Err(err) => err.to_compile_error().into(),
    }
}

fn gen_cube_type(input: TokenStream, with_launch: bool) -> TokenStream {
    let parsed = syn::parse(input);

    let input = match &parsed {
        Ok(val) => val,
        Err(err) => return err.to_compile_error().into(),
    };

    match generate_cube_type(input, with_launch) {
        Ok(val) => val.into(),
        Err(err) => err.to_compile_error().into(),
    }
}

/// Attribute macro to define a type that can be used as a kernel comptime
/// argument This derive Debug, Hash, `PartialEq`, Eq, Clone, Copy
#[proc_macro_attribute]
pub fn derive_cube_comptime(_metadata: TokenStream, input: TokenStream) -> TokenStream {
    let input: proc_macro2::TokenStream = input.into();
    quote! {
        #[derive(Debug, Hash, PartialEq, Eq, Clone, Copy)]
        #input
    }
    .into()
}

/// Attribute macro to derive cube traits for existing structs, without redefining that struct.
#[proc_macro_attribute]
pub fn derive_expand(metadata: TokenStream, input: TokenStream) -> TokenStream {
    match generate_derive_expand(input.into(), metadata.into()) {
        Ok(val) => val.into(),
        Err(err) => err.to_compile_error().into(),
    }
}

/// Mark the contents of this macro as compile time values, turning off all
/// expansion for this code and using it verbatim
///
/// # Example
/// ```ignored
/// #use cubecl_macros::cube;
/// #fn some_rust_function(a: u32) -> u32 {}
/// #[cube]
/// fn do_stuff(input: u32) -> u32 {
///     let comptime_value = comptime! { some_rust_function(3) };
///     input + comptime_value
/// }
/// ```
#[proc_macro]
pub fn comptime(input: TokenStream) -> TokenStream {
    let tokens: proc_macro2::TokenStream = input.into();
    quote![{ #tokens }].into()
}

/// Mark the contents of this macro as an intrinsic, turning off all expansion
/// for this code and calling it with the scope
///
/// # Example
/// ```ignored
/// #use cubecl_macros::cube;
/// #[cube]
/// fn do_stuff(input: u32) -> u32 {
///     let comptime_value = intrinsic! { |scope| u32::elem_size(scope) };
///     input + comptime_value
/// }
/// ```
#[proc_macro]
pub fn intrinsic(_input: TokenStream) -> TokenStream {
    quote![{ cubecl::unexpanded!() }].into()
}

/// GPU version of [`asm`](std::arch::asm). Currently parses all the same options, but most are not
/// applicable to GPU assembly architectures. Should validate and give proper errors at some point.
/// Also adds a new register spec: the inferred register specifier (`_`). This is because the
/// specifier isn't actually meaningful in PTX and is currently ignored, with constraints being
/// inferred from the value type. The reason it's still present is because we may want to add an
/// explicit `mem` specifier to allow fine-grained memory clobbering, or support for other assembly
/// formats that do use different register types.
///
/// # Example
/// ```ignored
/// #use cubecl_macros::cube;
/// #[cube]
/// fn do_stuff(input: u32) -> u32 {
///     let mut out: u32;
///     gpu_asm!("some.custom.ptx {} {}", out(_) out, in(_) input)
/// }
/// ```
#[proc_macro]
pub fn gpu_asm(input: TokenStream) -> TokenStream {
    match generate_asm_unexpanded(input.into()) {
        Ok(val) => val.into(),
        Err(err) => err.to_compile_error().into(),
    }
}

/// Makes the function return a compile time value
/// Useful in a cube trait to have a part of the trait return comptime values
///
/// # Example
/// ```ignored
/// #use cubecl_macros::cube;
/// #[cube]
/// fn do_stuff(#[comptime] input: u32) -> comptime_type!(u32) {
///     input + 5   
/// }
/// ```
///
/// TODO: calling a trait method returning `comptime_type` from
/// within another trait method does not work
#[proc_macro]
pub fn comptime_type(input: TokenStream) -> TokenStream {
    let tokens: proc_macro2::TokenStream = input.into();
    quote![ #tokens ].into()
}

/// Insert a literal comment into the kernel source code.
///
/// # Example
/// ```ignored
/// #use cubecl_macros::cube;
/// #[cube]
/// fn do_stuff(input: u32) -> u32 {
///     comment!("Add five to the input");
///     input + 5
/// }
/// ```
#[proc_macro]
pub fn comment(input: TokenStream) -> TokenStream {
    let tokens: proc_macro2::TokenStream = input.into();
    quote![{ #tokens }].into()
}

/// Terminate the execution of the kernel for the current unit.
///
/// This terminates the execution of the unit even if nested inside many
/// functions.
///
/// # Example
/// ```ignored
/// #use cubecl_macros::cube;
/// #[cube]
/// fn stop_if_more_than_ten(input: u32)  {
///     if input > 10 {
///         terminate!();
///     }
/// }
/// ```
#[proc_macro]
pub fn terminate(input: TokenStream) -> TokenStream {
    let tokens: proc_macro2::TokenStream = input.into();
    quote![{ #tokens }].into()
}

/// Implements display and initialization for autotune keys.
///
/// # Helper
///
/// Use the `#[autotune(anchor)]` helper attribute to anchor a numerical value.
/// This groups multiple numerical values into the same bucket.
///
/// For now, only an exponential function is supported, and it can be modified
/// with `exp`. By default, the base is '2' and there are no `min` or `max`
/// provided.
///
/// # Example
/// ```ignore
/// #[derive(AutotuneKey, Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
/// pub struct OperationKey {
///     #[autotune(name = "Batch Size")]
///     batch_size: usize,
///     channels: usize,
///     #[autotune(anchor(exp(min = 16, max = 1024, base = 2)))]
///     height: usize,
///     #[autotune(anchor)]
///     width: usize,
/// }
/// ```
#[proc_macro_derive(AutotuneKey, attributes(autotune))]
pub fn derive_autotune_key(input: TokenStream) -> TokenStream {
    let input = syn::parse(input).unwrap();
    match generate_autotune_key(input) {
        Ok(tokens) => tokens.into(),
        Err(e) => e.into_compile_error().into(),
    }
}

/// Implements `IntoRuntime` for a `CubeType`
#[proc_macro_derive(IntoRuntime, attributes(cube))]
pub fn derive_into_runtime(input: TokenStream) -> TokenStream {
    let input = syn::parse(input).unwrap();
    match generate_into_runtime(&input) {
        Ok(tokens) => tokens.into(),
        Err(e) => e.into_compile_error().into(),
    }
}

/// Implements mutability for a `CubeType`
#[proc_macro_derive(CubeTypeMut, attributes(cube))]
pub fn derive_assign(input: TokenStream) -> TokenStream {
    let input = syn::parse(input).unwrap();
    match generate_cube_type_mut(&input) {
        Ok(tokens) => tokens.into(),
        Err(e) => e.into_compile_error().into(),
    }
}