ruda_runtime/runtime/tune/
operation.rs

1use alloc::boxed::Box;
2use alloc::string::String;
3use alloc::sync::Arc;
4use alloc::vec::Vec;
5use core::fmt::{Debug, Display, Write};
6use core::hash::Hash;
7use spin::Once;
8
9use alloc::format;
10
11use super::{
12    AutotuneError, input_generator::InputGenerator, key_generator::KeyGenerator,
13    tune_inputs::TuneInputs,
14};
15use super::{Tunable, TunePlan};
16
17/// A type-erased delegate for a tunable function.
18///
19/// The lifetime `'inp` is the lifetime of the input data, the function must be defined such that
20/// it can be called for any lifetime `inp` and produce a `Result<Out, AutotuneError>`.
21type TuneDelegate<I, Out> =
22    dyn for<'inp> Fn(<I as TuneInputs>::At<'inp>) -> Result<Out, AutotuneError> + Send + Sync;
23
24/// A named, type-erased tunable function stored in a [`TunableSet`]. Constructed via
25/// [`Tunable::new`](super::Tunable::new); callers don't name this type directly.
26#[derive(new)]
27pub struct TuneFn<I: TuneInputs, Out> {
28    pub(crate) name: String,
29    func: Box<TuneDelegate<I, Out>>,
30}
31
32impl<I: TuneInputs, Out: 'static> TuneFn<I, Out> {
33    /// Run the wrapped function on the given inputs.
34    pub fn execute<'a>(&self, inputs: <I as TuneInputs>::At<'a>) -> Result<Out, AutotuneError> {
35        (self.func)(inputs)
36    }
37}
38
39/// A set of candidate tunable functions for autotune, sharing a key generator and an
40/// input generator. See [`TuneInputs`] for the `F` parameter.
41pub struct TunableSet<K: AutotuneKey, F: TuneInputs, Output: 'static> {
42    tunables: Vec<Tunable<K, F, Output>>,
43    key_gen: Arc<dyn KeyGenerator<K, F> + Send + Sync>,
44    input_gen: Arc<dyn InputGenerator<K, F> + Send + Sync>,
45    stack_reference: Option<usize>,
46    stack_revision: String,
47    stack_workload: Option<Arc<dyn for<'a> Fn(&F::At<'a>) -> String + Send + Sync>>,
48    checksum: Once<String>,
49    stack_manifest: Once<String>,
50    #[cfg(std_io)]
51    stack_plan_cache: Option<spin::Mutex<alloc::collections::VecDeque<(K, String, Arc<StaticStackPlan>)>>>,
52}
53
54#[cfg(std_io)]
55pub(crate) struct StaticStackPlan {
56    pub(crate) indices: Vec<usize>,
57    pub(crate) candidates: Vec<super::stack::Candidate>,
58    pub(crate) key: String,
59}
60
61impl<K: AutotuneKey, F: TuneInputs, Output: 'static> TunableSet<K, F, Output> {
62    /// The number of tunables in the set.
63    pub fn len(&self) -> usize {
64        self.tunables.len()
65    }
66
67    /// Whether this set contains no tunables.
68    pub fn is_empty(&self) -> bool {
69        self.tunables.is_empty()
70    }
71
72    /// Create a tunable set from a key generator and an input generator.
73    pub fn new(key_gen: impl KeyGenerator<K, F>, input_gen: impl InputGenerator<K, F>) -> Self {
74        Self {
75            tunables: Default::default(),
76            input_gen: Arc::new(input_gen),
77            key_gen: Arc::new(key_gen),
78            stack_reference: None,
79            stack_revision: String::new(),
80            stack_workload: None,
81            checksum: Once::new(),
82            stack_manifest: Once::new(),
83            #[cfg(std_io)]
84            stack_plan_cache: None,
85        }
86    }
87
88    /// Register an exact workload signature and a known-safe reference for full-stack tuning.
89    /// The input generator MUST isolate all mutable state for each benchmark invocation.
90    /// Include dtype, exact shape/strides, precision, optional inputs and semantic parameters.
91    pub fn with_stack_tuning(
92        mut self, reference: usize, revision: &str,
93        workload: impl for<'a> Fn(&F::At<'a>) -> String + Send + Sync + 'static,
94    ) -> Self {
95        self.stack_reference = Some(reference);
96        self.stack_revision = revision.into();
97        self.stack_workload = Some(Arc::new(workload));
98        self.stack_manifest = Once::new();
99        #[cfg(std_io)]
100        if let Some(cache) = &self.stack_plan_cache { cache.lock().clear(); }
101        self
102    }
103    pub fn stack_reference(&self) -> Option<usize> { self.stack_reference }
104    pub fn stack_workload<'a>(&self, inputs: &F::At<'a>) -> Option<String> {
105        self.stack_workload.as_ref().map(|f| f(inputs))
106    }
107
108    /// Cache plan metadata when priorities depend only on the generated key and immutable set.
109    /// This does not cache a winner or bypass normal selection, validation or expiration.
110    pub fn with_static_stack_plan(mut self) -> Self {
111        #[cfg(std_io)]
112        { self.stack_plan_cache = Some(spin::Mutex::new(Default::default())); }
113        self
114    }
115
116    #[cfg(std_io)]
117    pub(crate) fn static_stack_plan_enabled(&self) -> bool { self.stack_plan_cache.is_some() }
118
119    #[cfg(std_io)]
120    pub(crate) fn cached_stack_plan(&self, key: &K, identity: &str) -> Option<Arc<StaticStackPlan>> {
121        self.stack_plan_cache.as_ref()?.lock().iter()
122            .find(|(stored_key, stored_identity, _)| stored_key == key && stored_identity == identity)
123            .map(|(_, _, plan)| plan.clone())
124    }
125
126    #[cfg(std_io)]
127    pub(crate) fn cache_stack_plan(&self, key: K, identity: String, plan: Arc<StaticStackPlan>) {
128        if let Some(cache) = &self.stack_plan_cache {
129            let mut cache = cache.lock();
130            if cache.iter().any(|(stored_key, stored_identity, _)| stored_key == &key && stored_identity == &identity) { return; }
131            if cache.len() == 32 { cache.pop_front(); }
132            cache.push_back((key, identity, plan));
133        }
134    }
135
136    /// Shorthand for [`new`](Self::new) with a [`CloneInputGenerator`]: benchmarks run
137    /// on clones of the real call inputs.
138    pub fn new_cloning_inputs(key_gen: impl KeyGenerator<K, F>) -> Self {
139        Self::new(key_gen, super::CloneInputGenerator)
140    }
141
142    /// Register a tunable with this tunable set.
143    pub fn with(mut self, tunable: Tunable<K, F, Output>) -> Self {
144        self.tunables.push(tunable);
145        self.checksum = Once::new();
146        self.stack_manifest = Once::new();
147        #[cfg(std_io)]
148        if let Some(cache) = &self.stack_plan_cache { cache.lock().clear(); }
149        self
150    }
151
152    /// All candidate operations in this set, in registration order.
153    pub fn autotunables(&self) -> impl Iterator<Item = &TuneFn<F, Output>> {
154        self.tunables.iter().map(|tunable| &tunable.function)
155    }
156
157    /// Returns the [autotune plan](TunePlan) for the given set.
158    pub(crate) fn plan(&self, key: &K) -> TunePlan {
159        TunePlan::new(key, &self.tunables)
160    }
161
162    /// Returns the operation for the given index, matching the order returned by
163    /// `autotunables`. Tunables are tried in order, so index 0 should be a good default.
164    pub fn fastest(&self, fastest_index: usize) -> &TuneFn<F, Output> {
165        &self.tunables[fastest_index].function
166    }
167
168    /// Compute a checksum that invalidates outdated cached auto-tune results when the
169    /// set of tunable names changes.
170    pub fn compute_checksum(&self) -> String {
171        // Preserve legacy cache identity when the new controller is not enabled.
172        self.checksum.call_once(|| {
173            let mut checksum = String::new();
174            for tune in &self.tunables { checksum += &tune.function.name; }
175            format!("{:x}", md5::compute(checksum))
176        }).clone()
177    }
178
179    /// Separate, length-delimited manifest for the opt-in full-stack cache.
180    pub fn stack_checksum(&self) -> String {
181        String::from(self.stack_checksum_ref())
182    }
183
184    pub(crate) fn stack_checksum_ref(&self) -> &str {
185        self.stack_manifest.call_once(|| {
186            let mut checksum = format!("stack-v1:{}:{};reference={:?};", self.stack_revision.len(), self.stack_revision, self.stack_reference);
187            for tune in &self.tunables {
188                let _ = write!(checksum, "{}:{}", tune.function.name.len(), tune.function.name);
189            }
190            format!("{:x}", md5::compute(checksum))
191        }).as_str()
192    }
193
194    /// Generate a key from a set of inputs
195    pub fn generate_key<'a>(&self, inputs: &F::At<'a>) -> K {
196        self.key_gen.generate(inputs)
197    }
198
199    /// Generate a set of test inputs from a key and reference inputs.
200    pub fn generate_inputs<'a>(&self, key: &K, inputs: &F::At<'a>) -> F::At<'a> {
201        self.input_gen.generate(key, inputs)
202    }
203}
204
205#[cfg(std_io)]
206/// Trait alias with support for persistent caching
207pub trait AutotuneKey:
208    Clone
209    + Debug
210    + PartialEq
211    + Eq
212    + Hash
213    + Display
214    + serde::Serialize
215    + serde::de::DeserializeOwned
216    + Send
217    + Sync
218    + 'static
219{
220}
221#[cfg(not(std_io))]
222/// Trait alias
223pub trait AutotuneKey:
224    Clone + Debug + PartialEq + Eq + Hash + Display + Send + Sync + 'static
225{
226}
227
228impl AutotuneKey for String {}
229
230#[cfg(all(test, std_io))]
231mod static_plan_tests {
232    use super::*;
233    fn set() -> TunableSet<String, String, String> {
234        TunableSet::new_cloning_inputs(|value: &String| value.clone())
235            .with_static_stack_plan()
236    }
237    fn plan() -> Arc<StaticStackPlan> {
238        Arc::new(StaticStackPlan { indices: alloc::vec![0],
239            candidates: alloc::vec![super::super::stack::Candidate::new("reference")], key: "full-key".into() })
240    }
241    #[test]
242    fn cache_requires_exact_generated_key_and_full_identity() {
243        let set = set();
244        let plan = plan();
245        set.cache_stack_plan("shape".into(), "device/offset/stride".into(), plan.clone());
246        assert!(Arc::ptr_eq(&set.cached_stack_plan(&"shape".into(), "device/offset/stride").unwrap(), &plan));
247        assert!(set.cached_stack_plan(&"other-shape".into(), "device/offset/stride").is_none());
248        assert!(set.cached_stack_plan(&"shape".into(), "other-device/offset/stride").is_none());
249    }
250    #[test]
251    fn cache_is_bounded_and_duplicate_insert_does_not_evict() {
252        let set = set();
253        for index in 0..32 { set.cache_stack_plan(alloc::format!("{index}"), "binding".into(), plan()); }
254        set.cache_stack_plan("0".into(), "binding".into(), plan());
255        assert!(set.cached_stack_plan(&"0".into(), "binding").is_some());
256        set.cache_stack_plan("32".into(), "binding".into(), plan());
257        assert!(set.cached_stack_plan(&"0".into(), "binding").is_none());
258        assert!(set.cached_stack_plan(&"1".into(), "binding").is_some());
259    }
260    #[test]
261    fn changing_stack_contract_invalidates_plan_metadata() {
262        let set = set();
263        set.cache_stack_plan("shape".into(), "binding".into(), plan());
264        let set = set.with_stack_tuning(0, "new-revision", |value: &String| value.clone());
265        assert!(set.cached_stack_plan(&"shape".into(), "binding").is_none());
266    }
267}
ruda_runtime/runtime/tune/operation.rs

ruda_runtime/runtime/tune/
operation.rs
