Added all the other sampling functions. Also some fixes and cleanup

Author: oisyn

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

@@ -630,6 +630,7 @@ const SAMPLE_PARAM_COUNT: usize = 3;
 const SAMPLE_PARAM_TYPES: [&str; SAMPLE_PARAM_COUNT] = ["B", "L", "S"];
 const SAMPLE_PARAM_OPERANDS: [&str; SAMPLE_PARAM_COUNT] = ["Bias", "Lod", "Sample"];
 const SAMPLE_PARAM_NAMES: [&str; SAMPLE_PARAM_COUNT] = ["bias", "lod", "sample_index"];
+const SAMPLE_PARAM_EXPLICIT_LOD_MASK: usize = 0b010; // which params require the use of ExplicitLod rather than ImplicitLod
 
 struct SampleImplRewriter(usize, syn::Type);
 
@@ -653,9 +654,11 @@ impl SampleImplRewriter {
             }
             ty_str.push(',');
         }
-        ty_str.push_str(">");
+        ty_str.push('>');
         let ty: syn::Type = syn::parse(ty_str.parse().unwrap()).unwrap();
 
+        // use the type to insert it into the generic argument of the trait we're implementing
+        // e.g., `ImageWithMethods<Dummy>` becomes `ImageWithMethods<SampleParams<SomeTy<B>, NoneTy, NoneTy>>`
         if let Some(t) = &mut new_impl.trait_ {
             if let syn::PathArguments::AngleBracketed(a) =
                 &mut t.1.segments.last_mut().unwrap().arguments
@@ -665,10 +668,13 @@ impl SampleImplRewriter {
                 }
             }
         }
+
+        // rewrite the implemented functions
         SampleImplRewriter(mask, ty).visit_item_impl_mut(&mut new_impl);
         new_impl
     }
 
