add gpu picking

Co-authored-by: Kaylee Simmons <kay@the-simmons.net>

Co-authored-by: Nicola Papale <nicopap@users.noreply.github.com>

Cargo.toml

@@ -1443,6 +1443,16 @@ description = "Shows how to rumble a gamepad using force feedback"
 category = "Input"
 wasm = false
 
+[[example]]
+name = "gpu_picking"
+path = "examples/input/gpu_picking.rs"
+
+[package.metadata.example.gpu_picking]
+name = "GPU picking"
+description = "Mouse picking using the gpu"
+category = "Input"
+wasm = true
+
 [[example]]
 name = "keyboard_input"
 path = "examples/input/keyboard_input.rs"

assets/shaders/gpu_picking_material.wgsl

@@ -0,0 +1,28 @@
+// This shader shows how to enable the gpu picking feature for a material
+
+// You'll need the mesh binding because that's where the entity index is
+#import bevy_pbr::mesh_bindings mesh
+#import bevy_pbr::mesh_vertex_output MeshVertexOutput
+
+@group(1) @binding(0)
+var<uniform> color: vec4<f32>;
+
+// Gpu picking uses multiple fragment output
+struct FragmentOutput {
+    @location(0) color: vec4<f32>,
+// You can detect the feature with this flag
+#ifdef GPU_PICKING
+    @location(1) mesh_id: u32,
+#endif
+};
+
+@fragment
+fn fragment(in: MeshVertexOutput) -> FragmentOutput {
+    var out: FragmentOutput;
+    out.color = color;
+// make sure to output the entity index for gpu picking to work correctly
+#ifdef GPU_PICKING
+    out.mesh_id = mesh[in.instance_index].id;
+#endif
+    return out;
+}

crates/bevy_core_pipeline/src/core_3d/main_opaque_pass_3d_node.rs

