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