trueno/brick/
compute_brick.rs1use 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
16pub struct ComputeBrick<Op: ComputeOp> {
19 op: Op,
21 assertions: Vec<ComputeAssertion>,
23 budget: TokenBudget,
25 backend: Backend,
27 enforce_budget: bool,
29 _phantom: PhantomData<Op>,
31}
32
33impl<Op: ComputeOp> ComputeBrick<Op> {
34 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 #[must_use]
48 pub fn assert_equiv(mut self, baseline: Backend) -> Self {
49 self.assertions.push(ComputeAssertion::equiv(baseline));
50 self
51 }
52
53 #[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 #[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 #[must_use]
69 pub fn assert_finite(mut self) -> Self {
70 self.assertions.push(ComputeAssertion::finite());
71 self
72 }
73
74 #[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 #[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 #[must_use]
90 pub fn budget(mut self, budget: TokenBudget) -> Self {
91 self.budget = budget;
92 self
93 }
94
95 #[must_use]
97 pub fn backend(mut self, backend: Backend) -> Self {
98 self.backend = backend;
99 self
100 }
101
102 #[must_use]
104 pub fn enforce_budget(mut self, enforce: bool) -> Self {
105 self.enforce_budget = enforce;
106 self
107 }
108
109 pub fn name(&self) -> &'static str {
111 self.op.name()
112 }
113
114 pub fn get_budget(&self) -> TokenBudget {
116 self.budget
117 }
118
119 pub fn get_backend(&self) -> Backend {
121 self.backend
122 }
123
124 pub fn get_assertions(&self) -> &[ComputeAssertion] {
126 &self.assertions
127 }
128
129 pub fn run(&self, input: Op::Input) -> Result<TokenResult<Op::Output>, BrickError> {
131 let tokens = self.op.tokens(&input);
132
133 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 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 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 pub fn verify(&self) -> BrickVerification {
167 let start = Instant::now();
168
169 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 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#[derive(Debug, Default)]
237pub struct BrickLayer {
238 bricks: Vec<(String, f64)>, }
241
242impl BrickLayer {
243 pub fn new() -> Self {
245 Self::default()
246 }
247
248 #[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 #[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 pub fn throughput_ceiling(&self) -> f64 {
265 self.bricks.iter().map(|(_, tps)| *tps).fold(f64::INFINITY, f64::min)
266 }
267
268 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 pub fn bricks(&self) -> &[(String, f64)] {
278 &self.bricks
279 }
280}