diff --git a/node-graph/gcore-shaders/src/color/color_types.rs b/node-graph/gcore-shaders/src/color/color_types.rs index 2cb4d231b..63517341e 100644 --- a/node-graph/gcore-shaders/src/color/color_types.rs +++ b/node-graph/gcore-shaders/src/color/color_types.rs @@ -3,6 +3,7 @@ use super::discrete_srgb::{float_to_srgb_u8, srgb_u8_to_float}; use bytemuck::{Pod, Zeroable}; use core::fmt::Debug; use core::hash::Hash; +use glam::Vec4; use half::f16; #[cfg(not(feature = "std"))] use num_traits::Euclid; @@ -1075,6 +1076,21 @@ impl Color { ..*self } } + + #[inline(always)] + pub const fn from_vec4(vec: Vec4) -> Self { + Self { + red: vec.x, + green: vec.y, + blue: vec.z, + alpha: vec.w, + } + } + + #[inline(always)] + pub fn to_vec4(&self) -> Vec4 { + Vec4::new(self.red, self.green, self.blue, self.alpha) + } } #[cfg(test)] diff --git a/node-graph/graster-nodes/src/adjustments.rs b/node-graph/graster-nodes/src/adjustments.rs index b98c927bd..afbee0cbb 100644 --- a/node-graph/graster-nodes/src/adjustments.rs +++ b/node-graph/graster-nodes/src/adjustments.rs @@ -52,6 +52,7 @@ fn luminance>( Table, GradientStops, )] + #[gpu_image] mut input: T, luminance_calc: LuminanceCalculation, ) -> T { @@ -77,6 +78,7 @@ fn gamma_correction>( Table, GradientStops, )] + #[gpu_image] mut input: T, #[default(2.2)] #[range((0.01, 10.))] @@ -98,6 +100,7 @@ fn extract_channel>( Table, GradientStops, )] + #[gpu_image] mut input: T, channel: RedGreenBlueAlpha, ) -> T { @@ -122,6 +125,7 @@ fn make_opaque>( Table, GradientStops, )] + #[gpu_image] mut input: T, ) -> T { input.adjust(|color| { @@ -148,6 +152,7 @@ fn brightness_contrast>( Table, GradientStops, )] + #[gpu_image] mut input: T, brightness: SignedPercentage, contrast: SignedPercentage, @@ -238,6 +243,7 @@ fn levels>( Table, GradientStops, )] + #[gpu_image] mut image: T, #[default(0.)] shadows: Percentage, #[default(50.)] midtones: Percentage, @@ -306,6 +312,7 @@ fn black_and_white>( Table, GradientStops, )] + #[gpu_image] mut image: T, #[default(Color::BLACK)] tint: Color, #[default(40.)] @@ -379,6 +386,7 @@ fn hue_saturation>( Table, GradientStops, )] + #[gpu_image] mut input: T, hue_shift: Angle, saturation_shift: SignedPercentage, @@ -414,6 +422,7 @@ fn invert>( Table, GradientStops, )] + #[gpu_image] mut input: T, ) -> T { input.adjust(|color| { @@ -437,6 +446,7 @@ fn threshold>( Table, GradientStops, )] + #[gpu_image] mut image: T, #[default(50.)] min_luminance: Percentage, #[default(100.)] max_luminance: Percentage, @@ -483,6 +493,7 @@ fn vibrance>( Table, GradientStops, )] + #[gpu_image] mut image: T, vibrance: SignedPercentage, ) -> T { @@ -649,6 +660,7 @@ fn channel_mixer>( Table, GradientStops, )] + #[gpu_image] mut image: T, monochrome: bool, @@ -778,6 +790,7 @@ fn selective_color>( Table, GradientStops, )] + #[gpu_image] mut image: T, mode: RelativeAbsolute, @@ -921,6 +934,7 @@ fn posterize>( Table, GradientStops, )] + #[gpu_image] mut input: T, #[default(4)] #[hard_min(2.)] @@ -955,6 +969,7 @@ fn exposure>( Table, GradientStops, )] + #[gpu_image] mut input: T, exposure: f64, offset: f64, diff --git a/node-graph/graster-nodes/src/blending_nodes.rs b/node-graph/graster-nodes/src/blending_nodes.rs index 4036244ae..31ca3f469 100644 --- a/node-graph/graster-nodes/src/blending_nodes.rs +++ b/node-graph/graster-nodes/src/blending_nodes.rs @@ -141,6 +141,7 @@ fn blend + Send>( Table, GradientStops, )] + #[gpu_image] over: T, #[expose] #[implementations( @@ -149,6 +150,7 @@ fn blend + Send>( Table, GradientStops, )] + #[gpu_image] under: T, blend_mode: BlendMode, #[default(100.)] opacity: Percentage, @@ -165,6 +167,7 @@ fn color_overlay>( Table, GradientStops, )] + #[gpu_image] mut image: T, #[default(Color::BLACK)] color: Color, blend_mode: BlendMode, diff --git a/node-graph/node-macro/src/codegen.rs b/node-graph/node-macro/src/codegen.rs index 57a8b75cf..1c0c4f651 100644 --- a/node-graph/node-macro/src/codegen.rs +++ b/node-graph/node-macro/src/codegen.rs @@ -1,7 +1,7 @@ use crate::parsing::*; use convert_case::{Case, Casing}; use proc_macro_crate::FoundCrate; -use proc_macro2::TokenStream as TokenStream2; +use proc_macro2::{TokenStream as TokenStream2, TokenStream}; use quote::{ToTokens, format_ident, quote, quote_spanned}; use std::sync::atomic::AtomicU64; use syn::punctuated::Punctuated; @@ -295,6 +295,7 @@ pub(crate) fn generate_node_code(parsed: &ParsedNodeFn) -> syn::Result syn::Result, pub number_mode_range: Option, pub implementations: Punctuated, + pub gpu_image: bool, } +/// a param of `impl Node` with `#[implementation(in -> out)]` #[derive(Clone, Debug)] pub struct NodeParsedField { pub input_type: Type, @@ -529,6 +533,7 @@ fn parse_field(pat_ident: PatIdent, ty: Type, attrs: &[Attribute]) -> syn::Resul .map_err(|e| Error::new_spanned(attr, format!("Invalid `step` for argument '{ident}': {e}\nUSAGE EXAMPLE: #[step(2.)]"))) }) .transpose()?; + let gpu_image = extract_attribute(attrs, "gpu_image").is_some(); let (is_node, node_input_type, node_output_type) = parse_node_type(&ty); let description = attrs @@ -590,6 +595,7 @@ fn parse_field(pat_ident: PatIdent, ty: Type, attrs: &[Attribute]) -> syn::Resul ty, value_source, implementations, + gpu_image, }), name, description, diff --git a/node-graph/node-macro/src/shader_nodes.rs b/node-graph/node-macro/src/shader_nodes/mod.rs similarity index 50% rename from node-graph/node-macro/src/shader_nodes.rs rename to node-graph/node-macro/src/shader_nodes/mod.rs index 919d3ef87..0720869d0 100644 --- a/node-graph/node-macro/src/shader_nodes.rs +++ b/node-graph/node-macro/src/shader_nodes/mod.rs @@ -1,10 +1,13 @@ -use crate::parsing::NodeFnAttributes; +use crate::parsing::{NodeFnAttributes, ParsedNodeFn}; +use crate::shader_nodes::per_pixel_adjust::PerPixelAdjust; use proc_macro2::{Ident, TokenStream}; use quote::quote; -use strum::{EnumString, VariantNames}; +use strum::VariantNames; use syn::Error; use syn::parse::{Parse, ParseStream}; +pub mod per_pixel_adjust; + pub const STD_FEATURE_GATE: &str = "std"; pub fn modify_cfg(attributes: &NodeFnAttributes) -> TokenStream { @@ -16,17 +19,33 @@ pub fn modify_cfg(attributes: &NodeFnAttributes) -> TokenStream { } } -#[derive(Debug, EnumString, VariantNames)] +#[derive(Debug, VariantNames)] pub(crate) enum ShaderNodeType { - PerPixelAdjust, + PerPixelAdjust(PerPixelAdjust), } impl Parse for ShaderNodeType { fn parse(input: ParseStream) -> syn::Result { let ident: Ident = input.parse()?; Ok(match ident.to_string().as_str() { - "PerPixelAdjust" => ShaderNodeType::PerPixelAdjust, + "PerPixelAdjust" => ShaderNodeType::PerPixelAdjust(PerPixelAdjust::parse(input)?), _ => return Err(Error::new_spanned(&ident, format!("attr 'shader_node' must be one of {:?}", Self::VARIANTS))), }) } } + +pub trait CodegenShaderEntryPoint { + fn codegen_shader_entry_point(&self, parsed: &ParsedNodeFn) -> syn::Result; +} + +impl CodegenShaderEntryPoint for ShaderNodeType { + fn codegen_shader_entry_point(&self, parsed: &ParsedNodeFn) -> syn::Result { + if parsed.is_async { + return Err(Error::new_spanned(&parsed.fn_name, "Shader nodes must not be async")); + } + + match self { + ShaderNodeType::PerPixelAdjust(x) => x.codegen_shader_entry_point(parsed), + } + } +} diff --git a/node-graph/node-macro/src/shader_nodes/per_pixel_adjust.rs b/node-graph/node-macro/src/shader_nodes/per_pixel_adjust.rs new file mode 100644 index 000000000..0e220c8aa --- /dev/null +++ b/node-graph/node-macro/src/shader_nodes/per_pixel_adjust.rs @@ -0,0 +1,110 @@ +use crate::parsing::{ParsedFieldType, ParsedNodeFn, RegularParsedField}; +use crate::shader_nodes::CodegenShaderEntryPoint; +use proc_macro2::{Ident, TokenStream}; +use quote::{ToTokens, format_ident, quote}; +use std::borrow::Cow; +use syn::parse::{Parse, ParseStream}; + +#[derive(Debug)] +pub struct PerPixelAdjust {} + +impl Parse for PerPixelAdjust { + fn parse(_input: ParseStream) -> syn::Result { + Ok(Self {}) + } +} + +impl CodegenShaderEntryPoint for PerPixelAdjust { + fn codegen_shader_entry_point(&self, parsed: &ParsedNodeFn) -> syn::Result { + let fn_name = &parsed.fn_name; + let gpu_mod = format_ident!("{}_gpu", parsed.fn_name); + let spirv_image_ty = quote!(Image2d); + + // bindings for images start at 1 + let mut binding_cnt = 0; + let params = parsed + .fields + .iter() + .map(|f| { + let ident = &f.pat_ident; + match &f.ty { + ParsedFieldType::Node { .. } => Err(syn::Error::new_spanned(ident, "PerPixelAdjust shader nodes cannot accept other nodes as generics")), + ParsedFieldType::Regular(RegularParsedField { gpu_image: false, ty, .. }) => Ok(Param { + ident: Cow::Borrowed(&ident.ident), + ty: Cow::Owned(ty.to_token_stream()), + param_type: ParamType::Uniform, + }), + ParsedFieldType::Regular(RegularParsedField { gpu_image: true, .. }) => { + binding_cnt += 1; + Ok(Param { + ident: Cow::Owned(format_ident!("image_{}", &ident.ident)), + ty: Cow::Borrowed(&spirv_image_ty), + param_type: ParamType::Image { binding: binding_cnt }, + }) + } + } + }) + .collect::>>()?; + + let uniform_members = params + .iter() + .filter_map(|Param { ident, ty, param_type }| match param_type { + ParamType::Image { .. } => None, + ParamType::Uniform => Some(quote! {#ident: #ty}), + }) + .collect::>(); + let image_params = params + .iter() + .filter_map(|Param { ident, ty, param_type }| match param_type { + ParamType::Image { binding } => Some(quote! {#[spirv(descriptor_set = 0, binding = #binding)] #ident: &#ty}), + ParamType::Uniform => None, + }) + .collect::>(); + let call_args = params + .iter() + .map(|Param { ident, param_type, .. }| match param_type { + ParamType::Image { .. } => quote!(Color::from_vec4(#ident.fetch_with(texel_coord, lod(0)))), + ParamType::Uniform => quote!(uniform.#ident), + }) + .collect::>(); + let context = quote!(()); + + Ok(quote! { + pub mod #gpu_mod { + use super::*; + use graphene_core_shaders::color::Color; + use spirv_std::spirv; + use spirv_std::glam::{Vec4, Vec4Swizzles}; + use spirv_std::image::{Image2d, ImageWithMethods}; + use spirv_std::image::sample_with::lod; + + pub struct Uniform { + #(#uniform_members),* + } + + #[spirv(fragment)] + pub fn entry_point( + #[spirv(frag_coord)] frag_coord: Vec4, + color_out: &mut Vec4, + #[spirv(descriptor_set = 0, binding = 0, storage_buffer)] uniform: &Uniform, + #(#image_params),* + ) { + let texel_coord = frag_coord.xy().as_uvec2(); + let color: Color = #fn_name(#context, #(#call_args),*); + *color_out = color.to_vec4(); + } + } + }) + } +} + +struct Param<'a> { + ident: Cow<'a, Ident>, + ty: Cow<'a, TokenStream>, + param_type: ParamType, +} + +enum ParamType { + Image { binding: u32 }, + Uniform, +}