rlx_tensor/
cache.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//! Thread-local compile cache.
6//!
7//! Compiling a graph (fuse → memory-plan → backend lower) is far more
8//! expensive than running it, so every materialization (`Tensor::to_vec`,
9//! `Tensor::grad` + eval, `Func::run`) routes through here. A graph that was
10//! already compiled — same structure, same constant bytes, same device — gets
11//! its [`CompiledGraph`] back instead of recompiling. For [`crate::Func`] this
12//! turns repeated `run` calls into a real `jit`: compile once, execute many
13//! inputs.
14//!
15//! The cache is per-thread. Available with the `eval` feature.
16
17use std::cell::{Cell, RefCell};
18use std::collections::HashMap;
19use std::collections::hash_map::DefaultHasher;
20use std::hash::{Hash, Hasher};
21use std::rc::Rc;
22
23use rlx_ir::{Graph, NodeId, Op};
24use rlx_runtime::{CompiledGraph, Device, PrecisionPolicy, Session};
25
26/// Clear the cache once it grows past this many distinct graphs, a simple
27/// backstop against unbounded growth in long-lived threads.
28const MAX_ENTRIES: usize = 512;
29
30thread_local! {
31    static CACHE: RefCell<HashMap<u64, Rc<RefCell<CompiledGraph>>>> =
32        RefCell::new(HashMap::new());
33    static HITS: Cell<u64> = const { Cell::new(0) };
34    static MISSES: Cell<u64> = const { Cell::new(0) };
35}
36
37/// Structural fingerprint covering everything that affects the compiled
38/// result: op fields (via `Debug`), constant bytes, operands, shapes, output
39/// set, and target device.
40fn fingerprint(graph: &Graph, outputs: &[NodeId], device: Device) -> u64 {
41    let mut h = DefaultHasher::new();
42    device.hash(&mut h);
43    graph.name.hash(&mut h);
44    for node in graph.nodes() {
45        node.id.0.hash(&mut h);
46        match &node.op {
47            // Hash bytes directly — avoids a huge Debug string for weights.
48            Op::Constant { data } => {
49                0u8.hash(&mut h);
50                data.hash(&mut h);
51            }
52            other => format!("{other:?}").hash(&mut h),
53        }
54        for inp in &node.inputs {
55            inp.0.hash(&mut h);
56        }
57        format!("{:?}", node.shape).hash(&mut h);
58    }
59    for out in outputs {
60        out.0.hash(&mut h);
61    }
62    h.finish()
63}
64
65fn lookup(key: u64) -> Option<Rc<RefCell<CompiledGraph>>> {
66    let hit = CACHE.with(|c| c.borrow().get(&key).cloned());
67    if hit.is_some() {
68        HITS.with(|n| n.set(n.get() + 1));
69    } else {
70        MISSES.with(|n| n.set(n.get() + 1));
71    }
72    hit
73}
74
75/// Opt-in mixed-precision policy from `RLX_PRECISION` (default: F32, unchanged).
76/// `auto` = the Metal-safe policy (elementwise/reduction in F16, **matmul stays
77/// F32** — F16 matmul diverges on the residual/RoPE/gather paths); `f16`/`bf16`
78/// push matmul into low precision too (faster, but numerically riskier — bf16
79/// keeps F32's exponent range so it's the safer of the two).
80fn precision_policy_from_env() -> Option<PrecisionPolicy> {
81    match rlx_ir::env::var("RLX_PRECISION").as_deref() {
82        Some("auto") => Some(PrecisionPolicy::AutoMixed),
83        Some("f16") => Some(PrecisionPolicy::AutoMixedConservative),
84        Some("bf16") => Some(PrecisionPolicy::AutoMixedBf16),
85        Some("f16-all") => Some(PrecisionPolicy::AlwaysF16),
86        _ => None,
87    }
88}
89
90fn store(key: u64, device: Device, graph: Graph) -> Rc<RefCell<CompiledGraph>> {
91    let mut session = Session::new(device);
92    if let Some(policy) = precision_policy_from_env() {
93        session = session.with_policy(policy);
94    }
95    let compiled = Rc::new(RefCell::new(session.compile(graph)));
96    CACHE.with(|c| {
97        let mut map = c.borrow_mut();
98        if map.len() >= MAX_ENTRIES {
99            map.clear();
100        }
101        map.insert(key, compiled.clone());
102    });
103    compiled
104}
105
106/// Compiled graph for `(graph, device)` where `graph` already carries its
107/// outputs (the `Func` path). Cloned only on a cache miss.
108pub(crate) fn compiled(graph: &Graph, device: Device) -> Rc<RefCell<CompiledGraph>> {
109    let key = fingerprint(graph, &graph.outputs, device);
110    lookup(key).unwrap_or_else(|| store(key, device, graph.clone()))
111}
112
113/// Compiled graph that emits a single `output` node from a shared graph —
114/// the `Tensor::to_vec` path. **The graph is borrowed (not cloned) for the
115/// fingerprint**, so cache hits copy zero constant bytes; the clone +
116/// `set_outputs` happens only on a miss.
117pub(crate) fn compiled_output(
118    graph: &Graph,
119    output: NodeId,
120    device: Device,
121) -> Rc<RefCell<CompiledGraph>> {
122    let key = fingerprint(graph, &[output], device);
123    lookup(key).unwrap_or_else(|| {
124        let mut g = graph.clone();
125        g.set_outputs(vec![output]);
126        store(key, device, g)
127    })
128}
129
130/// `(hits, misses)` for the calling thread's cache. Useful for tests and
131/// confirming `jit`-style reuse.
132pub fn cache_stats() -> (u64, u64) {
133    (HITS.with(Cell::get), MISSES.with(Cell::get))
134}
135
136/// Drop all cached compiled graphs and reset stats on the calling thread.
137pub fn clear_cache() {
138    CACHE.with(|c| c.borrow_mut().clear());
139    HITS.with(|n| n.set(0));
140    MISSES.with(|n| n.set(0));
141}
rlx_tensor/cache.rs

rlx_tensor/
cache.rs
