Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Batch skinned meshes #10094

Open
wants to merge 3 commits into
base: main
Choose a base branch
from
Open

Batch skinned meshes #10094

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
32b950a
WIP: skinning storage buffer hacks
superdump Oct 9, 2023
3d3ff48
Use the entire storage buffer as the binding
superdump Oct 10, 2023
121caae
Generalise
superdump Oct 10, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 6 additions & 0 deletions crates/bevy_pbr/src/prepass/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -228,6 +228,7 @@ pub struct PrepassPipeline<M: Material> {
pub material_vertex_shader: Option<Handle<Shader>>,
pub material_fragment_shader: Option<Handle<Shader>>,
pub material_pipeline: MaterialPipeline<M>,
pub skinned_mesh_storage_buffer: bool,
_marker: PhantomData<M>,
}

Expand Down Expand Up @@ -323,6 +324,10 @@ impl<M: Material> FromWorld for PrepassPipeline<M> {
},
material_layout: M::bind_group_layout(render_device),
material_pipeline: world.resource::<MaterialPipeline<M>>().clone(),
skinned_mesh_storage_buffer: render_device
.limits()
.max_storage_buffers_per_shader_stage
> 0,
_marker: PhantomData,
}
}
Expand Down Expand Up @@ -424,6 +429,7 @@ where
&key.mesh_key,
&mut shader_defs,
&mut vertex_attributes,
self.skinned_mesh_storage_buffer,
);
bind_group_layouts.insert(2, bind_group);

