bevy_pbr: Use storage buffers for clustered forward rendering bindings

This enables using many more than 256 point lights in a scene. On an M1 Max
MacBook Pro, it can maintain 60fps with about 2150 point lights. With
non-pipelined rendering, the bottleneck seems to be assign_lights_to_clusters,
which was expected. There are possibly optimisations that can be done to reduce
the bottleneck CPU-side, but I have seen papers and presentations supporting
millions of lights when leveraging compute shaders for the light assignment to
clusters.

Author: superdump

crates/bevy_pbr/src/lib.rs

@@ -37,6 +37,7 @@ use bevy_asset::{load_internal_asset, Assets, Handle, HandleUntyped};
 use bevy_ecs::prelude::*;
 use bevy_reflect::TypeUuid;
 use bevy_render::{
+    options::WgpuOptions,
     prelude::Color,
     render_graph::RenderGraph,
     render_phase::{sort_phase_system, AddRenderCommand, DrawFunctions},
@@ -135,6 +136,16 @@ impl Plugin for PbrPlugin {
                 },
             );
 
+        // FIXME: What would be a good way of managing storage buffer limitations per
+        // stage?
+        let use_storage_buffers = app
+            .world
+            .get_resource::<WgpuOptions>()
+            .unwrap()
+            .limits
+            .max_storage_buffers_per_shader_stage
+            >= 3;
+
         let render_app = match app.get_sub_app_mut(RenderApp) {
             Ok(render_app) => render_app,
             Err(_) => return,
@@ -175,7 +186,7 @@ impl Plugin for PbrPlugin {
             .init_resource::<ShadowPipeline>()
             .init_resource::<DrawFunctions<Shadow>>()
             .init_resource::<LightMeta>()
-            .init_resource::<GlobalLightMeta>()
+            .insert_resource(GlobalLightMeta::new(use_storage_buffers))
             .init_resource::<SpecializedPipelines<ShadowPipeline>>();
 
         let shadow_pass_node = ShadowPassNode::new(&mut render_app.world);

crates/bevy_pbr/src/material.rs

@@ -16,6 +16,7 @@ use bevy_ecs::{
 };
 use bevy_render::{
     mesh::Mesh,
+    options::WgpuOptions,
     render_asset::{RenderAsset, RenderAssetPlugin, RenderAssets},
     render_component::ExtractComponentPlugin,
     render_phase::{
@@ -78,7 +79,11 @@ impl<M: Material> SpecializedMaterial for M {
     type Key = ();
 
     #[inline]
-    fn key(_material: &<Self as RenderAsset>::PreparedAsset) -> Self::Key {}
+    fn key(
+        _wgpu_options: &WgpuOptions,
+        _material: &<Self as RenderAsset>::PreparedAsset,
+    ) -> Self::Key {
+    }
 
     #[inline]
     fn specialize(_key: Self::Key, _descriptor: &mut RenderPipelineDescriptor) {}
@@ -127,7 +132,10 @@ pub trait SpecializedMaterial: Asset + RenderAsset {
     /// Extract the [`SpecializedMaterial::Key`] for the "prepared" version of this material. This key will be
     /// passed in to the [`SpecializedMaterial::specialize`] function when compiling the [`RenderPipeline`](bevy_render::render_resource::RenderPipeline)
     /// for a given entity's material.
-    fn key(material: &<Self as RenderAsset>::PreparedAsset) -> Self::Key;
+    fn key(
+        wgpu_options: &WgpuOptions,
+        material: &<Self as RenderAsset>::PreparedAsset,
+    ) -> Self::Key;
 
     /// Specializes the given `descriptor` according to the given `key`.
     fn specialize(key: Self::Key, descriptor: &mut RenderPipelineDescriptor);
@@ -271,6 +279,7 @@ impl<M: SpecializedMaterial, const I: usize> EntityRenderCommand for SetMaterial
 
 #[allow(clippy::too_many_arguments)]
 pub fn queue_material_meshes<M: SpecializedMaterial>(
+    wgpu_options: Res<WgpuOptions>,
     opaque_draw_functions: Res<DrawFunctions<Opaque3d>>,
     alpha_mask_draw_functions: Res<DrawFunctions<AlphaMask3d>>,
     transparent_draw_functions: Res<DrawFunctions<Transparent3d>>,
@@ -289,6 +298,7 @@ pub fn queue_material_meshes<M: SpecializedMaterial>(
         &mut RenderPhase<Transparent3d>,
     )>,
 ) {
+    let wgpu_options = wgpu_options.into_inner();
     for (view, visible_entities, mut opaque_phase, mut alpha_mask_phase, mut transparent_phase) in
         views.iter_mut()
     {
@@ -327,7 +337,7 @@ pub fn queue_material_meshes<M: SpecializedMaterial>(
                         mesh_key |= MeshPipelineKey::TRANSPARENT_MAIN_PASS;
                     }
 
-                    let specialized_key = M::key(material);
+                    let specialized_key = M::key(wgpu_options, material);
                     let pipeline_id = pipelines.specialize(
                         &mut pipeline_cache,
                         &material_pipeline,

crates/bevy_pbr/src/pbr_material.rs

@@ -5,6 +5,7 @@ use bevy_math::Vec4;
 use bevy_reflect::TypeUuid;
 use bevy_render::{
     color::Color,
+    options::WgpuOptions,
     prelude::Shader,
     render_asset::{PrepareAssetError, RenderAsset, RenderAssets},
     render_resource::{
@@ -327,14 +328,19 @@ impl RenderAsset for StandardMaterial {
 #[derive(Clone, PartialEq, Eq, Hash)]
 pub struct StandardMaterialKey {
     normal_map: bool,
+    storage_buffers: bool,
 }
 
 impl SpecializedMaterial for StandardMaterial {
     type Key = StandardMaterialKey;
 
-    fn key(render_asset: &<Self as RenderAsset>::PreparedAsset) -> Self::Key {
+    fn key(
+        wgpu_options: &WgpuOptions,
+        render_asset: &<Self as RenderAsset>::PreparedAsset,
+    ) -> Self::Key {
         StandardMaterialKey {
             normal_map: render_asset.has_normal_map,
+            storage_buffers: wgpu_options.limits.max_storage_buffers_per_shader_stage >= 3,
         }
     }
 
@@ -347,6 +353,14 @@ impl SpecializedMaterial for StandardMaterial {
                 .shader_defs
                 .push(String::from("STANDARDMATERIAL_NORMAL_MAP"));
         }
+        if !key.storage_buffers {
+            descriptor
+                .fragment
+                .as_mut()
+                .unwrap()
+                .shader_defs
+                .push(String::from("NO_STORAGE_BUFFERS_SUPPORT"));
+        }
         if let Some(label) = &mut descriptor.label {
             *label = format!("pbr_{}", *label).into();
         }