Tidy up storage alignment code (#326)

djeedai · May 19, 2024 · a80db80 · a80db80
1 parent 7b8f0dd
commit a80db80
Show file tree

Hide file tree

Showing 4 changed files with 106 additions and 242 deletions.
diff --git a/src/lib.rs b/src/lib.rs
@@ -1822,7 +1822,7 @@ else { return c1; }
 
             // Import bevy_hanabi::vfx_common
             {
-                let min_storage_buffer_offset_alignment = 256usize;
+                let min_storage_buffer_offset_alignment = 256;
                 let common_shader =
                     HanabiPlugin::make_common_shader(min_storage_buffer_offset_alignment);
                 let res = composer.add_composable_module((&common_shader).into());

diff --git a/src/plugin.rs b/src/plugin.rs
@@ -25,7 +25,7 @@ use crate::{
         EffectBindGroups, EffectCache, EffectsMeta, ExtractedEffects, GpuDispatchIndirect,
         GpuParticleGroup, GpuRenderEffectMetadata, GpuRenderGroupIndirect, GpuSpawnerParams,
         ParticlesInitPipeline, ParticlesRenderPipeline, ParticlesUpdatePipeline, ShaderCache,
-        SimParams, VfxSimulateDriverNode, VfxSimulateNode,
+        SimParams, StorageType as _, VfxSimulateDriverNode, VfxSimulateNode,
     },
     spawn::{self, Random},
     tick_spawners,
@@ -127,7 +127,7 @@ impl HanabiPlugin {
     /// This creates a new [`Shader`] from the `vfx_common.wgsl` code, by
     /// applying the given alignment for storage buffers. This produces a shader
     /// ready for the specific GPU device associated with that alignment.
-    pub(crate) fn make_common_shader(min_storage_buffer_offset_alignment: usize) -> Shader {
+    pub(crate) fn make_common_shader(min_storage_buffer_offset_alignment: u32) -> Shader {
         let spawner_padding_code =
             GpuSpawnerParams::padding_code(min_storage_buffer_offset_alignment);
         let dispatch_indirect_padding_code =
@@ -242,7 +242,7 @@ impl Plugin for HanabiPlugin {
         // for init/update/render shaders (but not the indirect one).
         {
             let common_shader = HanabiPlugin::make_common_shader(
-                render_device.limits().min_storage_buffer_offset_alignment as usize,
+                render_device.limits().min_storage_buffer_offset_alignment,
             );
             let mut assets = app.world.resource_mut::<Assets<Shader>>();
             assets.insert(HANABI_COMMON_TEMPLATE_HANDLE, common_shader);

diff --git a/src/render/effect_cache.rs b/src/render/effect_cache.rs
@@ -15,7 +15,9 @@ use std::{
 
 use crate::{
     asset::EffectAsset,
-    render::{GpuDispatchIndirect, GpuParticleGroup, GpuSpawnerParams, LayoutFlags},
+    render::{
+        GpuDispatchIndirect, GpuParticleGroup, GpuSpawnerParams, LayoutFlags, StorageType as _,
+    },
     ParticleLayout, PropertyLayout,
 };
 
@@ -239,10 +241,9 @@ impl EffectBuffer {
 
         // TODO - Cache particle_layout and associated bind group layout, instead of
         // creating one bind group layout per buffer using that layout...
-        let particle_group_size = NonZeroU64::new(GpuParticleGroup::aligned_size(
-            render_device.limits().min_storage_buffer_offset_alignment as usize,
-        ) as u64)
-        .unwrap();
+        let particle_group_size = GpuParticleGroup::aligned_size(
+            render_device.limits().min_storage_buffer_offset_alignment,
+        );
         let mut entries = vec![
             // @binding(0) var<storage, read_write> particle_buffer : ParticleBuffer
             BindGroupLayoutEntry {
@@ -301,10 +302,9 @@ impl EffectBuffer {
         let particles_buffer_layout_sim = render_device.create_bind_group_layout(label, &entries);
 
         // Create the render layout.
-        let dispatch_indirect_size = NonZeroU64::new(GpuDispatchIndirect::aligned_size(
-            render_device.limits().min_storage_buffer_offset_alignment as usize,
-        ) as u64)
-        .unwrap();
+        let dispatch_indirect_size = GpuDispatchIndirect::aligned_size(
+            render_device.limits().min_storage_buffer_offset_alignment,
+        );
         let mut entries = vec![
             BindGroupLayoutEntry {
                 binding: 0,