vyre_reference/execution/
op_count.rs

1//! Thread-local counting of arithmetic IR operations executed by the reference
2//! interpreter, a backend-agnostic dynamic operation count for roofline / complexity
3//! analysis.
4//!
5//! The reference interpreter executes the *same* vyre IR with the *same* data-dependent
6//! control flow that any backend (GPU or CPU) would, so the arithmetic-op count it
7//! reports for a `(program, inputs)` pair equals the dynamic IR-op count the GPU would
8//! execute for the same inputs, at the **vyre-IR** granularity, which is distinct from
9//! (and coarser than) hardware SASS instructions. This gives an honest, non-root
10//! operational-intensity measurement (`ops / bytes`) for the roofline without
11//! Nsight-Compute; the SASS-level dynamic count remains the ncu refinement.
12//!
13//! Counting is a no-op unless a [`count_ops`] scope is active on the current thread, so
14//! ordinary reference evaluation (the vast majority of interpreter use, all in tests)
15//! pays only one thread-local read per arithmetic op and no allocation.
16
17use std::cell::Cell;
18
19thread_local! {
20    /// `Some(n)` while a [`count_ops`] scope is active on this thread; `None` otherwise.
21    static OP_COUNTER: Cell<Option<u64>> = const { Cell::new(None) };
22}
23
24/// Record one arithmetic IR op if counting is active on this thread (a cheap
25/// thread-local read otherwise). Called by the interpreter's `BinOp` / `UnOp` / `Fma`
26/// evaluation arms (the arithmetic operations that make up the roofline's compute term).
27#[inline]
28pub(crate) fn record_op() {
29    OP_COUNTER.with(|counter| {
30        if let Some(count) = counter.get() {
31            counter.set(Some(count.saturating_add(1)));
32        }
33    });
34}
35
36/// Run `f` with arithmetic-op counting active and return `(f's result, arithmetic IR
37/// ops executed)`. The count is every `BinOp` / `UnOp` / `Fma` the reference interpreter
38/// evaluates during `f`: a backend-agnostic dynamic operation count.
39///
40/// Re-entrant-safe: a nested `count_ops` reports its own inner ops AND propagates them to
41/// the enclosing scope, so an outer scope's total still includes work done inside inner
42/// scopes.
43pub fn count_ops<R>(f: impl FnOnce() -> R) -> (R, u64) {
44    let saved = OP_COUNTER.with(|counter| counter.replace(Some(0)));
45    let result = f();
46    let count = OP_COUNTER
47        .with(|counter| counter.replace(saved))
48        .unwrap_or(0);
49    // Propagate our ops to an enclosing scope, if any, so nesting is additive.
50    if saved.is_some() {
51        OP_COUNTER.with(|counter| {
52            if let Some(outer) = counter.get() {
53                counter.set(Some(outer.saturating_add(count)));
54            }
55        });
56    }
57    (result, count)
58}
59
60#[cfg(test)]
61mod tests {
62    use super::*;
63
64    #[test]
65    fn count_ops_is_zero_when_no_ops_run() {
66        let (value, ops) = count_ops(|| 42u32);
67        assert_eq!(value, 42);
68        assert_eq!(ops, 0, "a closure that runs no interpreter ops counts zero");
69    }
70
71    #[test]
72    fn record_op_is_inert_outside_a_count_scope() {
73        // No active scope: record_op must not panic and must not accumulate anywhere.
74        record_op();
75        record_op();
76        let (_, ops) = count_ops(|| {
77            record_op();
78            record_op();
79            record_op();
80        });
81        assert_eq!(ops, 3, "only ops inside the scope are counted");
82    }
83
84    #[test]
85    fn nested_scopes_are_additive() {
86        let ((_, inner_ops), outer_ops) = count_ops(|| {
87            record_op(); // outer: 1
88            let inner = count_ops(|| {
89                record_op();
90                record_op(); // inner: 2
91            });
92            record_op(); // outer: another 1
93            inner
94        });
95        assert_eq!(inner_ops, 2, "inner scope counts only its own ops");
96        assert_eq!(
97            outer_ops, 4,
98            "outer scope counts its own 2 ops plus the 2 propagated from the inner scope"
99        );
100    }
101}
vyre_reference/execution/op_count.rs

vyre_reference/execution/
op_count.rs
