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

This commit is contained in:
Sylvester Hesp 2023-04-21 15:00:19 +02:00
parent c298bf2c56
commit 43a267bbb8
9 changed files with 286 additions and 20 deletions

View File

@ -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()
}

View File

@ -1135,6 +1135,7 @@ 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>,
@ -1142,6 +1143,64 @@ pub trait ImageWithMethods<
) -> 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

View File

@ -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,
}

View File

@ -9,9 +9,9 @@ error[E0277]: the trait bound `Image<f32, 0, 2, 0, 0, 1, 0, 4>: HasGather` is no
Image<SampledType, 3, DEPTH, ARRAYED, 0, SAMPLED, FORMAT, COMPONENTS>
Image<SampledType, 4, DEPTH, ARRAYED, 0, SAMPLED, FORMAT, COMPONENTS>
note: required by a bound in `Image::<SampledType, DIM, DEPTH, ARRAYED, spirv_std::::image::{impl#1}::{constant#0}, SAMPLED, FORMAT, COMPONENTS>::gather`
--> $SPIRV_STD_SRC/image.rs:195:15
--> $SPIRV_STD_SRC/image.rs:200:15
|
195 | Self: HasGather,
200 | Self: HasGather,
| ^^^^^^^^^ required by this bound in `Image::<SampledType, DIM, DEPTH, ARRAYED, spirv_std::::image::{impl#1}::{constant#0}, SAMPLED, FORMAT, COMPONENTS>::gather`
error[E0277]: the trait bound `Image<f32, 2, 2, 0, 0, 1, 0, 4>: HasGather` is not satisfied
@ -25,9 +25,9 @@ error[E0277]: the trait bound `Image<f32, 2, 2, 0, 0, 1, 0, 4>: HasGather` is no
Image<SampledType, 3, DEPTH, ARRAYED, 0, SAMPLED, FORMAT, COMPONENTS>
Image<SampledType, 4, DEPTH, ARRAYED, 0, SAMPLED, FORMAT, COMPONENTS>
note: required by a bound in `Image::<SampledType, DIM, DEPTH, ARRAYED, spirv_std::::image::{impl#1}::{constant#0}, SAMPLED, FORMAT, COMPONENTS>::gather`
--> $SPIRV_STD_SRC/image.rs:195:15
--> $SPIRV_STD_SRC/image.rs:200:15
|
195 | Self: HasGather,
200 | Self: HasGather,
| ^^^^^^^^^ required by this bound in `Image::<SampledType, DIM, DEPTH, ARRAYED, spirv_std::::image::{impl#1}::{constant#0}, SAMPLED, FORMAT, COMPONENTS>::gather`
error: aborting due to 2 previous errors

View File

@ -0,0 +1,17 @@
// build-pass
use spirv_std::spirv;
use spirv_std::{arch, image::sample, image::ImageWithMethods, Image, Sampler};
#[spirv(fragment)]
pub fn main(
#[spirv(descriptor_set = 0, binding = 0)] sampler: &Sampler,
#[spirv(descriptor_set = 0, binding = 1)] image1: &Image!(2D, type=f32, sampled, multisampled),
#[spirv(descriptor_set = 0, binding = 2)] image2: &Image!(2D, type=f32, sampled),
output: &mut glam::Vec4,
) {
let t1 = image1.fetch_with(glam::IVec2::new(0, 0), sample::sample_index(1));
let t2 = image2.sample_with(*sampler, glam::Vec2::new(0.5, 0.5), sample::bias(1.0));
let t3 = image2.sample_with(*sampler, glam::Vec2::new(0.5, 0.5), sample::lod(2.0));
*output = t1 + t2 + t3;
}

View File

