tessera_ui/renderer/compute/
pipeline.rs

1//! Custom GPU compute pipelines.
2//!
3//! ## Usage
4//!
5//! Implement [`ComputablePipeline`] and register it with [`ComputePipelineRegistry`].
6
7use std::{any::TypeId, collections::HashMap};
8
9use crate::{PxRect, PxSize, SendSync, render_graph::DrawGeometry};
10
11use super::command::ComputeCommand;
12
13/// Type-erased metadata describing a compute command within a batch.
14pub struct ErasedComputeBatchItem<'a> {
15    /// The compute command to execute.
16    pub command: &'a dyn ComputeCommand,
17    /// Local dimensions and their mapping into the pass textures.
18    pub geometry: DrawGeometry,
19    /// The rectangle written, in pass texture coordinates.
20    pub target_area: PxRect,
21}
22
23/// Strongly typed metadata describing a compute command within a batch.
24pub struct ComputeBatchItem<'a, C: ComputeCommand> {
25    /// The compute command to execute.
26    pub command: &'a C,
27    /// Local dimensions and their mapping into the pass textures.
28    pub geometry: DrawGeometry,
29    /// The rectangle written, in pass texture coordinates.
30    pub target_area: PxRect,
31}
32
33/// A batch of compute commands within an active compute pass.
34///
35/// The pass textures cover the scene region the batch reads and writes, so
36/// they can be larger than `target_size` and offset from the window.
37pub struct ComputeContext<'a, 'b, 'c, C: ComputeCommand> {
38    /// Renderer-owned resource creation and uploads.
39    pub graphics: &'a crate::GraphicsDevice,
40    /// WGPU device used to create and manage GPU resources.
41    pub device: &'a wgpu::Device,
42    /// Queue for submitting GPU workloads.
43    pub queue: &'a wgpu::Queue,
44    /// Surface configuration describing output formats and dimensions.
45    pub config: &'a wgpu::SurfaceConfiguration,
46    /// Size of the scene region covered by the pass textures.
47    pub target_size: PxSize,
48    /// Active compute pass encoder.
49    pub compute_pass: &'a mut wgpu::ComputePass<'b>,
50    /// Batch of typed compute items to process.
51    pub items: &'c [ComputeBatchItem<'c, C>],
52    /// Input texture view sampled by the compute pass.
53    pub input_view: &'a wgpu::TextureView,
54    /// Output texture view written by the compute pass.
55    pub output_view: &'a wgpu::TextureView,
56}
57
58/// Type-erased context used when dispatching compute pipelines.
59pub(crate) struct ErasedDispatchContext<'a, 'b> {
60    pub graphics: &'a crate::GraphicsDevice,
61    pub device: &'a wgpu::Device,
62    pub queue: &'a wgpu::Queue,
63    pub config: &'a wgpu::SurfaceConfiguration,
64    pub target_size: PxSize,
65    pub compute_pass: &'a mut wgpu::ComputePass<'b>,
66    pub input_view: &'a wgpu::TextureView,
67    pub output_view: &'a wgpu::TextureView,
68}
69
70/// Processes compute commands of type `C`.
71pub trait ComputablePipeline<C: ComputeCommand>: SendSync {
72    /// Dispatches the command batch over each item's target area. The output
73    /// starts as a copy of the input and uses a non-sRGB storage format; draws
74    /// that sample the scene next read the result.
75    fn dispatch(&mut self, context: &mut ComputeContext<C>);
76}
77
78pub(crate) trait ErasedComputablePipeline: SendSync {
79    /// Dispatches a type-erased compute command.
80    fn dispatch_erased(
81        &mut self,
82        context: ErasedDispatchContext<'_, '_>,
83        items: &[ErasedComputeBatchItem<'_>],
84    );
85}
86
87/// A wrapper to implement `ErasedComputablePipeline` for any
88/// `ComputablePipeline`.
89struct ComputablePipelineImpl<C: ComputeCommand, P: ComputablePipeline<C>> {
90    pipeline: P,
91    _command: std::marker::PhantomData<C>,
92}
93
94impl<C: ComputeCommand + 'static, P: ComputablePipeline<C>> ErasedComputablePipeline
95    for ComputablePipelineImpl<C, P>
96{
97    fn dispatch_erased(
98        &mut self,
99        context: ErasedDispatchContext<'_, '_>,
100        items: &[ErasedComputeBatchItem<'_>],
101    ) {
102        if items.is_empty() {
103            return;
104        }
105
106        let mut typed_items: Vec<ComputeBatchItem<'_, C>> = Vec::with_capacity(items.len());
107        for item in items {
108            let command = item
109                .command
110                .downcast_ref::<C>()
111                .expect("Compute batch contained command of unexpected type");
112            typed_items.push(ComputeBatchItem {
113                command,
114                geometry: item.geometry,
115                target_area: item.target_area,
116            });
117        }
118
119        self.pipeline.dispatch(&mut ComputeContext {
120            graphics: context.graphics,
121            device: context.device,
122            queue: context.queue,
123            config: context.config,
124            target_size: context.target_size,
125            compute_pass: context.compute_pass,
126            items: &typed_items,
127            input_view: context.input_view,
128            output_view: context.output_view,
129        });
130    }
131}
132
133/// Registered compute pipelines, selected by command type.
134#[derive(Default)]
135pub struct ComputePipelineRegistry {
136    pipelines: HashMap<TypeId, Box<dyn ErasedComputablePipeline>>,
137}
138
139impl ComputePipelineRegistry {
140    /// Creates an empty compute pipeline registry.
141    pub fn new() -> Self {
142        Self::default()
143    }
144
145    pub(crate) fn remove(&mut self, type_id: TypeId) {
146        self.pipelines.remove(&type_id);
147    }
148
149    /// Registers a pipeline for its command type, replacing any prior registration.
150    pub fn register<C: ComputeCommand + 'static>(
151        &mut self,
152        pipeline: impl ComputablePipeline<C> + 'static,
153    ) {
154        let erased_pipeline = Box::new(ComputablePipelineImpl {
155            pipeline,
156            _command: std::marker::PhantomData,
157        });
158        self.pipelines.insert(TypeId::of::<C>(), erased_pipeline);
159    }
160
161    /// Dispatches one or more commands to their corresponding registered
162    /// pipeline.
163    pub(crate) fn dispatch_erased(
164        &mut self,
165        context: ErasedDispatchContext<'_, '_>,
166        items: &[ErasedComputeBatchItem<'_>],
167    ) {
168        if items.is_empty() {
169            return;
170        }
171
172        let command_type_id = items[0].command.as_any().type_id();
173        if let Some(pipeline) = self.pipelines.get_mut(&command_type_id) {
174            pipeline.dispatch_erased(context, items);
175        } else {
176            panic!(
177                "No pipeline found for command {:?}",
178                std::any::type_name_of_val(items[0].command)
179            );
180        }
181    }
182}
tessera_ui/renderer/compute/pipeline.rs

tessera_ui/renderer/compute/
pipeline.rs
