rlx_flow/
side.rs

1// RLX — versatile ML compiler + runtime.
2// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
3
4//! Side-output collection (KV cache taps, auxiliary heads, …).
5
6use std::sync::{Arc, Mutex};
7
8use rlx_ir::HirNodeId;
9
10/// Collects extra graph outputs emitted by side-effect stages (e.g. KV taps).
11#[derive(Debug, Clone, Default)]
12pub struct SideOutputs {
13    inner: Arc<Mutex<Vec<HirNodeId>>>,
14}
15
16impl SideOutputs {
17    pub fn new() -> Self {
18        Self::default()
19    }
20
21    /// Shared handle for side-effect stages (KV taps, …).
22    pub fn inner(&self) -> Arc<Mutex<Vec<HirNodeId>>> {
23        Arc::clone(&self.inner)
24    }
25
26    pub fn drain(&self) -> Vec<HirNodeId> {
27        self.inner.lock().expect("side outputs").clone()
28    }
29
30    pub fn is_empty(&self) -> bool {
31        self.inner.lock().expect("side outputs").is_empty()
32    }
33}
rlx_flow/side.rs

rlx_flow/
side.rs
