metrale_circuit/memory/
workspace.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-10-02: The workspaces of a plan: every kernel family a plan group runs on
4//! ([`Families::of_group`]) with the workspaces it declares (`[[family.workspace]]`), evaluated at
5//! the plan's rows. One buffer serves every launch of a family, so each (family, workspace) is
6//! one term, attributed to every node whose group runs that family.
7//!
8//! Owner: metrale-circuit (memory).
9//! Invariants:
10//! - An expression reads the circuit's dims, `n` (the plan's rows), `k` (the node's input width)
11//!   and `sm_count` (the device's); any other name is an error, never 0. The term is the largest
12//!   value over the family's nodes.
13
14use std::collections::BTreeMap;
15
16use super::MemoryError;
17use crate::fuser::FusionPlan;
18use crate::ir::{Circuit, NodeIdx};
19use crate::venn::families::Families;
20
21/// 2026-10-02: One workspace of one family in a plan.
22#[derive(Debug, Clone, PartialEq, Eq)]
23pub struct WorkspaceTerm {
24    /// 2026-10-02: Family id.
25    pub family: String,
26    /// 2026-10-02: Workspace name.
27    pub name: String,
28    /// 2026-10-02: Bytes at the plan's rows.
29    pub bytes: u64,
30    /// 2026-10-02: Held inside the legacy buffer arena.
31    pub arena: bool,
32    /// 2026-10-02: The nodes whose groups run the family.
33    pub nodes: Vec<NodeIdx>,
34}
35
36/// 2026-10-02: One term per (family, workspace) over several plans' terms: the largest bytes,
37/// the nodes of the plan that needed them.
38pub(crate) fn widest(plans: Vec<Vec<WorkspaceTerm>>) -> Vec<WorkspaceTerm> {
39    let mut by: BTreeMap<(String, String), WorkspaceTerm> = BTreeMap::new();
40    for t in plans.into_iter().flatten() {
41        let key = (t.family.clone(), t.name.clone());
42        match by.get(&key) {
43            Some(old) if old.bytes >= t.bytes => {}
44            _ => {
45                by.insert(key, t);
46            }
47        }
48    }
49    by.into_values().collect()
50}
51
52/// 2026-10-02: The workspaces of `plan` at `rows` rows on a device of `sm_count` SMs.
53pub(crate) fn workspace_terms(
54    c: &Circuit,
55    plan: &FusionPlan,
56    fams: &Families,
57    rows: u64,
58    sm_count: u64,
59) -> Result<Vec<WorkspaceTerm>, MemoryError> {
60    let mut dims = c.dims.clone();
61    dims.insert("n".into(), rows);
62    dims.insert("sm_count".into(), sm_count);
63    let mut terms: BTreeMap<(String, String), WorkspaceTerm> = BTreeMap::new();
64    for g in &plan.groups {
65        for &n in &g.nodes {
66            let node = &c.nodes[n];
67            let Some(f) = fams.of_group(&g.kernels, &g.emitter, &node.op) else {
68                continue;
69            };
70            dims.insert(
71                "k".into(),
72                node.inputs.first().map_or(0, |&e| c.edges[e].dim_value),
73            );
74            for w in &f.workspace {
75                let mut bytes = 0;
76                for e in &w.bytes {
77                    bytes = bytes.max(e.eval(&dims).map_err(|e| MemoryError::Workspace {
78                        family: f.id.clone(),
79                        name: w.name.clone(),
80                        detail: e.to_string(),
81                    })?);
82                }
83                let t = terms
84                    .entry((f.id.clone(), w.name.clone()))
85                    .or_insert_with(|| WorkspaceTerm {
86                        family: f.id.clone(),
87                        name: w.name.clone(),
88                        bytes: 0,
89                        arena: w.arena,
90                        nodes: Vec::new(),
91                    });
92                t.bytes = t.bytes.max(bytes);
93                t.nodes.push(n);
94            }
95        }
96    }
97    Ok(terms.into_values().collect())
98}
metrale_circuit/memory/workspace.rs

metrale_circuit/memory/
workspace.rs
