rusticl: kernel caching

Signed-off-by: Karol Herbst <kherbst@redhat.com>
Acked-by: Alyssa Rosenzweig <alyssa.rosenzweig@collabora.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/15439>

src/gallium/frontends/rusticl/core/program.rs

@@ -5,6 +5,7 @@ use crate::impl_cl_type_trait;
 
 use mesa_rust::compiler::clc::*;
 use mesa_rust::compiler::nir::*;
+use mesa_rust::util::disk_cache::*;
 use mesa_rust_gen::*;
 use rusticl_opencl_gen::*;
 
@@ -17,6 +18,7 @@ use std::slice;
 use std::sync::Arc;
 use std::sync::Mutex;
 use std::sync::MutexGuard;
+use std::sync::Once;
 
 const BIN_HEADER_SIZE_V1: usize =
     // 1. format version
@@ -28,6 +30,19 @@ const BIN_HEADER_SIZE_V1: usize =
 
 const BIN_HEADER_SIZE: usize = BIN_HEADER_SIZE_V1;
 
+// kernel cache
+static mut DISK_CACHE: Option<DiskCache> = None;
+static DISK_CACHE_ONCE: Once = Once::new();
+
+fn get_disk_cache() -> &'static Option<DiskCache> {
+    unsafe {
+        DISK_CACHE_ONCE.call_once(|| {
+            DISK_CACHE = DiskCache::new("rusticl", "rusticl", 0);
+        });
+        &DISK_CACHE
+    }
+}
+
 #[repr(C)]
 pub struct Program {
     pub base: CLObjectBase<CL_INVALID_PROGRAM>,
@@ -287,8 +302,13 @@ impl Program {
         let lib = options.contains("-create-library");
 
         let args = prepare_options(&options, dev);
-        let (spirv, log) =
+        let (spirv, log) = spirv::SPIRVBin::from_clc(
-            spirv::SPIRVBin::from_clc(&self.src, &args, &Vec::new(), dev.cl_features());
+            &self.src,
+            &args,
+            &Vec::new(),
+            get_disk_cache(),
+            dev.cl_features(),
+        );
 
         d.log = log;
         d.options = options;
@@ -333,7 +353,13 @@ impl Program {
         let d = Self::dev_build_info(&mut info, dev);
         let args = prepare_options(&options, dev);
 
-        let (spirv, log) = spirv::SPIRVBin::from_clc(&self.src, &args, headers, dev.cl_features());
+        let (spirv, log) = spirv::SPIRVBin::from_clc(
+            &self.src,
+            &args,
+            headers,
+            get_disk_cache(),
+            dev.cl_features(),
+        );
 
         d.spirv = spirv;
         d.log = log;

src/gallium/frontends/rusticl/mesa/compiler/clc/spirv.rs

@@ -1,5 +1,6 @@
 use crate::compiler::nir::*;
 use crate::pipe::screen::*;
+use crate::util::disk_cache::*;
 
 use mesa_rust_gen::*;
 use mesa_rust_util::string::*;
@@ -41,8 +42,33 @@ impl SPIRVBin {
         source: &CString,
         args: &[CString],
         headers: &[CLCHeader],
+        cache: &Option<DiskCache>,
         features: clc_optional_features,
     ) -> (Option<Self>, String) {
+        let mut hash_key = None;
+        let has_includes = args.iter().any(|a| a.as_bytes()[0..2] == *b"-I");
+
+        if let Some(cache) = cache {
+            if !has_includes {
+                let mut key = Vec::new();
+
+                key.extend_from_slice(source.as_bytes());
+                args.iter()
+                    .for_each(|a| key.extend_from_slice(a.as_bytes()));
+                headers.iter().for_each(|h| {
+                    key.extend_from_slice(h.name.as_bytes());
+                    key.extend_from_slice(h.source.as_bytes());
+                });
+
+                let mut key = cache.gen_key(&key);
+                if let Some(data) = cache.get(&mut key) {
+                    return (Some(Self::from_bin(&data, false)), String::from(""));
+                }
+
+                hash_key = Some(key);
+            }
+        }
+
         let c_headers: Vec<_> = headers
             .iter()
             .map(|h| clc_named_value {
@@ -77,16 +103,27 @@ impl SPIRVBin {
         let res = unsafe { clc_compile_c_to_spirv(&args, &logger, &mut out) };
 
         let res = if res {
-            Some(SPIRVBin {
+            let spirv = SPIRVBin {
                 spirv: out,
                 info: None,
-            })
+            };
+
+            // add cache entry
+            if !has_includes {
+                if let Some(mut key) = hash_key {
+                    cache.as_ref().unwrap().put(spirv.to_bin(), &mut key);
+                }
+            }
+
+            Some(spirv)
         } else {
             None
         };
+
         (res, msgs.join("\n"))
     }
 
+    // TODO cache linking, parsing is around 25% of link time
     pub fn link(spirvs: &[&SPIRVBin], library: bool) -> (Option<Self>, String) {
         let bins: Vec<_> = spirvs.iter().map(|s| &s.spirv as *const _).collect();
 
@@ -229,7 +266,7 @@ impl SPIRVBin {
     pub fn get_lib_clc(screen: &PipeScreen) -> Option<NirShader> {
         let nir_options = screen.nir_shader_compiler_options(pipe_shader_type::PIPE_SHADER_COMPUTE);
         let spirv_options = Self::get_spirv_options(true, ptr::null());
-        let shader_cache = screen.shader_cache();
+        let shader_cache = DiskCacheBorrowed::as_ptr(&screen.shader_cache());
         NirShader::new(unsafe {
             nir_load_libclc_shader(64, shader_cache, &spirv_options, nir_options)
         })

src/gallium/frontends/rusticl/mesa/lib.rs

@@ -1,2 +1,3 @@
 pub mod compiler;
 pub mod pipe;
+pub mod util;

src/gallium/frontends/rusticl/mesa/pipe/screen.rs

@@ -2,6 +2,7 @@ use crate::compiler::nir::NirShader;
 use crate::pipe::context::*;
 use crate::pipe::device::*;
 use crate::pipe::resource::*;
+use crate::util::disk_cache::*;
 
 use mesa_rust_gen::*;
 use mesa_rust_util::string::*;
@@ -233,14 +234,16 @@ impl PipeScreen {
         }
     }
 
-    pub fn shader_cache(&self) -> *mut disk_cache {
+    pub fn shader_cache(&self) -> Option<DiskCacheBorrowed> {
         let s = &mut unsafe { *self.screen };
 
-        if let Some(func) = s.get_disk_shader_cache {
+        let ptr = if let Some(func) = s.get_disk_shader_cache {
             unsafe { func(self.screen) }
         } else {
             ptr::null_mut()
-        }
+        };
+
+        DiskCacheBorrowed::from_ptr(ptr)
     }
 
     pub fn finalize_nir(&self, nir: &NirShader) {

src/gallium/frontends/rusticl/mesa/util.rs

@@ -0,0 +1 @@
+pub mod disk_cache;

src/gallium/frontends/rusticl/mesa/util/disk_cache.rs

@@ -0,0 +1,120 @@
+use mesa_rust_gen::*;
+
+use std::ffi::CString;
+use std::ops::Deref;
+use std::ptr;
+use std::ptr::NonNull;
+use std::slice;
+
+pub struct DiskCacheBorrowed {
+    cache: NonNull<disk_cache>,
+}
+
+pub struct DiskCache {
+    inner: DiskCacheBorrowed,
+}
+
+// disk_cache is thread safe
+unsafe impl Sync for DiskCacheBorrowed {}
+
+impl DiskCacheBorrowed {
+    pub fn from_ptr(cache: *mut disk_cache) -> Option<Self> {
+        NonNull::new(cache).map(|c| Self { cache: c })
+    }
+
+    pub fn put(&self, data: &[u8], key: &mut cache_key) {
+        unsafe {
+            disk_cache_put(
+                self.cache.as_ptr(),
+                key,
+                data.as_ptr().cast(),
+                data.len(),
+                ptr::null_mut(),
+            );
+        }
+    }
+
+    pub fn get(&self, key: &mut cache_key) -> Option<DiskCacheEntry> {
+        let mut size = 0;
+
+        unsafe {
+            let data = disk_cache_get(self.cache.as_ptr(), key, &mut size);
+            if data.is_null() {
+                None
+            } else {
+                Some(DiskCacheEntry {
+                    data: slice::from_raw_parts_mut(data.cast(), size),
+                })
+            }
+        }
+    }
+
+    pub fn gen_key(&self, data: &[u8]) -> cache_key {
+        let mut key = cache_key::default();
+
+        unsafe {
+            disk_cache_compute_key(
+                self.cache.as_ptr(),
+                data.as_ptr().cast(),
+                data.len(),
+                &mut key,
+            );
+        }
+
+        key
+    }
+
+    pub fn as_ptr(s: &Option<Self>) -> *mut disk_cache {
+        if let Some(s) = s {
+            s.cache.as_ptr()
+        } else {
+            ptr::null_mut()
+        }
+    }
+}
+
+impl DiskCache {
+    pub fn new(name: &str, driver_id: &str, flags: u64) -> Option<Self> {
+        let c_name = CString::new(name).unwrap();
+        let c_id = CString::new(driver_id).unwrap();
+        let cache = unsafe { disk_cache_create(c_name.as_ptr(), c_id.as_ptr(), flags) };
+
+        DiskCacheBorrowed::from_ptr(cache).map(|c| Self { inner: c })
+    }
+}
+
+impl Deref for DiskCache {
+    type Target = DiskCacheBorrowed;
+
+    fn deref(&self) -> &Self::Target {
+        &self.inner
+    }
+}
+
+impl Drop for DiskCache {
+    fn drop(&mut self) {
+        unsafe {
+            disk_cache_destroy(self.cache.as_ptr());
+        }
+    }
+}
+
+pub struct DiskCacheEntry<'a> {
+    data: &'a mut [u8],
+}
+
+impl<'a> Deref for DiskCacheEntry<'a> {
+    type Target = [u8];
+
+    fn deref(&self) -> &Self::Target {
+        self.data
+    }
+}
+
+impl<'a> Drop for DiskCacheEntry<'a> {
+    fn drop(&mut self) {
+        unsafe {
+            free(self.data.as_mut_ptr().cast());
+        }
+    }
+}

src/gallium/frontends/rusticl/meson.build

@@ -199,6 +199,8 @@ rusticl_mesa_bindings_rs = rust.bindgen(
   args : [
     rusticl_bindgen_args,
     '--whitelist-function', 'clc_.*',
+    '--whitelist-function', 'disk_cache_.*',
+    '--whitelist-function', 'free',
     '--whitelist-function', 'glsl_.*',
     '--whitelist-function', 'malloc',
     '--whitelist-function', 'nir_.*',

src/gallium/frontends/rusticl/rusticl_mesa_bindings.h

@@ -12,6 +12,7 @@
 #include "pipe/p_state.h"
 #include "pipe-loader/pipe_loader.h"
 
+#include "util/disk_cache.h"
 #include "util/u_printf.h"
 #include "util/u_sampler.h"