Skip to main content

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}