Added Grad support to SampleParams

This commit is contained in:
Sylvester Hesp 2023-04-21 17:22:00 +02:00
parent 4edad8817c
commit 2fd0182352
3 changed files with 102 additions and 34 deletions

View File

@ -626,11 +626,17 @@ fn debug_printf_inner(input: DebugPrintfInput) -> TokenStream {
output.into()
}
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
const SAMPLE_PARAM_COUNT: usize = 4;
const SAMPLE_PARAM_GENERICS: [&str; SAMPLE_PARAM_COUNT] = ["B", "L", "G", "S"];
const SAMPLE_PARAM_TYPES: [&str; SAMPLE_PARAM_COUNT] = ["B", "L", "(G,G)", "S"];
const SAMPLE_PARAM_OPERANDS: [&str; SAMPLE_PARAM_COUNT] = ["Bias", "Lod", "Grad", "Sample"];
const SAMPLE_PARAM_NAMES: [&str; SAMPLE_PARAM_COUNT] = ["bias", "lod", "grad", "sample_index"];
const SAMPLE_PARAM_GRAD_INDEX: usize = 2; // Grad requires some special handling because it uses 2 arguments
const SAMPLE_PARAM_EXPLICIT_LOD_MASK: usize = 0b0110; // which params require the use of ExplicitLod rather than ImplicitLod
fn is_grad(i: usize) -> bool {
i == SAMPLE_PARAM_GRAD_INDEX
}
struct SampleImplRewriter(usize, syn::Type);
@ -644,7 +650,7 @@ impl SampleImplRewriter {
for i in 0..SAMPLE_PARAM_COUNT {
if mask & (1 << i) != 0 {
new_impl.generics.params.push(syn::GenericParam::Type(
syn::Ident::new(SAMPLE_PARAM_TYPES[i], Span::call_site()).into(),
syn::Ident::new(SAMPLE_PARAM_GENERICS[i], Span::call_site()).into(),
));
ty_str.push_str("SomeTy<");
ty_str.push_str(SAMPLE_PARAM_TYPES[i]);
@ -679,12 +685,16 @@ impl SampleImplRewriter {
let mut op = String::new();
for i in 0..SAMPLE_PARAM_COUNT {
if self.0 & (1 << i) != 0 {
if is_grad(i) {
op.push_str("Grad %grad_x %grad_y ");
} else {
op.push_str(SAMPLE_PARAM_OPERANDS[i]);
op.push_str(" %");
op.push_str(SAMPLE_PARAM_NAMES[i]);
op.push(' ');
}
}
}
op
}
@ -692,12 +702,29 @@ impl SampleImplRewriter {
fn add_loads(&self, t: &mut Vec<TokenTree>) {
for i in 0..SAMPLE_PARAM_COUNT {
if self.0 & (1 << i) != 0 {
if is_grad(i) {
t.push(TokenTree::Literal(proc_macro2::Literal::string(
"%grad_x = OpLoad _ {grad_x}",
)));
t.push(TokenTree::Punct(proc_macro2::Punct::new(
',',
proc_macro2::Spacing::Alone,
)));
t.push(TokenTree::Literal(proc_macro2::Literal::string(
"%grad_y = OpLoad _ {grad_y}",
)));
t.push(TokenTree::Punct(proc_macro2::Punct::new(
',',
proc_macro2::Spacing::Alone,
)));
} else {
let s = format!("%{0} = OpLoad _ {{{0}}}", SAMPLE_PARAM_NAMES[i]);
t.push(TokenTree::Literal(proc_macro2::Literal::string(s.as_str())));
t.push(TokenTree::Punct(proc_macro2::Punct::new(
',',
proc_macro2::Spacing::Alone,
)))
)));
}
}
}
}
@ -706,7 +733,11 @@ impl SampleImplRewriter {
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}.0,", SAMPLE_PARAM_NAMES[i]);
let s = if is_grad(i) {
String::from("grad_x=in(reg) &params.grad.0.0,grad_y=in(reg) &params.grad.0.1,")
} else {
format!("{0} = in(reg) &params.{0}.0,", SAMPLE_PARAM_NAMES[i])
};
let ts: proc_macro2::TokenStream = s.parse().unwrap();
t.extend(ts);
}
@ -780,7 +811,7 @@ impl VisitMut for SampleImplRewriter {
/// 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.
/// See `SAMPLE_PARAM_GENERICS` 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 {
@ -792,6 +823,6 @@ pub fn gen_sample_param_permutations(_attr: TokenStream, item: TokenStream) -> T
}
// uncomment to output generated tokenstream to stdout
// println!("{}", quote! { #(#fns)* }.to_string());
println!("{}", quote! { #(#fns)* }.to_string());
quote! { #(#fns)* }.into()
}

View File

@ -18,74 +18,106 @@ pub struct NoneTy;
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> {
/// struct, and then chain additional calls. No care is taken to avoid stating multiple operands that,
/// together, make no sense, such as Lod and Grad.
/// Example: `image.sample_with(coords, sample_with::bias(3.0).sample_index(1))`
pub struct SampleParams<B: OptionTy, L: OptionTy, G: OptionTy, S: OptionTy> {
/// 'Bias' image operand
pub bias: B,
/// 'Lod' image operand
pub lod: L,
/// 'Sample' image operand
/// 'Grad' image operand
pub grad: G,
/// 'Sample' image operandy
pub sample_index: S,
}
/// Sets the 'Bias' image operand
pub fn bias<B>(bias: B) -> SampleParams<SomeTy<B>, NoneTy, NoneTy> {
pub fn bias<B>(bias: B) -> SampleParams<SomeTy<B>, NoneTy, NoneTy, NoneTy> {
SampleParams {
bias: SomeTy(bias),
lod: NoneTy,
grad: NoneTy,
sample_index: NoneTy,
}
}
/// Sets the 'Lod' image operand
pub fn lod<L>(lod: L) -> SampleParams<NoneTy, SomeTy<L>, NoneTy> {
pub fn lod<L>(lod: L) -> SampleParams<NoneTy, SomeTy<L>, NoneTy, NoneTy> {
SampleParams {
bias: NoneTy,
lod: SomeTy(lod),
grad: NoneTy,
sample_index: NoneTy,
}
}
/// Sets the 'Grad' image operand
pub fn grad<T>(grad_x: T, grad_y: T) -> SampleParams<NoneTy, NoneTy, SomeTy<(T, T)>, NoneTy> {
SampleParams {
bias: NoneTy,
lod: NoneTy,
grad: SomeTy((grad_x, grad_y)),
sample_index: NoneTy,
}
}
/// Sets the 'Sample' image operand
pub fn sample_index<S>(sample_index: S) -> SampleParams<NoneTy, NoneTy, SomeTy<S>> {
pub fn sample_index<S>(sample_index: S) -> SampleParams<NoneTy, NoneTy, NoneTy, SomeTy<S>> {
SampleParams {
bias: NoneTy,
lod: NoneTy,
grad: NoneTy,
sample_index: SomeTy(sample_index),
}
}
impl<L: OptionTy, S: OptionTy> SampleParams<NoneTy, L, S> {
impl<L: OptionTy, G: OptionTy, S: OptionTy> SampleParams<NoneTy, L, G, S> {
/// Sets the 'Bias' image operand
pub fn bias<B>(self, bias: B) -> SampleParams<SomeTy<B>, L, S> {
pub fn bias<B>(self, bias: B) -> SampleParams<SomeTy<B>, L, G, S> {
SampleParams {
bias: SomeTy(bias),
lod: self.lod,
grad: self.grad,
sample_index: self.sample_index,
}
}
}
impl<B: OptionTy, S: OptionTy> SampleParams<B, NoneTy, S> {
impl<B: OptionTy, G: OptionTy, S: OptionTy> SampleParams<B, NoneTy, G, S> {
/// Sets the 'Lod' image operand
pub fn lod<L>(self, lod: L) -> SampleParams<B, SomeTy<L>, S> {
pub fn lod<L>(self, lod: L) -> SampleParams<B, SomeTy<L>, G, S> {
SampleParams {
bias: self.bias,
lod: SomeTy(lod),
grad: self.grad,
sample_index: self.sample_index,
}
}
}
impl<B: OptionTy, L: OptionTy> SampleParams<B, L, NoneTy> {
/// Sets the 'Sample' image operand
pub fn sample_index<S>(self, sample_index: S) -> SampleParams<B, L, SomeTy<S>> {
impl<B: OptionTy, L: OptionTy, S: OptionTy> SampleParams<B, L, NoneTy, S> {
/// Sets the 'Lod' image operand
pub fn grad<T>(self, grad_x: T, grad_y: T) -> SampleParams<B, L, SomeTy<(T, T)>, S> {
SampleParams {
bias: self.bias,
lod: self.lod,
grad: SomeTy((grad_x, grad_y)),
sample_index: self.sample_index,
}
}
}
impl<B: OptionTy, L: OptionTy, G: OptionTy> SampleParams<B, L, G, NoneTy> {
/// Sets the 'Sample' image operand
pub fn sample_index<S>(self, sample_index: S) -> SampleParams<B, L, G, SomeTy<S>> {
SampleParams {
bias: self.bias,
lod: self.lod,
grad: self.grad,
sample_index: SomeTy(sample_index),
}
}

View File

@ -13,5 +13,10 @@ pub fn main(
let t1 = image1.fetch_with(glam::IVec2::new(0, 0), sample_with::sample_index(1));
let t2 = image2.sample_with(*sampler, glam::Vec2::new(0.5, 0.5), sample_with::bias(1.0));
let t3 = image2.sample_with(*sampler, glam::Vec2::new(0.5, 0.5), sample_with::lod(2.0));
*output = t1 + t2 + t3;
let t4 = image2.sample_with(
*sampler,
glam::Vec2::new(0.5, 0.5),
sample_with::grad(glam::Vec2::new(0.5, 0.5), glam::Vec2::new(0.5, 0.5)),
);
*output = t1 + t2 + t3 + t4;
}