diff --git a/Cargo.lock b/Cargo.lock index 97edb1e0c7..7189e30817 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2191,6 +2191,7 @@ dependencies = [ "specta", "spirv-std", "tokio", + "wgpu-executor", ] [[package]] diff --git a/node-graph/gcore-shaders/src/blending.rs b/node-graph/gcore-shaders/src/blending.rs index c3701e2cc0..b305dd0910 100644 --- a/node-graph/gcore-shaders/src/blending.rs +++ b/node-graph/gcore-shaders/src/blending.rs @@ -66,7 +66,7 @@ impl AlphaBlending { } #[repr(i32)] -#[derive(Debug, Default, Clone, Copy, Eq, PartialEq, Hash)] +#[derive(Debug, Default, Clone, Copy, Eq, PartialEq, Hash, bytemuck::NoUninit)] #[cfg_attr(feature = "std", derive(dyn_any::DynAny, specta::Type, serde::Serialize, serde::Deserialize))] pub enum BlendMode { // Basic group diff --git a/node-graph/gcore/src/raster_types.rs b/node-graph/gcore/src/raster_types.rs index 97dd138145..7efae73fb8 100644 --- a/node-graph/gcore/src/raster_types.rs +++ b/node-graph/gcore/src/raster_types.rs @@ -137,7 +137,7 @@ mod gpu { #[derive(Clone, Debug, PartialEq, Hash)] pub struct GPU { - texture: wgpu::Texture, + pub texture: wgpu::Texture, } impl Sealed for Raster {} diff --git a/node-graph/graster-nodes/Cargo.toml b/node-graph/graster-nodes/Cargo.toml index a66dbb5eb3..7ed5357f00 100644 --- a/node-graph/graster-nodes/Cargo.toml +++ b/node-graph/graster-nodes/Cargo.toml @@ -17,6 +17,7 @@ default = ["std"] shader-nodes = [ "std", "dep:graphene-raster-nodes-shaders", + "dep:wgpu-executor", ] std = [ "dep:graphene-core", @@ -39,6 +40,7 @@ node-macro = { workspace = true } # Local std dependencies dyn-any = { workspace = true, optional = true } graphene-core = { workspace = true, optional = true } +wgpu-executor = { workspace = true, optional = true } graphene-raster-nodes-shaders = { path = "./shaders", optional = true } # Workspace dependencies diff --git a/node-graph/graster-nodes/src/adjustments.rs b/node-graph/graster-nodes/src/adjustments.rs index 35de34c76b..dc3bd15b82 100644 --- a/node-graph/graster-nodes/src/adjustments.rs +++ b/node-graph/graster-nodes/src/adjustments.rs @@ -30,7 +30,7 @@ use num_traits::float::Float; // https://www.adobe.com/devnet-apps/photoshop/fileformatashtml/#:~:text=%27clrL%27%20%3D%20Color%20Lookup // https://www.adobe.com/devnet-apps/photoshop/fileformatashtml/#:~:text=Color%20Lookup%20(Photoshop%20CS6 -#[derive(Debug, Default, Clone, Copy, Eq, PartialEq, Hash, node_macro::ChoiceType)] +#[derive(Debug, Default, Clone, Copy, Eq, PartialEq, Hash, node_macro::ChoiceType, bytemuck::NoUninit)] #[cfg_attr(feature = "std", derive(dyn_any::DynAny, specta::Type, serde::Serialize, serde::Deserialize))] #[widget(Dropdown)] #[repr(u32)] @@ -560,7 +560,7 @@ pub enum RedGreenBlue { } /// Color Channel -#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Hash, node_macro::ChoiceType)] +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Hash, node_macro::ChoiceType, bytemuck::NoUninit)] #[cfg_attr(feature = "std", derive(dyn_any::DynAny, specta::Type, serde::Serialize, serde::Deserialize))] #[widget(Radio)] #[repr(u32)] diff --git a/node-graph/graster-nodes/src/blending_nodes.rs b/node-graph/graster-nodes/src/blending_nodes.rs index 182f6a802d..13f93773d4 100644 --- a/node-graph/graster-nodes/src/blending_nodes.rs +++ b/node-graph/graster-nodes/src/blending_nodes.rs @@ -132,7 +132,7 @@ pub fn apply_blend_mode(foreground: Color, background: Color, blend_mode: BlendM } } -#[node_macro::node(category("Raster"), shader_node(PerPixelAdjust))] +#[node_macro::node(category("Raster"), cfg(feature = "std"))] fn blend + Send>( _: impl Ctx, #[implementations( diff --git a/node-graph/graster-nodes/src/fullscreen_vertex.rs b/node-graph/graster-nodes/src/fullscreen_vertex.rs new file mode 100644 index 0000000000..533efb03b3 --- /dev/null +++ b/node-graph/graster-nodes/src/fullscreen_vertex.rs @@ -0,0 +1,29 @@ +use glam::{Vec2, Vec4}; +use spirv_std::spirv; + +/// webgpu NDC is like OpenGL: (-1.0 .. 1.0, -1.0 .. 1.0, 0.0 .. 1.0) +/// https://www.w3.org/TR/webgpu/#coordinate-systems +/// +/// So to make a fullscreen triangle around a box at (-1..1): +/// +/// ```norun +/// 3 + +/// |\ +/// 2 | \ +/// | \ +/// 1 +-----+ +/// | |\ +/// 0 | 0 | \ +/// | | \ +/// -1 +-----+-----+ +/// -1 0 1 2 3 +/// ``` +const FULLSCREEN_VERTICES: [Vec2; 3] = [Vec2::new(-1., -1.), Vec2::new(-1., 3.), Vec2::new(3., -1.)]; + +#[spirv(vertex)] +pub fn fullscreen_vertex(#[spirv(vertex_index)] vertex_index: u32, #[spirv(position)] gl_position: &mut Vec4) { + // broken on edition 2024 branch + // let vertex = unsafe { *FULLSCREEN_VERTICES.index_unchecked(vertex_index as usize) }; + let vertex = FULLSCREEN_VERTICES[vertex_index as usize]; + *gl_position = Vec4::from((vertex, 0., 1.)); +} diff --git a/node-graph/graster-nodes/src/lib.rs b/node-graph/graster-nodes/src/lib.rs index 8dc169cea6..080504f9ea 100644 --- a/node-graph/graster-nodes/src/lib.rs +++ b/node-graph/graster-nodes/src/lib.rs @@ -4,6 +4,11 @@ pub mod adjust; pub mod adjustments; pub mod blending_nodes; pub mod cubic_spline; +pub mod fullscreen_vertex; + +/// required by shader macro +#[cfg(feature = "shader-nodes")] +pub use graphene_raster_nodes_shaders::WGSL_SHADER; #[cfg(feature = "std")] pub mod curve; diff --git a/node-graph/node-macro/src/codegen.rs b/node-graph/node-macro/src/codegen.rs index 05b8d19cc5..98bb5588a9 100644 --- a/node-graph/node-macro/src/codegen.rs +++ b/node-graph/node-macro/src/codegen.rs @@ -295,7 +295,8 @@ pub(crate) fn generate_node_code(parsed: &ParsedNodeFn) -> syn::Result syn::Result, pub(crate) display_name: Option, @@ -144,7 +144,7 @@ pub struct NodeParsedField { pub implementations: Punctuated, } -#[derive(Debug)] +#[derive(Clone, Debug)] pub(crate) struct Input { pub(crate) pat_ident: PatIdent, pub(crate) ty: Type, @@ -663,7 +663,7 @@ pub fn new_node_fn(attr: TokenStream2, item: TokenStream2) -> TokenStream2 { } impl ParsedNodeFn { - fn replace_impl_trait_in_input(&mut self) { + pub fn replace_impl_trait_in_input(&mut self) { if let Type::ImplTrait(impl_trait) = self.input.ty.clone() { let ident = Ident::new("_Input", impl_trait.span()); let mut bounds = impl_trait.bounds; diff --git a/node-graph/node-macro/src/shader_nodes/mod.rs b/node-graph/node-macro/src/shader_nodes/mod.rs index 0720869d01..d82ca8bfa8 100644 --- a/node-graph/node-macro/src/shader_nodes/mod.rs +++ b/node-graph/node-macro/src/shader_nodes/mod.rs @@ -3,24 +3,38 @@ use crate::shader_nodes::per_pixel_adjust::PerPixelAdjust; use proc_macro2::{Ident, TokenStream}; use quote::quote; use strum::VariantNames; -use syn::Error; use syn::parse::{Parse, ParseStream}; +use syn::punctuated::Punctuated; +use syn::{Error, Token}; pub mod per_pixel_adjust; pub const STD_FEATURE_GATE: &str = "std"; +pub const SHADER_NODES_FEATURE_GATE: &str = "shader-nodes"; pub fn modify_cfg(attributes: &NodeFnAttributes) -> TokenStream { - match (&attributes.cfg, &attributes.shader_node) { - (Some(cfg), Some(_)) => quote!(#[cfg(all(#cfg, feature = #STD_FEATURE_GATE))]), - (Some(cfg), None) => quote!(#[cfg(#cfg)]), - (None, Some(_)) => quote!(#[cfg(feature = #STD_FEATURE_GATE)]), - (None, None) => quote!(), - } + let feature_gate = match &attributes.shader_node { + // shader node cfg is done on the mod + Some(ShaderNodeType::ShaderNode) => quote!(), + Some(_) => quote!(feature = #STD_FEATURE_GATE), + None => quote!(), + }; + let cfgs: Punctuated<_, Token![,]> = match &attributes.cfg { + None => [&feature_gate].into_iter().collect(), + Some(cfg) => [cfg, &feature_gate].into_iter().collect(), + }; + quote!(#[cfg(all(#cfgs))]) } -#[derive(Debug, VariantNames)] +#[derive(Debug, Clone, VariantNames)] pub(crate) enum ShaderNodeType { + /// Marker for this node being in a gpu node crate, but not having a gpu implementation. This is distinct from not + /// declaring `shader_node` at all, as it will wrap the CPU node with a `#[cfg(feature = "std")]` feature gate. + None, + /// Marker for this node being a generated gpu node implementation, that should not emit anything to prevent + /// recursively generating more gpu nodes. But it still counts as a gpu node and will get the + /// `#[cfg(feature = "std")]` feature gate around it's impl. + ShaderNode, PerPixelAdjust(PerPixelAdjust), } @@ -28,24 +42,37 @@ impl Parse for ShaderNodeType { fn parse(input: ParseStream) -> syn::Result { let ident: Ident = input.parse()?; Ok(match ident.to_string().as_str() { + "None" => ShaderNodeType::None, "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; +pub trait ShaderCodegen { + fn codegen(&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")); +impl ShaderCodegen for ShaderNodeType { + fn codegen(&self, parsed: &ParsedNodeFn) -> syn::Result { + match self { + ShaderNodeType::None | ShaderNodeType::ShaderNode => (), + _ => { + 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), + ShaderNodeType::None | ShaderNodeType::ShaderNode => Ok(ShaderTokens::default()), + ShaderNodeType::PerPixelAdjust(x) => x.codegen(parsed), } } } + +#[derive(Clone, Default)] +pub struct ShaderTokens { + pub shader_entry_point: TokenStream, + pub gpu_node: TokenStream, +} 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 index 0e220c8aae..f01c364223 100644 --- a/node-graph/node-macro/src/shader_nodes/per_pixel_adjust.rs +++ b/node-graph/node-macro/src/shader_nodes/per_pixel_adjust.rs @@ -1,11 +1,15 @@ -use crate::parsing::{ParsedFieldType, ParsedNodeFn, RegularParsedField}; -use crate::shader_nodes::CodegenShaderEntryPoint; +use crate::parsing::{Input, NodeFnAttributes, ParsedField, ParsedFieldType, ParsedNodeFn, RegularParsedField}; +use crate::shader_nodes::{SHADER_NODES_FEATURE_GATE, ShaderCodegen, ShaderNodeType, ShaderTokens}; +use convert_case::{Case, Casing}; +use proc_macro_crate::FoundCrate; use proc_macro2::{Ident, TokenStream}; use quote::{ToTokens, format_ident, quote}; use std::borrow::Cow; use syn::parse::{Parse, ParseStream}; +use syn::punctuated::Punctuated; +use syn::{PatIdent, Type, parse_quote}; -#[derive(Debug)] +#[derive(Debug, Clone)] pub struct PerPixelAdjust {} impl Parse for PerPixelAdjust { @@ -14,53 +18,113 @@ impl Parse for PerPixelAdjust { } } -impl CodegenShaderEntryPoint for PerPixelAdjust { - fn codegen_shader_entry_point(&self, parsed: &ParsedNodeFn) -> syn::Result { +impl ShaderCodegen for PerPixelAdjust { + fn codegen(&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, .. }) => { + let mut params; + let has_uniform; + { + // categorize params + 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: ty.to_token_stream(), + param_type: ParamType::Uniform, + }), + ParsedFieldType::Regular(RegularParsedField { gpu_image: true, .. }) => { + let param = Param { + ident: Cow::Owned(format_ident!("image_{}", &ident.ident)), + ty: quote!(Image2d), + param_type: ParamType::Image { binding: 0 }, + }; + Ok(param) + } + } + }) + .collect::>>()?; + + has_uniform = params.iter().any(|p| matches!(p.param_type, ParamType::Uniform)); + + // assign image bindings + // if an arg_buffer exists, bindings for images start at 1 to leave 0 for arg buffer + let mut binding_cnt = if has_uniform { 1 } else { 0 }; + for p in params.iter_mut() { + match &mut p.param_type { + ParamType::Image { binding } => { + *binding = binding_cnt; 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 }, - }) } + ParamType::Uniform => {} } - }) - .collect::>>()?; + } + } + + let entry_point_mod = format_ident!("{}_gpu_entry_point", fn_name); + let entry_point_name_ident = format_ident!("ENTRY_POINT_NAME"); + let entry_point_name = quote!(#entry_point_mod::#entry_point_name_ident); + let uniform_struct_ident = format_ident!("Uniform"); + let uniform_struct = quote!(#entry_point_mod::#uniform_struct_ident); + let shader_node_mod = format_ident!("{}_shader_node", fn_name); - let uniform_members = params + let codegen = PerPixelAdjustCodegen { + parsed, + params, + has_uniform, + entry_point_mod, + entry_point_name_ident, + entry_point_name, + uniform_struct_ident, + uniform_struct, + shader_node_mod, + }; + + Ok(ShaderTokens { + shader_entry_point: codegen.codegen_shader_entry_point()?, + gpu_node: codegen.codegen_gpu_node()?, + }) + } +} + +pub struct PerPixelAdjustCodegen<'a> { + parsed: &'a ParsedNodeFn, + params: Vec>, + has_uniform: bool, + entry_point_mod: Ident, + entry_point_name_ident: Ident, + entry_point_name: TokenStream, + uniform_struct_ident: Ident, + uniform_struct: TokenStream, + shader_node_mod: Ident, +} + +impl PerPixelAdjustCodegen<'_> { + fn codegen_shader_entry_point(&self) -> syn::Result { + let fn_name = &self.parsed.fn_name; + let uniform_members = self + .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 + let image_params = self + .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 + let call_args = self + .params .iter() .map(|Param { ident, param_type, .. }| match param_type { ParamType::Image { .. } => quote!(Color::from_vec4(#ident.fetch_with(texel_coord, lod(0)))), @@ -69,8 +133,11 @@ impl CodegenShaderEntryPoint for PerPixelAdjust { .collect::>(); let context = quote!(()); + let entry_point_mod = &self.entry_point_mod; + let entry_point_name = &self.entry_point_name_ident; + let uniform_struct_ident = &self.uniform_struct_ident; Ok(quote! { - pub mod #gpu_mod { + pub mod #entry_point_mod { use super::*; use graphene_core_shaders::color::Color; use spirv_std::spirv; @@ -78,8 +145,12 @@ impl CodegenShaderEntryPoint for PerPixelAdjust { use spirv_std::image::{Image2d, ImageWithMethods}; use spirv_std::image::sample_with::lod; - pub struct Uniform { - #(#uniform_members),* + pub const #entry_point_name: &str = core::concat!(core::module_path!(), "::entry_point"); + + #[repr(C)] + #[derive(Copy, Clone, bytemuck::NoUninit)] + pub struct #uniform_struct_ident { + #(pub #uniform_members),* } #[spirv(fragment)] @@ -96,11 +167,162 @@ impl CodegenShaderEntryPoint for PerPixelAdjust { } }) } + + fn codegen_gpu_node(&self) -> syn::Result { + let gcore = match &self.parsed.crate_name { + FoundCrate::Itself => format_ident!("crate"), + FoundCrate::Name(name) => format_ident!("{name}"), + }; + + // adapt fields for gpu node + let raster_gpu: Type = parse_quote!(#gcore::table::Table<#gcore::raster_types::Raster<#gcore::raster_types::GPU>>); + let mut fields = self + .parsed + .fields + .iter() + .map(|f| match &f.ty { + ParsedFieldType::Regular(reg @ RegularParsedField { gpu_image: true, .. }) => Ok(ParsedField { + pat_ident: PatIdent { + mutability: None, + by_ref: None, + ..f.pat_ident.clone() + }, + ty: ParsedFieldType::Regular(RegularParsedField { + ty: raster_gpu.clone(), + implementations: Punctuated::default(), + ..reg.clone() + }), + ..f.clone() + }), + ParsedFieldType::Regular(RegularParsedField { gpu_image: false, .. }) => Ok(ParsedField { + pat_ident: PatIdent { + mutability: None, + by_ref: None, + ..f.pat_ident.clone() + }, + ..f.clone() + }), + ParsedFieldType::Node { .. } => Err(syn::Error::new_spanned(&f.pat_ident, "PerPixelAdjust shader nodes cannot accept other nodes as generics")), + }) + .collect::>>()?; + + // insert wgpu_executor field + let wgpu_executor = format_ident!("__wgpu_executor"); + fields.push(ParsedField { + pat_ident: PatIdent { + attrs: vec![], + by_ref: None, + mutability: None, + ident: parse_quote!(#wgpu_executor), + subpat: None, + }, + name: None, + description: "".to_string(), + widget_override: Default::default(), + ty: ParsedFieldType::Regular(RegularParsedField { + ty: parse_quote!(&'a WgpuExecutor), + exposed: false, + value_source: Default::default(), + number_soft_min: None, + number_soft_max: None, + number_hard_min: None, + number_hard_max: None, + number_mode_range: None, + implementations: Default::default(), + gpu_image: false, + }), + number_display_decimal_places: None, + number_step: None, + unit: None, + }); + + // find exactly one gpu_image field, runtime doesn't support more than 1 atm + let gpu_image_field = { + let mut iter = fields.iter().filter(|f| matches!(f.ty, ParsedFieldType::Regular(RegularParsedField { gpu_image: true, .. }))); + match (iter.next(), iter.next()) { + (Some(v), None) => Ok(v), + (Some(_), Some(more)) => Err(syn::Error::new_spanned(&more.pat_ident, "No more than one parameter must be annotated with `#[gpu_image]`")), + (None, _) => Err(syn::Error::new_spanned(&self.parsed.fn_name, "At least one parameter must be annotated with `#[gpu_image]`")), + }? + }; + let gpu_image = &gpu_image_field.pat_ident.ident; + + // uniform buffer struct construction + let has_uniform = self.has_uniform; + let uniform_buffer = if has_uniform { + let uniform_struct = &self.uniform_struct; + let uniform_members = self + .params + .iter() + .filter_map(|p| match p.param_type { + ParamType::Image { .. } => None, + ParamType::Uniform => Some(p.ident.as_ref()), + }) + .collect::>(); + quote!(Some(&super::#uniform_struct { + #(#uniform_members),* + })) + } else { + // explicit generics placed here cause it's easier than explicitly writing `run_per_pixel_adjust::<()>` + quote!(Option::<&()>::None) + }; + + // node function body + let entry_point_name = &self.entry_point_name; + let body = quote! { + { + #wgpu_executor.shader_runtime.run_per_pixel_adjust(&::wgpu_executor::shader_runtime::per_pixel_adjust_runtime::Shaders { + wgsl_shader: crate::WGSL_SHADER, + fragment_shader_name: super::#entry_point_name, + has_uniform: #has_uniform, + }, #gpu_image, #uniform_buffer).await + } + }; + + // call node codegen + let mut parsed_node_fn = ParsedNodeFn { + vis: self.parsed.vis.clone(), + attributes: NodeFnAttributes { + shader_node: Some(ShaderNodeType::ShaderNode), + ..self.parsed.attributes.clone() + }, + fn_name: self.shader_node_mod.clone(), + struct_name: format_ident!("{}", self.shader_node_mod.to_string().to_case(Case::Pascal)), + mod_name: self.shader_node_mod.clone(), + fn_generics: vec![parse_quote!('a: 'n)], + where_clause: None, + input: Input { + pat_ident: self.parsed.input.pat_ident.clone(), + ty: parse_quote!(impl #gcore::context::Ctx), + implementations: Default::default(), + }, + output_type: raster_gpu, + is_async: true, + fields, + body, + crate_name: self.parsed.crate_name.clone(), + description: "".to_string(), + }; + parsed_node_fn.replace_impl_trait_in_input(); + let gpu_node_impl = crate::codegen::generate_node_code(&parsed_node_fn)?; + + // wrap node in `mod #gpu_node_mod` + let shader_node_mod = &self.shader_node_mod; + Ok(quote! { + #[cfg(feature = #SHADER_NODES_FEATURE_GATE)] + mod #shader_node_mod { + use super::*; + use wgpu_executor::WgpuExecutor; + + #gpu_node_impl + } + }) + } } struct Param<'a> { ident: Cow<'a, Ident>, - ty: Cow<'a, TokenStream>, + ty: TokenStream, param_type: ParamType, } diff --git a/node-graph/wgpu-executor/src/lib.rs b/node-graph/wgpu-executor/src/lib.rs index 920a002c4e..0b42dd631e 100644 --- a/node-graph/wgpu-executor/src/lib.rs +++ b/node-graph/wgpu-executor/src/lib.rs @@ -1,6 +1,8 @@ mod context; +pub mod shader_runtime; pub mod texture_upload; +use crate::shader_runtime::ShaderRuntime; use anyhow::Result; pub use context::Context; use dyn_any::StaticType; @@ -18,6 +20,7 @@ use wgpu::{Origin3d, SurfaceConfiguration, TextureAspect}; pub struct WgpuExecutor { pub context: Context, vello_renderer: Mutex, + pub shader_runtime: ShaderRuntime, } impl std::fmt::Debug for WgpuExecutor { @@ -195,6 +198,7 @@ impl WgpuExecutor { .ok()?; Some(Self { + shader_runtime: ShaderRuntime::new(&context), context, vello_renderer: vello_renderer.into(), }) diff --git a/node-graph/wgpu-executor/src/shader_runtime/mod.rs b/node-graph/wgpu-executor/src/shader_runtime/mod.rs new file mode 100644 index 0000000000..e7e0df8d94 --- /dev/null +++ b/node-graph/wgpu-executor/src/shader_runtime/mod.rs @@ -0,0 +1,20 @@ +use crate::Context; +use crate::shader_runtime::per_pixel_adjust_runtime::PerPixelAdjustShaderRuntime; + +pub mod per_pixel_adjust_runtime; + +pub const FULLSCREEN_VERTEX_SHADER_NAME: &str = "fullscreen_vertexfullscreen_vertex"; + +pub struct ShaderRuntime { + context: Context, + per_pixel_adjust: PerPixelAdjustShaderRuntime, +} + +impl ShaderRuntime { + pub fn new(context: &Context) -> Self { + Self { + context: context.clone(), + per_pixel_adjust: PerPixelAdjustShaderRuntime::new(), + } + } +} diff --git a/node-graph/wgpu-executor/src/shader_runtime/per_pixel_adjust_runtime.rs b/node-graph/wgpu-executor/src/shader_runtime/per_pixel_adjust_runtime.rs new file mode 100644 index 0000000000..928f35a0b0 --- /dev/null +++ b/node-graph/wgpu-executor/src/shader_runtime/per_pixel_adjust_runtime.rs @@ -0,0 +1,243 @@ +use crate::Context; +use crate::shader_runtime::{FULLSCREEN_VERTEX_SHADER_NAME, ShaderRuntime}; +use bytemuck::NoUninit; +use futures::lock::Mutex; +use graphene_core::raster_types::{GPU, Raster}; +use graphene_core::table::{Table, TableRow}; +use std::borrow::Cow; +use std::collections::HashMap; +use wgpu::util::{BufferInitDescriptor, DeviceExt}; +use wgpu::{ + BindGroupDescriptor, BindGroupEntry, BindGroupLayoutDescriptor, BindGroupLayoutEntry, BindingResource, BindingType, Buffer, BufferBinding, BufferBindingType, BufferUsages, ColorTargetState, Face, + FragmentState, FrontFace, LoadOp, Operations, PipelineLayoutDescriptor, PolygonMode, PrimitiveState, PrimitiveTopology, RenderPassColorAttachment, RenderPassDescriptor, RenderPipelineDescriptor, + ShaderModuleDescriptor, ShaderSource, ShaderStages, StoreOp, TextureDescriptor, TextureDimension, TextureFormat, TextureSampleType, TextureViewDescriptor, TextureViewDimension, VertexState, +}; + +pub struct PerPixelAdjustShaderRuntime { + // TODO: PerPixelAdjustGraphicsPipeline already contains the key as `name` + pipeline_cache: Mutex>, +} + +impl PerPixelAdjustShaderRuntime { + pub fn new() -> Self { + Self { + pipeline_cache: Mutex::new(HashMap::new()), + } + } +} + +impl ShaderRuntime { + pub async fn run_per_pixel_adjust(&self, shaders: &Shaders<'_>, textures: Table>, args: Option<&T>) -> Table> { + let mut cache = self.per_pixel_adjust.pipeline_cache.lock().await; + let pipeline = cache + .entry(shaders.fragment_shader_name.to_owned()) + .or_insert_with(|| PerPixelAdjustGraphicsPipeline::new(&self.context, &shaders)); + + let arg_buffer = args.map(|args| { + let device = &self.context.device; + device.create_buffer_init(&BufferInitDescriptor { + label: Some(&format!("{} arg buffer", pipeline.name.as_str())), + usage: BufferUsages::STORAGE, + contents: bytemuck::bytes_of(args), + }) + }); + pipeline.dispatch(&self.context, textures, arg_buffer) + } +} + +pub struct Shaders<'a> { + pub wgsl_shader: &'a str, + pub fragment_shader_name: &'a str, + pub has_uniform: bool, +} + +pub struct PerPixelAdjustGraphicsPipeline { + name: String, + has_uniform: bool, + pipeline: wgpu::RenderPipeline, +} + +impl PerPixelAdjustGraphicsPipeline { + pub fn new(context: &Context, info: &Shaders) -> Self { + let device = &context.device; + let name = info.fragment_shader_name.to_owned(); + + let fragment_name = &name; + let fragment_name = &fragment_name[(fragment_name.find("::").unwrap() + 2)..]; + // TODO workaround to naga removing `:` + let fragment_name = fragment_name.replace(":", ""); + let shader_module = device.create_shader_module(ShaderModuleDescriptor { + label: Some(&format!("PerPixelAdjust {} wgsl shader", name)), + source: ShaderSource::Wgsl(Cow::Borrowed(info.wgsl_shader)), + }); + + let entries: &[_] = if info.has_uniform { + &[ + BindGroupLayoutEntry { + binding: 0, + visibility: ShaderStages::FRAGMENT, + ty: BindingType::Buffer { + ty: BufferBindingType::Storage { read_only: true }, + has_dynamic_offset: false, + min_binding_size: None, + }, + count: None, + }, + BindGroupLayoutEntry { + binding: 1, + visibility: ShaderStages::FRAGMENT, + ty: BindingType::Texture { + sample_type: TextureSampleType::Float { filterable: false }, + view_dimension: TextureViewDimension::D2, + multisampled: false, + }, + count: None, + }, + ] + } else { + &[BindGroupLayoutEntry { + binding: 0, + visibility: ShaderStages::FRAGMENT, + ty: BindingType::Texture { + sample_type: TextureSampleType::Float { filterable: false }, + view_dimension: TextureViewDimension::D2, + multisampled: false, + }, + count: None, + }] + }; + let pipeline_layout = device.create_pipeline_layout(&PipelineLayoutDescriptor { + label: Some(&format!("PerPixelAdjust {} PipelineLayout", name)), + bind_group_layouts: &[&device.create_bind_group_layout(&BindGroupLayoutDescriptor { + label: Some(&format!("PerPixelAdjust {} BindGroupLayout 0", name)), + entries, + })], + push_constant_ranges: &[], + }); + + let pipeline = device.create_render_pipeline(&RenderPipelineDescriptor { + label: Some(&format!("PerPixelAdjust {} Pipeline", name)), + layout: Some(&pipeline_layout), + vertex: VertexState { + module: &shader_module, + entry_point: Some(FULLSCREEN_VERTEX_SHADER_NAME), + compilation_options: Default::default(), + buffers: &[], + }, + primitive: PrimitiveState { + topology: PrimitiveTopology::TriangleList, + strip_index_format: None, + front_face: FrontFace::Ccw, + cull_mode: Some(Face::Back), + unclipped_depth: false, + polygon_mode: PolygonMode::Fill, + conservative: false, + }, + depth_stencil: None, + multisample: Default::default(), + fragment: Some(FragmentState { + module: &shader_module, + entry_point: Some(&fragment_name), + compilation_options: Default::default(), + targets: &[Some(ColorTargetState { + format: TextureFormat::Rgba8UnormSrgb, + blend: None, + write_mask: Default::default(), + })], + }), + multiview: None, + cache: None, + }); + Self { + pipeline, + name, + has_uniform: info.has_uniform, + } + } + + pub fn dispatch(&self, context: &Context, textures: Table>, arg_buffer: Option) -> Table> { + assert_eq!(self.has_uniform, arg_buffer.is_some()); + let device = &context.device; + let name = self.name.as_str(); + + let mut cmd = device.create_command_encoder(&wgpu::CommandEncoderDescriptor { + label: Some(&format!("{name} cmd encoder")), + }); + let out = textures + .iter() + .map(|instance| { + let tex_in = &instance.element.texture; + let view_in = tex_in.create_view(&TextureViewDescriptor::default()); + let format = tex_in.format(); + + let entries: &[_] = if let Some(arg_buffer) = arg_buffer.as_ref() { + &[ + BindGroupEntry { + binding: 0, + resource: BindingResource::Buffer(BufferBinding { + buffer: arg_buffer, + offset: 0, + size: None, + }), + }, + BindGroupEntry { + binding: 1, + resource: BindingResource::TextureView(&view_in), + }, + ] + } else { + &[BindGroupEntry { + binding: 0, + resource: BindingResource::TextureView(&view_in), + }] + }; + let bind_group = device.create_bind_group(&BindGroupDescriptor { + label: Some(&format!("{name} bind group")), + // `get_bind_group_layout` allocates unnecessary memory, we could create it manually to not do that + layout: &self.pipeline.get_bind_group_layout(0), + entries, + }); + + let tex_out = device.create_texture(&TextureDescriptor { + label: Some(&format!("{name} texture out")), + size: tex_in.size(), + mip_level_count: 1, + sample_count: 1, + dimension: TextureDimension::D2, + format, + usage: wgpu::TextureUsages::TEXTURE_BINDING | wgpu::TextureUsages::COPY_DST | wgpu::TextureUsages::COPY_SRC | wgpu::TextureUsages::RENDER_ATTACHMENT, + view_formats: &[format], + }); + + let view_out = tex_out.create_view(&TextureViewDescriptor::default()); + let mut rp = cmd.begin_render_pass(&RenderPassDescriptor { + label: Some(&format!("{name} render pipeline")), + color_attachments: &[Some(RenderPassColorAttachment { + view: &view_out, + resolve_target: None, + ops: Operations { + // should be dont_care but wgpu doesn't expose that + load: LoadOp::Clear(wgpu::Color::BLACK), + store: StoreOp::Store, + }, + })], + depth_stencil_attachment: None, + timestamp_writes: None, + occlusion_query_set: None, + }); + rp.set_pipeline(&self.pipeline); + rp.set_bind_group(0, Some(&bind_group), &[]); + rp.draw(0..3, 0..1); + + TableRow { + element: Raster::new(GPU { texture: tex_out }), + transform: *instance.transform, + alpha_blending: *instance.alpha_blending, + source_node_id: *instance.source_node_id, + } + }) + .collect::>(); + context.queue.submit([cmd.finish()]); + out + } +}