trueno/lib.rs
1// ============================================================================
2// Development-phase lint allows - to be addressed incrementally
3// ============================================================================
4// Allow manual_div_ceil - clearer for block calculations
5#![allow(clippy::manual_div_ceil)]
6// Allow manual_is_multiple_of - clearer alignment checks
7#![allow(clippy::manual_is_multiple_of)]
8// Allow needless_range_loop - index access is clearer in some SIMD algorithms
9#![allow(clippy::needless_range_loop)]
10// Allow empty line after doc comments - formatting preference
11#![allow(clippy::empty_line_after_doc_comments)]
12// Allow similar names - semantic distinction is clear
13#![allow(clippy::similar_names)]
14// Allow many single char names - standard math/matrix notation
15#![allow(clippy::many_single_char_names)]
16// Allow too many arguments - SIMD/compute APIs require many parameters
17#![allow(clippy::too_many_arguments)]
18// Allow type complexity - complex SIMD types
19#![allow(clippy::type_complexity)]
20// Allow macro metavars in unsafe - necessary for SIMD dispatch macros
21#![allow(clippy::macro_metavars_in_unsafe)]
22// Allow missing panics doc - will be added incrementally
23#![allow(clippy::missing_panics_doc)]
24// Allow uninit_vec - intentional pattern for perf-critical paths where
25// every element is SET (not accumulated) before any read. Each use has
26// a SAFETY comment documenting the write-before-read invariant.
27#![allow(clippy::uninit_vec)]
28// Allow missing errors doc - will be added incrementally
29#![allow(clippy::missing_errors_doc)]
30// Allow missing safety doc - will be added incrementally
31#![allow(clippy::missing_safety_doc)]
32// Allow excessive precision - SIMD math constants need specific precision
33#![allow(clippy::excessive_precision)]
34// Allow unnecessary cast - clearer type annotations in some cases
35#![allow(clippy::unnecessary_cast)]
36// Allow cast_possible_truncation - handled in SIMD code
37#![allow(clippy::cast_possible_truncation)]
38// Allow cast_sign_loss - handled in SIMD code
39#![allow(clippy::cast_sign_loss)]
40// Allow cast_precision_loss - handled in SIMD code
41#![allow(clippy::cast_precision_loss)]
42// Allow large stack arrays - SIMD/GPU test data and proptest expansions
43#![allow(clippy::large_stack_arrays)]
44// Allow unwrap/float_cmp in test code — safe in assertions, banned in production
45#![cfg_attr(test, allow(clippy::disallowed_methods, clippy::float_cmp))]
46
47//! Trueno: Multi-Target High-Performance Compute Library
48//!
49//! **Trueno** (Spanish: "thunder") provides unified, high-performance compute primitives
50//! across three execution targets:
51//!
52//! 1. **CPU SIMD** - x86 (SSE2/AVX/AVX2/AVX-512), ARM (NEON), WASM (SIMD128)
53//! 2. **GPU** - Vulkan/Metal/DX12/WebGPU via `wgpu`
54//! 3. **WebAssembly** - Portable SIMD128 for browser/edge deployment
55//!
56//! # Design Principles
57//!
58//! - **Write once, optimize everywhere**: Single algorithm, multiple backends
59//! - **Runtime dispatch**: Auto-select best implementation based on CPU features
60//! - **Zero unsafe in public API**: Safety via type system, `unsafe` isolated in backends
61//! - **Benchmarked performance**: Every optimization must prove ≥10% speedup
62//! - **Extreme TDD**: >90% test coverage, mutation testing, property-based tests
63//!
64//! # Quick Start
65//!
66//! ```rust
67//! use trueno::Vector;
68//!
69//! let a = Vector::from_slice(&[1.0, 2.0, 3.0, 4.0]);
70//! let b = Vector::from_slice(&[5.0, 6.0, 7.0, 8.0]);
71//!
72//! // Auto-selects best backend (AVX2/GPU/WASM)
73//! let result = a.add(&b).unwrap();
74//! assert_eq!(result.as_slice(), &[6.0, 8.0, 10.0, 12.0]);
75//! ```
76
77// Contract assertions from YAML (pv codegen)
78#[macro_use]
79#[allow(unused_macros)]
80mod generated_contracts;
81
82// Fallback macros for contracts not yet in codegen
83macro_rules! contract_pre_add { () => {{}}; ($($x:expr),+ $(,)?) => {{ $(let _ = &$x;)+ }}; }
84macro_rules! contract_pre_gemv { () => {{}}; ($($x:expr),+ $(,)?) => {{ $(let _ = &$x;)+ }}; }
85
86pub mod activations;
87pub mod backends;
88pub mod blis;
89pub mod brick;
90pub mod chaos;
91pub mod contracts;
92pub mod eigen;
93pub mod error;
94pub mod hardware;
95pub mod hash;
96pub mod inference;
97pub mod matrix;
98pub mod monitor;
99/// Backend discovery: probe → enumerate → print (PP-066 R-0a).
100pub mod registry;
101pub mod simulation;
102pub mod tiling;
103pub mod tuner;
104pub mod vector;
105
106// Canonical scalar activation functions (UCBD §4, trueno #103)
107pub use activations::{
108 f16_to_f32, f32_to_f16, gelu_scalar, relu_scalar, sigmoid_scalar, silu_scalar, tanh_scalar,
109};
110pub use eigen::SymmetricEigen;
111pub use error::{Result, TruenoError};
112pub use hash::{hash_bytes, hash_key, hash_keys_batch, hash_keys_batch_with_backend};
113pub use matrix::Matrix;
114pub use monitor::{
115 cuda_monitor_available, GpuBackend, GpuClockMetrics, GpuDeviceInfo, GpuMemoryMetrics,
116 GpuMetrics, GpuMonitor, GpuPcieMetrics, GpuPowerMetrics, GpuThermalMetrics, GpuUtilization,
117 GpuVendor, MonitorConfig, MonitorError,
118};
119#[cfg(feature = "cuda-monitor")]
120pub use monitor::{enumerate_cuda_devices, query_cuda_device_info, query_cuda_memory};
121pub use vector::Vector;
122
123// ComputeBrick exports
124pub use brick::{
125 fnv1a_f32_checksum,
126 AddOp,
127 AssertionResult,
128 AttentionOp,
129 // QUANT-Q5K: Q5_K and Q6_K quantization formats (llama.cpp compatible)
130 BlockQ5K,
131 BlockQ6K,
132 BrickBottleneck,
133 BrickCategory,
134 BrickError,
135 // PAR-200: BrickProfiler v2 types
136 BrickId,
137 BrickIdTimer,
138 BrickLayer,
139 BrickProfiler,
140 BrickSample,
141 BrickStats,
142 BrickTimer,
143 BrickVerification,
144 ByteBudget,
145 CategoryStats,
146 ComputeAssertion,
147 ComputeBackend,
148 ComputeBrick,
149 ComputeOp,
150 DivergenceInfo,
151 DotOp,
152 DotQ5KOp,
153 DotQ6KOp,
154 EdgeType,
155 ExecutionEdge,
156 ExecutionGraph,
157 ExecutionNode,
158 // PAR-201: Execution path graph types
159 ExecutionNodeId,
160 FusedGateUpOp,
161 FusedGateUpWeights,
162 FusedQKVOp,
163 FusedQKVWeights,
164 // CORRECTNESS-011: Divergence detection types
165 KernelChecksum,
166 MatmulOp,
167 PtxRegistry,
168 SoftmaxOp,
169 SyncMode,
170 // TILING-SPEC-001: Tile-level profiling types
171 TileLevel,
172 TileStats,
173 TileTimer,
174 TokenBudget,
175 TokenResult,
176};
177
178// Hardware capability exports (PMAT-447)
179pub use hardware::{
180 default_hardware_path, Bottleneck, CpuCapability, GpuBackend as HardwareGpuBackend,
181 GpuCapability, HardwareCapability, RooflineParams, SimdWidth,
182};
183
184// ML Tuner exports (T-TUNER-003 through T-TUNER-007, GH#80-84)
185pub use tuner::{
186 BottleneckClass, BottleneckPrediction, BrickTuner, ConceptDriftStatus, ExperimentSuggestion,
187 FeatureExtractor, KernelClassifier, KernelRecommendation, KernelType, QuantType, RunConfig,
188 ThroughputPrediction, ThroughputRegressor, TrainingSample, TrainingStats, TunerDataCollector,
189 TunerError, TunerFeatures, TunerRecommendation, UserFeedback,
190};
191
192// Tiling Compute Blocks exports (TILING-SPEC-001)
193pub use tiling::{
194 optimal_prefetch_distance, pack_a_index, pack_b_index, swizzle_index, PackingLayout,
195 PrefetchLocality, TcbGeometry, TcbIndexCalculator, TcbLevel, TiledQ4KMatvec, TilingBackend,
196 TilingConfig, TilingError, TilingStats, Q4K_SUPERBLOCK_BYTES, Q4K_SUPERBLOCK_SIZE,
197};
198
199/// Backend execution target
200#[derive(Debug, Clone, Copy, PartialEq, Eq)]
201pub enum Backend {
202 /// Scalar fallback (no SIMD)
203 Scalar,
204 /// SSE2 (x86_64 baseline)
205 SSE2,
206 /// AVX (256-bit)
207 AVX,
208 /// AVX2 (256-bit with FMA)
209 AVX2,
210 /// AVX-512 (512-bit)
211 AVX512,
212 /// ARM NEON
213 NEON,
214 /// WebAssembly SIMD128
215 WasmSIMD,
216 /// GPU compute (wgpu)
217 GPU,
218 /// Auto-select best available
219 Auto,
220}
221
222impl Backend {
223 /// Select the best available backend for the current platform
224 ///
225 /// This is a convenience wrapper around `select_best_available_backend()`
226 pub fn select_best() -> Self {
227 select_best_available_backend()
228 }
229}
230
231/// Operation complexity for GPU dispatch eligibility
232#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
233pub enum OpComplexity {
234 /// Simple operations (add, mul) - prefer SIMD unless very large
235 Low = 0,
236 /// Moderate operations (dot, reduce) - GPU beneficial at 100K+
237 Medium = 1,
238 /// Complex operations (matmul, convolution) - GPU beneficial at 10K+
239 High = 2,
240}
241
242/// Operation type for SIMD backend selection
243///
244/// Based on AVX-512 performance analysis (see AVX512_ANALYSIS.md), operations are
245/// categorized by their memory vs compute characteristics to guide optimal backend selection.
246#[derive(Debug, Clone, Copy, PartialEq, Eq)]
247pub enum OperationType {
248 /// Memory-bound operations (add, sub, mul, scale, div)
249 ///
250 /// These operations perform minimal computation per memory access (arithmetic intensity < 1 op/byte).
251 /// Prefer AVX2 over AVX-512 due to memory bandwidth bottleneck.
252 ///
253 /// AVX-512 performance: 0.67-1.20x scalar (often slower!)
254 /// AVX2 performance: 1.0-1.2x scalar
255 MemoryBound,
256
257 /// Compute-bound operations (dot, max, min, argmax, argmin)
258 ///
259 /// These operations perform significant computation per memory access (arithmetic intensity > 1 op/byte).
260 /// AVX-512 excels due to wider SIMD parallelism.
261 ///
262 /// AVX-512 performance: 7-14x scalar (validated)
263 /// AVX2 performance: 4-12x scalar (validated)
264 ComputeBound,
265
266 /// Mixed operations (fma, sqrt, exp, sigmoid, activations)
267 ///
268 /// Performance depends on data size and hardware.
269 /// Use size-based heuristics or default to AVX2 for safety.
270 Mixed,
271}
272
273/// Detect best SIMD backend for x86/x86_64 platforms
274///
275/// **IMPORTANT**: Prefers AVX2 over AVX-512 by default based on performance analysis.
276///
277/// AVX-512 is **NOT** universally faster - it causes 10-33% slowdown for memory-bound
278/// operations (add, mul, sub) due to memory bandwidth bottleneck and thermal throttling.
279/// See AVX512_ANALYSIS.md for detailed benchmarking results.
280///
281/// For operation-specific backend selection, use `select_backend_for_operation()`.
282#[cfg(any(target_arch = "x86_64", target_arch = "x86"))]
283fn detect_x86_backend() -> Backend {
284 // Prefer AVX2 over AVX-512 for safety (AVX-512 causes regressions for memory-bound ops)
285 if is_x86_feature_detected!("avx2") && is_x86_feature_detected!("fma") {
286 return Backend::AVX2;
287 }
288 // Note: AVX-512 is intentionally NOT checked here
289 // Use select_backend_for_operation(OperationType::ComputeBound) for AVX-512
290 if is_x86_feature_detected!("avx") {
291 return Backend::AVX;
292 }
293 if is_x86_feature_detected!("sse2") {
294 return Backend::SSE2;
295 }
296 Backend::Scalar
297}
298
299/// Detect best SIMD backend for ARM platforms
300#[cfg(any(target_arch = "aarch64", target_arch = "arm"))]
301fn detect_arm_backend() -> Backend {
302 #[cfg(target_feature = "neon")]
303 {
304 Backend::NEON
305 }
306 #[cfg(not(target_feature = "neon"))]
307 {
308 Backend::Scalar
309 }
310}
311
312/// Detect best SIMD backend for WebAssembly
313#[cfg(target_arch = "wasm32")]
314fn detect_wasm_backend() -> Backend {
315 #[cfg(target_feature = "simd128")]
316 {
317 Backend::WasmSIMD
318 }
319 #[cfg(not(target_feature = "simd128"))]
320 {
321 Backend::Scalar
322 }
323}
324
325/// Select the best available backend for the current platform
326///
327/// This function performs runtime CPU feature detection and selects the most
328/// optimized backend available. The selection follows this priority:
329///
330/// **x86/x86_64**:
331/// 1. AVX-512 (if `avx512f` feature detected)
332/// 2. AVX2 (if `avx2` and `fma` features detected)
333/// 3. AVX (if `avx` feature detected)
334/// 4. SSE2 (baseline for x86_64)
335/// 5. Scalar (fallback)
336///
337/// **ARM**:
338/// 1. NEON (if available)
339/// 2. Scalar (fallback)
340///
341/// **WASM**: SIMD128 (if available), else Scalar
342///
343/// **Other platforms**: Scalar
344///
345/// # Returns
346///
347/// The most optimized backend available on this CPU/platform
348///
349/// # Examples
350///
351/// ```
352/// use trueno::select_best_available_backend;
353///
354/// let backend = select_best_available_backend();
355/// println!("Using backend: {:?}", backend);
356/// ```
357pub fn select_best_available_backend() -> Backend {
358 // Cache backend selection using OnceLock to avoid repeated CPU feature detection
359 // This eliminates 3-5% overhead from calling is_x86_feature_detected!() repeatedly
360 static BEST_BACKEND: std::sync::OnceLock<Backend> = std::sync::OnceLock::new();
361
362 *BEST_BACKEND.get_or_init(|| {
363 #[cfg(any(target_arch = "x86_64", target_arch = "x86"))]
364 {
365 detect_x86_backend()
366 }
367
368 #[cfg(any(target_arch = "aarch64", target_arch = "arm"))]
369 {
370 detect_arm_backend()
371 }
372
373 #[cfg(target_arch = "wasm32")]
374 {
375 detect_wasm_backend()
376 }
377
378 #[cfg(not(any(
379 target_arch = "x86_64",
380 target_arch = "x86",
381 target_arch = "aarch64",
382 target_arch = "arm",
383 target_arch = "wasm32"
384 )))]
385 {
386 Backend::Scalar
387 }
388 })
389}
390
391/// Select the optimal backend for a specific operation type
392///
393/// This function considers the memory vs compute characteristics of operations
394/// to select the backend that will provide the best performance. Based on
395/// comprehensive benchmarking (see AVX512_ANALYSIS.md), AVX-512 is avoided
396/// for memory-bound operations where it causes 10-33% performance degradation.
397///
398/// # Operation Classification
399///
400/// - **MemoryBound**: add, sub, mul, div, scale, abs, clamp, lerp, relu
401/// - Prefer AVX2 (1.0-1.2x scalar) over AVX-512 (0.67-1.20x scalar)
402/// - Memory bandwidth bottleneck limits wider SIMD benefit
403///
404/// - **ComputeBound**: dot, max, min, argmax, argmin, norm_l1, norm_l2, norm_linf
405/// - Prefer AVX-512 (7-14x scalar) over AVX2 (4-12x scalar)
406/// - High arithmetic intensity benefits from wider SIMD
407///
408/// - **Mixed**: fma, sqrt, exp, ln, sigmoid, tanh, gelu, swish
409/// - Default to AVX2 for safety (avoids AVX-512 thermal throttling)
410/// - Size-based heuristics could improve this in future
411///
412/// # Backend Selection Priority
413///
414/// **For MemoryBound operations**:
415/// 1. AVX2 (if available) - BEST for memory-bound
416/// 2. SSE2 (x86_64 baseline)
417/// 3. AVX-512 (AVOIDED - causes slowdown)
418/// 4. NEON (ARM)
419/// 5. WASM SIMD128
420/// 6. Scalar (fallback)
421///
422/// **For ComputeBound operations**:
423/// 1. AVX-512 (if available) - BEST for compute-bound
424/// 2. AVX2
425/// 3. SSE2
426/// 4. NEON (ARM)
427/// 5. WASM SIMD128
428/// 6. Scalar (fallback)
429///
430/// # Arguments
431///
432/// * `op_type` - The type of operation being performed
433///
434/// # Returns
435///
436/// The optimal backend for the given operation type
437///
438/// # Examples
439///
440/// ```
441/// use trueno::{select_backend_for_operation, OperationType};
442///
443/// // Memory-bound operation - prefers AVX2 over AVX-512
444/// let backend = select_backend_for_operation(OperationType::MemoryBound);
445///
446/// // Compute-bound operation - uses AVX-512 if available
447/// let backend = select_backend_for_operation(OperationType::ComputeBound);
448/// ```
449///
450/// # Performance Impact
451///
452/// Using operation-aware backend selection fixes performance regressions:
453/// - mul with AVX-512: 0.67x → 1.0x (use AVX2 instead)
454/// - sub with AVX-512: 0.87x → 1.0x (use AVX2 instead)
455/// - dot with AVX-512: 7.89x (keep AVX-512)
456pub fn select_backend_for_operation(op_type: OperationType) -> Backend {
457 // Allow unused on non-x86 architectures
458 let _ = &op_type;
459
460 #[cfg(any(target_arch = "x86_64", target_arch = "x86"))]
461 {
462 select_x86_backend_for_operation(op_type)
463 }
464
465 #[cfg(any(target_arch = "aarch64", target_arch = "arm"))]
466 {
467 detect_arm_backend()
468 }
469
470 #[cfg(target_arch = "wasm32")]
471 {
472 detect_wasm_backend()
473 }
474
475 #[cfg(not(any(
476 target_arch = "x86_64",
477 target_arch = "x86",
478 target_arch = "aarch64",
479 target_arch = "arm",
480 target_arch = "wasm32"
481 )))]
482 {
483 Backend::Scalar
484 }
485}
486
487/// Select the best x86 backend based on operation type and available features.
488///
489/// Separated from `select_backend_for_operation` to reduce cyclomatic complexity.
490#[cfg(any(target_arch = "x86_64", target_arch = "x86"))]
491fn select_x86_backend_for_operation(op_type: OperationType) -> Backend {
492 use std::arch::is_x86_feature_detected;
493
494 // Check for AVX-512 (only for compute-bound operations)
495 let use_avx512 = op_type == OperationType::ComputeBound && is_x86_feature_detected!("avx512f");
496 if use_avx512 {
497 return Backend::AVX512;
498 }
499
500 // AVX2 with FMA is preferred for most operations
501 if is_x86_feature_detected!("avx2") && is_x86_feature_detected!("fma") {
502 return Backend::AVX2;
503 }
504
505 // Fallback chain: AVX -> SSE2 -> Scalar
506 if is_x86_feature_detected!("avx") {
507 return Backend::AVX;
508 }
509 if is_x86_feature_detected!("sse2") {
510 return Backend::SSE2;
511 }
512
513 Backend::Scalar
514}
515
516#[cfg(test)]
517mod contract_tests;
518
519#[cfg(test)]
520mod contract_tests_image;
521
522#[cfg(test)]
523mod contract_tests_linalg;
524
525#[cfg(test)]
526mod tests;