1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
// RLX — versatile ML compiler + runtime.
// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
// SPDX-License-Identifier: MIT OR Apache-2.0
//! Demonstrates sub-graph compilation + execution via SubgraphCache.
//!
//! Builds two simple graphs and runs them through the same backend
//! at different times. The key point: the backend doesn't need to
//! know about Op::If / Op::While ahead of time — the runtime helper
//! recursively compiles the sub-graphs as needed.
//!
//! Once Op::If / Op::While are wired into the thunk executor, the
//! pattern below will be hidden inside the backend; this example
//! just shows the mechanism is in place.
//!
//! cargo run --example control_flow --release \
//! --features "cpu,blas-accelerate"
#[cfg(feature = "cpu")]
fn main() {
use rlx_ir::*;
use rlx_runtime::backend::Backend;
use rlx_runtime::{CompileOptions, Device, Session, SubgraphCache, run_if};
// Constants encoded as [4]-shaped vectors — the runtime's
// BinaryFull thunk doesn't yet broadcast scalars, so we use
// full-shape constants (broadcast support is a separate task).
let const4 = |v: f32| -> Vec<u8> {
let mut out = Vec::with_capacity(16);
for _ in 0..4 {
out.extend_from_slice(&v.to_le_bytes());
}
out
};
// ── "then" sub-graph: y = x + 10 (broadcast via [4]) ───
let then_branch = {
let mut g = Graph::new("then");
let x = g.input("x", Shape::new(&[4], DType::F32));
let c = g.add_node(
Op::Constant { data: const4(10.0) },
vec![],
Shape::new(&[4], DType::F32),
);
let y = g.binary(op::BinaryOp::Add, x, c, Shape::new(&[4], DType::F32));
g.set_outputs(vec![y]);
g
};
// ── "else" sub-graph: y = x * 2 ────────────────────────
let else_branch = {
let mut g = Graph::new("else");
let x = g.input("x", Shape::new(&[4], DType::F32));
let c = g.add_node(
Op::Constant { data: const4(2.0) },
vec![],
Shape::new(&[4], DType::F32),
);
let y = g.binary(op::BinaryOp::Mul, x, c, Shape::new(&[4], DType::F32));
g.set_outputs(vec![y]);
g
};
let _session = Session::new(Device::Cpu);
let backend: Box<dyn Backend> = Box::new(rlx_runtime::backend::cpu_backend::CpuBackend);
let mut cache = SubgraphCache::new(CompileOptions::new());
let x_data = vec![1.0, 2.0, 3.0, 4.0];
let then_result = run_if(
&mut cache,
backend.as_ref(),
1.0,
&then_branch,
&else_branch,
&[("x", &x_data)],
);
let else_result = run_if(
&mut cache,
backend.as_ref(),
0.0,
&then_branch,
&else_branch,
&[("x", &x_data)],
);
println!("then (predicate=1): x + 10 = {:?}", then_result[0]);
println!("else (predicate=0): x * 2 = {:?}", else_result[0]);
assert_eq!(then_result[0], vec![11.0, 12.0, 13.0, 14.0]);
assert_eq!(else_result[0], vec![2.0, 4.0, 6.0, 8.0]);
println!("✓ sub-graph execution works through SubgraphCache");
// Run again — sub-graphs are cached, no recompile.
let r2 = run_if(
&mut cache,
backend.as_ref(),
1.0,
&then_branch,
&else_branch,
&[("x", &x_data)],
);
assert_eq!(r2, then_result);
println!("✓ cached sub-graph reuse works (second invocation)");
}
#[cfg(not(feature = "cpu"))]
fn main() {
eprintln!("requires --features cpu");
}