rlx_runtime/
subgraph.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//! Sub-graph execution helper.
6//!
7//! Op::If branches and Op::While body/cond are sub-graphs nested inside
8//! a parent graph. To execute them, the backend recursively compiles
9//! and runs the inner graph with bound inputs.
10//!
11//! Strategy: compile the sub-graph lazily on first encounter, cache the
12//! `ExecutableGraph` for repeated invocations (loops). A future
13//! optimization: hoist the compile to the parent's compile-time once
14//! we have a stable IR for sub-graphs.
15
16use crate::CompileOptions;
17use crate::backend::{Backend, ExecutableGraph};
18use rlx_ir::Graph;
19use std::collections::HashMap;
20
21/// Lazily-compiled sub-graph cache.
22/// Keyed by sub-graph name (caller must ensure names are unique within
23/// the parent graph). Backend-agnostic: stores boxed ExecutableGraphs.
24pub struct SubgraphCache {
25    cache: HashMap<String, Box<dyn ExecutableGraph>>,
26    options: CompileOptions,
27}
28
29impl SubgraphCache {
30    pub fn new(options: CompileOptions) -> Self {
31        Self {
32            cache: HashMap::new(),
33            options,
34        }
35    }
36
37    /// Compile a sub-graph if not cached, return mutable executable handle.
38    pub fn get_or_compile<'a>(
39        &'a mut self,
40        backend: &dyn Backend,
41        graph: &Graph,
42    ) -> &'a mut Box<dyn ExecutableGraph> {
43        let key = graph.name.clone();
44        self.cache
45            .entry(key)
46            .or_insert_with(|| backend.compile(graph.clone(), &self.options))
47    }
48
49    /// Run a sub-graph with named inputs, returning its outputs.
50    pub fn run(
51        &mut self,
52        backend: &dyn Backend,
53        graph: &Graph,
54        inputs: &[(&str, &[f32])],
55    ) -> Vec<Vec<f32>> {
56        let exe = self.get_or_compile(backend, graph);
57        exe.run(inputs)
58    }
59}
60
61/// Helper: evaluate an Op::If by running one of two sub-graphs.
62pub fn run_if(
63    cache: &mut SubgraphCache,
64    backend: &dyn Backend,
65    predicate: f32,
66    then_branch: &Graph,
67    else_branch: &Graph,
68    inputs: &[(&str, &[f32])],
69) -> Vec<Vec<f32>> {
70    let chosen = if predicate != 0.0 {
71        then_branch
72    } else {
73        else_branch
74    };
75    cache.run(backend, chosen, inputs)
76}
77
78/// Helper: evaluate an Op::While by repeatedly running cond + body.
79/// `loop_carried` are the values flowing through iterations.
80pub fn run_while(
81    cache: &mut SubgraphCache,
82    backend: &dyn Backend,
83    cond: &Graph,
84    body: &Graph,
85    initial: Vec<Vec<f32>>,
86    input_names: &[&str],
87    max_iterations: Option<usize>,
88) -> Vec<Vec<f32>> {
89    let mut state = initial;
90    let limit = max_iterations.unwrap_or(usize::MAX);
91    for _ in 0..limit {
92        // Build named-input slice for cond + body
93        let bindings: Vec<(&str, &[f32])> = input_names
94            .iter()
95            .zip(state.iter())
96            .map(|(n, v)| (*n, v.as_slice()))
97            .collect();
98        let cond_out = cache.run(backend, cond, &bindings);
99        // Cond is a scalar bool: stop if it's zero / false
100        if cond_out
101            .first()
102            .map(|v| v.first().copied().unwrap_or(0.0))
103            .unwrap_or(0.0)
104            == 0.0
105        {
106            break;
107        }
108        state = cache.run(backend, body, &bindings);
109    }
110    state
111}
rlx_runtime/subgraph.rs

rlx_runtime/
subgraph.rs
