Reduce compute shading ceremony, remove branches in compute shaders

Author: friendlymatthew

src/renderer/app_state.rs

@@ -1,13 +1,15 @@
 use crate::{
     image::grammar::Image,
     renderer::{
+        compute_effect::ComputeEffect,
         draw_uniform::DrawUniform,
         feature_uniform::{FeatureUniform, TransformAction},
         gpu_state::GpuResourceAllocator,
         mouse_state::MouseState,
         shader::{Shader, TextureResource},
         shape::{compute_distance, Circle, EditorState},
         shape_uniform::{CircleData, ShapeUniform, MAX_CIRCLES},
+        Texture,
     },
 };
 use anyhow::Result;
@@ -39,9 +41,9 @@ pub struct AppState<'a> {
     pub shape_uniform: ShapeUniform,
     pub circle_storage_buffer: wgpu::Buffer,
 
-    pub color_correct_pipeline: wgpu::ComputePipeline,
-    pub color_correct_bind_group: wgpu::BindGroup,
-    pub output_texture: crate::renderer::Texture,
+    pub gamma_effect: ComputeEffect,
+    pub grayscale_effect: ComputeEffect,
+    pub invert_effect: ComputeEffect,
 }
 
 impl<'a> AppState<'a> {
@@ -72,43 +74,36 @@ impl<'a> AppState<'a> {
         let circle_storage_buffer =
             gpu_allocator.create_storage_buffer("circle_storage", &empty_circles)?;
 
-        let color_correct_pipeline = gpu_allocator.create_compute_pipeline(
-            "color_correct_compute",
-            include_str!("color_correct_compute.wgsl"),
-            "cs_main",
-        );
-
-        let output_texture =
-            gpu_allocator.create_storage_texture("processed_texture", size.width, size.height);
-
-        let color_correct_bind_group = {
-            let bind_group_layout = color_correct_pipeline.get_bind_group_layout(0);
-            gpu_allocator
-                .device
-                .create_bind_group(&wgpu::BindGroupDescriptor {
-                    layout: &bind_group_layout,
-                    entries: &[
-                        wgpu::BindGroupEntry {
-                            binding: 0,
-                            resource: wgpu::BindingResource::TextureView(
-                                &image_texture_resource.resource.view,
-                            ),
-                        },
-                        wgpu::BindGroupEntry {
-                            binding: 1,
-                            resource: wgpu::BindingResource::TextureView(&output_texture.view),
-                        },
-                        wgpu::BindGroupEntry {
-                            binding: 2,
-                            resource: feature_uniform_resource.resource.as_entire_binding(),
-                        },
-                    ],
-                    label: Some("color_correct_bind_group"),
-                })
-        };
+        let texture_a = gpu_allocator.create_storage_texture("texture_a", size.width, size.height);
+
+        // individual compute effects
+        let gamma_effect = ComputeEffect::builder("gamma")
+            .with_shader(include_str!("gamma_correct_compute.wgsl"))
+            .with_uniform(feature_uniform.gamma())
+            .build(
+                &gpu_allocator.device,
+                &image_texture_resource.resource.view,
+                &texture_a.view,
+            )?;
+
+        let grayscale_effect = ComputeEffect::builder("grayscale")
+            .with_shader(include_str!("grayscale_compute.wgsl"))
+            .build(
+                &gpu_allocator.device,
+                &image_texture_resource.resource.view,
+                &texture_a.view,
+            )?;
+
+        let invert_effect = ComputeEffect::builder("invert")
+            .with_shader(include_str!("invert_compute.wgsl"))
+            .build(
+                &gpu_allocator.device,
+                &image_texture_resource.resource.view,
+                &texture_a.view,
+            )?;
 
         let processed_texture_resource = gpu_allocator
-            .create_texture_resource_from_existing("processed_texture_ref", &output_texture);
+            .create_texture_resource_from_existing("processed_texture_ref", &texture_a);
 
         let shape_texture_for_image = gpu_allocator.create_texture_resource_from_existing(
             "shape_texture_ref",
@@ -147,9 +142,9 @@ impl<'a> AppState<'a> {
             shape_render_texture,
             shape_uniform,
             circle_storage_buffer,
-            color_correct_pipeline,
-            color_correct_bind_group,
-            output_texture,
+            gamma_effect,
+            grayscale_effect,
+            invert_effect,
         })
     }
 
@@ -428,6 +423,10 @@ impl<'a> AppState<'a> {
     }
 
     pub(crate) fn update(&mut self) {
+        // Update gamma effect uniform
+        self.gamma_effect
+            .update_uniform(&self.gpu_allocator.queue, self.feature_uniform.gamma());
+
         // Update image shader uniforms
         let uniform_resources = &self.image_shader.uniform_resources;
         self.gpu_allocator
@@ -472,19 +471,23 @@ impl<'a> AppState<'a> {
     pub(crate) fn render(&self) -> Result<(), wgpu::SurfaceError> {
         let (output, view, mut encoder) = self.gpu_allocator.begin_frame()?;
 
-        // compute pass 1: apply color corrections
+        // Compute shader pass
         {
-            let mut compute_pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
-                label: Some("color_correct_compute_pass"),
-                timestamp_writes: None,
-            });
-
-            compute_pass.set_pipeline(&self.color_correct_pipeline);
-            compute_pass.set_bind_group(0, &self.color_correct_bind_group, &[]);
-
+            // todo, these color effects work independently from each other
+            // how do we pass intermediate textures?
             let workgroup_count_x = (self.size.width + 15) / 16;
             let workgroup_count_y = (self.size.height + 15) / 16;
-            compute_pass.dispatch_workgroups(workgroup_count_x, workgroup_count_y, 1);
+
+            if self.feature_uniform.invert() {
+                self.invert_effect
+                    .dispatch(&mut encoder, workgroup_count_x, workgroup_count_y);
+            } else if self.feature_uniform.grayscale() {
+                self.grayscale_effect
+                    .dispatch(&mut encoder, workgroup_count_x, workgroup_count_y);
+            } else {
+                self.gamma_effect
+                    .dispatch(&mut encoder, workgroup_count_x, workgroup_count_y);
+            }
         }
 
         // First pass: Render shapes to shape texture

src/renderer/compute_effect.rs

@@ -0,0 +1,178 @@
+use anyhow::Result;
+use wgpu::util::DeviceExt;
+
+// a compute shader effect that can be applied to textures
+#[derive(Debug)]
+pub struct ComputeEffect {
+    pub pipeline: wgpu::ComputePipeline,
+    pub bind_group: wgpu::BindGroup,
+    uniform_buffer: Option<wgpu::Buffer>,
+}
+
+impl ComputeEffect {
+    pub fn builder<'a>(label: &'a str) -> ComputeEffectBuilder<'a> {
+        ComputeEffectBuilder {
+            label,
+            shader_source: None,
+            entry_point: "main",
+            uniform_data: None,
+        }
+    }
+
+    pub fn update_uniform<T: bytemuck::Pod>(&self, queue: &wgpu::Queue, data: T) {
+        if let Some(buffer) = &self.uniform_buffer {
+            queue.write_buffer(buffer, 0, bytemuck::cast_slice(&[data]));
+        }
+    }
+
+    pub fn dispatch(
+        &self,
+        encoder: &mut wgpu::CommandEncoder,
+        workgroup_count_x: u32,
+        workgroup_count_y: u32,
+    ) {
+        let mut compute_pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
+            label: None,
+            timestamp_writes: None,
+        });
+
+        compute_pass.set_pipeline(&self.pipeline);
+        compute_pass.set_bind_group(0, &self.bind_group, &[]);
+        compute_pass.dispatch_workgroups(workgroup_count_x, workgroup_count_y, 1);
+    }
+}
+
+pub struct ComputeEffectBuilder<'a> {
+    label: &'a str,
+    shader_source: Option<&'a str>,
+    entry_point: &'a str,
+    uniform_data: Option<Vec<u8>>,
+}
+
+impl<'a> ComputeEffectBuilder<'a> {
+    pub fn with_shader(mut self, source: &'a str) -> Self {
+        self.shader_source = Some(source);
+        self
+    }
+
+    pub fn with_entry_point(mut self, entry_point: &'a str) -> Self {
+        self.entry_point = entry_point;
+        self
+    }
+
+    pub fn with_uniform<T: bytemuck::Pod>(mut self, data: T) -> Self {
+        self.uniform_data = Some(bytemuck::cast_slice(&[data]).to_vec());
+        self
+    }
+
+    pub fn build(
+        self,
+        device: &wgpu::Device,
+        input_texture: &wgpu::TextureView,
+        output_texture: &wgpu::TextureView,
+    ) -> Result<ComputeEffect> {
+        let shader_source = self
+            .shader_source
+            .ok_or_else(|| anyhow::anyhow!("Shader source is required"))?;
+
+        let shader = device.create_shader_module(wgpu::ShaderModuleDescriptor {
+            label: Some(self.label),
+            source: wgpu::ShaderSource::Wgsl(shader_source.into()),
+        });
+
+        let uniform_buffer = self.uniform_data.as_ref().map(|data| {
+            device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
+                label: Some(&format!("{}_uniform", self.label)),
+                contents: data,
+                usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
+            })
+        });
+
+        let mut bind_group_layout_entries = vec![
+            wgpu::BindGroupLayoutEntry {
+                binding: 0,
+                visibility: wgpu::ShaderStages::COMPUTE,
+                ty: wgpu::BindingType::Texture {
+                    multisampled: false,
+                    view_dimension: wgpu::TextureViewDimension::D2,
+                    sample_type: wgpu::TextureSampleType::Float { filterable: false },
+                },
+                count: None,
+            },
+            // output storage texture
+            wgpu::BindGroupLayoutEntry {
+                binding: 1,
+                visibility: wgpu::ShaderStages::COMPUTE,
+                ty: wgpu::BindingType::StorageTexture {
+                    access: wgpu::StorageTextureAccess::WriteOnly,
+                    format: wgpu::TextureFormat::Rgba8Unorm,
+                    view_dimension: wgpu::TextureViewDimension::D2,
+                },
+                count: None,
+            },
+        ];
+
+        if uniform_buffer.is_some() {
+            bind_group_layout_entries.push(wgpu::BindGroupLayoutEntry {
+                binding: 2,
+                visibility: wgpu::ShaderStages::COMPUTE,
+                ty: wgpu::BindingType::Buffer {
+                    ty: wgpu::BufferBindingType::Uniform,
+                    has_dynamic_offset: false,
+                    min_binding_size: None,
+                },
+                count: None,
+            });
+        }
+
+        let bind_group_layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
+            entries: &bind_group_layout_entries,
+            label: Some(&format!("{}_bind_group_layout", self.label)),
+        });
+
+        let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
+            label: Some(&format!("{}_pipeline_layout", self.label)),
+            bind_group_layouts: &[&bind_group_layout],
+            push_constant_ranges: &[],
+        });
+
+        let pipeline = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
+            label: Some(self.label),
+            layout: Some(&pipeline_layout),
+            module: &shader,
+            entry_point: self.entry_point,
+            compilation_options: Default::default(),
+            cache: None,
+        });
+
+        let mut bind_group_entries = vec![
+            wgpu::BindGroupEntry {
+                binding: 0,
+                resource: wgpu::BindingResource::TextureView(input_texture),
+            },
+            wgpu::BindGroupEntry {
+                binding: 1,
+                resource: wgpu::BindingResource::TextureView(output_texture),
+            },
+        ];
+
+        if let Some(buffer) = &uniform_buffer {
+            bind_group_entries.push(wgpu::BindGroupEntry {
+                binding: 2,
+                resource: buffer.as_entire_binding(),
+            });
+        }
+
+        let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
+            layout: &bind_group_layout,
+            entries: &bind_group_entries,
+            label: Some(&format!("{}_bind_group", self.label)),
+        });
+
+        Ok(ComputeEffect {
+            pipeline,
+            bind_group,
+            uniform_buffer,
+        })
+    }
+}

