rlx_flow/side.rs
1// RLX — versatile ML compiler + runtime.
2// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
3//
4// This program is free software: you can redistribute it and/or modify
5// it under the terms of the GNU General Public License as published by
6// the Free Software Foundation, version 3.
7//
8// This program is distributed in the hope that it will be useful,
9// but WITHOUT ANY WARRANTY; without even the implied warranty of
10// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
11// GNU General Public License for more details.
12//
13// You should have received a copy of the GNU General Public License
14// along with this program. If not, see <https://www.gnu.org/licenses/>.
15
16//! Side-output collection (KV cache taps, auxiliary heads, …).
17
18use std::sync::{Arc, Mutex};
19
20use rlx_ir::HirNodeId;
21
22/// Collects extra graph outputs emitted by side-effect stages (e.g. KV taps).
23#[derive(Debug, Clone, Default)]
24pub struct SideOutputs {
25 inner: Arc<Mutex<Vec<HirNodeId>>>,
26}
27
28impl SideOutputs {
29 pub fn new() -> Self {
30 Self::default()
31 }
32
33 /// Shared handle for side-effect stages (KV taps, …).
34 pub fn inner(&self) -> Arc<Mutex<Vec<HirNodeId>>> {
35 Arc::clone(&self.inner)
36 }
37
38 pub fn drain(&self) -> Vec<HirNodeId> {
39 self.inner.lock().expect("side outputs").clone()
40 }
41
42 pub fn is_empty(&self) -> bool {
43 self.inner.lock().expect("side outputs").is_empty()
44 }
45}