cubecl_macros_internal/
lib.rs

1use proc_macro::TokenStream;
2use syn::{DeriveInput, ItemImpl, parse_macro_input};
3use type_hash::type_hash_impl;
4
5use crate::{
6    generate::{
7        const_eval::generate_const_eval,
8        cube_op::{generate_cube_op, generate_op_traits},
9        pass_name::{generate_pass_name, generate_rewrite_name},
10        simplify::generate_simplify,
11    },
12    parse::{
13        cube_op::{CubeOp, CubeOpArgs, PathList},
14        from_meta_tokens,
15    },
16};
17
18mod enum_set;
19mod generate;
20mod parse;
21mod type_hash;
22
23macro_rules! macro_try {
24    ($op: expr) => {
25        match $op {
26            Ok(res) => res,
27            Err(e) => return e.into_compile_error().into(),
28        }
29    };
30}
31
32/// Derives the machinery an enum needs to live in a `cubecl_ir::EnumSet`.
33///
34/// The enum must be a plain C-like enum of at most 64 unit variants with no explicit
35/// discriminants, and must also derive `Clone`, `Copy`, `PartialEq` and `Eq`, which the
36/// `EnumSetType` trait requires.
37#[proc_macro_derive(EnumSetType)]
38pub fn derive_enum_set_type(input: TokenStream) -> TokenStream {
39    let input = parse_macro_input!(input as DeriveInput);
40    macro_try!(crate::enum_set::enum_set_type_impl(input)).into()
41}
42
43#[proc_macro_derive(TypeHash, attributes(type_hash))]
44pub fn derive_type_hash(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
45    let input = syn::parse(input).unwrap();
46    type_hash_impl(input).into()
47}
48
49/// Version of `pliron_op` that automatically derives shape ops, accessors, and a constructor from
50/// struct fields. Also enables `format` by default, and sets the verifier to `"succ"` unless
51/// overridden.
52#[proc_macro_attribute]
53pub fn cube_op(args: TokenStream, input: TokenStream) -> TokenStream {
54    let args = macro_try!(from_meta_tokens(args.into()));
55    let input = macro_try!(syn::parse(input));
56    macro_try!(generate_cube_op(input, args)).into()
57}
58
59#[proc_macro_attribute]
60pub fn op_traits(args: TokenStream, input: TokenStream) -> TokenStream {
61    let traits = parse_macro_input!(args as PathList);
62    let input = parse_macro_input!(input as DeriveInput);
63    macro_try!(generate_op_traits(input, traits)).into()
64}
65
66/// Implements `Pass::name` on the annotated `impl Pass for X` block, returning the short name of
67/// the pass struct (e.g. `"InsertEntrypointPass"`) instead of the fully-qualified type path.
68#[proc_macro_attribute]
69pub fn pass_name(_args: TokenStream, input: TokenStream) -> TokenStream {
70    let input = parse_macro_input!(input as ItemImpl);
71    macro_try!(generate_pass_name(input)).into()
72}
73
74/// Implements `NamedRewrite::name` on the annotated struct, returning the short name of
75/// the rewrite struct (e.g. `"InsertEntrypointPass"`) instead of the fully-qualified type path.
76#[proc_macro_derive(NamedRewrite)]
77pub fn derive_rewrite_name(input: TokenStream) -> TokenStream {
78    let input = parse_macro_input!(input as DeriveInput);
79    macro_try!(generate_rewrite_name(input)).into()
80}
81
82#[proc_macro]
83pub fn const_eval(input: TokenStream) -> TokenStream {
84    macro_try!(generate_const_eval(input.into())).into()
85}
86
87#[proc_macro]
88pub fn simplify(input: TokenStream) -> TokenStream {
89    macro_try!(generate_simplify(input.into())).into()
90}
cubecl_macros_internal/lib.rs

cubecl_macros_internal/
lib.rs