src/renderer/feature_uniform.rs

@@ -174,3 +174,9 @@ impl FeatureUniform {
         }
     }
 }
+
+impl FeatureUniform {
+    pub fn gamma(&self) -> u32 {
+        self.gamma
+    }
+}

src/renderer/gamma_correct_compute.wgsl

@@ -0,0 +1,43 @@
+// a compute shader for gamma correction
+
+@group(0) 
+@binding(0)
+var input_texture: texture_2d<f32>;
+
+@group(0) 
+@binding(1)
+var output_texture: texture_storage_2d<rgba8unorm, write>;
+
+@group(0)
+@binding(2)
+var<uniform> gamma: u32;
+
+@compute
+@workgroup_size(16, 16)
+fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
+    let coords = vec2<i32>(global_id.xy);
+    let dimensions = textureDimensions(input_texture);
+
+    if (coords.x >= i32(dimensions.x) || coords.y >= i32(dimensions.y)) {
+        return;
+    }
+
+    var color = textureLoad(input_texture, coords, 0);
+
+    // If gamma is 0, just pass through (no correction)
+    if (gamma != 0u) {
+        let gamma_value = f32(gamma) / 100000.0;
+        let inv_gamma = 1.0 / gamma_value;
+
+        color = vec4<f32>(
+            pow(color.r, inv_gamma),
+            pow(color.g, inv_gamma),
+            pow(color.b, inv_gamma),
+            color.a
+        );
+    }
+
+    textureStore(output_texture, coords, color);
+
+}
+