/home/noah/src/trueno/src/brick/mod.rs
Line | Count | Source |
1 | | //! ComputeBrick: Token-Centric Compute Units |
2 | | //! |
3 | | //! A **ComputeBrick** is a self-verifying, token-centric compute unit that bundles: |
4 | | //! - **Operation**: The compute operation (matmul, dot, softmax, etc.) |
5 | | //! - **Assertions**: Falsifiable claims about the output (equivalence, bounds) |
6 | | //! - **Budget**: Performance target in µs/token or tokens/sec |
7 | | //! - **Backend**: Execution target (Scalar, AVX2, CUDA, etc.) |
8 | | //! |
9 | | //! # Core Insight |
10 | | //! |
11 | | //! A **token** is the unit of data; a **ComputeBrick** is the unit of compute. |
12 | | //! |
13 | | //! ```text |
14 | | //! Token ──▶ [ComputeBrick] ──▶ Token |
15 | | //! (matmul, softmax, attention) |
16 | | //! ``` |
17 | | //! |
18 | | //! # Example |
19 | | //! |
20 | | //! ```rust,ignore |
21 | | //! use trueno::brick::{ComputeBrick, ComputeBackend, MatmulOp}; |
22 | | //! |
23 | | //! let matmul = ComputeBrick::new(MatmulOp::new(1024, 1024, 1024)) |
24 | | //! .assert_equiv(ComputeBackend::Scalar) |
25 | | //! .budget_tok_per_sec(50_000.0) |
26 | | //! .backend(ComputeBackend::Avx2); |
27 | | //! |
28 | | //! let result = matmul.run((a, b))?; |
29 | | //! println!("Throughput: {:.0} tok/s", result.tokens_per_sec); |
30 | | //! ``` |
31 | | //! |
32 | | //! # Scientific Basis |
33 | | //! |
34 | | //! Per Popper (1959), a theory that makes no falsifiable predictions is not scientific. |
35 | | //! A ComputeBrick with no assertions makes no testable claims and is therefore invalid. |
36 | | |
37 | | // Submodules |
38 | | mod batch; |
39 | | mod buffer; |
40 | | mod circuit; |
41 | | mod connection; |
42 | | mod memory; |
43 | | mod perf_metrics; |
44 | | mod profiling; |
45 | | mod rate_limit; |
46 | | mod resource_pool; |
47 | | mod shutdown; |
48 | | |
49 | | // Re-export profiling functions |
50 | | pub use profiling::{ |
51 | | cached_nanos, cached_nanos_or_now, cpu_cycles, get_page_faults, init_time_service, |
52 | | with_page_fault_tracking, |
53 | | }; |
54 | | |
55 | | // Re-export perf_metrics types |
56 | | pub use perf_metrics::{InferencePhase, PerfMetrics}; |
57 | | |
58 | | // Re-export memory types |
59 | | #[cfg(not(target_arch = "wasm32"))] |
60 | | pub use memory::AlignedBuffer; |
61 | | pub use memory::{ |
62 | | is_direct_io_aligned, madvise_region, prefetch_for_inference, prefetch_ptr, prefetch_slice, |
63 | | CacheAligned, MemoryAdvice, PrefetchLocality, CACHE_LINE_SIZE, CACHE_LINE_SIZE_F32, |
64 | | DIRECT_IO_ALIGNMENT, |
65 | | }; |
66 | | |
67 | | // Re-export buffer types |
68 | | pub use buffer::{BufferWatermarks, WatermarkedBuffer}; |
69 | | |
70 | | // Re-export circuit breaker types |
71 | | pub use circuit::{CircuitBreaker, CircuitState}; |
72 | | |
73 | | // Re-export shutdown types |
74 | | pub use shutdown::{GracefulShutdown, ShutdownGuard, ShutdownResult}; |
75 | | |
76 | | // Re-export resource pool types |
77 | | pub use resource_pool::{PooledResource, ResourcePool}; |
78 | | |
79 | | // Re-export rate limiting types |
80 | | pub use rate_limit::{LimitError, ServeLimits}; |
81 | | |
82 | | // Re-export connection types |
83 | | pub use connection::{ConnectionState, KeepAliveConfig, ManagedConnection}; |
84 | | |
85 | | // Re-export batch types |
86 | | pub use batch::{balance211, Balance211Iter, BatchSplitStrategy, split_batch}; |
87 | | |
88 | | // KV cache management |
89 | | mod kv_cache; |
90 | | pub use kv_cache::{KvCacheManager, KvCacheSlotInfo, SequentialBatchOrderer}; |
91 | | |
92 | | // SIMD configuration |
93 | | mod simd_config; |
94 | | pub use simd_config::{ |
95 | | unroll_tail_process, AmxTileConfig, LazySimdConfig, SimdBackendState, UnrollFactor, |
96 | | UnrollTailIterator, |
97 | | }; |
98 | | |
99 | | // Execution graph and brick profiling types (PAR-073, PAR-200, PAR-201) |
100 | | mod exec_graph; |
101 | | pub use exec_graph::{ |
102 | | BrickBottleneck, BrickCategory, BrickId, BrickSample, BrickStats, CategoryStats, EdgeType, |
103 | | ExecutionEdge, ExecutionGraph, ExecutionNode, ExecutionNodeId, PtxRegistry, SyncMode, |
104 | | TransferDirection, |
105 | | }; |
106 | | |
107 | | // BrickProfiler and tile profiling (TILING-SPEC-001) |
108 | | mod profiler; |
109 | | pub use profiler::{ |
110 | | fnv1a_f32_checksum, BrickIdTimer, BrickProfiler, BrickTimer, DivergenceInfo, KernelChecksum, |
111 | | TileLevel, TileStats, TileTimer, |
112 | | }; |
113 | | |
114 | | // Model-level inference tracing (Phase 13, E.11) |
115 | | mod tracing; |
116 | | pub use tracing::{ |
117 | | AttentionTraceConfig, AttentionWeightTrace, KvCacheSessionTrace, KvCacheStateTrace, |
118 | | LayerActivationTrace, LogitEvolutionTrace, ModelActivationTrace, ModelQuantizationError, |
119 | | ModelTracer, ModelTracerConfig, ModelTracerSummary, QuantType, QuantizationErrorTrace, |
120 | | TensorStats, TokenLogitEvolution, |
121 | | }; |
122 | | |
123 | | // Async and buffer patterns (Phase 12, E.10) |
124 | | mod patterns; |
125 | | pub use patterns::{ |
126 | | reserve_capacity, AsyncResult, BoundedQueue, DualWakerState, FlowControlError, |
127 | | GraphReuseCounter, ReserveStrategy, StrategicBuffer, StreamCapacity, WakeDecision, |
128 | | WakeSkipState, |
129 | | }; |
130 | | |
131 | | // Built-in compute operations |
132 | | mod ops; |
133 | | pub use ops::{AddOp, DotOp, MatmulOp, SoftmaxOp}; |
134 | | |
135 | | // Fused operations for transformer inference (PMAT-PERF-009) |
136 | | mod fused_ops; |
137 | | pub use fused_ops::{FusedGateUpOp, FusedGateUpWeights, FusedQKVOp, FusedQKVWeights}; |
138 | | |
139 | | // SIMD-optimized attention operation (PMAT-017) |
140 | | mod attention; |
141 | | pub use attention::AttentionOp; |
142 | | |
143 | | // Q5_K and Q6_K quantization operations (llama.cpp compatible) |
144 | | mod quant_ops; |
145 | | pub use quant_ops::{BlockQ5K, BlockQ6K, DotQ5KOp, DotQ6KOp}; |
146 | | |
147 | | // Tests (7,400+ lines extracted for TDG compliance) |
148 | | #[cfg(test)] |
149 | | mod tests; |
150 | | |
151 | | use crate::error::TruenoError; |
152 | | use std::fmt; |
153 | | use std::marker::PhantomData; |
154 | | use std::time::Instant; |
155 | | |
156 | | // ============================================================================ |
157 | | // Async Task Profiler (Pattern 3 from actix-web) |
158 | | // ============================================================================ |
159 | | |
160 | | /// Async task profiler for measuring poll efficiency (Phase 11, E.9.4). |
161 | | /// |
162 | | /// Tracks how many times a future is polled before completion. |
163 | | /// High poll counts indicate inefficient async code or spurious wakeups. |
164 | | /// |
165 | | /// # Example |
166 | | /// ```rust,ignore |
167 | | /// let mut profiler = AsyncTaskProfiler::new("inference_request"); |
168 | | /// |
169 | | /// profiler.on_poll_start(); |
170 | | /// // ... poll the future ... |
171 | | /// profiler.on_poll_end(is_ready); |
172 | | /// |
173 | | /// println!("Poll efficiency: {:.1}%", profiler.efficiency() * 100.0); |
174 | | /// ``` |
175 | | #[derive(Debug, Clone)] |
176 | | pub struct AsyncTaskProfiler { |
177 | | /// Task name for identification |
178 | | pub name: String, |
179 | | /// Number of times poll() was called |
180 | | pub poll_count: u64, |
181 | | /// Number of times poll() returned Pending |
182 | | pub yield_count: u64, |
183 | | /// Total time spent in poll() (nanoseconds) |
184 | | pub total_poll_ns: u64, |
185 | | /// Start time of current poll |
186 | | last_poll_start: u64, |
187 | | /// CPU cycles at poll start |
188 | | last_poll_cycles: u64, |
189 | | /// Total CPU cycles in poll() |
190 | | pub total_poll_cycles: u64, |
191 | | } |
192 | | |
193 | | impl AsyncTaskProfiler { |
194 | | /// Create a new async task profiler. |
195 | 0 | pub fn new(name: impl Into<String>) -> Self { |
196 | 0 | Self { |
197 | 0 | name: name.into(), |
198 | 0 | poll_count: 0, |
199 | 0 | yield_count: 0, |
200 | 0 | total_poll_ns: 0, |
201 | 0 | last_poll_start: 0, |
202 | 0 | last_poll_cycles: 0, |
203 | 0 | total_poll_cycles: 0, |
204 | 0 | } |
205 | 0 | } |
206 | | |
207 | | /// Call at the start of each poll() invocation. |
208 | | #[inline] |
209 | 0 | pub fn on_poll_start(&mut self) { |
210 | 0 | self.poll_count += 1; |
211 | 0 | self.last_poll_start = cached_nanos_or_now(); |
212 | 0 | self.last_poll_cycles = cpu_cycles(); |
213 | 0 | } |
214 | | |
215 | | /// Call at the end of each poll() invocation. |
216 | | /// |
217 | | /// # Arguments |
218 | | /// - `is_ready`: true if the future returned Poll::Ready |
219 | | #[inline] |
220 | 0 | pub fn on_poll_end(&mut self, is_ready: bool) { |
221 | 0 | let now = cached_nanos_or_now(); |
222 | 0 | let cycles = cpu_cycles(); |
223 | | |
224 | 0 | self.total_poll_ns += now.saturating_sub(self.last_poll_start); |
225 | 0 | self.total_poll_cycles += cycles.saturating_sub(self.last_poll_cycles); |
226 | | |
227 | 0 | if !is_ready { |
228 | 0 | self.yield_count += 1; |
229 | 0 | } |
230 | 0 | } |
231 | | |
232 | | /// Poll efficiency ratio (0.0 to 1.0). |
233 | | /// |
234 | | /// - 1.0 = Perfect (ready on first poll) |
235 | | /// - 0.5 = 2 polls required |
236 | | /// - Lower = more wakeups/polls needed |
237 | | #[must_use] |
238 | 0 | pub fn efficiency(&self) -> f64 { |
239 | 0 | if self.poll_count == 0 { |
240 | 0 | 0.0 |
241 | | } else { |
242 | 0 | 1.0 / self.poll_count as f64 |
243 | | } |
244 | 0 | } |
245 | | |
246 | | /// Average time per poll in microseconds. |
247 | | #[must_use] |
248 | 0 | pub fn avg_poll_us(&self) -> f64 { |
249 | 0 | if self.poll_count == 0 { |
250 | 0 | 0.0 |
251 | | } else { |
252 | 0 | self.total_poll_ns as f64 / self.poll_count as f64 / 1000.0 |
253 | | } |
254 | 0 | } |
255 | | |
256 | | /// Yield ratio (Pending / total polls). |
257 | | /// |
258 | | /// High yield ratio indicates the task is often not ready when polled. |
259 | | #[must_use] |
260 | 0 | pub fn yield_ratio(&self) -> f64 { |
261 | 0 | if self.poll_count == 0 { |
262 | 0 | 0.0 |
263 | | } else { |
264 | 0 | self.yield_count as f64 / self.poll_count as f64 |
265 | | } |
266 | 0 | } |
267 | | |
268 | | /// Convert to ExecutionNode for graph integration. |
269 | 0 | pub fn to_execution_node(&self) -> ExecutionNode { |
270 | 0 | ExecutionNode::AsyncTask { |
271 | 0 | name: self.name.clone(), |
272 | 0 | poll_count: self.poll_count, |
273 | 0 | yield_count: self.yield_count, |
274 | 0 | total_poll_ns: self.total_poll_ns, |
275 | 0 | } |
276 | 0 | } |
277 | | } |
278 | | |
279 | | impl Default for AsyncTaskProfiler { |
280 | 0 | fn default() -> Self { |
281 | 0 | Self::new("unnamed") |
282 | 0 | } |
283 | | } |
284 | | |
285 | | |
286 | | /// Execution backend for compute operations. |
287 | | /// This is the brick-specific backend enum with additional GPU backends. |
288 | | #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Default)] |
289 | | pub enum ComputeBackend { |
290 | | /// Pure Rust scalar fallback (always available, baseline for correctness) |
291 | | Scalar, |
292 | | /// SSE2 SIMD (x86_64 baseline) |
293 | | Sse2, |
294 | | /// AVX2 256-bit SIMD with FMA |
295 | | #[default] |
296 | | Avx2, |
297 | | /// AVX-512 512-bit SIMD |
298 | | Avx512, |
299 | | /// ARM NEON SIMD |
300 | | Neon, |
301 | | /// WebAssembly SIMD128 |
302 | | Wasm, |
303 | | /// NVIDIA CUDA via PTX |
304 | | Cuda, |
305 | | /// Cross-platform GPU via wgpu |
306 | | Wgpu, |
307 | | /// Auto-select best available backend |
308 | | Auto, |
309 | | } |
310 | | |
311 | | impl fmt::Display for ComputeBackend { |
312 | 0 | fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { |
313 | 0 | match self { |
314 | 0 | ComputeBackend::Scalar => write!(f, "Scalar"), |
315 | 0 | ComputeBackend::Sse2 => write!(f, "SSE2"), |
316 | 0 | ComputeBackend::Avx2 => write!(f, "AVX2"), |
317 | 0 | ComputeBackend::Avx512 => write!(f, "AVX-512"), |
318 | 0 | ComputeBackend::Neon => write!(f, "NEON"), |
319 | 0 | ComputeBackend::Wasm => write!(f, "WASM"), |
320 | 0 | ComputeBackend::Cuda => write!(f, "CUDA"), |
321 | 0 | ComputeBackend::Wgpu => write!(f, "wgpu"), |
322 | 0 | ComputeBackend::Auto => write!(f, "Auto"), |
323 | | } |
324 | 0 | } |
325 | | } |
326 | | |
327 | | /// Type alias for backward compatibility |
328 | | pub type Backend = ComputeBackend; |
329 | | |
330 | | /// Performance budget expressed in token terms. |
331 | | /// Aligns compute costs with LLM inference metrics. |
332 | | #[derive(Debug, Clone, Copy)] |
333 | | pub struct TokenBudget { |
334 | | /// Latency budget per token (microseconds) |
335 | | pub us_per_token: f64, |
336 | | /// Throughput target (tokens/second) |
337 | | pub tokens_per_sec: f64, |
338 | | /// Batch size for amortization |
339 | | pub batch_size: usize, |
340 | | } |
341 | | |
342 | | /// Performance budget for byte-oriented operations (compression, I/O). |
343 | | /// Use this for trueno-zram, disk I/O, network throughput, etc. |
344 | | /// |
345 | | /// PMAT-452: Serializable for hardware.toml export. |
346 | | #[derive(Debug, Clone, Copy, serde::Serialize, serde::Deserialize)] |
347 | | pub struct ByteBudget { |
348 | | /// Latency budget per page (microseconds) |
349 | | pub us_per_page: f64, |
350 | | /// Throughput target (GB/s) |
351 | | pub gb_per_sec: f64, |
352 | | /// Page size in bytes (default 4096) |
353 | | pub page_size: usize, |
354 | | } |
355 | | |
356 | | impl Default for ByteBudget { |
357 | 0 | fn default() -> Self { |
358 | | // Default: 25 GB/s (trueno-zram ZSTD target) |
359 | 0 | Self::from_throughput(25.0) |
360 | 0 | } |
361 | | } |
362 | | |
363 | | impl ByteBudget { |
364 | | /// Create budget from throughput target (GB/s). |
365 | | /// 25 GB/s = 0.16µs per 4KB page |
366 | 0 | pub fn from_throughput(gb_per_sec: f64) -> Self { |
367 | 0 | let bytes_per_sec = gb_per_sec * 1e9; |
368 | 0 | let pages_per_sec = bytes_per_sec / 4096.0; |
369 | 0 | Self { |
370 | 0 | us_per_page: 1_000_000.0 / pages_per_sec, |
371 | 0 | gb_per_sec, |
372 | 0 | page_size: 4096, |
373 | 0 | } |
374 | 0 | } |
375 | | |
376 | | /// Create budget from latency target (µs per page). |
377 | 0 | pub fn from_latency(us_per_page: f64) -> Self { |
378 | 0 | let pages_per_sec = 1_000_000.0 / us_per_page; |
379 | 0 | let bytes_per_sec = pages_per_sec * 4096.0; |
380 | 0 | Self { |
381 | 0 | us_per_page, |
382 | 0 | gb_per_sec: bytes_per_sec / 1e9, |
383 | 0 | page_size: 4096, |
384 | 0 | } |
385 | 0 | } |
386 | | |
387 | | /// Set custom page size (e.g., 64KB for huge pages). |
388 | | #[must_use] |
389 | 0 | pub fn with_page_size(mut self, page_size: usize) -> Self { |
390 | | // Recalculate us_per_page based on new page size |
391 | 0 | let bytes_per_sec = self.gb_per_sec * 1e9; |
392 | 0 | let pages_per_sec = bytes_per_sec / page_size as f64; |
393 | 0 | self.us_per_page = 1_000_000.0 / pages_per_sec; |
394 | 0 | self.page_size = page_size; |
395 | 0 | self |
396 | 0 | } |
397 | | |
398 | | /// Convert to TokenBudget (1 token = 1 page). |
399 | | /// Useful for integrating byte workloads with token-centric monitoring. |
400 | 0 | pub fn to_token_budget(&self) -> TokenBudget { |
401 | 0 | TokenBudget { |
402 | 0 | us_per_token: self.us_per_page, |
403 | 0 | tokens_per_sec: 1_000_000.0 / self.us_per_page, |
404 | 0 | batch_size: 1, |
405 | 0 | } |
406 | 0 | } |
407 | | |
408 | | /// Check if actual performance meets budget. |
409 | 0 | pub fn is_met(&self, actual_us_per_page: f64) -> bool { |
410 | 0 | actual_us_per_page <= self.us_per_page |
411 | 0 | } |
412 | | |
413 | | /// Calculate budget utilization. |
414 | 0 | pub fn utilization(&self, actual_us_per_page: f64) -> f64 { |
415 | 0 | actual_us_per_page / self.us_per_page |
416 | 0 | } |
417 | | |
418 | | /// Calculate actual throughput from latency. |
419 | 0 | pub fn throughput_from_latency(us_per_page: f64, page_size: usize) -> f64 { |
420 | 0 | let pages_per_sec = 1_000_000.0 / us_per_page; |
421 | 0 | pages_per_sec * page_size as f64 / 1e9 |
422 | 0 | } |
423 | | } |
424 | | |
425 | | impl Default for TokenBudget { |
426 | 0 | fn default() -> Self { |
427 | | // Default: 50µs/token = 20,000 tokens/sec |
428 | 0 | Self::from_latency(50.0) |
429 | 0 | } |
430 | | } |
431 | | |
432 | | impl TokenBudget { |
433 | | /// Create budget from latency target. |
434 | | /// 50µs/token = 20,000 tokens/sec |
435 | 0 | pub fn from_latency(us_per_token: f64) -> Self { |
436 | 0 | Self { |
437 | 0 | us_per_token, |
438 | 0 | tokens_per_sec: 1_000_000.0 / us_per_token, |
439 | 0 | batch_size: 1, |
440 | 0 | } |
441 | 0 | } |
442 | | |
443 | | /// Create budget from throughput target. |
444 | | /// 20,000 tokens/sec = 50µs/token |
445 | 0 | pub fn from_throughput(tokens_per_sec: f64) -> Self { |
446 | 0 | Self { |
447 | 0 | us_per_token: 1_000_000.0 / tokens_per_sec, |
448 | 0 | tokens_per_sec, |
449 | 0 | batch_size: 1, |
450 | 0 | } |
451 | 0 | } |
452 | | |
453 | | /// Set batch size for amortization. |
454 | | #[must_use] |
455 | 0 | pub fn with_batch_size(mut self, batch_size: usize) -> Self { |
456 | 0 | self.batch_size = batch_size.max(1); |
457 | 0 | self |
458 | 0 | } |
459 | | |
460 | | /// Check if actual performance meets budget. |
461 | 0 | pub fn is_met(&self, actual_us_per_token: f64) -> bool { |
462 | 0 | actual_us_per_token <= self.us_per_token |
463 | 0 | } |
464 | | |
465 | | /// Calculate budget utilization (0.0 = unused, 1.0 = exactly at budget, >1.0 = over budget). |
466 | 0 | pub fn utilization(&self, actual_us_per_token: f64) -> f64 { |
467 | 0 | actual_us_per_token / self.us_per_token |
468 | 0 | } |
469 | | } |
470 | | |
471 | | /// Result of ComputeBrick execution with token metrics. |
472 | | #[derive(Debug, Clone)] |
473 | | pub struct TokenResult<T> { |
474 | | /// Computed output |
475 | | pub output: T, |
476 | | /// Number of tokens processed |
477 | | pub tokens_processed: usize, |
478 | | /// Actual latency (microseconds/token) |
479 | | pub us_per_token: f64, |
480 | | /// Actual throughput (tokens/second) |
481 | | pub tokens_per_sec: f64, |
482 | | /// Did we meet the budget? |
483 | | pub budget_met: bool, |
484 | | /// Budget utilization (0.0-1.0+ where 1.0 = exactly at budget) |
485 | | pub budget_utilization: f64, |
486 | | } |
487 | | |
488 | | impl<T> TokenResult<T> { |
489 | | /// Map the output to a new type. |
490 | 0 | pub fn map<U, F: FnOnce(T) -> U>(self, f: F) -> TokenResult<U> { |
491 | 0 | TokenResult { |
492 | 0 | output: f(self.output), |
493 | 0 | tokens_processed: self.tokens_processed, |
494 | 0 | us_per_token: self.us_per_token, |
495 | 0 | tokens_per_sec: self.tokens_per_sec, |
496 | 0 | budget_met: self.budget_met, |
497 | 0 | budget_utilization: self.budget_utilization, |
498 | 0 | } |
499 | 0 | } |
500 | | } |
501 | | |
502 | | /// Errors from ComputeBrick execution. |
503 | | /// Tells you exactly what failed (Jidoka: stop and signal). |
504 | | #[derive(Debug, thiserror::Error)] |
505 | | pub enum BrickError { |
506 | | /// Assertion failed during verification |
507 | | #[error("Assertion failed: {name} - expected {expected}, got {actual}")] |
508 | | AssertionFailed { |
509 | | name: String, |
510 | | expected: String, |
511 | | actual: String, |
512 | | }, |
513 | | |
514 | | /// Performance budget exceeded |
515 | | #[error("Budget exceeded: {limit_us:.1}µs/tok limit, {actual_us:.1}µs/tok actual ({utilization:.0}% of budget)")] |
516 | | BudgetExceeded { |
517 | | limit_us: f64, |
518 | | actual_us: f64, |
519 | | utilization: f64, |
520 | | }, |
521 | | |
522 | | /// Underlying compute error |
523 | | #[error("Compute error: {0}")] |
524 | | ComputeError(#[from] TruenoError), |
525 | | |
526 | | /// No assertions defined (violates Popperian falsifiability) |
527 | | #[error("Brick has no assertions - violates Popperian falsifiability requirement")] |
528 | | NoAssertions, |
529 | | |
530 | | /// Backend not available |
531 | | #[error("Backend {0} not available on this system")] |
532 | | BackendUnavailable(Backend), |
533 | | } |
534 | | |
535 | | /// Type of assertion for compute verification. |
536 | | #[derive(Debug, Clone)] |
537 | | pub enum ComputeAssertion { |
538 | | /// Output must match baseline backend within tolerance |
539 | | Equivalence { |
540 | | baseline: Backend, |
541 | | tolerance: f64, |
542 | | }, |
543 | | /// Output values must be within bounds |
544 | | Bounds { |
545 | | min: f64, |
546 | | max: f64, |
547 | | }, |
548 | | /// Output must not contain NaN or infinity |
549 | | Finite, |
550 | | /// Custom assertion with name and check function index |
551 | | Custom { |
552 | | name: String, |
553 | | }, |
554 | | } |
555 | | |
556 | | impl ComputeAssertion { |
557 | | /// Create equivalence assertion with default tolerance (1e-5). |
558 | 0 | pub fn equiv(baseline: Backend) -> Self { |
559 | 0 | Self::Equivalence { |
560 | 0 | baseline, |
561 | 0 | tolerance: 1e-5, |
562 | 0 | } |
563 | 0 | } |
564 | | |
565 | | /// Create equivalence assertion with custom tolerance. |
566 | 0 | pub fn equiv_with_tolerance(baseline: Backend, tolerance: f64) -> Self { |
567 | 0 | Self::Equivalence { baseline, tolerance } |
568 | 0 | } |
569 | | |
570 | | /// Create bounds assertion. |
571 | 0 | pub fn bounds(min: f64, max: f64) -> Self { |
572 | 0 | Self::Bounds { min, max } |
573 | 0 | } |
574 | | |
575 | | /// Create finite assertion (no NaN/Inf). |
576 | 0 | pub fn finite() -> Self { |
577 | 0 | Self::Finite |
578 | 0 | } |
579 | | } |
580 | | |
581 | | /// Verification result from ComputeBrick. |
582 | | #[derive(Debug, Clone)] |
583 | | pub struct BrickVerification { |
584 | | /// Overall pass/fail |
585 | | pub passed: bool, |
586 | | /// Individual assertion results |
587 | | pub assertion_results: Vec<AssertionResult>, |
588 | | /// Verification time in microseconds |
589 | | pub verification_us: f64, |
590 | | } |
591 | | |
592 | | impl BrickVerification { |
593 | | /// Check if all assertions passed. |
594 | 0 | pub fn is_valid(&self) -> bool { |
595 | 0 | self.passed |
596 | 0 | } |
597 | | |
598 | | /// Get failed assertions. |
599 | 0 | pub fn failures(&self) -> impl Iterator<Item = &AssertionResult> { |
600 | 0 | self.assertion_results.iter().filter(|r| !r.passed) |
601 | 0 | } |
602 | | } |
603 | | |
604 | | /// Result of a single assertion check. |
605 | | #[derive(Debug, Clone)] |
606 | | pub struct AssertionResult { |
607 | | /// Assertion that was checked |
608 | | pub assertion: ComputeAssertion, |
609 | | /// Did it pass? |
610 | | pub passed: bool, |
611 | | /// Error message if failed |
612 | | pub error: Option<String>, |
613 | | } |
614 | | |
615 | | /// Trait for compute operations that can be wrapped in a ComputeBrick. |
616 | | pub trait ComputeOp: Send + Sync { |
617 | | /// Input type for this operation |
618 | | type Input; |
619 | | /// Output type for this operation |
620 | | type Output; |
621 | | |
622 | | /// Operation name for identification |
623 | | fn name(&self) -> &'static str; |
624 | | |
625 | | /// Execute the operation on the given backend |
626 | | fn execute(&self, input: Self::Input, backend: Backend) -> Result<Self::Output, TruenoError>; |
627 | | |
628 | | /// Number of tokens this operation processes (for budget calculation) |
629 | | fn tokens(&self, input: &Self::Input) -> usize; |
630 | | |
631 | | /// Clone the input for verification (if needed) |
632 | 0 | fn clone_input(&self, input: &Self::Input) -> Option<Self::Input> |
633 | 0 | where |
634 | 0 | Self::Input: Clone, |
635 | | { |
636 | 0 | Some(input.clone()) |
637 | 0 | } |
638 | | } |
639 | | |
640 | | /// Self-verifying, token-centric compute unit. |
641 | | /// Bundles: operation + assertions + budget + verification |
642 | | pub struct ComputeBrick<Op: ComputeOp> { |
643 | | /// The compute operation |
644 | | op: Op, |
645 | | /// Falsifiable assertions |
646 | | assertions: Vec<ComputeAssertion>, |
647 | | /// Token-centric performance budget |
648 | | budget: TokenBudget, |
649 | | /// Execution backend |
650 | | backend: Backend, |
651 | | /// Enforce budget (fail if exceeded) |
652 | | enforce_budget: bool, |
653 | | /// Phantom for variance |
654 | | _phantom: PhantomData<Op>, |
655 | | } |
656 | | |
657 | | impl<Op: ComputeOp> ComputeBrick<Op> { |
658 | | /// Create a new compute brick with the given operation. |
659 | 0 | pub fn new(op: Op) -> Self { |
660 | 0 | Self { |
661 | 0 | op, |
662 | 0 | assertions: Vec::new(), |
663 | 0 | budget: TokenBudget::default(), |
664 | 0 | backend: Backend::Auto, |
665 | 0 | enforce_budget: false, |
666 | 0 | _phantom: PhantomData, |
667 | 0 | } |
668 | 0 | } |
669 | | |
670 | | /// Add equivalence assertion (output must match baseline backend). |
671 | | #[must_use] |
672 | 0 | pub fn assert_equiv(mut self, baseline: Backend) -> Self { |
673 | 0 | self.assertions.push(ComputeAssertion::equiv(baseline)); |
674 | 0 | self |
675 | 0 | } |
676 | | |
677 | | /// Add equivalence assertion with custom tolerance. |
678 | | #[must_use] |
679 | 0 | pub fn assert_equiv_with_tolerance(mut self, baseline: Backend, tolerance: f64) -> Self { |
680 | 0 | self.assertions |
681 | 0 | .push(ComputeAssertion::equiv_with_tolerance(baseline, tolerance)); |
682 | 0 | self |
683 | 0 | } |
684 | | |
685 | | /// Add bounds assertion (output values within range). |
686 | | #[must_use] |
687 | 0 | pub fn assert_bounds(mut self, min: f64, max: f64) -> Self { |
688 | 0 | self.assertions.push(ComputeAssertion::bounds(min, max)); |
689 | 0 | self |
690 | 0 | } |
691 | | |
692 | | /// Add finite assertion (no NaN/Inf in output). |
693 | | #[must_use] |
694 | 0 | pub fn assert_finite(mut self) -> Self { |
695 | 0 | self.assertions.push(ComputeAssertion::finite()); |
696 | 0 | self |
697 | 0 | } |
698 | | |
699 | | /// Set token throughput budget (tokens/second). |
700 | | #[must_use] |
701 | 0 | pub fn budget_tok_per_sec(mut self, tps: f64) -> Self { |
702 | 0 | self.budget = TokenBudget::from_throughput(tps); |
703 | 0 | self |
704 | 0 | } |
705 | | |
706 | | /// Set token latency budget (microseconds/token). |
707 | | #[must_use] |
708 | 0 | pub fn budget_us_per_tok(mut self, us: f64) -> Self { |
709 | 0 | self.budget = TokenBudget::from_latency(us); |
710 | 0 | self |
711 | 0 | } |
712 | | |
713 | | /// Set full budget configuration. |
714 | | #[must_use] |
715 | 0 | pub fn budget(mut self, budget: TokenBudget) -> Self { |
716 | 0 | self.budget = budget; |
717 | 0 | self |
718 | 0 | } |
719 | | |
720 | | /// Set execution backend. |
721 | | #[must_use] |
722 | 0 | pub fn backend(mut self, backend: Backend) -> Self { |
723 | 0 | self.backend = backend; |
724 | 0 | self |
725 | 0 | } |
726 | | |
727 | | /// Enforce budget (fail if exceeded). Default is false (just report). |
728 | | #[must_use] |
729 | 0 | pub fn enforce_budget(mut self, enforce: bool) -> Self { |
730 | 0 | self.enforce_budget = enforce; |
731 | 0 | self |
732 | 0 | } |
733 | | |
734 | | /// Get the brick name (from operation). |
735 | 0 | pub fn name(&self) -> &'static str { |
736 | 0 | self.op.name() |
737 | 0 | } |
738 | | |
739 | | /// Get current budget. |
740 | 0 | pub fn get_budget(&self) -> TokenBudget { |
741 | 0 | self.budget |
742 | 0 | } |
743 | | |
744 | | /// Get current backend. |
745 | 0 | pub fn get_backend(&self) -> Backend { |
746 | 0 | self.backend |
747 | 0 | } |
748 | | |
749 | | /// Get assertions. |
750 | 0 | pub fn get_assertions(&self) -> &[ComputeAssertion] { |
751 | 0 | &self.assertions |
752 | 0 | } |
753 | | |
754 | | /// Run the compute brick with full verification (Jidoka gate). |
755 | 0 | pub fn run(&self, input: Op::Input) -> Result<TokenResult<Op::Output>, BrickError> { |
756 | 0 | let tokens = self.op.tokens(&input); |
757 | | |
758 | | // Execute with timing |
759 | 0 | let start = Instant::now(); |
760 | 0 | let output = self.op.execute(input, self.backend)?; |
761 | 0 | let elapsed_us = start.elapsed().as_secs_f64() * 1_000_000.0; |
762 | | |
763 | | // Calculate metrics |
764 | 0 | let us_per_token = if tokens > 0 { |
765 | 0 | elapsed_us / tokens as f64 |
766 | | } else { |
767 | 0 | elapsed_us |
768 | | }; |
769 | 0 | let tokens_per_sec = if elapsed_us > 0.0 { |
770 | 0 | tokens as f64 * 1_000_000.0 / elapsed_us |
771 | | } else { |
772 | 0 | f64::INFINITY |
773 | | }; |
774 | 0 | let budget_met = self.budget.is_met(us_per_token); |
775 | 0 | let budget_utilization = self.budget.utilization(us_per_token); |
776 | | |
777 | | // Check budget enforcement |
778 | 0 | if self.enforce_budget && !budget_met { |
779 | 0 | return Err(BrickError::BudgetExceeded { |
780 | 0 | limit_us: self.budget.us_per_token, |
781 | 0 | actual_us: us_per_token, |
782 | 0 | utilization: budget_utilization * 100.0, |
783 | 0 | }); |
784 | 0 | } |
785 | | |
786 | 0 | Ok(TokenResult { |
787 | 0 | output, |
788 | 0 | tokens_processed: tokens, |
789 | 0 | us_per_token, |
790 | 0 | tokens_per_sec, |
791 | 0 | budget_met, |
792 | 0 | budget_utilization, |
793 | 0 | }) |
794 | 0 | } |
795 | | |
796 | | /// Verify assertions without full execution. |
797 | | /// Returns verification status. |
798 | 0 | pub fn verify(&self) -> BrickVerification { |
799 | 0 | let start = Instant::now(); |
800 | | |
801 | | // Check if we have assertions (Popperian requirement) |
802 | 0 | if self.assertions.is_empty() { |
803 | 0 | return BrickVerification { |
804 | 0 | passed: false, |
805 | 0 | assertion_results: vec![AssertionResult { |
806 | 0 | assertion: ComputeAssertion::Custom { |
807 | 0 | name: "popperian_falsifiability".to_string(), |
808 | 0 | }, |
809 | 0 | passed: false, |
810 | 0 | error: Some("No assertions defined - violates Popperian falsifiability".to_string()), |
811 | 0 | }], |
812 | 0 | verification_us: start.elapsed().as_secs_f64() * 1_000_000.0, |
813 | 0 | }; |
814 | 0 | } |
815 | | |
816 | | // For now, just validate assertion structure |
817 | | // Full verification requires input data |
818 | 0 | let results: Vec<AssertionResult> = self |
819 | 0 | .assertions |
820 | 0 | .iter() |
821 | 0 | .map(|a| AssertionResult { |
822 | 0 | assertion: a.clone(), |
823 | | passed: true, |
824 | 0 | error: None, |
825 | 0 | }) |
826 | 0 | .collect(); |
827 | | |
828 | 0 | let passed = results.iter().all(|r| r.passed); |
829 | | |
830 | 0 | BrickVerification { |
831 | 0 | passed, |
832 | 0 | assertion_results: results, |
833 | 0 | verification_us: start.elapsed().as_secs_f64() * 1_000_000.0, |
834 | 0 | } |
835 | 0 | } |
836 | | } |
837 | | |
838 | | impl<Op: ComputeOp + Clone> Clone for ComputeBrick<Op> { |
839 | 0 | fn clone(&self) -> Self { |
840 | 0 | Self { |
841 | 0 | op: self.op.clone(), |
842 | 0 | assertions: self.assertions.clone(), |
843 | 0 | budget: self.budget, |
844 | 0 | backend: self.backend, |
845 | 0 | enforce_budget: self.enforce_budget, |
846 | 0 | _phantom: PhantomData, |
847 | 0 | } |
848 | 0 | } |
849 | | } |
850 | | |
851 | | impl<Op: ComputeOp> fmt::Debug for ComputeBrick<Op> { |
852 | 0 | fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { |
853 | 0 | f.debug_struct("ComputeBrick") |
854 | 0 | .field("name", &self.op.name()) |
855 | 0 | .field("backend", &self.backend) |
856 | 0 | .field("budget", &self.budget) |
857 | 0 | .field("assertions", &self.assertions.len()) |
858 | 0 | .field("enforce_budget", &self.enforce_budget) |
859 | 0 | .finish() |
860 | 0 | } |
861 | | } |
862 | | |
863 | | // ============================================================================ |
864 | | // LLM Transformer Fused Operations (PMAT-PERF-009) |
865 | | // BrickLayer: Compose multiple bricks |
866 | | // ============================================================================ |
867 | | |
868 | | /// A layer of compute bricks that execute sequentially. |
869 | | /// Throughput ceiling = min(component throughputs). |
870 | | #[derive(Debug, Default)] |
871 | | pub struct BrickLayer { |
872 | | /// Named bricks in this layer |
873 | | bricks: Vec<(String, f64)>, // (name, budget_tok_per_sec) |
874 | | } |
875 | | |
876 | | impl BrickLayer { |
877 | | /// Create a new empty layer. |
878 | 0 | pub fn new() -> Self { |
879 | 0 | Self::default() |
880 | 0 | } |
881 | | |
882 | | /// Add a brick to the layer. |
883 | | #[must_use] |
884 | 0 | pub fn with_brick<Op: ComputeOp>(mut self, brick: &ComputeBrick<Op>) -> Self { |
885 | 0 | self.bricks |
886 | 0 | .push((brick.name().to_string(), brick.budget.tokens_per_sec)); |
887 | 0 | self |
888 | 0 | } |
889 | | |
890 | | /// Add a named entry with throughput budget. |
891 | | #[must_use] |
892 | 0 | pub fn with_named(mut self, name: &str, budget_tok_per_sec: f64) -> Self { |
893 | 0 | self.bricks.push((name.to_string(), budget_tok_per_sec)); |
894 | 0 | self |
895 | 0 | } |
896 | | |
897 | | /// Get the throughput ceiling (bottleneck). |
898 | | /// Layer throughput = min(component throughputs). |
899 | 0 | pub fn throughput_ceiling(&self) -> f64 { |
900 | 0 | self.bricks |
901 | 0 | .iter() |
902 | 0 | .map(|(_, tps)| *tps) |
903 | 0 | .fold(f64::INFINITY, f64::min) |
904 | 0 | } |
905 | | |
906 | | /// Get the bottleneck brick name. |
907 | 0 | pub fn bottleneck(&self) -> Option<&str> { |
908 | 0 | self.bricks |
909 | 0 | .iter() |
910 | 0 | .min_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal)) |
911 | 0 | .map(|(name, _)| name.as_str()) |
912 | 0 | } |
913 | | |
914 | | /// Get all bricks with their budgets. |
915 | 0 | pub fn bricks(&self) -> &[(String, f64)] { |
916 | 0 | &self.bricks |
917 | 0 | } |
918 | | } |