diff --git a/crates/bevy_core_pipeline/src/upscaling/mod.rs b/crates/bevy_core_pipeline/src/upscaling/mod.rs index 4ce91de393..37ae95e3cb 100644 --- a/crates/bevy_core_pipeline/src/upscaling/mod.rs +++ b/crates/bevy_core_pipeline/src/upscaling/mod.rs @@ -42,6 +42,7 @@ fn prepare_view_upscaling_pipelines( mut pipelines: ResMut>, blit_pipeline: Res, view_targets: Query<(Entity, &ViewTarget, Option<&ExtractedCamera>)>, + persistent_pipeline_cache: Option>, ) { let mut output_textures = >::default(); for (entity, view_target, camera) in view_targets.iter() { @@ -81,7 +82,7 @@ fn prepare_view_upscaling_pipelines( let pipeline = pipelines.specialize(&pipeline_cache, &blit_pipeline, key); // Ensure the pipeline is loaded before continuing the frame to prevent frames without any GPU work submitted - pipeline_cache.block_on_render_pipeline(pipeline); + pipeline_cache.block_on_render_pipeline(pipeline, persistent_pipeline_cache.as_deref()); commands .entity(entity) diff --git a/crates/bevy_render/src/render_resource/mod.rs b/crates/bevy_render/src/render_resource/mod.rs index f156b0ecb0..cc8ec94f87 100644 --- a/crates/bevy_render/src/render_resource/mod.rs +++ b/crates/bevy_render/src/render_resource/mod.rs @@ -7,6 +7,7 @@ mod bindless; mod buffer; mod buffer_vec; mod gpu_array_buffer; +mod persistent_pipeline_cache; mod pipeline; mod pipeline_cache; mod pipeline_specializer; @@ -25,6 +26,7 @@ pub use bindless::*; pub use buffer::*; pub use buffer_vec::*; pub use gpu_array_buffer::*; +pub use persistent_pipeline_cache::*; pub use pipeline::*; pub use pipeline_cache::*; pub use pipeline_specializer::*; diff --git a/crates/bevy_render/src/render_resource/persistent_pipeline_cache.rs b/crates/bevy_render/src/render_resource/persistent_pipeline_cache.rs new file mode 100644 index 0000000000..bca4aa4d81 --- /dev/null +++ b/crates/bevy_render/src/render_resource/persistent_pipeline_cache.rs @@ -0,0 +1,346 @@ +use crate::{renderer::RenderAdapterInfo, ExtractSchedule, RenderSystems}; +use alloc::sync::Arc; +use bevy_app::{App, Plugin}; +use bevy_ecs::{ + change_detection::{Res, ResMut}, + error::BevyError, + prelude::{not, resource_exists, IntoScheduleConfigs}, + resource::Resource, + system::{Commands, Local}, +}; +use bevy_platform::hash::FixedHasher; +use bevy_render::{render_resource::PipelineCache, renderer::RenderDevice, Extract, Render}; +use bevy_utils::WgpuWrapper; +use core::hash::BuildHasher; +use std::{ + fs, + fs::OpenOptions, + io, + io::Write, + path::{Path, PathBuf}, + sync::Mutex, + thread::JoinHandle, +}; +use thiserror::Error; +use tracing::{debug, warn}; +use wgpu::{hal::PipelineCacheError, Backend, PipelineCacheDescriptor}; + +/// Plugin for managing [`wgpu::PipelineCache`] resources across application runs. +/// +/// When pipelines are compiled by [`crate::PipelineCache`], if this plugin is enabled, it will +/// persist the pipeline cache to disk, allowing for faster startup times in subsequent runs. +/// +/// Note: This plugin is currently only supported on the Vulkan backend. +pub struct PersistentPipelineCachePlugin { + /// A unique key for the application, used to identify the cache directory. Should change + /// if the application is updated or if the cache should be invalidated. + pub application_key: &'static str, + /// The directory where the pipeline cache will be stored. + pub data_dir: PathBuf, + /// The eviction policy for the cache. + pub eviction_policy: EvictionPolicy, +} + +impl PersistentPipelineCachePlugin { + /// Creates a new instance of the `PersistentPipelineCachePlugin` with the specified + /// application key. + pub fn new(application_key: &'static str, data_dir: PathBuf) -> Self { + Self { + application_key, + data_dir, + eviction_policy: EvictionPolicy::Never, + } + } +} + +impl Plugin for PersistentPipelineCachePlugin { + fn build(&self, _app: &mut App) {} + + fn finish(&self, app: &mut App) { + if !self.data_dir.exists() || !self.data_dir.is_dir() { + warn!( + "PersistentPipelineCachePlugin data directory does not exist or is not a directory: {:?}", + self.data_dir + ); + return; + } + + if let Some(render_app) = app.get_sub_app_mut(bevy_render::RenderApp) { + let adapter_debug = render_app.world().resource::(); + if adapter_debug.backend != Backend::Vulkan { + warn!("PersistentPipelineCachePlugin is only supported on Vulkan backend.."); + return; + } + render_app + .add_systems( + ExtractSchedule, + extract_persistent_pipeline_cache + .run_if(not(resource_exists::)), + ) + .add_systems( + Render, + write_persistent_pipeline_cache + .run_if(resource_exists::) + .in_set(RenderSystems::Cleanup), + ); + }; + + app.insert_resource(PersistentPipelineCacheConfig { + application_key: self.application_key, + data_dir: self.data_dir.clone(), + eviction_policy: self.eviction_policy, + }); + } +} + +pub fn extract_persistent_pipeline_cache( + mut commands: Commands, + persistent_pipeline_cache_config: Extract>>, + adapter_debug: Res, + render_device: Res, +) -> Result<(), BevyError> { + let Some(persistent_pipeline_cache_config) = &*persistent_pipeline_cache_config else { + return Ok(()); + }; + + debug!( + "Extracting persistent pipeline cache with application key: {}", + persistent_pipeline_cache_config.application_key + ); + let cache_path = persistent_pipeline_cache_config + .data_dir + .join(persistent_pipeline_cache_config.application_key); + + match persistent_pipeline_cache_config.eviction_policy { + EvictionPolicy::Always => { + // Evict all existing data + if cache_path.exists() { + fs::remove_dir_all(&cache_path).map_err(PersistentPipelineCacheError::Io)?; + } + } + EvictionPolicy::Stale => { + // Evict all but matching our application key + if cache_path.exists() { + for entry in fs::read_dir(&cache_path).map_err(PersistentPipelineCacheError::Io)? { + // Check if the entry is a directory and doesn't match the cache path + let entry = entry.map_err(PersistentPipelineCacheError::Io)?; + if entry + .file_type() + .map_err(PersistentPipelineCacheError::Io)? + .is_dir() + && entry.file_name() != cache_path + { + fs::remove_dir_all(entry.path()) + .map_err(PersistentPipelineCacheError::Io)?; + debug!("Evicted stale pipeline cache at {:?}", entry.path()); + } + } + } + } + EvictionPolicy::Never => {} + } + + let cache_key = wgpu::util::pipeline_cache_key(&adapter_debug) + .ok_or(PersistentPipelineCacheError::InvalidAdapterKey)?; + let cache_path = cache_path.join(cache_key); + + // Ensure the cache directory exists + if let Some(parent) = cache_path.parent() { + if !parent.exists() { + debug!( + "Creating persistent pipeline cache directory at {:?}", + parent + ); + fs::create_dir_all(parent).map_err(PersistentPipelineCacheError::Io)?; + } + } + + let persistent_pipeline_cache = PersistentPipelineCache::new( + &render_device, + persistent_pipeline_cache_config.application_key, + &cache_path, + )?; + + commands.insert_resource(persistent_pipeline_cache); + Ok(()) +} + +pub fn write_persistent_pipeline_cache( + mut persistent_pipeline_cache: ResMut, + pipeline_cache: Res, + mut pipeline_cache_size: Local, +) -> Result<(), BevyError> { + let cache_size = pipeline_cache.size(); + if cache_size > *pipeline_cache_size { + persistent_pipeline_cache.write()?; + *pipeline_cache_size = cache_size; + } + + Ok(()) +} + +/// Configuration for the persistent pipeline cache. +#[derive(Resource)] +pub struct PersistentPipelineCacheConfig { + /// A unique key for the application, used to identify the cache directory. + pub application_key: &'static str, + /// The directory where the pipeline cache will be stored. + pub data_dir: PathBuf, + /// The eviction policy for the cache. + pub eviction_policy: EvictionPolicy, +} + +/// Resource for managing [`wgpu::PipelineCache`]. +#[derive(Resource)] +pub struct PersistentPipelineCache { + cache: Arc>, + write_lock: Arc>, + write_tasks: Vec>>, + cache_path: PathBuf, + data_key: u64, +} + +impl PersistentPipelineCache { + /// Create a new instance of the persistent pipeline cache with the given application key and + /// cache path. + pub fn new( + render_device: &RenderDevice, + app_key: &'static str, + cache_path: &Path, + ) -> Result { + // Get data if the cache file exists + let cache_data = if cache_path.exists() { + let data = fs::read(cache_path).map_err(PersistentPipelineCacheError::Io)?; + debug!( + "Loaded persistent pipeline cache from {:?}, size: {}", + cache_path, + data.len() + ); + Some(data) + } else { + debug!("Creating new persistent pipeline cache at {:?}", cache_path); + None + }; + // SAFETY: Data was created with a cache key that matches the adapter. + let cache = unsafe { + render_device + .wgpu_device() + .create_pipeline_cache(&PipelineCacheDescriptor { + data: cache_data.as_deref(), + label: app_key.into(), + fallback: true, + }) + }; + + let data_key = { + let hasher = FixedHasher; + hasher.hash_one(&cache_data) + }; + + Ok(PersistentPipelineCache { + cache: Arc::new(WgpuWrapper::new(cache)), + write_lock: Arc::new(Mutex::new(())), + write_tasks: vec![], + cache_path: cache_path.to_path_buf(), + data_key, + }) + } + + /// Get the cached data if it has changed since the last call. + pub fn get_data(&mut self) -> Option> { + let data = self.cache.get_data(); + let hasher = FixedHasher; + let data_key = hasher.hash_one(&data); + if self.data_key != data_key { + self.data_key = data_key; + return data; + } + + None + } + + /// Write the cached data to disk, if it has changed since the last write. + pub fn write(&mut self) -> Result<(), PersistentPipelineCacheError> { + // Process existing tasks + let mut pending_tasks = vec![]; + let mut error = None; + for task in self.write_tasks.drain(..) { + if task.is_finished() { + match task.join() { + Ok(Ok(())) => { + debug!("Persistent pipeline cache write task completed successfully."); + } + Ok(Err(err)) => { + warn!("Persistent pipeline cache write task failed: {}", err); + error = Some(err); + } + Err(err) => { + warn!("Persistent pipeline cache write task panicked: {:?}", err); + error = Some(PersistentPipelineCacheError::Io(io::Error::other( + "Persistent pipeline cache write task panicked", + ))); + } + } + } else { + pending_tasks.push(task); + } + } + + if let Some(err) = error { + return Err(err); + } + + if let Some(data) = self.get_data() { + let temp = self.cache_path.with_extension("tmp"); + let cache_path = self.cache_path.clone(); + let lock = self.write_lock.clone(); + let join_handle = std::thread::spawn(move || { + let _guard = lock + .lock() + .or(Err(PersistentPipelineCacheError::LockError))?; + let mut file = OpenOptions::new() + .write(true) + .create_new(true) + .open(&temp) + .map_err(PersistentPipelineCacheError::Io)?; + file.write_all(&data) + .map_err(PersistentPipelineCacheError::Io)?; + fs::rename(&temp, &cache_path).map_err(PersistentPipelineCacheError::Io)?; + Ok(()) + }); + self.write_tasks.push(join_handle); + } + + Ok(()) + } + + /// Get the underlying wgpu pipeline cache. + pub fn get_cache(&self) -> Arc> { + self.cache.clone() + } +} + +/// Describes the eviction policy for the persistent pipeline cache. +#[derive(Debug, Copy, Clone)] +pub enum EvictionPolicy { + /// Evict all existing data on startup. + Always, + /// Evict all but the data matching the application key on startup. + Stale, + /// Never evict any data. + Never, +} + +/// Error type for persistent pipeline cache operations. +#[derive(Debug, Error)] +#[error("Error while handling persistent pipeline cache")] +pub enum PersistentPipelineCacheError { + #[error(transparent)] + Io(#[from] io::Error), + #[error("Failed to create pipeline cache: {0}")] + DeviceError(#[from] PipelineCacheError), + #[error("Could not create cache key from adapter")] + InvalidAdapterKey, + #[error("Failed to acquire write lock for persistent pipeline cache")] + LockError, +} diff --git a/crates/bevy_render/src/render_resource/pipeline_cache.rs b/crates/bevy_render/src/render_resource/pipeline_cache.rs index 328c5e5600..a1474f6e8f 100644 --- a/crates/bevy_render/src/render_resource/pipeline_cache.rs +++ b/crates/bevy_render/src/render_resource/pipeline_cache.rs @@ -1,3 +1,4 @@ +use crate::render_resource::persistent_pipeline_cache::PersistentPipelineCache; use crate::{ render_resource::*, renderer::{RenderAdapter, RenderDevice}, @@ -695,9 +696,13 @@ impl PipelineCache { /// Wait for a render pipeline to finish compiling. #[inline] - pub fn block_on_render_pipeline(&mut self, id: CachedRenderPipelineId) { + pub fn block_on_render_pipeline( + &mut self, + id: CachedRenderPipelineId, + persistent_cache: Option<&PersistentPipelineCache>, + ) { if self.pipelines.len() <= id.0 { - self.process_queue(); + self.process_queue(persistent_cache); } let state = &mut self.pipelines[id.0].state; @@ -807,10 +812,12 @@ impl PipelineCache { &mut self, id: CachedPipelineId, descriptor: RenderPipelineDescriptor, + persistent_cache: Option<&PersistentPipelineCache>, ) -> CachedPipelineState { let device = self.device.clone(); let shader_cache = self.shader_cache.clone(); let layout_cache = self.layout_cache.clone(); + let cache = persistent_cache.map(PersistentPipelineCache::get_cache); create_pipeline_task( async move { @@ -903,7 +910,7 @@ impl PipelineCache { // TODO: Should this be the same as the vertex compilation options? compilation_options, }), - cache: None, + cache: cache.as_deref().map(|v| &**v), }; Ok(Pipeline::RenderPipeline( @@ -918,10 +925,12 @@ impl PipelineCache { &mut self, id: CachedPipelineId, descriptor: ComputePipelineDescriptor, + persistent_cache: Option<&PersistentPipelineCache>, ) -> CachedPipelineState { let device = self.device.clone(); let shader_cache = self.shader_cache.clone(); let layout_cache = self.layout_cache.clone(); + let cache = persistent_cache.map(PersistentPipelineCache::get_cache); create_pipeline_task( async move { @@ -962,7 +971,7 @@ impl PipelineCache { zero_initialize_workgroup_memory: descriptor .zero_initialize_workgroup_memory, }, - cache: None, + cache: cache.as_deref().map(|v| &**v), }; Ok(Pipeline::ComputePipeline( @@ -979,7 +988,7 @@ impl PipelineCache { /// be called manually to force creation at a different time. /// /// [`RenderSystems::Render`]: crate::RenderSystems::Render - pub fn process_queue(&mut self) { + pub fn process_queue(&mut self, persistent_cache: Option<&PersistentPipelineCache>) { let mut waiting_pipelines = mem::take(&mut self.waiting_pipelines); let mut pipelines = mem::take(&mut self.pipelines); @@ -996,22 +1005,26 @@ impl PipelineCache { } for id in waiting_pipelines { - self.process_pipeline(&mut pipelines[id], id); + self.process_pipeline(&mut pipelines[id], id, persistent_cache); } self.pipelines = pipelines; } - fn process_pipeline(&mut self, cached_pipeline: &mut CachedPipeline, id: usize) { + fn process_pipeline( + &mut self, + cached_pipeline: &mut CachedPipeline, + id: usize, + persistent_cache: Option<&PersistentPipelineCache>, + ) { match &mut cached_pipeline.state { CachedPipelineState::Queued => { cached_pipeline.state = match &cached_pipeline.descriptor { PipelineDescriptor::RenderPipelineDescriptor(descriptor) => { - self.start_create_render_pipeline(id, *descriptor.clone()) - } - PipelineDescriptor::ComputePipelineDescriptor(descriptor) => { - self.start_create_compute_pipeline(id, *descriptor.clone()) + self.start_create_render_pipeline(id, *descriptor.clone(), persistent_cache) } + PipelineDescriptor::ComputePipelineDescriptor(descriptor) => self + .start_create_compute_pipeline(id, *descriptor.clone(), persistent_cache), }; } @@ -1051,8 +1064,11 @@ impl PipelineCache { self.waiting_pipelines.insert(id); } - pub(crate) fn process_pipeline_queue_system(mut cache: ResMut) { - cache.process_queue(); + pub(crate) fn process_pipeline_queue_system( + mut cache: ResMut, + persistent_pipeline_cache: Option>, + ) { + cache.process_queue(persistent_pipeline_cache.as_deref()); } pub(crate) fn extract_shaders( @@ -1080,6 +1096,10 @@ impl PipelineCache { } } } + + pub fn size(&self) -> usize { + self.pipelines.len() + } } #[cfg(all(