ruda_runtime/runtime/tune/stack/
runtime_adapter.rs

1use super::*;
2use super::policy::{CandidateView, ProblemView};
3use crate::runtime::{backend::Runtime, client::ComputeClient, tune::{AutotuneKey, AutotuneOutput, TunableSet, TuneInputs, AutotuneError}};
4use std::{any::type_name, path::PathBuf, string::{String, ToString}, sync::{Arc, OnceLock}, time::{Duration, Instant, SystemTime, UNIX_EPOCH}, vec::Vec};
5
6static GLOBAL: OnceLock<StackTuner> = OnceLock::new();
7static SESSION: OnceLock<String> = OnceLock::new();
8/// Install once, before worker threads/model loading. No implicit environment mutation or I/O.
9pub fn enable_stack_autotune(policy: StackPolicy, cache_directory: Option<PathBuf>) -> Result<&'static StackTuner, TuneFailure> {
10    enable_stack_autotune_with_memory_eviction(policy, cache_directory, MemoryEviction::Fifo)
11}
12/// Install a controller with explicit memory eviction; disk behavior and trials stay unchanged.
13pub fn enable_stack_autotune_with_memory_eviction(policy: StackPolicy, cache_directory: Option<PathBuf>, memory_eviction: MemoryEviction) -> Result<&'static StackTuner, TuneFailure> {
14    let tuner = StackTuner::new_with_memory_eviction(policy, cache_directory, memory_eviction)?;
15    GLOBAL.set(tuner).map_err(|_| TuneFailure::invalid("stack autotuning was already configured"))?;
16    Ok(GLOBAL.get().expect("controller was initialized"))
17}
18/// Return the installed controller, or None to preserve the legacy tuner route.
19/// Merely querying it does not create a controller, inspect files or run trials.
20pub fn stack_autotuner() -> Option<&'static StackTuner> { GLOBAL.get() }
21
22/// Supply extra tags for a deployment's compiler flags and load/topology/power regime.
23/// Set these BEFORE constructing devices/starting inference; changing them live is unsupported.
24#[derive(Debug, Clone)]
25pub struct RuntimeEnvironment {
26    /// Canonical backend/device/driver/build/runtime-options signature.
27    pub fingerprint: String,
28    /// Whether a driver identity and compiled build identity permit disk reuse.
29    pub persistent: bool,
30    /// Deployment-supplied regime, defaulting to isolated-single-device.
31    /// The default does not establish actual isolation or topology.
32    pub execution_context: String
33}
34/// Probe and memoize immutable environment identity for this runtime/client.
35/// Deployment tags must be set before first use; live mutation is not supported.
36/// Missing identity yields session-only caching rather than a fabricated driver version.
37pub fn runtime_environment<R: Runtime>(client: &ComputeClient<R>) -> RuntimeEnvironment {
38    runtime_environment_shared(client).as_ref().clone()
39}
40fn runtime_environment_shared<R: Runtime>(client: &ComputeClient<R>) -> Arc<RuntimeEnvironment> {
41    use std::{any::TypeId, collections::BTreeMap, sync::RwLock};
42    type Key = (TypeId, u16, u16, u64);
43    static ENVIRONMENTS: OnceLock<RwLock<BTreeMap<Key, Arc<RuntimeEnvironment>>>> = OnceLock::new();
44    let id = client.device_id();
45    let key = (TypeId::of::<R>(), id.type_id, id.index_id, client.properties_fingerprint());
46    let cache = ENVIRONMENTS.get_or_init(|| RwLock::new(BTreeMap::new()));
47    if let Some(env) = cache.read().unwrap_or_else(|p| p.into_inner()).get(&key).cloned() { return env; }
48    let env = Arc::new(probe_environment(client));
49    let mut cache = cache.write().unwrap_or_else(|p| p.into_inner());
50    // Device/runtime metadata is immutable after initialization. Keep this memo bounded.
51    if cache.len() < 256 { cache.insert(key, env.clone()); }
52    env
53}
54fn probe_environment<R: Runtime>(client: &ComputeClient<R>) -> RuntimeEnvironment {
55    let driver_override = std::env::var("RUDA_AUTOTUNE_DRIVER_TAG").ok().filter(|s| !s.is_empty());
56    let driver = match (R::autotune_driver_fingerprint(client), driver_override) {
57        (Some(probed), Some(extra)) => Some(fields(&[&probed, &extra])),
58        (Some(probed), None) => Some(probed),
59        (None, Some(asserted)) => Some(asserted),
60        (None, None) => None,
61    };
62    let persistent = driver.is_some() && env!("RUDA_STACK_BUILD_ID") != "unavailable";
63    let driver = driver.unwrap_or_else(|| SESSION.get_or_init(|| std::format!("session-only:{}:{}", std::process::id(), SystemTime::now().duration_since(UNIX_EPOCH).unwrap_or_default().as_nanos())).clone());
64    let build_extra = std::env::var("RUDA_AUTOTUNE_BUILD_TAG").unwrap_or_default();
65    let execution_context = std::env::var("RUDA_AUTOTUNE_CONTEXT_TAG").unwrap_or_else(|_| "isolated-single-device".into());
66    use crate::runtime::config::{RudaRuntimeConfig, RuntimeConfig};
67    let config = RudaRuntimeConfig::get();
68    let runtime_options = std::format!("compilation={:?};streaming={:?};memory={:?}", config.compilation, config.streaming, config.memory);
69    let fingerprint = fields(&[R::name(client), &std::format!("{:?}", client.device_id()), &std::format!("{:016x}", client.properties_fingerprint()),
70        &driver, type_name::<R::Compiler>(), env!("RUDA_STACK_BUILD_ID"), &build_extra,
71        &std::format!("{:?}", client.info()), &runtime_options]);
72    RuntimeEnvironment { fingerprint, persistent, execution_context }
73}
74
75fn complete<R: Runtime>(client: &ComputeClient<R>) -> Result<(), TuneFailure> {
76    let submitted = client.flush();
77    // Completion errors are never demoted to a failed candidate followed by another launch.
78    ruda_core::future::block_on(client.sync()).map_err(|e| TuneFailure::device(std::format!("autotune synchronization failed: {e:?}")))?;
79    submitted.map_err(|e| TuneFailure::rejected(std::format!("candidate submission failed after a successful completion fence: {e:?}")))
80}
81fn run_output<R: Runtime, K: AutotuneKey, I: TuneInputs, O: AutotuneOutput>(
82    client: &ComputeClient<R>, set: &TunableSet<K,I,O>, index: usize, input: I::At<'_>,
83) -> Result<O, TuneFailure> {
84    let result = set.fastest(index).execute(input);
85    complete(client)?;
86    result.map_err(|e| TuneFailure::rejected(std::format!("candidate rejected: {e:?}")))
87}
88struct RuntimeTrials<'set, 'inp, R: Runtime, K: AutotuneKey, I: TuneInputs, O: AutotuneOutput> {
89    client: ComputeClient<R>, set: &'set TunableSet<K,I,O>, key: &'set K,
90    input: &'set I::At<'inp>, indices: &'set [usize], timing: Timing,
91}
92impl<R: Runtime, K: AutotuneKey, I: TuneInputs, O: AutotuneOutput> TrialRunner for RuntimeTrials<'_, '_, R,K,I,O> {
93    fn validate(&mut self, reference: usize, candidate: usize, tolerance: Tolerance) -> Result<Validation, TuneFailure> {
94        let a = self.set.generate_inputs(self.key, self.input);
95        let b = self.set.generate_inputs(self.key, self.input);
96        complete(&self.client)?;
97        let client = self.client.clone(); let set = self.set;
98        let ri = self.indices[reference]; let ci = self.indices[candidate];
99        self.client.exclusive(move || {
100            let _depth = super::engine::DepthGuard::enter();
101            let expected = run_output(&client, set, ri, a)?;
102            let actual = run_output(&client, set, ci, b)?;
103            let result = expected.validate_for_tuning(&actual, tolerance.absolute, tolerance.relative, tolerance.max_bytes);
104            // Readback/contiguous conversion can also launch kernels. Fence before accepting a
105            // validation mismatch or freeing trial state.
106            complete(&client)?;
107            match result {
108                Ok(true) => Ok(Validation::Passed), Ok(false) => Ok(Validation::Unsupported),
109                Err(e) => Err(TuneFailure::rejected(e)),
110            }
111        }).map_err(|e| TuneFailure::device(std::format!("exclusive validation failed: {e:?}")))?
112    }
113    fn measure(&mut self, candidate: usize) -> Result<Duration, TuneFailure> {
114        let input = self.set.generate_inputs(self.key, self.input);
115        complete(&self.client)?;
116        let client = self.client.clone(); let set = self.set;
117        let index = self.indices[candidate]; let timing = self.timing;
118        self.client.exclusive(move || {
119            let _depth = super::engine::DepthGuard::enter();
120            // Preparation/sandbox copies are outside the measured region. All allocations and
121            // relayouts made INSIDE the actual candidate operation are part of EndToEnd timing.
122            complete(&client)?;
123            match timing {
124                Timing::EndToEnd => {
125                    let start = Instant::now();
126                    let out = run_output(&client, set, index, input)?;
127                    let elapsed = start.elapsed();
128                    std::hint::black_box(&out); drop(out);
129                    Ok(elapsed)
130                }
131                Timing::Device => {
132                    let profiled = client.profile(move || set.fastest(index).execute(input), "stack-autotune");
133                    complete(&client)?;
134                    let (output, profile) = profiled.map_err(|e| TuneFailure::rejected(std::format!("device profile failed: {e:?}")))?;
135                    let output = output.map_err(|e| TuneFailure::rejected(std::format!("candidate rejected: {e:?}")))?;
136                    let elapsed = ruda_core::future::block_on(profile.resolve()).duration();
137                    std::hint::black_box(&output); drop(output);
138                    Ok(elapsed)
139                }
140            }
141        }).map_err(|e| TuneFailure::device(std::format!("exclusive measurement failed: {e:?}")))?
142    }
143}
144
145/// Fallible full-stack path. Actual request execution is performed ONCE after selection; an
146/// execution error invalidates future choices but is never silently replayed on another kernel.
147/// Requires an installed controller, an in-range explicit reference and an exact workload
148/// signature on the set. Missing metadata returns AutotuneError before live execution.
149pub fn try_execute_stack<'a, R: Runtime, K: AutotuneKey, I: TuneInputs, O: AutotuneOutput>(
150    name: &str, device_id: &str, client: &ComputeClient<R>, set: Arc<TunableSet<K,I,O>>, input: I::At<'a>,
151) -> Result<O, AutotuneError> {
152    let tuner = stack_autotuner().ok_or_else(|| autotune_error(name, "stack autotuning is not enabled"))?;
153    let reference = set.stack_reference().ok_or_else(|| autotune_error(name, "no explicit reference/workload signature was registered"))?;
154    if reference >= set.len() { return Err(autotune_error(name, "reference index out of range")); }
155    let key = set.generate_key(&input);
156    let env = runtime_environment_shared(client);
157    let scope = if name.contains("fusion") { Scope::Graph } else { Scope::Operator };
158    let operation = fields(&[name, device_id, set.stack_checksum_ref()]);
159    let workload = set.stack_workload(&input).ok_or_else(|| autotune_error(name, "no exact workload signature"))?;
160    let problem = ProblemView { scope, operation: &operation, environment: &env.fingerprint,
161        workload: &workload, execution_context: &env.execution_context, persistent: env.persistent };
162    if set.static_stack_plan_enabled() {
163        let identity = fields(&[&std::format!("{:?}", problem.scope), problem.operation,
164            problem.environment, problem.workload, problem.execution_context,
165            if problem.persistent { "persistent" } else { "session" }]);
166        let plan = if let Some(plan) = set.cached_stack_plan(&key, &identity) { plan } else {
167            let mut indices = std::vec![reference];
168            let mut sequence = set.plan(&key);
169            let mut included = std::vec![false; set.len()];
170            included[reference] = true;
171            loop {
172                let batch = sequence.next(None);
173                if batch.is_empty() { break; }
174                for index in batch { if !included[index] { indices.push(index); included[index] = true; } }
175            }
176            let candidates: Vec<_> = indices.iter().map(|&index| Candidate::new(set.fastest(index).name.clone())).collect();
177            let prepared_key = tuner.prepare_key(problem, &candidates, 0).map_err(|error| autotune_error(name, &error.to_string()))?;
178            let plan = Arc::new(crate::runtime::tune::operation::StaticStackPlan { indices, candidates, key: prepared_key });
179            set.cache_stack_plan(key.clone(), identity, plan.clone());
180            plan
181        };
182        let mut trials = RuntimeTrials { client: client.clone(), set: &set, key: &key,
183            input: &input, indices: &plan.indices, timing: tuner.policy().timing };
184        let decision = tuner.select_prepared_candidates(problem, &plan.candidates, 0, &plan.key, &mut trials)
185            .map_err(|error| autotune_error(name, &error.to_string()))?;
186        let output = set.fastest(plan.indices[decision.index]).execute(input);
187        if output.is_err() { tuner.invalidate(&decision, true); }
188        return output;
189    }
190    // Enumerate EVERY eligible priority group. Group order remains a useful bounded-search
191    // heuristic but a first valid group no longer terminates the search.
192    let mut indices = std::vec![reference]; let mut plan = set.plan(&key);
193    let mut included = std::vec![false; set.len()];
194    included[reference] = true;
195    loop { let batch = plan.next(None); if batch.is_empty() { break; }
196        for index in batch { if !included[index] { indices.push(index); included[index] = true; } }
197    }
198    let candidates: Vec<_> = indices.iter().map(|&i| CandidateView::new(&set.fastest(i).name)).collect();
199    // Cache hits must retain normal asynchronous execution. TrialRunner fences before every
200    // validation/measurement; do NOT synchronize ordinary requests just to read this cache.
201    let mut trials = RuntimeTrials { client: client.clone(), set: &set, key: &key, input: &input, indices: &indices, timing: tuner.policy().timing };
202    let decision = tuner.select_candidates(problem, &candidates, 0, &mut trials).map_err(|e| autotune_error(name, &e.to_string()))?;
203    let output = set.fastest(indices[decision.index]).execute(input);
204    if output.is_err() { tuner.invalidate(&decision, true); }
205    output
206}
207fn autotune_error(name: &str, message: &str) -> AutotuneError {
208    AutotuneError::Unknown { name: name.to_string(), err: message.to_string() }
209}
210
211pub struct PreparedStackOperation<K: AutotuneKey, I: TuneInputs, O: AutotuneOutput> {
212    set: Arc<TunableSet<K,I,O>>, key: K, problem: Problem,
213    indices: Vec<usize>, candidates: Vec<Candidate>, cache_key: String,
214    runtime: std::any::TypeId, device: (u16, u16, u64),
215}
216
217pub fn prepare_stack_operation<R: Runtime, K: AutotuneKey, I: TuneInputs, O: AutotuneOutput>(
218    name: &str, device_id: &str, client: &ComputeClient<R>, set: Arc<TunableSet<K,I,O>>, input: &I::At<'_>,
219) -> Result<PreparedStackOperation<K,I,O>, AutotuneError> {
220    let tuner = stack_autotuner().ok_or_else(|| autotune_error(name, "stack autotuning is not enabled"))?;
221    let reference = set.stack_reference().ok_or_else(|| autotune_error(name, "no explicit reference/workload signature was registered"))?;
222    if reference >= set.len() { return Err(autotune_error(name, "reference index out of range")); }
223    let key = set.generate_key(input);
224    let env = runtime_environment(client);
225    let problem = Problem { scope: if name.contains("fusion") { Scope::Graph } else { Scope::Operator },
226        operation: fields(&[name, device_id, set.stack_checksum_ref()]), environment: env.fingerprint,
227        workload: set.stack_workload(input).ok_or_else(|| autotune_error(name, "no exact workload signature"))?,
228        execution_context: env.execution_context, persistent: env.persistent };
229    let mut indices = std::vec![reference]; let mut plan = set.plan(&key);
230    let mut included = std::vec![false; set.len()]; included[reference] = true;
231    loop { let batch = plan.next(None); if batch.is_empty() { break; }
232        for index in batch { if !included[index] { indices.push(index); included[index] = true; } }
233    }
234    let candidates: Vec<_> = indices.iter().map(|&i| Candidate::new(set.fastest(i).name.clone())).collect();
235    let cache_key = tuner.prepare_key((&problem).into(), &candidates, 0).map_err(|e| autotune_error(name, &e.to_string()))?;
236    let id = client.device_id();
237    Ok(PreparedStackOperation { set, key, problem, indices, candidates, cache_key,
238        runtime: std::any::TypeId::of::<R>(), device: (id.type_id, id.index_id, client.properties_fingerprint()) })
239}
240
241pub fn try_execute_prepared_stack<'a, R: Runtime, K: AutotuneKey, I: TuneInputs, O: AutotuneOutput>(
242    name: &str, client: &ComputeClient<R>, prepared: &PreparedStackOperation<K,I,O>, input: I::At<'a>,
243) -> Result<O, AutotuneError> {
244    let tuner = stack_autotuner().ok_or_else(|| autotune_error(name, "stack autotuning is not enabled"))?;
245    let id = client.device_id();
246    if prepared.runtime != std::any::TypeId::of::<R>()
247        || prepared.device != (id.type_id, id.index_id, client.properties_fingerprint()) {
248        return Err(autotune_error(name, "prepared operation runtime/device mismatch"));
249    }
250    let mut trials = RuntimeTrials { client: client.clone(), set: &prepared.set, key: &prepared.key,
251        input: &input, indices: &prepared.indices, timing: tuner.policy().timing };
252    let decision = tuner.select_prepared_candidates((&prepared.problem).into(), &prepared.candidates, 0,
253        &prepared.cache_key, &mut trials).map_err(|e| autotune_error(name, &e.to_string()))?;
254    let output = prepared.set.fastest(prepared.indices[decision.index]).execute(input);
255    if output.is_err() { tuner.invalidate(&decision, true); }
256    output
257}
ruda_runtime/runtime/tune/stack/runtime_adapter.rs

ruda_runtime/runtime/tune/stack/
runtime_adapter.rs
