accel_derive/lib.rs
1#![recursion_limit = "128"]
2
3//! Get compiled PTX as `String`
4//! ----------------------------
5//!
6//! The proc-macro `#[kernel]` creates a submodule `add::` in addition to a function `add`.
7//! Kernel Rust code is compiled into PTX string using rustc's `nvptx64-nvidia-cuda` toolchain.
8//! Generated PTX string is embedded into proc-macro output as `{kernel_name}::PTX_STR`.
9//!
10//! ```
11//! use accel_derive::kernel;
12//!
13//! #[kernel]
14//! unsafe fn add(a: *const f64, b: *const f64, c: *mut f64, n: usize) {
15//! let i = accel_core::index();
16//! if (i as usize) < n {
17//! *c.offset(i) = *a.offset(i) + *b.offset(i);
18//! }
19//! }
20//!
21//! // PTX assembler code is embedded as `add::PTX_STR`
22//! println!("{}", add::PTX_STR);
23//! ```
24
25mod builder;
26mod host;
27mod parser;
28
29use proc_macro::TokenStream;
30
31#[proc_macro_attribute]
32pub fn kernel(_attr: TokenStream, func: TokenStream) -> TokenStream {
33 let func: syn::ItemFn = syn::parse(func).expect("Not a function");
34 let ptx_str = builder::compile_tokens(&func).expect("Failed to compile to PTX");
35 host::func2caller(&ptx_str, &func).into()
36}