Skip to main content

viewport_lib/runtime/
gpu_plugin.rs

1//! GPU plugin trait and lifecycle hooks for runtime extensions.
2//!
3//! `GpuPlugin` is the GPU-side counterpart to [`RuntimePlugin`](super::RuntimePlugin).
4//! Where `RuntimePlugin` mutates the scene each frame in `step`, `GpuPlugin`
5//! encodes wgpu command buffers at well-defined lifecycle points around the
6//! renderer's own work. The two traits are independent: a plugin may implement
7//! one, the other, or both and register itself separately with
8//! [`ViewportRuntime::with_plugin`](super::ViewportRuntime::with_plugin) and
9//! [`ViewportRuntime::with_gpu_plugin`](super::ViewportRuntime::with_gpu_plugin).
10//!
11//! # Frame integration
12//!
13//! ```text
14//! output = runtime.step(scene, selection, &frame_ctx);
15//! // host applies output (write_mesh_positions_normals, set_skin_palette, ...)
16//!
17//! let plugin_bufs  = runtime.pre_prepare(device, queue, &gpu_ctx);
18//! let prepare_bufs = renderer.pass().prepare(device, queue, &frame);
19//! queue.submit(plugin_bufs.into_iter().chain(prepare_bufs));
20//! ```
21//!
22//! wgpu submit ordering guarantees plugin command buffers complete before
23//! `prepare()`'s, so a storage buffer written by a plugin is observable by
24//! the standard render passes in the same frame.
25
26use crate::camera::Camera;
27use crate::renderer::ViewportId;
28
29/// Priority bands for GPU plugin lifecycle points.
30///
31/// Mirrors [`super::plugin::phase`] for CPU plugins. Use a band constant as
32/// the base of [`GpuPlugin::priority`] and offset within it (e.g.
33/// `gpu_phase::PRE_PREPARE + 10`) to order plugins inside the same band.
34pub mod gpu_phase {
35    /// Compute that produces inputs read by `renderer.prepare()` and the
36    /// standard render passes: cloth, hair, GPU particles, morph blends,
37    /// audio-reactive displacement, probe capture.
38    pub const PRE_PREPARE: i32 = 100;
39
40    /// Reserved priority band between `PRE_PREPARE` and `POST_PAINT`.
41    /// Currently unused; plugins cannot slot work between the renderer's own
42    /// internal passes.
43    pub const _RESERVED_INTERNAL: i32 = 500;
44
45    /// Compute that samples rendered targets (custom AO, motion blur,
46    /// screen-space outline, color grading).
47    pub const POST_PAINT: i32 = 900;
48}
49
50/// Rendered target views handed to [`GpuPlugin::post_paint`].
51///
52/// The host owns these textures and is responsible for keeping them alive for
53/// the duration of the `post_paint` call. A plugin that wants to *modify*
54/// `color_view` should encode into a sibling target it owns; the host then
55/// composites that overlay into the final image. The lib does not compose
56/// plugin output back into the rendered color target for you.
57pub struct PostPaintTargets<'a> {
58    /// View of the just-rendered color target.
59    pub color_view: &'a wgpu::TextureView,
60    /// View of the depth target produced during paint.
61    pub depth_view: &'a wgpu::TextureView,
62    /// View of the pick-id target (`R32Uint`) if the host renders one. `None`
63    /// when the active path does not produce a pick-id texture (e.g. an eframe
64    /// callback that did not request GPU picking this frame).
65    pub pick_id_view: Option<&'a wgpu::TextureView>,
66    /// Format of `color_view`. Useful when a plugin builds a pipeline whose
67    /// output format must match.
68    pub color_format: wgpu::TextureFormat,
69}
70
71/// Per-frame, read-only context passed to a [`GpuPlugin`].
72///
73/// Mirrors the render-relevant fields of [`RuntimeFrameContext`](super::RuntimeFrameContext)
74/// without the input or pick state, which GPU plugins do not consume.
75pub struct GpuFrameContext<'a> {
76    /// Active camera for this frame.
77    pub camera: &'a Camera,
78    /// Viewport size in physical pixels.
79    pub viewport_size: glam::Vec2,
80    /// Wall-clock seconds since the previous frame.
81    pub dt: f32,
82    /// Monotonically increasing frame counter, supplied by the host.
83    pub frame_index: u64,
84    /// Viewport this invocation is associated with, if the host renders
85    /// multiple viewports per frame and calls `pre_prepare` / `post_paint`
86    /// once per viewport. `None` for single-viewport hosts.
87    ///
88    /// Plugins registered via
89    /// [`ViewportRuntime::with_gpu_plugin_for_viewport`](super::ViewportRuntime::with_gpu_plugin_for_viewport)
90    /// only execute when this field matches the viewport they were scoped
91    /// to. Unscoped plugins run on every call regardless of this field.
92    pub viewport_id: Option<ViewportId>,
93}
94
95/// A plugin that encodes GPU work into the per-frame command stream.
96///
97/// Register with [`ViewportRuntime::with_gpu_plugin`](super::ViewportRuntime::with_gpu_plugin).
98/// Each frame, after `runtime.step()` and before `renderer.prepare()`, the
99/// host calls [`ViewportRuntime::pre_prepare`](super::ViewportRuntime::pre_prepare),
100/// which invokes every registered plugin's `pre_prepare` in ascending priority
101/// order and returns the concatenated command buffers for submission.
102///
103/// # Example
104///
105/// ```rust,ignore
106/// use viewport_lib::runtime::{GpuFrameContext, GpuPlugin, gpu_phase};
107///
108/// struct WarpPlugin {
109///     pipeline: Option<wgpu::ComputePipeline>,
110///     bind_group: Option<wgpu::BindGroup>,
111/// }
112///
113/// impl GpuPlugin for WarpPlugin {
114///     fn priority(&self) -> i32 { gpu_phase::PRE_PREPARE }
115///
116///     fn init_gpu(&mut self, device: &wgpu::Device) {
117///         // build pipeline and bind group, store on self
118///     }
119///
120///     fn pre_prepare(
121///         &mut self,
122///         device: &wgpu::Device,
123///         _queue: &wgpu::Queue,
124///         _ctx: &GpuFrameContext<'_>,
125///     ) -> Vec<wgpu::CommandBuffer> {
126///         let mut enc = device.create_command_encoder(&Default::default());
127///         {
128///             let mut pass = enc.begin_compute_pass(&Default::default());
129///             pass.set_pipeline(self.pipeline.as_ref().unwrap());
130///             pass.set_bind_group(0, self.bind_group.as_ref().unwrap(), &[]);
131///             pass.dispatch_workgroups(64, 1, 1);
132///         }
133///         vec![enc.finish()]
134///     }
135/// }
136/// ```
137pub trait GpuPlugin: Send + 'static {
138    /// Ascending priority. Lower runs first. Use [`gpu_phase`] constants as
139    /// base values. Defaults to [`gpu_phase::PRE_PREPARE`] since that is the
140    /// only active band today.
141    fn priority(&self) -> i32 {
142        gpu_phase::PRE_PREPARE
143    }
144
145    /// Build pipelines, persistent buffers, and bind group layouts.
146    ///
147    /// Called once before the plugin's first `pre_prepare`. If a new GPU
148    /// plugin is registered after the runtime has already run, every plugin's
149    /// `init_gpu` is invoked again on the next frame; implementations should
150    /// either be idempotent or guard their own one-time setup.
151    fn init_gpu(&mut self, _device: &wgpu::Device) {}
152
153    /// Called when the wgpu device is recreated, e.g. after device loss or a
154    /// host-driven reset. Every cached buffer, texture, bind group, or
155    /// pipeline the plugin built against the old device is now invalid and
156    /// must be rebuilt against `device` / `queue`.
157    ///
158    /// The host invokes this via
159    /// [`ViewportRuntime::notify_device_recreated`](super::ViewportRuntime::notify_device_recreated);
160    /// the runtime does not detect device loss on its own. After this call
161    /// returns, [`init_gpu`](Self::init_gpu) is also re-invoked on every
162    /// plugin before the next `pre_prepare`, so a typical implementation
163    /// can simply drop its cached resources here and let `init_gpu` rebuild
164    /// them.
165    fn on_device_recreated(&mut self, _device: &wgpu::Device, _queue: &wgpu::Queue) {}
166
167    /// Encode work that runs before `renderer.prepare()`.
168    ///
169    /// CPU plugins have already stepped this frame; their outputs are visible
170    /// through any shared state the consumer wired. The renderer has not
171    /// started its own work yet, so anything written here is observable by
172    /// `prepare()`'s passes (cluster build, shadow render, etc.).
173    fn pre_prepare(
174        &mut self,
175        _device: &wgpu::Device,
176        _queue: &wgpu::Queue,
177        _ctx: &GpuFrameContext<'_>,
178    ) -> Vec<wgpu::CommandBuffer> {
179        Vec::new()
180    }
181
182    /// Encode work that runs after `renderer.paint_to()`.
183    ///
184    /// The plugin receives views of the just-rendered color, depth, and
185    /// (optionally) pick-id targets and may sample them in a compute or
186    /// fullscreen render pass. Typical uses: custom AO, motion blur,
187    /// screen-space outlines, color grading LUTs.
188    ///
189    /// To *modify* the color target, write to a sibling texture the plugin
190    /// owns and let the host composite it during the final blit. The lib
191    /// does not loop a plugin's output back into the rendered color.
192    fn post_paint(
193        &mut self,
194        _device: &wgpu::Device,
195        _queue: &wgpu::Queue,
196        _targets: &PostPaintTargets<'_>,
197        _ctx: &GpuFrameContext<'_>,
198    ) -> Vec<wgpu::CommandBuffer> {
199        Vec::new()
200    }
201}
202
203/// Internal adapter that filters a [`GpuPlugin`] to a single viewport.
204///
205/// Constructed by
206/// [`ViewportRuntime::with_gpu_plugin_for_viewport`](super::ViewportRuntime::with_gpu_plugin_for_viewport).
207/// The wrapper forwards `priority`, `init_gpu`, and `on_device_recreated`
208/// unconditionally, but only forwards `pre_prepare` / `post_paint` when
209/// `ctx.viewport_id == Some(scoped_viewport)`.
210pub(super) struct ViewportScopedPlugin<P: GpuPlugin> {
211    pub(super) viewport: ViewportId,
212    pub(super) inner: P,
213}
214
215impl<P: GpuPlugin> GpuPlugin for ViewportScopedPlugin<P> {
216    fn priority(&self) -> i32 {
217        self.inner.priority()
218    }
219
220    fn init_gpu(&mut self, device: &wgpu::Device) {
221        self.inner.init_gpu(device);
222    }
223
224    fn on_device_recreated(&mut self, device: &wgpu::Device, queue: &wgpu::Queue) {
225        self.inner.on_device_recreated(device, queue);
226    }
227
228    fn pre_prepare(
229        &mut self,
230        device: &wgpu::Device,
231        queue: &wgpu::Queue,
232        ctx: &GpuFrameContext<'_>,
233    ) -> Vec<wgpu::CommandBuffer> {
234        if ctx.viewport_id == Some(self.viewport) {
235            self.inner.pre_prepare(device, queue, ctx)
236        } else {
237            Vec::new()
238        }
239    }
240
241    fn post_paint(
242        &mut self,
243        device: &wgpu::Device,
244        queue: &wgpu::Queue,
245        targets: &PostPaintTargets<'_>,
246        ctx: &GpuFrameContext<'_>,
247    ) -> Vec<wgpu::CommandBuffer> {
248        if ctx.viewport_id == Some(self.viewport) {
249            self.inner.post_paint(device, queue, targets, ctx)
250        } else {
251            Vec::new()
252        }
253    }
254}
255
256#[cfg(test)]
257mod tests {
258    use super::*;
259    use std::sync::{Arc, Mutex};
260
261    struct RecordingPlugin {
262        prio: i32,
263        log: Arc<Mutex<Vec<i32>>>,
264        post_log: Arc<Mutex<Vec<i32>>>,
265        init_count: Arc<Mutex<u32>>,
266    }
267
268    impl GpuPlugin for RecordingPlugin {
269        fn priority(&self) -> i32 {
270            self.prio
271        }
272
273        fn init_gpu(&mut self, _device: &wgpu::Device) {
274            *self.init_count.lock().unwrap() += 1;
275        }
276
277        fn pre_prepare(
278            &mut self,
279            _device: &wgpu::Device,
280            _queue: &wgpu::Queue,
281            _ctx: &GpuFrameContext<'_>,
282        ) -> Vec<wgpu::CommandBuffer> {
283            self.log.lock().unwrap().push(self.prio);
284            Vec::new()
285        }
286
287        fn post_paint(
288            &mut self,
289            _device: &wgpu::Device,
290            _queue: &wgpu::Queue,
291            _targets: &PostPaintTargets<'_>,
292            _ctx: &GpuFrameContext<'_>,
293        ) -> Vec<wgpu::CommandBuffer> {
294            self.post_log.lock().unwrap().push(self.prio);
295            Vec::new()
296        }
297    }
298
299    fn make_dummy_view(
300        device: &wgpu::Device,
301        format: wgpu::TextureFormat,
302        usage: wgpu::TextureUsages,
303    ) -> (wgpu::Texture, wgpu::TextureView) {
304        let tex = device.create_texture(&wgpu::TextureDescriptor {
305            label: Some("post_paint_test_target"),
306            size: wgpu::Extent3d { width: 1, height: 1, depth_or_array_layers: 1 },
307            mip_level_count: 1,
308            sample_count: 1,
309            dimension: wgpu::TextureDimension::D2,
310            format,
311            usage,
312            view_formats: &[],
313        });
314        let view = tex.create_view(&wgpu::TextureViewDescriptor::default());
315        (tex, view)
316    }
317
318    fn try_make_device() -> Option<(wgpu::Device, wgpu::Queue)> {
319        let instance = wgpu::Instance::new(&wgpu::InstanceDescriptor::default());
320        let adapter = pollster::block_on(instance.request_adapter(&wgpu::RequestAdapterOptions {
321            power_preference: wgpu::PowerPreference::LowPower,
322            compatible_surface: None,
323            force_fallback_adapter: false,
324        }))
325        .ok()?;
326        pollster::block_on(adapter.request_device(&wgpu::DeviceDescriptor::default())).ok()
327    }
328
329    #[test]
330    fn priority_order_and_init_once() {
331        let Some((device, queue)) = try_make_device() else {
332            eprintln!("skipping: no wgpu adapter available");
333            return;
334        };
335
336        let log = Arc::new(Mutex::new(Vec::new()));
337        let post_log = Arc::new(Mutex::new(Vec::new()));
338        let init_a = Arc::new(Mutex::new(0));
339        let init_b = Arc::new(Mutex::new(0));
340
341        let mut runtime = crate::runtime::ViewportRuntime::new()
342            .with_gpu_plugin(RecordingPlugin {
343                prio: 200,
344                log: log.clone(),
345                post_log: post_log.clone(),
346                init_count: init_a.clone(),
347            })
348            .with_gpu_plugin(RecordingPlugin {
349                prio: 100,
350                log: log.clone(),
351                post_log: post_log.clone(),
352                init_count: init_b.clone(),
353            });
354
355        let camera = Camera::default();
356        let ctx = GpuFrameContext {
357            camera: &camera,
358            viewport_size: glam::Vec2::new(800.0, 600.0),
359            dt: 1.0 / 60.0,
360            frame_index: 0,
361            viewport_id: None,
362        };
363
364        let bufs = runtime.pre_prepare(&device, &queue, &ctx);
365        assert!(bufs.is_empty(), "empty plugin returns should concatenate cleanly");
366        assert_eq!(*log.lock().unwrap(), vec![100, 200], "ascending priority order");
367        assert_eq!(*init_a.lock().unwrap(), 1);
368        assert_eq!(*init_b.lock().unwrap(), 1);
369
370        // Second frame: init_gpu is not called again.
371        let _ = runtime.pre_prepare(&device, &queue, &ctx);
372        assert_eq!(*init_a.lock().unwrap(), 1);
373        assert_eq!(*init_b.lock().unwrap(), 1);
374        assert_eq!(*log.lock().unwrap(), vec![100, 200, 100, 200]);
375
376        // post_paint: same priority contract.
377        let (_color_tex, color_view) = make_dummy_view(
378            &device,
379            wgpu::TextureFormat::Rgba8UnormSrgb,
380            wgpu::TextureUsages::RENDER_ATTACHMENT | wgpu::TextureUsages::TEXTURE_BINDING,
381        );
382        let (_depth_tex, depth_view) = make_dummy_view(
383            &device,
384            wgpu::TextureFormat::Depth32Float,
385            wgpu::TextureUsages::RENDER_ATTACHMENT | wgpu::TextureUsages::TEXTURE_BINDING,
386        );
387        let targets = PostPaintTargets {
388            color_view: &color_view,
389            depth_view: &depth_view,
390            pick_id_view: None,
391            color_format: wgpu::TextureFormat::Rgba8UnormSrgb,
392        };
393
394        let bufs = runtime.post_paint(&device, &queue, &targets, &ctx);
395        assert!(bufs.is_empty());
396        assert_eq!(*post_log.lock().unwrap(), vec![100, 200]);
397    }
398
399    struct DeviceLossRecorder {
400        recreated: Arc<Mutex<u32>>,
401        init_count: Arc<Mutex<u32>>,
402    }
403
404    impl GpuPlugin for DeviceLossRecorder {
405        fn init_gpu(&mut self, _device: &wgpu::Device) {
406            *self.init_count.lock().unwrap() += 1;
407        }
408        fn on_device_recreated(&mut self, _device: &wgpu::Device, _queue: &wgpu::Queue) {
409            *self.recreated.lock().unwrap() += 1;
410        }
411    }
412
413    #[test]
414    fn notify_device_recreated_invokes_hook_and_reinits() {
415        let Some((device, queue)) = try_make_device() else {
416            eprintln!("skipping: no wgpu adapter available");
417            return;
418        };
419
420        let recreated = Arc::new(Mutex::new(0));
421        let init_count = Arc::new(Mutex::new(0));
422
423        let mut runtime = crate::runtime::ViewportRuntime::new().with_gpu_plugin(DeviceLossRecorder {
424            recreated: recreated.clone(),
425            init_count: init_count.clone(),
426        });
427
428        let camera = Camera::default();
429        let ctx = GpuFrameContext {
430            camera: &camera,
431            viewport_size: glam::Vec2::new(1.0, 1.0),
432            dt: 0.0,
433            frame_index: 0,
434            viewport_id: None,
435        };
436
437        let _ = runtime.pre_prepare(&device, &queue, &ctx);
438        assert_eq!(*init_count.lock().unwrap(), 1);
439        assert_eq!(*recreated.lock().unwrap(), 0);
440
441        runtime.notify_device_recreated(&device, &queue);
442        assert_eq!(*recreated.lock().unwrap(), 1);
443
444        // Next pre_prepare re-runs init_gpu.
445        let _ = runtime.pre_prepare(&device, &queue, &ctx);
446        assert_eq!(*init_count.lock().unwrap(), 2);
447    }
448
449    #[test]
450    fn viewport_scoped_plugin_only_runs_for_matching_viewport() {
451        let Some((device, queue)) = try_make_device() else {
452            eprintln!("skipping: no wgpu adapter available");
453            return;
454        };
455
456        let log = Arc::new(Mutex::new(Vec::new()));
457        let post_log = Arc::new(Mutex::new(Vec::new()));
458        let init_count = Arc::new(Mutex::new(0));
459
460        // The scoped viewport is constructed directly: ViewportId is opaque,
461        // so we use `unsafe { std::mem::transmute }` would be wrong. Instead
462        // rely on the `pub(crate)` constructor by going through the renderer's
463        // `create_viewport` would need a renderer; for this test we synthesize
464        // a ViewportId via the `Default` representation: it's a newtype around
465        // `usize`, so `ViewportId(0)` is what `create_viewport` returns first.
466        // ViewportId's inner field is pub(crate); in-crate test code can
467        // construct synthetic ids without spinning up a full renderer.
468        let vp_a = crate::renderer::ViewportId(0);
469        let vp_b = crate::renderer::ViewportId(1);
470
471        let mut runtime = crate::runtime::ViewportRuntime::new().with_gpu_plugin_for_viewport(
472            vp_a,
473            RecordingPlugin {
474                prio: 100,
475                log: log.clone(),
476                post_log: post_log.clone(),
477                init_count: init_count.clone(),
478            },
479        );
480
481        let camera = Camera::default();
482        let ctx_a = GpuFrameContext {
483            camera: &camera,
484            viewport_size: glam::Vec2::new(1.0, 1.0),
485            dt: 0.0,
486            frame_index: 0,
487            viewport_id: Some(vp_a),
488        };
489        let ctx_b = GpuFrameContext { viewport_id: Some(vp_b), ..ctx_a };
490        let ctx_none = GpuFrameContext { viewport_id: None, ..ctx_a };
491
492        let _ = runtime.pre_prepare(&device, &queue, &ctx_a);
493        let _ = runtime.pre_prepare(&device, &queue, &ctx_b);
494        let _ = runtime.pre_prepare(&device, &queue, &ctx_none);
495
496        assert_eq!(*log.lock().unwrap(), vec![100], "only ran for vp_a");
497        // init_gpu still runs unconditionally on first frame.
498        assert_eq!(*init_count.lock().unwrap(), 1);
499    }
500}