cranelift_codegen/
incremental_cache.rs

1//! This module provides a set of primitives that allow implementing an incremental cache on top of
2//! Cranelift, making it possible to reuse previous compiled artifacts for functions that have been
3//! compiled previously.
4//!
5//! This set of operation is experimental and can be enabled using the Cargo feature
6//! `incremental-cache`.
7//!
8//! This can bring speedups in different cases: change-code-and-immediately-recompile iterations
9//! get faster, modules sharing lots of code can reuse each other's artifacts, etc.
10//!
11//! The three main primitives are the following:
12//! - `compute_cache_key` is used to compute the cache key associated to a `Function`. This is
13//!   basically the content of the function, modulo a few things the caching system is resilient to.
14//! - `serialize_compiled` is used to serialize the result of a compilation, so it can be reused
15//!   later on by...
16//! - `try_finish_recompile`, which reads binary blobs serialized with `serialize_compiled`,
17//!   re-creating the compilation artifact from those.
18//!
19//! The `CacheStore` trait and `Context::compile_with_cache` method are provided as
20//! high-level, easy-to-use facilities to make use of that cache, and show an example of how to use
21//! the above three primitives to form a full incremental caching system.
22
23use core::fmt;
24use core::hash::{Hash, Hasher};
25
26use crate::alloc::vec::Vec;
27use crate::ir::Function;
28use crate::ir::function::{FunctionStencil, VersionMarker};
29use crate::machinst::{CompiledCode, CompiledCodeStencil};
30use crate::result::CompileResult;
31use crate::{CompileError, Context, trace};
32use crate::{isa::TargetIsa, timing};
33use alloc::borrow::Cow;
34use cranelift_control::ControlPlane;
35
36impl Context {
37    /// Compile the function, as in `compile`, but tries to reuse compiled artifacts from former
38    /// compilations using the provided cache store.
39    pub fn compile_with_cache(
40        &mut self,
41        isa: &dyn TargetIsa,
42        cache_store: &mut dyn CacheKvStore,
43        ctrl_plane: &mut ControlPlane,
44    ) -> CompileResult<'_, (&CompiledCode, bool)> {
45        let cache_key_hash = {
46            let _tt = timing::try_incremental_cache();
47
48            let cache_key_hash = compute_cache_key(isa, &self.func);
49
50            if let Some(blob) = cache_store.get(&cache_key_hash.0) {
51                match try_finish_recompile(&self.func, &blob) {
52                    Ok(compiled_code) => {
53                        let info = compiled_code.code_info();
54
55                        if isa.flags().enable_incremental_compilation_cache_checks() {
56                            let actual_result = self.compile(isa, ctrl_plane)?;
57                            assert_eq!(*actual_result, compiled_code);
58                            assert_eq!(actual_result.code_info(), info);
59                            // no need to set `compiled_code` here, it's set by `compile()`.
60                            return Ok((actual_result, true));
61                        }
62
63                        let compiled_code = self.compiled_code.insert(compiled_code);
64                        return Ok((compiled_code, true));
65                    }
66                    Err(err) => {
67                        trace!("error when finishing recompilation: {err}");
68                    }
69                }
70            }
71
72            cache_key_hash
73        };
74
75        let stencil = self
76            .compile_stencil(isa, ctrl_plane)
77            .map_err(|err| CompileError {
78                inner: err,
79                func: &self.func,
80            })?;
81
82        let stencil = {
83            let _tt = timing::store_incremental_cache();
84            let (stencil, res) = serialize_compiled(stencil);
85            if let Ok(blob) = res {
86                cache_store.insert(&cache_key_hash.0, blob);
87            }
88            stencil
89        };
90
91        let compiled_code = self
92            .compiled_code
93            .insert(stencil.apply_params(&self.func.params));
94
95        Ok((compiled_code, false))
96    }
97}
98
99/// Backing storage for an incremental compilation cache, when enabled.
100pub trait CacheKvStore {
101    /// Given a cache key hash, retrieves the associated opaque serialized data.
102    fn get(&self, key: &[u8]) -> Option<Cow<'_, [u8]>>;
103
104    /// Given a new cache key and a serialized blob obtained from `serialize_compiled`, stores it
105    /// in the cache store.
106    fn insert(&mut self, key: &[u8], val: Vec<u8>);
107}
108
109/// Hashed `CachedKey`, to use as an identifier when looking up whether a function has already been
110/// compiled or not.
111#[derive(Clone, Hash, PartialEq, Eq)]
112pub struct CacheKeyHash([u8; 32]);
113
114impl std::fmt::Display for CacheKeyHash {
115    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
116        write!(f, "CacheKeyHash:{:?}", self.0)
117    }
118}
119
120#[derive(serde_derive::Serialize, serde_derive::Deserialize)]
121struct CachedFunc {
122    // Note: The version marker must be first to ensure deserialization stops in case of a version
123    // mismatch before attempting to deserialize the actual compiled code.
124    version_marker: VersionMarker,
125    stencil: CompiledCodeStencil,
126}
127
128/// Key for caching a single function's compilation.
129///
130/// If two functions get the same `CacheKey`, then we can reuse the compiled artifacts, modulo some
131/// fixups.
132///
133/// Note: the key will be invalidated across different versions of cranelift, as the
134/// `FunctionStencil` contains a `VersionMarker` itself.
135struct CacheKey<'a> {
136    stencil: &'a FunctionStencil,
137    isa: &'a dyn TargetIsa,
138}
139
140impl<'a> Hash for CacheKey<'a> {
141    fn hash<H: Hasher>(&self, state: &mut H) {
142        self.stencil.hash(state);
143        self.isa.name().hash(state);
144        self.isa.triple().hash(state);
145        self.isa.flags().hash(state);
146        self.isa.isa_flags_hash_key().hash(state);
147    }
148}
149
150impl<'a> CacheKey<'a> {
151    /// Creates a new cache store key for a function.
152    ///
153    /// This is a bit expensive to compute, so it should be cached and reused as much as possible.
154    fn new(isa: &'a dyn TargetIsa, f: &'a Function) -> Self {
155        CacheKey {
156            stencil: &f.stencil,
157            isa,
158        }
159    }
160}
161
162/// Compute a cache key, and hash it on your behalf.
163///
164/// Since computing the `CacheKey` is a bit expensive, it should be done as least as possible.
165pub fn compute_cache_key(isa: &dyn TargetIsa, func: &Function) -> CacheKeyHash {
166    use core::hash::{Hash as _, Hasher};
167    use sha2::Digest as _;
168
169    struct Sha256Hasher(sha2::Sha256);
170
171    impl Hasher for Sha256Hasher {
172        fn finish(&self) -> u64 {
173            panic!("Sha256Hasher doesn't support finish!");
174        }
175        fn write(&mut self, bytes: &[u8]) {
176            self.0.update(bytes);
177        }
178    }
179
180    let cache_key = CacheKey::new(isa, func);
181
182    let mut hasher = Sha256Hasher(sha2::Sha256::new());
183    cache_key.hash(&mut hasher);
184    let hash: [u8; 32] = hasher.0.finalize().into();
185
186    CacheKeyHash(hash)
187}
188
189/// Given a function that's been successfully compiled, serialize it to a blob that the caller may
190/// store somewhere for future use by `try_finish_recompile`.
191///
192/// As this function requires ownership on the `CompiledCodeStencil`, it gives it back at the end
193/// of the function call. The value is left untouched.
194pub fn serialize_compiled(
195    result: CompiledCodeStencil,
196) -> (CompiledCodeStencil, Result<Vec<u8>, postcard::Error>) {
197    let cached = CachedFunc {
198        version_marker: VersionMarker,
199        stencil: result,
200    };
201    let result = postcard::to_allocvec(&cached);
202    (cached.stencil, result)
203}
204
205/// An error returned when recompiling failed.
206#[derive(Debug)]
207pub enum RecompileError {
208    /// The version embedded in the cache entry isn't the same as cranelift's current version.
209    VersionMismatch,
210    /// An error occurred while deserializing the cache entry.
211    Deserialize(postcard::Error),
212}
213
214impl fmt::Display for RecompileError {
215    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
216        match self {
217            RecompileError::VersionMismatch => write!(f, "cranelift version mismatch",),
218            RecompileError::Deserialize(err) => {
219                write!(f, "postcard failed during deserialization: {err}")
220            }
221        }
222    }
223}
224
225/// Given a function that's been precompiled and its entry in the caching storage, try to shortcut
226/// compilation of the given function.
227///
228/// Precondition: the bytes must have retrieved from a cache store entry which hash value
229/// is strictly the same as the `Function`'s computed hash retrieved from `compute_cache_key`.
230pub fn try_finish_recompile(func: &Function, bytes: &[u8]) -> Result<CompiledCode, RecompileError> {
231    match postcard::from_bytes::<CachedFunc>(bytes) {
232        Ok(result) => {
233            if result.version_marker != func.stencil.version_marker {
234                Err(RecompileError::VersionMismatch)
235            } else {
236                Ok(result.stencil.apply_params(&func.params))
237            }
238        }
239        Err(err) => Err(RecompileError::Deserialize(err)),
240    }
241}
cranelift_codegen/incremental_cache.rs

cranelift_codegen/
incremental_cache.rs
