rlx_runtime/
graph_io.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//! Static IO / sync profile for compiled graphs (Phase 0 — fusion planning).
6
7use rlx_ir::{Graph, Node, Op, OpKind};
8use serde::{Deserialize, Serialize};
9use std::collections::HashSet;
10
11/// Tuning for static IO analysis.
12#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
13pub struct GraphIoOptions {
14    /// Count each `Op::Fft` as a host-sync boundary (non-native fallback).
15    pub fft_host_sync: bool,
16}
17
18/// Host-visible traffic and sync points for one forward pass.
19#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
20pub struct GraphIoProfile {
21    /// Kernel / thunk dispatches (one per non-view executable node).
22    pub kernel_launches: usize,
23    /// GPU flush + host-side thunk boundaries (Metal LogMel, host FFT fallback, …).
24    pub sync_points: usize,
25    /// Bytes returned to the caller via graph outputs (`CompiledGraph::run`).
26    pub host_output_bytes: u64,
27    /// Bytes moved inside the device arena (read inputs + write outputs per node).
28    pub device_traffic_bytes: u64,
29}
30
31impl GraphIoProfile {
32    pub fn host_readback_bytes(&self, unified_memory: bool) -> u64 {
33        if unified_memory {
34            self.host_output_bytes
35        } else {
36            self.host_output_bytes
37                .saturating_add(self.device_traffic_bytes / 4)
38        }
39    }
40}
41
42/// Ops that force a GPU sync + host thunk on Metal today.
43pub fn metal_host_sync_kinds() -> &'static [OpKind] {
44    &[
45        OpKind::LogMel,
46        OpKind::LogMelBackward,
47        OpKind::Custom,
48        OpKind::WelchPeaks,
49    ]
50}
51
52/// Profile a graph before compile (conservative static estimate).
53pub fn profile_graph_io(graph: &Graph) -> GraphIoProfile {
54    profile_graph_io_with_options(graph, GraphIoOptions::default())
55}
56
57pub fn profile_graph_io_with_options(graph: &Graph, opts: GraphIoOptions) -> GraphIoProfile {
58    let mut profile = GraphIoProfile::default();
59    let output_nodes: HashSet<_> = graph.outputs.iter().copied().collect();
60
61    for node in graph.nodes() {
62        if is_metadata_op(&node.op) {
63            continue;
64        }
65        profile.kernel_launches += 1;
66        profile.device_traffic_bytes += node_io_bytes(node, graph);
67
68        let kind = node.op.kind();
69        if metal_host_sync_kinds().contains(&kind) {
70            profile.sync_points += 1;
71        }
72        if opts.fft_host_sync && kind == OpKind::Fft {
73            profile.sync_points += 1;
74        }
75
76        if output_nodes.contains(&node.id) {
77            profile.host_output_bytes += tensor_bytes(&node.shape);
78        }
79    }
80
81    profile
82}
83
84/// Profile with only selected outputs materialized on the host (peaks-only, logits-only, …).
85pub fn profile_graph_io_outputs(graph: &Graph, output_indices: &[usize]) -> GraphIoProfile {
86    let mut profile = profile_graph_io(graph);
87    profile.host_output_bytes = graph
88        .outputs
89        .iter()
90        .enumerate()
91        .filter(|(i, _)| output_indices.contains(i))
92        .filter_map(|(_, id)| graph.node(*id).shape.num_elements())
93        .map(|n| (n * 4) as u64)
94        .sum();
95    profile
96}
97
98fn is_metadata_op(op: &Op) -> bool {
99    matches!(
100        op,
101        Op::Input { .. }
102            | Op::Param { .. }
103            | Op::Constant { .. }
104            | Op::Reshape { .. }
105            | Op::Transpose { .. }
106            | Op::Narrow { .. }
107    )
108}
109
110fn node_io_bytes(node: &Node, graph: &Graph) -> u64 {
111    let out = tensor_bytes(&node.shape);
112    let inputs: u64 = node
113        .inputs
114        .iter()
115        .map(|&id| tensor_bytes(&graph.node(id).shape))
116        .sum();
117    inputs.saturating_add(out)
118}
119
120fn tensor_bytes(shape: &rlx_ir::Shape) -> u64 {
121    shape
122        .num_elements()
123        .map(|n| (n * shape.dtype().size_bytes()) as u64)
124        .unwrap_or(0)
125}
126
127#[cfg(test)]
128mod tests {
129    use super::*;
130    use rlx_ir::infer::GraphExt;
131    use rlx_ir::{DType, Shape};
132
133    #[test]
134    fn fft_graph_io_profile() {
135        let mut g = Graph::new("fft");
136        let x = g.input("x", Shape::new(&[8, 512], DType::F32));
137        let zeros = g.sub(x, x);
138        let block = g.concat_(vec![x, zeros], 1);
139        let y = g.fft(block, false);
140        g.set_outputs(vec![y]);
141        let p = profile_graph_io(&g);
142        assert!(p.kernel_launches >= 3);
143        assert_eq!(p.host_output_bytes, (8 * 512 * 2 * 4) as u64);
144    }
145
146    #[test]
147    fn peaks_only_output_smaller_readback() {
148        let mut g = Graph::new("peaks");
149        let spec = g.input("spec", Shape::new(&[4, 512], DType::F32));
150        let peaks = g.welch_peaks(spec, 16, 2);
151        g.set_outputs(vec![peaks]);
152        let full = profile_graph_io(&g);
153        assert_eq!(full.host_output_bytes, (2 * 16 * 2 * 4) as u64);
154        assert!(full.device_traffic_bytes > full.host_output_bytes);
155    }
156}
rlx_runtime/graph_io.rs

rlx_runtime/
graph_io.rs
