Add gpu_only proc macro (#392)

EmbarkStudios · Jan 25, 2021 · 353bb96 · 353bb96
1 parent 19cf332
commit 353bb96
Show file tree

Hide file tree

Showing 6 changed files with 39 additions and 25 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/crates/spirv-std-macros/Cargo.toml b/crates/spirv-std-macros/Cargo.toml
@@ -8,3 +8,7 @@ repository = "https://github.com/EmbarkStudios/rust-gpu"
 
 [lib]
 proc-macro = true
+
+[dependencies]
+quote = "1.0.8"
+syn = { version = "1.0.58", features=["full"] }
diff --git a/crates/spirv-std-macros/src/lib.rs b/crates/spirv-std-macros/src/lib.rs
@@ -29,3 +29,28 @@ pub fn spirv(_attr: TokenStream, item: TokenStream) -> TokenStream {
     }
     tokens.into_iter().collect()
 }
+
+#[proc_macro_attribute]
+pub fn gpu_only(_attr: TokenStream, item: TokenStream) -> TokenStream {
+    let syn::ItemFn {
+        attrs,
+        vis,
+        sig,
+        block,
+    } = syn::parse_macro_input!(item as syn::ItemFn);
+
+    let fn_name = sig.ident.clone();
+
+    let output = quote::quote! {
+        // Don't warn on unused arguments on the CPU side.
+        #[cfg_attr(not(target_arch = "spirv"), allow(unused_variables))]
+        #(#attrs)* #vis #sig {
+            #[cfg(target_arch="spirv")] { #block }
+            #[cfg(not(target_arch="spirv"))] {
+                unimplemented!(concat!("`", stringify!(#fn_name), "` is only available on SPIR-V platforms."))
+            }
+        }
+    };
+
+    output.into()
+}
diff --git a/crates/spirv-std/src/derivative.rs b/crates/spirv-std/src/derivative.rs
@@ -56,10 +56,8 @@ macro_rules! deriv_caps {
 
 macro_rules! deriv_fn {
     ($name:ident, $inst:ident, $needs_caps:tt) => {
+        #[spirv_std_macros::gpu_only]
         fn $name(self) -> Self {
-            #[cfg(not(target_arch = "spirv"))]
-            panic!(concat!(stringify!($name), " is not supported on the CPU"));
-            #[cfg(target_arch = "spirv")]
             unsafe {
                 let mut o = Default::default();
                 deriv_caps!($needs_caps);

diff --git a/crates/spirv-std/src/lib.rs b/crates/spirv-std/src/lib.rs
@@ -52,8 +52,8 @@ pub use num_traits;
 pub use textures::*;
 
 /// Calls the `OpDemoteToHelperInvocationEXT` instruction, which corresponds to discard() in HLSL
+#[spirv_std_macros::gpu_only]
 pub fn demote_to_helper_invocation() {
-    #[cfg(target_arch = "spirv")]
     unsafe {
         asm!(
             "OpExtension \"SPV_EXT_demote_to_helper_invocation\"",
@@ -64,8 +64,8 @@ pub fn demote_to_helper_invocation() {
 }
 
 /// Calls the `OpKill` instruction, which corresponds to discard() in GLSL
+#[spirv_std_macros::gpu_only]
 pub fn discard() {
-    #[cfg(target_arch = "spirv")]
     unsafe {
         asm!("OpKill", "%unused = OpLabel");
     }

diff --git a/crates/spirv-std/src/textures.rs b/crates/spirv-std/src/textures.rs
@@ -23,14 +23,8 @@ pub struct Image2d {
 }
 
 impl Image2d {
+    #[spirv_std_macros::gpu_only]
     pub fn sample(&self, sampler: Sampler, coord: Vec2) -> Vec4 {
-        #[cfg(not(target_arch = "spirv"))]
-        {
-            let _ = sampler;
-            let _ = coord;
-            panic!("Image sampling not supported on CPU");
-        }
-        #[cfg(target_arch = "spirv")]
         unsafe {
             let mut result = Default::default();
             asm!(
@@ -66,14 +60,8 @@ pub struct Image2dArray {
 }
 
 impl Image2dArray {
+    #[spirv_std_macros::gpu_only]
     pub fn sample(&self, sampler: Sampler, coord: Vec3A) -> Vec4 {
-        #[cfg(not(target_arch = "spirv"))]
-        {
-            let _ = sampler;
-            let _ = coord;
-            panic!("Image sampling not supported on CPU");
-        }
-        #[cfg(target_arch = "spirv")]
         unsafe {
             let mut result = Default::default();
             asm!(
@@ -101,13 +89,8 @@ pub struct SampledImage<I> {
 }
 
 impl SampledImage<Image2d> {
+    #[spirv_std_macros::gpu_only]
     pub fn sample(&self, coord: Vec2) -> Vec4 {
-        #[cfg(not(target_arch = "spirv"))]
-        {
-            let _ = coord;
-            panic!("Image sampling not supported on CPU");
-        }
-        #[cfg(target_arch = "spirv")]
         unsafe {
             let mut result = Default::default();
             asm!(