av_denoise/
cache.rs

1use std::path::PathBuf;
2
3use cubecl::config::cache::CacheConfig;
4use cubecl::config::{CubeClRuntimeConfig, RuntimeConfig};
5
6/// Name of the environment variable that, when set, redirects CubeCL's
7/// compilation and autotune caches to the given directory.
8pub const COMPILATION_CACHE_ENV: &str = "AV_DENOISE_COMPILATION_CACHE";
9
10/// The CubeCL global config was already initialized before this
11/// helper ran, so the override can no longer be installed.
12#[derive(Debug, thiserror::Error)]
13#[error(
14    "CubeCL global config already initialized. Call apply_compilation_cache_env() before any Denoiser::create"
15)]
16pub struct CacheAlreadyInitialisedError;
17
18/// Apply the `AV_DENOISE_COMPILATION_CACHE` override if set.
19///
20/// Returns `Ok(Some(path))` if the override was installed, `Ok(None)`
21/// if the env var was unset, or `Err` if the global config was already
22/// read by something else (e.g. a previously-created CubeCL client).
23///
24/// Must be called before any CubeCL client is created.
25pub fn apply_compilation_cache_env() -> Result<Option<PathBuf>, CacheAlreadyInitialisedError> {
26    let Some(raw) = std::env::var_os(COMPILATION_CACHE_ENV) else {
27        return Ok(None);
28    };
29    let path = PathBuf::from(raw);
30
31    let mut cfg = CubeClRuntimeConfig::from_current_dir().override_from_env();
32    cfg.compilation.cache = Some(CacheConfig::File(path.clone()));
33    cfg.autotune.cache = CacheConfig::File(path.clone());
34
35    // `RuntimeConfig::set` panics if the singleton has already been
36    // initialized. Catch that so callers get a typed error instead of
37    // an abort.
38    //
39    // Unfortunately, CubeCL doesn't currently expose a graceful variant
40    // of this function, so we're left with catching the panic.
41    std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
42        CubeClRuntimeConfig::set(cfg);
43    }))
44    .map_err(|_| CacheAlreadyInitialisedError)?;
45
46    Ok(Some(path))
47}
av_denoise/cache.rs

av_denoise/
cache.rs
