tessera_ui/renderer/compute/
pipeline.rs

1//! GPU compute pipeline system for Tessera UI framework.
2//!
3//! This module provides the infrastructure for GPU compute operations in Tessera,
4//! enabling advanced visual effects and post-processing operations that would be
5//! inefficient or impossible to achieve with traditional CPU-based approaches.
6//!
7//! # Architecture Overview
8//!
9//! The compute pipeline system is designed to work seamlessly with the rendering
10//! pipeline, using a ping-pong buffer approach for efficient multi-pass operations.
11//! Each compute pipeline processes a specific type of compute command and operates
12//! on texture data using GPU compute shaders.
13//!
14//! ## Key Components
15//!
16//! - [`ComputablePipeline<C>`]: The main trait for implementing custom compute pipelines
17//! - [`ComputePipelineRegistry`]: Manages and dispatches commands to registered compute pipelines
18//! - [`ComputeResourceManager`]: Manages GPU buffers and resources for compute operations
19//!
20//! # Design Philosophy
21//!
22//! The compute pipeline system embraces WGPU's compute shader capabilities to enable:
23//!
24//! - **Advanced Post-Processing**: Blur, contrast adjustment, color grading, and other image effects
25//! - **Parallel Processing**: Leverage GPU parallelism for computationally intensive operations
26//! - **Real-Time Effects**: Achieve complex visual effects at interactive frame rates
27//! - **Memory Efficiency**: Use GPU memory directly without CPU roundtrips
28//!
29//! # Ping-Pong Rendering
30//!
31//! The system uses a ping-pong approach where:
32//!
33//! 1. **Input Texture**: Contains the result from previous rendering or compute pass
34//! 2. **Output Texture**: Receives the processed result from the current compute operation
35//! 3. **Format Convention**: All textures use `wgpu::TextureFormat::Rgba8Unorm` for compatibility
36//!
37//! This approach enables efficient chaining of multiple compute operations without
38//! intermediate CPU involvement.
39//!
40//! # Implementation Guide
41//!
42//! ## Creating a Custom Compute Pipeline
43//!
44//! To create a custom compute pipeline:
45//!
46//! 1. Define your compute command struct implementing [`ComputeCommand`]
47//! 2. Create a pipeline struct implementing [`ComputablePipeline<YourCommand>`]
48//! 3. Write a compute shader in WGSL
49//! 4. Register the pipeline with [`ComputePipelineRegistry::register`]
50//!
51//! ## Example: Simple Brightness Adjustment Pipeline
52//!
53//! ```rust,ignore
54//! use tessera_ui::{ComputeCommand, ComputablePipeline, compute::resource::ComputeResourceManager};
55//! use wgpu;
56//!
57//! // 1. Define the compute command
58//! #[derive(Debug)]
59//! struct BrightnessCommand {
60//!     brightness: f32,
61//! }
62//!
63//! impl ComputeCommand for BrightnessCommand {}
64//!
65//! // 2. Implement the pipeline
66//! struct BrightnessPipeline {
67//!     compute_pipeline: wgpu::ComputePipeline,
68//!     bind_group_layout: wgpu::BindGroupLayout,
69//! }
70//!
71//! impl BrightnessPipeline {
72//!     fn new(device: &wgpu::Device) -> Self {
73//!         // Create compute shader and pipeline
74//!         let shader = device.create_shader_module(wgpu::ShaderModuleDescriptor {
75//!             label: Some("Brightness Shader"),
76//!             source: wgpu::ShaderSource::Wgsl(include_str!("brightness.wgsl").into()),
77//!         });
78//!         
79//!         // ... setup bind group layout and pipeline ...
80//!         # unimplemented!()
81//!     }
82//! }
83//!
84//! impl ComputablePipeline<BrightnessCommand> for BrightnessPipeline {
85//!     fn dispatch(
86//!         &mut self,
87//!         device: &wgpu::Device,
88//!         queue: &wgpu::Queue,
89//!         config: &wgpu::SurfaceConfiguration,
90//!         compute_pass: &mut wgpu::ComputePass<'_>,
91//!         command: &BrightnessCommand,
92//!         resource_manager: &mut ComputeResourceManager,
93//!         input_view: &wgpu::TextureView,
94//!         output_view: &wgpu::TextureView,
95//!     ) {
96//!         // Create uniforms buffer with brightness value
97//!         let uniforms = [command.brightness];
98//!         let uniform_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
99//!             label: Some("Brightness Uniforms"),
100//!             contents: bytemuck::cast_slice(&uniforms),
101//!             usage: wgpu::BufferUsages::UNIFORM,
102//!         });
103//!         
104//!         // Create bind group with input/output textures and uniforms
105//!         let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
106//!             layout: &self.bind_group_layout,
107//!             entries: &[
108//!                 wgpu::BindGroupEntry { binding: 0, resource: uniform_buffer.as_entire_binding() },
109//!                 wgpu::BindGroupEntry { binding: 1, resource: wgpu::BindingResource::TextureView(input_view) },
110//!                 wgpu::BindGroupEntry { binding: 2, resource: wgpu::BindingResource::TextureView(output_view) },
111//!             ],
112//!             label: Some("brightness_bind_group"),
113//!         });
114//!         
115//!         // Dispatch compute shader
116//!         compute_pass.set_pipeline(&self.compute_pipeline);
117//!         compute_pass.set_bind_group(0, &bind_group, &[]);
118//!         compute_pass.dispatch_workgroups(
119//!             (config.width + 7) / 8,
120//!             (config.height + 7) / 8,
121//!             1
122//!         );
123//!     }
124//! }
125//!
126//! // 3. Register the pipeline
127//! let mut registry = ComputePipelineRegistry::new();
128//! let brightness_pipeline = BrightnessPipeline::new(&device);
129//! registry.register(brightness_pipeline);
130//! ```
131//!
132//! ## Example WGSL Compute Shader
133//!
134//! ```wgsl
135//! @group(0) @binding(0) var<uniform> brightness: f32;
136//! @group(0) @binding(1) var input_texture: texture_2d<f32>;
137//! @group(0) @binding(2) var output_texture: texture_storage_2d<rgba8unorm, write>;
138//!
139//! @compute @workgroup_size(8, 8)
140//! fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
141//!     let coords = vec2<i32>(global_id.xy);
142//!     let input_color = textureLoad(input_texture, coords, 0);
143//!     let output_color = vec4<f32>(input_color.rgb * brightness, input_color.a);
144//!     textureStore(output_texture, coords, output_color);
145//! }
146//! ```
147//!
148//! # Integration with Basic Components
149//!
150//! The `tessera_basic_components` crate provides several compute pipeline implementations:
151//!
152//! - **BlurPipeline**: Gaussian blur effects for backgrounds and UI elements
153//! - **MeanPipeline**: Average color calculation for adaptive UI themes
154//! - **ContrastPipeline**: Contrast and saturation adjustments
155//!
156//! These pipelines demonstrate real-world usage patterns and can serve as references
157//! for implementing custom compute operations.
158//!
159//! # Performance Considerations
160//!
161//! - **Workgroup Size**: Choose workgroup sizes that align with GPU architecture (typically 8x8 or 16x16)
162//! - **Memory Access**: Optimize memory access patterns in shaders for better cache utilization
163//! - **Resource Reuse**: Use the [`ComputeResourceManager`] to reuse buffers across frames
164//! - **Batch Operations**: Combine multiple similar operations when possible
165//!
166//! # Texture Format Requirements
167//!
168//! Due to WGPU limitations, compute shaders require specific texture formats:
169//!
170//! - **Input Textures**: Can be any readable format, typically from render passes
171//! - **Output Textures**: Must use `wgpu::TextureFormat::Rgba8Unorm` for storage binding
172//! - **sRGB Limitation**: sRGB formats cannot be used as storage textures
173//!
174//! The framework automatically handles format conversions when necessary.
175
176use std::{any::TypeId, collections::HashMap};
177
178use crate::{PxRect, compute::resource::ComputeResourceManager, renderer::command::AsAny};
179
180use super::command::ComputeCommand;
181
182/// Core trait for implementing GPU compute pipelines.
183///
184/// This trait defines the interface for compute pipelines that process specific types
185/// of compute commands using GPU compute shaders. Each pipeline is responsible for
186/// setting up compute resources, managing shader dispatch, and processing texture data.
187///
188/// # Type Parameters
189///
190/// * `C` - The specific [`ComputeCommand`] type this pipeline can handle
191///
192/// # Design Principles
193///
194/// - **Single Responsibility**: Each pipeline handles one specific type of compute operation
195/// - **Stateless Operation**: Pipelines should not maintain state between dispatch calls
196/// - **Resource Efficiency**: Reuse GPU resources when possible through the resource manager
197/// - **Thread Safety**: All implementations must be `Send + Sync` for parallel execution
198///
199/// # Integration with Rendering
200///
201/// Compute pipelines operate within the broader rendering pipeline, typically:
202///
203/// 1. **After Rendering**: Process the rendered scene for post-effects
204/// 2. **Between Passes**: Transform data between different rendering stages
205/// 3. **Before Rendering**: Prepare data or textures for subsequent render operations
206///
207/// # Example Implementation Pattern
208///
209/// ```rust,ignore
210/// impl ComputablePipeline<MyCommand> for MyPipeline {
211///     fn dispatch(&mut self, device, queue, config, compute_pass, command,
212///                 resource_manager, input_view, output_view) {
213///         // 1. Create or retrieve uniform buffer
214///         let uniforms = create_uniforms_from_command(command);
215///         let uniform_buffer = device.create_buffer_init(...);
216///         
217///         // 2. Create bind group with textures and uniforms
218///         let bind_group = device.create_bind_group(...);
219///         
220///         // 3. Set pipeline and dispatch
221///         compute_pass.set_pipeline(&self.compute_pipeline);
222///         compute_pass.set_bind_group(0, &bind_group, &[]);
223///         compute_pass.dispatch_workgroups(workgroup_x, workgroup_y, 1);
224///     }
225/// }
226/// ```
227pub trait ComputablePipeline<C: ComputeCommand>: Send + Sync + 'static {
228    /// Dispatches the compute command within an active compute pass.
229    ///
230    /// This method is called once for each compute command that needs to be processed.
231    /// It should set up the necessary GPU resources, bind them to the compute pipeline,
232    /// and dispatch the appropriate number of workgroups to process the input texture.
233    ///
234    /// # Parameters
235    ///
236    /// * `device` - The WGPU device for creating GPU resources
237    /// * `queue` - The WGPU queue for submitting commands and updating buffers
238    /// * `config` - Current surface configuration containing dimensions and format info
239    /// * `compute_pass` - The active compute pass to record commands into
240    /// * `command` - The specific compute command containing operation parameters
241    /// * `resource_manager` - Manager for reusing GPU buffers across operations
242    /// * `input_view` - View of the input texture (result from previous pass)
243    /// * `output_view` - View of the output texture (target for this operation)
244    ///
245    /// # Texture Format Requirements
246    ///
247    /// Due to WGPU limitations, storage textures have specific format requirements:
248    ///
249    /// - **Input Texture**: Can be any readable format, typically from render passes
250    /// - **Output Texture**: Must use `wgpu::TextureFormat::Rgba8Unorm` format
251    /// - **sRGB Limitation**: sRGB formats cannot be used as storage textures
252    ///
253    /// The framework ensures that `output_view` always uses a compatible format
254    /// for storage binding operations.
255    ///
256    /// # Workgroup Dispatch Guidelines
257    ///
258    /// When dispatching workgroups, consider:
259    ///
260    /// - **Workgroup Size**: Match your shader's `@workgroup_size` declaration
261    /// - **Coverage**: Ensure all pixels are processed by calculating appropriate dispatch dimensions
262    /// - **Alignment**: Round up dispatch dimensions to cover the entire texture
263    ///
264    /// Common dispatch pattern:
265    /// ```rust,ignore
266    /// let workgroup_size = 8; // Match shader @workgroup_size(8, 8)
267    /// let dispatch_x = (config.width + workgroup_size - 1) / workgroup_size;
268    /// let dispatch_y = (config.height + workgroup_size - 1) / workgroup_size;
269    /// compute_pass.dispatch_workgroups(dispatch_x, dispatch_y, 1);
270    /// ```
271    ///
272    /// # Resource Management
273    ///
274    /// Use the `resource_manager` to:
275    /// - Store persistent buffers that can be reused across frames
276    /// - Avoid recreating expensive GPU resources
277    /// - Manage buffer lifetimes efficiently
278    ///
279    /// # Error Handling
280    ///
281    /// This method should handle errors gracefully:
282    /// - Validate command parameters before use
283    /// - Ensure texture dimensions are compatible
284    /// - Handle resource creation failures appropriately
285    fn dispatch(
286        &mut self,
287        device: &wgpu::Device,
288        queue: &wgpu::Queue,
289        config: &wgpu::SurfaceConfiguration,
290        compute_pass: &mut wgpu::ComputePass<'_>,
291        command: &C,
292        resource_manager: &mut ComputeResourceManager,
293        target_area: PxRect,
294        input_view: &wgpu::TextureView,
295        output_view: &wgpu::TextureView,
296    );
297}
298
299/// Internal trait for type erasure of computable pipelines.
300///
301/// This trait enables dynamic dispatch of compute commands to their corresponding pipelines
302/// without knowing the specific command type at compile time. It's used internally by
303/// the [`ComputePipelineRegistry`] and should not be implemented directly by users.
304///
305/// The type erasure is achieved through the [`AsAny`] trait, which allows downcasting
306/// from `&dyn ComputeCommand` to concrete command types.
307///
308/// # Implementation Note
309///
310/// This trait is automatically implemented for any type that implements
311/// [`ComputablePipeline<C>`] through the [`ComputablePipelineImpl`] wrapper.
312pub(crate) trait ErasedComputablePipeline: Send + Sync {
313    /// Dispatches a type-erased compute command.
314    fn dispatch_erased(
315        &mut self,
316        device: &wgpu::Device,
317        queue: &wgpu::Queue,
318        config: &wgpu::SurfaceConfiguration,
319        compute_pass: &mut wgpu::ComputePass<'_>,
320        command: &dyn ComputeCommand,
321        resource_manager: &mut ComputeResourceManager,
322        target_area: PxRect,
323        input_view: &wgpu::TextureView,
324        output_view: &wgpu::TextureView,
325    );
326}
327
328/// A wrapper to implement `ErasedComputablePipeline` for any `ComputablePipeline`.
329struct ComputablePipelineImpl<C: ComputeCommand, P: ComputablePipeline<C>> {
330    pipeline: P,
331    _command: std::marker::PhantomData<C>,
332}
333
334impl<C: ComputeCommand + 'static, P: ComputablePipeline<C>> ErasedComputablePipeline
335    for ComputablePipelineImpl<C, P>
336{
337    fn dispatch_erased(
338        &mut self,
339        device: &wgpu::Device,
340        queue: &wgpu::Queue,
341        config: &wgpu::SurfaceConfiguration,
342        compute_pass: &mut wgpu::ComputePass<'_>,
343        command: &dyn ComputeCommand,
344        resource_manager: &mut ComputeResourceManager,
345        target_area: PxRect,
346        input_view: &wgpu::TextureView,
347        output_view: &wgpu::TextureView,
348    ) {
349        if let Some(command) = AsAny::as_any(command).downcast_ref::<C>() {
350            self.pipeline.dispatch(
351                device,
352                queue,
353                config,
354                compute_pass,
355                command,
356                resource_manager,
357                target_area,
358                input_view,
359                output_view,
360            );
361        }
362    }
363}
364
365/// Registry for managing and dispatching compute pipelines.
366///
367/// The `ComputePipelineRegistry` serves as the central hub for all compute pipelines
368/// in the Tessera framework. It maintains a collection of registered pipelines and
369/// handles the dispatch of compute commands to their appropriate pipelines.
370///
371/// # Architecture
372///
373/// The registry uses type erasure to store pipelines of different types in a single
374/// collection. When a compute command needs to be processed, the registry attempts
375/// to dispatch it to all registered pipelines until one handles it successfully.
376///
377/// # Usage Pattern
378///
379/// 1. Create a new registry
380/// 2. Register all required compute pipelines during application initialization
381/// 3. The renderer uses the registry to dispatch commands during frame rendering
382///
383/// # Example
384///
385/// ```rust,ignore
386/// use tessera_ui::renderer::compute::ComputePipelineRegistry;
387///
388/// // Create registry and register pipelines
389/// let mut registry = ComputePipelineRegistry::new();
390/// registry.register(blur_pipeline);
391/// registry.register(contrast_pipeline);
392/// registry.register(brightness_pipeline);
393///
394/// // Registry is now ready for use by the renderer
395/// ```
396///
397/// # Performance Considerations
398///
399/// - Pipeline lookup is O(1) on average due to HashMap implementation.
400///
401/// # Thread Safety
402///
403/// The registry and all registered pipelines must be `Send + Sync` to support
404/// parallel execution in the rendering system.
405#[derive(Default)]
406pub struct ComputePipelineRegistry {
407    pipelines: HashMap<TypeId, Box<dyn ErasedComputablePipeline>>,
408}
409
410impl ComputePipelineRegistry {
411    /// Creates a new empty compute pipeline registry.
412    ///
413    /// # Example
414    ///
415    /// ```
416    /// use tessera_ui::renderer::compute::ComputePipelineRegistry;
417    ///
418    /// let registry = ComputePipelineRegistry::new();
419    /// ```
420    pub fn new() -> Self {
421        Self::default()
422    }
423
424    /// Registers a new compute pipeline for a specific command type.
425    ///
426    /// This method takes ownership of the pipeline and wraps it in a type-erased
427    /// container that can be stored alongside other pipelines of different types.
428    ///
429    /// # Type Parameters
430    ///
431    /// * `C` - The [`ComputeCommand`] type this pipeline handles
432    ///
433    /// # Parameters
434    ///
435    /// * `pipeline` - The pipeline instance to register
436    ///
437    /// # Example
438    ///
439    /// ```rust,ignore
440    /// use tessera_ui::renderer::compute::ComputePipelineRegistry;
441    ///
442    /// let mut registry = ComputePipelineRegistry::new();
443    ///
444    /// // Register custom compute pipelines
445    /// let blur_pipeline = BlurPipeline::new(&device);
446    /// registry.register(blur_pipeline);
447    ///
448    /// let contrast_pipeline = ContrastPipeline::new(&device);
449    /// registry.register(contrast_pipeline);
450    ///
451    /// // Register multiple pipelines for different effects
452    /// registry.register(BrightnessAdjustmentPipeline::new(&device));
453    /// registry.register(ColorGradingPipeline::new(&device));
454    /// ```
455    ///
456    /// # Thread Safety
457    ///
458    /// The pipeline must implement `Send + Sync` to be compatible with Tessera's
459    /// parallel rendering architecture.
460    pub fn register<C: ComputeCommand + 'static>(
461        &mut self,
462        pipeline: impl ComputablePipeline<C> + 'static,
463    ) {
464        let erased_pipeline = Box::new(ComputablePipelineImpl {
465            pipeline,
466            _command: std::marker::PhantomData,
467        });
468        self.pipelines.insert(TypeId::of::<C>(), erased_pipeline);
469    }
470
471    /// Dispatches a command to its corresponding registered pipeline.
472    pub(crate) fn dispatch_erased(
473        &mut self,
474        device: &wgpu::Device,
475        queue: &wgpu::Queue,
476        config: &wgpu::SurfaceConfiguration,
477        compute_pass: &mut wgpu::ComputePass<'_>,
478        command: &dyn ComputeCommand,
479        resource_manager: &mut ComputeResourceManager,
480        target_area: PxRect,
481        input_view: &wgpu::TextureView,
482        output_view: &wgpu::TextureView,
483    ) {
484        let command_type_id = AsAny::as_any(command).type_id();
485        if let Some(pipeline) = self.pipelines.get_mut(&command_type_id) {
486            pipeline.dispatch_erased(
487                device,
488                queue,
489                config,
490                compute_pass,
491                command,
492                resource_manager,
493                target_area,
494                input_view,
495                output_view,
496            );
497        } else {
498            panic!(
499                "No pipeline found for command {:?}",
500                std::any::type_name_of_val(command)
501            );
502        }
503    }
504}