Render assets diagnostics

crates/bevy_pbr/src/diagnostic.rs

@@ -0,0 +1,124 @@
+use core::{
+    any::{type_name, Any, TypeId},
+    marker::PhantomData,
+};
+
+use bevy_app::{Plugin, PreUpdate};
+use bevy_diagnostic::{Diagnostic, DiagnosticPath, Diagnostics, RegisterDiagnostic};
+use bevy_ecs::{resource::Resource, system::Res};
+use bevy_platform::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
+use bevy_render::{Extract, ExtractSchedule, RenderApp};
+
+use crate::{Material, MaterialBindGroupAllocators};
+
+pub struct MaterialAllocatorDiagnosticPlugin<M: Material> {
+    suffix: &'static str,
+    _phantom: PhantomData<M>,
+}
+
+impl<M: Material> MaterialAllocatorDiagnosticPlugin<M> {
+    pub fn new(suffix: &'static str) -> Self {
+        Self {
+            suffix,
+            _phantom: PhantomData,
+        }
+    }
+}
+
+impl<M: Material> Default for MaterialAllocatorDiagnosticPlugin<M> {
+    fn default() -> Self {
+        Self {
+            suffix: " materials",
+            _phantom: PhantomData,
+        }
+    }
+}
+
+impl<M: Material> MaterialAllocatorDiagnosticPlugin<M> {
+    /// Get the [`DiagnosticPath`] for slab count
+    pub fn slabs_diagnostic_path() -> DiagnosticPath {
+        DiagnosticPath::from_components(["material_allocator_slabs", type_name::<M>()])
+    }
+    /// Get the [`DiagnosticPath`] for total slabs size
+    pub fn slabs_size_diagnostic_path() -> DiagnosticPath {
+        DiagnosticPath::from_components(["material_allocator_slabs_size", type_name::<M>()])
+    }
+    /// Get the [`DiagnosticPath`] for material allocations
+    pub fn allocations_diagnostic_path() -> DiagnosticPath {
+        DiagnosticPath::from_components(["material_allocator_allocations", type_name::<M>()])
+    }
+}
+
+impl<M: Material> Plugin for MaterialAllocatorDiagnosticPlugin<M> {
+    fn build(&self, app: &mut bevy_app::App) {
+        app.register_diagnostic(
+            Diagnostic::new(Self::slabs_diagnostic_path()).with_suffix(" slabs"),
+        )
+        .register_diagnostic(
+            Diagnostic::new(Self::slabs_size_diagnostic_path()).with_suffix(" bytes"),
+        )
+        .register_diagnostic(
+            Diagnostic::new(Self::allocations_diagnostic_path()).with_suffix(self.suffix),
+        )
+        .init_resource::<MaterialAllocatorMeasurements<M>>()
+        .add_systems(PreUpdate, add_material_allocator_measurement::<M>);
+
+        if let Some(render_app) = app.get_sub_app_mut(RenderApp) {
+            render_app.add_systems(ExtractSchedule, measure_allocator::<M>);
+        }
+    }
+}
+
+#[derive(Debug, Resource)]
+struct MaterialAllocatorMeasurements<M: Material> {
+    slabs: AtomicUsize,
+    slabs_size: AtomicUsize,
+    allocations: AtomicU64,
+    _phantom: PhantomData<M>,
+}
+
+impl<M: Material> Default for MaterialAllocatorMeasurements<M> {
+    fn default() -> Self {
+        Self {
+            slabs: AtomicUsize::default(),
+            slabs_size: AtomicUsize::default(),
+            allocations: AtomicU64::default(),
+            _phantom: PhantomData,
+        }
+    }
+}
+
+fn add_material_allocator_measurement<M: Material>(
+    mut diagnostics: Diagnostics,
+    measurements: Res<MaterialAllocatorMeasurements<M>>,
+) {
+    diagnostics.add_measurement(
+        &MaterialAllocatorDiagnosticPlugin::<M>::slabs_diagnostic_path(),
+        || measurements.slabs.load(Ordering::Relaxed) as f64,
+    );
+    diagnostics.add_measurement(
+        &MaterialAllocatorDiagnosticPlugin::<M>::slabs_size_diagnostic_path(),
+        || measurements.slabs_size.load(Ordering::Relaxed) as f64,
+    );
+    diagnostics.add_measurement(
+        &MaterialAllocatorDiagnosticPlugin::<M>::allocations_diagnostic_path(),
+        || measurements.allocations.load(Ordering::Relaxed) as f64,
+    );
+}
+
+fn measure_allocator<M: Material + Any>(
+    measurements: Extract<Res<MaterialAllocatorMeasurements<M>>>,
+    allocators: Res<MaterialBindGroupAllocators>,
+) {
+    if let Some(allocator) = allocators.get(&TypeId::of::<M>()) {
+        measurements
+            .slabs
+            .store(allocator.slab_count(), Ordering::Relaxed);
+        measurements
+            .slabs_size
+            .store(allocator.slabs_size(), Ordering::Relaxed);
+        measurements
+            .allocations
+            .store(allocator.allocations(), Ordering::Relaxed);
+    }
+}

