rlx_runtime/
flexible_session.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//! Session that defers backend choice until compile time.
6
7use rlx_driver::Device;
8use rlx_ir::{Graph, GraphModule, HirModule};
9use rlx_opt::PrecisionPolicy;
10
11use crate::compiled::CompiledGraph;
12use crate::device_policy::{DevicePolicy, resolve_device};
13use crate::precision::Precision;
14use crate::session::Session;
15
16/// Compile-time settings without a fixed [`Device`].
17///
18/// Pick the backend per graph via [`Self::compile_resolved`] or
19/// [`Self::compile_on`].
20pub struct FlexibleSession {
21    device_policy: DevicePolicy,
22    precision: Precision,
23    op_policy: Option<PrecisionPolicy>,
24}
25
26impl Default for FlexibleSession {
27    fn default() -> Self {
28        Self::new()
29    }
30}
31
32impl FlexibleSession {
33    pub fn new() -> Self {
34        Self {
35            device_policy: DevicePolicy::default(),
36            precision: Precision::F32,
37            op_policy: None,
38        }
39    }
40
41    pub fn from_env() -> Self {
42        Self {
43            device_policy: DevicePolicy::from_env(),
44            ..Self::new()
45        }
46    }
47
48    pub fn with_device_policy(mut self, policy: DevicePolicy) -> Self {
49        self.device_policy = policy;
50        self
51    }
52
53    pub fn with_precision(mut self, precision: Precision) -> Self {
54        self.precision = precision;
55        self
56    }
57
58    pub fn with_op_policy(mut self, policy: PrecisionPolicy) -> Self {
59        self.op_policy = Some(policy);
60        self
61    }
62
63    pub fn device_policy(&self) -> &DevicePolicy {
64        &self.device_policy
65    }
66
67    pub fn precision(&self) -> Precision {
68        self.precision
69    }
70
71    fn session_on(&self, device: Device) -> Session {
72        let mut s = Session::new_with_precision(device, self.precision);
73        if let Some(p) = &self.op_policy {
74            s = s.with_policy(p.clone());
75        }
76        s
77    }
78
79    pub fn compile_on(&self, graph: Graph, device: Device) -> Result<CompiledGraph, String> {
80        Ok(self.session_on(device).compile(graph))
81    }
82
83    pub fn compile_with_on(
84        &self,
85        graph: Graph,
86        device: Device,
87        options: &crate::CompileOptions,
88    ) -> Result<CompiledGraph, String> {
89        Ok(self.session_on(device).compile_with(graph, options))
90    }
91
92    pub fn compile_resolved(
93        &self,
94        graph: Graph,
95        hint: Option<Device>,
96    ) -> Result<CompiledGraph, String> {
97        let device = resolve_device(&graph, hint, &self.device_policy)?;
98        self.compile_on(graph, device)
99    }
100
101    pub fn compile_resolved_with(
102        &self,
103        graph: Graph,
104        hint: Option<Device>,
105        options: &crate::CompileOptions,
106    ) -> Result<CompiledGraph, String> {
107        let device = resolve_device(&graph, hint, &self.device_policy)?;
108        self.compile_with_on(graph, device, options)
109    }
110}
111
112impl FlexibleSession {
113    pub fn compile_hir_on(
114        &self,
115        hir: HirModule,
116        device: Device,
117    ) -> Result<CompiledGraph, rlx_ir::hir::LowerError> {
118        self.session_on(device).compile_hir(hir)
119    }
120
121    pub fn compile_module_on(
122        &self,
123        module: GraphModule,
124        device: Device,
125    ) -> Result<CompiledGraph, rlx_ir::hir::LowerError> {
126        self.session_on(device).compile_module(module)
127    }
128}
129
130#[cfg(test)]
131mod tests {
132    use super::*;
133    use rlx_ir::{DType, Shape};
134
135    #[test]
136    fn compile_resolved_picks_cpu() {
137        let mut g = Graph::new("id");
138        let x = g.input("x", Shape::new(&[2], DType::F32));
139        g.set_outputs(vec![x]);
140        let session = FlexibleSession::new().with_device_policy(DevicePolicy::only([Device::Cpu]));
141        let compiled = session.compile_resolved(g, None).expect("compile");
142        assert_eq!(compiled.device(), Device::Cpu);
143    }
144
145    #[test]
146    fn compile_resolved_with_matches_compile() {
147        let mut g = Graph::new("id");
148        let x = g.input("x", Shape::new(&[2], DType::F32));
149        let y = g.input("y", Shape::new(&[2], DType::F32));
150        let s = g.add_node(
151            rlx_ir::Op::Binary(rlx_ir::op::BinaryOp::Add),
152            vec![x, y],
153            Shape::new(&[2], DType::F32),
154        );
155        g.set_outputs(vec![s]);
156
157        let session = FlexibleSession::new().with_device_policy(DevicePolicy::only([Device::Cpu]));
158        let g1 = g.clone();
159        let g2 = g;
160        let a = session.compile_resolved(g1, None).expect("compile");
161        let b = session
162            .compile_resolved_with(g2, None, &crate::CompileOptions::new())
163            .expect("compile_with");
164        assert_eq!(a.device(), b.device());
165    }
166}
rlx_runtime/flexible_session.rs

rlx_runtime/
flexible_session.rs
