rlx_runtime/
graph_devices.rs

1// RLX — versatile ML compiler + runtime.
2// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
3// SPDX-License-Identifier: MIT OR Apache-2.0
4
5//! Multi-backend execution — compile once per device, run on any of them.
6
7use std::collections::HashMap;
8
9use rlx_driver::Device;
10use rlx_ir::{DType, Graph, Op};
11
12use crate::compiled::CompiledGraph;
13use crate::cost::fastest_device_for_with_policy;
14use crate::device_bench::{DeviceBenchResult, benchmark_devices, warm_all};
15use crate::device_ext::is_available;
16use crate::device_policy::{
17    DeviceCandidate, DeviceFallbackError, DevicePickStrategy, DevicePolicy, device_chain_from_env,
18    device_report, devices_for_with_policy, resolve_device, resolve_device_chain,
19};
20use crate::session::Session;
21
22/// Param names declared in `graph` (`Op::Param`).
23pub fn graph_param_names(graph: &Graph) -> Vec<String> {
24    graph
25        .nodes()
26        .iter()
27        .filter_map(|n| match &n.op {
28            Op::Param { name } => Some(name.clone()),
29            _ => None,
30        })
31        .collect()
32}
33
34#[derive(Debug, Clone)]
35enum CachedParam {
36    F32(Vec<f32>),
37    Typed { bytes: Vec<u8>, dtype: DType },
38}
39
40fn apply_cached_params(compiled: &mut CompiledGraph, params: &HashMap<String, CachedParam>) {
41    for (name, param) in params {
42        match param {
43            CachedParam::F32(data) => compiled.set_param(name, data),
44            CachedParam::Typed { bytes, dtype } => compiled.set_param_typed(name, bytes, *dtype),
45        }
46    }
47}
48
49/// A graph plus lazy per-device compiled executables.
50pub struct GraphDevices {
51    graph: Graph,
52    policy: DevicePolicy,
53    pick: DevicePickStrategy,
54    supported: Vec<Device>,
55    params: HashMap<String, CachedParam>,
56    benchmark_winner: Option<Device>,
57    cache: HashMap<Device, CompiledGraph>,
58}
59
60impl GraphDevices {
61    pub fn new(graph: Graph) -> Self {
62        Self::with_policy(graph, DevicePolicy::default())
63    }
64
65    pub fn with_policy(graph: Graph, policy: DevicePolicy) -> Self {
66        let pick = policy.pick_strategy();
67        let supported = devices_for_with_policy(&graph, &policy);
68        Self {
69            graph,
70            policy,
71            pick,
72            supported,
73            params: HashMap::new(),
74            benchmark_winner: None,
75            cache: HashMap::new(),
76        }
77    }
78
79    pub fn from_env(graph: Graph) -> Self {
80        Self::with_policy(graph, DevicePolicy::from_env())
81    }
82
83    pub fn policy(&self) -> &DevicePolicy {
84        &self.policy
85    }
86
87    pub fn graph(&self) -> &Graph {
88        &self.graph
89    }
90
91    pub fn devices(&self) -> &[Device] {
92        &self.supported
93    }
94
95    pub fn report(&self) -> Vec<DeviceCandidate> {
96        device_report(&self.graph, &self.policy)
97    }
98
99    pub fn fastest(&self) -> Device {
100        fastest_device_for_with_policy(&self.graph, &self.policy)
101    }
102
103    pub fn resolve(&self, hint: Option<Device>) -> Result<Device, String> {
104        resolve_device(&self.graph, hint, &self.policy)
105    }
106
107    /// Resolve using `RLX_DEVICE_CHAIN` when set, else [`Self::resolve`].
108    pub fn resolve_chain(&self, hint: Option<Device>) -> Result<Device, String> {
109        if let Some(device) = hint {
110            return self.resolve(Some(device));
111        }
112        let chain = device_chain_from_env();
113        if chain.is_empty() {
114            return self.resolve(None);
115        }
116        resolve_device_chain(&self.graph, &chain, &self.policy)
117    }
118
119    /// Upload a param to every cached executor and future compilations.
120    pub fn set_param(&mut self, name: &str, data: &[f32]) {
121        self.params
122            .insert(name.to_string(), CachedParam::F32(data.to_vec()));
123        for compiled in self.cache.values_mut() {
124            compiled.set_param(name, data);
125        }
126    }
127
128    /// Typed param upload — mirrored to all cached backends.
129    pub fn set_param_typed(&mut self, name: &str, data: &[u8], dtype: DType) {
130        self.params.insert(
131            name.to_string(),
132            CachedParam::Typed {
133                bytes: data.to_vec(),
134                dtype,
135            },
136        );
137        for compiled in self.cache.values_mut() {
138            compiled.set_param_typed(name, data, dtype);
139        }
140    }
141
142    /// Re-apply stored params to all cached backends (after manual cache changes).
143    pub fn sync_params_to_all(&mut self) {
144        for compiled in self.cache.values_mut() {
145            apply_cached_params(compiled, &self.params);
146        }
147    }
148
149    /// Hint → env → cost model, or micro-benchmark when policy requests it.
150    pub fn resolve_with_inputs(
151        &mut self,
152        hint: Option<Device>,
153        inputs: &[(&str, &[f32])],
154    ) -> Result<Device, String> {
155        if hint.is_some() {
156            return self.resolve(hint);
157        }
158        match self.pick {
159            DevicePickStrategy::CostModel => self.resolve(None),
160            DevicePickStrategy::Benchmark { runs } => {
161                if let Some(device) = self.benchmark_winner {
162                    return Ok(device);
163                }
164                let ranked = self.benchmark(inputs, runs)?;
165                let device = ranked
166                    .first()
167                    .map(|r| r.device)
168                    .unwrap_or_else(|| self.fastest());
169                self.benchmark_winner = Some(device);
170                Ok(device)
171            }
172        }
173    }
174
175    pub fn compile(&mut self, device: Device) -> Result<&mut CompiledGraph, String> {
176        Self::ensure_supported(&self.supported, device)?;
177        if !self.cache.contains_key(&device) {
178            let mut compiled = Session::new(device).compile(self.graph.clone());
179            apply_cached_params(&mut compiled, &self.params);
180            self.cache.insert(device, compiled);
181        }
182        Ok(self.cache.get_mut(&device).expect("just inserted"))
183    }
184
185    pub fn compile_fastest(&mut self) -> Result<&mut CompiledGraph, String> {
186        self.compile(self.fastest())
187    }
188
189    pub fn compile_resolved(&mut self, hint: Option<Device>) -> Result<&mut CompiledGraph, String> {
190        self.compile(self.resolve(hint)?)
191    }
192
193    pub fn compile_chain(&mut self, hint: Option<Device>) -> Result<&mut CompiledGraph, String> {
194        self.compile(self.resolve_chain(hint)?)
195    }
196
197    pub fn warm_all(&mut self) -> Result<Vec<Device>, String> {
198        warm_all(self)
199    }
200
201    pub fn benchmark(
202        &mut self,
203        inputs: &[(&str, &[f32])],
204        runs: usize,
205    ) -> Result<Vec<DeviceBenchResult>, String> {
206        benchmark_devices(self, inputs, runs)
207    }
208
209    pub fn run(
210        &mut self,
211        device: Device,
212        inputs: &[(&str, &[f32])],
213    ) -> Result<Vec<Vec<f32>>, String> {
214        Ok(self.compile(device)?.run(inputs))
215    }
216
217    pub fn run_resolved(
218        &mut self,
219        hint: Option<Device>,
220        inputs: &[(&str, &[f32])],
221    ) -> Result<Vec<Vec<f32>>, String> {
222        Ok(self.compile_resolved(hint)?.run(inputs))
223    }
224
225    pub fn run_fastest(&mut self, inputs: &[(&str, &[f32])]) -> Result<Vec<Vec<f32>>, String> {
226        Ok(self.compile_fastest()?.run(inputs))
227    }
228
229    /// Try `chain` in order until one backend compiles and runs successfully.
230    pub fn run_try(
231        &mut self,
232        chain: &[Device],
233        inputs: &[(&str, &[f32])],
234    ) -> Result<(Device, Vec<Vec<f32>>), DeviceFallbackError> {
235        let viable: Vec<Device> = self.devices().to_vec();
236        let mut attempts = Vec::new();
237        for &device in chain {
238            if !viable.contains(&device) {
239                attempts.push((device, "not viable for this graph under policy".into()));
240                continue;
241            }
242            match self.run(device, inputs) {
243                Ok(value) => return Ok((device, value)),
244                Err(err) => attempts.push((device, err)),
245            }
246        }
247        if attempts.is_empty() {
248            attempts.push((Device::Cpu, "empty fallback chain".into()));
249        }
250        Err(DeviceFallbackError { attempts })
251    }
252
253    /// Like [`Self::run_try`] using `RLX_DEVICE_CHAIN` when set.
254    pub fn run_chain(
255        &mut self,
256        hint: Option<Device>,
257        inputs: &[(&str, &[f32])],
258    ) -> Result<(Device, Vec<Vec<f32>>), DeviceFallbackError> {
259        if let Some(device) = hint {
260            self.run(device, inputs)
261                .map(|v| (device, v))
262                .map_err(|e| DeviceFallbackError {
263                    attempts: vec![(device, e)],
264                })
265        } else {
266            let chain = device_chain_from_env();
267            if chain.is_empty() {
268                let device = self.resolve(None).map_err(|e| DeviceFallbackError {
269                    attempts: vec![(Device::Cpu, e)],
270                })?;
271                self.run(device, inputs)
272                    .map(|v| (device, v))
273                    .map_err(|e| DeviceFallbackError {
274                        attempts: vec![(device, e)],
275                    })
276            } else {
277                self.run_try(&chain, inputs)
278            }
279        }
280    }
281
282    pub fn compile_resolved_with_inputs(
283        &mut self,
284        hint: Option<Device>,
285        inputs: &[(&str, &[f32])],
286    ) -> Result<&mut CompiledGraph, String> {
287        let device = self.resolve_with_inputs(hint, inputs)?;
288        self.compile(device)
289    }
290
291    pub fn run_resolved_with_inputs(
292        &mut self,
293        hint: Option<Device>,
294        inputs: &[(&str, &[f32])],
295    ) -> Result<Vec<Vec<f32>>, String> {
296        Ok(self.compile_resolved_with_inputs(hint, inputs)?.run(inputs))
297    }
298
299    pub fn invalidate_cache(&mut self) {
300        self.cache.clear();
301        self.benchmark_winner = None;
302        self.supported = devices_for_with_policy(&self.graph, &self.policy);
303    }
304
305    pub fn set_policy(&mut self, policy: DevicePolicy) {
306        self.policy = policy.clone();
307        self.pick = policy.pick_strategy();
308        self.invalidate_cache();
309    }
310
311    fn ensure_supported(supported: &[Device], device: Device) -> Result<(), String> {
312        if !is_available(device) {
313            return Err(format!(
314                "device {device} is not available — enable the matching Cargo feature"
315            ));
316        }
317        if !supported.contains(&device) {
318            return Err(format!(
319                "device {device} cannot lower this graph under the active policy"
320            ));
321        }
322        Ok(())
323    }
324}
325
326#[cfg(test)]
327mod tests {
328    use super::*;
329    use rlx_ir::{DType, Shape};
330
331    fn identity_graph() -> Graph {
332        let mut g = Graph::new("id");
333        let x = g.input("x", Shape::new(&[4], DType::F32));
334        g.set_outputs(vec![x]);
335        g
336    }
337
338    #[test]
339    fn set_param_applies_to_new_compile() {
340        let mut g = Graph::new("p");
341        let x = g.input("x", Shape::new(&[2], DType::F32));
342        let w = g.param("w", Shape::new(&[2], DType::F32));
343        let y = g.binary(
344            rlx_ir::op::BinaryOp::Add,
345            x,
346            w,
347            Shape::new(&[2], DType::F32),
348        );
349        g.set_outputs(vec![y]);
350
351        let mut runner = GraphDevices::new(g);
352        runner.set_param("w", &[1.0, 2.0]);
353        let out = runner.run(Device::Cpu, &[("x", &[3.0, 4.0])]).unwrap();
354        assert_eq!(out[0], vec![4.0, 6.0]);
355    }
356
357    #[test]
358    fn run_on_cpu_roundtrip() {
359        let mut runner = GraphDevices::new(identity_graph());
360        let out = runner
361            .run(Device::Cpu, &[("x", &[1.0, 2.0, 3.0, 4.0])])
362            .expect("cpu run");
363        assert_eq!(out[0], vec![1.0, 2.0, 3.0, 4.0]);
364    }
365
366    #[test]
367    fn run_try_falls_back_to_cpu() {
368        let mut runner = GraphDevices::new(identity_graph());
369        // Lead with a device THIS BUILD genuinely cannot run, so the fallback
370        // path is exercised on every host. Hard-coding `Device::Cuda` only
371        // passed where CUDA was absent — on a CUDA host `run_try` correctly
372        // selects Cuda and an unconditional `== Cpu` assertion fails.
373        let chain: Vec<Device> = match [Device::Cuda, Device::Rocm, Device::Vulkan, Device::Metal]
374            .into_iter()
375            .find(|d| !crate::is_available(*d))
376        {
377            Some(unavailable) => vec![unavailable, Device::Cpu],
378            None => vec![Device::Cpu],
379        };
380        let (dev, out) = runner
381            .run_try(&chain, &[("x", &[1.0, 2.0, 3.0, 4.0])])
382            .expect("fallback");
383        assert_eq!(dev, Device::Cpu);
384        assert_eq!(out[0][0], 1.0);
385    }
386}
rlx_runtime/graph_devices.rs

rlx_runtime/
graph_devices.rs
