Merge pull request #10 from szymon-zadworny/macros

Macros

Author: bratpiorka

Cargo.lock

@@ -1,3 +1,3 @@
 [workspace]
 resolver = "3"
-members = ["oneapi-rs","oneapi-rs-sys"]
+members = ["oneapi-rs", "oneapi-rs-derive","oneapi-rs-sys"]

Cargo.toml

@@ -0,0 +1,13 @@
+[package]
+name = "oneapi-rs-derive"
+version = "0.1.0"
+edition = "2024"
+
+[lib]
+proc-macro = true
+
+[dependencies]
+proc-macro-crate = "3.5.0"
+proc-macro2 = "1.0.107"
+quote = "1.0.47"
+syn = "3.0.3"

oneapi-rs-derive/Cargo.toml

@@ -0,0 +1,85 @@
+use proc_macro::TokenStream;
+use proc_macro_crate::{FoundCrate, crate_name};
+use quote::{format_ident, quote};
+use syn::{
+    Data, DataStruct, DeriveInput, Error, Field, Ident, LitInt, WhereClause, parse_macro_input,
+    parse_quote,
+};
+
+fn find_oneapi() -> Ident {
+    let crate_name =
+        crate_name("oneapi-rs").expect("Expected oneapi-rs to be present in Cargo.toml");
+    match crate_name {
+        FoundCrate::Itself => format_ident!("oneapi_rs"),
+        FoundCrate::Name(name) => format_ident!("{name}"),
+    }
+}
+
+/// Derive macro generating an impl of the `KernelArgumentList` trait for a given struct.
+#[proc_macro_derive(KernelArgumentList)]
+pub fn derive_kernel_argument_list(input: TokenStream) -> TokenStream {
+    let mut input = parse_macro_input!(input as DeriveInput);
+    let oneapi = find_oneapi();
+
+    let Data::Struct(data) = &input.data else {
+        return Error::new_spanned(input, "This derive macro only works on structs.")
+            .into_compile_error()
+            .into();
+    };
+    expand_where_clause(input.generics.make_where_clause(), data, &oneapi);
+    let (impl_generics, ty_generics, where_clause) = input.generics.split_for_impl();
+
+    let ident = input.ident;
+    let argc = data.fields.len();
+    let members = data.fields.members();
+
+    let expanded = quote! {
+        unsafe impl #impl_generics #oneapi::kernel::KernelArgumentList<#argc>
+        for #ident #ty_generics #where_clause {
+            unsafe fn as_raw_arg_list(&self) -> [&[u8]; #argc] {
+                [ #(unsafe { self.#members.as_raw_arg() }),* ]
+            }
+        }
+    };
+
+    TokenStream::from(expanded)
+}
+
+fn expand_where_clause(where_clause: &mut WhereClause, data: &DataStruct, oneapi: &Ident) {
+    for Field { ty, .. } in &data.fields {
+        where_clause
+            .predicates
+            .push(parse_quote!(#ty: #oneapi::kernel::KernelArgument));
+    }
+}
+
+fn get_single_tuple_impl(argc: usize) -> proc_macro2::TokenStream {
+    let iter = { 0..argc }.map(syn::Index::from);
+    let types = { 0..argc }
+        .map(|i| format_ident!("T{i}"))
+        .collect::<Vec<_>>();
+
+    quote! {
+        unsafe impl<#(#types),*> crate::kernel::KernelArgumentList<#argc> for (#(#types),*)
+        where #(#types: crate::kernel::KernelArgument),* {
+            unsafe fn as_raw_arg_list(&self) -> [&[u8]; #argc] {
+                [ #(unsafe { self.#iter.as_raw_arg() }),* ]
+            }
+        }
+    }
+}
+
+/// A macro that generates tuples from 2..N which implement the `KernelArgumentList` trait.
+#[proc_macro]
+pub fn impl_arg_list_for_tuples(input: TokenStream) -> TokenStream {
+    let input = parse_macro_input!(input as LitInt);
+    let argc = input.base10_parse::<usize>().unwrap();
+
+    let impls = { 2..=argc }.map(get_single_tuple_impl);
+
+    let expanded = quote! {
+        #(#impls)*
+    };
+
+    TokenStream::from(expanded)
+}

oneapi-rs-derive/src/lib.rs

@@ -9,6 +9,7 @@ allocator-api2 = "0.4.0"
 bytemuck = "1.25.1"
 cxx = "1.0.194"
 oneapi-rs-sys = { path = "../oneapi-rs-sys" }
+oneapi-rs-derive = { path = "../oneapi-rs-derive" }
 pin-project = "1.1.13"
 thiserror = "2.0.18"
 

oneapi-rs/Cargo.toml

@@ -6,13 +6,7 @@
 // SPDX-License-Identifier: MIT OR Apache-2.0
 //
 
-use oneapi_rs::{
-    buffer::Buffer,
-    kernel::{KernelArgument, KernelArgumentList},
-    queue::Queue,
-    range::NdRange,
-    usm::{SharedAllocator, UsmAllocator},
-};
+use oneapi_rs::{queue::Queue, range::NdRange};
 
 static IOTA_SRC: &str = r#"
 #include <sycl/sycl.hpp>
@@ -21,46 +15,23 @@ namespace syclexp = sycl::ext::oneapi::experimental;
 
 extern "C"
 SYCL_EXT_ONEAPI_FUNCTION_PROPERTY((syclexp::nd_range_kernel<1>))
-void iota(float start, float *ptr) {
+void iota(double start, double *ptr) {
     size_t id = syclext::this_work_item::get_nd_item<1>().get_global_linear_id();
-    ptr[id] = start + static_cast<float>(id);
+    ptr[id] = start + static_cast<double>(id);
 }
 "#;
 
-struct IotaArgs<'a> {
-    start: f32,
-    buffer: &'a mut Buffer<f32, UsmAllocator<SharedAllocator>>,
-}
-
-unsafe impl<'a> KernelArgumentList<2> for IotaArgs<'a> {
-    unsafe fn as_raw_arg_list(&self) -> [&[u8]; 2] {
-        return [unsafe { self.start.as_raw_arg() }, unsafe {
-            self.buffer.as_raw_arg()
-        }];
-    }
-}
-
 fn main() {
     let mut queue = Queue::new();
-    let mut buffer = queue.alloc_shared::<f32>(1024).wait();
+    let mut buffer = queue.alloc_shared::<f64>(1024).wait();
 
     let kernel = queue
         .get_context()
         .create_kernel_bundle_from_source(IOTA_SRC)
         .build()
         .get_kernel("iota");
 
-    unsafe {
-        queue.launch(
-            NdRange::new([1024], [16]),
-            &kernel,
-            IotaArgs {
-                start: 3.14,
-                buffer: &mut buffer,
-            },
-        )
-    }
-    .wait();
+    unsafe { queue.launch(NdRange::new([1024], [16]), &kernel, (3.14, &mut buffer)) }.wait();
 
     for e in buffer.iter() {
         print!("{e} ");