rlx_opscope/
dataflow.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//! Dataflow repeated-pattern mining. Where `motifs` walks *linear* single-
6//! consumer op chains, this canonicalizes each node's **input cone** (the
7//! branching sub-DAG that produces it, up to a depth) into a structural
8//! signature and counts recurrences. A sub-DAG that recurs across the graph is
9//! a unit of repeated computation — a candidate to be **decomposed once and
10//! shared**, or **fused** into a specialized op. This is the "track the data
11//! flow that can be additionally decomposed, especially repeated patterns" seam.
12//!
13//! The signature is a merkle hash of the cone: `sig(n, d)` = the op name of `n`
14//! joined with `sig` of each operand at depth `d-1`; at `d = 0` or a leaf the
15//! cone is truncated to just the op name. Two nodes share a signature iff their
16//! depth-`d` input cones are op-isomorphic. Cheap (memoized), exact for DAGs.
17//!
18//! The miner runs on a plain node list `(op_name, input_indices)` so it works on
19//! both an in-memory [`Graph`] and a graph dumped from another workspace (see
20//! [`repeated_flow_patterns_on`] and the `opscope-graph` bin).
21
22use rlx_ir::{Graph, Op};
23use std::collections::{HashMap, HashSet};
24
25/// A recurring rooted dataflow pattern.
26#[derive(Clone, Debug)]
27pub struct FlowPattern {
28    /// Cone depth this pattern was found at.
29    pub depth: usize,
30    /// Number of distinct roots whose depth-`depth` cone matches.
31    pub count: usize,
32    /// Human-readable cone, e.g. `Relu(Add(MatMul(·,w),w))`.
33    pub tree: String,
34    /// Root node indices (a few examples).
35    pub sites: Vec<u32>,
36}
37
38impl FlowPattern {
39    /// Decomposition payoff proxy: recurrence × cone size.
40    pub fn score(&self) -> usize {
41        self.count * (self.tree.matches('(').count() + 1)
42    }
43}
44
45/// A readable op label that keeps the discriminating detail (activation kind,
46/// binary op) instead of collapsing everything to `Activation`/`Binary`.
47pub fn op_name(op: &Op) -> String {
48    match op {
49        Op::Activation(a) => format!("{a:?}"),
50        Op::Binary(b) => format!("{b:?}"),
51        Op::Compare(c) => format!("Cmp{c:?}"),
52        Op::Input { .. } => "in".into(),
53        Op::Param { .. } => "w".into(),
54        Op::Constant { .. } => "k".into(),
55        other => format!("{:?}", other.kind()),
56    }
57}
58
59fn sig(
60    ops: &[String],
61    inputs: &[Vec<usize>],
62    node: usize,
63    depth: usize,
64    memo: &mut HashMap<(usize, usize), String>,
65) -> String {
66    if let Some(s) = memo.get(&(node, depth)) {
67        return s.clone();
68    }
69    let name = ops[node].clone();
70    let s = if depth == 0 || inputs[node].is_empty() {
71        name
72    } else {
73        let kids: Vec<String> = inputs[node]
74            .iter()
75            .map(|&i| sig(ops, inputs, i, depth - 1, memo))
76            .collect();
77        format!("{name}({})", kids.join(","))
78    };
79    memo.insert((node, depth), s.clone());
80    s
81}
82
83/// Mine repeated input-cone patterns over a plain node list. `ops[i]` is node
84/// `i`'s op label, `inputs[i]` its operand indices. Keeps cones recurring at
85/// least `min_count` times, ranked by payoff; bare leaves are dropped.
86pub fn repeated_flow_patterns_on(
87    ops: &[String],
88    inputs: &[Vec<usize>],
89    min_depth: usize,
90    max_depth: usize,
91    min_count: usize,
92) -> Vec<FlowPattern> {
93    let mut memo: HashMap<(usize, usize), String> = HashMap::new();
94    let mut out: Vec<FlowPattern> = Vec::new();
95    let mut seen_tree: HashSet<String> = HashSet::new();
96
97    for depth in min_depth..=max_depth {
98        let mut groups: HashMap<String, Vec<u32>> = HashMap::new();
99        for node in 0..ops.len() {
100            if inputs[node].is_empty() {
101                continue; // root cones only at compute nodes
102            }
103            let s = sig(ops, inputs, node, depth, &mut memo);
104            if !s.contains('(') {
105                continue; // must actually branch/recurse
106            }
107            groups.entry(s).or_default().push(node as u32);
108        }
109        for (tree, sites) in groups {
110            if sites.len() >= min_count && seen_tree.insert(tree.clone()) {
111                out.push(FlowPattern {
112                    depth,
113                    count: sites.len(),
114                    sites: sites.into_iter().take(6).collect(),
115                    tree,
116                });
117            }
118        }
119    }
120
121    out.sort_by(|a, b| b.score().cmp(&a.score()).then(b.count.cmp(&a.count)));
122    out
123}
124
125/// Convenience over an in-memory [`Graph`]. Node ids are dense insertion order,
126/// so `NodeId.0` doubles as the list index.
127pub fn repeated_flow_patterns(
128    g: &Graph,
129    min_depth: usize,
130    max_depth: usize,
131    min_count: usize,
132) -> Vec<FlowPattern> {
133    let ops: Vec<String> = g.nodes().iter().map(|n| op_name(&n.op)).collect();
134    let inputs: Vec<Vec<usize>> = g
135        .nodes()
136        .iter()
137        .map(|n| n.inputs.iter().map(|i| i.0 as usize).collect())
138        .collect();
139    repeated_flow_patterns_on(&ops, &inputs, min_depth, max_depth, min_count)
140}
141
142/// A short "how to exploit this repeat" hint for a mined cone.
143pub fn decomposition_hint(p: &FlowPattern) -> &'static str {
144    let t = &p.tree;
145    let has_act = t.contains("Relu") || t.contains("Silu") || t.contains("Gelu");
146    let has_attn = t.contains("Attention") || t.contains("Softmax") || t.contains("Rope");
147    let has_moe = t.contains("GroupedMatMul") || t.contains("TopK");
148    if has_moe {
149        "→ MoE expert block: fuse grouped-matmul + gating"
150    } else if has_attn {
151        "→ attention block: fuse (FusedAttentionBlock)"
152    } else if t.contains("MatMul(") && t.contains("Add(") && has_act {
153        if t.starts_with("Add(") {
154            "→ residual linear+bias+act: fuse (FusedMatMulBiasAct + residual)"
155        } else {
156            "→ linear+bias+act: fuse into one kernel (FusedMatMulBiasAct)"
157        }
158    } else if t.contains("MatMul(") {
159        "→ recurring matmul cone: share/prepack"
160    } else {
161        "→ recurring sub-DAG: decompose once + reuse"
162    }
163}
rlx_opscope/dataflow.rs

rlx_opscope/
dataflow.rs