Expand Down
3 changes: 2 additions & 1 deletion crates/bevy_pbr/src/prepass/prepass.wgsl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -88,7 +88,8 @@ fn vertex(vertex_no_morph: Vertex) -> VertexOutput {
#endif

#ifdef SKINNED
var model = bevy_pbr::skinning::skin_model(vertex.joint_indices, vertex.joint_weights);
let skin_index = mesh[get_instance_index(vertex_no_morph.instance_index)].skin_index;
var model = bevy_pbr::skinning::skin_model(skin_index, vertex.joint_indices, vertex.joint_weights);
#else // SKINNED
// Use vertex_no_morph.instance_index instead of vertex.instance_index to work around a wgpu dx12 bug.
// See https://github.com/gfx-rs/naga/issues/2416
Expand Down
56 changes: 39 additions & 17 deletions crates/bevy_pbr/src/render/mesh.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,8 +2,9 @@ use crate::{
environment_map, prepass, EnvironmentMapLight, FogMeta, GlobalLightMeta, GpuFog, GpuLights,
GpuPointLights, LightMeta, MaterialBindGroupId, NotShadowCaster, NotShadowReceiver,
PreviousGlobalTransform, ScreenSpaceAmbientOcclusionTextures, Shadow, ShadowSamplers,
ViewClusterBindings, ViewFogUniformOffset, ViewLightsUniformOffset, ViewShadowBindings,
CLUSTERED_FORWARD_STORAGE_BUFFER_COUNT, MAX_CASCADES_PER_LIGHT, MAX_DIRECTIONAL_LIGHTS,
SkinIndex, ViewClusterBindings, ViewFogUniformOffset, ViewLightsUniformOffset,
ViewShadowBindings, CLUSTERED_FORWARD_STORAGE_BUFFER_COUNT, MAX_CASCADES_PER_LIGHT,
MAX_DIRECTIONAL_LIGHTS,
};
use bevy_app::{Plugin, PostUpdate};
use bevy_asset::{load_internal_asset, AssetId, Handle};
Expand Down Expand Up @@ -121,7 +122,6 @@ impl Plugin for MeshRenderPlugin {
render_app
.init_resource::<RenderMeshInstances>()
.init_resource::<MeshBindGroups>()
.init_resource::<SkinUniform>()
.init_resource::<SkinIndices>()
.init_resource::<MorphUniform>()
.init_resource::<MorphIndices>()
Expand Down Expand Up @@ -154,19 +154,21 @@ impl Plugin for MeshRenderPlugin {
let mut mesh_bindings_shader_defs = Vec::with_capacity(1);

if let Ok(render_app) = app.get_sub_app_mut(RenderApp) {
if let Some(per_object_buffer_batch_size) = GpuArrayBuffer::<MeshUniform>::batch_size(
render_app.world.resource::<RenderDevice>(),
) {
let render_device = render_app.world.resource::<RenderDevice>();
if let Some(per_object_buffer_batch_size) =
GpuArrayBuffer::<MeshUniform>::batch_size(render_device)
{
mesh_bindings_shader_defs.push(ShaderDefVal::UInt(
"PER_OBJECT_BUFFER_BATCH_SIZE".into(),
per_object_buffer_batch_size,
));
}

let mesh_uniform_buffer = GpuArrayBuffer::<MeshUniform>::new(render_device);
let skin_uniform_buffer = SkinUniform::new(render_device);
render_app
.insert_resource(GpuArrayBuffer::<MeshUniform>::new(
render_app.world.resource::<RenderDevice>(),
))
.insert_resource(mesh_uniform_buffer)
.insert_resource(skin_uniform_buffer)
.init_resource::<MeshPipeline>();
}

Expand Down Expand Up @@ -201,10 +203,11 @@ pub struct MeshUniform {
pub inverse_transpose_model_a: [Vec4; 2],
pub inverse_transpose_model_b: f32,
pub flags: u32,
pub skin_index: u32,
}

impl From<&MeshTransforms> for MeshUniform {
fn from(mesh_transforms: &MeshTransforms) -> Self {
impl MeshUniform {
pub fn from(mesh_transforms: &MeshTransforms, skin_index: u32) -> Self {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
pub fn from(mesh_transforms: &MeshTransforms, skin_index: u32) -> Self {
pub fn new(mesh_transforms: &MeshTransforms, skin_index: u32) -> Self {

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Agree with this change, as from has a special meaning in Rust wrt the From-trait (I am aware that you're aware, but potential new contributors might not)

let (inverse_transpose_model_a, inverse_transpose_model_b) =
mesh_transforms.transform.inverse_transpose_3x3();
Self {
Expand All @@ -213,6 +216,7 @@ impl From<&MeshTransforms> for MeshUniform {
inverse_transpose_model_a,
inverse_transpose_model_b,
flags: mesh_transforms.flags,
skin_index,
}
}
}
Expand Down Expand Up @@ -339,6 +343,7 @@ pub struct MeshPipeline {
/// ##endif // PER_OBJECT_BUFFER_BATCH_SIZE
/// ```
pub per_object_buffer_batch_size: Option<u32>,
pub skinned_mesh_storage_buffer: bool,
}

impl FromWorld for MeshPipeline {
Expand Down Expand Up @@ -579,6 +584,10 @@ impl FromWorld for MeshPipeline {
dummy_white_gpu_image,
mesh_layouts: MeshLayouts::new(&render_device),
per_object_buffer_batch_size: GpuArrayBuffer::<MeshUniform>::batch_size(&render_device),
skinned_mesh_storage_buffer: render_device
.limits()
.max_storage_buffers_per_shader_stage
> 0,
}
}
}
Expand All @@ -602,21 +611,29 @@ impl MeshPipeline {
}

impl GetBatchData for MeshPipeline {
type Param = SRes<RenderMeshInstances>;
type Param = (SRes<RenderMeshInstances>, SRes<SkinIndices>);
type Query = Entity;
type QueryFilter = With<Mesh3d>;
type CompareData = (MaterialBindGroupId, AssetId<Mesh>);
type BufferData = MeshUniform;

fn get_batch_data(
mesh_instances: &SystemParamItem<Self::Param>,
(mesh_instances, skin_indices): &SystemParamItem<Self::Param>,
entity: &QueryItem<Self::Query>,
) -> (Self::BufferData, Option<Self::CompareData>) {
let mesh_instance = mesh_instances
.get(entity)
.expect("Failed to find render mesh instance");
(
(&mesh_instance.transforms).into(),
MeshUniform::from(
&mesh_instance.transforms,
skin_indices
.get(entity)
.map_or(u32::MAX, |skin_index| match skin_index {
SkinIndex::Index(index) => *index,
SkinIndex::DynamicOffset(_) => u32::MAX,
}),
),
mesh_instance.automatic_batching.then_some((
mesh_instance.material_bind_group_id,
mesh_instance.mesh_asset_id,
Expand Down Expand Up @@ -737,9 +754,13 @@ pub fn setup_morph_and_skinning_defs(
key: &MeshPipelineKey,
shader_defs: &mut Vec<ShaderDefVal>,
vertex_attributes: &mut Vec<VertexAttributeDescriptor>,
skinned_mesh_storage_buffer: bool,
) -> BindGroupLayout {
let mut add_skin_data = || {
shader_defs.push("SKINNED".into());
if skinned_mesh_storage_buffer {
shader_defs.push("SKINNED_MESH_STORAGE_BUFFER".into());
}
vertex_attributes.push(Mesh::ATTRIBUTE_JOINT_INDEX.at_shader_location(offset));
vertex_attributes.push(Mesh::ATTRIBUTE_JOINT_WEIGHT.at_shader_location(offset + 1));
};
Expand Down Expand Up @@ -820,6 +841,7 @@ impl SpecializedMeshPipeline for MeshPipeline {
&key,
&mut shader_defs,
&mut vertex_attributes,
self.skinned_mesh_storage_buffer,
));

if key.contains(MeshPipelineKey::SCREEN_SPACE_AMBIENT_OCCLUSION) {
Expand Down Expand Up @@ -1047,7 +1069,7 @@ pub fn prepare_mesh_bind_group(
};
groups.model_only = Some(layouts.model_only(&render_device, &model));

let skin = skins_uniform.buffer.buffer();
let skin = skins_uniform.buffer();
if let Some(skin) = skin {
groups.skinned = Some(layouts.skinned(&render_device, &model, skin));
}
Expand Down Expand Up @@ -1317,8 +1339,8 @@ impl<P: PhaseItem, const I: usize> RenderCommand<P> for SetMeshBindGroup<I> {
dynamic_offsets[offset_count] = dynamic_offset.get();
offset_count += 1;
}
if let Some(skin_index) = skin_index {
dynamic_offsets[offset_count] = skin_index.index;
if let Some(SkinIndex::DynamicOffset(dynamic_offset)) = skin_index {
dynamic_offsets[offset_count] = *dynamic_offset;
offset_count += 1;
}
if let Some(morph_index) = morph_index {
Expand Down
3 changes: 2 additions & 1 deletion crates/bevy_pbr/src/render/mesh.wgsl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -64,7 +64,8 @@ fn vertex(vertex_no_morph: Vertex) -> MeshVertexOutput {
#endif

#ifdef SKINNED
var model = bevy_pbr::skinning::skin_model(vertex.joint_indices, vertex.joint_weights);
let skin_index = mesh[get_instance_index(vertex_no_morph.instance_index)].skin_index;
var model = bevy_pbr::skinning::skin_model(skin_index, vertex.joint_indices, vertex.joint_weights);
#else
// Use vertex_no_morph.instance_index instead of vertex.instance_index to work around a wgpu dx12 bug.
// See https://github.com/gfx-rs/naga/issues/2416 .
Expand Down
72 changes: 57 additions & 15 deletions crates/bevy_pbr/src/render/mesh_bindings.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,7 @@ pub(crate) const JOINT_BUFFER_SIZE: usize = MAX_JOINTS * JOINT_SIZE;

/// Individual layout entries.
mod layout_entry {
use super::{JOINT_BUFFER_SIZE, MORPH_BUFFER_SIZE};
use super::{JOINT_BUFFER_SIZE, JOINT_SIZE, MORPH_BUFFER_SIZE};
use crate::MeshUniform;
use bevy_render::{
render_resource::{
Expand All @@ -30,14 +30,23 @@ mod layout_entry {
renderer::RenderDevice,
};

fn buffer(binding: u32, size: u64, visibility: ShaderStages) -> BindGroupLayoutEntry {
fn buffer(
binding: u32,
size: u64,
visibility: ShaderStages,
is_storage: bool,
) -> BindGroupLayoutEntry {
BindGroupLayoutEntry {
binding,
visibility,
count: None,
ty: BindingType::Buffer {
ty: BufferBindingType::Uniform,
has_dynamic_offset: true,
ty: if is_storage {
BufferBindingType::Storage { read_only: true }
} else {
BufferBindingType::Uniform
},
has_dynamic_offset: !is_storage,
min_binding_size: BufferSize::new(size),
},
}
Expand All @@ -49,11 +58,27 @@ mod layout_entry {
render_device,
)
}
pub(super) fn skinning(binding: u32) -> BindGroupLayoutEntry {
buffer(binding, JOINT_BUFFER_SIZE as u64, ShaderStages::VERTEX)
pub(super) fn skinning(render_device: &RenderDevice, binding: u32) -> BindGroupLayoutEntry {
let is_storage = render_device.limits().max_storage_buffers_per_shader_stage > 0;
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
let is_storage = render_device.limits().max_storage_buffers_per_shader_stage > 0;
let is_storage = render_device.limits().max_storage_buffers_per_shader_stage > 0;

At this point, I think we should add a method like "storage_buffers_supported()" to RenderDevice.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What does your suggestion change? The lines look identical to me.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I didn't change anything sorry. This was meant to be a regular comment on that line.

let min_binding_size = if is_storage {
JOINT_SIZE
} else {
JOINT_BUFFER_SIZE
};
buffer(
binding,
min_binding_size as u64,
ShaderStages::VERTEX,
is_storage,
)
}
pub(super) fn weights(binding: u32) -> BindGroupLayoutEntry {
buffer(binding, MORPH_BUFFER_SIZE as u64, ShaderStages::VERTEX)
buffer(
binding,
MORPH_BUFFER_SIZE as u64,
ShaderStages::VERTEX,
false,
)
}
pub(super) fn targets(binding: u32) -> BindGroupLayoutEntry {
BindGroupLayoutEntry {
Expand All @@ -72,8 +97,11 @@ mod layout_entry {
/// for bind groups.
mod entry {
use super::{JOINT_BUFFER_SIZE, MORPH_BUFFER_SIZE};
use bevy_render::render_resource::{
BindGroupEntry, BindingResource, Buffer, BufferBinding, BufferSize, TextureView,
use bevy_render::{
render_resource::{
BindGroupEntry, BindingResource, Buffer, BufferBinding, BufferSize, TextureView,
},
renderer::RenderDevice,
};

fn entry(binding: u32, size: u64, buffer: &Buffer) -> BindGroupEntry {
Expand All @@ -89,8 +117,19 @@ mod entry {
pub(super) fn model(binding: u32, resource: BindingResource) -> BindGroupEntry {
BindGroupEntry { binding, resource }
}
pub(super) fn skinning(binding: u32, buffer: &Buffer) -> BindGroupEntry {
entry(binding, JOINT_BUFFER_SIZE as u64, buffer)
pub(super) fn skinning<'b>(
render_device: &RenderDevice,
binding: u32,
buffer: &'b Buffer,
) -> BindGroupEntry<'b> {
if render_device.limits().max_storage_buffers_per_shader_stage > 0 {
BindGroupEntry {
binding,
resource: buffer.as_entire_binding(),
}
} else {
entry(binding, JOINT_BUFFER_SIZE as u64, buffer)
}
}
pub(super) fn weights(binding: u32, buffer: &Buffer) -> BindGroupEntry {
entry(binding, MORPH_BUFFER_SIZE as u64, buffer)
Expand Down Expand Up @@ -149,7 +188,7 @@ impl MeshLayouts {
render_device.create_bind_group_layout(&BindGroupLayoutDescriptor {
entries: &[
layout_entry::model(render_device, 0),
layout_entry::skinning(1),
layout_entry::skinning(render_device, 1),
],
label: Some("skinned_mesh_layout"),
})
Expand All @@ -168,7 +207,7 @@ impl MeshLayouts {
render_device.create_bind_group_layout(&BindGroupLayoutDescriptor {
entries: &[
layout_entry::model(render_device, 0),
layout_entry::skinning(1),
layout_entry::skinning(render_device, 1),
layout_entry::weights(2),
layout_entry::targets(3),
],
Expand All @@ -192,7 +231,10 @@ impl MeshLayouts {
skin: &Buffer,
) -> BindGroup {
render_device.create_bind_group(&BindGroupDescriptor {
entries: &[entry::model(0, model.clone()), entry::skinning(1, skin)],
entries: &[
entry::model(0, model.clone()),
entry::skinning(render_device, 1, skin),
],
layout: &self.skinned,
label: Some("skinned_mesh_bind_group"),
})
Expand Down Expand Up @@ -225,7 +267,7 @@ impl MeshLayouts {
render_device.create_bind_group(&BindGroupDescriptor {
entries: &[
entry::model(0, model.clone()),
entry::skinning(1, skin),
entry::skinning(render_device, 1, skin),
entry::weights(2, weights),
entry::targets(3, targets),
],
Expand Down
5 changes: 5 additions & 0 deletions crates/bevy_pbr/src/render/mesh_types.wgsl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,11 +14,16 @@ struct Mesh {
inverse_transpose_model_b: f32,
// 'flags' is a bit field indicating various options. u32 is 32 bits so we have up to 32 options.
flags: u32,
skin_index: u32,
};

#ifdef SKINNED
struct SkinnedMesh {
#ifdef SKINNED_MESH_STORAGE_BUFFER
data: array<mat4x4<f32>>,
#else
data: array<mat4x4<f32>, 256u>,
#endif
};
#endif

Expand Down
Loading