cubecl-macros-internal 0.11.2

Internal 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
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90

use proc_macro::TokenStream;
use syn::{DeriveInput, ItemImpl, parse_macro_input};
use type_hash::type_hash_impl;

use crate::{
    generate::{
        const_eval::generate_const_eval,
        cube_op::{generate_cube_op, generate_op_traits},
        pass_name::{generate_pass_name, generate_rewrite_name},
        simplify::generate_simplify,
    },
    parse::{
        cube_op::{CubeOp, CubeOpArgs, PathList},
        from_meta_tokens,
    },
};

mod enum_set;
mod generate;
mod parse;
mod type_hash;

macro_rules! macro_try {
    ($op: expr) => {
        match $op {
            Ok(res) => res,
            Err(e) => return e.into_compile_error().into(),
        }
    };
}

/// Derives the machinery an enum needs to live in a `cubecl_ir::EnumSet`.
///
/// The enum must be a plain C-like enum of at most 64 unit variants with no explicit
/// discriminants, and must also derive `Clone`, `Copy`, `PartialEq` and `Eq`, which the
/// `EnumSetType` trait requires.
#[proc_macro_derive(EnumSetType)]
pub fn derive_enum_set_type(input: TokenStream) -> TokenStream {
    let input = parse_macro_input!(input as DeriveInput);
    macro_try!(crate::enum_set::enum_set_type_impl(input)).into()
}

#[proc_macro_derive(TypeHash, attributes(type_hash))]
pub fn derive_type_hash(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
    let input = syn::parse(input).unwrap();
    type_hash_impl(input).into()
}

/// Version of `pliron_op` that automatically derives shape ops, accessors, and a constructor from
/// struct fields. Also enables `format` by default, and sets the verifier to `"succ"` unless
/// overridden.
#[proc_macro_attribute]
pub fn cube_op(args: TokenStream, input: TokenStream) -> TokenStream {
    let args = macro_try!(from_meta_tokens(args.into()));
    let input = macro_try!(syn::parse(input));
    macro_try!(generate_cube_op(input, args)).into()
}

#[proc_macro_attribute]
pub fn op_traits(args: TokenStream, input: TokenStream) -> TokenStream {
    let traits = parse_macro_input!(args as PathList);
    let input = parse_macro_input!(input as DeriveInput);
    macro_try!(generate_op_traits(input, traits)).into()
}

/// Implements `Pass::name` on the annotated `impl Pass for X` block, returning the short name of
/// the pass struct (e.g. `"InsertEntrypointPass"`) instead of the fully-qualified type path.
#[proc_macro_attribute]
pub fn pass_name(_args: TokenStream, input: TokenStream) -> TokenStream {
    let input = parse_macro_input!(input as ItemImpl);
    macro_try!(generate_pass_name(input)).into()
}

/// Implements `NamedRewrite::name` on the annotated struct, returning the short name of
/// the rewrite struct (e.g. `"InsertEntrypointPass"`) instead of the fully-qualified type path.
#[proc_macro_derive(NamedRewrite)]
pub fn derive_rewrite_name(input: TokenStream) -> TokenStream {
    let input = parse_macro_input!(input as DeriveInput);
    macro_try!(generate_rewrite_name(input)).into()
}

#[proc_macro]
pub fn const_eval(input: TokenStream) -> TokenStream {
    macro_try!(generate_const_eval(input.into())).into()
}

#[proc_macro]
pub fn simplify(input: TokenStream) -> TokenStream {
    macro_try!(generate_simplify(input.into())).into()
}