crates/bevy_pbr/src/lib.rs

@@ -29,6 +29,7 @@ mod cluster;
 mod components;
 pub mod decal;
 pub mod deferred;
+pub mod diagnostic;
 mod extended_material;
 mod fog;
 mod light;

crates/bevy_pbr/src/material_bind_groups.rs

@@ -590,6 +590,45 @@ impl MaterialBindGroupAllocator {
             }
         }
     }
+
+    /// Get number of allocated slabs for bindless material, returns 0 if it is
+    /// [`Self::NonBindless`].
+    pub fn slab_count(&self) -> usize {
+        match self {
+            Self::Bindless(bless) => bless.slabs.len(),
+            Self::NonBindless(_) => 0,
+        }
+    }
+
+    /// Get total size of slabs allocated for bindless material, returns 0 if it is
+    /// [`Self::NonBindless`].
+    pub fn slabs_size(&self) -> usize {
+        match self {
+            Self::Bindless(bless) => bless
+                .slabs
+                .iter()
+                .flat_map(|slab| {
+                    slab.data_buffers
+                        .iter()
+                        .map(|(_, buffer)| buffer.buffer.len())
+                })
+                .sum(),
+            Self::NonBindless(_) => 0,
+        }
+    }
+
+    /// Get number of bindless material allocations in slabs, returns 0 if it is
+    /// [`Self::NonBindless`].
+    pub fn allocations(&self) -> u64 {
+        match self {
+            Self::Bindless(bless) => bless
+                .slabs
+                .iter()
+                .map(|slab| u64::from(slab.allocated_resource_count))
+                .sum(),
+            Self::NonBindless(_) => 0,
+        }
+    }
 }
 
 impl MaterialBindlessIndexTable {

crates/bevy_render/src/diagnostic/mesh_allocator_diagnostic_plugin.rs

@@ -0,0 +1,91 @@
+use bevy_app::{Plugin, PreUpdate};
+use bevy_diagnostic::{Diagnostic, DiagnosticPath, Diagnostics, RegisterDiagnostic};
+use bevy_ecs::{resource::Resource, system::Res};
+use bevy_platform::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
+
+use crate::{mesh::allocator::MeshAllocator, Extract, ExtractSchedule, RenderApp};
+
+/// Number of meshes allocated by the allocator
+static MESH_ALLOCATOR_SLABS: DiagnosticPath = DiagnosticPath::const_new("mesh_allocator_slabs");
+
+/// Total size of all slabs
+static MESH_ALLOCATOR_SLABS_SIZE: DiagnosticPath =
+    DiagnosticPath::const_new("mesh_allocator_slabs_size");
+
+/// Number of meshes allocated into slabs
+static MESH_ALLOCATOR_ALLOCATIONS: DiagnosticPath =
+    DiagnosticPath::const_new("mesh_allocator_allocations");
+
+pub struct MeshAllocatorDiagnosticPlugin;
+
+impl MeshAllocatorDiagnosticPlugin {
+    /// Get the [`DiagnosticPath`] for slab count
+    pub fn slabs_diagnostic_path() -> &'static DiagnosticPath {
+        &MESH_ALLOCATOR_SLABS
+    }
+    /// Get the [`DiagnosticPath`] for total slabs size
+    pub fn slabs_size_diagnostic_path() -> &'static DiagnosticPath {
+        &MESH_ALLOCATOR_SLABS_SIZE
+    }
+    /// Get the [`DiagnosticPath`] for mesh allocations
+    pub fn allocations_diagnostic_path() -> &'static DiagnosticPath {
+        &MESH_ALLOCATOR_ALLOCATIONS
+    }
+}
+
+impl Plugin for MeshAllocatorDiagnosticPlugin {
+    fn build(&self, app: &mut bevy_app::App) {
+        app.register_diagnostic(
+            Diagnostic::new(MESH_ALLOCATOR_SLABS.clone()).with_suffix(" slabs"),
+        )
+        .register_diagnostic(
+            Diagnostic::new(MESH_ALLOCATOR_SLABS_SIZE.clone()).with_suffix(" bytes"),
+        )
+        .register_diagnostic(
+            Diagnostic::new(MESH_ALLOCATOR_ALLOCATIONS.clone()).with_suffix(" meshes"),
+        )
+        .init_resource::<MeshAllocatorMeasurements>()
+        .add_systems(PreUpdate, add_mesh_allocator_measurement);
+
+        if let Some(render_app) = app.get_sub_app_mut(RenderApp) {
+            render_app.add_systems(ExtractSchedule, measure_allocator);
+        }
+    }
+}
+
+#[derive(Debug, Default, Resource)]
+struct MeshAllocatorMeasurements {
+    slabs: AtomicUsize,
+    slabs_size: AtomicU64,
+    allocations: AtomicUsize,
+}
+
+fn add_mesh_allocator_measurement(
+    mut diagnostics: Diagnostics,
+    measurements: Res<MeshAllocatorMeasurements>,
+) {
+    diagnostics.add_measurement(&MESH_ALLOCATOR_SLABS, || {
+        measurements.slabs.load(Ordering::Relaxed) as f64
+    });
+    diagnostics.add_measurement(&MESH_ALLOCATOR_SLABS_SIZE, || {
+        measurements.slabs_size.load(Ordering::Relaxed) as f64
+    });
+    diagnostics.add_measurement(&MESH_ALLOCATOR_ALLOCATIONS, || {
+        measurements.allocations.load(Ordering::Relaxed) as f64
+    });
+}
+
+fn measure_allocator(
+    measurements: Extract<Res<MeshAllocatorMeasurements>>,
+    allocator: Res<MeshAllocator>,
+) {
+    measurements
+        .slabs
+        .store(allocator.slab_count(), Ordering::Relaxed);
+    measurements
+        .slabs_size
+        .store(allocator.slabs_size(), Ordering::Relaxed);
+    measurements
+        .allocations
+        .store(allocator.allocations(), Ordering::Relaxed);
+}

