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