From 02ff96ba26a88bcb59106aba72166e2220cd7cbb Mon Sep 17 00:00:00 2001 From: mosure Date: Mon, 1 Jan 2024 19:20:42 -0600 Subject: [PATCH] fix: sorted texture binding --- Cargo.toml | 2 +- src/morph/particle.rs | 10 ++++- src/render/gaussian.wgsl | 34 ++++++++++++++-- src/render/mod.rs | 61 ++++++++++++++++++++++++++++- src/render/packed.rs | 10 ++++- src/render/planar.rs | 10 ++++- src/render/texture.rs | 60 +++++++++++++++++++++++------ src/render/texture.wgsl | 15 +++----- src/sort/mod.rs | 83 ++++++++++++++++++++++++++++++++++------ src/sort/radix.rs | 33 +++++++++++++++- 10 files changed, 276 insertions(+), 42 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index bc9f29b7..0b4a6b5b 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -49,7 +49,7 @@ default = [ "query_select", "query_sparse", - "sort_radix", + # "sort_radix", "sort_rayon", "sort_std", diff --git a/src/morph/particle.rs b/src/morph/particle.rs index 75f7e9ca..988790fb 100644 --- a/src/morph/particle.rs +++ b/src/morph/particle.rs @@ -24,7 +24,15 @@ use bevy::{ RenderAssets, RenderAssetPlugin, }, - render_resource::*, + render_resource::{ + Buffer, + BufferInitDescriptor, + BufferUsages, + Extent3d, + ShaderType, + TextureDimension, + TextureFormat, + }, renderer::{ RenderContext, RenderDevice, diff --git a/src/render/gaussian.wgsl b/src/render/gaussian.wgsl index 0b7f1c11..0adf5a96 100644 --- a/src/render/gaussian.wgsl +++ b/src/render/gaussian.wgsl @@ -48,12 +48,36 @@ get_scale, get_opacity, get_visibility, + location, } #endif +#ifdef BUFFER_STORAGE @group(3) @binding(0) var sorted_entries: array; +fn get_entry(index: u32) -> Entry { + return sorted_entries[index]; +} +#endif + +#ifdef BUFFER_TEXTURE +@group(3) @binding(0) var sorted_entries: texture_2d; + +fn get_entry(index: u32) -> Entry { + let sample = textureLoad( + sorted_entries, + location(index), + 0, + ); + + return Entry( + sample.r, + sample.g, + ); +} +#endif + struct GaussianVertexOutput { @builtin(position) position: vec4, @location(0) @interpolate(flat) color: vec4, @@ -239,11 +263,13 @@ fn vs_points( @builtin(vertex_index) vertex_index: u32, ) -> GaussianVertexOutput { var output: GaussianVertexOutput; - let splat_index = sorted_entries[instance_index][1]; + + let entry = get_entry(instance_index); + let splat_index = entry.value; var discard_quad = false; - discard_quad |= sorted_entries[instance_index][0] == 0xFFFFFFFFu; // || splat_index == 0u; + discard_quad |= entry.key == 0xFFFFFFFFu; // || splat_index == 0u; let position = vec4(get_position(splat_index), 1.0); @@ -277,8 +303,8 @@ fn vs_points( var rgb = vec3(0.0); #ifdef VISUALIZE_DEPTH - let first_position = vec4(get_position(sorted_entries[1][1]), 1.0); - let last_position = vec4(get_position(sorted_entries[gaussian_uniforms.count - 1u][1]), 1.0); + let first_position = vec4(get_position(get_entry(1u).value), 1.0); + let last_position = vec4(get_position(get_entry(gaussian_uniforms.count - 1u).value), 1.0); let min_position = (gaussian_uniforms.global_transform * first_position).xyz; let max_position = (gaussian_uniforms.global_transform * last_position).xyz; diff --git a/src/render/mod.rs b/src/render/mod.rs index fbec1840..f7b91673 100644 --- a/src/render/mod.rs +++ b/src/render/mod.rs @@ -41,7 +41,43 @@ use bevy::{ SetItemPipeline, TrackedRenderPass, }, - render_resource::*, + render_resource::{ + BindGroup, + BindGroupEntry, + BindGroupLayout, + BindGroupLayoutDescriptor, + BindGroupLayoutEntry, + BindingResource, + BindingType, + BlendState, + Buffer, + BufferBinding, + BufferBindingType, + BufferInitDescriptor, + BufferUsages, + ColorTargetState, + ColorWrites, + CompareFunction, + DepthBiasState, + DepthStencilState, + FragmentState, + FrontFace, + MultisampleState, + PipelineCache, + PolygonMode, + PrimitiveState, + PrimitiveTopology, + RenderPipelineDescriptor, + ShaderDefVal, + ShaderStages, + ShaderType, + SpecializedRenderPipeline, + SpecializedRenderPipelines, + StencilFaceState, + StencilState, + TextureFormat, + VertexState, + }, renderer::RenderDevice, Render, RenderApp, @@ -393,7 +429,7 @@ impl FromWorld for GaussianCloudPipeline { #[cfg(feature = "buffer_texture")] let gaussian_cloud_layout = texture::get_bind_group_layout(&render_device, read_only); - // TODO: support sorted layout as a texture + #[cfg(feature = "buffer_storage")] let sorted_layout = render_device.create_bind_group_layout(&BindGroupLayoutDescriptor { label: Some("sorted_layout"), entries: &[ @@ -409,6 +445,8 @@ impl FromWorld for GaussianCloudPipeline { }, ], }); + #[cfg(feature = "buffer_texture")] + let sorted_layout = texture::get_sorted_bind_group_layout(&render_device); GaussianCloudPipeline { gaussian_cloud_layout, @@ -742,6 +780,9 @@ fn queue_gaussian_bind_group( &Handle, &texture::GpuTextureBuffers, )>, + + #[cfg(feature = "buffer_texture")] + gpu_images: Res>, ) { let Some(model) = gaussian_uniforms.buffer() else { return; @@ -788,7 +829,9 @@ fn queue_gaussian_bind_group( continue; } + #[cfg(not(feature = "buffer_texture"))] let cloud: &GpuGaussianCloud = gaussian_cloud_res.get(cloud_handle).unwrap(); + let sorted_entries = sorted_entries_res.get(sorted_entries_handle).unwrap(); #[cfg(feature = "packed")] @@ -798,6 +841,7 @@ fn queue_gaussian_bind_group( #[cfg(feature = "buffer_texture")] let cloud_bind_group = texture_buffers.bind_group.clone(); + #[cfg(feature = "buffer_storage")] let sorted_bind_group = render_device.create_bind_group( "render_sorted_bind_group", &gaussian_cloud_pipeline.sorted_layout, @@ -812,6 +856,19 @@ fn queue_gaussian_bind_group( }, ], ); + #[cfg(feature = "buffer_texture")] + let sorted_bind_group = render_device.create_bind_group( + Some("render_sorted_bind_group"), + &gaussian_cloud_pipeline.sorted_layout, + &[ + BindGroupEntry { + binding: 0, + resource: BindingResource::TextureView( + &gpu_images.get(&sorted_entries.texture).unwrap().texture_view + ), + }, + ], + ); commands.entity(entity).insert(GaussianCloudBindGroup { cloud_bind_group, diff --git a/src/render/packed.rs b/src/render/packed.rs index a8396f8c..99e8d692 100644 --- a/src/render/packed.rs +++ b/src/render/packed.rs @@ -1,5 +1,13 @@ use bevy::render::{ - render_resource::*, + render_resource::{ + Buffer, + BufferInitDescriptor, + BufferUsages, + Extent3d, + ShaderType, + TextureDimension, + TextureFormat, + }, renderer::RenderDevice, }; diff --git a/src/render/planar.rs b/src/render/planar.rs index 3bdfd799..529499d2 100644 --- a/src/render/planar.rs +++ b/src/render/planar.rs @@ -1,5 +1,13 @@ use bevy::render::{ - render_resource::*, + render_resource::{ + Buffer, + BufferInitDescriptor, + BufferUsages, + Extent3d, + ShaderType, + TextureDimension, + TextureFormat, + }, renderer::RenderDevice, }; diff --git a/src/render/texture.rs b/src/render/texture.rs index 35d656d2..903892c2 100644 --- a/src/render/texture.rs +++ b/src/render/texture.rs @@ -11,7 +11,22 @@ use bevy::{ RenderApp, RenderSet, render_asset::RenderAssets, - render_resource::*, + render_resource::{ + BindGroup, + BindGroupLayout, + BindGroupLayoutDescriptor, + BindGroupLayoutEntry, + BindGroupEntry, + BindingType, + BindingResource, + Extent3d, + TextureDimension, + TextureFormat, + TextureSampleType, + TextureUsages, + TextureViewDimension, + ShaderStages, + }, renderer::RenderDevice, }, }; @@ -27,15 +42,15 @@ use crate::{ ScaleOpacity, }, }, - render::{ - GaussianCloudPipeline, - GpuGaussianCloud, - }, material::spherical_harmonics::{ SH_COEFF_COUNT, SH_VEC4_PLANES, SphericalHarmonicCoefficients, }, + render::{ + GaussianCloudPipeline, + GpuGaussianCloud, + }, }; @@ -275,6 +290,27 @@ fn queue_textures( } +pub fn get_sorted_bind_group_layout( + render_device: &RenderDevice, +) -> BindGroupLayout { + render_device.create_bind_group_layout(&BindGroupLayoutDescriptor { + label: Some("texture_sorted_layout"), + entries: &[ + BindGroupLayoutEntry { + binding: 0, + visibility: ShaderStages::all(), + ty: BindingType::Texture { + view_dimension: TextureViewDimension::D2, + sample_type: TextureSampleType::Uint, + multisampled: false, + }, + count: None, + }, + ], + }) +} + + #[cfg(feature = "f16")] pub fn get_bind_group_layout( render_device: &RenderDevice, @@ -287,21 +323,23 @@ pub fn get_bind_group_layout( }; render_device.create_bind_group_layout(&BindGroupLayoutDescriptor { - label: Some("planar_f16_gaussian_cloud_layout"), + label: Some("texture_f16_gaussian_cloud_layout"), entries: &[ BindGroupLayoutEntry { binding: 0, - visibility: ShaderStages::VERTEX_FRAGMENT, + visibility: ShaderStages::all(), ty: BindingType::Texture { view_dimension: TextureViewDimension::D2, - sample_type: TextureSampleType::Uint, + sample_type: TextureSampleType::Float { + filterable: false, + }, multisampled: false, }, count: None, }, BindGroupLayoutEntry { binding: 1, - visibility: ShaderStages::VERTEX_FRAGMENT, + visibility: ShaderStages::all(), ty: BindingType::Texture { view_dimension: sh_view_dimension, sample_type: TextureSampleType::Uint, @@ -311,7 +349,7 @@ pub fn get_bind_group_layout( }, BindGroupLayoutEntry { binding: 2, - visibility: ShaderStages::VERTEX_FRAGMENT, + visibility: ShaderStages::all(), ty: BindingType::Texture { view_dimension: TextureViewDimension::D2, sample_type: TextureSampleType::Uint, @@ -330,7 +368,7 @@ pub fn get_bind_group_layout( read_only: bool ) -> BindGroupLayout { render_device.create_bind_group_layout(&BindGroupLayoutDescriptor { - label: Some("planar_f32_gaussian_cloud_layout"), + label: Some("texture_f32_gaussian_cloud_layout"), entries: &[ BindGroupLayoutEntry { binding: 0, diff --git a/src/render/texture.wgsl b/src/render/texture.wgsl index 5ca8dd84..297cb839 100644 --- a/src/render/texture.wgsl +++ b/src/render/texture.wgsl @@ -24,6 +24,7 @@ fn get_position(index: u32) -> vec3 { let sample = textureLoad( position_visibility, location(index), + 0, ); return sample.xyz; @@ -33,19 +34,11 @@ fn get_spherical_harmonics(index: u32) -> array { var coefficients: array; for (var i = 0u; i < #{SH_VEC4_PLANES}u; i = i + 1u) { - -#if SH_VEC4_PLANES == 1 - let sample = textureLoad( - spherical_harmonics, - location(index), - ); -#else let sample = textureLoad( spherical_harmonics, location(index), - i, + i32(i), ); -#endif let v0 = unpack2x16float(sample.x); let v1 = unpack2x16float(sample.y); @@ -73,6 +66,7 @@ fn get_rotation(index: u32) -> vec4 { let sample = textureLoad( rotation_scale_opacity, location(index), + 0, ); let q0 = unpack2x16float(sample.x); @@ -88,6 +82,7 @@ fn get_scale(index: u32) -> vec3 { let sample = textureLoad( rotation_scale_opacity, location(index), + 0, ); let s0 = unpack2x16float(sample.z); @@ -103,6 +98,7 @@ fn get_opacity(index: u32) -> f32 { let sample = textureLoad( rotation_scale_opacity, location(index), + 0, ); return unpack2x16float(sample.w).x; @@ -112,6 +108,7 @@ fn get_visibility(index: u32) -> f32 { let sample = textureLoad( position_visibility, location(index), + 0, ); return sample.w; diff --git a/src/sort/mod.rs b/src/sort/mod.rs index 992a7b96..61cf5b3c 100644 --- a/src/sort/mod.rs +++ b/src/sort/mod.rs @@ -7,12 +7,20 @@ use bevy::{ }, reflect::TypeUuid, render::{ + render_resource::{ + Buffer, + BufferInitDescriptor, + BufferUsages, + Extent3d, + ShaderType, + TextureDimension, + TextureFormat, + }, render_asset::{ RenderAsset, RenderAssetPlugin, PrepareAssetError, }, - render_resource::*, renderer::RenderDevice, }, }; @@ -107,6 +115,31 @@ impl Plugin for SortPlugin { app.add_plugins(RenderAssetPlugin::::default()); app.add_systems(Update, auto_insert_sorted_entries); + + #[cfg(feature = "buffer_texture")] + app.add_systems(PostUpdate, update_textures_on_change); + } +} + + +#[cfg(feature = "buffer_texture")] +fn update_textures_on_change( + mut images: ResMut>, + mut ev_asset: EventReader>, + sorted_entries_res: Res>, +) { + for ev in ev_asset.read() { + match ev { + AssetEvent::Modified { id } => { + let sorted_entries = sorted_entries_res.get(*id).unwrap(); + let image = images.get_mut(&sorted_entries.texture).unwrap(); + + image.data = bytemuck::cast_slice(sorted_entries.sorted.as_slice()).to_vec(); + }, + AssetEvent::Added { id: _ } => {}, + AssetEvent::Removed { id: _ } => {}, + AssetEvent::LoadedWithDependencies { id: _ } => {}, + } } } @@ -123,6 +156,9 @@ fn auto_insert_sorted_entries( &GaussianCloudSettings, Without>, )>, + + #[cfg(feature = "buffer_texture")] + mut images: ResMut>, ) { for ( entity, @@ -145,16 +181,34 @@ fn auto_insert_sorted_entries( } let cloud = cloud.unwrap(); + let sorted: Vec = (0..cloud.len()) + .map(|idx| { + SortEntry { + key: 1, + index: idx as u32, + } + }) + .collect(); + // TODO: move gaussian_cloud and sorted_entry assets into an asset bundle + #[cfg(feature = "buffer_storage")] + let sorted_entries = sorted_entries_res.add(SortedEntries { + sorted, + }); + + #[cfg(feature = "buffer_texture")] let sorted_entries = sorted_entries_res.add(SortedEntries { - sorted: (0..cloud.len()) - .map(|idx| { - SortEntry { - key: 1, - index: idx as u32, - } - }) - .collect(), + texture: images.add(Image::new( + Extent3d { + width: cloud.len_sqrt_ceil() as u32, + height: cloud.len_sqrt_ceil() as u32, + depth_or_array_layers: 1, + }, + TextureDimension::D2, + bytemuck::cast_slice(sorted.as_slice()).to_vec(), + TextureFormat::Rg32Uint, + )), + sorted, }); commands.entity(entity) @@ -180,8 +234,6 @@ pub struct SortEntry { pub index: u32, } -// TODO: add RenderAssetPlugin for SortedEntries & auto-insert to GaussianCloudBundles if their sort mode is not None -// supports pre-sorting or CPU sorting in main world, initializes the sorting_entry_buffer #[derive( Clone, Asset, @@ -194,6 +246,9 @@ pub struct SortEntry { #[uuid = "ac2f08eb-fa13-ccdd-ea11-51571ea332d5"] pub struct SortedEntries { pub sorted: Vec, + + #[cfg(feature = "buffer_texture")] + pub texture: Handle, } impl RenderAsset for SortedEntries { @@ -220,6 +275,9 @@ impl RenderAsset for SortedEntries { Ok(GpuSortedEntry { sorted_entry_buffer, count, + + #[cfg(feature = "buffer_texture")] + texture: sorted_entries.texture, }) } } @@ -231,4 +289,7 @@ impl RenderAsset for SortedEntries { pub struct GpuSortedEntry { pub sorted_entry_buffer: Buffer, pub count: usize, + + #[cfg(feature = "buffer_texture")] + pub texture: Handle, } diff --git a/src/sort/radix.rs b/src/sort/radix.rs index 05fc0786..85003014 100644 --- a/src/sort/radix.rs +++ b/src/sort/radix.rs @@ -9,7 +9,28 @@ use bevy::{ core_pipeline::core_3d::CORE_3D, render::{ render_asset::RenderAssets, - render_resource::*, + render_resource::{ + BindGroup, + BindGroupEntry, + BindGroupLayout, + BindGroupLayoutDescriptor, + BindGroupLayoutEntry, + BindingResource, + BindingType, + Buffer, + BufferBindingType, + BufferDescriptor, + BufferInitDescriptor, + BufferBinding, + BufferSize, + BufferUsages, + CachedComputePipelineId, + CachedPipelineState, + ComputePassDescriptor, + ComputePipelineDescriptor, + PipelineCache, + ShaderStages, + }, renderer::{ RenderContext, RenderDevice, @@ -26,6 +47,7 @@ use bevy::{ view::ViewUniformOffset, }, }; +use static_assertions::assert_cfg; use crate::{ gaussian::cloud::GaussianCloud, @@ -47,6 +69,15 @@ use crate::{ }; +assert_cfg!( + not(all( + feature = "sort_radix", + feature = "buffer_texture", + )), + "sort_radix and buffer_texture are incompatible", +); + + const RADIX_SHADER_HANDLE: Handle = Handle::weak_from_u128(6234673214); const TEMPORAL_SORT_SHADER_HANDLE: Handle = Handle::weak_from_u128(1634543224);