crates/bevy_render/src/diagnostic/mod.rs

@@ -3,6 +3,8 @@
 //! For more info, see [`RenderDiagnosticsPlugin`].
 
 pub(crate) mod internal;
+mod mesh_allocator_diagnostic_plugin;
+mod render_asset_diagnostic_plugin;
 #[cfg(feature = "tracing-tracy")]
 mod tracy_gpu;
 
@@ -16,6 +18,10 @@ use crate::{renderer::RenderAdapterInfo, RenderApp};
 use self::internal::{
     sync_diagnostics, DiagnosticsRecorder, Pass, RenderDiagnosticsMutex, WriteTimestamp,
 };
+pub use self::{
+    mesh_allocator_diagnostic_plugin::MeshAllocatorDiagnosticPlugin,
+    render_asset_diagnostic_plugin::RenderAssetDiagnosticPlugin,
+};
 
 use super::{RenderDevice, RenderQueue};
 

crates/bevy_render/src/diagnostic/render_asset_diagnostic_plugin.rs

@@ -0,0 +1,77 @@
+use core::{any::type_name, marker::PhantomData};
+
+use bevy_app::{Plugin, PreUpdate};
+use bevy_diagnostic::{Diagnostic, DiagnosticPath, Diagnostics, RegisterDiagnostic};
+use bevy_ecs::{resource::Resource, system::Res};
+use bevy_platform::sync::atomic::{AtomicUsize, Ordering};
+
+use crate::{
+    render_asset::{RenderAsset, RenderAssets},
+    Extract, ExtractSchedule, RenderApp,
+};
+
+pub struct RenderAssetDiagnosticPlugin<A: RenderAsset> {
+    suffix: &'static str,
+    _phantom: PhantomData<A>,
+}
+
+impl<A: RenderAsset> RenderAssetDiagnosticPlugin<A> {
+    pub fn new(suffix: &'static str) -> Self {
+        Self {
+            suffix,
+            _phantom: PhantomData,
+        }
+    }
+
+    pub fn render_asset_diagnostic_path() -> DiagnosticPath {
+        DiagnosticPath::from_components(["render_asset", type_name::<A>()])
+    }
+}
+
+impl<A: RenderAsset> Plugin for RenderAssetDiagnosticPlugin<A> {
+    fn build(&self, app: &mut bevy_app::App) {
+        app.register_diagnostic(
+            Diagnostic::new(Self::render_asset_diagnostic_path()).with_suffix(self.suffix),
+        )
+        .init_resource::<RenderAssetMeasurements<A>>()
+        .add_systems(PreUpdate, add_render_asset_measurement::<A>);
+
+        if let Some(render_app) = app.get_sub_app_mut(RenderApp) {
+            render_app.add_systems(ExtractSchedule, measure_render_asset::<A>);
+        }
+    }
+}
+
+#[derive(Debug, Resource)]
+struct RenderAssetMeasurements<A: RenderAsset> {
+    assets: AtomicUsize,
+    _phantom: PhantomData<A>,
+}
+
+impl<A: RenderAsset> Default for RenderAssetMeasurements<A> {
+    fn default() -> Self {
+        Self {
+            assets: AtomicUsize::default(),
+            _phantom: PhantomData,
+        }
+    }
+}
+
+fn add_render_asset_measurement<A: RenderAsset>(
+    mut diagnostics: Diagnostics,
+    measurements: Res<RenderAssetMeasurements<A>>,
+) {
+    diagnostics.add_measurement(
+        &RenderAssetDiagnosticPlugin::<A>::render_asset_diagnostic_path(),
+        || measurements.assets.load(Ordering::Relaxed) as f64,
+    );
+}
+
+fn measure_render_asset<A: RenderAsset>(
+    measurements: Extract<Res<RenderAssetMeasurements<A>>>,
+    assets: Res<RenderAssets<A>>,
+) {
+    measurements
+        .assets
+        .store(assets.iter().count(), Ordering::Relaxed);
+}