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

re_renderer/draw_phases/
picking_layer.rs
