cubecl-macros 0.11.0-pre.1

Procedural macros for CubeCL
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58

use proc_macro2::TokenStream;
use quote::{ToTokens, quote};
use syn::Type;

use crate::{
    parse::{kernel::expand_kernel_ty, signature::*},
    paths::prelude_type,
};

impl ToTokens for KernelSignature {
    fn to_tokens(&self, tokens: &mut TokenStream) {
        let scope = prelude_type("Scope");

        let name = &self.name;
        let generics = &self.generics;
        let where_clause = &generics.where_clause;
        let scope_lifetime = &self.scope_lifetime;

        let return_type = match &self.returns {
            KernelReturns::ExpandType(ty) => {
                let normalized_ty = expand_kernel_ty(ty.clone(), false)
                    .unwrap_or_else(|err| Type::Verbatim(err.into_compile_error()));
                quote![#normalized_ty]
            }
            KernelReturns::Plain(ty) => quote![#ty],
        };
        let out = if let Some(receiver) = &self.receiver_arg {
            let args = self.parameters.iter().skip(1);

            quote! {
                fn #name #generics(
                    #receiver,
                    scope: &#scope_lifetime #scope,
                    #(#args),*
                ) -> #return_type #where_clause
            }
        } else {
            let args = &self.parameters;
            quote! {
                fn #name #generics(
                    scope: &#scope_lifetime #scope,
                    #(#args),*
                ) -> #return_type #where_clause
            }
        };

        tokens.extend(out);
    }
}

impl ToTokens for KernelParam {
    fn to_tokens(&self, tokens: &mut TokenStream) {
        let name = &self.name;
        let ty = &self.normalized_ty;
        let mut_ = &self.mutability;
        tokens.extend(quote![#mut_ #name: #ty]);
    }
}