/home/noah/src/trueno/src/brick/ops.rs
Line | Count | Source |
1 | | //! Built-in Compute Operations |
2 | | //! |
3 | | //! Pre-defined operations that implement the ComputeOp trait: |
4 | | //! - DotOp: Vector dot product |
5 | | //! - AddOp: Element-wise vector addition |
6 | | //! - MatmulOp: Matrix multiplication (SIMD-optimized) |
7 | | //! - SoftmaxOp: Softmax with SIMD exp approximation (SIMD-EXP) |
8 | | |
9 | | use super::{Backend, ComputeOp}; |
10 | | use crate::error::TruenoError; |
11 | | |
12 | | // ============================================================================ |
13 | | // DotOp: Dot Product |
14 | | // ============================================================================ |
15 | | |
16 | | /// Dot product operation. |
17 | | #[derive(Debug, Clone)] |
18 | | pub struct DotOp { |
19 | | /// Expected vector length |
20 | | pub len: usize, |
21 | | } |
22 | | |
23 | | impl DotOp { |
24 | 0 | pub fn new(len: usize) -> Self { |
25 | 0 | Self { len } |
26 | 0 | } |
27 | | } |
28 | | |
29 | | impl ComputeOp for DotOp { |
30 | | type Input = (Vec<f32>, Vec<f32>); |
31 | | type Output = f32; |
32 | | |
33 | 0 | fn name(&self) -> &'static str { |
34 | 0 | "dot" |
35 | 0 | } |
36 | | |
37 | 0 | fn execute(&self, input: Self::Input, _backend: Backend) -> Result<Self::Output, TruenoError> { |
38 | 0 | let (a, b) = input; |
39 | 0 | if a.len() != b.len() { |
40 | 0 | return Err(TruenoError::SizeMismatch { |
41 | 0 | expected: a.len(), |
42 | 0 | actual: b.len(), |
43 | 0 | }); |
44 | 0 | } |
45 | | // Simple scalar implementation for now |
46 | 0 | let sum: f32 = a.iter().zip(b.iter()).map(|(x, y)| x * y).sum(); |
47 | 0 | Ok(sum) |
48 | 0 | } |
49 | | |
50 | 0 | fn tokens(&self, input: &Self::Input) -> usize { |
51 | | // Each element pair is roughly 1 "token" of work |
52 | 0 | input.0.len() |
53 | 0 | } |
54 | | } |
55 | | |
56 | | // ============================================================================ |
57 | | // AddOp: Element-wise Addition |
58 | | // ============================================================================ |
59 | | |
60 | | /// Element-wise add operation. |
61 | | #[derive(Debug, Clone)] |
62 | | pub struct AddOp { |
63 | | /// Expected vector length |
64 | | pub len: usize, |
65 | | } |
66 | | |
67 | | impl AddOp { |
68 | 0 | pub fn new(len: usize) -> Self { |
69 | 0 | Self { len } |
70 | 0 | } |
71 | | } |
72 | | |
73 | | impl ComputeOp for AddOp { |
74 | | type Input = (Vec<f32>, Vec<f32>); |
75 | | type Output = Vec<f32>; |
76 | | |
77 | 0 | fn name(&self) -> &'static str { |
78 | 0 | "add" |
79 | 0 | } |
80 | | |
81 | 0 | fn execute(&self, input: Self::Input, _backend: Backend) -> Result<Self::Output, TruenoError> { |
82 | 0 | let (a, b) = input; |
83 | 0 | if a.len() != b.len() { |
84 | 0 | return Err(TruenoError::SizeMismatch { |
85 | 0 | expected: a.len(), |
86 | 0 | actual: b.len(), |
87 | 0 | }); |
88 | 0 | } |
89 | 0 | Ok(a.iter().zip(b.iter()).map(|(x, y)| x + y).collect()) |
90 | 0 | } |
91 | | |
92 | 0 | fn tokens(&self, input: &Self::Input) -> usize { |
93 | 0 | input.0.len() |
94 | 0 | } |
95 | | } |
96 | | |
97 | | // ============================================================================ |
98 | | // MatmulOp: Matrix Multiplication |
99 | | // ============================================================================ |
100 | | |
101 | | /// Matrix multiplication operation. |
102 | | #[derive(Debug, Clone)] |
103 | | pub struct MatmulOp { |
104 | | /// M dimension (rows of A) |
105 | | pub m: usize, |
106 | | /// K dimension (cols of A = rows of B) |
107 | | pub k: usize, |
108 | | /// N dimension (cols of B) |
109 | | pub n: usize, |
110 | | } |
111 | | |
112 | | impl MatmulOp { |
113 | 0 | pub fn new(m: usize, k: usize, n: usize) -> Self { |
114 | 0 | Self { m, k, n } |
115 | 0 | } |
116 | | } |
117 | | |
118 | | impl ComputeOp for MatmulOp { |
119 | | type Input = (Vec<f32>, Vec<f32>); |
120 | | type Output = Vec<f32>; |
121 | | |
122 | 0 | fn name(&self) -> &'static str { |
123 | 0 | "matmul" |
124 | 0 | } |
125 | | |
126 | 0 | fn execute(&self, input: Self::Input, _backend: Backend) -> Result<Self::Output, TruenoError> { |
127 | 0 | let (a, b) = input; |
128 | 0 | let expected_a = self.m * self.k; |
129 | 0 | let expected_b = self.k * self.n; |
130 | | |
131 | 0 | if a.len() != expected_a { |
132 | 0 | return Err(TruenoError::SizeMismatch { |
133 | 0 | expected: expected_a, |
134 | 0 | actual: a.len(), |
135 | 0 | }); |
136 | 0 | } |
137 | 0 | if b.len() != expected_b { |
138 | 0 | return Err(TruenoError::SizeMismatch { |
139 | 0 | expected: expected_b, |
140 | 0 | actual: b.len(), |
141 | 0 | }); |
142 | 0 | } |
143 | | |
144 | | // SIMD-optimized matrix multiplication via Matrix type |
145 | | // Uses AVX2/AVX-512 with cache blocking for ~10-50x speedup |
146 | 0 | let simd_backend = crate::Backend::select_best(); |
147 | 0 | let mat_a = crate::Matrix::from_vec_with_backend(self.m, self.k, a, simd_backend); |
148 | 0 | let mat_b = crate::Matrix::from_vec_with_backend(self.k, self.n, b, simd_backend); |
149 | | |
150 | 0 | let result = mat_a.matmul(&mat_b)?; |
151 | 0 | Ok(result.as_slice().to_vec()) |
152 | 0 | } |
153 | | |
154 | 0 | fn tokens(&self, _input: &Self::Input) -> usize { |
155 | | // For matmul, "tokens" = number of output elements |
156 | | // Each output requires K multiply-adds |
157 | 0 | self.m * self.n |
158 | 0 | } |
159 | | } |
160 | | |
161 | | // ============================================================================ |
162 | | // SoftmaxOp: Softmax with SIMD Exp (SIMD-EXP) |
163 | | // ============================================================================ |
164 | | |
165 | | /// Softmax operation. |
166 | | #[derive(Debug, Clone)] |
167 | | pub struct SoftmaxOp { |
168 | | /// Expected vector length |
169 | | pub len: usize, |
170 | | } |
171 | | |
172 | | impl SoftmaxOp { |
173 | 0 | pub fn new(len: usize) -> Self { |
174 | 0 | Self { len } |
175 | 0 | } |
176 | | } |
177 | | |
178 | | impl ComputeOp for SoftmaxOp { |
179 | | type Input = Vec<f32>; |
180 | | type Output = Vec<f32>; |
181 | | |
182 | 0 | fn name(&self) -> &'static str { |
183 | 0 | "softmax" |
184 | 0 | } |
185 | | |
186 | 0 | fn execute(&self, input: Self::Input, backend: Backend) -> Result<Self::Output, TruenoError> { |
187 | 0 | if input.is_empty() { |
188 | 0 | return Ok(vec![]); |
189 | 0 | } |
190 | | |
191 | | // SIMD-EXP: Use SIMD backends for 2-3x speedup on softmax |
192 | | // The exp() is the bottleneck in softmax - SIMD polynomial approximation |
193 | | // matches llama.cpp's ggml_v_expf performance. |
194 | | |
195 | | // Step 1: Find max for numerical stability (SIMD max) |
196 | 0 | let max = Self::simd_max(&input, backend); |
197 | | |
198 | | // Step 2: Subtract max and compute exp (SIMD exp) |
199 | 0 | let shifted: Vec<f32> = input.iter().map(|x| x - max).collect(); |
200 | 0 | let mut exp_vals = vec![0.0f32; shifted.len()]; |
201 | 0 | Self::simd_exp(&shifted, &mut exp_vals, backend); |
202 | | |
203 | | // Step 3: Sum (SIMD sum) |
204 | 0 | let exp_sum = Self::simd_sum(&exp_vals, backend); |
205 | | |
206 | | // Step 4: Normalize (SIMD scale) |
207 | 0 | let inv_sum = 1.0 / exp_sum; |
208 | 0 | let mut result = vec![0.0f32; exp_vals.len()]; |
209 | 0 | Self::simd_scale(&exp_vals, inv_sum, &mut result, backend); |
210 | | |
211 | 0 | Ok(result) |
212 | 0 | } |
213 | | |
214 | 0 | fn tokens(&self, input: &Self::Input) -> usize { |
215 | 0 | input.len() |
216 | 0 | } |
217 | | } |
218 | | |
219 | | impl SoftmaxOp { |
220 | | /// Check if backend supports SIMD acceleration |
221 | | #[inline] |
222 | 0 | pub fn is_simd_backend(backend: Backend) -> bool { |
223 | 0 | matches!( |
224 | 0 | backend, |
225 | | Backend::Avx2 | Backend::Avx512 | Backend::Sse2 | Backend::Neon | Backend::Auto |
226 | | ) |
227 | 0 | } |
228 | | |
229 | | /// SIMD-accelerated max reduction |
230 | | #[inline] |
231 | 0 | fn simd_max(input: &[f32], backend: Backend) -> f32 { |
232 | | #[cfg(target_arch = "x86_64")] |
233 | | { |
234 | 0 | if Self::is_simd_backend(backend) && is_x86_feature_detected!("avx2") { |
235 | 0 | return unsafe { Self::avx2_max(input) }; |
236 | 0 | } |
237 | | } |
238 | 0 | let _ = backend; // suppress warning on non-x86 |
239 | | // Scalar fallback |
240 | 0 | input.iter().cloned().fold(f32::NEG_INFINITY, f32::max) |
241 | 0 | } |
242 | | |
243 | | /// SIMD-accelerated exp using polynomial approximation (SIMD-EXP) |
244 | | /// |
245 | | /// Uses 6th-degree Remez minimax polynomial matching llama.cpp's ggml_v_expf. |
246 | | /// Range reduction: exp(x) = 2^k * e^r where r in [-ln(2)/2, ln(2)/2] |
247 | | #[inline] |
248 | 0 | fn simd_exp(input: &[f32], output: &mut [f32], backend: Backend) { |
249 | | #[cfg(target_arch = "x86_64")] |
250 | | { |
251 | 0 | if Self::is_simd_backend(backend) && is_x86_feature_detected!("avx2") { |
252 | 0 | unsafe { Self::avx2_exp(input, output) }; |
253 | 0 | return; |
254 | 0 | } |
255 | | } |
256 | 0 | let _ = backend; // suppress warning on non-x86 |
257 | | // Scalar fallback |
258 | 0 | for (i, &x) in input.iter().enumerate() { |
259 | 0 | output[i] = x.exp(); |
260 | 0 | } |
261 | 0 | } |
262 | | |
263 | | /// SIMD-accelerated sum reduction |
264 | | #[inline] |
265 | 0 | fn simd_sum(input: &[f32], backend: Backend) -> f32 { |
266 | | #[cfg(target_arch = "x86_64")] |
267 | | { |
268 | 0 | if Self::is_simd_backend(backend) && is_x86_feature_detected!("avx2") { |
269 | 0 | return unsafe { Self::avx2_sum(input) }; |
270 | 0 | } |
271 | | } |
272 | 0 | let _ = backend; // suppress warning on non-x86 |
273 | | // Scalar fallback |
274 | 0 | input.iter().sum() |
275 | 0 | } |
276 | | |
277 | | /// SIMD-accelerated scale |
278 | | #[inline] |
279 | 0 | fn simd_scale(input: &[f32], scalar: f32, output: &mut [f32], backend: Backend) { |
280 | | #[cfg(target_arch = "x86_64")] |
281 | | { |
282 | 0 | if Self::is_simd_backend(backend) && is_x86_feature_detected!("avx2") { |
283 | 0 | unsafe { Self::avx2_scale(input, scalar, output) }; |
284 | 0 | return; |
285 | 0 | } |
286 | | } |
287 | 0 | let _ = backend; // suppress warning on non-x86 |
288 | | // Scalar fallback |
289 | 0 | for (i, &x) in input.iter().enumerate() { |
290 | 0 | output[i] = x * scalar; |
291 | 0 | } |
292 | 0 | } |
293 | | |
294 | | // AVX2 implementations |
295 | | |
296 | | #[cfg(target_arch = "x86_64")] |
297 | | #[target_feature(enable = "avx2")] |
298 | 0 | unsafe fn avx2_max(input: &[f32]) -> f32 { |
299 | | use std::arch::x86_64::*; |
300 | 0 | let len = input.len(); |
301 | 0 | let mut i = 0; |
302 | 0 | let mut vmax = _mm256_set1_ps(f32::NEG_INFINITY); |
303 | | |
304 | 0 | while i + 8 <= len { |
305 | 0 | let v = _mm256_loadu_ps(input.as_ptr().add(i)); |
306 | 0 | vmax = _mm256_max_ps(vmax, v); |
307 | 0 | i += 8; |
308 | 0 | } |
309 | | |
310 | | // Horizontal max |
311 | 0 | let high = _mm256_extractf128_ps(vmax, 1); |
312 | 0 | let low = _mm256_castps256_ps128(vmax); |
313 | 0 | let max128 = _mm_max_ps(high, low); |
314 | 0 | let max64 = _mm_max_ps(max128, _mm_movehl_ps(max128, max128)); |
315 | 0 | let max32 = _mm_max_ss(max64, _mm_shuffle_ps(max64, max64, 1)); |
316 | 0 | let mut result = _mm_cvtss_f32(max32); |
317 | | |
318 | | // Handle remainder |
319 | 0 | for &val in &input[i..] { |
320 | 0 | result = result.max(val); |
321 | 0 | } |
322 | 0 | result |
323 | 0 | } |
324 | | |
325 | | #[cfg(target_arch = "x86_64")] |
326 | | #[target_feature(enable = "avx2", enable = "fma")] |
327 | 0 | unsafe fn avx2_exp(input: &[f32], output: &mut [f32]) { |
328 | | use std::arch::x86_64::*; |
329 | | |
330 | 0 | let len = input.len(); |
331 | 0 | let mut i = 0; |
332 | | |
333 | | // Constants for range reduction (matches llama.cpp ggml_v_expf) |
334 | 0 | let log2e = _mm256_set1_ps(std::f32::consts::LOG2_E); |
335 | 0 | let ln2 = _mm256_set1_ps(std::f32::consts::LN_2); |
336 | 0 | let half = _mm256_set1_ps(0.5); |
337 | 0 | let one = _mm256_set1_ps(1.0); |
338 | | |
339 | | // Remez minimax polynomial coefficients for e^r on [-ln(2)/2, ln(2)/2] |
340 | 0 | let c1 = _mm256_set1_ps(1.0); |
341 | 0 | let c2 = _mm256_set1_ps(0.5); |
342 | 0 | let c3 = _mm256_set1_ps(0.166_666_67); |
343 | 0 | let c4 = _mm256_set1_ps(0.041_666_668); |
344 | 0 | let c5 = _mm256_set1_ps(0.008_333_334); |
345 | 0 | let c6 = _mm256_set1_ps(0.001_388_889); |
346 | | |
347 | 0 | let exp_hi = _mm256_set1_ps(88.376_26); |
348 | 0 | let exp_lo = _mm256_set1_ps(-87.336_55); |
349 | | |
350 | 0 | while i + 8 <= len { |
351 | 0 | let x = _mm256_loadu_ps(input.as_ptr().add(i)); |
352 | 0 | let x = _mm256_max_ps(_mm256_min_ps(x, exp_hi), exp_lo); |
353 | 0 |
|
354 | 0 | // Range reduction: x' = x * log2(e), k = round(x'), r = (x' - k) * ln2 |
355 | 0 | let fx = _mm256_fmadd_ps(x, log2e, half); |
356 | 0 | let fx = _mm256_floor_ps(fx); |
357 | 0 | let r = _mm256_fnmadd_ps(fx, ln2, x); |
358 | 0 |
|
359 | 0 | // Polynomial: e^r ≈ 1 + r + r²/2 + r³/6 + r⁴/24 + r⁵/120 + r⁶/720 |
360 | 0 | // Using Horner's method for efficient evaluation |
361 | 0 | let p = _mm256_fmadd_ps(c6, r, c5); |
362 | 0 | let p = _mm256_fmadd_ps(p, r, c4); |
363 | 0 | let p = _mm256_fmadd_ps(p, r, c3); |
364 | 0 | let p = _mm256_fmadd_ps(p, r, c2); |
365 | 0 | let p = _mm256_fmadd_ps(p, r, c1); |
366 | 0 | let p = _mm256_fmadd_ps(p, r, one); |
367 | 0 |
|
368 | 0 | // Scale by 2^k using integer exponent manipulation |
369 | 0 | let k = _mm256_cvtps_epi32(fx); |
370 | 0 | let k = _mm256_add_epi32(k, _mm256_set1_epi32(127)); |
371 | 0 | let k = _mm256_slli_epi32(k, 23); |
372 | 0 | let pow2k = _mm256_castsi256_ps(k); |
373 | 0 | let result = _mm256_mul_ps(p, pow2k); |
374 | 0 |
|
375 | 0 | _mm256_storeu_ps(output.as_mut_ptr().add(i), result); |
376 | 0 | i += 8; |
377 | 0 | } |
378 | | |
379 | | // Scalar remainder |
380 | 0 | for j in i..len { |
381 | 0 | output[j] = input[j].exp(); |
382 | 0 | } |
383 | 0 | } |
384 | | |
385 | | #[cfg(target_arch = "x86_64")] |
386 | | #[target_feature(enable = "avx2")] |
387 | 0 | unsafe fn avx2_sum(input: &[f32]) -> f32 { |
388 | | use std::arch::x86_64::*; |
389 | 0 | let len = input.len(); |
390 | 0 | let mut i = 0; |
391 | 0 | let mut acc = _mm256_setzero_ps(); |
392 | | |
393 | 0 | while i + 8 <= len { |
394 | 0 | let v = _mm256_loadu_ps(input.as_ptr().add(i)); |
395 | 0 | acc = _mm256_add_ps(acc, v); |
396 | 0 | i += 8; |
397 | 0 | } |
398 | | |
399 | | // Horizontal sum |
400 | 0 | let high = _mm256_extractf128_ps(acc, 1); |
401 | 0 | let low = _mm256_castps256_ps128(acc); |
402 | 0 | let sum128 = _mm_add_ps(high, low); |
403 | 0 | let sum64 = _mm_add_ps(sum128, _mm_movehl_ps(sum128, sum128)); |
404 | 0 | let sum32 = _mm_add_ss(sum64, _mm_shuffle_ps(sum64, sum64, 1)); |
405 | 0 | let mut result = _mm_cvtss_f32(sum32); |
406 | | |
407 | | // Handle remainder |
408 | 0 | for &val in &input[i..] { |
409 | 0 | result += val; |
410 | 0 | } |
411 | 0 | result |
412 | 0 | } |
413 | | |
414 | | #[cfg(target_arch = "x86_64")] |
415 | | #[target_feature(enable = "avx2")] |
416 | 0 | unsafe fn avx2_scale(input: &[f32], scalar: f32, output: &mut [f32]) { |
417 | | use std::arch::x86_64::*; |
418 | 0 | let len = input.len(); |
419 | 0 | let mut i = 0; |
420 | 0 | let vscalar = _mm256_set1_ps(scalar); |
421 | | |
422 | 0 | while i + 8 <= len { |
423 | 0 | let v = _mm256_loadu_ps(input.as_ptr().add(i)); |
424 | 0 | let result = _mm256_mul_ps(v, vscalar); |
425 | 0 | _mm256_storeu_ps(output.as_mut_ptr().add(i), result); |
426 | 0 | i += 8; |
427 | 0 | } |
428 | | |
429 | | // Scalar remainder |
430 | 0 | for j in i..len { |
431 | 0 | output[j] = input[j] * scalar; |
432 | 0 | } |
433 | 0 | } |
434 | | } |
435 | | |
436 | | // ============================================================================ |
437 | | // Tests |
438 | | // ============================================================================ |
439 | | |
440 | | #[cfg(test)] |
441 | | mod tests { |
442 | | use super::*; |
443 | | |
444 | | #[test] |
445 | | fn test_dot_op() { |
446 | | let op = DotOp::new(4); |
447 | | let a = vec![1.0, 2.0, 3.0, 4.0]; |
448 | | let b = vec![5.0, 6.0, 7.0, 8.0]; |
449 | | let result = op.execute((a, b), Backend::Scalar).unwrap(); |
450 | | assert!((result - 70.0).abs() < 0.001); // 1*5 + 2*6 + 3*7 + 4*8 = 70 |
451 | | } |
452 | | |
453 | | #[test] |
454 | | fn test_dot_op_mismatch() { |
455 | | let op = DotOp::new(4); |
456 | | let a = vec![1.0, 2.0, 3.0]; |
457 | | let b = vec![4.0, 5.0, 6.0, 7.0]; |
458 | | assert!(op.execute((a, b), Backend::Scalar).is_err()); |
459 | | } |
460 | | |
461 | | #[test] |
462 | | fn test_add_op() { |
463 | | let op = AddOp::new(4); |
464 | | let a = vec![1.0, 2.0, 3.0, 4.0]; |
465 | | let b = vec![5.0, 6.0, 7.0, 8.0]; |
466 | | let result = op.execute((a, b), Backend::Scalar).unwrap(); |
467 | | assert_eq!(result, vec![6.0, 8.0, 10.0, 12.0]); |
468 | | } |
469 | | |
470 | | #[test] |
471 | | fn test_add_op_mismatch() { |
472 | | let op = AddOp::new(4); |
473 | | let a = vec![1.0, 2.0]; |
474 | | let b = vec![3.0, 4.0, 5.0]; |
475 | | assert!(op.execute((a, b), Backend::Scalar).is_err()); |
476 | | } |
477 | | |
478 | | #[test] |
479 | | fn test_matmul_op() { |
480 | | // 2x3 * 3x2 = 2x2 |
481 | | let op = MatmulOp::new(2, 3, 2); |
482 | | let a = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]; |
483 | | let b = vec![7.0, 8.0, 9.0, 10.0, 11.0, 12.0]; |
484 | | let result = op.execute((a, b), Backend::Scalar).unwrap(); |
485 | | |
486 | | // Expected: [[1*7+2*9+3*11, 1*8+2*10+3*12], [4*7+5*9+6*11, 4*8+5*10+6*12]] |
487 | | // = [[7+18+33, 8+20+36], [28+45+66, 32+50+72]] |
488 | | // = [[58, 64], [139, 154]] |
489 | | assert!((result[0] - 58.0).abs() < 0.001); |
490 | | assert!((result[1] - 64.0).abs() < 0.001); |
491 | | assert!((result[2] - 139.0).abs() < 0.001); |
492 | | assert!((result[3] - 154.0).abs() < 0.001); |
493 | | } |
494 | | |
495 | | #[test] |
496 | | fn test_softmax_op() { |
497 | | let op = SoftmaxOp::new(4); |
498 | | let input = vec![1.0, 2.0, 3.0, 4.0]; |
499 | | let result = op.execute(input, Backend::Scalar).unwrap(); |
500 | | |
501 | | // Check sum = 1 |
502 | | let sum: f32 = result.iter().sum(); |
503 | | assert!((sum - 1.0).abs() < 0.001); |
504 | | |
505 | | // Check monotonicity |
506 | | assert!(result[0] < result[1]); |
507 | | assert!(result[1] < result[2]); |
508 | | assert!(result[2] < result[3]); |
509 | | } |
510 | | |
511 | | #[test] |
512 | | fn test_softmax_op_empty() { |
513 | | let op = SoftmaxOp::new(0); |
514 | | let result = op.execute(vec![], Backend::Scalar).unwrap(); |
515 | | assert!(result.is_empty()); |
516 | | } |
517 | | |
518 | | #[test] |
519 | | fn test_softmax_numerical_stability() { |
520 | | let op = SoftmaxOp::new(3); |
521 | | // Large values that would overflow naive exp |
522 | | let input = vec![1000.0, 1001.0, 1002.0]; |
523 | | let result = op.execute(input, Backend::Scalar).unwrap(); |
524 | | |
525 | | // Should still be valid probabilities |
526 | | let sum: f32 = result.iter().sum(); |
527 | | assert!((sum - 1.0).abs() < 0.001); |
528 | | assert!(result.iter().all(|&x| x.is_finite())); |
529 | | } |
530 | | |
531 | | /// FALSIFICATION: Dot product is commutative |
532 | | #[test] |
533 | | fn test_falsify_dot_commutative() { |
534 | | let op = DotOp::new(5); |
535 | | let a = vec![1.0, 2.0, 3.0, 4.0, 5.0]; |
536 | | let b = vec![5.0, 4.0, 3.0, 2.0, 1.0]; |
537 | | |
538 | | let result1 = op.execute((a.clone(), b.clone()), Backend::Scalar).unwrap(); |
539 | | let result2 = op.execute((b, a), Backend::Scalar).unwrap(); |
540 | | |
541 | | assert!( |
542 | | (result1 - result2).abs() < 1e-6, |
543 | | "FALSIFICATION FAILED: dot product not commutative" |
544 | | ); |
545 | | } |
546 | | |
547 | | /// FALSIFICATION: Softmax probabilities sum to 1 |
548 | | #[test] |
549 | | fn test_falsify_softmax_sum_to_one() { |
550 | | for len in [1, 5, 10, 100] { |
551 | | let op = SoftmaxOp::new(len); |
552 | | let input: Vec<f32> = (0..len).map(|i| i as f32).collect(); |
553 | | let result = op.execute(input, Backend::Scalar).unwrap(); |
554 | | let sum: f32 = result.iter().sum(); |
555 | | assert!( |
556 | | (sum - 1.0).abs() < 1e-5, |
557 | | "FALSIFICATION FAILED: softmax sum {} != 1.0 for len {}", |
558 | | sum, |
559 | | len |
560 | | ); |
561 | | } |
562 | | } |
563 | | } |