+    // generates an operands string for use in the assembly, e.g. "Bias %bias Lod %lod", based on the mask
     fn get_operands(&self) -> String {
         let mut op = String::new();
         for i in 0..SAMPLE_PARAM_COUNT {
@@ -682,6 +688,7 @@ impl SampleImplRewriter {
         op
     }
 
+    // generates list of assembly loads for the data, e.g. "%bias = OpLoad _ {bias}", etc.
     fn add_loads(&self, t: &mut Vec<TokenTree>) {
         for i in 0..SAMPLE_PARAM_COUNT {
             if self.0 & (1 << i) != 0 {
@@ -695,10 +702,11 @@ impl SampleImplRewriter {
         }
     }
 
+    // generates list of register specifications, e.g. `bias = in(reg) &params.bias.0, ...` as separate tokens
     fn add_regs(&self, t: &mut Vec<TokenTree>) {
         for i in 0..SAMPLE_PARAM_COUNT {
             if self.0 & (1 << i) != 0 {
-                let s = format!("{0} = in(reg) &params.{0},", SAMPLE_PARAM_NAMES[i]);
+                let s = format!("{0} = in(reg) &params.{0}.0,", SAMPLE_PARAM_NAMES[i]);
                 let ts: proc_macro2::TokenStream = s.parse().unwrap();
                 t.extend(ts);
             }
@@ -708,6 +716,7 @@ impl SampleImplRewriter {
 
 impl VisitMut for SampleImplRewriter {
     fn visit_impl_item_method_mut(&mut self, item: &mut syn::ImplItemMethod) {
+        // rewrite the last parameter of this method to be of type `SampleParams<...>` we generated earlier
         if let Some(syn::FnArg::Typed(p)) = item.sig.inputs.last_mut() {
             *p.ty.as_mut() = self.1.clone();
         }
@@ -716,6 +725,7 @@ impl VisitMut for SampleImplRewriter {
 
     fn visit_macro_mut(&mut self, m: &mut syn::Macro) {
         if m.path.is_ident("asm") {
+            // this is where the asm! block is manipulated
             let t = m.tokens.clone();
             let mut new_t = Vec::new();
             let mut altered = false;
@@ -724,11 +734,21 @@ impl VisitMut for SampleImplRewriter {
                 match tt {
                     TokenTree::Literal(l) => {
                         if let Ok(l) = syn::parse::<syn::LitStr>(l.to_token_stream().into()) {
+                            // found a string literal
                             let s = l.value();
                             if s.contains("$PARAMS") {
                                 altered = true;
+                                // add load instructions before the sampling instruction
                                 self.add_loads(&mut new_t);
+                                // and insert image operands
                                 let s = s.replace("$PARAMS", &self.get_operands());
+                                let lod_type = if self.0 & SAMPLE_PARAM_EXPLICIT_LOD_MASK != 0 {
+                                    "ExplicitLod"
+                                } else {
+                                    "ImplicitLod "
+                                };
+                                let s = s.replace("$LOD", lod_type);
+
                                 new_t.push(TokenTree::Literal(proc_macro2::Literal::string(
                                     s.as_str(),
                                 )));
@@ -746,18 +766,21 @@ impl VisitMut for SampleImplRewriter {
             }
 
             if altered {
+                // finally, add register specs
                 self.add_regs(&mut new_t);
             }
 
+            // replace all tokens within the asm! block with our new list
             m.tokens = new_t.into_iter().collect();
         }
     }
 }
 
-/// Generates permutations of a sampling function containing inline assembly with an image
-/// instruction ending with a placeholder `$PARAMS` operand. The last parameter must be named
-/// `params`, its type will be rewritten. Relevant generic arguments are added to the function
-/// signature. See `SAMPLE_PARAM_TYPES` for a list of names you cannot use as generic arguments.
+/// Generates permutations of an `ImageWithMethods` implementation containing sampling functions
+/// that have asm instruction ending with a placeholder `$PARAMS` operand. The last parameter
+/// of each function must be named `params`, its type will be rewritten. Relevant generic
+/// arguments are added to the impl generics.
+/// See `SAMPLE_PARAM_TYPES` for a list of names you cannot use as generic arguments.
 #[proc_macro_attribute]
 #[doc(hidden)]
 pub fn gen_sample_param_permutations(_attr: TokenStream, item: TokenStream) -> TokenStream {
@@ -768,6 +791,7 @@ pub fn gen_sample_param_permutations(_attr: TokenStream, item: TokenStream) -> T
         fns.push(SampleImplRewriter::rewrite(m, &item_impl));
     }
 
-    println!("{}", quote! { #(#fns)* }.to_string());
+    // uncomment to output generated tokenstream to stdout
+    // println!("{}", quote! { #(#fns)* }.to_string());
     quote! { #(#fns)* }.into()
 }

crates/spirv-std/src/image.rs

@@ -1135,13 +1135,72 @@ pub trait ImageWithMethods<
 >
 {
     /// Fetch a single texel with a sampler set at compile time
+    #[doc(alias = "OpImageFetch")]
     fn fetch_with<I>(
         &self,
         coordinate: impl ImageCoordinate<I, DIM, ARRAYED>,
         params: Params,
     ) -> SampledType::SampleResult
     where
         I: Integer;
+
+    /// Gathers the requested component from four texels.
+    #[doc(alias = "OpImageGather")]
+    fn gather_with<F>(
+        &self,
+        sampler: Sampler,
+        coordinate: impl ImageCoordinate<F, DIM, ARRAYED>,
+        component: u32,
+        params: Params,
+    ) -> SampledType::Vec4
+    where
+        Self: HasGather,
+        F: Float;
+
+    /// Sample texels at `coord` from the image using `sampler`.
+    fn sample_with<F>(
+        &self,
+        sampler: Sampler,
+        coord: impl ImageCoordinate<F, DIM, ARRAYED>,
+        params: Params,
+    ) -> SampledType::SampleResult
+    where
+        F: Float;
+
+    /// Sample the image's depth reference
+    #[doc(alias = "OpImageSampleDrefImplicitLod")]
+    fn sample_depth_reference_with<F>(
+        &self,
+        sampler: Sampler,
+        coordinate: impl ImageCoordinate<F, DIM, ARRAYED>,
+        depth_reference: f32,
+        params: Params,
+    ) -> SampledType
+    where
+        F: Float;
+
+    /// Sample the image with a project coordinate
+    #[doc(alias = "OpImageSampleProjImplicitLod")]
+    fn sample_with_project_coordinate_with<F>(
+        &self,
+        sampler: Sampler,
+        project_coordinate: impl ImageCoordinate<F, DIM, { Arrayed::True as u32 }>,
+        params: Params,
+    ) -> SampledType::SampleResult
+    where
+        F: Float;
+
+    /// Sample the image's depth reference with the project coordinate
+    #[doc(alias = "OpImageSampleProjDrefImplicitLod")]
+    fn sample_depth_reference_with_project_coordinate_with<F>(
+        &self,
+        sampler: Sampler,
+        project_coordinate: impl ImageCoordinate<F, DIM, { Arrayed::True as u32 }>,
+        depth_reference: f32,
+        params: Params,
+    ) -> SampledType
+    where
+        F: Float;
 }
 
 #[crate::macros::gen_sample_param_permutations]
@@ -1191,6 +1250,167 @@ impl<
         }
         result.truncate_into()
     }
+
+    /// Gathers the requested component from four texels.
+    #[crate::macros::gpu_only]
+    #[doc(alias = "OpImageGather")]
+    #[inline]
+    fn gather_with<F>(
+        &self,
+        sampler: Sampler,
+        coordinate: impl ImageCoordinate<F, DIM, ARRAYED>,
+        component: u32,
+        params: SampleParams,
+    ) -> SampledType::Vec4
+    where
+        Self: HasGather,
+        F: Float,
+    {
+        let mut result = SampledType::Vec4::default();
+        unsafe {
+            asm! {
+                "%typeSampledImage = OpTypeSampledImage typeof*{this}",
+                "%image = OpLoad _ {this}",
+                "%sampler = OpLoad _ {sampler}",
+                "%coordinate = OpLoad _ {coordinate}",
+                "%sampledImage = OpSampledImage %typeSampledImage %image %sampler",
+                "%result = OpImageGather typeof*{result} %sampledImage %coordinate {component} $PARAMS",
+                "OpStore {result} %result",
+                result = in(reg) &mut result,
+                this = in(reg) self,
+                sampler = in(reg) &sampler,
+                coordinate = in(reg) &coordinate,
+                component = in(reg) component,
+            }
+        }
+        result
+    }
+
+    /// Sample texels at `coord` from the image using `sampler`.
+    #[crate::macros::gpu_only]
+    fn sample_with<F>(
+        &self,
+        sampler: Sampler,
+        coord: impl ImageCoordinate<F, DIM, ARRAYED>,
+        params: SampleParams,
+    ) -> SampledType::SampleResult
+    where
+        F: Float,
+    {
+        unsafe {
+            let mut result = SampledType::Vec4::default();
+            asm!(
+                "%typeSampledImage = OpTypeSampledImage typeof*{this}",
+                "%image = OpLoad _ {this}",
+                "%sampler = OpLoad _ {sampler}",
+                "%coord = OpLoad _ {coord}",
+                "%sampledImage = OpSampledImage %typeSampledImage %image %sampler",
+                "%result = OpImageSample$LOD typeof*{result} %sampledImage %coord $PARAMS",
+                "OpStore {result} %result",
+                result = in(reg) &mut result,
+                this = in(reg) self,
+                sampler = in(reg) &sampler,
+                coord = in(reg) &coord,
+            );
+            result.truncate_into()
+        }
+    }
+
+    /// Sample the image's depth reference
+    #[crate::macros::gpu_only]
+    #[doc(alias = "OpImageSampleDrefImplicitLod")]
+    fn sample_depth_reference_with<F>(
+        &self,
+        sampler: Sampler,
+        coordinate: impl ImageCoordinate<F, DIM, ARRAYED>,
+        depth_reference: f32,
+        params: SampleParams,
+    ) -> SampledType
+    where
+        F: Float,
+    {
+        let mut result = Default::default();
+        unsafe {
+            asm!(
+                "%image = OpLoad _ {this}",
+                "%sampler = OpLoad _ {sampler}",
+                "%coordinate = OpLoad _ {coordinate}",
+                "%depth_reference = OpLoad _ {depth_reference}", // not required to do this way, but done for consistency
+                "%sampledImage = OpSampledImage _ %image %sampler",
+                "%result = OpImageSampleDref$LOD _ %sampledImage %coordinate %depth_reference $PARAMS",
+                "OpStore {result} %result",
+                result = in(reg) &mut result,
+                this = in(reg) self,
+                sampler = in(reg) &sampler,
+                coordinate = in(reg) &coordinate,
+                depth_reference = in(reg) &depth_reference,
+            );
+        }
+        result
+    }
+
+    /// Sample the image with a project coordinate
+    #[crate::macros::gpu_only]
+    #[doc(alias = "OpImageSampleProjImplicitLod")]
+    fn sample_with_project_coordinate_with<F>(
+        &self,
+        sampler: Sampler,
+        project_coordinate: impl ImageCoordinate<F, DIM, { Arrayed::True as u32 }>,
+        params: SampleParams,
+    ) -> SampledType::SampleResult
+    where
+        F: Float,
+    {
+        unsafe {
+            let mut result = SampledType::Vec4::default();
+            asm!(
+                "%image = OpLoad _ {this}",
+                "%sampler = OpLoad _ {sampler}",
+                "%project_coordinate = OpLoad _ {project_coordinate}",
+                "%sampledImage = OpSampledImage _ %image %sampler",
+                "%result = OpImageSampleProj$LOD _ %sampledImage %project_coordinate $PARAMS",
+                "OpStore {result} %result",
+                result = in(reg) &mut result,
+                this = in(reg) self,
+                sampler = in(reg) &sampler,
+                project_coordinate = in(reg) &project_coordinate,
+            );
+            result.truncate_into()
+        }
+    }
+
+    /// Sample the image's depth reference with the project coordinate
+    #[crate::macros::gpu_only]
+    #[doc(alias = "OpImageSampleProjDrefImplicitLod")]
+    fn sample_depth_reference_with_project_coordinate_with<F>(
+        &self,
+        sampler: Sampler,
+        project_coordinate: impl ImageCoordinate<F, DIM, { Arrayed::True as u32 }>,
+        depth_reference: f32,
+        params: SampleParams,
+    ) -> SampledType
+    where
+        F: Float,
+    {
+        let mut result = Default::default();
+        unsafe {
+            asm!(
+                "%image = OpLoad _ {this}",
+                "%sampler = OpLoad _ {sampler}",
+                "%project_coordinate = OpLoad _ {project_coordinate}",
+                "%depth_reference = OpLoad _ {depth_reference}", // not required to do this way, but done for consistency
+                "%sampledImage = OpSampledImage _ %image %sampler",
+                "%result = OpImageSampleProjDref$LOD _ %sampledImage %project_coordinate %depth_reference $PARAMS",
+                "OpStore {result} %result",
+                result = in(reg) &mut result,
+                this = in(reg) self,
+                sampler = in(reg) &sampler,
+                project_coordinate = in(reg) &project_coordinate,
+                depth_reference = in(reg) &depth_reference,
+            );
+        }
+        result
+    }
 }
 
 /// This is a marker trait to represent the constraints on `OpImageGather` too complex to be

crates/spirv-std/src/image/sample.rs

@@ -15,14 +15,19 @@ impl<T> OptionTy for SomeTy<T> {
 pub struct NoneTy;
 
 /// Helper struct that denotes that the type does exist and is of type T, analog to `Option::Some(T)`
-pub struct SomeTy<T>(T);
+pub struct SomeTy<T>(pub T);
 
 /// Helper struct that allows building image operands. Start with a global function that returns this
 /// struct, and then chain additional calls.
 /// Example: `image.sample_with(coords, params::bias(3.0).sample_index(1))`
 pub struct SampleParams<B: OptionTy, L: OptionTy, S: OptionTy> {
+    /// 'Bias' image operand
     pub bias: B,
+
+    /// 'Lod' image operand
     pub lod: L,
+
+    /// 'Sample' image operand
     pub sample_index: S,
 }