@@ -7,6 +7,7 @@ use crate::{
 use bevy_ecs::{prelude::*, query::QueryItem};
 use bevy_render::{
     camera::ExtractedCamera,
+    picking::{ExtractedGpuPickingCamera, VisibleMeshIdTextures},
     render_graph::{NodeRunError, RenderGraphContext, ViewNode},
     render_phase::RenderPhase,
     render_resource::{
@@ -34,8 +35,10 @@ impl ViewNode for MainOpaquePass3dNode {
         Option<&'static DepthPrepass>,
         Option<&'static NormalPrepass>,
         Option<&'static MotionVectorPrepass>,
+        Option<&'static ExtractedGpuPickingCamera>,
         Option<&'static SkyboxPipelineId>,
         Option<&'static SkyboxBindGroup>,
+        Option<&'static VisibleMeshIdTextures>,
         &'static ViewUniformOffset,
     );
 
@@ -53,8 +56,10 @@ impl ViewNode for MainOpaquePass3dNode {
             depth_prepass,
             normal_prepass,
             motion_vector_prepass,
+            gpu_picking_camera,
             skybox_pipeline,
             skybox_bind_group,
+            mesh_id_textures,
             view_uniform_offset,
         ): QueryItem<Self::ViewQuery>,
         world: &World,
@@ -64,21 +69,34 @@ impl ViewNode for MainOpaquePass3dNode {
         #[cfg(feature = "trace")]
         let _main_opaque_pass_3d_span = info_span!("main_opaque_pass_3d").entered();
 
+        let mut color_attachments = vec![Some(target.get_color_attachment(Operations {
+            load: match camera_3d.clear_color {
+                ClearColorConfig::Default => LoadOp::Clear(world.resource::<ClearColor>().0.into()),
+                ClearColorConfig::Custom(color) => LoadOp::Clear(color.into()),
+                ClearColorConfig::None => LoadOp::Load,
+            },
+            store: true,
+        }))];
+
+        if gpu_picking_camera.is_some() {
+            if let Some(mesh_id_textures) = mesh_id_textures {
+                color_attachments.push(Some(mesh_id_textures.get_color_attachment(Operations {
+                    load: match camera_3d.clear_color {
+                        ClearColorConfig::None => LoadOp::Load,
+                        // TODO clear this earlier?
+                        _ => LoadOp::Clear(VisibleMeshIdTextures::clear_color()),
+                    },
+                    store: true,
+                })));
+            }
+        }
+
         // Setup render pass
         let mut render_pass = render_context.begin_tracked_render_pass(RenderPassDescriptor {
             label: Some("main_opaque_pass_3d"),
             // NOTE: The opaque pass loads the color
             // buffer as well as writing to it.
-            color_attachments: &[Some(target.get_color_attachment(Operations {
-                load: match camera_3d.clear_color {
-                    ClearColorConfig::Default => {
-                        LoadOp::Clear(world.resource::<ClearColor>().0.into())
-                    }
-                    ClearColorConfig::Custom(color) => LoadOp::Clear(color.into()),
-                    ClearColorConfig::None => LoadOp::Load,
-                },
-                store: true,
-            }))],
+            color_attachments: &color_attachments,
             depth_stencil_attachment: Some(RenderPassDepthStencilAttachment {
                 view: &depth.view,
                 // NOTE: The opaque main pass loads the depth buffer and possibly overwrites it

crates/bevy_core_pipeline/src/core_3d/main_transparent_pass_3d_node.rs

@@ -2,6 +2,7 @@ use crate::core_3d::Transparent3d;
 use bevy_ecs::{prelude::*, query::QueryItem};
 use bevy_render::{
     camera::ExtractedCamera,
+    picking::{ExtractedGpuPickingCamera, VisibleMeshIdTextures},
     render_graph::{NodeRunError, RenderGraphContext, ViewNode},
     render_phase::RenderPhase,
     render_resource::{LoadOp, Operations, RenderPassDepthStencilAttachment, RenderPassDescriptor},
@@ -21,12 +22,16 @@ impl ViewNode for MainTransparentPass3dNode {
         &'static RenderPhase<Transparent3d>,
         &'static ViewTarget,
         &'static ViewDepthTexture,
+        Option<&'static ExtractedGpuPickingCamera>,
+        Option<&'static VisibleMeshIdTextures>,
     );
     fn run(
         &self,
         graph: &mut RenderGraphContext,
         render_context: &mut RenderContext,
-        (camera, transparent_phase, target, depth): QueryItem<Self::ViewQuery>,
+        (camera, transparent_phase, target, depth, gpu_picking_camera, mesh_id_textures): QueryItem<
+            Self::ViewQuery,
+        >,
         world: &World,
     ) -> Result<(), NodeRunError> {
         let view_entity = graph.view_entity();
@@ -37,13 +42,27 @@ impl ViewNode for MainTransparentPass3dNode {
             #[cfg(feature = "trace")]
             let _main_transparent_pass_3d_span = info_span!("main_transparent_pass_3d").entered();
 
+            let mut color_attachments = vec![Some(target.get_color_attachment(Operations {
+                load: LoadOp::Load,
+                store: true,
+            }))];
+
+            if gpu_picking_camera.is_some() {
+                if let Some(mesh_id_textures) = mesh_id_textures {
+                    color_attachments.push(Some(mesh_id_textures.get_color_attachment(
+                        Operations {
+                            // The texture is already cleared in the opaque pass
+                            load: LoadOp::Load,
+                            store: true,
+                        },
+                    )));
+                }
+            }
+
             let mut render_pass = render_context.begin_tracked_render_pass(RenderPassDescriptor {
                 label: Some("main_transparent_pass_3d"),
                 // NOTE: The transparent pass loads the color buffer as well as overwriting it where appropriate.
-                color_attachments: &[Some(target.get_color_attachment(Operations {
-                    load: LoadOp::Load,
-                    store: true,
-                }))],
+                color_attachments: &color_attachments,
                 depth_stencil_attachment: Some(RenderPassDepthStencilAttachment {
                     view: &depth.view,
                     // NOTE: For the transparent pass we load the depth buffer. There should be no

crates/bevy_core_pipeline/src/core_3d/mod.rs

@@ -14,6 +14,7 @@ pub mod graph {
         pub const MAIN_OPAQUE_PASS: &str = "main_opaque_pass";
         pub const MAIN_TRANSPARENT_PASS: &str = "main_transparent_pass";
         pub const END_MAIN_PASS: &str = "end_main_pass";
+        pub const ENTITY_INDEX_BUFFER_COPY: &str = "entity_index_buffer_copy";
         pub const BLOOM: &str = "bloom";
         pub const TONEMAPPING: &str = "tonemapping";
         pub const FXAA: &str = "fxaa";
@@ -35,6 +36,7 @@ use bevy_ecs::prelude::*;
 use bevy_render::{
     camera::{Camera, ExtractedCamera},
     extract_component::ExtractComponentPlugin,
+    picking::{ExtractedGpuPickingCamera, VisibleMeshIdTextures, MESH_ID_TEXTURE_FORMAT},
     prelude::Msaa,
     render_graph::{EmptyNode, RenderGraphApp, ViewNodeRunner},
     render_phase::{
@@ -94,6 +96,7 @@ impl Plugin for Core3dPlugin {
                     sort_phase_system::<AlphaMask3dPrepass>.in_set(RenderSet::PhaseSort),
                     prepare_core_3d_depth_textures.in_set(RenderSet::PrepareResources),
                     prepare_prepass_textures.in_set(RenderSet::PrepareResources),
+                    prepare_entity_textures.in_set(RenderSet::PrepareResources),
                 ),
             );
 
@@ -493,3 +496,63 @@ pub fn prepare_prepass_textures(
         });
     }
 }
+
+/// Create the required textures based on the camera size
+pub fn prepare_entity_textures(
+    mut commands: Commands,
+    mut texture_cache: ResMut<TextureCache>,
+    msaa: Res<Msaa>,
+    render_device: Res<RenderDevice>,
+    views_3d: Query<
+        (Entity, &ExtractedCamera),
+        (
+            With<ExtractedGpuPickingCamera>,
+            With<RenderPhase<Opaque3d>>,
+            With<RenderPhase<AlphaMask3d>>,
+        ),
+    >,
+) {
+    for (entity, camera) in &views_3d {
+        let Some(physical_target_size) = camera.physical_target_size else {
+            continue;
+        };
+
+        let size = Extent3d {
+            depth_or_array_layers: 1,
+            width: physical_target_size.x,
+            height: physical_target_size.y,
+        };
+
+        let descriptor = TextureDescriptor {
+            label: None,
+            size,
+            mip_level_count: 1,
+            sample_count: 1,
+            dimension: TextureDimension::D2,
+            format: MESH_ID_TEXTURE_FORMAT,
+            usage: TextureUsages::RENDER_ATTACHMENT | TextureUsages::COPY_SRC,
+            view_formats: &[],
+        };
+
+        let mesh_id_textures = VisibleMeshIdTextures {
+            main: texture_cache.get(
+                &render_device,
+                TextureDescriptor {
+                    label: Some("main_entity_texture"),
+                    ..descriptor
+                },
+            ),
+            sampled: (msaa.samples() > 1).then(|| {
+                texture_cache.get(
+                    &render_device,
+                    TextureDescriptor {
+                        label: Some("main_entity_texture_sampled"),
+                        sample_count: msaa.samples(),
+                        ..descriptor
+                    },
+                )
+            }),
+        };
+        commands.entity(entity).insert(mesh_id_textures);
+    }
+}

crates/bevy_core_pipeline/src/entity_index_buffer_copy/mod.rs

@@ -0,0 +1,53 @@
+use bevy_app::Plugin;
+use bevy_ecs::{query::QueryItem, world::World};
+use bevy_render::{
+    picking::{CurrentGpuPickingBufferIndex, ExtractedGpuPickingCamera, VisibleMeshIdTextures},
+    render_graph::{RenderGraphApp, RenderGraphContext, ViewNode, ViewNodeRunner},
+    renderer::RenderContext,
+    RenderApp,
+};
+
+use crate::core_3d::CORE_3D;
+
+#[derive(Default)]
+pub struct EntityIndexBufferCopyNode;
+impl ViewNode for EntityIndexBufferCopyNode {
+    type ViewQuery = (
+        &'static VisibleMeshIdTextures,
+        &'static ExtractedGpuPickingCamera,
+    );
+
+    fn run(
+        &self,
+        _graph: &mut RenderGraphContext,
+        render_context: &mut RenderContext,
+        (mesh_id_textures, gpu_picking_camera): QueryItem<Self::ViewQuery>,
+        world: &World,
+    ) -> Result<(), bevy_render::render_graph::NodeRunError> {
+        let current_buffer_index = world.resource::<CurrentGpuPickingBufferIndex>();
+        gpu_picking_camera.run_node(
+            render_context.command_encoder(),
+            &mesh_id_textures.main.texture,
+            current_buffer_index,
+        );
+        Ok(())
+    }
+}
+
+pub struct EntityIndexBufferCopyPlugin;
+impl Plugin for EntityIndexBufferCopyPlugin {
+    fn build(&self, app: &mut bevy_app::App) {
+        let Ok(render_app) = app.get_sub_app_mut(RenderApp) else {
+            return;
+        };
+
+        // 3D
+        use crate::core_3d::graph::node::*;
+        render_app
+            .add_render_graph_node::<ViewNodeRunner<EntityIndexBufferCopyNode>>(
+                CORE_3D,
+                ENTITY_INDEX_BUFFER_COPY,
+            )
+            .add_render_graph_edge(CORE_3D, UPSCALING, ENTITY_INDEX_BUFFER_COPY);
+    }
+}

crates/bevy_core_pipeline/src/lib.rs

@@ -6,6 +6,7 @@ pub mod clear_color;
 pub mod contrast_adaptive_sharpening;
 pub mod core_2d;
 pub mod core_3d;
+pub mod entity_index_buffer_copy;
 pub mod fullscreen_vertex_shader;
 pub mod fxaa;
 pub mod msaa_writeback;
@@ -40,6 +41,7 @@ use crate::{
     contrast_adaptive_sharpening::CASPlugin,
     core_2d::Core2dPlugin,
     core_3d::Core3dPlugin,
+    entity_index_buffer_copy::EntityIndexBufferCopyPlugin,
     fullscreen_vertex_shader::FULLSCREEN_SHADER_HANDLE,
     fxaa::FxaaPlugin,
     msaa_writeback::MsaaWritebackPlugin,
@@ -79,6 +81,7 @@ impl Plugin for CorePipelinePlugin {
                 BloomPlugin,
                 FxaaPlugin,
                 CASPlugin,
+                EntityIndexBufferCopyPlugin,
             ));
     }
 }