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    /// A stable name for this plugin, used to key per-plugin timing in
199    /// [`RuntimeStats`](crate::RuntimeStats). Defaults to the concrete type
200    /// name; override to give a friendlier label.
201    fn type_name(&self) -> &'static str {
202        std::any::type_name::<Self>()
203    }
204
205    /// Build pipelines, persistent buffers, and bind group layouts.
206    ///
207    /// Called once before the plugin's first `pre_prepare`. If a new GPU
208    /// plugin is registered after the runtime has already run, every plugin's
209    /// `init_gpu` is invoked again on the next frame; implementations should
210    /// either be idempotent or guard their own one-time setup.
211    fn init_gpu(&mut self, _device: &wgpu::Device) {}
212
213    /// Called when the wgpu device is recreated, e.g. after device loss or a
214    /// host-driven reset. Every cached buffer, texture, bind group, or
215    /// pipeline the plugin built against the old device is now invalid and
216    /// must be rebuilt against `device` / `queue`.
217    ///
218    /// The host invokes this via
219    /// [`ViewportRuntime::notify_device_recreated`](super::ViewportRuntime::notify_device_recreated);
220    /// the runtime does not detect device loss on its own. After this call
221    /// returns, [`init_gpu`](Self::init_gpu) is also re-invoked on every
222    /// plugin before the next `pre_prepare`, so a typical implementation
223    /// can simply drop its cached resources here and let `init_gpu` rebuild
224    /// them.
225    fn on_device_recreated(&mut self, _device: &wgpu::Device, _queue: &wgpu::Queue) {}
226
227    /// Encode work that runs before `renderer.prepare()`.
228    ///
229    /// CPU plugins have already stepped this frame; their outputs are visible
230    /// through any shared state the consumer wired. The renderer has not
231    /// started its own work yet, so anything written here is observable by
232    /// `prepare()`'s passes (cluster build, shadow render, etc.).
233    fn pre_prepare(
234        &mut self,
235        _device: &wgpu::Device,
236        _queue: &wgpu::Queue,
237        _ctx: &GpuFrameContext<'_>,
238    ) -> Vec<wgpu::CommandBuffer> {
239        Vec::new()
240    }
241
242    /// Encode work that runs after `renderer.paint_to()`.
243    ///
244    /// The plugin receives views of the just-rendered color, depth, and
245    /// (optionally) pick-id targets and may sample them in a compute or
246    /// fullscreen render pass. Typical uses: custom AO, motion blur,
247    /// screen-space outlines, color grading LUTs.
248    ///
249    /// To *modify* the color target, write to a sibling texture the plugin
250    /// owns and let the host composite it during the final blit. The lib
251    /// does not loop a plugin's output back into the rendered color.
252    fn post_paint(
253        &mut self,
254        _device: &wgpu::Device,
255        _queue: &wgpu::Queue,
256        _targets: &PostPaintTargets<'_>,
257        _ctx: &GpuFrameContext<'_>,
258    ) -> Vec<wgpu::CommandBuffer> {
259        Vec::new()
260    }
261}
262
263/// Internal adapter that filters a [`GpuPlugin`] to a single viewport.
264///
265/// Constructed by
266/// [`ViewportRuntime::with_gpu_plugin_for_viewport`](super::ViewportRuntime::with_gpu_plugin_for_viewport).
267/// The wrapper forwards `priority`, `init_gpu`, and `on_device_recreated`
268/// unconditionally, but only forwards `pre_prepare` / `post_paint` when
269/// `ctx.viewport_id == Some(scoped_viewport)`.
270pub(super) struct ViewportScopedPlugin<P: GpuPlugin> {
271    pub(super) viewport: ViewportId,
272    pub(super) inner: P,
273}
274
275impl<P: GpuPlugin> GpuPlugin for ViewportScopedPlugin<P> {
276    fn priority(&self) -> i32 {
277        self.inner.priority()
278    }
279
280    fn init_gpu(&mut self, device: &wgpu::Device) {
281        self.inner.init_gpu(device);
282    }
283
284    fn on_device_recreated(&mut self, device: &wgpu::Device, queue: &wgpu::Queue) {
285        self.inner.on_device_recreated(device, queue);
286    }
287
288    fn pre_prepare(
289        &mut self,
290        device: &wgpu::Device,
291        queue: &wgpu::Queue,
292        ctx: &GpuFrameContext<'_>,
293    ) -> Vec<wgpu::CommandBuffer> {
294        if ctx.viewport_id == Some(self.viewport) {
295            self.inner.pre_prepare(device, queue, ctx)
296        } else {
297            Vec::new()
298        }
299    }
300
301    fn post_paint(
302        &mut self,
303        device: &wgpu::Device,
304        queue: &wgpu::Queue,
305        targets: &PostPaintTargets<'_>,
306        ctx: &GpuFrameContext<'_>,
307    ) -> Vec<wgpu::CommandBuffer> {
308        if ctx.viewport_id == Some(self.viewport) {
309            self.inner.post_paint(device, queue, targets, ctx)
310        } else {
311            Vec::new()
312        }
313    }
314}
315
316#[cfg(test)]
317mod tests {
318    use super::*;
319    use std::sync::{Arc, Mutex};
320
321    struct RecordingPlugin {
322        prio: i32,
323        log: Arc<Mutex<Vec<i32>>>,
324        post_log: Arc<Mutex<Vec<i32>>>,
325        init_count: Arc<Mutex<u32>>,
326    }
327
328    impl GpuPlugin for RecordingPlugin {
329        fn priority(&self) -> i32 {
330            self.prio
331        }
332
333        fn init_gpu(&mut self, _device: &wgpu::Device) {
334            *self.init_count.lock().unwrap() += 1;
335        }
336
337        fn pre_prepare(
338            &mut self,
339            _device: &wgpu::Device,
340            _queue: &wgpu::Queue,
341            _ctx: &GpuFrameContext<'_>,
342        ) -> Vec<wgpu::CommandBuffer> {
343            self.log.lock().unwrap().push(self.prio);
344            Vec::new()
345        }
346
347        fn post_paint(
348            &mut self,
349            _device: &wgpu::Device,
350            _queue: &wgpu::Queue,
351            _targets: &PostPaintTargets<'_>,
352            _ctx: &GpuFrameContext<'_>,
353        ) -> Vec<wgpu::CommandBuffer> {
354            self.post_log.lock().unwrap().push(self.prio);
355            Vec::new()
356        }
357    }
358
359    fn make_dummy_view(
360        device: &wgpu::Device,
361        format: wgpu::TextureFormat,
362        usage: wgpu::TextureUsages,
363    ) -> (wgpu::Texture, wgpu::TextureView) {
364        let tex = device.create_texture(&wgpu::TextureDescriptor {
365            label: Some("post_paint_test_target"),
366            size: wgpu::Extent3d {
367                width: 1,
368                height: 1,
369                depth_or_array_layers: 1,
370            },
371            mip_level_count: 1,
372            sample_count: 1,
373            dimension: wgpu::TextureDimension::D2,
374            format,
375            usage,
376            view_formats: &[],
377        });
378        let view = tex.create_view(&wgpu::TextureViewDescriptor::default());
379        (tex, view)
380    }
381
382    fn try_make_device() -> Option<(wgpu::Device, wgpu::Queue)> {
383        let instance = wgpu::Instance::new(&wgpu::InstanceDescriptor::default());
384        let adapter = pollster::block_on(instance.request_adapter(&wgpu::RequestAdapterOptions {
385            power_preference: wgpu::PowerPreference::LowPower,
386            compatible_surface: None,
387            force_fallback_adapter: false,
388        }))
389        .ok()?;
390        pollster::block_on(adapter.request_device(&wgpu::DeviceDescriptor::default())).ok()
391    }
392
393    #[test]
394    fn priority_order_and_init_once() {
395        let Some((device, queue)) = try_make_device() else {
396            eprintln!("skipping: no wgpu adapter available");
397            return;
398        };
399
400        let log = Arc::new(Mutex::new(Vec::new()));
401        let post_log = Arc::new(Mutex::new(Vec::new()));
402        let init_a = Arc::new(Mutex::new(0));
403        let init_b = Arc::new(Mutex::new(0));
404
405        let mut runtime = crate::runtime::ViewportRuntime::new()
406            .with_gpu_plugin(RecordingPlugin {
407                prio: 200,
408                log: log.clone(),
409                post_log: post_log.clone(),
410                init_count: init_a.clone(),
411            })
412            .with_gpu_plugin(RecordingPlugin {
413                prio: 100,
414                log: log.clone(),
415                post_log: post_log.clone(),
416                init_count: init_b.clone(),
417            });
418
419        let camera = Camera::default();
420        let ctx = GpuFrameContext::new(&camera, glam::Vec2::new(800.0, 600.0), 1.0 / 60.0, 0);
421
422        let bufs = runtime.pre_prepare(&device, &queue, &ctx);
423        assert!(
424            bufs.is_empty(),
425            "empty plugin returns should concatenate cleanly"
426        );
427        assert_eq!(
428            *log.lock().unwrap(),
429            vec![100, 200],
430            "ascending priority order"
431        );
432        assert_eq!(*init_a.lock().unwrap(), 1);
433        assert_eq!(*init_b.lock().unwrap(), 1);
434
435        // Second frame: init_gpu is not called again.
436        let _ = runtime.pre_prepare(&device, &queue, &ctx);
437        assert_eq!(*init_a.lock().unwrap(), 1);
438        assert_eq!(*init_b.lock().unwrap(), 1);
439        assert_eq!(*log.lock().unwrap(), vec![100, 200, 100, 200]);
440
441        // post_paint: same priority contract.
442        let (_color_tex, color_view) = make_dummy_view(
443            &device,
444            wgpu::TextureFormat::Rgba8UnormSrgb,
445            wgpu::TextureUsages::RENDER_ATTACHMENT | wgpu::TextureUsages::TEXTURE_BINDING,
446        );
447        let (_depth_tex, depth_view) = make_dummy_view(
448            &device,
449            wgpu::TextureFormat::Depth32Float,
450            wgpu::TextureUsages::RENDER_ATTACHMENT | wgpu::TextureUsages::TEXTURE_BINDING,
451        );
452        let targets = PostPaintTargets::new(
453            &color_view,
454            &depth_view,
455            wgpu::TextureFormat::Rgba8UnormSrgb,
456        );
457
458        let bufs = runtime.post_paint(&device, &queue, &targets, &ctx);
459        assert!(bufs.is_empty());
460        assert_eq!(*post_log.lock().unwrap(), vec![100, 200]);
461    }
462
463    struct DeviceLossRecorder {
464        recreated: Arc<Mutex<u32>>,
465        init_count: Arc<Mutex<u32>>,
466    }
467
468    impl GpuPlugin for DeviceLossRecorder {
469        fn init_gpu(&mut self, _device: &wgpu::Device) {
470            *self.init_count.lock().unwrap() += 1;
471        }
472        fn on_device_recreated(&mut self, _device: &wgpu::Device, _queue: &wgpu::Queue) {
473            *self.recreated.lock().unwrap() += 1;
474        }
475    }
476
477    #[test]
478    fn notify_device_recreated_invokes_hook_and_reinits() {
479        let Some((device, queue)) = try_make_device() else {
480            eprintln!("skipping: no wgpu adapter available");
481            return;
482        };
483
484        let recreated = Arc::new(Mutex::new(0));
485        let init_count = Arc::new(Mutex::new(0));
486
487        let mut runtime =
488            crate::runtime::ViewportRuntime::new().with_gpu_plugin(DeviceLossRecorder {
489                recreated: recreated.clone(),
490                init_count: init_count.clone(),
491            });
492
493        let camera = Camera::default();
494        let ctx = GpuFrameContext::new(&camera, glam::Vec2::new(1.0, 1.0), 0.0, 0);
495
496        let _ = runtime.pre_prepare(&device, &queue, &ctx);
497        assert_eq!(*init_count.lock().unwrap(), 1);
498        assert_eq!(*recreated.lock().unwrap(), 0);
499
500        runtime.notify_device_recreated(&device, &queue);
501        assert_eq!(*recreated.lock().unwrap(), 1);
502
503        // Next pre_prepare re-runs init_gpu.
504        let _ = runtime.pre_prepare(&device, &queue, &ctx);
505        assert_eq!(*init_count.lock().unwrap(), 2);
506    }
507
508    #[test]
509    fn viewport_scoped_plugin_only_runs_for_matching_viewport() {
510        let Some((device, queue)) = try_make_device() else {
511            eprintln!("skipping: no wgpu adapter available");
512            return;
513        };
514
515        let log = Arc::new(Mutex::new(Vec::new()));
516        let post_log = Arc::new(Mutex::new(Vec::new()));
517        let init_count = Arc::new(Mutex::new(0));
518
519        // The scoped viewport is constructed directly: ViewportId is opaque,
520        // so we use `unsafe { std::mem::transmute }` would be wrong. Instead
521        // rely on the `pub(crate)` constructor by going through the renderer's
522        // `create_viewport` would need a renderer; for this test we synthesize
523        // a ViewportId via the `Default` representation: it's a newtype around
524        // `usize`, so `ViewportId(0)` is what `create_viewport` returns first.
525        // ViewportId's inner field is pub(crate); in-crate test code can
526        // construct synthetic ids without spinning up a full renderer.
527        let vp_a = crate::renderer::ViewportId(0);
528        let vp_b = crate::renderer::ViewportId(1);
529
530        let mut runtime = crate::runtime::ViewportRuntime::new().with_gpu_plugin_for_viewport(
531            vp_a,
532            RecordingPlugin {
533                prio: 100,
534                log: log.clone(),
535                post_log: post_log.clone(),
536                init_count: init_count.clone(),
537            },
538        );
539
540        let camera = Camera::default();
541        let sz = glam::Vec2::new(1.0, 1.0);
542        let ctx_a = GpuFrameContext::new(&camera, sz, 0.0, 0).with_viewport(vp_a);
543        let ctx_b = GpuFrameContext::new(&camera, sz, 0.0, 0).with_viewport(vp_b);
544        let ctx_none = GpuFrameContext::new(&camera, sz, 0.0, 0);
545
546        let _ = runtime.pre_prepare(&device, &queue, &ctx_a);
547        let _ = runtime.pre_prepare(&device, &queue, &ctx_b);
548        let _ = runtime.pre_prepare(&device, &queue, &ctx_none);
549
550        assert_eq!(*log.lock().unwrap(), vec![100], "only ran for vp_a");
551        // init_gpu still runs unconditionally on first frame.
552        assert_eq!(*init_count.lock().unwrap(), 1);
553    }
554}