Skip to main content

trueno/brick/
compute_brick.rs

1//! ComputeBrick: Self-verifying, token-centric compute unit.
2//!
3//! Bundles an operation with assertions, budget, and verification.
4//! Also includes BrickLayer for composing multiple bricks.
5
6use std::fmt;
7use std::marker::PhantomData;
8use std::time::Instant;
9
10use super::budget::TokenBudget;
11use super::types::{
12    AssertionResult, Backend, BrickError, BrickVerification, ComputeAssertion, ComputeOp,
13};
14use super::TokenResult;
15
16/// Self-verifying, token-centric compute unit.
17/// Bundles: operation + assertions + budget + verification
18pub struct ComputeBrick<Op: ComputeOp> {
19    /// The compute operation
20    op: Op,
21    /// Falsifiable assertions
22    assertions: Vec<ComputeAssertion>,
23    /// Token-centric performance budget
24    budget: TokenBudget,
25    /// Execution backend
26    backend: Backend,
27    /// Enforce budget (fail if exceeded)
28    enforce_budget: bool,
29    /// Phantom for variance
30    _phantom: PhantomData<Op>,
31}
32
33impl<Op: ComputeOp> ComputeBrick<Op> {
34    /// Create a new compute brick with the given operation.
35    pub fn new(op: Op) -> Self {
36        Self {
37            op,
38            assertions: Vec::new(),
39            budget: TokenBudget::default(),
40            backend: Backend::Auto,
41            enforce_budget: false,
42            _phantom: PhantomData,
43        }
44    }
45
46    /// Add equivalence assertion (output must match baseline backend).
47    #[must_use]
48    pub fn assert_equiv(mut self, baseline: Backend) -> Self {
49        self.assertions.push(ComputeAssertion::equiv(baseline));
50        self
51    }
52
53    /// Add equivalence assertion with custom tolerance.
54    #[must_use]
55    pub fn assert_equiv_with_tolerance(mut self, baseline: Backend, tolerance: f64) -> Self {
56        self.assertions.push(ComputeAssertion::equiv_with_tolerance(baseline, tolerance));
57        self
58    }
59
60    /// Add bounds assertion (output values within range).
61    #[must_use]
62    pub fn assert_bounds(mut self, min: f64, max: f64) -> Self {
63        self.assertions.push(ComputeAssertion::bounds(min, max));
64        self
65    }
66
67    /// Add finite assertion (no NaN/Inf in output).
68    #[must_use]
69    pub fn assert_finite(mut self) -> Self {
70        self.assertions.push(ComputeAssertion::finite());
71        self
72    }
73
74    /// Set token throughput budget (tokens/second).
75    #[must_use]
76    pub fn budget_tok_per_sec(mut self, tps: f64) -> Self {
77        self.budget = TokenBudget::from_throughput(tps);
78        self
79    }
80
81    /// Set token latency budget (microseconds/token).
82    #[must_use]
83    pub fn budget_us_per_tok(mut self, us: f64) -> Self {
84        self.budget = TokenBudget::from_latency(us);
85        self
86    }
87
88    /// Set full budget configuration.
89    #[must_use]
90    pub fn budget(mut self, budget: TokenBudget) -> Self {
91        self.budget = budget;
92        self
93    }
94
95    /// Set execution backend.
96    #[must_use]
97    pub fn backend(mut self, backend: Backend) -> Self {
98        self.backend = backend;
99        self
100    }
101
102    /// Enforce budget (fail if exceeded). Default is false (just report).
103    #[must_use]
104    pub fn enforce_budget(mut self, enforce: bool) -> Self {
105        self.enforce_budget = enforce;
106        self
107    }
108
109    /// Get the brick name (from operation).
110    pub fn name(&self) -> &'static str {
111        self.op.name()
112    }
113
114    /// Get current budget.
115    pub fn get_budget(&self) -> TokenBudget {
116        self.budget
117    }
118
119    /// Get current backend.
120    pub fn get_backend(&self) -> Backend {
121        self.backend
122    }
123
124    /// Get assertions.
125    pub fn get_assertions(&self) -> &[ComputeAssertion] {
126        &self.assertions
127    }
128
129    /// Run the compute brick with full verification (Jidoka gate).
130    pub fn run(&self, input: Op::Input) -> Result<TokenResult<Op::Output>, BrickError> {
131        let tokens = self.op.tokens(&input);
132
133        // Execute with timing
134        let start = Instant::now();
135        let output = self.op.execute(input, self.backend)?;
136        let elapsed_us = start.elapsed().as_secs_f64() * 1_000_000.0;
137
138        // Calculate metrics
139        let us_per_token = if tokens > 0 { elapsed_us / tokens as f64 } else { elapsed_us };
140        let tokens_per_sec =
141            if elapsed_us > 0.0 { tokens as f64 * 1_000_000.0 / elapsed_us } else { f64::INFINITY };
142        let budget_met = self.budget.is_met(us_per_token);
143        let budget_utilization = self.budget.utilization(us_per_token);
144
145        // Check budget enforcement
146        if self.enforce_budget && !budget_met {
147            return Err(BrickError::BudgetExceeded {
148                limit_us: self.budget.us_per_token,
149                actual_us: us_per_token,
150                utilization: budget_utilization * 100.0,
151            });
152        }
153
154        Ok(TokenResult {
155            output,
156            tokens_processed: tokens,
157            us_per_token,
158            tokens_per_sec,
159            budget_met,
160            budget_utilization,
161        })
162    }
163
164    /// Verify assertions without full execution.
165    /// Returns verification status.
166    pub fn verify(&self) -> BrickVerification {
167        let start = Instant::now();
168
169        // Check if we have assertions (Popperian requirement)
170        if self.assertions.is_empty() {
171            return BrickVerification {
172                passed: false,
173                assertion_results: vec![AssertionResult {
174                    assertion: ComputeAssertion::Custom {
175                        name: "popperian_falsifiability".to_string(),
176                    },
177                    passed: false,
178                    error: Some(
179                        "No assertions defined - violates Popperian falsifiability".to_string(),
180                    ),
181                }],
182                verification_us: start.elapsed().as_secs_f64() * 1_000_000.0,
183            };
184        }
185
186        // For now, just validate assertion structure
187        // Full verification requires input data
188        let results: Vec<AssertionResult> = self
189            .assertions
190            .iter()
191            .map(|a| AssertionResult { assertion: a.clone(), passed: true, error: None })
192            .collect();
193
194        let passed = results.iter().all(|r| r.passed);
195
196        BrickVerification {
197            passed,
198            assertion_results: results,
199            verification_us: start.elapsed().as_secs_f64() * 1_000_000.0,
200        }
201    }
202}
203
204impl<Op: ComputeOp + Clone> Clone for ComputeBrick<Op> {
205    fn clone(&self) -> Self {
206        Self {
207            op: self.op.clone(),
208            assertions: self.assertions.clone(),
209            budget: self.budget,
210            backend: self.backend,
211            enforce_budget: self.enforce_budget,
212            _phantom: PhantomData,
213        }
214    }
215}
216
217impl<Op: ComputeOp> fmt::Debug for ComputeBrick<Op> {
218    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
219        f.debug_struct("ComputeBrick")
220            .field("name", &self.op.name())
221            .field("backend", &self.backend)
222            .field("budget", &self.budget)
223            .field("assertions", &self.assertions.len())
224            .field("enforce_budget", &self.enforce_budget)
225            .finish()
226    }
227}
228
229// ============================================================================
230// LLM Transformer Fused Operations (PMAT-PERF-009)
231// BrickLayer: Compose multiple bricks
232// ============================================================================
233
234/// A layer of compute bricks that execute sequentially.
235/// Throughput ceiling = min(component throughputs).
236#[derive(Debug, Default)]
237pub struct BrickLayer {
238    /// Named bricks in this layer
239    bricks: Vec<(String, f64)>, // (name, budget_tok_per_sec)
240}
241
242impl BrickLayer {
243    /// Create a new empty layer.
244    pub fn new() -> Self {
245        Self::default()
246    }
247
248    /// Add a brick to the layer.
249    #[must_use]
250    pub fn with_brick<Op: ComputeOp>(mut self, brick: &ComputeBrick<Op>) -> Self {
251        self.bricks.push((brick.name().to_string(), brick.budget.tokens_per_sec));
252        self
253    }
254
255    /// Add a named entry with throughput budget.
256    #[must_use]
257    pub fn with_named(mut self, name: &str, budget_tok_per_sec: f64) -> Self {
258        self.bricks.push((name.to_string(), budget_tok_per_sec));
259        self
260    }
261
262    /// Get the throughput ceiling (bottleneck).
263    /// Layer throughput = min(component throughputs).
264    pub fn throughput_ceiling(&self) -> f64 {
265        self.bricks.iter().map(|(_, tps)| *tps).fold(f64::INFINITY, f64::min)
266    }
267
268    /// Get the bottleneck brick name.
269    pub fn bottleneck(&self) -> Option<&str> {
270        self.bricks
271            .iter()
272            .min_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal))
273            .map(|(name, _)| name.as_str())
274    }
275
276    /// Get all bricks with their budgets.
277    pub fn bricks(&self) -> &[(String, f64)] {
278        &self.bricks
279    }
280}