re_renderer/draw_phases/
picking_layer.rs

1//! GPU driven picking.
2//!
3//! This module provides the [`PickingLayerProcessor`] which is responsible for rendering & processing the picking layer.
4//! Picking is done in a separate pass to a as-small-as needed render target (size is user configurable).
5//!
6//! The picking layer is a RGBA texture with 32bit per channel, the red & green channel are used for the [`PickingLayerObjectId`],
7//! the blue & alpha channel are used for the [`PickingLayerInstanceId`].
8//! (Keep in mind that GPUs are little endian, so R will have the lower bytes and G the higher ones)
9//!
10//! In order to accomplish small render targets, the projection matrix is cropped to only render the area of interest.
11
12use re_log::debug_assert_eq;
13use re_mutex::Mutex;
14use smallvec::smallvec;
15
16use crate::allocator::create_and_fill_uniform_buffer;
17use crate::global_bindings::FrameUniformBuffer;
18use crate::rect::RectF32;
19use crate::texture_info::Texture2DBufferInfo;
20use crate::transform::{RectTransform, ndc_from_pixel};
21use crate::view_builder::ViewBuilder;
22use crate::wgpu_resources::{
23    BindGroupDesc, BindGroupEntry, BindGroupLayoutDesc, GpuBindGroup, GpuRenderPipelineHandle,
24    GpuRenderPipelinePoolAccessor, GpuTexture, GpuTextureHandle, PipelineLayoutDesc, PoolError,
25    RenderPipelineDesc, TextureDesc,
26};
27use crate::{
28    GpuReadbackBuffer, GpuReadbackIdentifier, Label, RectInt, RenderContext, include_shader_module,
29};
30
31/// GPU retrieved & processed picking data result.
32pub struct PickingResult {
33    /// Picking rect supplied on picking request.
34    /// Describes the area of the picking layer that was read back.
35    pub rect: RectInt,
36
37    /// Picking id data for the requested rectangle.
38    ///
39    /// GPU internal row padding has already been removed from this buffer.
40    /// Pixel data is stored in the normal fashion - row wise, left to right, top to bottom.
41    pub picking_id_data: Vec<PickingLayerId>,
42
43    /// Picking depth data for the requested rectangle.
44    ///
45    /// Use [`PickingResult::picked_world_position`] for easy interpretation of the data.
46    ///
47    /// GPU internal row padding has already been removed from this buffer.
48    /// Pixel data is stored in the normal fashion - row wise, left to right, top to bottom.
49    pub picking_depth_data: Vec<f32>,
50
51    /// Transforms a NDC position on the picking rect to a world position.
52    world_from_cropped_projection: glam::Mat4,
53}
54
55impl PickingResult {
56    /// Returns the picked world position.
57    ///
58    /// Panics if the position is outside of the picking rect.
59    ///
60    /// Keep in mind that the picked position may be (negative) infinity if nothing was picked.
61    #[inline]
62    pub fn picked_world_position(&self, pos_on_picking_rect: glam::UVec2) -> glam::Vec3 {
63        let raw_depth = self.picking_depth_data
64            [(pos_on_picking_rect.y * self.rect.width() + pos_on_picking_rect.x) as usize];
65
66        self.world_from_cropped_projection.project_point3(
67            ndc_from_pixel(pos_on_picking_rect.as_vec2(), self.rect.extent).extend(raw_depth),
68        )
69    }
70
71    /// Returns the picked picking id.
72    ///
73    /// Panics if the position is outside of the picking rect.
74    #[inline]
75    pub fn picked_id(&self, pos_on_picking_rect: glam::UVec2) -> PickingLayerId {
76        self.picking_id_data
77            [(pos_on_picking_rect.y * self.rect.width() + pos_on_picking_rect.x) as usize]
78    }
79}
80
81/// Type used as user data on the gpu readback belt.
82struct ReadbackBeltMetadata {
83    picking_rect: RectInt,
84    world_from_cropped_projection: glam::Mat4,
85
86    depth_readback_workaround_in_use: bool,
87}
88
89/// The first 64bit of the picking layer.
90///
91/// Typically used to identify higher level objects
92/// Some renderers might allow to change this part of the picking identifier only at a coarse grained level.
93#[repr(C)]
94#[derive(Clone, Copy, bytemuck::Zeroable, bytemuck::Pod, Default, Debug, PartialEq, Eq)]
95pub struct PickingLayerObjectId(pub u64);
96
97/// The second 64bit of the picking layer.
98///
99/// Typically used to identify instances.
100/// Some renderers might allow to change only this part of the picking identifier at a fine grained level.
101#[repr(C)]
102#[derive(Clone, Copy, bytemuck::Zeroable, bytemuck::Pod, Default, Debug, PartialEq, Eq)]
103pub struct PickingLayerInstanceId(pub u64);
104
105impl re_byte_size::SizeBytes for PickingLayerInstanceId {
106    #[inline]
107    fn heap_size_bytes(&self) -> u64 {
108        0
109    }
110
111    #[inline]
112    fn is_pod() -> bool {
113        true
114    }
115}
116
117/// Combination of `PickingLayerObjectId` and `PickingLayerInstanceId`.
118///
119/// This is the same memory order as it is found in the GPU picking layer texture.
120#[repr(C)]
121#[derive(Clone, Copy, bytemuck::Zeroable, bytemuck::Pod, Default, Debug, PartialEq, Eq)]
122pub struct PickingLayerId {
123    pub object: PickingLayerObjectId,
124    pub instance: PickingLayerInstanceId,
125}
126
127impl From<PickingLayerId> for [u32; 4] {
128    fn from(val: PickingLayerId) -> Self {
129        [
130            val.object.0 as u32,
131            (val.object.0 >> 32) as u32,
132            val.instance.0 as u32,
133            (val.instance.0 >> 32) as u32,
134        ]
135    }
136}
137
138#[derive(thiserror::Error, Debug)]
139pub enum PickingLayerError {
140    #[error(transparent)]
141    ReadbackError(#[from] crate::allocator::GpuReadbackError),
142
143    #[error(transparent)]
144    ResourcePoolError(#[from] PoolError),
145}
146
147/// Manages the rendering of the picking layer pass, its render targets & readback buffer.
148///
149/// The view builder creates this for every frame that requests a picking result.
150pub struct PickingLayerProcessor {
151    pub picking_target: GpuTexture,
152    picking_depth_target: GpuTexture,
153    readback_buffer: Mutex<GpuReadbackBuffer>,
154    bind_group_0: GpuBindGroup,
155
156    depth_readback_workaround: Option<DepthReadbackWorkaround>,
157}
158
159impl PickingLayerProcessor {
160    /// The texture format used for the picking layer.
161    pub const PICKING_LAYER_FORMAT: wgpu::TextureFormat = wgpu::TextureFormat::Rgba32Uint;
162
163    /// The depth format used for the picking layer - f32 makes it easiest to deal with retrieved depth and is guaranteed to be copyable.
164    pub const PICKING_LAYER_DEPTH_FORMAT: wgpu::TextureFormat = wgpu::TextureFormat::Depth32Float;
165
166    pub const PICKING_LAYER_MSAA_STATE: wgpu::MultisampleState = wgpu::MultisampleState {
167        count: 1,
168        mask: !0,
169        alpha_to_coverage_enabled: false,
170    };
171
172    pub const PICKING_LAYER_DEPTH_STATE: Option<wgpu::DepthStencilState> =
173        Some(ViewBuilder::MAIN_TARGET_DEFAULT_DEPTH_STATE);
174
175    /// New picking layer for a given screen.
176    ///
177    /// Note that out-of-bounds rectangles *are* allowed, the picking layer will *not* be clipped to the screen.
178    /// This means that the content of the picking layer rectangle will behave as-if the screen was bigger,
179    /// containing valid picking data.
180    /// It's up to the user when interpreting the picking data to do any required clipping.
181    ///
182    /// `enable_picking_target_sampling` should be enabled only for debugging purposes.
183    /// It allows to sample the picking layer texture in a shader.
184    pub fn new(
185        ctx: &RenderContext,
186        view_name: &Label,
187        screen_resolution: glam::UVec2,
188        picking_rect: RectInt,
189        frame_uniform_buffer_content: &FrameUniformBuffer,
190        enable_picking_target_sampling: bool,
191        readback_identifier: GpuReadbackIdentifier,
192    ) -> Self {
193        let mut picking_target_usage =
194            wgpu::TextureUsages::RENDER_ATTACHMENT | wgpu::TextureUsages::COPY_SRC;
195        picking_target_usage.set(
196            wgpu::TextureUsages::TEXTURE_BINDING,
197            enable_picking_target_sampling,
198        );
199
200        let picking_target = ctx.gpu_resources.textures.alloc(
201            &ctx.device,
202            &TextureDesc {
203                label: format!("{view_name} - PickingLayerProcessor").into(),
204                size: picking_rect.wgpu_extent(),
205                mip_level_count: 1,
206                sample_count: 1,
207                dimension: wgpu::TextureDimension::D2,
208                format: Self::PICKING_LAYER_FORMAT,
209                usage: picking_target_usage,
210            },
211        );
212
213        let direct_depth_readback = ctx.device_caps().tier.support_depth_readback();
214
215        let picking_depth_target = ctx.gpu_resources.textures.alloc(
216            &ctx.device,
217            &TextureDesc {
218                label: format!("{view_name} - picking_layer depth target").into(),
219                format: Self::PICKING_LAYER_DEPTH_FORMAT,
220                usage: if direct_depth_readback {
221                    wgpu::TextureUsages::RENDER_ATTACHMENT | wgpu::TextureUsages::COPY_SRC
222                } else {
223                    wgpu::TextureUsages::RENDER_ATTACHMENT | wgpu::TextureUsages::TEXTURE_BINDING
224                },
225                ..picking_target.creation_desc
226            },
227        );
228
229        let depth_readback_workaround = (!direct_depth_readback).then(|| {
230            DepthReadbackWorkaround::new(ctx, picking_rect.extent, picking_depth_target.handle)
231        });
232
233        let cropped_projection_from_projection = RectTransform {
234            region_of_interest: picking_rect.into(),
235            region: RectF32 {
236                min: glam::Vec2::ZERO,
237                extent: screen_resolution.as_vec2(),
238            },
239        }
240        .to_ndc_scale_and_translation();
241
242        // Setup frame uniform buffer
243        let previous_projection_from_world: glam::Mat4 =
244            frame_uniform_buffer_content.projection_from_world.into();
245        let cropped_projection_from_world =
246            cropped_projection_from_projection * previous_projection_from_world;
247        let previous_projection_from_view: glam::Mat4 =
248            frame_uniform_buffer_content.projection_from_view.into();
249        let cropped_projection_from_view =
250            cropped_projection_from_projection * previous_projection_from_view;
251
252        let frame_uniform_buffer_content = FrameUniformBuffer {
253            projection_from_world: cropped_projection_from_world.into(),
254            projection_from_view: cropped_projection_from_view.into(),
255            ..*frame_uniform_buffer_content
256        };
257
258        let frame_uniform_buffer = create_and_fill_uniform_buffer(
259            ctx,
260            format!("{view_name} - picking_layer frame uniform buffer").into(),
261            frame_uniform_buffer_content,
262        );
263
264        let bind_group_0 = ctx.global_bindings.create_bind_group(
265            &ctx.gpu_resources,
266            &ctx.device,
267            frame_uniform_buffer,
268        );
269
270        let row_info_id =
271            Texture2DBufferInfo::new(Self::PICKING_LAYER_FORMAT, picking_rect.wgpu_extent());
272        let row_info_depth = Texture2DBufferInfo::new(
273            if direct_depth_readback {
274                Self::PICKING_LAYER_DEPTH_FORMAT
275            } else {
276                DepthReadbackWorkaround::READBACK_FORMAT
277            },
278            picking_rect.wgpu_extent(),
279        );
280
281        // Offset of the depth buffer in the readback buffer needs to be aligned to size of a depth pixel.
282        // This is "trivially true" if the size of the depth format is a multiple of the size of the id format.
283        re_log::debug_assert!(
284            Self::PICKING_LAYER_FORMAT.block_copy_size(None).unwrap()
285                % Self::PICKING_LAYER_DEPTH_FORMAT
286                    .block_copy_size(Some(wgpu::TextureAspect::DepthOnly))
287                    .unwrap()
288                == 0
289        );
290        let buffer_size = row_info_id.buffer_size_padded + row_info_depth.buffer_size_padded;
291
292        let readback_buffer = Mutex::new(ctx.gpu_readback_belt.lock().allocate(
293            &ctx.device,
294            &ctx.gpu_resources.buffers,
295            buffer_size,
296            readback_identifier,
297            Box::new(ReadbackBeltMetadata {
298                picking_rect,
299                world_from_cropped_projection: cropped_projection_from_world.inverse(),
300                depth_readback_workaround_in_use: depth_readback_workaround.is_some(),
301            }),
302        ));
303
304        Self {
305            picking_target,
306            picking_depth_target,
307            readback_buffer,
308            bind_group_0,
309            depth_readback_workaround,
310        }
311    }
312
313    pub fn begin_render_pass<'a>(
314        &'a self,
315        view_name: &Label,
316        encoder: &'a mut wgpu::CommandEncoder,
317    ) -> wgpu::RenderPass<'a> {
318        re_tracing::profile_function!();
319
320        let mut pass = encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
321            label: Label::from(format!("{view_name} - picking_layer pass")).wgpu_label(),
322            color_attachments: &[Some(wgpu::RenderPassColorAttachment {
323                view: &self.picking_target.default_view,
324                depth_slice: None,
325                resolve_target: None,
326                ops: wgpu::Operations {
327                    load: wgpu::LoadOp::Clear(wgpu::Color::TRANSPARENT),
328                    store: wgpu::StoreOp::Store, // Store for readback!
329                },
330            })],
331            depth_stencil_attachment: Some(wgpu::RenderPassDepthStencilAttachment {
332                view: &self.picking_depth_target.default_view,
333                depth_ops: Some(wgpu::Operations {
334                    load: ViewBuilder::DEFAULT_DEPTH_CLEAR,
335                    store: wgpu::StoreOp::Store, // Store for readback!
336                }),
337                stencil_ops: None,
338            }),
339            timestamp_writes: None,
340            occlusion_query_set: None,
341            multiview_mask: None,
342        });
343
344        pass.set_bind_group(0, &self.bind_group_0, &[]);
345
346        pass
347    }
348
349    pub fn end_render_pass(
350        &self,
351        encoder: &mut wgpu::CommandEncoder,
352        render_pipelines: &GpuRenderPipelinePoolAccessor<'_>,
353    ) -> Result<(), PickingLayerError> {
354        let extent = self.picking_target.texture.size();
355
356        let readable_depth_texture =
357            if let Some(depth_copy_workaround) = self.depth_readback_workaround.as_ref() {
358                depth_copy_workaround.copy_to_readable_texture(
359                    encoder,
360                    render_pipelines,
361                    &self.bind_group_0,
362                )?
363            } else {
364                &self.picking_depth_target
365            };
366
367        self.readback_buffer.lock().read_multiple_texture2d(
368            encoder,
369            &[
370                (
371                    wgpu::TexelCopyTextureInfo {
372                        texture: &self.picking_target.texture,
373                        mip_level: 0,
374                        origin: wgpu::Origin3d::ZERO,
375                        aspect: wgpu::TextureAspect::All,
376                    },
377                    extent,
378                ),
379                (
380                    wgpu::TexelCopyTextureInfo {
381                        texture: &readable_depth_texture.texture,
382                        mip_level: 0,
383                        origin: wgpu::Origin3d::ZERO,
384                        aspect: if self.depth_readback_workaround.is_some() {
385                            wgpu::TextureAspect::All
386                        } else {
387                            wgpu::TextureAspect::DepthOnly
388                        },
389                    },
390                    extent,
391                ),
392            ],
393        )?;
394
395        Ok(())
396    }
397
398    /// Returns the latest available picking result for a given identifier.
399    ///
400    /// Ready data that hasn't been retrieved for more than a frame will be discarded.
401    ///
402    /// See also [`crate::view_builder::ViewPickingConfiguration`]
403    pub fn readback_result(
404        ctx: &RenderContext,
405        identifier: GpuReadbackIdentifier,
406    ) -> Option<PickingResult> {
407        ctx.gpu_readback_belt
408            .lock()
409            .readback_newest_available(identifier, |data, metadata: Box<ReadbackBeltMetadata>| {
410                // Assert that our texture data reinterpretation works out from a pixel size point of view.
411                debug_assert_eq!(
412                    Self::PICKING_LAYER_DEPTH_FORMAT
413                        .block_copy_size(Some(wgpu::TextureAspect::DepthOnly))
414                        .unwrap(),
415                    std::mem::size_of::<f32>() as u32
416                );
417                debug_assert_eq!(
418                    Self::PICKING_LAYER_FORMAT.block_copy_size(None).unwrap() as usize,
419                    std::mem::size_of::<PickingLayerId>()
420                );
421
422                let buffer_info_id = Texture2DBufferInfo::new(
423                    Self::PICKING_LAYER_FORMAT,
424                    metadata.picking_rect.wgpu_extent(),
425                );
426                let buffer_info_depth = Texture2DBufferInfo::new(
427                    if metadata.depth_readback_workaround_in_use {
428                        DepthReadbackWorkaround::READBACK_FORMAT
429                    } else {
430                        Self::PICKING_LAYER_DEPTH_FORMAT
431                    },
432                    metadata.picking_rect.wgpu_extent(),
433                );
434
435                let picking_id_data = buffer_info_id
436                    .remove_padding_and_convert(&data[..buffer_info_id.buffer_size_padded as _]);
437                let mut picking_depth_data = buffer_info_depth
438                    .remove_padding_and_convert(&data[buffer_info_id.buffer_size_padded as _..]);
439
440                if metadata.depth_readback_workaround_in_use {
441                    // Can't read back depth textures & can't read back R32Float textures either!
442                    // See https://github.com/gfx-rs/wgpu/issues/3644
443                    debug_assert_eq!(
444                        DepthReadbackWorkaround::READBACK_FORMAT
445                            .block_copy_size(None)
446                            .unwrap() as usize,
447                        std::mem::size_of::<f32>() * 4
448                    );
449                    picking_depth_data = picking_depth_data.into_iter().step_by(4).collect();
450                }
451
452                Some(PickingResult {
453                    picking_id_data,
454                    picking_depth_data,
455                    rect: metadata.picking_rect,
456                    world_from_cropped_projection: metadata.world_from_cropped_projection,
457                })
458            })
459            .flatten()
460    }
461}
462
463/// Utility for copying a depth texture when it can't be read-back directly to a [`wgpu::TextureFormat::R32Float`] which is readable texture.
464///
465/// Implementation note:
466/// This is a plain & simple "sample in shader and write to texture" utility.
467/// It might be worth abstracting this further into a general purpose operator.
468/// There is not much in here that is specific to the depth usecase!
469struct DepthReadbackWorkaround {
470    render_pipeline: GpuRenderPipelineHandle,
471    bind_group: GpuBindGroup,
472    readable_texture: GpuTexture,
473}
474
475impl DepthReadbackWorkaround {
476    /// There's two layers of workarounds here:
477    /// * WebGL (via spec) not being able to read back depth textures
478    /// * unclear behavior for any readback that isn't RGBA.
479    ///   Furthermore, integer textures also seemed to be problematic,
480    ///   but it seems to work fine for [`wgpu::TextureFormat::Rgba32Uint`] which we use for our picking ID.
481    ///   Details see [wgpu#3644](https://github.com/gfx-rs/wgpu/issues/3644)
482    const READBACK_FORMAT: wgpu::TextureFormat = wgpu::TextureFormat::Rgba32Float;
483
484    fn new(
485        ctx: &RenderContext,
486        extent: glam::UVec2,
487        depth_target_handle: GpuTextureHandle,
488    ) -> Self {
489        let readable_texture = ctx.gpu_resources.textures.alloc(
490            &ctx.device,
491            &TextureDesc {
492                label: "DepthCopyWorkaround::readable_texture".into(),
493                format: Self::READBACK_FORMAT,
494                usage: wgpu::TextureUsages::COPY_SRC | wgpu::TextureUsages::RENDER_ATTACHMENT,
495                size: wgpu::Extent3d {
496                    width: extent.x,
497                    height: extent.y,
498                    depth_or_array_layers: 1,
499                },
500                mip_level_count: 1,
501                sample_count: 1,
502                dimension: wgpu::TextureDimension::D2,
503            },
504        );
505
506        let bind_group_layout = ctx.gpu_resources.bind_group_layouts.get_or_create(
507            &ctx.device,
508            &BindGroupLayoutDesc {
509                label: "DepthCopyWorkaround::bind_group_layout".into(),
510                entries: vec![wgpu::BindGroupLayoutEntry {
511                    binding: 0,
512                    visibility: wgpu::ShaderStages::FRAGMENT,
513                    ty: wgpu::BindingType::Texture {
514                        sample_type: wgpu::TextureSampleType::Float { filterable: false },
515                        view_dimension: wgpu::TextureViewDimension::D2,
516                        multisampled: false,
517                    },
518                    count: None,
519                }],
520            },
521        );
522
523        let bind_group = ctx.gpu_resources.bind_groups.alloc(
524            &ctx.device,
525            &ctx.gpu_resources,
526            &BindGroupDesc {
527                label: "DepthCopyWorkaround::bind_group".into(),
528                entries: smallvec![BindGroupEntry::DefaultTextureView(depth_target_handle)],
529                layout: bind_group_layout,
530            },
531        );
532
533        let render_pipeline = ctx.gpu_resources.render_pipelines.get_or_create(
534            ctx,
535            &RenderPipelineDesc {
536                label: "DepthCopyWorkaround::render_pipeline".into(),
537                pipeline_layout: ctx.gpu_resources.pipeline_layouts.get_or_create(
538                    ctx,
539                    &PipelineLayoutDesc {
540                        label: "DepthCopyWorkaround::render_pipeline".into(),
541                        entries: vec![ctx.global_bindings.layout, bind_group_layout],
542                    },
543                ),
544                vertex_entrypoint: "main".into(),
545                vertex_handle: ctx.gpu_resources.shader_modules.get_or_create(
546                    ctx,
547                    &include_shader_module!("../../shader/screen_triangle.wgsl"),
548                ),
549                fragment_entrypoint: "main".into(),
550                fragment_handle: ctx.gpu_resources.shader_modules.get_or_create(
551                    ctx,
552                    &include_shader_module!("../../shader/copy_texture.wgsl"),
553                ),
554                vertex_buffers: smallvec![],
555                render_targets: smallvec![Some(readable_texture.texture.format().into())],
556                primitive: wgpu::PrimitiveState {
557                    topology: wgpu::PrimitiveTopology::TriangleStrip,
558                    cull_mode: None,
559                    ..Default::default()
560                },
561                depth_stencil: None,
562                multisample: wgpu::MultisampleState::default(),
563            },
564        );
565
566        Self {
567            render_pipeline,
568            bind_group,
569            readable_texture,
570        }
571    }
572
573    fn copy_to_readable_texture(
574        &self,
575        encoder: &mut wgpu::CommandEncoder,
576        render_pipelines: &GpuRenderPipelinePoolAccessor<'_>,
577        global_binding_bind_group: &GpuBindGroup,
578    ) -> Result<&GpuTexture, PoolError> {
579        // Copy depth texture to a readable (color) texture with a screen filling triangle.
580        let mut pass = encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
581            label: Label::from("Depth copy workaround").wgpu_label(),
582            color_attachments: &[Some(wgpu::RenderPassColorAttachment {
583                view: &self.readable_texture.default_view,
584                depth_slice: None,
585                resolve_target: None,
586                ops: wgpu::Operations {
587                    load: wgpu::LoadOp::Clear(wgpu::Color::TRANSPARENT),
588                    store: wgpu::StoreOp::Store, // Store for readback!
589                },
590            })],
591            depth_stencil_attachment: None,
592            timestamp_writes: None,
593            occlusion_query_set: None,
594            multiview_mask: None,
595        });
596
597        let pipeline = render_pipelines.get(self.render_pipeline)?;
598        pass.set_pipeline(pipeline);
599        pass.set_bind_group(0, global_binding_bind_group, &[]);
600        pass.set_bind_group(1, &self.bind_group, &[]);
601        pass.draw(0..3, 0..1);
602
603        Ok(&self.readable_texture)
604    }
605}
re_renderer/draw_phases/picking_layer.rs

re_renderer/draw_phases/
picking_layer.rs
