cranpose_render_wgpu/
pipeline_disk_cache.rs

1use std::{
2    hash::Hasher,
3    io::Write,
4    path::{Path, PathBuf},
5    sync::{
6        Arc, Condvar, Mutex, MutexGuard, OnceLock, PoisonError,
7        atomic::{AtomicBool, Ordering},
8    },
9};
10
11use cranpose_ui_graphics::FxHasher;
12use web_time::{Duration, Instant};
13
14use crate::{
15    debug_toggles::DebugToggle,
16    pipeline_recorder::PipelineRecorder,
17    pipeline_records::{self, PipelineRecords, Recorded, ShaderPipelineRecord},
18};
19
20static DISK_CACHE: DebugToggle = DebugToggle::new("CRANPOSE_PIPELINE_DISK_CACHE");
21
22fn disk_cache_enabled() -> bool {
23    !DISK_CACHE.equals("0")
24}
25
26static HOST_FILE: OnceLock<PathBuf> = OnceLock::new();
27
28/// Names the file the driver's compiled pipelines are kept in between runs.
29///
30/// This crate renders and does not know where an application may write, so the
31/// platform backend passes the path in once, before the first device is
32/// created. Without it the cache lives and dies with the process, and every
33/// launch pays the driver's compile again. A later call is ignored: the
34/// loaded blob and the watcher that writes it back must name one file.
35pub fn set_file_path(path: PathBuf) {
36    let _ = HOST_FILE.set(path);
37}
38
39pub(crate) fn file_path() -> Option<PathBuf> {
40    if !disk_cache_enabled() {
41        return None;
42    }
43    match crate::debug_toggles::debug_toggle_os("CRANPOSE_PIPELINE_CACHE_FILE") {
44        Some(path) if !path.is_empty() => Some(PathBuf::from(path)),
45        _ => HOST_FILE.get().cloned(),
46    }
47}
48
49/// What the last launches left for this one: the driver's compiled
50/// pipelines, where the device keeps them, and the pipelines those launches
51/// drew with.
52pub(crate) struct Loaded {
53    pub(crate) cache: Option<wgpu::PipelineCache>,
54    pub(crate) records: Arc<PipelineRecords>,
55    /// Whether another build wrote the records: an update, whose compiled
56    /// pipelines are gone.
57    pub(crate) updated: bool,
58    /// Whether the driver cache was filled from this build's compiled
59    /// pipelines, so a recorded pipeline builds from it as a cache hit.
60    pub(crate) holds_recorded: bool,
61}
62
63pub(crate) fn load(device: &wgpu::Device) -> Loaded {
64    let path = file_path();
65    let file = path.as_deref().and_then(|path| match std::fs::read(path) {
66        Ok(bytes) => Some(bytes),
67        Err(error) if error.kind() == std::io::ErrorKind::NotFound => None,
68        Err(error) => {
69            log::warn!("[pipeline-cache] unreadable {path:?}: {error}");
70            None
71        }
72    });
73    let contents = file.as_deref().and_then(current_contents);
74    let updated = contents
75        .as_ref()
76        .is_some_and(|contents| !contents.this_build);
77    let data = contents
78        .as_ref()
79        .and_then(|contents| contents.blob)
80        .filter(|blob| !blob.is_empty());
81    let cache = driver_cache(device, data, file.as_deref(), updated);
82    Loaded {
83        holds_recorded: data.is_some() && cache.is_some(),
84        cache,
85        records: Arc::new(contents.map(Contents::into_records).unwrap_or_default()),
86        updated,
87    }
88}
89
90/// Starts the cache file at `path` from `seed`, a cache file an app ships
91/// for its fresh installs, unless the file exists. A seed another GPU or
92/// driver wrote counts as another build's: its records name what the app
93/// draws, but its compiled pipelines are that driver's, so it is written
94/// under a build key no build has and loads without them.
95pub fn seed(path: &Path, mut bytes: Vec<u8>, same_driver: bool) {
96    if path.exists() {
97        return;
98    }
99    if !same_driver && let Some(build) = bytes.get_mut(..blob_key().len()) {
100        build.fill(0);
101    }
102    if let Some(parent) = path.parent()
103        && let Err(error) = std::fs::create_dir_all(parent)
104    {
105        log::warn!("[pipeline-cache] create_dir_all {parent:?}: {error}");
106        return;
107    }
108    match write_file(path, &bytes) {
109        Ok(()) => log::info!(
110            "[pipeline-cache] seeded {} B from the app{}",
111            bytes.len(),
112            if same_driver { "" } else { ", records only" },
113        ),
114        Err(error) => log::warn!("[pipeline-cache] seeding {path:?}: {error}"),
115    }
116}
117
118/// The device's pipeline cache, filled from `data` when the file had this
119/// build's.
120fn driver_cache(
121    device: &wgpu::Device,
122    data: Option<&[u8]>,
123    file: Option<&[u8]>,
124    updated: bool,
125) -> Option<wgpu::PipelineCache> {
126    if !device.features().contains(wgpu::Features::PIPELINE_CACHE) {
127        log::info!(
128            "[pipeline-cache] not offered by {:?}; compiled pipelines persist only as far \
129             as the driver's own cache does",
130            device.adapter_info().backend
131        );
132        return None;
133    }
134    // SAFETY: `data` is this build's own `get_data` output, and `fallback:
135    // true` has wgpu validate the header and fall back to an empty cache.
136    #[expect(unsafe_code)]
137    let cache = unsafe {
138        device.create_pipeline_cache(&wgpu::PipelineCacheDescriptor {
139            label: Some("cranpose pipeline disk cache"),
140            data,
141            fallback: true,
142        })
143    };
144    match (data, file) {
145        (Some(data), _) => log::info!("[pipeline-cache] loaded {} B from disk", data.len()),
146        (None, Some(file)) if updated => log::info!(
147            "[pipeline-cache] cold: dropped {} B compiled from other shaders",
148            file.len()
149        ),
150        (None, _) => log::info!("[pipeline-cache] cold (no blob on disk)"),
151    }
152    Some(cache)
153}
154
155const FILE_LAYOUT: u32 = 4;
156
157/// Names what fills a file: its layout, the framework's WGSL sources, and
158/// this crate's version, which changes with each release of the shader
159/// rewrites, pipeline layouts and translator between those sources and the
160/// driver, and with every pipeline key's layout.
161///
162/// The driver keeps every pipeline of the blob it loads in the cache it saves,
163/// so a blob kept across shader changes only grows, and the driver holds all
164/// of it resident. A file under another key loads cold and is replaced.
165fn blob_key() -> [u8; 8] {
166    let mut hasher = FxHasher::default();
167    hasher.write_u32(FILE_LAYOUT);
168    hasher.write_u64(cranpose_ui_graphics::framework_shaders::SOURCES_KEY);
169    hasher.write(env!("CARGO_PKG_VERSION").as_bytes());
170    hasher.finish().to_le_bytes()
171}
172
173struct Contents<'a> {
174    first_screen: FirstScreenKeys<'a>,
175    records: PipelineRecords,
176    /// Whether this build wrote the file.
177    this_build: bool,
178    blob: Option<&'a [u8]>,
179}
180
181impl Contents<'_> {
182    /// The records, the last first screen's shapes among them.
183    fn into_records(self) -> PipelineRecords {
184        let mut records = self.records;
185        records.shapes.splice(
186            0..0,
187            self.first_screen.map(|key| Recorded {
188                entry: key,
189                age: 0,
190                first_screen: true,
191            }),
192        );
193        records
194    }
195}
196
197#[derive(Clone)]
198struct FirstScreenKeys<'a>(&'a [u8]);
199
200impl Iterator for FirstScreenKeys<'_> {
201    type Item = u64;
202
203    fn next(&mut self) -> Option<u64> {
204        let (key, rest) = self.0.split_first_chunk::<8>()?;
205        self.0 = rest;
206        Some(u64::from_le_bytes(*key))
207    }
208
209    fn size_hint(&self) -> (usize, Option<usize>) {
210        (self.len(), Some(self.len()))
211    }
212}
213
214impl ExactSizeIterator for FirstScreenKeys<'_> {
215    fn len(&self) -> usize {
216        self.0.len() / 8
217    }
218}
219
220fn current_contents(file: &[u8]) -> Option<Contents<'_>> {
221    let (build, rest) = file.split_first_chunk::<8>()?;
222    let rest = rest.strip_prefix(crate::render::ShapePipelineKey::DISK_LAYOUT.as_slice())?;
223    let (count, rest) = rest.split_first_chunk::<4>()?;
224    let keys_len = usize::try_from(u32::from_le_bytes(*count))
225        .ok()?
226        .checked_mul(8)?;
227    let (keys, rest) = rest.split_at_checked(keys_len)?;
228    // The shape keys stand on their own: records that do not decode cost
229    // only themselves and the driver blob behind them.
230    let (records, blob) = pipeline_records::decode(rest).map_or_else(
231        || (PipelineRecords::default(), None),
232        |(records, blob)| (records, Some(blob)),
233    );
234    let this_build = *build == blob_key();
235    Some(Contents {
236        first_screen: FirstScreenKeys(keys),
237        records,
238        this_build,
239        blob: blob.filter(|_| this_build),
240    })
241}
242
243/// The file holding `shapes`, the first screen's ahead of the records, the
244/// other records, and `blob`.
245fn file_bytes<'a>(
246    shapes: impl Iterator<Item = Recorded<u64>> + Clone,
247    shaders: impl Iterator<Item = Recorded<&'a ShaderPipelineRecord>>,
248    fixed: impl Iterator<Item = Recorded<&'a str>>,
249    blob: &[u8],
250) -> Option<Vec<u8>> {
251    let mut bytes = Vec::with_capacity(4096 + blob.len());
252    bytes.extend_from_slice(&blob_key());
253    bytes.extend_from_slice(&crate::render::ShapePipelineKey::DISK_LAYOUT);
254    let first_screen = shapes.clone().filter(|shape| shape.first_screen);
255    bytes.extend_from_slice(
256        &u32::try_from(first_screen.clone().count())
257            .ok()?
258            .to_le_bytes(),
259    );
260    for shape in first_screen {
261        bytes.extend_from_slice(&shape.entry.to_le_bytes());
262    }
263    pipeline_records::encode(
264        shapes.filter(|shape| !shape.first_screen),
265        shaders,
266        fixed,
267        &mut bytes,
268    )?;
269    bytes.extend_from_slice(blob);
270    Some(bytes)
271}
272
273fn write_file(path: &Path, bytes: &[u8]) -> std::io::Result<()> {
274    let mut file = std::fs::File::create(path)?;
275    file.write_all(bytes)
276}
277
278pub(crate) fn persist(
279    cache: Option<&wgpu::PipelineCache>,
280    path: &Path,
281    previous: &PipelineRecords,
282    recorder: &PipelineRecorder,
283) {
284    let started = Instant::now();
285    let blob = cache
286        .and_then(wgpu::PipelineCache::get_data)
287        .unwrap_or_default();
288    let existing = std::fs::read(path).ok();
289    // A launch closed before its first frame keeps the last launch's records.
290    let (bytes, key_count) = if recorder.frame_drawn() {
291        let drawn = recorder.drawn();
292        let shapes = recent(
293            drawn.shapes.iter().copied(),
294            previous.shapes.iter().copied(),
295        );
296        let shaders = recent(
297            drawn.shaders.iter().map(Recorded::as_ref),
298            previous.shaders.iter().map(Recorded::as_ref),
299        );
300        let fixed = recent(
301            drawn.fixed.iter().copied(),
302            previous
303                .fixed
304                .iter()
305                .map(|label| label.as_ref().map(String::as_str)),
306        );
307        let count = shapes.clone().count() + shaders.clone().count() + fixed.clone().count();
308        (file_bytes(shapes, shaders, fixed, &blob), count)
309    } else {
310        (
311            file_bytes(
312                previous.shapes.iter().copied(),
313                previous.shaders.iter().map(Recorded::as_ref),
314                previous
315                    .fixed
316                    .iter()
317                    .map(|label| label.as_ref().map(String::as_str)),
318                &blob,
319            ),
320            previous.shapes.len() + previous.shaders.len() + previous.fixed.len(),
321        )
322    };
323    let Some(bytes) = bytes else {
324        return;
325    };
326    if existing.is_some_and(|existing| existing == bytes) {
327        return;
328    }
329    if let Some(parent) = path.parent()
330        && let Err(error) = std::fs::create_dir_all(parent)
331    {
332        log::warn!("[pipeline-cache] create_dir_all {parent:?}: {error}");
333        return;
334    }
335    let tmp = path.with_extension("tmp");
336    let written = write_file(&tmp, &bytes).and_then(|()| std::fs::rename(&tmp, path));
337    match written {
338        Ok(()) => log::info!(
339            "[pipeline-cache] persisted {} B and {} recorded pipelines in {:.1} ms",
340            blob.len(),
341            key_count,
342            crate::render::instant_ms(started, Instant::now()),
343        ),
344        Err(error) => log::warn!("[pipeline-cache] write {path:?}: {error}"),
345    }
346}
347
348/// How many launches a pipeline stays recorded after the last that drew
349/// with it.
350const RECORDED_LAUNCHES: u8 = 4;
351/// The most pipelines of each kind the file records.
352const MAX_RECORDS: usize = 128;
353
354/// This launch's pipelines, then those of the previous launches it did not
355/// draw with, a launch older: each while it is younger than
356/// [`RECORDED_LAUNCHES`], at most [`MAX_RECORDS`] of them.
357fn recent<T: Copy + PartialEq>(
358    session: impl Iterator<Item = Recorded<T>> + Clone,
359    previous: impl Iterator<Item = Recorded<T>> + Clone,
360) -> impl Iterator<Item = Recorded<T>> + Clone {
361    let drawn = session.clone();
362    session
363        .chain(previous.filter_map(move |record| {
364            let unseen = !drawn.clone().any(|drawn| drawn.entry == record.entry);
365            (unseen && record.age + 1 < RECORDED_LAUNCHES).then_some(Recorded {
366                age: record.age + 1,
367                first_screen: false,
368                ..record
369            })
370        }))
371        .take(MAX_RECORDS)
372}
373
374/// How long pipeline builds must pause before the cache is written. A burst
375/// of builds, such as a screen's first frames, is written once after its last
376/// build, and a session closed a moment after launch still keeps what the
377/// launch compiled, so the next launch skips those compiles.
378const PERSIST_QUIET: Duration = Duration::from_millis(500);
379
380static CHANGES: PersistSignal = PersistSignal::new();
381
382/// Counts a change the cache file should keep, such as a pipeline build,
383/// and wakes the watchers waiting to write it.
384pub(crate) fn note_change() {
385    CHANGES.note_change();
386}
387
388struct PersistSignal {
389    revision: Mutex<u64>,
390    changed: Condvar,
391}
392
393/// Why a watcher woke, with the build count at that moment.
394#[derive(Debug, PartialEq, Eq)]
395enum Wake {
396    /// Builds past the written count paused for the quiet period.
397    Quiet(u64),
398    /// The watcher's owner let it go.
399    Stopped(u64),
400}
401
402impl PersistSignal {
403    const fn new() -> Self {
404        Self {
405            revision: Mutex::new(0),
406            changed: Condvar::new(),
407        }
408    }
409
410    fn lock(&self) -> MutexGuard<'_, u64> {
411        self.revision.lock().unwrap_or_else(PoisonError::into_inner)
412    }
413
414    fn note_change(&self) {
415        *self.lock() += 1;
416        self.changed.notify_all();
417    }
418
419    fn stop(&self, stopped: &AtomicBool) {
420        stopped.store(true, Ordering::Release);
421        // Holding the lock orders the flag against a watcher between its
422        // check and its wait, so the wake below cannot be lost.
423        let _held = self.lock();
424        self.changed.notify_all();
425    }
426
427    fn wait_for_quiet(&self, written: u64, quiet: Duration, stopped: &AtomicBool) -> Wake {
428        let is_stopped = || stopped.load(Ordering::Acquire);
429        let mut revision = self
430            .changed
431            .wait_while(self.lock(), |revision| {
432                *revision == written && !is_stopped()
433            })
434            .unwrap_or_else(PoisonError::into_inner);
435        loop {
436            if is_stopped() {
437                return Wake::Stopped(*revision);
438            }
439            let seen = *revision;
440            let (next, waited) = self
441                .changed
442                .wait_timeout_while(revision, quiet, |revision| {
443                    *revision == seen && !is_stopped()
444                })
445                .unwrap_or_else(PoisonError::into_inner);
446            revision = next;
447            if waited.timed_out() {
448                return Wake::Quiet(*revision);
449            }
450        }
451    }
452}
453
454/// Keeps the cache's watcher writing new pipelines to disk. Dropping it
455/// writes what is still unwritten and ends the watcher.
456pub(crate) struct PersistWatcher {
457    stopped: Arc<AtomicBool>,
458    thread: Option<std::thread::JoinHandle<()>>,
459}
460
461impl Drop for PersistWatcher {
462    fn drop(&mut self) {
463        CHANGES.stop(&self.stopped);
464        if let Some(thread) = self.thread.take()
465            && thread.join().is_err()
466        {
467            log::warn!("[pipeline-cache] persist thread panicked during shutdown");
468        }
469    }
470}
471
472pub(crate) fn spawn_persist_watcher(
473    cache: Option<wgpu::PipelineCache>,
474    previous: Arc<PipelineRecords>,
475    recorder: PipelineRecorder,
476) -> Option<PersistWatcher> {
477    let path = file_path()?;
478    let stopped = Arc::new(AtomicBool::new(false));
479    let watcher_stopped = Arc::clone(&stopped);
480    let mut written = *CHANGES.lock();
481    let spawned = std::thread::Builder::new()
482        .name("cranpose-pl-cache".into())
483        .spawn(move || {
484            loop {
485                match CHANGES.wait_for_quiet(written, PERSIST_QUIET, &watcher_stopped) {
486                    Wake::Quiet(revision) => {
487                        persist(cache.as_ref(), &path, &previous, &recorder);
488                        written = revision;
489                    }
490                    Wake::Stopped(revision) => {
491                        if revision != written {
492                            persist(cache.as_ref(), &path, &previous, &recorder);
493                        }
494                        return;
495                    }
496                }
497            }
498        });
499    match spawned {
500        Ok(thread) => Some(PersistWatcher {
501            stopped,
502            thread: Some(thread),
503        }),
504        Err(error) => {
505            log::warn!("[pipeline-cache] persist thread failed to spawn: {error}");
506            None
507        }
508    }
509}
510
511#[cfg(test)]
512#[path = "tests/pipeline_disk_cache_tests.rs"]
513mod tests;
cranpose_render_wgpu/pipeline_disk_cache.rs

cranpose_render_wgpu/
pipeline_disk_cache.rs