@ -10,9 +10,9 @@ error[E0277]: the trait bound `Image<f32, 4, 2, 0, 0, 1, 0, 4>: HasQueryLevels`
Image<SampledType, 2, DEPTH, ARRAYED, MULTISAMPLED, SAMPLED, FORMAT, COMPONENTS>
Image<SampledType, 3, DEPTH, ARRAYED, MULTISAMPLED, SAMPLED, FORMAT, COMPONENTS>
note: required by a bound in `Image::<SampledType, DIM, DEPTH, ARRAYED, MULTISAMPLED, SAMPLED, FORMAT, COMPONENTS>::query_levels`
--> $SPIRV_STD_SRC/image.rs:879:15
--> $SPIRV_STD_SRC/image.rs:884:15
|
879 | Self: HasQueryLevels,
884 | Self: HasQueryLevels,
| ^^^^^^^^^^^^^^ required by this bound in `Image::<SampledType, DIM, DEPTH, ARRAYED, MULTISAMPLED, SAMPLED, FORMAT, COMPONENTS>::query_levels`
error: aborting due to previous error

View File

@ -10,9 +10,9 @@ error[E0277]: the trait bound `Image<f32, 4, 2, 0, 0, 1, 0, 4>: HasQueryLevels`
Image<SampledType, 2, DEPTH, ARRAYED, MULTISAMPLED, SAMPLED, FORMAT, COMPONENTS>
Image<SampledType, 3, DEPTH, ARRAYED, MULTISAMPLED, SAMPLED, FORMAT, COMPONENTS>
note: required by a bound in `Image::<SampledType, DIM, DEPTH, ARRAYED, MULTISAMPLED, SAMPLED, FORMAT, COMPONENTS>::query_lod`
--> $SPIRV_STD_SRC/image.rs:905:15
--> $SPIRV_STD_SRC/image.rs:910:15
|
905 | Self: HasQueryLevels,
910 | Self: HasQueryLevels,
| ^^^^^^^^^^^^^^ required by this bound in `Image::<SampledType, DIM, DEPTH, ARRAYED, MULTISAMPLED, SAMPLED, FORMAT, COMPONENTS>::query_lod`
error: aborting due to previous error

View File

@ -15,9 +15,9 @@ error[E0277]: the trait bound `Image<f32, 1, 2, 0, 0, 1, 0, 4>: HasQuerySize` is
Image<SampledType, 2, DEPTH, ARRAYED, 0, 2, FORMAT, COMPONENTS>
and 6 others
note: required by a bound in `Image::<SampledType, DIM, DEPTH, ARRAYED, MULTISAMPLED, SAMPLED, FORMAT, COMPONENTS>::query_size`
--> $SPIRV_STD_SRC/image.rs:936:15
--> $SPIRV_STD_SRC/image.rs:941:15
|
936 | Self: HasQuerySize,
941 | Self: HasQuerySize,
| ^^^^^^^^^^^^ required by this bound in `Image::<SampledType, DIM, DEPTH, ARRAYED, MULTISAMPLED, SAMPLED, FORMAT, COMPONENTS>::query_size`
error: aborting due to previous error

View File

@ -10,9 +10,9 @@ error[E0277]: the trait bound `Image<f32, 4, 2, 0, 0, 1, 0, 4>: HasQuerySizeLod`
Image<SampledType, 2, DEPTH, ARRAYED, 0, SAMPLED, FORMAT, COMPONENTS>
Image<SampledType, 3, DEPTH, ARRAYED, 0, SAMPLED, FORMAT, COMPONENTS>
note: required by a bound in `Image::<SampledType, DIM, DEPTH, ARRAYED, spirv_std::::image::{impl#7}::{constant#0}, SAMPLED, FORMAT, COMPONENTS>::query_size_lod`
--> $SPIRV_STD_SRC/image.rs:980:15
--> $SPIRV_STD_SRC/image.rs:985:15
|
980 | Self: HasQuerySizeLod,
985 | Self: HasQuerySizeLod,
| ^^^^^^^^^^^^^^^ required by this bound in `Image::<SampledType, DIM, DEPTH, ARRAYED, spirv_std::::image::{impl#7}::{constant#0}, SAMPLED, FORMAT, COMPONENTS>::query_size_lod`
error: aborting due to previous error