scalar or vector: add #[derive(ScalarOrVectorComposite)]

Author: Firestar99

crates/spirv-std/macros/src/lib.rs

@@ -74,6 +74,7 @@
 mod debug_printf;
 mod image;
 mod sample_param_permutations;
+mod scalar_or_vector_composite;
 
 use crate::debug_printf::{DebugPrintfInput, debug_printf_inner};
 use proc_macro::TokenStream;
@@ -311,3 +312,10 @@ pub fn debug_printfln(input: TokenStream) -> TokenStream {
 pub fn gen_sample_param_permutations(_attr: TokenStream, item: TokenStream) -> TokenStream {
     sample_param_permutations::gen_sample_param_permutations(item)
 }
+
+#[proc_macro_derive(ScalarOrVectorComposite)]
+pub fn derive_scalar_or_vector_composite(item: TokenStream) -> TokenStream {
+    scalar_or_vector_composite::derive(item.into())
+        .unwrap_or_else(syn::Error::into_compile_error)
+        .into()
+}

crates/spirv-std/macros/src/scalar_or_vector_composite.rs

@@ -0,0 +1,55 @@
+use proc_macro2::TokenStream;
+use quote::{ToTokens, quote};
+use syn::punctuated::Punctuated;
+use syn::{Fields, FieldsNamed, FieldsUnnamed, GenericParam, Token};
+
+pub fn derive(item: TokenStream) -> syn::Result<TokenStream> {
+    // Whenever we'll properly resolve the crate symbol, replace this.
+    let spirv_std = quote!(spirv_std);
+
+    // Defer all validation to our codegen backend. Rather than erroring here, emit garbage.
+    let item = syn::parse2::<syn::ItemStruct>(item)?;
+    let struct_ident = &item.ident;
+    let gens = &item.generics.params;
+    let gen_refs = &item
+        .generics
+        .params
+        .iter()
+        .map(|p| match p {
+            GenericParam::Lifetime(p) => p.lifetime.to_token_stream(),
+            GenericParam::Type(p) => p.ident.to_token_stream(),
+            GenericParam::Const(p) => p.ident.to_token_stream(),
+        })
+        .collect::<Punctuated<_, Token![,]>>();
+    let where_clause = &item.generics.where_clause;
+
+    let content =
+        match item.fields {
+            Fields::Named(FieldsNamed { named, .. }) => {
+                let content = named.iter().map(|f| {
+                let ident = &f.ident;
+                quote!(#ident: #spirv_std::ScalarOrVectorComposite::transform(self.#ident, f))
+            }).collect::<Punctuated<_, Token![,]>>();
+                quote!(Self { #content })
+            }
+            Fields::Unnamed(FieldsUnnamed { unnamed, .. }) => {
+                let content = (0..unnamed.len())
+                    .map(|i| {
+                        let i = syn::Index::from(i);
+                        quote!(#spirv_std::ScalarOrVectorComposite::transform(self.#i, f))
+                    })
+                    .collect::<Punctuated<_, Token![,]>>();
+                quote!(Self(#content))
+            }
+            Fields::Unit => quote!(Self),
+        };
+
+    Ok(quote! {
+        impl<#gens> #spirv_std::ScalarOrVectorComposite for #struct_ident<#gen_refs> #where_clause {
+            #[inline]
+            fn transform<F: #spirv_std::ScalarOrVectorTransform>(self, f: &mut F) -> Self {
+                #content
+            }
+        }
+    })
+}

crates/spirv-std/src/lib.rs

@@ -87,6 +87,7 @@
 /// Public re-export of the `spirv-std-macros` crate.
 #[macro_use]
 pub extern crate spirv_std_macros as macros;
+pub use macros::ScalarOrVectorComposite;
 pub use macros::spirv;
 pub use macros::{debug_printf, debug_printfln};
 

tests/compiletests/ui/arch/subgroup/subgroup_composite.rs

@@ -0,0 +1,54 @@
+// build-pass
+// compile-flags: -C target-feature=+GroupNonUniform,+GroupNonUniformBallot,+GroupNonUniformShuffle,+GroupNonUniformShuffleRelative,+ext:SPV_KHR_vulkan_memory_model
+// normalize-stderr-test "OpLine .*\n" -> ""
+// ignore-vulkan1.0
+// ignore-vulkan1.1
+// ignore-spv1.0
+// ignore-spv1.1
+// ignore-spv1.2
+// ignore-spv1.3
+// ignore-spv1.4
+
+use glam::*;
+use spirv_std::ScalarOrVectorComposite;
+use spirv_std::arch::*;
+use spirv_std::spirv;
+
+#[derive(Copy, Clone, ScalarOrVectorComposite)]
+pub struct MyStruct {
+    a: f32,
+    b: UVec3,
+    c: Nested,
+    d: Zst,
+}
+
+#[derive(Copy, Clone, ScalarOrVectorComposite)]
+pub struct Nested(i32);
+
+#[derive(Copy, Clone, ScalarOrVectorComposite)]
+pub struct Zst;
+
+#[spirv(compute(threads(32)))]
+pub fn main(
+    #[spirv(local_invocation_index)] inv_id: UVec3,
+    #[spirv(descriptor_set = 0, binding = 0, storage_buffer)] output: &mut UVec3,
+) {
+    unsafe {
+        let my_struct = MyStruct {
+            a: 1.,
+            b: inv_id,
+            c: Nested(-42),
+            d: Zst,
+        };
+
+        let mut out = UVec3::ZERO;
+        // before spv1.5 / vulkan1.2, this id = 19 must be a constant
+        out += subgroup_broadcast(my_struct, 19).b;
+        out += subgroup_broadcast_first(my_struct).b;
+        out += subgroup_shuffle(my_struct, 2).b;
+        out += subgroup_shuffle_xor(my_struct, 4).b;
+        out += subgroup_shuffle_up(my_struct, 5).b;
+        out += subgroup_shuffle_down(my_struct, 7).b;
+        *output = out;
+    }
+}

tests/compiletests/ui/arch/subgroup/subgroup_composite_shuffle.rs

@@ -0,0 +1,45 @@
+// build-pass
+// compile-flags: -C target-feature=+GroupNonUniform,+GroupNonUniformShuffle,+ext:SPV_KHR_vulkan_memory_model
+// compile-flags: -C llvm-args=--disassemble-fn=subgroup_composite_shuffle::disassembly
+// normalize-stderr-test "OpLine .*\n" -> ""
+
+use glam::*;
+use spirv_std::ScalarOrVectorComposite;
+use spirv_std::arch::*;
+use spirv_std::spirv;
+
+#[derive(Copy, Clone, ScalarOrVectorComposite)]
+pub struct MyStruct {
+    a: f32,
+    b: UVec3,
+    c: Nested,
+    d: Zst,
+}
+
+#[derive(Copy, Clone, ScalarOrVectorComposite)]
+pub struct Nested(i32);
+
+#[derive(Copy, Clone, ScalarOrVectorComposite)]
+pub struct Zst;
+
+/// this should be 3 `subgroup_shuffle` instructions, with all calls inlined
+fn disassembly(my_struct: MyStruct, id: u32) -> MyStruct {
+    subgroup_shuffle(my_struct, id)
+}
+
+#[spirv(compute(threads(32)))]
+pub fn main(
+    #[spirv(local_invocation_index)] inv_id: UVec3,
+    #[spirv(descriptor_set = 0, binding = 0, storage_buffer)] output: &mut MyStruct,
+) {
+    unsafe {
+        let my_struct = MyStruct {
+            a: inv_id.x as f32,
+            b: inv_id,
+            c: Nested(5i32 - inv_id.x as i32),
+            d: Zst,
+        };
+
+        *output = disassembly(my_struct, 5);
+    }
+}

tests/compiletests/ui/arch/subgroup/subgroup_composite_shuffle.stderr

@@ -0,0 +1,16 @@
+%1 = OpFunction  %2  None %3
+%4 = OpFunctionParameter  %2
+%5 = OpFunctionParameter  %6
+%7 = OpLabel
+%9 = OpCompositeExtract  %10  %4 0
+%12 = OpGroupNonUniformShuffle  %10  %13 %9 %5
+%14 = OpCompositeExtract  %15  %4 1
+%16 = OpGroupNonUniformShuffle  %15  %13 %14 %5
+%17 = OpCompositeExtract  %18  %4 2
+%19 = OpGroupNonUniformShuffle  %18  %13 %17 %5
+%20 = OpCompositeInsert  %2  %12 %21 0
+%22 = OpCompositeInsert  %2  %16 %20 1
+%23 = OpCompositeInsert  %2  %19 %22 2
+OpNoLine
+OpReturnValue %23
+OpFunctionEnd