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, SkinningPlugin::attach_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 {
307                width: 1,
308                height: 1,
309                depth_or_array_layers: 1,
310            },
311            mip_level_count: 1,
312            sample_count: 1,
313            dimension: wgpu::TextureDimension::D2,
314            format,
315            usage,
316            view_formats: &[],
317        });
318        let view = tex.create_view(&wgpu::TextureViewDescriptor::default());
319        (tex, view)
320    }
321
322    fn try_make_device() -> Option<(wgpu::Device, wgpu::Queue)> {
323        let instance = wgpu::Instance::new(&wgpu::InstanceDescriptor::default());
324        let adapter = pollster::block_on(instance.request_adapter(&wgpu::RequestAdapterOptions {
325            power_preference: wgpu::PowerPreference::LowPower,
326            compatible_surface: None,
327            force_fallback_adapter: false,
328        }))
329        .ok()?;
330        pollster::block_on(adapter.request_device(&wgpu::DeviceDescriptor::default())).ok()
331    }
332
333    #[test]
334    fn priority_order_and_init_once() {
335        let Some((device, queue)) = try_make_device() else {
336            eprintln!("skipping: no wgpu adapter available");
337            return;
338        };
339
340        let log = Arc::new(Mutex::new(Vec::new()));
341        let post_log = Arc::new(Mutex::new(Vec::new()));
342        let init_a = Arc::new(Mutex::new(0));
343        let init_b = Arc::new(Mutex::new(0));
344
345        let mut runtime = crate::runtime::ViewportRuntime::new()
346            .with_gpu_plugin(RecordingPlugin {
347                prio: 200,
348                log: log.clone(),
349                post_log: post_log.clone(),
350                init_count: init_a.clone(),
351            })
352            .with_gpu_plugin(RecordingPlugin {
353                prio: 100,
354                log: log.clone(),
355                post_log: post_log.clone(),
356                init_count: init_b.clone(),
357            });
358
359        let camera = Camera::default();
360        let ctx = GpuFrameContext {
361            camera: &camera,
362            viewport_size: glam::Vec2::new(800.0, 600.0),
363            dt: 1.0 / 60.0,
364            frame_index: 0,
365            viewport_id: None,
366        };
367
368        let bufs = runtime.pre_prepare(&device, &queue, &ctx);
369        assert!(
370            bufs.is_empty(),
371            "empty plugin returns should concatenate cleanly"
372        );
373        assert_eq!(
374            *log.lock().unwrap(),
375            vec![100, 200],
376            "ascending priority order"
377        );
378        assert_eq!(*init_a.lock().unwrap(), 1);
379        assert_eq!(*init_b.lock().unwrap(), 1);
380
381        // Second frame: init_gpu is not called again.
382        let _ = runtime.pre_prepare(&device, &queue, &ctx);
383        assert_eq!(*init_a.lock().unwrap(), 1);
384        assert_eq!(*init_b.lock().unwrap(), 1);
385        assert_eq!(*log.lock().unwrap(), vec![100, 200, 100, 200]);
386
387        // post_paint: same priority contract.
388        let (_color_tex, color_view) = make_dummy_view(
389            &device,
390            wgpu::TextureFormat::Rgba8UnormSrgb,
391            wgpu::TextureUsages::RENDER_ATTACHMENT | wgpu::TextureUsages::TEXTURE_BINDING,
392        );
393        let (_depth_tex, depth_view) = make_dummy_view(
394            &device,
395            wgpu::TextureFormat::Depth32Float,
396            wgpu::TextureUsages::RENDER_ATTACHMENT | wgpu::TextureUsages::TEXTURE_BINDING,
397        );
398        let targets = PostPaintTargets {
399            color_view: &color_view,
400            depth_view: &depth_view,
401            pick_id_view: None,
402            color_format: wgpu::TextureFormat::Rgba8UnormSrgb,
403        };
404
405        let bufs = runtime.post_paint(&device, &queue, &targets, &ctx);
406        assert!(bufs.is_empty());
407        assert_eq!(*post_log.lock().unwrap(), vec![100, 200]);
408    }
409
410    struct DeviceLossRecorder {
411        recreated: Arc<Mutex<u32>>,
412        init_count: Arc<Mutex<u32>>,
413    }
414
415    impl GpuPlugin for DeviceLossRecorder {
416        fn init_gpu(&mut self, _device: &wgpu::Device) {
417            *self.init_count.lock().unwrap() += 1;
418        }
419        fn on_device_recreated(&mut self, _device: &wgpu::Device, _queue: &wgpu::Queue) {
420            *self.recreated.lock().unwrap() += 1;
421        }
422    }
423
424    #[test]
425    fn notify_device_recreated_invokes_hook_and_reinits() {
426        let Some((device, queue)) = try_make_device() else {
427            eprintln!("skipping: no wgpu adapter available");
428            return;
429        };
430
431        let recreated = Arc::new(Mutex::new(0));
432        let init_count = Arc::new(Mutex::new(0));
433
434        let mut runtime =
435            crate::runtime::ViewportRuntime::new().with_gpu_plugin(DeviceLossRecorder {
436                recreated: recreated.clone(),
437                init_count: init_count.clone(),
438            });
439
440        let camera = Camera::default();
441        let ctx = GpuFrameContext {
442            camera: &camera,
443            viewport_size: glam::Vec2::new(1.0, 1.0),
444            dt: 0.0,
445            frame_index: 0,
446            viewport_id: None,
447        };
448
449        let _ = runtime.pre_prepare(&device, &queue, &ctx);
450        assert_eq!(*init_count.lock().unwrap(), 1);
451        assert_eq!(*recreated.lock().unwrap(), 0);
452
453        runtime.notify_device_recreated(&device, &queue);
454        assert_eq!(*recreated.lock().unwrap(), 1);
455
456        // Next pre_prepare re-runs init_gpu.
457        let _ = runtime.pre_prepare(&device, &queue, &ctx);
458        assert_eq!(*init_count.lock().unwrap(), 2);
459    }
460
461    #[test]
462    fn viewport_scoped_plugin_only_runs_for_matching_viewport() {
463        let Some((device, queue)) = try_make_device() else {
464            eprintln!("skipping: no wgpu adapter available");
465            return;
466        };
467
468        let log = Arc::new(Mutex::new(Vec::new()));
469        let post_log = Arc::new(Mutex::new(Vec::new()));
470        let init_count = Arc::new(Mutex::new(0));
471
472        // The scoped viewport is constructed directly: ViewportId is opaque,
473        // so we use `unsafe { std::mem::transmute }` would be wrong. Instead
474        // rely on the `pub(crate)` constructor by going through the renderer's
475        // `create_viewport` would need a renderer; for this test we synthesize
476        // a ViewportId via the `Default` representation: it's a newtype around
477        // `usize`, so `ViewportId(0)` is what `create_viewport` returns first.
478        // ViewportId's inner field is pub(crate); in-crate test code can
479        // construct synthetic ids without spinning up a full renderer.
480        let vp_a = crate::renderer::ViewportId(0);
481        let vp_b = crate::renderer::ViewportId(1);
482
483        let mut runtime = crate::runtime::ViewportRuntime::new().with_gpu_plugin_for_viewport(
484            vp_a,
485            RecordingPlugin {
486                prio: 100,
487                log: log.clone(),
488                post_log: post_log.clone(),
489                init_count: init_count.clone(),
490            },
491        );
492
493        let camera = Camera::default();
494        let ctx_a = GpuFrameContext {
495            camera: &camera,
496            viewport_size: glam::Vec2::new(1.0, 1.0),
497            dt: 0.0,
498            frame_index: 0,
499            viewport_id: Some(vp_a),
500        };
501        let ctx_b = GpuFrameContext {
502            viewport_id: Some(vp_b),
503            ..ctx_a
504        };
505        let ctx_none = GpuFrameContext {
506            viewport_id: None,
507            ..ctx_a
508        };
509
510        let _ = runtime.pre_prepare(&device, &queue, &ctx_a);
511        let _ = runtime.pre_prepare(&device, &queue, &ctx_b);
512        let _ = runtime.pre_prepare(&device, &queue, &ctx_none);
513
514        assert_eq!(*log.lock().unwrap(), vec![100], "only ran for vp_a");
515        // init_gpu still runs unconditionally on first frame.
516        assert_eq!(*init_count.lock().unwrap(), 1);
517    }
518}