llvm_mca_macros/
lib.rs

1#![doc = include_str!("../README.md")]
2#![feature(proc_macro_quote)]
3use proc_macro::{Span, TokenStream};
4use quote::{quote, ToTokens as _};
5use syn::{
6    parse::{Parse, ParseStream},
7    parse_quote, Item, ItemFn,
8};
9
10struct MacroArgs {
11    allow_inline: bool,
12}
13
14impl Parse for MacroArgs {
15    fn parse(input: ParseStream) -> syn::Result<Self> {
16        if input.is_empty() {
17            return Ok(Self {
18                allow_inline: false,
19            });
20        }
21
22        // TODO: Allow specifying regions names in arguments, to match
23        // `llvm_mca_begin!(...)` and `llvm_mca_end!(...)`
24        let arg = input.parse::<syn::Ident>()?;
25        if arg == "allow_inline" {
26            Ok(Self { allow_inline: true })
27        } else {
28            Err(syn::Error::new(arg.span(), "expected `allow_inline`"))
29        }
30    }
31}
32
33/// Wrap the body of a function with `LLVM-MCA-BEGIN` and `LLVM-MCA-END` markers.
34///
35/// The markers are inserted as inline assembly, after the function prologue and
36/// before the function epilogue.
37///
38/// # Examples
39///
40/// This:
41/// ```
42/// use llvm_mca_macros::llvm_mca;
43/// #[llvm_mca]
44/// fn quadruple(x: u32) -> u32 {
45///     let doubled = x + x;
46///     doubled + doubled
47/// }
48/// ```
49///
50/// is equivalent to:
51/// ```
52/// #[inline(never)]
53/// fn quadruple(x: u32) -> u32 {
54///     // emit LLVM-MCA-BEGIN marker
55///     let ret = {
56///         let doubled = x + x;
57///         doubled + doubled
58///     };
59///     // emit LLVM-MCA-END marker
60///     ret
61/// }
62/// ```
63///
64/// If inlining is desired, specify the `allow_inline` argument:
65/// ```
66/// use llvm_mca_macros::llvm_mca;
67/// #[llvm_mca(allow_inline)]
68/// fn quadruple(x: u32) -> u32 {
69///     let doubled = x + x;
70///     doubled + doubled
71/// }
72/// ```
73///
74/// which is equivalent to:
75/// ```
76/// fn quadruple(x: u32) -> u32 {
77///     // emit LLVM-MCA-BEGIN marker
78///     let ret = {
79///         let doubled = x + x;
80///         doubled + doubled
81///     };
82///     // emit LLVM-MCA-END marker
83///     ret
84/// }
85/// ```
86#[proc_macro_attribute]
87pub fn llvm_mca(attrs: TokenStream, input: TokenStream) -> TokenStream {
88    let function = match syn::parse(input) {
89        Ok(Item::Fn(function)) => function,
90        _ => {
91            return syn::Error::new(Span::call_site().into(), "expected function")
92                .to_compile_error()
93                .into()
94        }
95    };
96
97    // Take the original block and wedge it between the two markers. By default,
98    //
99    // `rustc` assumes that an `asm!(..)` block requires a stack frame so it
100    // includes stack-frame setup/teardown in the function prologue/epilogue. We
101    // can avoid this by adding the `options(nostack)` attribute to the
102    // `asm!(..)` block.
103    //
104    //Use the `options(nostack)` attribute to prevent this
105    let original_block = function.block;
106    let block = syn::parse(
107        quote! {{
108            unsafe {
109                std::arch::asm!(";# LLVM-MCA-BEGIN", options(nostack));
110            }
111            let ret = #original_block;
112            unsafe {
113                std::arch::asm!(";# LLVM-MCA-END", options(nostack));
114            }
115            ret
116        }}
117        .into(),
118    )
119    .unwrap();
120
121    let args = match syn::parse::<MacroArgs>(attrs) {
122        Ok(args) => args,
123        Err(err) => return err.to_compile_error().into(),
124    };
125
126    // Add `#[inline(never)]` to the function attributes if `allow_inline` is
127    // _not_ specified
128    let attrs = if args.allow_inline {
129        function.attrs
130    } else {
131        function
132            .attrs
133            .into_iter()
134            .chain(std::iter::once(parse_quote! {
135                #[inline(never)]
136            }))
137            .collect()
138    };
139
140    let result = ItemFn {
141        attrs,
142        block,
143        ..function
144    };
145
146    result.into_token_stream().into()
147}
llvm_mca_macros/lib.rs

llvm_mca_macros/
lib.rs
