metrale_model_layers/circuit_exec/
compile.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-28: Plan to program. [`layout`] states the storage the emitters' kernels need (the
4//! residual stream and the logits bound to the model's buffers, in-place outputs, packed
5//! outputs); the caller places the rest in one workspace allocated at boot; [`compile`] asks
6//! each group's emitter for its launches.
7//!
8//! Owner: model-layers circuit executor.
9//! Invariants:
10//! - Every stream edge (a block's `stream_in` or `stream_out`) is the model's `hidden` buffer,
11//!   updated in place, as the legacy layers update it; every declared output is the model
12//!   buffer its producer writes (`logits`, `tokens`).
13//! - The program has exactly `plan.launches()` launches; any other count is a compile error.
14//! - A group's weights are checked against the format its nodes were resolved to.
15
16use anyhow::{Context, Result, bail, ensure};
17use metrale_circuit::planner::{BufferPlan, Layout};
18use metrale_circuit::{Circuit, FusionPlan, Group, Mode};
19use metrale_config::ModelConfig;
20use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
21
22use super::bindings::{BoundWeight, CircuitLayer, HeadBinding, WeightSlot};
23use super::emitters::emitter;
24pub use super::fixed::{DraftFixed, Fixed};
25use super::kernels::KernelTable;
26use super::program::{Launch, LaunchKind, Program, RunFn};
27use crate::layer::AttnMetadataDev;
28
29/// 2026-09-28: A group of the plan being compiled.
30pub(crate) struct GroupRef<'a> {
31    pub circuit: &'a Circuit,
32    pub group: &'a Group,
33    pub index: usize,
34}
35
36impl<'a> GroupRef<'a> {
37    /// 2026-09-28: Member `i`, in pattern order.
38    pub fn node(&self, i: usize) -> &'a metrale_circuit::ir::Node {
39        &self.circuit.nodes[self.group.nodes[i]]
40    }
41
42    /// 2026-09-28: Input `j` of member `i`.
43    pub fn input(&self, i: usize, j: usize) -> Result<usize> {
44        let n = self.node(i);
45        n.inputs
46            .get(j)
47            .copied()
48            .with_context(|| format!("`{}` has no input {j}", n.id))
49    }
50
51    /// 2026-09-28: Output `j` of member `i`.
52    pub fn output(&self, i: usize, j: usize) -> Result<usize> {
53        let n = self.node(i);
54        n.outputs
55            .get(j)
56            .copied()
57            .with_context(|| format!("`{}` has no output {j}", n.id))
58    }
59
60    /// 2026-09-28: Refuse a group whose members are not exactly `ops`, in order.
61    pub fn expect_ops(&self, emitter: &str, ops: &[&str]) -> Result<()> {
62        let got: Vec<String> = self
63            .group
64            .nodes
65            .iter()
66            .map(|&n| self.circuit.nodes[n].op.name())
67            .collect();
68        ensure!(
69            got.len() == ops.len() && got.iter().zip(ops).all(|(g, w)| g.starts_with(w)),
70            "emitter `{emitter}` cannot launch group {} ({:?}); it takes {ops:?}",
71            self.index,
72            got
73        );
74        Ok(())
75    }
76}
77
78/// 2026-09-28: What an emitter reads while compiling one group.
79pub(crate) struct Cx<'a> {
80    pub g: GroupRef<'a>,
81    pub gpu: &'a dyn GpuBackend,
82    pub mode: Mode,
83    pub rows: u64,
84    pub config: &'a ModelConfig,
85    pub fixed: &'a Fixed,
86    pub layers: &'a [CircuitLayer],
87    pub head: &'a HeadBinding,
88    pub draft: Option<&'a CircuitLayer>,
89    table: Option<&'a metrale_circuit::RowTable>,
90    ptrs: &'a [Option<DevicePtr>],
91    strides: &'a [Option<u64>],
92    formats: &'a [metrale_circuit::Format],
93    handles: Vec<KernelHandle>,
94    launches: Vec<Launch>,
95}
96
97impl<'a> Cx<'a> {
98    fn buffer(&self, edge: usize) -> Result<DevicePtr> {
99        self.ptrs[edge].with_context(|| {
100            format!(
101                "edge `{}` has no buffer in this plan (fused away or outside the section)",
102                self.g.circuit.edges[edge].id
103            )
104        })
105    }
106
107    fn row_bytes(&self, edge: usize) -> Result<u64> {
108        let e = &self.g.circuit.edges[edge];
109        self.formats[edge]
110            .bytes(1, e.dim_value)
111            .with_context(|| format!("edge `{}` has no row size", e.id))
112    }
113
114    /// 2026-09-28: Bytes from one row of an edge to the next.
115    fn stride_bytes(&self, edge: usize) -> Result<u64> {
116        match self.strides[edge] {
117            Some(s) => Ok(s),
118            None => self.row_bytes(edge),
119        }
120    }
121
122    /// 2026-09-28: The buffer of a materialised edge whose rows are contiguous; an edge laid out
123    /// in a row pack is refused, since a kernel reading it here would step rows by its width.
124    pub fn ptr(&self, edge: usize) -> Result<DevicePtr> {
125        let p = self.buffer(edge)?;
126        ensure!(
127            self.rows <= 1 || self.stride_bytes(edge)? == self.row_bytes(edge)?,
128            "edge `{}` is laid out in a row pack; it must be read with its row stride",
129            self.g.circuit.edges[edge].id
130        );
131        Ok(p)
132    }
133
134    /// 2026-09-28: The buffer of a materialised edge and its row stride in elements, for a
135    /// kernel that takes the stride.
136    pub fn strided(&self, edge: usize) -> Result<(DevicePtr, u32)> {
137        let e = &self.g.circuit.edges[edge];
138        let (stride, row) = (self.stride_bytes(edge)?, self.row_bytes(edge)?);
139        ensure!(
140            e.dim_value > 0 && row % e.dim_value == 0 && stride % (row / e.dim_value) == 0,
141            "edge `{}`: a {stride}-byte row stride is not whole elements",
142            e.id
143        );
144        Ok((
145            self.buffer(edge)?,
146            u32::try_from(stride / (row / e.dim_value))?,
147        ))
148    }
149
150    /// 2026-09-28: Row `row` of a materialised edge: its buffer plus `row` row strides.
151    pub fn row_ptr(&self, edge: usize, row: usize) -> Result<DevicePtr> {
152        Ok(self
153            .buffer(edge)?
154            .offset(row * self.stride_bytes(edge)? as usize))
155    }
156
157    /// 2026-09-28: The attention metadata this plan's steps upload.
158    pub fn meta(&self) -> Result<AttnMetadataDev> {
159        Ok(match self.mode {
160            Mode::Draft => self.draft_fixed()?.meta_rows(self.rows)?,
161            Mode::Decode => self.fixed.meta,
162            Mode::MultiSeq => self.fixed.batch_meta,
163            Mode::Verify => self.fixed.verify_meta,
164            Mode::VerifyBatch => self.fixed.verify_batch_meta,
165        })
166    }
167
168    /// 2026-09-30: The plan's row table (a batched verify's); `None` in other modes.
169    pub fn table(&self) -> Option<&'a metrale_circuit::RowTable> {
170        self.table
171    }
172
173    /// 2026-09-30: The handle of `k`, one of the group's kernels (a per-run launch).
174    pub fn run_handle(&self, k: &metrale_circuit::KernelId) -> Result<KernelHandle> {
175        let i = self
176            .g
177            .group
178            .kernels
179            .iter()
180            .position(|x| x == k)
181            .with_context(|| format!("group {} does not list {k}", self.g.index))?;
182        self.handle(i)
183    }
184
185    /// 2026-09-30: Queue one launch of `k`, one of the group's kernels.
186    pub fn push_kernel(&mut self, k: &metrale_circuit::KernelId, run: RunFn) -> Result<()> {
187        let i = self
188            .g
189            .group
190            .kernels
191            .iter()
192            .position(|x| x == k)
193            .with_context(|| format!("group {} does not list {k}", self.g.index))?;
194        self.push(i, run)
195    }
196
197    /// 2026-09-29: The draft head's buffers; an error without a bound draft head.
198    pub fn draft_fixed(&self) -> Result<&'a DraftFixed> {
199        self.fixed
200            .draft
201            .as_ref()
202            .context("a draft plan without the draft head's buffers")
203    }
204
205    /// 2026-09-29: The GDN state row `row` reads: its own sequence's in a multi-sequence
206    /// step, the one sequence's in a verify or a decode.
207    pub fn state_row(&self, row: usize) -> usize {
208        if self.mode == Mode::MultiSeq { row } else { 0 }
209    }
210
211    /// 2026-09-28: Launches of each kernel this group makes: its repeat at the plan's rows.
212    pub fn reps(&self) -> usize {
213        self.g.group.repeat.count(self.rows).unwrap_or(0) as usize
214    }
215
216    /// 2026-09-28: The group's `i`-th kernel.
217    pub fn handle(&self, i: usize) -> Result<KernelHandle> {
218        self.handles
219            .get(i)
220            .copied()
221            .with_context(|| format!("group {} has no kernel {i}", self.g.index))
222    }
223
224    /// 2026-09-28: The binding of the layer member `i` belongs to.
225    pub fn layer(&self, i: usize) -> Result<&'a CircuitLayer> {
226        let n = self.g.node(i);
227        if self.mode == Mode::Draft {
228            return self
229                .draft
230                .with_context(|| format!("`{}`: no draft head is bound", n.id));
231        }
232        let l = n
233            .layer
234            .with_context(|| format!("`{}` is outside the layers", n.id))?;
235        self.layers
236            .get(l)
237            .with_context(|| format!("no binding for layer {l}"))
238    }
239
240    /// 2026-09-28: Member `i`'s weight in `slot`, checked against the node's resolved format when
241    /// the node reads a linear weight.
242    pub fn weight(&self, i: usize, slot: WeightSlot) -> Result<BoundWeight> {
243        let n = self.g.node(i);
244        let w = *self
245            .layer(i)?
246            .weights
247            .get(&slot)
248            .with_context(|| format!("`{}`: its layer binds no {slot:?}", n.id))?;
249        if let Some(f) = n.weight {
250            ensure!(
251                f.name().starts_with(w.family()),
252                "`{}`: the plan resolved {} but the layer holds {}",
253                n.id,
254                f.name(),
255                w.family()
256            );
257        }
258        Ok(w)
259    }
260
261    /// 2026-09-28: Queue one launch of kernel `k` of this group.
262    pub fn push(&mut self, k: usize, run: RunFn) -> Result<()> {
263        let kid = self
264            .g
265            .group
266            .kernels
267            .get(k)
268            .with_context(|| format!("group {} has no kernel {k}", self.g.index))?;
269        self.launches.push(Launch {
270            group: self.g.index,
271            kernel: format!("{}::{}", kid.module, kid.func),
272            kind: LaunchKind::Kernel,
273            run,
274        });
275        Ok(())
276    }
277
278    /// 2026-09-29: Queue one copy-engine transfer of this group (its rule's `copies`).
279    pub fn push_copy(&mut self, run: RunFn) -> Result<()> {
280        ensure!(
281            self.g.group.copies.is_some() || self.g.group.runs.iter().any(|r| r.copies.is_some()),
282            "group {} declares no copies",
283            self.g.index
284        );
285        self.launches.push(Launch {
286            group: self.g.index,
287            kernel: "copy".to_string(),
288            kind: LaunchKind::Copy,
289            run,
290        });
291        Ok(())
292    }
293}
294
295/// 2026-09-28: How an emitter id turns a group into launches.
296pub(crate) trait OpEmitter: Sync {
297    /// 2026-09-28: The emitter id FUSIONS.toml names.
298    fn id(&self) -> &'static str;
299
300    /// 2026-09-28: Storage its kernels require of the group's edges (in-place, packed).
301    fn constrain(&self, _g: &GroupRef<'_>, _layout: &mut Layout) -> Result<()> {
302        Ok(())
303    }
304
305    /// 2026-09-28: Queue the group's launches.
306    fn emit(&self, cx: &mut Cx<'_>) -> Result<()>;
307}
308
309/// 2026-09-28: The storage constraints of `plan`.
310pub fn layout(circuit: &Circuit, plan: &FusionPlan) -> Result<Layout> {
311    let mut layout = Layout::default();
312    let materialized =
313        |e: usize| plan.edge_states[e] == Some(metrale_circuit::EdgeState::Materialized);
314    for b in &circuit.blocks {
315        for e in b.stream_in.into_iter().chain(b.stream_out) {
316            if materialized(e) {
317                layout.external.insert(e);
318            }
319        }
320    }
321    // 2026-09-30: The declared outputs (`Edge::is_output`); `external_buffer` places each by
322    // the model buffer it binds to and refuses one that binds none.
323    for (e, edge) in circuit.edges.iter().enumerate() {
324        if edge.is_output && materialized(e) {
325            layout.external.insert(e);
326        }
327    }
328    for (index, group) in plan.groups.iter().enumerate() {
329        let g = GroupRef {
330            circuit,
331            group,
332            index,
333        };
334        emitter(&group.emitter)?.constrain(&g, &mut layout)?;
335    }
336    Ok(layout)
337}
338
339/// 2026-09-29: The model buffer an external edge is: a stream edge is `hidden`; otherwise
340/// (2026-09-30) the buffer its declared output binds (`Edge::binds`, LIFECYCLE-DESIGN.md 3.4).
341/// An output that binds none is refused.
342pub(super) fn external_buffer(
343    circuit: &Circuit,
344    e: usize,
345    fixed: &Fixed,
346    mode: Mode,
347) -> Result<DevicePtr> {
348    let edge = &circuit.edges[e];
349    let stream = circuit
350        .blocks
351        .iter()
352        .any(|b| b.stream_in == Some(e) || b.stream_out == Some(e));
353    if stream {
354        return Ok(fixed.hidden);
355    }
356    use metrale_circuit::model_buffer::ModelBuffer;
357    match edge.binds {
358        Some(ModelBuffer::Logits) => Ok(fixed.logits),
359        Some(ModelBuffer::Tokens) if mode == Mode::VerifyBatch => Ok(fixed.verify_batch_tokens),
360        Some(ModelBuffer::Tokens) => Ok(fixed.tokens),
361        Some(ModelBuffer::DraftEmbed) => fixed
362            .draft
363            .as_ref()
364            .map(|d| d.embed)
365            .with_context(|| format!("`{}`: no draft embedding buffer", edge.id)),
366        None => bail!(
367            "`{}` is read outside the program but binds no model buffer: an unbound output",
368            edge.id
369        ),
370    }
371}
372
373/// 2026-09-28: Everything [`compile`] reads besides the plan.
374pub struct Inputs<'a> {
375    pub gpu: &'a dyn GpuBackend,
376    pub config: &'a ModelConfig,
377    pub kernels: &'a KernelTable,
378    pub fixed: &'a Fixed,
379    pub layers: &'a [CircuitLayer],
380    pub head: &'a HeadBinding,
381    pub draft: Option<&'a CircuitLayer>,
382}
383
384/// 2026-09-28: Compile `plan` with its buffers placed by `buffers` at `workspace`.
385pub fn compile(
386    circuit: &Circuit,
387    plan: &FusionPlan,
388    layout: &Layout,
389    buffers: &BufferPlan,
390    workspace: DevicePtr,
391    inp: &Inputs<'_>,
392) -> Result<Program> {
393    let mut ptrs: Vec<Option<DevicePtr>> = vec![None; circuit.edges.len()];
394    let mut strides: Vec<Option<u64>> = vec![None; circuit.edges.len()];
395    for s in &buffers.slots {
396        ptrs[s.edge] = Some(workspace.offset(s.offset as usize));
397        strides[s.edge] = Some(s.row_stride);
398    }
399    for &e in &layout.external {
400        ptrs[e] = Some(external_buffer(circuit, e, inp.fixed, plan.mode)?);
401    }
402    let mut launches = Vec::with_capacity(plan.launches() as usize);
403    for (index, group) in plan.groups.iter().enumerate() {
404        let g = GroupRef {
405            circuit,
406            group,
407            index,
408        };
409        let handles = group
410            .kernels
411            .iter()
412            .map(|k| inp.kernels.handle(k))
413            .collect::<Result<Vec<_>>>()?;
414        let mut cx = Cx {
415            g,
416            gpu: inp.gpu,
417            mode: plan.mode,
418            rows: plan.rows,
419            config: inp.config,
420            fixed: inp.fixed,
421            layers: inp.layers,
422            head: inp.head,
423            draft: inp.draft,
424            table: plan.table.as_ref(),
425            ptrs: &ptrs,
426            strides: &strides,
427            formats: &plan.edge_formats,
428            handles,
429            launches: Vec::new(),
430        };
431        emitter(&group.emitter)?
432            .emit(&mut cx)
433            .with_context(|| format!("group {index} (rule `{}`)", group.rule))?;
434        let want = group.launch_count(plan.rows);
435        let want_copies = group.copy_count(plan.rows);
436        let copies = cx
437            .launches
438            .iter()
439            .filter(|l| l.kind == LaunchKind::Copy)
440            .count() as u64;
441        let kernels = cx.launches.len() as u64 - copies;
442        if (kernels, copies) != (want, want_copies) {
443            bail!(
444                "emitter `{}` queued {kernels} launches and {copies} copies for group {index}; \
445                 the plan counts {want} and {want_copies}",
446                group.emitter,
447            );
448        }
449        launches.append(&mut cx.launches);
450    }
451    ensure!(launches.len() as u64 == plan.launches() + plan.copies());
452    Ok(Program {
453        mode: plan.mode,
454        rows: plan.rows,
455        plan_digest: plan.digest.clone(),
456        launches,
457    })
458}
metrale_model_layers/circuit_exec/compile.rs

metrale_model_layers/circuit_exec/
compile.rs
