Skip to content

Commit

Permalink
clean up
Browse files Browse the repository at this point in the history
  • Loading branch information
IceSentry committed Jun 12, 2023
1 parent 1b1920a commit 143fb61
Show file tree
Hide file tree
Showing 9 changed files with 97 additions and 172 deletions.
2 changes: 2 additions & 0 deletions Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -85,6 +85,8 @@ bevy_gltf = ["bevy_internal/bevy_gltf", "bevy_asset", "bevy_scene", "bevy_pbr"]
# Adds PBR rendering
bevy_pbr = ["bevy_internal/bevy_pbr", "bevy_asset", "bevy_render", "bevy_core_pipeline"]

gpu_picking = []

# Provides rendering functionality
bevy_render = ["bevy_internal/bevy_render"]

Expand Down
1 change: 1 addition & 0 deletions crates/bevy_core_pipeline/src/core_3d/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -86,6 +86,7 @@ impl Plugin for Core3dPlugin {
prepare_core_3d_depth_textures
.in_set(RenderSet::Prepare)
.after(bevy_render::view::prepare_windows),
// #[cfg(feature = "gpu_picking")]
prepare_entity_textures
.in_set(RenderSet::Prepare)
.after(bevy_render::view::prepare_windows),
Expand Down
2 changes: 0 additions & 2 deletions crates/bevy_core_pipeline/src/entity_index_buffer_copy/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -25,11 +25,9 @@ impl ViewNode for EntityIndexBufferCopyNode {
return Ok(());
};

// copy entity index texture
buffers.copy_texture_to_buffer(
render_context.command_encoder(),
&entity_index_textures.main.texture,
&buffers.entity_buffer,
);

Ok(())
Expand Down
5 changes: 4 additions & 1 deletion crates/bevy_gizmos/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -37,14 +37,16 @@ use bevy_reflect::{
use bevy_render::{
color::Color,
mesh::Mesh,
picking,
primitives::Aabb,
render_phase::AddRenderCommand,
render_resource::{PrimitiveTopology, Shader, SpecializedMeshPipelines},
Extract, ExtractSchedule, Render, RenderApp, RenderSet,
};
use bevy_transform::components::{GlobalTransform, Transform};

// #[cfg(feature = "gpu_picking")]
use bevy_render::picking;

#[cfg(feature = "bevy_pbr")]
use bevy_pbr::MeshUniform;
#[cfg(feature = "bevy_sprite")]
Expand Down Expand Up @@ -326,6 +328,7 @@ fn extract_gizmo_data(
transform,
previous_transform: transform,
inverse_transpose_model,
// #[cfg(feature = "gpu_picking")]
entity: picking::entity_as_uvec2(Entity::PLACEHOLDER),
},
),
Expand Down
1 change: 1 addition & 0 deletions crates/bevy_pbr/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@ keywords = ["bevy"]

[features]
webgl = []
gpu_picking = []

[dependencies]
# bevy
Expand Down
3 changes: 2 additions & 1 deletion crates/bevy_pbr/src/render/mesh.rs
Original file line number Diff line number Diff line change
Expand Up @@ -132,7 +132,7 @@ pub struct MeshUniform {
pub transform: Mat4,
pub previous_transform: Mat4,
pub inverse_transpose_model: Mat4,
// TODO should probably be a separate uniform in case gpu picking is disabled
// #[cfg(feature = "gpu_picking")]
pub entity: UVec2,
pub flags: u32,
}
Expand Down Expand Up @@ -187,6 +187,7 @@ pub fn extract_meshes(
flags: flags.bits(),
transform,
previous_transform,
// #[cfg(feature = "gpu_picking")]
entity: picking::entity_as_uvec2(entity),
inverse_transpose_model: transform.inverse().transpose(),
};
Expand Down
4 changes: 1 addition & 3 deletions crates/bevy_render/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -45,7 +45,6 @@ pub mod prelude {
use bevy_window::{PrimaryWindow, RawHandleWrapper};
use globals::GlobalsPlugin;
pub use once_cell;
use picking::GpuPickingPlugin;
use renderer::{RenderAdapter, RenderAdapterInfo, RenderDevice, RenderQueue};
use wgpu::Instance;

Expand Down Expand Up @@ -336,8 +335,7 @@ impl Plugin for RenderPlugin {
.add_plugin(CameraPlugin)
.add_plugin(ViewPlugin)
.add_plugin(MeshPlugin)
.add_plugin(GlobalsPlugin)
.add_plugin(GpuPickingPlugin);
.add_plugin(GlobalsPlugin);

app.register_type::<color::Color>()
.register_type::<primitives::Aabb>()
Expand Down
196 changes: 74 additions & 122 deletions crates/bevy_render/src/picking.rs
Original file line number Diff line number Diff line change
@@ -1,132 +1,106 @@
use crate::{
camera::ExtractedCamera,
extract_component::ExtractComponentPlugin,
render_resource::{Buffer, BufferSlice, Texture},
render_resource::{Buffer, Texture},
renderer::RenderDevice,
texture::{CachedTexture, TextureFormatPixelInfo},
Render, RenderApp, RenderSet,
};
use async_channel::{Receiver, Sender, TrySendError};
use async_channel::{Receiver, Sender};
use bevy_app::{Plugin, Update};
use bevy_ecs::prelude::*;

use bevy_math::UVec2;
use bevy_render_macros::ExtractComponent;
use bevy_tasks::AsyncComputeTaskPool;
use bevy_utils::{default, HashMap};
use bevy_utils::default;
use wgpu::{
BufferDescriptor, BufferUsages, Color, CommandEncoder, Extent3d, ImageDataLayout, Maintain,
MapMode, Operations, RenderPassColorAttachment, TextureFormat,
BufferDescriptor, BufferUsages, Color, CommandEncoder, Extent3d, ImageDataLayout, MapMode,
Operations, RenderPassColorAttachment, TextureFormat,
};

pub const ENTITY_TEXTURE_FORMAT: TextureFormat = TextureFormat::Rg32Uint;

/// This plugin enables the gpu picking feature of bevy.
///
/// Gpu picking let's you know which entity is currently under the mouse.
///
/// # How this works:
///
/// - For each entity being rendered, it will output it's entity id to a texture.
/// - Once everything is rendered it will copy that texture to the cpu and send it to the main world
/// - From the main world you can give it a position like the current mouse position and
/// know exactly which entity was rendered at that specific screen location.
/// - This works at the [`Camera`] level, so it will work with multiple windows or split screen
///
/// # Api Overview:
///
/// You need to add the [`GpuPickingCamera`] to any `Camera` that will be used for picking.
/// Then add the [`GpuPickingMesh`] comnponent to any `Mesh` that will need to be picked.
///
/// Once those components are added, you can query for [`GpuPickingCamera`]
/// and use `GpuPickingCamera::get_entity(position)` to know which entity is at the given position on screen
pub struct GpuPickingPlugin;
impl Plugin for GpuPickingPlugin {
fn build(&self, app: &mut bevy_app::App) {
app.add_plugin(ExtractComponentPlugin::<GpuPickingMesh>::default())
.add_plugin(ExtractComponentPlugin::<GpuPickingCamera>::default())
.add_systems(Update, update_entity_buffer);
.add_systems(Update, receive_buffer);

let Ok(render_app) = app.get_sub_app_mut(RenderApp) else { return; };

render_app.add_systems(
Render,
(
prepare_gpu_picking_buffers.in_set(RenderSet::Prepare),
send_buffer_to_main_world.in_set(RenderSet::RenderFlush),
map_and_send_buffer_async.in_set(RenderSet::RenderFlush),
),
);
}
}

/// Copies the entity buffer to the cpu and sends it to the main world
fn send_buffer_to_main_world(
query: Query<&ExtractedGpuPickingCamera, With<ExtractedCamera>>,
render_device: Res<RenderDevice>,
) {
/// Maps the entity buffer and sends it to the main world asynchronously
fn map_and_send_buffer_async(query: Query<&ExtractedGpuPickingCamera, With<ExtractedCamera>>) {
for gpu_picking_camera in &query {
let Some(buffers) = gpu_picking_camera.buffers.as_ref() else {
return;
};
let buffers = buffers.clone();
let sender = gpu_picking_camera.sender.clone();

// Send the data to the main world
fn send(
sender: Sender<GpuPickingData>,
padded_bytes_per_row: usize,
entity_buffer: &Buffer,
buffer_slice: BufferSlice,
) {
let buffer_view = buffer_slice.get_mapped_range();
// immediately copy the data to CPU to avoid holding the mapped view for long
let entity_data = Vec::from(&*buffer_view);
drop(buffer_view);
entity_buffer.unmap();
// Mapping the buffer is an asynchronous process.
// This means we need to wait until the buffer is mapped before sending it to the main world
let task = async move {
let (tx, rx) = async_channel::bounded(1);

if let Err(err) = sender.try_send(GpuPickingData {
padded_bytes_per_row,
entity_data,
}) {
match err {
TrySendError::Full(_) => bevy_log::error!("GPU Picking channel is full"),
TrySendError::Closed(_) => {
bevy_log::error!("GPU Picking channel is closed");
}
}
}

// This can only fail if the sender is full or closed
// and we can't do anything if either of those things happen
// let _ = sender.try_send(GpuPickingData {
// padded_bytes_per_row,
// entity_data,
// });
}

if true {
// sync
// map entity buffer
let buffer_slice = buffers.entity_buffer.slice(..);
buffer_slice.map_async(MapMode::Read, move |result| match result {
Ok(_) => {}
Ok(_) => tx.try_send(()).unwrap(),
Err(err) => bevy_log::error!("Failed to map entity buffer: {err}"),
});

// WARN This is blocking
render_device.poll(Maintain::Wait);
// Buffer is mapped and ready to be sent
rx.recv().await.unwrap();

// Send the buffer to the main world
send(
gpu_picking_camera.sender.clone(),
buffers.buffer_dimensions.padded_bytes_per_row,
&buffers.entity_buffer,
buffer_slice,
);
} else {
// async
let entity_buffer = buffers.entity_buffer.clone();
let sender = gpu_picking_camera.sender.clone();
let padded_bytes_per_row = buffers.buffer_dimensions.padded_bytes_per_row;

// Mapping the buffer is an asynchronous process.
// This means we need to wait until the buffer is mapped before sending it to the main world
let task = async move {
let (tx, rx) = async_channel::bounded(1);

// map entity buffer
let buffer_slice = entity_buffer.slice(..);
buffer_slice.map_async(MapMode::Read, move |result| match result {
Ok(_) => tx.try_send(()).unwrap(),
Err(err) => bevy_log::error!("Failed to map entity buffer: {err}"),
});

// Buffer is mapped and ready to be sent
rx.recv().await.unwrap();

// Send the buffer to the main world
send(sender, padded_bytes_per_row, &entity_buffer, buffer_slice);
};
AsyncComputeTaskPool::get().spawn(task).detach();
}
let entity_buffer = &buffers.entity_buffer;
let buffer_view = buffer_slice.get_mapped_range();
// immediately copy the data to CPU to avoid holding the mapped view for long
let entity_data = Vec::from(&*buffer_view);
drop(buffer_view);
entity_buffer.unmap();

// Because the channel is bounded to 1 entry, it's pretty common to have the channel full.
// This isn't ideal but not blocking makes it faster which is preferred.
//
// The other possible error is for the channel to be closed and in that case we can't do anything
let _ = sender.try_send(GpuPickingData {
padded_bytes_per_row: buffers.buffer_dimensions.padded_bytes_per_row,
entity_data,
});
};
AsyncComputeTaskPool::get().spawn(task).detach();
}
}

Expand Down Expand Up @@ -166,7 +140,7 @@ impl Default for GpuPickingCamera {
impl GpuPickingCamera {
pub fn new() -> Self {
Self {
channel: async_channel::unbounded(),
channel: async_channel::bounded(1),
data: GpuPickingData::default(),
}
}
Expand Down Expand Up @@ -218,24 +192,21 @@ impl crate::extract_component::ExtractComponent for GpuPickingCamera {
}

/// Contains the buffer and it's dimension required for gpu picking
#[derive(Clone)]
pub struct GpuPickingCameraBuffers {
pub entity_buffer: Buffer,
buffer_dimensions: BufferDimensions,
}

impl GpuPickingCameraBuffers {
pub fn copy_texture_to_buffer(
&self,
encoder: &mut CommandEncoder,
texture: &Texture,
buffer: &Buffer,
) {
/// Copies the given texture to the entity buffer
pub fn copy_texture_to_buffer(&self, encoder: &mut CommandEncoder, texture: &Texture) {
// This can't be in the Node because it needs access to wgpu but
// bevy_core_pipeline doesn't depend on wgpu
encoder.copy_texture_to_buffer(
texture.as_image_copy(),
wgpu::ImageCopyBuffer {
buffer,
buffer: &self.entity_buffer,
layout: ImageDataLayout {
offset: 0,
bytes_per_row: Some(self.buffer_dimensions.padded_bytes_per_row as u32),
Expand All @@ -251,16 +222,11 @@ impl GpuPickingCameraBuffers {
}
}

fn update_entity_buffer(mut q: Query<&mut GpuPickingCamera>) {
fn receive_buffer(mut q: Query<&mut GpuPickingCamera>) {
for mut cam in &mut q {
let (_, receiver) = cam.channel.clone();
loop {
let Ok(buffer) = receiver.try_recv() else { break; };
if receiver.is_empty() {
cam.data = buffer;
break;
}
}
let Ok(data) = receiver.try_recv() else { continue; };
cam.data = data;
}
}

Expand Down Expand Up @@ -316,39 +282,25 @@ impl EntityTextures {
fn prepare_gpu_picking_buffers(
render_device: Res<RenderDevice>,
mut cameras: Query<
(Entity, &ExtractedCamera, &mut ExtractedGpuPickingCamera),
(&ExtractedCamera, &mut ExtractedGpuPickingCamera),
Changed<ExtractedCamera>,
>,
mut buffer_cache: Local<HashMap<Entity, (BufferDimensions, Buffer)>>,
) {
for (entity, camera, mut gpu_picking_camera) in &mut cameras {
for (camera, mut gpu_picking_camera) in &mut cameras {
let Some(size) = camera.physical_target_size else { continue; };

// TODO create 2 buffers and altenate between each buffer each frame

// Only create a buffer if it doesn't already exist or the size is different
let mut create_buffer = true;
if let Some((buffer_dimensions, _)) = buffer_cache.get(&entity) {
create_buffer = buffer_dimensions.width != size.x as usize
|| buffer_dimensions.height != size.y as usize;
}

if create_buffer {
let buffer_dimensions =
BufferDimensions::new(size.x as usize, size.y as usize, ENTITY_TEXTURE_FORMAT);
let entity_buffer = render_device.create_buffer(&BufferDescriptor {
label: Some("Entity Buffer"),
size: (buffer_dimensions.padded_bytes_per_row * buffer_dimensions.height) as u64,
usage: BufferUsages::COPY_DST | BufferUsages::MAP_READ,
mapped_at_creation: false,
});
buffer_cache.insert(entity, (buffer_dimensions, entity_buffer));
}
let buffer_dimensions =
BufferDimensions::new(size.x as usize, size.y as usize, ENTITY_TEXTURE_FORMAT);
let entity_buffer = render_device.create_buffer(&BufferDescriptor {
label: Some("Entity Buffer"),
size: (buffer_dimensions.padded_bytes_per_row * buffer_dimensions.height) as u64,
usage: BufferUsages::COPY_DST | BufferUsages::MAP_READ,
mapped_at_creation: false,
});

let Some((buffer_dimensions, entity_buffer)) = buffer_cache.get(&entity) else { continue; };
gpu_picking_camera.buffers = Some(GpuPickingCameraBuffers {
entity_buffer: entity_buffer.clone(),
buffer_dimensions: *buffer_dimensions,
buffer_dimensions,
});
}
}
Expand Down
Loading

0 comments on commit 143fb61

Please sign in to comment.