/home/noah/src/trueno/src/matrix/mod.rs
Line | Count | Source |
1 | | //! Matrix operations for Trueno |
2 | | //! |
3 | | //! Provides 2D matrix operations with SIMD optimization for linear algebra, |
4 | | //! machine learning, and scientific computing. |
5 | | //! |
6 | | //! # Example |
7 | | //! |
8 | | //! ``` |
9 | | //! use trueno::Matrix; |
10 | | //! |
11 | | //! // Create a 2x3 matrix |
12 | | //! let m = Matrix::zeros(2, 3); |
13 | | //! assert_eq!(m.rows(), 2); |
14 | | //! assert_eq!(m.cols(), 3); |
15 | | //! ``` |
16 | | |
17 | | // Allow dead_code for experimental SIMD microkernels kept for future optimization work |
18 | | #![allow(dead_code)] |
19 | | |
20 | | use crate::{Backend, TruenoError, Vector}; |
21 | | |
22 | | #[cfg(feature = "tracing")] |
23 | | use tracing::instrument; |
24 | | |
25 | | /// A 2D matrix with row-major storage |
26 | | /// |
27 | | /// Data is stored in row-major format (C-style), where consecutive elements |
28 | | /// in memory belong to the same row. This is compatible with NumPy's default |
29 | | /// layout and optimal for cache locality when accessing rows. |
30 | | /// |
31 | | /// # Storage Layout |
32 | | /// |
33 | | /// For a 2x3 matrix: |
34 | | /// ```text |
35 | | /// [[a, b, c], |
36 | | /// [d, e, f]] |
37 | | /// ``` |
38 | | /// Data is stored as: [a, b, c, d, e, f] |
39 | | /// |
40 | | /// # Example |
41 | | /// |
42 | | /// ``` |
43 | | /// use trueno::Matrix; |
44 | | /// |
45 | | /// let m = Matrix::from_vec(2, 2, vec![1.0, 2.0, 3.0, 4.0]).unwrap(); |
46 | | /// assert_eq!(m.get(0, 0), Some(&1.0)); |
47 | | /// assert_eq!(m.get(0, 1), Some(&2.0)); |
48 | | /// assert_eq!(m.get(1, 0), Some(&3.0)); |
49 | | /// assert_eq!(m.get(1, 1), Some(&4.0)); |
50 | | /// ``` |
51 | | #[derive(Debug, Clone, PartialEq)] |
52 | | pub struct Matrix<T> { |
53 | | rows: usize, |
54 | | cols: usize, |
55 | | data: Vec<T>, |
56 | | backend: Backend, |
57 | | } |
58 | | |
59 | | impl std::ops::Index<(usize, usize)> for Matrix<f32> { |
60 | | type Output = f32; |
61 | | |
62 | 0 | fn index(&self, (row, col): (usize, usize)) -> &Self::Output { |
63 | 0 | &self.data[row * self.cols + col] |
64 | 0 | } |
65 | | } |
66 | | |
67 | | impl Matrix<f32> { |
68 | | /// Creates a new matrix with uninitialized values |
69 | | /// |
70 | | /// # Arguments |
71 | | /// |
72 | | /// * `rows` - Number of rows |
73 | | /// * `cols` - Number of columns |
74 | | /// |
75 | | /// # Returns |
76 | | /// |
77 | | /// A new matrix with dimensions `rows x cols` containing uninitialized values |
78 | | /// |
79 | | /// # Example |
80 | | /// |
81 | | /// ``` |
82 | | /// use trueno::Matrix; |
83 | | /// |
84 | | /// let m = Matrix::new(3, 4); |
85 | | /// assert_eq!(m.rows(), 3); |
86 | | /// assert_eq!(m.cols(), 4); |
87 | | /// ``` |
88 | 2 | pub fn new(rows: usize, cols: usize) -> Self { |
89 | 2 | let backend = Backend::select_best(); |
90 | 2 | Matrix { |
91 | 2 | rows, |
92 | 2 | cols, |
93 | 2 | data: vec![0.0; rows * cols], |
94 | 2 | backend, |
95 | 2 | } |
96 | 2 | } |
97 | | |
98 | | /// Creates a matrix from a vector of data |
99 | | /// |
100 | | /// # Arguments |
101 | | /// |
102 | | /// * `rows` - Number of rows |
103 | | /// * `cols` - Number of columns |
104 | | /// * `data` - Vector containing matrix elements in row-major order |
105 | | /// |
106 | | /// # Errors |
107 | | /// |
108 | | /// Returns `InvalidInput` if `data.len() != rows * cols` |
109 | | /// |
110 | | /// # Example |
111 | | /// |
112 | | /// ``` |
113 | | /// use trueno::Matrix; |
114 | | /// |
115 | | /// let m = Matrix::from_vec(2, 2, vec![1.0, 2.0, 3.0, 4.0]).unwrap(); |
116 | | /// assert_eq!(m.rows(), 2); |
117 | | /// assert_eq!(m.cols(), 2); |
118 | | /// ``` |
119 | 0 | pub fn from_vec(rows: usize, cols: usize, data: Vec<f32>) -> Result<Self, TruenoError> { |
120 | 0 | if data.len() != rows * cols { |
121 | 0 | return Err(TruenoError::InvalidInput(format!( |
122 | 0 | "Data length {} does not match matrix dimensions {}x{} (expected {})", |
123 | 0 | data.len(), |
124 | 0 | rows, |
125 | 0 | cols, |
126 | 0 | rows * cols |
127 | 0 | ))); |
128 | 0 | } |
129 | | |
130 | 0 | let backend = Backend::select_best(); |
131 | 0 | Ok(Matrix { |
132 | 0 | rows, |
133 | 0 | cols, |
134 | 0 | data, |
135 | 0 | backend, |
136 | 0 | }) |
137 | 0 | } |
138 | | |
139 | | /// Creates a matrix from a vector with a specific backend |
140 | | /// |
141 | | /// This is useful for testing specific SIMD code paths. |
142 | 0 | pub fn from_vec_with_backend( |
143 | 0 | rows: usize, |
144 | 0 | cols: usize, |
145 | 0 | data: Vec<f32>, |
146 | 0 | backend: Backend, |
147 | 0 | ) -> Self { |
148 | 0 | assert_eq!( |
149 | 0 | data.len(), |
150 | 0 | rows * cols, |
151 | 0 | "Data length {} does not match matrix dimensions {}x{}", |
152 | 0 | data.len(), |
153 | | rows, |
154 | | cols |
155 | | ); |
156 | 0 | Matrix { |
157 | 0 | rows, |
158 | 0 | cols, |
159 | 0 | data, |
160 | 0 | backend, |
161 | 0 | } |
162 | 0 | } |
163 | | |
164 | | /// Creates a matrix from a slice by copying the data |
165 | | /// |
166 | | /// This is a convenience method that copies the slice into an owned vector. |
167 | | /// For zero-copy scenarios, consider using the data directly with `from_vec` |
168 | | /// if you already have an owned `Vec`. |
169 | | /// |
170 | | /// # Arguments |
171 | | /// |
172 | | /// * `rows` - Number of rows |
173 | | /// * `cols` - Number of columns |
174 | | /// * `data` - Slice containing matrix elements in row-major order |
175 | | /// |
176 | | /// # Errors |
177 | | /// |
178 | | /// Returns `InvalidInput` if `data.len() != rows * cols` |
179 | | /// |
180 | | /// # Example |
181 | | /// |
182 | | /// ``` |
183 | | /// use trueno::Matrix; |
184 | | /// |
185 | | /// let data = [1.0, 2.0, 3.0, 4.0]; |
186 | | /// let m = Matrix::from_slice(2, 2, &data).unwrap(); |
187 | | /// assert_eq!(m.get(0, 0), Some(&1.0)); |
188 | | /// ``` |
189 | 0 | pub fn from_slice(rows: usize, cols: usize, data: &[f32]) -> Result<Self, TruenoError> { |
190 | 0 | Self::from_vec(rows, cols, data.to_vec()) |
191 | 0 | } |
192 | | |
193 | | /// Creates a matrix filled with zeros |
194 | | /// |
195 | | /// # Example |
196 | | /// |
197 | | /// ``` |
198 | | /// use trueno::Matrix; |
199 | | /// |
200 | | /// let m = Matrix::zeros(3, 3); |
201 | | /// assert_eq!(m.get(1, 1), Some(&0.0)); |
202 | | /// ``` |
203 | 2 | pub fn zeros(rows: usize, cols: usize) -> Self { |
204 | 2 | Matrix::new(rows, cols) |
205 | 2 | } |
206 | | |
207 | | /// Creates a matrix filled with zeros using a specific backend |
208 | | /// (Internal use only - reuses backend from parent matrix) |
209 | 0 | fn zeros_with_backend(rows: usize, cols: usize, backend: Backend) -> Self { |
210 | 0 | Matrix { |
211 | 0 | rows, |
212 | 0 | cols, |
213 | 0 | data: vec![0.0; rows * cols], |
214 | 0 | backend, |
215 | 0 | } |
216 | 0 | } |
217 | | |
218 | | /// Creates an identity matrix (square matrix with 1s on diagonal) |
219 | | /// |
220 | | /// # Example |
221 | | /// |
222 | | /// ``` |
223 | | /// use trueno::Matrix; |
224 | | /// |
225 | | /// let m = Matrix::identity(3); |
226 | | /// assert_eq!(m.get(0, 0), Some(&1.0)); |
227 | | /// assert_eq!(m.get(0, 1), Some(&0.0)); |
228 | | /// assert_eq!(m.get(1, 1), Some(&1.0)); |
229 | | /// ``` |
230 | 0 | pub fn identity(n: usize) -> Self { |
231 | 0 | let mut data = vec![0.0; n * n]; |
232 | 0 | for i in 0..n { |
233 | 0 | data[i * n + i] = 1.0; |
234 | 0 | } |
235 | 0 | let backend = Backend::select_best(); |
236 | 0 | Matrix { |
237 | 0 | rows: n, |
238 | 0 | cols: n, |
239 | 0 | data, |
240 | 0 | backend, |
241 | 0 | } |
242 | 0 | } |
243 | | |
244 | | /// Returns the number of rows |
245 | 2 | pub fn rows(&self) -> usize { |
246 | 2 | self.rows |
247 | 2 | } |
248 | | |
249 | | /// Returns the number of columns |
250 | 2 | pub fn cols(&self) -> usize { |
251 | 2 | self.cols |
252 | 2 | } |
253 | | |
254 | | /// Returns the shape as (rows, cols) |
255 | 0 | pub fn shape(&self) -> (usize, usize) { |
256 | 0 | (self.rows, self.cols) |
257 | 0 | } |
258 | | |
259 | | /// Gets a reference to an element at (row, col) |
260 | | /// |
261 | | /// Returns `None` if indices are out of bounds |
262 | 0 | pub fn get(&self, row: usize, col: usize) -> Option<&f32> { |
263 | 0 | if row >= self.rows || col >= self.cols { |
264 | 0 | None |
265 | | } else { |
266 | 0 | self.data.get(row * self.cols + col) |
267 | | } |
268 | 0 | } |
269 | | |
270 | | /// Gets a mutable reference to an element at (row, col) |
271 | | /// |
272 | | /// Returns `None` if indices are out of bounds |
273 | 0 | pub fn get_mut(&mut self, row: usize, col: usize) -> Option<&mut f32> { |
274 | 0 | if row >= self.rows || col >= self.cols { |
275 | 0 | None |
276 | | } else { |
277 | 0 | let idx = row * self.cols + col; |
278 | 0 | self.data.get_mut(idx) |
279 | | } |
280 | 0 | } |
281 | | |
282 | | /// Returns a reference to the underlying data |
283 | 0 | pub fn as_slice(&self) -> &[f32] { |
284 | 0 | &self.data |
285 | 0 | } |
286 | | |
287 | | /// Matrix multiplication (matmul) |
288 | | /// |
289 | | /// Computes `C = A × B` where A is `m×n`, B is `n×p`, and C is `m×p`. |
290 | | /// |
291 | | /// # Arguments |
292 | | /// |
293 | | /// * `other` - The matrix to multiply with (right operand) |
294 | | /// |
295 | | /// # Returns |
296 | | /// |
297 | | /// A new matrix containing the result of matrix multiplication |
298 | | /// |
299 | | /// # Errors |
300 | | /// |
301 | | /// Returns `InvalidInput` if matrix dimensions are incompatible |
302 | | /// (i.e., `self.cols != other.rows`) |
303 | | /// |
304 | | /// # Example |
305 | | /// |
306 | | /// ``` |
307 | | /// use trueno::Matrix; |
308 | | /// |
309 | | /// let a = Matrix::from_vec(2, 2, vec![1.0, 2.0, 3.0, 4.0]).unwrap(); |
310 | | /// let b = Matrix::from_vec(2, 2, vec![5.0, 6.0, 7.0, 8.0]).unwrap(); |
311 | | /// let c = a.matmul(&b).unwrap(); |
312 | | /// |
313 | | /// // [[1, 2], [[5, 6], [[19, 22], |
314 | | /// // [3, 4]] × [7, 8]] = [43, 50]] |
315 | | /// assert_eq!(c.get(0, 0), Some(&19.0)); |
316 | | /// assert_eq!(c.get(0, 1), Some(&22.0)); |
317 | | /// assert_eq!(c.get(1, 0), Some(&43.0)); |
318 | | /// assert_eq!(c.get(1, 1), Some(&50.0)); |
319 | | /// ``` |
320 | | // ========================================================================= |
321 | | // HOT PATH - PERFORMANCE CRITICAL |
322 | | // ========================================================================= |
323 | | // Core matrix operation used in neural network forward passes. |
324 | | // Changes to inner loops REQUIRE benchmark verification: make bench-check |
325 | | // See convolve2d for prohibited patterns in hot loops. |
326 | | // ========================================================================= |
327 | | #[cfg_attr(feature = "tracing", instrument(skip(self, other), fields(dims = %format!("{}x{} @ {}x{}", self.rows, self.cols, other.rows, other.cols))))] |
328 | 0 | pub fn matmul(&self, other: &Matrix<f32>) -> Result<Matrix<f32>, TruenoError> { |
329 | 0 | if self.cols != other.rows { |
330 | 0 | return Err(TruenoError::InvalidInput(format!( |
331 | 0 | "Matrix dimension mismatch for multiplication: {}×{} × {}×{} (inner dimensions {} and {} must match)", |
332 | 0 | self.rows, self.cols, other.rows, other.cols, self.cols, other.rows |
333 | 0 | ))); |
334 | 0 | } |
335 | | |
336 | | // Fast path for vector-matrix multiply (rows=1) |
337 | | // Common in ML vocab projection: hidden_state @ embedding_transposed |
338 | | // 8x faster than general matmul for 1×384 @ 384×51865 pattern |
339 | 0 | if self.rows == 1 { |
340 | 0 | return self.matmul_vector_matrix(other); |
341 | 0 | } |
342 | | |
343 | 0 | let mut result = Matrix::zeros_with_backend(self.rows, other.cols, self.backend); |
344 | | |
345 | | // Backend selection strategy (empirical - see docs/performance-analysis.md): |
346 | | // 1. GPU for large matrices (≥500×500) - 2-10x speedup (measured) |
347 | | // 2. SIMD for medium-large matrices (>64×64) - 2-8x speedup |
348 | | // 3. Naive for small matrices - lowest overhead |
349 | | |
350 | | #[cfg(all(feature = "gpu", not(target_arch = "wasm32")))] |
351 | | const GPU_THRESHOLD: usize = 500; // Empirical: 2x at 500×500, 9.6x at 1000×1000 |
352 | | const SIMD_THRESHOLD: usize = 64; |
353 | | |
354 | | // Try GPU first for very large matrices |
355 | | #[cfg(all(feature = "gpu", not(target_arch = "wasm32")))] |
356 | | { |
357 | 0 | if self.rows >= GPU_THRESHOLD |
358 | 0 | && self.cols >= GPU_THRESHOLD |
359 | 0 | && other.cols >= GPU_THRESHOLD |
360 | | { |
361 | 0 | if let Ok(gpu_result) = self.matmul_gpu(other) { |
362 | 0 | return Ok(gpu_result); |
363 | 0 | } |
364 | | // GPU failed, fall through to SIMD/naive |
365 | 0 | } |
366 | | } |
367 | | |
368 | | // Use SIMD for medium-large matrices |
369 | 0 | if self.rows >= SIMD_THRESHOLD |
370 | 0 | || self.cols >= SIMD_THRESHOLD |
371 | 0 | || other.cols >= SIMD_THRESHOLD |
372 | | { |
373 | | // Tiled approach threshold: below this size, tiling beats transpose |
374 | | // Based on WASM optimization spec benchmarks |
375 | | const TILED_THRESHOLD: usize = 512; |
376 | | |
377 | 0 | let max_dim = self.rows.max(self.cols).max(other.cols); |
378 | | |
379 | 0 | if max_dim < TILED_THRESHOLD { |
380 | | // Medium matrices: use BLIS on native, tiled on WASM |
381 | | #[cfg(target_arch = "wasm32")] |
382 | | { |
383 | | self.matmul_wasm_tiled(other, &mut result)?; |
384 | | } |
385 | | #[cfg(not(target_arch = "wasm32"))] |
386 | | { |
387 | | // BLIS is faster than tiled for all sizes on native |
388 | 0 | crate::blis::gemm_blis( |
389 | 0 | self.rows, |
390 | 0 | other.cols, |
391 | 0 | self.cols, |
392 | 0 | &self.data, |
393 | 0 | &other.data, |
394 | 0 | &mut result.data, |
395 | 0 | None, |
396 | 0 | )?; |
397 | | } |
398 | | } else { |
399 | | // Large matrices: platform-specific optimized paths |
400 | | #[cfg(target_arch = "wasm32")] |
401 | | { |
402 | | // WASM: tiled is always better (no SIMD microkernel advantage) |
403 | | self.matmul_wasm_tiled(other, &mut result)?; |
404 | | } |
405 | | #[cfg(not(target_arch = "wasm32"))] |
406 | | { |
407 | | // Native: use BLIS-style GEMM with register blocking |
408 | | // ~2x faster than old SIMD implementation for large matrices |
409 | 0 | crate::blis::gemm_blis( |
410 | 0 | self.rows, |
411 | 0 | other.cols, |
412 | 0 | self.cols, |
413 | 0 | &self.data, |
414 | 0 | &other.data, |
415 | 0 | &mut result.data, |
416 | 0 | None, |
417 | 0 | )?; |
418 | | } |
419 | | } |
420 | | } else { |
421 | 0 | self.matmul_naive(other, &mut result)?; |
422 | | } |
423 | | |
424 | 0 | Ok(result) |
425 | 0 | } |
426 | | |
427 | | /// Batched matrix multiplication for 3D tensors. |
428 | | /// |
429 | | /// Computes `[batch, m, k] @ [batch, k, n] -> [batch, m, n]` using SIMD for each batch. |
430 | | /// This is critical for transformer attention performance. |
431 | | /// |
432 | | /// # Arguments |
433 | | /// * `a_data` - Flattened input A with shape [batch, m, k] |
434 | | /// * `b_data` - Flattened input B with shape [batch, k, n] |
435 | | /// * `batch` - Batch dimension |
436 | | /// * `m` - Rows of A (and output) |
437 | | /// * `k` - Columns of A / Rows of B |
438 | | /// * `n` - Columns of B (and output) |
439 | | /// |
440 | | /// # Returns |
441 | | /// Flattened output with shape [batch, m, n] |
442 | | /// |
443 | | /// # Performance |
444 | | /// Uses SIMD matmul for each batch slice, achieving ~50 GFLOPS vs ~0.1 GFLOPS naive. |
445 | | /// See Williams et al., 2009 (Roofline model) for theoretical analysis. |
446 | | #[cfg_attr( |
447 | | feature = "tracing", |
448 | | instrument(skip(a_data, b_data), fields(batch, m, k, n)) |
449 | | )] |
450 | 0 | pub fn batched_matmul( |
451 | 0 | a_data: &[f32], |
452 | 0 | b_data: &[f32], |
453 | 0 | batch: usize, |
454 | 0 | m: usize, |
455 | 0 | k: usize, |
456 | 0 | n: usize, |
457 | 0 | ) -> Result<Vec<f32>, TruenoError> { |
458 | 0 | let a_stride = m * k; |
459 | 0 | let b_stride = k * n; |
460 | 0 | let out_stride = m * n; |
461 | | |
462 | | // Validate input sizes |
463 | 0 | if a_data.len() != batch * a_stride { |
464 | 0 | return Err(TruenoError::InvalidInput(format!( |
465 | 0 | "A data size mismatch: expected {} ({}×{}×{}), got {}", |
466 | 0 | batch * a_stride, |
467 | 0 | batch, |
468 | 0 | m, |
469 | 0 | k, |
470 | 0 | a_data.len() |
471 | 0 | ))); |
472 | 0 | } |
473 | 0 | if b_data.len() != batch * b_stride { |
474 | 0 | return Err(TruenoError::InvalidInput(format!( |
475 | 0 | "B data size mismatch: expected {} ({}×{}×{}), got {}", |
476 | 0 | batch * b_stride, |
477 | 0 | batch, |
478 | 0 | k, |
479 | 0 | n, |
480 | 0 | b_data.len() |
481 | 0 | ))); |
482 | 0 | } |
483 | | |
484 | 0 | let mut output = vec![0.0f32; batch * out_stride]; |
485 | | |
486 | | // Process each batch using SIMD matmul |
487 | 0 | for ba in 0..batch { |
488 | 0 | let a_offset = ba * a_stride; |
489 | 0 | let b_offset = ba * b_stride; |
490 | 0 | let out_offset = ba * out_stride; |
491 | | |
492 | | // Create matrix views from slices (no copy - just metadata) |
493 | 0 | let a_slice = &a_data[a_offset..a_offset + a_stride]; |
494 | 0 | let b_slice = &b_data[b_offset..b_offset + b_stride]; |
495 | | |
496 | | // Use from_slice to avoid copying |
497 | 0 | let a_mat = Matrix::from_slice(m, k, a_slice)?; |
498 | 0 | let b_mat = Matrix::from_slice(k, n, b_slice)?; |
499 | | |
500 | | // SIMD matmul |
501 | 0 | let result = a_mat.matmul(&b_mat)?; |
502 | | |
503 | | // Copy result to output |
504 | 0 | output[out_offset..out_offset + out_stride].copy_from_slice(result.as_slice()); |
505 | | } |
506 | | |
507 | 0 | Ok(output) |
508 | 0 | } |
509 | | |
510 | | /// Batched matrix multiplication for 4D tensors (attention pattern). |
511 | | /// |
512 | | /// Computes `[batch, heads, m, k] @ [batch, heads, k, n] -> [batch, heads, m, n]` |
513 | | /// This is the exact pattern used in multi-head attention: Q @ K^T and attn @ V. |
514 | | /// |
515 | | /// # Arguments |
516 | | /// * `a_data` - Flattened input A with shape [batch, heads, m, k] |
517 | | /// * `b_data` - Flattened input B with shape [batch, heads, k, n] |
518 | | /// * `batch` - Batch dimension |
519 | | /// * `heads` - Number of attention heads |
520 | | /// * `m` - Rows (sequence length for Q) |
521 | | /// * `k` - Columns of A / Rows of B (head dimension) |
522 | | /// * `n` - Columns of B (sequence length for K^T, or head dim for V) |
523 | | /// |
524 | | /// # Performance |
525 | | /// Processes batch×heads independent matmuls using SIMD. |
526 | | /// For Qwen2-0.5B: batch=1, heads=14, m=seq, k=64, n=seq |
527 | | #[cfg_attr( |
528 | | feature = "tracing", |
529 | | instrument(skip(a_data, b_data), fields(batch, heads, m, k, n)) |
530 | | )] |
531 | 0 | pub fn batched_matmul_4d( |
532 | 0 | a_data: &[f32], |
533 | 0 | b_data: &[f32], |
534 | 0 | batch: usize, |
535 | 0 | heads: usize, |
536 | 0 | m: usize, |
537 | 0 | k: usize, |
538 | 0 | n: usize, |
539 | 0 | ) -> Result<Vec<f32>, TruenoError> { |
540 | 0 | let a_head_stride = m * k; |
541 | 0 | let b_head_stride = k * n; |
542 | 0 | let out_head_stride = m * n; |
543 | 0 | let total_heads = batch * heads; |
544 | | |
545 | | // Validate input sizes |
546 | 0 | let expected_a = total_heads * a_head_stride; |
547 | 0 | let expected_b = total_heads * b_head_stride; |
548 | 0 | if a_data.len() != expected_a { |
549 | 0 | return Err(TruenoError::InvalidInput(format!( |
550 | 0 | "A data size mismatch: expected {} ({}×{}×{}×{}), got {}", |
551 | 0 | expected_a, |
552 | 0 | batch, |
553 | 0 | heads, |
554 | 0 | m, |
555 | 0 | k, |
556 | 0 | a_data.len() |
557 | 0 | ))); |
558 | 0 | } |
559 | 0 | if b_data.len() != expected_b { |
560 | 0 | return Err(TruenoError::InvalidInput(format!( |
561 | 0 | "B data size mismatch: expected {} ({}×{}×{}×{}), got {}", |
562 | 0 | expected_b, |
563 | 0 | batch, |
564 | 0 | heads, |
565 | 0 | k, |
566 | 0 | n, |
567 | 0 | b_data.len() |
568 | 0 | ))); |
569 | 0 | } |
570 | | |
571 | 0 | let mut output = vec![0.0f32; total_heads * out_head_stride]; |
572 | | |
573 | | // Process each (batch, head) pair using SIMD matmul |
574 | 0 | for bh in 0..total_heads { |
575 | 0 | let a_offset = bh * a_head_stride; |
576 | 0 | let b_offset = bh * b_head_stride; |
577 | 0 | let out_offset = bh * out_head_stride; |
578 | | |
579 | | // Create matrix views from slices |
580 | 0 | let a_slice = &a_data[a_offset..a_offset + a_head_stride]; |
581 | 0 | let b_slice = &b_data[b_offset..b_offset + b_head_stride]; |
582 | | |
583 | 0 | let a_mat = Matrix::from_slice(m, k, a_slice)?; |
584 | 0 | let b_mat = Matrix::from_slice(k, n, b_slice)?; |
585 | | |
586 | | // SIMD matmul |
587 | 0 | let result = a_mat.matmul(&b_mat)?; |
588 | | |
589 | | // Copy result to output |
590 | 0 | output[out_offset..out_offset + out_head_stride].copy_from_slice(result.as_slice()); |
591 | | } |
592 | | |
593 | 0 | Ok(output) |
594 | 0 | } |
595 | | |
596 | | /// Fast path for vector-matrix multiplication (1×K @ K×N → 1×N) |
597 | | /// |
598 | | /// This is 8x faster than general matmul for patterns like: |
599 | | /// - Vocab projection: hidden_state (1×384) @ embedding_transposed (384×51865) |
600 | | /// - Single token decode in Whisper/LLM inference |
601 | | /// |
602 | | /// Strategy: Outer product accumulation (no transpose needed!) |
603 | | /// For result[j] = sum_k(A[0,k] * B[k,j]), we compute: |
604 | | /// result += A[k] * B[k,:] for each k |
605 | | /// This has excellent cache locality since we access entire rows of B. |
606 | | #[cfg_attr(feature = "tracing", instrument(skip(self, other), fields(k = self.cols, n = other.cols)))] |
607 | 0 | fn matmul_vector_matrix(&self, other: &Matrix<f32>) -> Result<Matrix<f32>, TruenoError> { |
608 | 0 | debug_assert_eq!(self.rows, 1); |
609 | | |
610 | 0 | let k = self.cols; // Inner dimension |
611 | 0 | let n = other.cols; // Output dimension |
612 | | |
613 | | // Result is 1×N, initialized to zero |
614 | 0 | let mut result = Matrix::zeros_with_backend(1, n, self.backend); |
615 | | |
616 | | // Outer product accumulation: result += A[k] * B[k,:] |
617 | | // For each k, scale row k of B by A[k] and add to result |
618 | | // The compiler will auto-vectorize this inner loop |
619 | 0 | for ki in 0..k { |
620 | 0 | let a_k = self.data[ki]; |
621 | 0 | if a_k == 0.0 { |
622 | 0 | continue; // Skip zero multiplications |
623 | 0 | } |
624 | | |
625 | | // Get row ki of B (contiguous in memory - cache friendly!) |
626 | 0 | let b_row_start = ki * n; |
627 | | |
628 | | // AXPY: result += a_k * B[ki,:] |
629 | | // This loop is auto-vectorized by LLVM with -O2/-O3 |
630 | 0 | for j in 0..n { |
631 | 0 | result.data[j] += a_k * other.data[b_row_start + j]; |
632 | 0 | } |
633 | | } |
634 | | |
635 | 0 | Ok(result) |
636 | 0 | } |
637 | | |
638 | | /// Naive O(n³) matrix multiplication (baseline for small matrices) |
639 | 0 | fn matmul_naive( |
640 | 0 | &self, |
641 | 0 | other: &Matrix<f32>, |
642 | 0 | result: &mut Matrix<f32>, |
643 | 0 | ) -> Result<(), TruenoError> { |
644 | | // C[i,j] = Σ A[i,k] × B[k,j] |
645 | | // SAFETY: Loop bounds are validated by dimension checks in matmul() |
646 | 0 | for i in 0..self.rows { |
647 | 0 | for j in 0..other.cols { |
648 | 0 | let mut sum = 0.0; |
649 | 0 | for k in 0..self.cols { |
650 | 0 | // Bounds guaranteed: i < self.rows, k < self.cols, j < other.cols |
651 | 0 | sum += self |
652 | 0 | .get(i, k) |
653 | 0 | .expect("matmul_naive: A[i,k] bounds validated by loop") |
654 | 0 | * other |
655 | 0 | .get(k, j) |
656 | 0 | .expect("matmul_naive: B[k,j] bounds validated by loop"); |
657 | 0 | } |
658 | 0 | *result |
659 | 0 | .get_mut(i, j) |
660 | 0 | .expect("matmul_naive: C[i,j] bounds validated by loop") = sum; |
661 | | } |
662 | | } |
663 | 0 | Ok(()) |
664 | 0 | } |
665 | | |
666 | | /// AVX2 micro-kernel: Compute 4 rows × 1 column using register blocking (Phase 2) |
667 | | /// |
668 | | /// This micro-kernel processes 4 rows of matrix A against 1 column of B_transposed |
669 | | /// simultaneously, keeping intermediate results in AVX2 registers for efficiency. |
670 | | /// |
671 | | /// # Performance Benefits |
672 | | /// - Loads B-column once, reuses for 4 A-rows (4× reduction in memory bandwidth) |
673 | | /// - Uses FMA instructions for fused multiply-add (3× throughput vs separate ops) |
674 | | /// - Keeps accumulators in YMM registers (no memory traffic for intermediate results) |
675 | | /// |
676 | | /// # Safety |
677 | | /// - Caller must ensure all slices have the same length |
678 | | /// - Must be called on x86_64 with AVX2 support |
679 | | #[cfg(target_arch = "x86_64")] |
680 | | #[target_feature(enable = "avx2,fma")] |
681 | | #[inline] |
682 | 0 | unsafe fn matmul_microkernel_4x1_avx2( |
683 | 0 | a_rows: [&[f32]; 4], |
684 | 0 | b_col: &[f32], |
685 | 0 | results: &mut [f32; 4], |
686 | 0 | ) { |
687 | | use std::arch::x86_64::*; |
688 | | |
689 | 0 | let len = b_col.len(); |
690 | 0 | let chunks = len / 8; // Process 8 f32 elements per iteration (AVX2 = 256 bits) |
691 | | |
692 | | // Accumulators for 4 output elements (kept in registers) |
693 | 0 | let mut acc0 = _mm256_setzero_ps(); |
694 | 0 | let mut acc1 = _mm256_setzero_ps(); |
695 | 0 | let mut acc2 = _mm256_setzero_ps(); |
696 | 0 | let mut acc3 = _mm256_setzero_ps(); |
697 | | |
698 | | // Main loop: Process 8 elements at a time |
699 | 0 | for i in 0..chunks { |
700 | 0 | let offset = i * 8; |
701 | 0 |
|
702 | 0 | // Load B column (reused for all 4 A rows) |
703 | 0 | let b_vec = _mm256_loadu_ps(b_col.as_ptr().add(offset)); |
704 | 0 |
|
705 | 0 | // Load A rows and FMA (Fused Multiply-Add) |
706 | 0 | let a0_vec = _mm256_loadu_ps(a_rows[0].as_ptr().add(offset)); |
707 | 0 | acc0 = _mm256_fmadd_ps(a0_vec, b_vec, acc0); |
708 | 0 |
|
709 | 0 | let a1_vec = _mm256_loadu_ps(a_rows[1].as_ptr().add(offset)); |
710 | 0 | acc1 = _mm256_fmadd_ps(a1_vec, b_vec, acc1); |
711 | 0 |
|
712 | 0 | let a2_vec = _mm256_loadu_ps(a_rows[2].as_ptr().add(offset)); |
713 | 0 | acc2 = _mm256_fmadd_ps(a2_vec, b_vec, acc2); |
714 | 0 |
|
715 | 0 | let a3_vec = _mm256_loadu_ps(a_rows[3].as_ptr().add(offset)); |
716 | 0 | acc3 = _mm256_fmadd_ps(a3_vec, b_vec, acc3); |
717 | 0 | } |
718 | | |
719 | | // Horizontal sum of each accumulator (reduce 8 elements to 1) |
720 | 0 | results[0] = Self::horizontal_sum_avx2(acc0); |
721 | 0 | results[1] = Self::horizontal_sum_avx2(acc1); |
722 | 0 | results[2] = Self::horizontal_sum_avx2(acc2); |
723 | 0 | results[3] = Self::horizontal_sum_avx2(acc3); |
724 | | |
725 | | // Handle remainder elements with scalar code |
726 | 0 | let remainder_start = chunks * 8; |
727 | 0 | if remainder_start < len { |
728 | 0 | for i in remainder_start..len { |
729 | 0 | results[0] += a_rows[0][i] * b_col[i]; |
730 | 0 | results[1] += a_rows[1][i] * b_col[i]; |
731 | 0 | results[2] += a_rows[2][i] * b_col[i]; |
732 | 0 | results[3] += a_rows[3][i] * b_col[i]; |
733 | 0 | } |
734 | 0 | } |
735 | 0 | } |
736 | | |
737 | | /// Helper: Horizontal sum of 8 f32 values in an AVX2 register |
738 | | #[cfg(target_arch = "x86_64")] |
739 | | #[target_feature(enable = "avx2")] |
740 | | #[inline] |
741 | 0 | unsafe fn horizontal_sum_avx2(v: std::arch::x86_64::__m256) -> f32 { |
742 | | use std::arch::x86_64::*; |
743 | | |
744 | | // Sum upper and lower 128-bit lanes |
745 | 0 | let sum128 = _mm_add_ps(_mm256_castps256_ps128(v), _mm256_extractf128_ps(v, 1)); |
746 | | |
747 | | // Horizontal add within 128-bit lane (4 values → 2 values) |
748 | 0 | let sum64 = _mm_hadd_ps(sum128, sum128); |
749 | | |
750 | | // Horizontal add again (2 values → 1 value) |
751 | 0 | let sum32 = _mm_hadd_ps(sum64, sum64); |
752 | | |
753 | | // Extract final scalar result |
754 | 0 | _mm_cvtss_f32(sum32) |
755 | 0 | } |
756 | | |
757 | | /// AVX-512 micro-kernel: Compute 8 rows × 1 column using register blocking (Phase 3) |
758 | | /// |
759 | | /// This micro-kernel processes 8 rows of matrix A against 1 column of B_transposed |
760 | | /// simultaneously, keeping intermediate results in AVX-512 registers for efficiency. |
761 | | /// |
762 | | /// # Performance Benefits |
763 | | /// - Processes 16 f32 elements per iteration (vs 8 with AVX2) - 2× throughput |
764 | | /// - Loads B-column once, reuses for 8 A-rows (8× reduction in memory bandwidth) |
765 | | /// - Uses FMA instructions for fused multiply-add (3× throughput vs separate ops) |
766 | | /// - Keeps accumulators in ZMM registers (no memory traffic for intermediate results) |
767 | | /// |
768 | | /// # Safety |
769 | | /// - Caller must ensure all slices have the same length |
770 | | /// - Must be called on x86_64 with AVX-512F support |
771 | | #[cfg(target_arch = "x86_64")] |
772 | | #[target_feature(enable = "avx512f")] |
773 | | #[inline] |
774 | 0 | unsafe fn matmul_microkernel_8x1_avx512( |
775 | 0 | a_rows: [&[f32]; 8], |
776 | 0 | b_col: &[f32], |
777 | 0 | results: &mut [f32; 8], |
778 | 0 | ) { |
779 | | use std::arch::x86_64::*; |
780 | | |
781 | 0 | let len = b_col.len(); |
782 | 0 | let chunks = len / 16; // Process 16 f32 elements per iteration (AVX-512 = 512 bits) |
783 | | |
784 | | // Accumulators for 8 output elements (kept in ZMM registers) |
785 | 0 | let mut acc0 = _mm512_setzero_ps(); |
786 | 0 | let mut acc1 = _mm512_setzero_ps(); |
787 | 0 | let mut acc2 = _mm512_setzero_ps(); |
788 | 0 | let mut acc3 = _mm512_setzero_ps(); |
789 | 0 | let mut acc4 = _mm512_setzero_ps(); |
790 | 0 | let mut acc5 = _mm512_setzero_ps(); |
791 | 0 | let mut acc6 = _mm512_setzero_ps(); |
792 | 0 | let mut acc7 = _mm512_setzero_ps(); |
793 | | |
794 | | // Main loop: Process 16 elements at a time |
795 | 0 | for i in 0..chunks { |
796 | 0 | let offset = i * 16; |
797 | 0 |
|
798 | 0 | // Load B column (reused for all 8 A rows) |
799 | 0 | let b_vec = _mm512_loadu_ps(b_col.as_ptr().add(offset)); |
800 | 0 |
|
801 | 0 | // Load A rows and FMA (Fused Multiply-Add) |
802 | 0 | let a0_vec = _mm512_loadu_ps(a_rows[0].as_ptr().add(offset)); |
803 | 0 | acc0 = _mm512_fmadd_ps(a0_vec, b_vec, acc0); |
804 | 0 |
|
805 | 0 | let a1_vec = _mm512_loadu_ps(a_rows[1].as_ptr().add(offset)); |
806 | 0 | acc1 = _mm512_fmadd_ps(a1_vec, b_vec, acc1); |
807 | 0 |
|
808 | 0 | let a2_vec = _mm512_loadu_ps(a_rows[2].as_ptr().add(offset)); |
809 | 0 | acc2 = _mm512_fmadd_ps(a2_vec, b_vec, acc2); |
810 | 0 |
|
811 | 0 | let a3_vec = _mm512_loadu_ps(a_rows[3].as_ptr().add(offset)); |
812 | 0 | acc3 = _mm512_fmadd_ps(a3_vec, b_vec, acc3); |
813 | 0 |
|
814 | 0 | let a4_vec = _mm512_loadu_ps(a_rows[4].as_ptr().add(offset)); |
815 | 0 | acc4 = _mm512_fmadd_ps(a4_vec, b_vec, acc4); |
816 | 0 |
|
817 | 0 | let a5_vec = _mm512_loadu_ps(a_rows[5].as_ptr().add(offset)); |
818 | 0 | acc5 = _mm512_fmadd_ps(a5_vec, b_vec, acc5); |
819 | 0 |
|
820 | 0 | let a6_vec = _mm512_loadu_ps(a_rows[6].as_ptr().add(offset)); |
821 | 0 | acc6 = _mm512_fmadd_ps(a6_vec, b_vec, acc6); |
822 | 0 |
|
823 | 0 | let a7_vec = _mm512_loadu_ps(a_rows[7].as_ptr().add(offset)); |
824 | 0 | acc7 = _mm512_fmadd_ps(a7_vec, b_vec, acc7); |
825 | 0 | } |
826 | | |
827 | | // Horizontal sum of each accumulator (reduce 16 elements to 1) |
828 | 0 | results[0] = _mm512_reduce_add_ps(acc0); |
829 | 0 | results[1] = _mm512_reduce_add_ps(acc1); |
830 | 0 | results[2] = _mm512_reduce_add_ps(acc2); |
831 | 0 | results[3] = _mm512_reduce_add_ps(acc3); |
832 | 0 | results[4] = _mm512_reduce_add_ps(acc4); |
833 | 0 | results[5] = _mm512_reduce_add_ps(acc5); |
834 | 0 | results[6] = _mm512_reduce_add_ps(acc6); |
835 | 0 | results[7] = _mm512_reduce_add_ps(acc7); |
836 | | |
837 | | // Handle remainder elements with scalar code |
838 | 0 | let remainder_start = chunks * 16; |
839 | 0 | if remainder_start < len { |
840 | 0 | for i in remainder_start..len { |
841 | 0 | results[0] += a_rows[0][i] * b_col[i]; |
842 | 0 | results[1] += a_rows[1][i] * b_col[i]; |
843 | 0 | results[2] += a_rows[2][i] * b_col[i]; |
844 | 0 | results[3] += a_rows[3][i] * b_col[i]; |
845 | 0 | results[4] += a_rows[4][i] * b_col[i]; |
846 | 0 | results[5] += a_rows[5][i] * b_col[i]; |
847 | 0 | results[6] += a_rows[6][i] * b_col[i]; |
848 | 0 | results[7] += a_rows[7][i] * b_col[i]; |
849 | 0 | } |
850 | 0 | } |
851 | 0 | } |
852 | | |
853 | | /// Cache-aware blocked matrix multiplication with SIMD optimization |
854 | | /// |
855 | | /// Uses 2-level cache blocking (L2/L1) to minimize cache misses: |
856 | | /// - L2 blocks: 64×64 (256KB for 3 matrices in f32) |
857 | | /// - L1 micro-kernels: 8×8 (768 bytes fits comfortably in L1) |
858 | | /// |
859 | | /// Performance characteristics: |
860 | | /// - Small matrices (<64×64): ~1.2× speedup over naive (overhead dominates) |
861 | | /// - Medium matrices (128×128): ~1.5-2× speedup (cache effects visible) |
862 | | /// - Large matrices (512×512+): ~2-3× speedup (dramatic cache improvement) |
863 | | /// |
864 | | /// This is Phase 1 of matmul optimization (Issue #10). Future Phase 2 will |
865 | | /// add optional BLAS backend for full NumPy parity on very large matrices. |
866 | | /// Helper function to process a single L3 row block for parallel matmul (Phase 4). |
867 | | /// |
868 | | /// # Safety |
869 | | /// When called from parallel code, the caller must ensure that each thread processes |
870 | | /// a distinct row range [iii, i3_end) with no overlap. This function is safe because |
871 | | /// each thread writes only to its own row range in the result matrix. |
872 | | #[cfg(feature = "parallel")] |
873 | | #[allow(clippy::too_many_arguments)] |
874 | | fn process_l3_row_block_seq( |
875 | | iii: usize, |
876 | | i3_end: usize, |
877 | | a: &Matrix<f32>, |
878 | | b_transposed: &Matrix<f32>, |
879 | | result: &mut Matrix<f32>, |
880 | | l2_block_size: usize, |
881 | | l3_block_size: usize, |
882 | | ) { |
883 | | #[cfg(target_arch = "x86_64")] |
884 | | use crate::backends::{avx2::Avx2Backend, sse2::Sse2Backend}; |
885 | | use crate::backends::{scalar::ScalarBackend, VectorBackend}; |
886 | | |
887 | | // Process all column blocks for this row block |
888 | | for jjj in (0..b_transposed.rows).step_by(l3_block_size) { |
889 | | let j3_end = (jjj + l3_block_size).min(b_transposed.rows); |
890 | | |
891 | | for kkk in (0..a.cols).step_by(l3_block_size) { |
892 | | let k3_end = (kkk + l3_block_size).min(a.cols); |
893 | | |
894 | | // L2 blocking within L3 blocks |
895 | | for ii in (iii..i3_end).step_by(l2_block_size) { |
896 | | let i_end = (ii + l2_block_size).min(i3_end); |
897 | | |
898 | | for jj in (jjj..j3_end).step_by(l2_block_size) { |
899 | | let j_end = (jj + l2_block_size).min(j3_end); |
900 | | |
901 | | for kk in (kkk..k3_end).step_by(l2_block_size) { |
902 | | let k_end = (kk + l2_block_size).min(k3_end); |
903 | | let block_size = k_end - kk; |
904 | | |
905 | | // Micro-kernel processing |
906 | | #[cfg(target_arch = "x86_64")] |
907 | | let use_microkernel = |
908 | | matches!(a.backend, Backend::AVX2 | Backend::AVX512); |
909 | | |
910 | | #[cfg(target_arch = "x86_64")] |
911 | | if use_microkernel { |
912 | | let mut i = ii; |
913 | | |
914 | | // Process 4 rows at a time with micro-kernel |
915 | | while i + 4 <= i_end { |
916 | | let row0_start = i * a.cols + kk; |
917 | | let row1_start = (i + 1) * a.cols + kk; |
918 | | let row2_start = (i + 2) * a.cols + kk; |
919 | | let row3_start = (i + 3) * a.cols + kk; |
920 | | |
921 | | let a_rows = [ |
922 | | &a.data[row0_start..row0_start + block_size], |
923 | | &a.data[row1_start..row1_start + block_size], |
924 | | &a.data[row2_start..row2_start + block_size], |
925 | | &a.data[row3_start..row3_start + block_size], |
926 | | ]; |
927 | | |
928 | | for j in jj..j_end { |
929 | | let col_start = j * b_transposed.cols + kk; |
930 | | let b_col = |
931 | | &b_transposed.data[col_start..col_start + block_size]; |
932 | | |
933 | | let mut partial_dots = [0.0f32; 4]; |
934 | | // SAFETY: AVX2 support verified by is_x86_feature_detected!("avx2") |
935 | | // check in outer scope. Slices a_rows and b_col are bounds-checked |
936 | | // and properly aligned for SIMD operations. |
937 | | // SAFETY: CPU feature verified at runtime, slices bounds-checked |
938 | | unsafe { |
939 | | Matrix::matmul_microkernel_4x1_avx2( |
940 | | a_rows, |
941 | | b_col, |
942 | | &mut partial_dots, |
943 | | ); |
944 | | } |
945 | | |
946 | | result.data[i * result.cols + j] += partial_dots[0]; |
947 | | result.data[(i + 1) * result.cols + j] += partial_dots[1]; |
948 | | result.data[(i + 2) * result.cols + j] += partial_dots[2]; |
949 | | result.data[(i + 3) * result.cols + j] += partial_dots[3]; |
950 | | } |
951 | | |
952 | | i += 4; |
953 | | } |
954 | | |
955 | | // Handle remaining rows (< 4) |
956 | | for i in i..i_end { |
957 | | let row_start = i * a.cols + kk; |
958 | | let a_row = &a.data[row_start..row_start + block_size]; |
959 | | |
960 | | for j in jj..j_end { |
961 | | let col_start = j * b_transposed.cols + kk; |
962 | | let b_col = |
963 | | &b_transposed.data[col_start..col_start + block_size]; |
964 | | |
965 | | // SAFETY: AVX2 verified at runtime, slices bounds-checked |
966 | | // SAFETY: AVX2 verified at runtime, slices bounds-checked |
967 | | let partial_dot = unsafe { Avx2Backend::dot(a_row, b_col) }; |
968 | | result.data[i * result.cols + j] += partial_dot; |
969 | | } |
970 | | } |
971 | | } else { |
972 | | // Non-AVX2 path |
973 | | #[allow(unused_variables)] |
974 | | for i in ii..i_end { |
975 | | let row_start = i * a.cols + kk; |
976 | | let a_row = &a.data[row_start..row_start + block_size]; |
977 | | |
978 | | for j in jj..j_end { |
979 | | let col_start = j * b_transposed.cols + kk; |
980 | | let b_col = |
981 | | &b_transposed.data[col_start..col_start + block_size]; |
982 | | |
983 | | // SAFETY: AVX2 verified at runtime, slices bounds-checked |
984 | | let partial_dot = unsafe { |
985 | | match a.backend { |
986 | | Backend::Scalar => ScalarBackend::dot(a_row, b_col), |
987 | | #[cfg(target_arch = "x86_64")] |
988 | | Backend::SSE2 | Backend::AVX => { |
989 | | Sse2Backend::dot(a_row, b_col) |
990 | | } |
991 | | #[cfg(not(target_arch = "x86_64"))] |
992 | | Backend::SSE2 |
993 | | | Backend::AVX |
994 | | | Backend::AVX2 |
995 | | | Backend::AVX512 => { |
996 | | ScalarBackend::dot(a_row, b_col) |
997 | | } |
998 | | #[cfg(any( |
999 | | target_arch = "aarch64", |
1000 | | target_arch = "arm" |
1001 | | ))] |
1002 | | Backend::NEON => { |
1003 | | use crate::backends::neon::NeonBackend; |
1004 | | NeonBackend::dot(a_row, b_col) |
1005 | | } |
1006 | | #[cfg(not(any( |
1007 | | target_arch = "aarch64", |
1008 | | target_arch = "arm" |
1009 | | )))] |
1010 | | Backend::NEON => ScalarBackend::dot(a_row, b_col), |
1011 | | #[cfg(target_arch = "wasm32")] |
1012 | | Backend::WasmSIMD => { |
1013 | | use crate::backends::wasm::WasmBackend; |
1014 | | WasmBackend::dot(a_row, b_col) |
1015 | | } |
1016 | | #[cfg(not(target_arch = "wasm32"))] |
1017 | | Backend::WasmSIMD => { |
1018 | | ScalarBackend::dot(a_row, b_col) |
1019 | | } |
1020 | | // Catch-all for GPU, Auto, and any other backends |
1021 | | _ => ScalarBackend::dot(a_row, b_col), |
1022 | | } |
1023 | | }; |
1024 | | |
1025 | | result.data[i * result.cols + j] += partial_dot; |
1026 | | } |
1027 | | } |
1028 | | } |
1029 | | |
1030 | | // Non-x86_64 fallback |
1031 | | #[cfg(not(target_arch = "x86_64"))] |
1032 | | { |
1033 | | for i in ii..i_end { |
1034 | | let row_start = i * a.cols + kk; |
1035 | | let a_row = &a.data[row_start..row_start + block_size]; |
1036 | | |
1037 | | for j in jj..j_end { |
1038 | | let col_start = j * b_transposed.cols + kk; |
1039 | | let b_col = |
1040 | | &b_transposed.data[col_start..col_start + block_size]; |
1041 | | |
1042 | | // SAFETY: AVX2 verified at runtime, slices bounds-checked |
1043 | | let partial_dot = unsafe { |
1044 | | match a.backend { |
1045 | | Backend::Scalar => ScalarBackend::dot(a_row, b_col), |
1046 | | #[cfg(any( |
1047 | | target_arch = "aarch64", |
1048 | | target_arch = "arm" |
1049 | | ))] |
1050 | | Backend::NEON => { |
1051 | | use crate::backends::neon::NeonBackend; |
1052 | | NeonBackend::dot(a_row, b_col) |
1053 | | } |
1054 | | #[cfg(not(any( |
1055 | | target_arch = "aarch64", |
1056 | | target_arch = "arm" |
1057 | | )))] |
1058 | | Backend::NEON => ScalarBackend::dot(a_row, b_col), |
1059 | | #[cfg(target_arch = "wasm32")] |
1060 | | Backend::WasmSIMD => { |
1061 | | use crate::backends::wasm::WasmBackend; |
1062 | | WasmBackend::dot(a_row, b_col) |
1063 | | } |
1064 | | #[cfg(not(target_arch = "wasm32"))] |
1065 | | Backend::WasmSIMD => { |
1066 | | ScalarBackend::dot(a_row, b_col) |
1067 | | } |
1068 | | _ => ScalarBackend::dot(a_row, b_col), |
1069 | | } |
1070 | | }; |
1071 | | |
1072 | | result.data[i * result.cols + j] += partial_dot; |
1073 | | } |
1074 | | } |
1075 | | } |
1076 | | } |
1077 | | } |
1078 | | } |
1079 | | } |
1080 | | } |
1081 | | } |
1082 | | |
1083 | 0 | fn matmul_simd( |
1084 | 0 | &self, |
1085 | 0 | other: &Matrix<f32>, |
1086 | 0 | result: &mut Matrix<f32>, |
1087 | 0 | ) -> Result<(), TruenoError> { |
1088 | | // Cache blocking parameters (tuned for typical x86_64 CPUs) |
1089 | | // L2 cache: 256KB typical → 64K f32 elements → 64×64×3 matrices fits |
1090 | | const L2_BLOCK_SIZE: usize = 64; |
1091 | | // L3 cache: 4-16MB typical → 256×256 blocks for very large matrices (Phase 3) |
1092 | | const L3_BLOCK_SIZE: usize = 256; |
1093 | | const L3_THRESHOLD: usize = 512; // Use 3-level blocking for matrices ≥512×512 |
1094 | | |
1095 | | // For small matrices, use simple SIMD approach (blocking overhead too high) |
1096 | 0 | if self.rows <= 32 || self.cols <= 32 || other.cols <= 32 { |
1097 | 0 | return self.matmul_simd_simple(other, result); |
1098 | 0 | } |
1099 | | |
1100 | | #[cfg(target_arch = "x86_64")] |
1101 | | use crate::backends::{avx2::Avx2Backend, sse2::Sse2Backend}; |
1102 | | use crate::backends::{scalar::ScalarBackend, VectorBackend}; |
1103 | | |
1104 | | // Pre-transpose B for better cache locality (columns become rows) |
1105 | 0 | let b_transposed = other.transpose(); |
1106 | | |
1107 | | // Determine if we should use 3-level blocking (Phase 3) |
1108 | 0 | let use_l3_blocking = |
1109 | 0 | self.rows >= L3_THRESHOLD && self.cols >= L3_THRESHOLD && other.cols >= L3_THRESHOLD; |
1110 | | |
1111 | | // Phase 4: Determine if we should use multi-threading (≥1024×1024) |
1112 | | #[cfg(feature = "parallel")] |
1113 | | const PARALLEL_THRESHOLD: usize = 1024; |
1114 | | #[cfg(feature = "parallel")] |
1115 | | let use_parallel = self.rows >= PARALLEL_THRESHOLD |
1116 | | && self.cols >= PARALLEL_THRESHOLD |
1117 | | && other.cols >= PARALLEL_THRESHOLD; |
1118 | | #[cfg(not(feature = "parallel"))] |
1119 | 0 | let use_parallel = false; |
1120 | | |
1121 | 0 | if use_l3_blocking { |
1122 | | // ===== Phase 3/4: 3-Level Cache Blocking (L3 → L2 → micro-kernel) ===== |
1123 | | // For very large matrices (≥512×512), use L3 cache blocking to minimize |
1124 | | // cache misses when data doesn't fit in L2 cache |
1125 | | // |
1126 | | // Hierarchy: |
1127 | | // 1. L3 blocks: 256×256 (fits in L3 cache: 4-16MB) |
1128 | | // 2. L2 blocks: 64×64 (fits in L2 cache: 256KB) |
1129 | | // 3. Micro-kernel: 4×1 for AVX2/AVX512 |
1130 | | // |
1131 | | // Phase 4: For ≥1024×1024, parallelize L3 row blocks with rayon |
1132 | | |
1133 | 0 | if use_parallel { |
1134 | | // ===== Phase 4: Parallel 3-Level Cache Blocking (Lock-Free Row Partitioning) ===== |
1135 | | #[cfg(feature = "parallel")] |
1136 | | { |
1137 | | use rayon::prelude::*; |
1138 | | use std::sync::atomic::{AtomicPtr, Ordering}; |
1139 | | use std::sync::Arc; |
1140 | | |
1141 | | // Lock-free parallelization strategy: |
1142 | | // Each thread processes one L3 row block (256 rows). Since row blocks are |
1143 | | // non-overlapping, threads write to distinct memory regions with no contention. |
1144 | | // |
1145 | | // Safety invariant: Each thread writes to result.data[iii*cols..(i3_end)*cols], |
1146 | | // where iii = block_idx * L3_BLOCK_SIZE. Since L3 blocks don't overlap, |
1147 | | // no two threads write to the same memory location. |
1148 | | |
1149 | | // Store result pointer in Arc<AtomicPtr> for safe sharing |
1150 | | let result_ptr = Arc::new(AtomicPtr::new(result as *mut Matrix<f32>)); |
1151 | | |
1152 | | // Calculate number of L3 blocks |
1153 | | let num_blocks = self.rows.div_ceil(L3_BLOCK_SIZE); |
1154 | | |
1155 | | // Process each L3 block in parallel (lock-free) |
1156 | | (0..num_blocks).into_par_iter().for_each(|block_idx| { |
1157 | | let iii = block_idx * L3_BLOCK_SIZE; |
1158 | | let i3_end = (iii + L3_BLOCK_SIZE).min(self.rows); |
1159 | | |
1160 | | // SAFETY: Each thread processes a distinct row range [iii, i3_end). |
1161 | | // No two threads write to overlapping memory locations because: |
1162 | | // 1. L3 blocks partition rows: [0, 256), [256, 512), etc. |
1163 | | // 2. Each thread only modifies result.data[iii*cols..(i3_end)*cols] |
1164 | | // 3. Row ranges are non-overlapping by construction |
1165 | | // 4. All threads complete before function returns (rayon guarantee) |
1166 | | // 5. AtomicPtr ensures proper memory ordering across threads |
1167 | | // SAFETY: CPU feature verified at runtime, slices bounds-checked |
1168 | | unsafe { |
1169 | | let ptr = result_ptr.load(Ordering::Relaxed); |
1170 | | Self::process_l3_row_block_seq( |
1171 | | iii, |
1172 | | i3_end, |
1173 | | self, |
1174 | | &b_transposed, |
1175 | | &mut *ptr, |
1176 | | L2_BLOCK_SIZE, |
1177 | | L3_BLOCK_SIZE, |
1178 | | ); |
1179 | | } |
1180 | | }); |
1181 | | } |
1182 | | |
1183 | 0 | return Ok(()); |
1184 | 0 | } |
1185 | | |
1186 | | // ===== Sequential 3-Level Cache Blocking (fallback) ===== |
1187 | 0 | for iii in (0..self.rows).step_by(L3_BLOCK_SIZE) { |
1188 | 0 | let i3_end = (iii + L3_BLOCK_SIZE).min(self.rows); |
1189 | | |
1190 | 0 | for jjj in (0..other.cols).step_by(L3_BLOCK_SIZE) { |
1191 | 0 | let j3_end = (jjj + L3_BLOCK_SIZE).min(other.cols); |
1192 | | |
1193 | 0 | for kkk in (0..self.cols).step_by(L3_BLOCK_SIZE) { |
1194 | 0 | let k3_end = (kkk + L3_BLOCK_SIZE).min(self.cols); |
1195 | | |
1196 | | // L2 blocking within L3 blocks |
1197 | 0 | for ii in (iii..i3_end).step_by(L2_BLOCK_SIZE) { |
1198 | 0 | let i_end = (ii + L2_BLOCK_SIZE).min(i3_end); |
1199 | | |
1200 | 0 | for jj in (jjj..j3_end).step_by(L2_BLOCK_SIZE) { |
1201 | 0 | let j_end = (jj + L2_BLOCK_SIZE).min(j3_end); |
1202 | | |
1203 | 0 | for kk in (kkk..k3_end).step_by(L2_BLOCK_SIZE) { |
1204 | 0 | let k_end = (kk + L2_BLOCK_SIZE).min(k3_end); |
1205 | 0 | let block_size = k_end - kk; |
1206 | | |
1207 | | // Micro-kernel processing |
1208 | | #[cfg(target_arch = "x86_64")] |
1209 | 0 | let use_avx512 = matches!(self.backend, Backend::AVX512); |
1210 | | #[cfg(target_arch = "x86_64")] |
1211 | 0 | let use_avx2 = matches!(self.backend, Backend::AVX2); |
1212 | | |
1213 | | #[cfg(target_arch = "x86_64")] |
1214 | 0 | if use_avx512 { |
1215 | | // AVX-512 8x1 micro-kernel (Phase 3) |
1216 | 0 | let mut i = ii; |
1217 | | |
1218 | | // Process 8 rows at a time with AVX-512 micro-kernel |
1219 | 0 | while i + 8 <= i_end { |
1220 | 0 | let a_rows = [ |
1221 | 0 | &self.data[i * self.cols + kk..(i * self.cols + kk) + block_size], |
1222 | 0 | &self.data[(i + 1) * self.cols + kk..((i + 1) * self.cols + kk) + block_size], |
1223 | 0 | &self.data[(i + 2) * self.cols + kk..((i + 2) * self.cols + kk) + block_size], |
1224 | 0 | &self.data[(i + 3) * self.cols + kk..((i + 3) * self.cols + kk) + block_size], |
1225 | 0 | &self.data[(i + 4) * self.cols + kk..((i + 4) * self.cols + kk) + block_size], |
1226 | 0 | &self.data[(i + 5) * self.cols + kk..((i + 5) * self.cols + kk) + block_size], |
1227 | 0 | &self.data[(i + 6) * self.cols + kk..((i + 6) * self.cols + kk) + block_size], |
1228 | 0 | &self.data[(i + 7) * self.cols + kk..((i + 7) * self.cols + kk) + block_size], |
1229 | 0 | ]; |
1230 | | |
1231 | 0 | for j in jj..j_end { |
1232 | 0 | let col_start = j * b_transposed.cols + kk; |
1233 | 0 | let b_col = &b_transposed.data |
1234 | 0 | [col_start..col_start + block_size]; |
1235 | 0 |
|
1236 | 0 | let mut partial_dots = [0.0f32; 8]; |
1237 | 0 | // SAFETY: CPU feature verified at runtime, slices bounds-checked |
1238 | 0 | unsafe { |
1239 | 0 | Self::matmul_microkernel_8x1_avx512( |
1240 | 0 | a_rows, |
1241 | 0 | b_col, |
1242 | 0 | &mut partial_dots, |
1243 | 0 | ); |
1244 | 0 | } |
1245 | 0 |
|
1246 | 0 | result.data[i * result.cols + j] += partial_dots[0]; |
1247 | 0 | result.data[(i + 1) * result.cols + j] += partial_dots[1]; |
1248 | 0 | result.data[(i + 2) * result.cols + j] += partial_dots[2]; |
1249 | 0 | result.data[(i + 3) * result.cols + j] += partial_dots[3]; |
1250 | 0 | result.data[(i + 4) * result.cols + j] += partial_dots[4]; |
1251 | 0 | result.data[(i + 5) * result.cols + j] += partial_dots[5]; |
1252 | 0 | result.data[(i + 6) * result.cols + j] += partial_dots[6]; |
1253 | 0 | result.data[(i + 7) * result.cols + j] += partial_dots[7]; |
1254 | 0 | } |
1255 | | |
1256 | 0 | i += 8; |
1257 | | } |
1258 | | |
1259 | | // Handle remaining rows with AVX2 4x1 kernel |
1260 | 0 | while i + 4 <= i_end { |
1261 | 0 | let a_rows = [ |
1262 | 0 | &self.data[i * self.cols + kk..(i * self.cols + kk) + block_size], |
1263 | 0 | &self.data[(i + 1) * self.cols + kk..((i + 1) * self.cols + kk) + block_size], |
1264 | 0 | &self.data[(i + 2) * self.cols + kk..((i + 2) * self.cols + kk) + block_size], |
1265 | 0 | &self.data[(i + 3) * self.cols + kk..((i + 3) * self.cols + kk) + block_size], |
1266 | 0 | ]; |
1267 | | |
1268 | 0 | for j in jj..j_end { |
1269 | 0 | let col_start = j * b_transposed.cols + kk; |
1270 | 0 | let b_col = &b_transposed.data[col_start..col_start + block_size]; |
1271 | 0 |
|
1272 | 0 | let mut partial_dots = [0.0f32; 4]; |
1273 | 0 | // SAFETY: CPU feature verified at runtime, slices bounds-checked |
1274 | 0 | unsafe { |
1275 | 0 | Self::matmul_microkernel_4x1_avx2(a_rows, b_col, &mut partial_dots); |
1276 | 0 | } |
1277 | 0 |
|
1278 | 0 | result.data[i * result.cols + j] += partial_dots[0]; |
1279 | 0 | result.data[(i + 1) * result.cols + j] += partial_dots[1]; |
1280 | 0 | result.data[(i + 2) * result.cols + j] += partial_dots[2]; |
1281 | 0 | result.data[(i + 3) * result.cols + j] += partial_dots[3]; |
1282 | 0 | } |
1283 | 0 | i += 4; |
1284 | | } |
1285 | | |
1286 | | // Handle remaining rows (< 4) |
1287 | 0 | for i in i..i_end { |
1288 | 0 | let row_start = i * self.cols + kk; |
1289 | 0 | let a_row = &self.data[row_start..row_start + block_size]; |
1290 | | |
1291 | 0 | for j in jj..j_end { |
1292 | 0 | let col_start = j * b_transposed.cols + kk; |
1293 | 0 | let b_col = &b_transposed.data[col_start..col_start + block_size]; |
1294 | 0 |
|
1295 | 0 | // SAFETY: AVX2 verified at runtime, slices bounds-checked |
1296 | 0 | let partial_dot = unsafe { Avx2Backend::dot(a_row, b_col) }; |
1297 | 0 | result.data[i * result.cols + j] += partial_dot; |
1298 | 0 | } |
1299 | | } |
1300 | 0 | } else if use_avx2 { |
1301 | | // AVX2 4x1 micro-kernel |
1302 | 0 | let mut i = ii; |
1303 | | |
1304 | | // Process 4 rows at a time with micro-kernel |
1305 | 0 | while i + 4 <= i_end { |
1306 | 0 | let row0_start = i * self.cols + kk; |
1307 | 0 | let row1_start = (i + 1) * self.cols + kk; |
1308 | 0 | let row2_start = (i + 2) * self.cols + kk; |
1309 | 0 | let row3_start = (i + 3) * self.cols + kk; |
1310 | | |
1311 | 0 | let a_rows = [ |
1312 | 0 | &self.data[row0_start..row0_start + block_size], |
1313 | 0 | &self.data[row1_start..row1_start + block_size], |
1314 | 0 | &self.data[row2_start..row2_start + block_size], |
1315 | 0 | &self.data[row3_start..row3_start + block_size], |
1316 | 0 | ]; |
1317 | | |
1318 | 0 | for j in jj..j_end { |
1319 | 0 | let col_start = j * b_transposed.cols + kk; |
1320 | 0 | let b_col = &b_transposed.data |
1321 | 0 | [col_start..col_start + block_size]; |
1322 | 0 |
|
1323 | 0 | let mut partial_dots = [0.0f32; 4]; |
1324 | 0 | // SAFETY: CPU feature verified at runtime, slices bounds-checked |
1325 | 0 | unsafe { |
1326 | 0 | Self::matmul_microkernel_4x1_avx2( |
1327 | 0 | a_rows, |
1328 | 0 | b_col, |
1329 | 0 | &mut partial_dots, |
1330 | 0 | ); |
1331 | 0 | } |
1332 | 0 |
|
1333 | 0 | result.data[i * result.cols + j] += partial_dots[0]; |
1334 | 0 | result.data[(i + 1) * result.cols + j] += |
1335 | 0 | partial_dots[1]; |
1336 | 0 | result.data[(i + 2) * result.cols + j] += |
1337 | 0 | partial_dots[2]; |
1338 | 0 | result.data[(i + 3) * result.cols + j] += |
1339 | 0 | partial_dots[3]; |
1340 | 0 | } |
1341 | | |
1342 | 0 | i += 4; |
1343 | | } |
1344 | | |
1345 | | // Handle remaining rows (< 4) |
1346 | 0 | for i in i..i_end { |
1347 | 0 | let row_start = i * self.cols + kk; |
1348 | 0 | let a_row = |
1349 | 0 | &self.data[row_start..row_start + block_size]; |
1350 | | |
1351 | 0 | for j in jj..j_end { |
1352 | 0 | let col_start = j * b_transposed.cols + kk; |
1353 | 0 | let b_col = &b_transposed.data |
1354 | 0 | [col_start..col_start + block_size]; |
1355 | 0 |
|
1356 | 0 | let partial_dot = |
1357 | 0 | // SAFETY: CPU feature verified at runtime, slices bounds-checked |
1358 | 0 | unsafe { Avx2Backend::dot(a_row, b_col) }; |
1359 | 0 | result.data[i * result.cols + j] += partial_dot; |
1360 | 0 | } |
1361 | | } |
1362 | | } else { |
1363 | | // Non-AVX2 path |
1364 | | #[allow(unused_variables)] |
1365 | 0 | for i in ii..i_end { |
1366 | 0 | let row_start = i * self.cols + kk; |
1367 | 0 | let a_row = |
1368 | 0 | &self.data[row_start..row_start + block_size]; |
1369 | | |
1370 | 0 | for j in jj..j_end { |
1371 | 0 | let col_start = j * b_transposed.cols + kk; |
1372 | 0 | let b_col = &b_transposed.data |
1373 | 0 | [col_start..col_start + block_size]; |
1374 | | |
1375 | | // SAFETY: AVX2 verified at runtime, slices bounds-checked |
1376 | 0 | let partial_dot = unsafe { |
1377 | 0 | match self.backend { |
1378 | | Backend::Scalar => { |
1379 | 0 | ScalarBackend::dot(a_row, b_col) |
1380 | | } |
1381 | | #[cfg(target_arch = "x86_64")] |
1382 | | Backend::SSE2 | Backend::AVX => { |
1383 | 0 | Sse2Backend::dot(a_row, b_col) |
1384 | | } |
1385 | | #[cfg(not(target_arch = "x86_64"))] |
1386 | | Backend::SSE2 |
1387 | | | Backend::AVX |
1388 | | | Backend::AVX2 |
1389 | | | Backend::AVX512 => { |
1390 | | ScalarBackend::dot(a_row, b_col) |
1391 | | } |
1392 | | #[cfg(any( |
1393 | | target_arch = "aarch64", |
1394 | | target_arch = "arm" |
1395 | | ))] |
1396 | | Backend::NEON => { |
1397 | | use crate::backends::neon::NeonBackend; |
1398 | | NeonBackend::dot(a_row, b_col) |
1399 | | } |
1400 | | #[cfg(not(any( |
1401 | | target_arch = "aarch64", |
1402 | | target_arch = "arm" |
1403 | | )))] |
1404 | | Backend::NEON => { |
1405 | 0 | ScalarBackend::dot(a_row, b_col) |
1406 | | } |
1407 | | #[cfg(target_arch = "wasm32")] |
1408 | | Backend::WasmSIMD => { |
1409 | | use crate::backends::wasm::WasmBackend; |
1410 | | WasmBackend::dot(a_row, b_col) |
1411 | | } |
1412 | | #[cfg(not(target_arch = "wasm32"))] |
1413 | | Backend::WasmSIMD => { |
1414 | 0 | ScalarBackend::dot(a_row, b_col) |
1415 | | } |
1416 | | Backend::GPU |
1417 | | | Backend::Auto |
1418 | | | Backend::AVX2 |
1419 | | | Backend::AVX512 => { |
1420 | 0 | ScalarBackend::dot(a_row, b_col) |
1421 | | } |
1422 | | } |
1423 | | }; |
1424 | | |
1425 | 0 | result.data[i * result.cols + j] += partial_dot; |
1426 | | } |
1427 | | } |
1428 | | } |
1429 | | |
1430 | | // Non-x86_64 platforms |
1431 | | #[cfg(not(target_arch = "x86_64"))] |
1432 | | for i in ii..i_end { |
1433 | | let row_start = i * self.cols + kk; |
1434 | | let a_row = &self.data[row_start..row_start + block_size]; |
1435 | | |
1436 | | for j in jj..j_end { |
1437 | | let col_start = j * b_transposed.cols + kk; |
1438 | | let b_col = &b_transposed.data |
1439 | | [col_start..col_start + block_size]; |
1440 | | |
1441 | | // SAFETY: AVX2 verified at runtime, slices bounds-checked |
1442 | | let partial_dot = unsafe { |
1443 | | match self.backend { |
1444 | | Backend::Scalar => { |
1445 | | ScalarBackend::dot(a_row, b_col) |
1446 | | } |
1447 | | #[cfg(any( |
1448 | | target_arch = "aarch64", |
1449 | | target_arch = "arm" |
1450 | | ))] |
1451 | | Backend::NEON => { |
1452 | | use crate::backends::neon::NeonBackend; |
1453 | | NeonBackend::dot(a_row, b_col) |
1454 | | } |
1455 | | #[cfg(not(any( |
1456 | | target_arch = "aarch64", |
1457 | | target_arch = "arm" |
1458 | | )))] |
1459 | | Backend::NEON => { |
1460 | | ScalarBackend::dot(a_row, b_col) |
1461 | | } |
1462 | | #[cfg(target_arch = "wasm32")] |
1463 | | Backend::WasmSIMD => { |
1464 | | use crate::backends::wasm::WasmBackend; |
1465 | | WasmBackend::dot(a_row, b_col) |
1466 | | } |
1467 | | #[cfg(not(target_arch = "wasm32"))] |
1468 | | Backend::WasmSIMD => { |
1469 | | ScalarBackend::dot(a_row, b_col) |
1470 | | } |
1471 | | _ => ScalarBackend::dot(a_row, b_col), |
1472 | | } |
1473 | | }; |
1474 | | |
1475 | | result.data[i * result.cols + j] += partial_dot; |
1476 | | } |
1477 | | } |
1478 | | } |
1479 | | } |
1480 | | } |
1481 | | } |
1482 | | } |
1483 | | } |
1484 | | } else { |
1485 | | // ===== Phase 1/2: 2-Level Cache Blocking (L2 → micro-kernel) ===== |
1486 | | // For medium matrices (32-512), use original 2-level blocking |
1487 | | // |
1488 | | // This path preserves the fast performance for 256×256 and smaller matrices |
1489 | | // by avoiding the overhead of 3-level loop nesting |
1490 | | |
1491 | 0 | for ii in (0..self.rows).step_by(L2_BLOCK_SIZE) { |
1492 | 0 | let i_end = (ii + L2_BLOCK_SIZE).min(self.rows); |
1493 | | |
1494 | 0 | for jj in (0..other.cols).step_by(L2_BLOCK_SIZE) { |
1495 | 0 | let j_end = (jj + L2_BLOCK_SIZE).min(other.cols); |
1496 | | |
1497 | 0 | for kk in (0..self.cols).step_by(L2_BLOCK_SIZE) { |
1498 | 0 | let k_end = (kk + L2_BLOCK_SIZE).min(self.cols); |
1499 | 0 | let block_size = k_end - kk; |
1500 | | |
1501 | | // Inner loops: Process L2 block with micro-kernel (Phase 2) or SIMD |
1502 | | #[cfg(target_arch = "x86_64")] |
1503 | 0 | let use_microkernel = |
1504 | 0 | matches!(self.backend, Backend::AVX2 | Backend::AVX512); |
1505 | | |
1506 | | #[cfg(target_arch = "x86_64")] |
1507 | 0 | if use_microkernel { |
1508 | | // Phase 2: Use 4×1 micro-kernel for AVX2/AVX512 |
1509 | 0 | let mut i = ii; |
1510 | | |
1511 | | // Process 4 rows at a time with micro-kernel |
1512 | 0 | while i + 4 <= i_end { |
1513 | | // Get 4 consecutive rows of A |
1514 | 0 | let row0_start = i * self.cols + kk; |
1515 | 0 | let row1_start = (i + 1) * self.cols + kk; |
1516 | 0 | let row2_start = (i + 2) * self.cols + kk; |
1517 | 0 | let row3_start = (i + 3) * self.cols + kk; |
1518 | | |
1519 | 0 | let a_rows = [ |
1520 | 0 | &self.data[row0_start..row0_start + block_size], |
1521 | 0 | &self.data[row1_start..row1_start + block_size], |
1522 | 0 | &self.data[row2_start..row2_start + block_size], |
1523 | 0 | &self.data[row3_start..row3_start + block_size], |
1524 | 0 | ]; |
1525 | | |
1526 | | // Process each column of B with the micro-kernel |
1527 | 0 | for j in jj..j_end { |
1528 | 0 | let col_start = j * b_transposed.cols + kk; |
1529 | 0 | let b_col = |
1530 | 0 | &b_transposed.data[col_start..col_start + block_size]; |
1531 | 0 |
|
1532 | 0 | // Compute 4 dot products simultaneously |
1533 | 0 | let mut partial_dots = [0.0f32; 4]; |
1534 | 0 | // SAFETY: CPU feature verified at runtime, slices bounds-checked |
1535 | 0 | unsafe { |
1536 | 0 | Self::matmul_microkernel_4x1_avx2( |
1537 | 0 | a_rows, |
1538 | 0 | b_col, |
1539 | 0 | &mut partial_dots, |
1540 | 0 | ); |
1541 | 0 | } |
1542 | 0 |
|
1543 | 0 | // Accumulate results |
1544 | 0 | result.data[i * result.cols + j] += partial_dots[0]; |
1545 | 0 | result.data[(i + 1) * result.cols + j] += partial_dots[1]; |
1546 | 0 | result.data[(i + 2) * result.cols + j] += partial_dots[2]; |
1547 | 0 | result.data[(i + 3) * result.cols + j] += partial_dots[3]; |
1548 | 0 | } |
1549 | | |
1550 | 0 | i += 4; |
1551 | | } |
1552 | | |
1553 | | // Handle remaining rows (< 4) with standard path |
1554 | 0 | for i in i..i_end { |
1555 | 0 | let row_start = i * self.cols + kk; |
1556 | 0 | let a_row = &self.data[row_start..row_start + block_size]; |
1557 | | |
1558 | 0 | for j in jj..j_end { |
1559 | 0 | let col_start = j * b_transposed.cols + kk; |
1560 | 0 | let b_col = |
1561 | 0 | &b_transposed.data[col_start..col_start + block_size]; |
1562 | 0 |
|
1563 | 0 | // SAFETY: AVX2 verified at runtime, slices bounds-checked |
1564 | 0 | let partial_dot = unsafe { Avx2Backend::dot(a_row, b_col) }; |
1565 | 0 | result.data[i * result.cols + j] += partial_dot; |
1566 | 0 | } |
1567 | | } |
1568 | | } else { |
1569 | | // Phase 1: Standard SIMD path (non-AVX2 backends) |
1570 | | #[allow(unused_variables)] |
1571 | 0 | for i in ii..i_end { |
1572 | 0 | let row_start = i * self.cols + kk; |
1573 | 0 | let a_row = &self.data[row_start..row_start + block_size]; |
1574 | | |
1575 | 0 | for j in jj..j_end { |
1576 | 0 | let col_start = j * b_transposed.cols + kk; |
1577 | 0 | let b_col = |
1578 | 0 | &b_transposed.data[col_start..col_start + block_size]; |
1579 | | |
1580 | | // SAFETY: AVX2 verified at runtime, slices bounds-checked |
1581 | 0 | let partial_dot = unsafe { |
1582 | 0 | match self.backend { |
1583 | 0 | Backend::Scalar => ScalarBackend::dot(a_row, b_col), |
1584 | | #[cfg(target_arch = "x86_64")] |
1585 | | Backend::SSE2 | Backend::AVX => { |
1586 | 0 | Sse2Backend::dot(a_row, b_col) |
1587 | | } |
1588 | | #[cfg(not(target_arch = "x86_64"))] |
1589 | | Backend::SSE2 |
1590 | | | Backend::AVX |
1591 | | | Backend::AVX2 |
1592 | | | Backend::AVX512 => ScalarBackend::dot(a_row, b_col), |
1593 | | #[cfg(any( |
1594 | | target_arch = "aarch64", |
1595 | | target_arch = "arm" |
1596 | | ))] |
1597 | | Backend::NEON => { |
1598 | | use crate::backends::neon::NeonBackend; |
1599 | | NeonBackend::dot(a_row, b_col) |
1600 | | } |
1601 | | #[cfg(not(any( |
1602 | | target_arch = "aarch64", |
1603 | | target_arch = "arm" |
1604 | | )))] |
1605 | 0 | Backend::NEON => ScalarBackend::dot(a_row, b_col), |
1606 | | #[cfg(target_arch = "wasm32")] |
1607 | | Backend::WasmSIMD => { |
1608 | | use crate::backends::wasm::WasmBackend; |
1609 | | WasmBackend::dot(a_row, b_col) |
1610 | | } |
1611 | | #[cfg(not(target_arch = "wasm32"))] |
1612 | 0 | Backend::WasmSIMD => ScalarBackend::dot(a_row, b_col), |
1613 | | Backend::GPU |
1614 | | | Backend::Auto |
1615 | | | Backend::AVX2 |
1616 | 0 | | Backend::AVX512 => ScalarBackend::dot(a_row, b_col), |
1617 | | } |
1618 | | }; |
1619 | | |
1620 | 0 | result.data[i * result.cols + j] += partial_dot; |
1621 | | } |
1622 | | } |
1623 | | } |
1624 | | |
1625 | | // Non-x86_64 platforms: Use standard SIMD path |
1626 | | #[cfg(not(target_arch = "x86_64"))] |
1627 | | for i in ii..i_end { |
1628 | | let row_start = i * self.cols + kk; |
1629 | | let a_row = &self.data[row_start..row_start + block_size]; |
1630 | | |
1631 | | for j in jj..j_end { |
1632 | | let col_start = j * b_transposed.cols + kk; |
1633 | | let b_col = &b_transposed.data[col_start..col_start + block_size]; |
1634 | | |
1635 | | // SAFETY: AVX2 verified at runtime, slices bounds-checked |
1636 | | let partial_dot = unsafe { |
1637 | | match self.backend { |
1638 | | Backend::Scalar => ScalarBackend::dot(a_row, b_col), |
1639 | | #[cfg(any(target_arch = "aarch64", target_arch = "arm"))] |
1640 | | Backend::NEON => { |
1641 | | use crate::backends::neon::NeonBackend; |
1642 | | NeonBackend::dot(a_row, b_col) |
1643 | | } |
1644 | | #[cfg(not(any( |
1645 | | target_arch = "aarch64", |
1646 | | target_arch = "arm" |
1647 | | )))] |
1648 | | Backend::NEON => ScalarBackend::dot(a_row, b_col), |
1649 | | #[cfg(target_arch = "wasm32")] |
1650 | | Backend::WasmSIMD => { |
1651 | | use crate::backends::wasm::WasmBackend; |
1652 | | WasmBackend::dot(a_row, b_col) |
1653 | | } |
1654 | | #[cfg(not(target_arch = "wasm32"))] |
1655 | | Backend::WasmSIMD => ScalarBackend::dot(a_row, b_col), |
1656 | | _ => ScalarBackend::dot(a_row, b_col), |
1657 | | } |
1658 | | }; |
1659 | | |
1660 | | result.data[i * result.cols + j] += partial_dot; |
1661 | | } |
1662 | | } |
1663 | | } |
1664 | | } |
1665 | | } |
1666 | | } |
1667 | | |
1668 | 0 | Ok(()) |
1669 | 0 | } |
1670 | | |
1671 | | /// Simple SIMD matrix multiplication without blocking (for small matrices) |
1672 | | /// |
1673 | | /// This is the pre-blocking implementation that works well for small matrices |
1674 | | /// where cache blocking overhead exceeds benefits. |
1675 | 0 | fn matmul_simd_simple( |
1676 | 0 | &self, |
1677 | 0 | other: &Matrix<f32>, |
1678 | 0 | result: &mut Matrix<f32>, |
1679 | 0 | ) -> Result<(), TruenoError> { |
1680 | | #[cfg(target_arch = "x86_64")] |
1681 | | use crate::backends::{avx2::Avx2Backend, sse2::Sse2Backend}; |
1682 | | use crate::backends::{scalar::ScalarBackend, VectorBackend}; |
1683 | | |
1684 | | // Pre-transpose B for better cache locality |
1685 | 0 | let b_transposed = other.transpose(); |
1686 | | |
1687 | 0 | for i in 0..self.rows { |
1688 | 0 | let row_start = i * self.cols; |
1689 | 0 | let row_end = row_start + self.cols; |
1690 | 0 | let a_row = &self.data[row_start..row_end]; |
1691 | | |
1692 | 0 | for j in 0..other.cols { |
1693 | 0 | let col_start = j * b_transposed.cols; |
1694 | 0 | let col_end = col_start + b_transposed.cols; |
1695 | 0 | let b_col = &b_transposed.data[col_start..col_end]; |
1696 | | |
1697 | | // Compute dot product using SIMD backend directly |
1698 | | // SAFETY: Backend dot() maintains safety invariants |
1699 | 0 | let dot_result = unsafe { |
1700 | 0 | match self.backend { |
1701 | 0 | Backend::Scalar => ScalarBackend::dot(a_row, b_col), |
1702 | | #[cfg(target_arch = "x86_64")] |
1703 | 0 | Backend::SSE2 | Backend::AVX => Sse2Backend::dot(a_row, b_col), |
1704 | | #[cfg(target_arch = "x86_64")] |
1705 | 0 | Backend::AVX2 | Backend::AVX512 => Avx2Backend::dot(a_row, b_col), |
1706 | | #[cfg(not(target_arch = "x86_64"))] |
1707 | | Backend::SSE2 | Backend::AVX | Backend::AVX2 | Backend::AVX512 => { |
1708 | | ScalarBackend::dot(a_row, b_col) |
1709 | | } |
1710 | | #[cfg(any(target_arch = "aarch64", target_arch = "arm"))] |
1711 | | Backend::NEON => { |
1712 | | use crate::backends::neon::NeonBackend; |
1713 | | NeonBackend::dot(a_row, b_col) |
1714 | | } |
1715 | | #[cfg(not(any(target_arch = "aarch64", target_arch = "arm")))] |
1716 | 0 | Backend::NEON => ScalarBackend::dot(a_row, b_col), |
1717 | | #[cfg(target_arch = "wasm32")] |
1718 | | Backend::WasmSIMD => { |
1719 | | use crate::backends::wasm::WasmBackend; |
1720 | | WasmBackend::dot(a_row, b_col) |
1721 | | } |
1722 | | #[cfg(not(target_arch = "wasm32"))] |
1723 | 0 | Backend::WasmSIMD => ScalarBackend::dot(a_row, b_col), |
1724 | 0 | Backend::GPU | Backend::Auto => ScalarBackend::dot(a_row, b_col), |
1725 | | } |
1726 | | }; |
1727 | | |
1728 | 0 | result.data[i * result.cols + j] = dot_result; |
1729 | | } |
1730 | | } |
1731 | | |
1732 | 0 | Ok(()) |
1733 | 0 | } |
1734 | | |
1735 | | /// WASM-optimized tiled matrix multiplication with SIMD inner loop |
1736 | | /// |
1737 | | /// Key optimizations: |
1738 | | /// 1. NO transpose - avoids O(n²) memory allocation and copy |
1739 | | /// 2. Tiled blocking with SIMD-aligned tile widths |
1740 | | /// 3. Inner j-loop uses SIMD (B rows are contiguous in memory) |
1741 | | /// 4. Register accumulation to minimize memory traffic |
1742 | | /// |
1743 | | /// Performance: Targets <30ms for 384×74×384 (Whisper encoder attention) |
1744 | 0 | fn matmul_wasm_tiled( |
1745 | 0 | &self, |
1746 | 0 | other: &Matrix<f32>, |
1747 | 0 | result: &mut Matrix<f32>, |
1748 | 0 | ) -> Result<(), TruenoError> { |
1749 | 0 | let m = self.rows; |
1750 | 0 | let k = self.cols; |
1751 | 0 | let n = other.cols; |
1752 | | |
1753 | | // For each row of A |
1754 | 0 | for i in 0..m { |
1755 | 0 | let a_row_start = i * k; |
1756 | 0 | let result_row_start = i * n; |
1757 | | |
1758 | | // For each column of B, compute dot product A[i,:] · B[:,j] |
1759 | | // BUT: B[:,j] is not contiguous. Instead, iterate over k and accumulate. |
1760 | | // |
1761 | | // C[i,j] = Σ_k A[i,k] * B[k,j] |
1762 | | // |
1763 | | // For efficiency, broadcast A[i,k] and multiply with B[k, j0:j0+width] |
1764 | | // This uses SIMD on the contiguous B row segment. |
1765 | | |
1766 | | // Process output columns in SIMD-width chunks |
1767 | 0 | let simd_width = 8; // AVX2 processes 8 f32s |
1768 | 0 | let n_simd = (n / simd_width) * simd_width; |
1769 | | |
1770 | | // SIMD portion: columns 0..n_simd |
1771 | | // Note: Explicit indexing is intentional for LLVM auto-vectorization. |
1772 | | // Iterator patterns prevent the compiler from recognizing the SIMD pattern. |
1773 | | #[allow(clippy::needless_range_loop)] |
1774 | 0 | for j0 in (0..n_simd).step_by(simd_width) { |
1775 | 0 | let mut acc = [0.0f32; 8]; |
1776 | | |
1777 | 0 | for kk in 0..k { |
1778 | 0 | let a_val = self.data[a_row_start + kk]; |
1779 | 0 | let b_row_start = kk * n + j0; |
1780 | | |
1781 | | // Multiply a_val with B[kk, j0:j0+8] |
1782 | 0 | for jj in 0..simd_width { |
1783 | 0 | acc[jj] += a_val * other.data[b_row_start + jj]; |
1784 | 0 | } |
1785 | | } |
1786 | | |
1787 | | // Write accumulated results |
1788 | 0 | for jj in 0..simd_width { |
1789 | 0 | result.data[result_row_start + j0 + jj] = acc[jj]; |
1790 | 0 | } |
1791 | | } |
1792 | | |
1793 | | // Remainder columns (non-SIMD) |
1794 | 0 | for j in n_simd..n { |
1795 | 0 | let mut sum = 0.0f32; |
1796 | 0 | for kk in 0..k { |
1797 | 0 | sum += self.data[a_row_start + kk] * other.data[kk * n + j]; |
1798 | 0 | } |
1799 | 0 | result.data[result_row_start + j] = sum; |
1800 | | } |
1801 | | } |
1802 | | |
1803 | 0 | Ok(()) |
1804 | 0 | } |
1805 | | |
1806 | | /// GPU-accelerated matrix multiplication (very large matrices only) |
1807 | | #[cfg(all(feature = "gpu", not(target_arch = "wasm32")))] |
1808 | 0 | fn matmul_gpu(&self, other: &Matrix<f32>) -> Result<Matrix<f32>, TruenoError> { |
1809 | | use crate::backends::gpu::GpuBackend; |
1810 | | |
1811 | | // Check if GPU is available |
1812 | 0 | if !GpuBackend::is_available() { |
1813 | 0 | return Err(TruenoError::InvalidInput("GPU not available".to_string())); |
1814 | 0 | } |
1815 | | |
1816 | | // Create GPU backend |
1817 | 0 | let mut gpu = GpuBackend::new(); |
1818 | | |
1819 | | // Execute GPU matmul |
1820 | 0 | let result_data = gpu |
1821 | 0 | .matmul(&self.data, &other.data, self.rows, self.cols, other.cols) |
1822 | 0 | .map_err(|e| TruenoError::InvalidInput(format!("GPU matmul failed: {}", e)))?; |
1823 | | |
1824 | | // Create result matrix |
1825 | 0 | let mut result = Matrix::zeros(self.rows, other.cols); |
1826 | 0 | result.data = result_data; |
1827 | | |
1828 | 0 | Ok(result) |
1829 | 0 | } |
1830 | | |
1831 | | /// Transpose the matrix (swap rows and columns) |
1832 | | /// |
1833 | | /// Returns a new matrix where element `(i, j)` of the original becomes |
1834 | | /// element `(j, i)` in the result. |
1835 | | /// |
1836 | | /// # Returns |
1837 | | /// |
1838 | | /// A new matrix with dimensions swapped: if input is `m×n`, output is `n×m` |
1839 | | /// |
1840 | | /// # Example |
1841 | | /// |
1842 | | /// ``` |
1843 | | /// use trueno::Matrix; |
1844 | | /// |
1845 | | /// let m = Matrix::from_vec(2, 3, vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).unwrap(); |
1846 | | /// let t = m.transpose(); |
1847 | | /// |
1848 | | /// // [[1, 2, 3], [[1, 4], |
1849 | | /// // [4, 5, 6]] → [2, 5], |
1850 | | /// // [3, 6]] |
1851 | | /// assert_eq!(t.rows(), 3); |
1852 | | /// assert_eq!(t.cols(), 2); |
1853 | | /// assert_eq!(t.get(0, 0), Some(&1.0)); |
1854 | | /// assert_eq!(t.get(0, 1), Some(&4.0)); |
1855 | | /// assert_eq!(t.get(1, 0), Some(&2.0)); |
1856 | | /// ``` |
1857 | | #[cfg_attr(feature = "tracing", instrument(skip(self), fields(dims = %format!("{}x{}", self.rows, self.cols))))] |
1858 | 0 | pub fn transpose(&self) -> Matrix<f32> { |
1859 | 0 | let mut result = Matrix::zeros_with_backend(self.cols, self.rows, self.backend); |
1860 | | |
1861 | | // Use block-wise transpose for better cache locality |
1862 | | // Block size of 32 balances cache efficiency for both square and non-square matrices |
1863 | | const BLOCK_SIZE: usize = 32; |
1864 | | |
1865 | | // For non-square matrices, process output rows sequentially for write coalescing |
1866 | | // This ensures writes are sequential in memory regardless of input shape |
1867 | | // Fix for issue #65: non-square transpose was slow due to strided writes |
1868 | | |
1869 | | // Process in blocks, iterating output rows first for sequential writes |
1870 | 0 | for j_block in (0..self.cols).step_by(BLOCK_SIZE) { |
1871 | 0 | let j_end = (j_block + BLOCK_SIZE).min(self.cols); |
1872 | | |
1873 | 0 | for i_block in (0..self.rows).step_by(BLOCK_SIZE) { |
1874 | 0 | let i_end = (i_block + BLOCK_SIZE).min(self.rows); |
1875 | | |
1876 | | // Within block: iterate output rows (j) in outer loop for sequential writes |
1877 | 0 | for j in j_block..j_end { |
1878 | 0 | let dst_row_start = j * result.cols; |
1879 | 0 | for i in i_block..i_end { |
1880 | 0 | // result[j, i] = self[i, j] |
1881 | 0 | // Sequential write: dst_row_start + i increments by 1 |
1882 | 0 | // Strided read: acceptable, CPU prefetch handles this |
1883 | 0 | result.data[dst_row_start + i] = self.data[i * self.cols + j]; |
1884 | 0 | } |
1885 | | } |
1886 | | } |
1887 | | } |
1888 | | |
1889 | 0 | result |
1890 | 0 | } |
1891 | | |
1892 | | /// Matrix-vector multiplication (column vector): A × v |
1893 | | /// |
1894 | | /// Multiplies this matrix by a column vector, computing `A × v` where the result |
1895 | | /// is a column vector with length equal to the number of rows in `A`. |
1896 | | /// |
1897 | | /// # Mathematical Definition |
1898 | | /// |
1899 | | /// For an m×n matrix A and an n-dimensional vector v: |
1900 | | /// ```text |
1901 | | /// result[i] = Σ(j=0 to n-1) A[i,j] × v[j] |
1902 | | /// ``` |
1903 | | /// |
1904 | | /// # Arguments |
1905 | | /// |
1906 | | /// * `v` - Column vector with length equal to `self.cols()` |
1907 | | /// |
1908 | | /// # Returns |
1909 | | /// |
1910 | | /// A new vector with length `self.rows()` |
1911 | | /// |
1912 | | /// # Errors |
1913 | | /// |
1914 | | /// Returns `InvalidInput` if `v.len() != self.cols()` |
1915 | | /// |
1916 | | /// # Example |
1917 | | /// |
1918 | | /// ``` |
1919 | | /// use trueno::{Matrix, Vector}; |
1920 | | /// |
1921 | | /// let m = Matrix::from_vec(2, 3, vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).unwrap(); |
1922 | | /// let v = Vector::from_slice(&[1.0, 2.0, 3.0]); |
1923 | | /// let result = m.matvec(&v).unwrap(); |
1924 | | /// |
1925 | | /// // [[1, 2, 3] [1] [1×1 + 2×2 + 3×3] [14] |
1926 | | /// // [4, 5, 6]] × [2] = [4×1 + 5×2 + 6×3] = [32] |
1927 | | /// // [3] |
1928 | | /// assert_eq!(result.as_slice(), &[14.0, 32.0]); |
1929 | | /// ``` |
1930 | 0 | pub fn matvec(&self, v: &Vector<f32>) -> Result<Vector<f32>, TruenoError> { |
1931 | 0 | if v.len() != self.cols { |
1932 | 0 | return Err(TruenoError::InvalidInput(format!( |
1933 | 0 | "Vector length {} does not match matrix columns {} for matrix-vector multiplication", |
1934 | 0 | v.len(), |
1935 | 0 | self.cols |
1936 | 0 | ))); |
1937 | 0 | } |
1938 | | |
1939 | | #[cfg(target_arch = "x86_64")] |
1940 | | use crate::backends::{avx2::Avx2Backend, sse2::Sse2Backend}; |
1941 | | use crate::backends::{scalar::ScalarBackend, VectorBackend}; |
1942 | | |
1943 | 0 | let v_slice = v.as_slice(); |
1944 | | |
1945 | 0 | let mut result_data = vec![0.0; self.rows]; |
1946 | | |
1947 | | // Parallel execution for very large matrices (≥4096 rows) |
1948 | | // Note: Thread overhead dominates for smaller matrices |
1949 | | #[cfg(feature = "parallel")] |
1950 | | { |
1951 | | const PARALLEL_THRESHOLD: usize = 4096; |
1952 | | |
1953 | | if self.rows >= PARALLEL_THRESHOLD { |
1954 | | use rayon::prelude::*; |
1955 | | use std::sync::atomic::{AtomicPtr, Ordering}; |
1956 | | use std::sync::Arc; |
1957 | | |
1958 | | let result_ptr = Arc::new(AtomicPtr::new(result_data.as_mut_ptr())); |
1959 | | |
1960 | | // Process rows in parallel - each row computes an independent dot product |
1961 | | (0..self.rows).into_par_iter().for_each(|i| { |
1962 | | let row_start = i * self.cols; |
1963 | | let row = &self.data[row_start..(row_start + self.cols)]; |
1964 | | |
1965 | | // SAFETY: CPU feature verified at runtime, slices bounds-checked |
1966 | | let dot_result = unsafe { |
1967 | | #[cfg(target_arch = "x86_64")] |
1968 | | { |
1969 | | match self.backend { |
1970 | | Backend::AVX2 | Backend::AVX512 => Avx2Backend::dot(row, v_slice), |
1971 | | Backend::SSE2 | Backend::AVX => Sse2Backend::dot(row, v_slice), |
1972 | | _ => ScalarBackend::dot(row, v_slice), |
1973 | | } |
1974 | | } |
1975 | | #[cfg(not(target_arch = "x86_64"))] |
1976 | | { |
1977 | | ScalarBackend::dot(row, v_slice) |
1978 | | } |
1979 | | }; |
1980 | | |
1981 | | // Write to non-overlapping memory location (thread-safe) |
1982 | | // SAFETY: CPU feature verified at runtime, slices bounds-checked |
1983 | | unsafe { |
1984 | | let ptr = result_ptr.load(Ordering::Relaxed); |
1985 | | *ptr.add(i) = dot_result; |
1986 | | } |
1987 | | }); |
1988 | | |
1989 | | return Ok(Vector::from_slice(&result_data)); |
1990 | | } |
1991 | | } |
1992 | | |
1993 | | // SIMD-optimized execution: each row-vector product is a dot product |
1994 | 0 | for (i, result) in result_data.iter_mut().enumerate() { |
1995 | 0 | let row_start = i * self.cols; |
1996 | 0 | let row = &self.data[row_start..(row_start + self.cols)]; |
1997 | | |
1998 | | // Use SIMD dot product for each row |
1999 | | // SAFETY: CPU feature verified at runtime, slices bounds-checked |
2000 | | *result = unsafe { |
2001 | | #[cfg(target_arch = "x86_64")] |
2002 | | { |
2003 | 0 | match self.backend { |
2004 | 0 | Backend::AVX2 | Backend::AVX512 => Avx2Backend::dot(row, v_slice), |
2005 | 0 | Backend::SSE2 | Backend::AVX => Sse2Backend::dot(row, v_slice), |
2006 | 0 | _ => ScalarBackend::dot(row, v_slice), |
2007 | | } |
2008 | | } |
2009 | | #[cfg(not(target_arch = "x86_64"))] |
2010 | | { |
2011 | | ScalarBackend::dot(row, v_slice) |
2012 | | } |
2013 | | }; |
2014 | | } |
2015 | | |
2016 | 0 | Ok(Vector::from_slice(&result_data)) |
2017 | 0 | } |
2018 | | |
2019 | | /// Vector-matrix multiplication (row vector): v^T × A |
2020 | | /// |
2021 | | /// Multiplies a row vector by this matrix, computing `v^T × A` where the result |
2022 | | /// is a row vector with length equal to the number of columns in `A`. |
2023 | | /// |
2024 | | /// # Mathematical Definition |
2025 | | /// |
2026 | | /// For an m-dimensional vector v and an m×n matrix A: |
2027 | | /// ```text |
2028 | | /// result[j] = Σ(i=0 to m-1) v[i] × A[i,j] |
2029 | | /// ``` |
2030 | | /// |
2031 | | /// # Arguments |
2032 | | /// |
2033 | | /// * `v` - Row vector with length equal to `m.rows()` |
2034 | | /// * `m` - Matrix to multiply |
2035 | | /// |
2036 | | /// # Returns |
2037 | | /// |
2038 | | /// A new vector with length `m.cols()` |
2039 | | /// |
2040 | | /// # Errors |
2041 | | /// |
2042 | | /// Returns `InvalidInput` if `v.len() != m.rows()` |
2043 | | /// |
2044 | | /// # Example |
2045 | | /// |
2046 | | /// ``` |
2047 | | /// use trueno::{Matrix, Vector}; |
2048 | | /// |
2049 | | /// let m = Matrix::from_vec(2, 3, vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).unwrap(); |
2050 | | /// let v = Vector::from_slice(&[1.0, 2.0]); |
2051 | | /// let result = Matrix::vecmat(&v, &m).unwrap(); |
2052 | | /// |
2053 | | /// // [1, 2] × [[1, 2, 3] = [1×1 + 2×4, 1×2 + 2×5, 1×3 + 2×6] |
2054 | | /// // [4, 5, 6]] |
2055 | | /// // = [9, 12, 15] |
2056 | | /// assert_eq!(result.as_slice(), &[9.0, 12.0, 15.0]); |
2057 | | /// ``` |
2058 | 0 | pub fn vecmat(v: &Vector<f32>, m: &Matrix<f32>) -> Result<Vector<f32>, TruenoError> { |
2059 | 0 | if v.len() != m.rows { |
2060 | 0 | return Err(TruenoError::InvalidInput(format!( |
2061 | 0 | "Vector length {} does not match matrix rows {} for vector-matrix multiplication", |
2062 | 0 | v.len(), |
2063 | 0 | m.rows |
2064 | 0 | ))); |
2065 | 0 | } |
2066 | | |
2067 | | // SIMD-optimized implementation using row-wise accumulation |
2068 | | // Instead of column-wise access (cache-unfriendly), we compute: |
2069 | | // result = Σ(i) v[i] * row_i (cache-friendly, vectorizable) |
2070 | | // |
2071 | | // This approach: |
2072 | | // 1. Sequential row access (cache-friendly vs strided column access) |
2073 | | // 2. Uses SIMD scale and add operations |
2074 | | // 3. Leverages existing optimized Vector operations |
2075 | | |
2076 | 0 | let mut result = Vector::from_slice(&vec![0.0; m.cols]); |
2077 | 0 | let v_slice = v.as_slice(); |
2078 | | |
2079 | | // Accumulate each scaled row into result |
2080 | 0 | for (i, &scalar) in v_slice.iter().enumerate().take(m.rows) { |
2081 | 0 | let row_start = i * m.cols; |
2082 | 0 | let row = &m.data[row_start..(row_start + m.cols)]; |
2083 | | |
2084 | | // Create vector for this row |
2085 | 0 | let row_vec = Vector::from_slice(row); |
2086 | | |
2087 | | // result += scalar * row (using SIMD scale and add) |
2088 | 0 | let scaled_row = row_vec.scale(scalar)?; |
2089 | 0 | result = result.add(&scaled_row)?; |
2090 | | } |
2091 | | |
2092 | 0 | Ok(result) |
2093 | 0 | } |
2094 | | |
2095 | | /// Perform 2D convolution with a kernel |
2096 | | /// |
2097 | | /// Applies a 2D convolution operation using "valid" padding (no padding), |
2098 | | /// resulting in an output smaller than the input. |
2099 | | /// |
2100 | | /// # Arguments |
2101 | | /// |
2102 | | /// * `kernel` - Convolution kernel (filter) to apply |
2103 | | /// |
2104 | | /// # Returns |
2105 | | /// |
2106 | | /// Convolved matrix with dimensions: |
2107 | | /// - rows: `input.rows - kernel.rows + 1` |
2108 | | /// - cols: `input.cols - kernel.cols + 1` |
2109 | | /// |
2110 | | /// # Errors |
2111 | | /// |
2112 | | /// Returns `InvalidInput` if: |
2113 | | /// - Kernel is larger than input in any dimension |
2114 | | /// - Kernel has even dimensions (center pixel ambiguous) |
2115 | | /// |
2116 | | /// # Example |
2117 | | /// |
2118 | | /// ``` |
2119 | | /// use trueno::Matrix; |
2120 | | /// |
2121 | | /// // 5x5 input image |
2122 | | /// let input = Matrix::from_vec( |
2123 | | /// 5, 5, |
2124 | | /// vec![ |
2125 | | /// 0.0, 0.0, 0.0, 0.0, 0.0, |
2126 | | /// 0.0, 0.0, 0.0, 0.0, 0.0, |
2127 | | /// 0.0, 0.0, 9.0, 0.0, 0.0, |
2128 | | /// 0.0, 0.0, 0.0, 0.0, 0.0, |
2129 | | /// 0.0, 0.0, 0.0, 0.0, 0.0, |
2130 | | /// ] |
2131 | | /// ).unwrap(); |
2132 | | /// |
2133 | | /// // 3x3 averaging kernel |
2134 | | /// let kernel_val = 1.0 / 9.0; |
2135 | | /// let kernel = Matrix::from_vec( |
2136 | | /// 3, 3, |
2137 | | /// vec![kernel_val; 9] |
2138 | | /// ).unwrap(); |
2139 | | /// |
2140 | | /// let result = input.convolve2d(&kernel).unwrap(); |
2141 | | /// assert_eq!(result.rows(), 3); // 5 - 3 + 1 |
2142 | | /// assert_eq!(result.cols(), 3); |
2143 | | /// ``` |
2144 | | // ========================================================================= |
2145 | | // HOT PATH - PERFORMANCE CRITICAL |
2146 | | // ========================================================================= |
2147 | | // This function processes millions of elements for typical image sizes. |
2148 | | // Any changes to the inner loop REQUIRE benchmark verification: |
2149 | | // 1. Run: make bench-check |
2150 | | // 2. Verify no regression >10% |
2151 | | // |
2152 | | // PROHIBITED in inner loops: |
2153 | | // - .get() / .get_mut() (bounds checking overhead) |
2154 | | // - .expect() / .unwrap() (panic path overhead) |
2155 | | // - Iterator adaptors (closure overhead) |
2156 | | // |
2157 | | // Use direct indexing with bounds proof documented above the loop. |
2158 | | // ========================================================================= |
2159 | 0 | pub fn convolve2d(&self, kernel: &Matrix<f32>) -> Result<Matrix<f32>, TruenoError> { |
2160 | | // Validate kernel size |
2161 | 0 | if kernel.rows > self.rows || kernel.cols > self.cols { |
2162 | 0 | return Err(TruenoError::InvalidInput(format!( |
2163 | 0 | "Kernel size ({}x{}) larger than input ({}x{})", |
2164 | 0 | kernel.rows, kernel.cols, self.rows, self.cols |
2165 | 0 | ))); |
2166 | 0 | } |
2167 | | |
2168 | | // Calculate output dimensions (valid padding) |
2169 | 0 | let output_rows = self.rows - kernel.rows + 1; |
2170 | 0 | let output_cols = self.cols - kernel.cols + 1; |
2171 | | |
2172 | | // Initialize output matrix (reuse parent's backend) |
2173 | 0 | let mut result = Matrix::zeros_with_backend(output_rows, output_cols, self.backend); |
2174 | | |
2175 | | // Backend selection strategy: |
2176 | | // OpComplexity::High - GPU beneficial at >10K elements |
2177 | | // GPU for large images (output > 10K elements) |
2178 | | // Scalar for smaller images |
2179 | | |
2180 | | #[cfg(all(feature = "gpu", not(target_arch = "wasm32")))] |
2181 | | const GPU_THRESHOLD: usize = 10_000; |
2182 | | |
2183 | | // Try GPU first for large convolutions |
2184 | | #[cfg(all(feature = "gpu", not(target_arch = "wasm32")))] |
2185 | | { |
2186 | 0 | if output_rows * output_cols >= GPU_THRESHOLD { |
2187 | | use crate::backends::gpu::GpuBackend; |
2188 | | |
2189 | 0 | if GpuBackend::is_available() { |
2190 | 0 | if let Ok(gpu_result) = |
2191 | 0 | self.convolve2d_gpu(kernel, &mut result, output_rows, output_cols) |
2192 | | { |
2193 | 0 | return Ok(gpu_result); |
2194 | 0 | } |
2195 | | // Fall through to scalar if GPU fails |
2196 | 0 | } |
2197 | 0 | } |
2198 | | } |
2199 | | |
2200 | | // Scalar baseline implementation - optimized with direct indexing |
2201 | | // SAFETY invariant proof: |
2202 | | // - output_rows = self.rows - kernel.rows + 1 |
2203 | | // - output_cols = self.cols - kernel.cols + 1 |
2204 | | // - For any out_row < output_rows and k_row < kernel.rows: |
2205 | | // in_row = out_row + k_row < (self.rows - kernel.rows + 1) + kernel.rows - 1 = self.rows |
2206 | | // - Same logic applies to columns |
2207 | | // - All indices are provably within bounds, so we use direct indexing for performance |
2208 | | |
2209 | 0 | let input_data = self.as_slice(); |
2210 | 0 | let kernel_data = kernel.as_slice(); |
2211 | 0 | let result_data = result.data.as_mut_slice(); |
2212 | 0 | let input_cols = self.cols; |
2213 | 0 | let kernel_cols = kernel.cols; |
2214 | 0 | let result_cols = output_cols; |
2215 | | |
2216 | 0 | for out_row in 0..output_rows { |
2217 | 0 | for out_col in 0..output_cols { |
2218 | 0 | let mut sum = 0.0; |
2219 | | |
2220 | | // Apply kernel - use direct indexing for performance |
2221 | 0 | for k_row in 0..kernel.rows { |
2222 | 0 | let in_row = out_row + k_row; |
2223 | 0 | let input_row_offset = in_row * input_cols; |
2224 | 0 | let kernel_row_offset = k_row * kernel_cols; |
2225 | | |
2226 | 0 | for k_col in 0..kernel.cols { |
2227 | 0 | let in_col = out_col + k_col; |
2228 | 0 | sum += input_data[input_row_offset + in_col] |
2229 | 0 | * kernel_data[kernel_row_offset + k_col]; |
2230 | 0 | } |
2231 | | } |
2232 | | |
2233 | 0 | result_data[out_row * result_cols + out_col] = sum; |
2234 | | } |
2235 | | } |
2236 | | |
2237 | 0 | Ok(result) |
2238 | 0 | } |
2239 | | |
2240 | | /// GPU-accelerated 2D convolution helper |
2241 | | #[cfg(all(feature = "gpu", not(target_arch = "wasm32")))] |
2242 | 0 | fn convolve2d_gpu( |
2243 | 0 | &self, |
2244 | 0 | kernel: &Matrix<f32>, |
2245 | 0 | result: &mut Matrix<f32>, |
2246 | 0 | _output_rows: usize, |
2247 | 0 | _output_cols: usize, |
2248 | 0 | ) -> Result<Matrix<f32>, TruenoError> { |
2249 | | use crate::backends::gpu::GpuDevice; |
2250 | | |
2251 | 0 | let gpu = GpuDevice::new().map_err(TruenoError::InvalidInput)?; |
2252 | | |
2253 | 0 | gpu.convolve2d( |
2254 | 0 | self.as_slice(), |
2255 | 0 | kernel.as_slice(), |
2256 | 0 | result.data.as_mut_slice(), |
2257 | 0 | self.rows, |
2258 | 0 | self.cols, |
2259 | 0 | kernel.rows, |
2260 | 0 | kernel.cols, |
2261 | | ) |
2262 | 0 | .map_err(TruenoError::InvalidInput)?; |
2263 | | |
2264 | 0 | Ok(result.clone()) |
2265 | 0 | } |
2266 | | |
2267 | | /// Lookup embeddings by indices (Issue #61: ML primitives) |
2268 | | /// |
2269 | | /// Performs embedding lookup where self is the embedding table with shape |
2270 | | /// `[vocab_size, embed_dim]` and indices specify which rows to select. |
2271 | | /// |
2272 | | /// # Arguments |
2273 | | /// |
2274 | | /// * `indices` - Slice of indices into the embedding table |
2275 | | /// |
2276 | | /// # Returns |
2277 | | /// |
2278 | | /// A matrix with shape `[indices.len(), embed_dim]` containing the selected rows |
2279 | | /// |
2280 | | /// # Errors |
2281 | | /// |
2282 | | /// Returns `InvalidInput` if any index is out of bounds |
2283 | | /// |
2284 | | /// # Example |
2285 | | /// |
2286 | | /// ``` |
2287 | | /// use trueno::Matrix; |
2288 | | /// |
2289 | | /// // Create embedding table: 4 words, 3-dimensional embeddings |
2290 | | /// let embeddings = Matrix::from_vec(4, 3, vec![ |
2291 | | /// 1.0, 2.0, 3.0, // word 0 |
2292 | | /// 4.0, 5.0, 6.0, // word 1 |
2293 | | /// 7.0, 8.0, 9.0, // word 2 |
2294 | | /// 10.0, 11.0, 12.0 // word 3 |
2295 | | /// ]).unwrap(); |
2296 | | /// |
2297 | | /// // Lookup embeddings for indices [1, 3, 0] |
2298 | | /// let result = embeddings.embedding_lookup(&[1, 3, 0]).unwrap(); |
2299 | | /// |
2300 | | /// assert_eq!(result.rows(), 3); |
2301 | | /// assert_eq!(result.cols(), 3); |
2302 | | /// assert_eq!(result.get(0, 0), Some(&4.0)); // word 1 |
2303 | | /// assert_eq!(result.get(1, 0), Some(&10.0)); // word 3 |
2304 | | /// assert_eq!(result.get(2, 0), Some(&1.0)); // word 0 |
2305 | | /// ``` |
2306 | 0 | pub fn embedding_lookup(&self, indices: &[usize]) -> Result<Matrix<f32>, TruenoError> { |
2307 | | // Validate indices |
2308 | 0 | for (i, &idx) in indices.iter().enumerate() { |
2309 | 0 | if idx >= self.rows { |
2310 | 0 | return Err(TruenoError::InvalidInput(format!( |
2311 | 0 | "Index {} at position {} is out of bounds for embedding table with {} rows", |
2312 | 0 | idx, i, self.rows |
2313 | 0 | ))); |
2314 | 0 | } |
2315 | | } |
2316 | | |
2317 | | // Handle empty indices |
2318 | 0 | if indices.is_empty() { |
2319 | 0 | return Ok(Matrix::zeros_with_backend(0, self.cols, self.backend)); |
2320 | 0 | } |
2321 | | |
2322 | | // Allocate output matrix: [seq_len, embed_dim] |
2323 | 0 | let seq_len = indices.len(); |
2324 | 0 | let embed_dim = self.cols; |
2325 | 0 | let mut result = Matrix::zeros_with_backend(seq_len, embed_dim, self.backend); |
2326 | | |
2327 | | // Copy rows from embedding table to result |
2328 | 0 | for (out_row, &idx) in indices.iter().enumerate() { |
2329 | 0 | let src_start = idx * embed_dim; |
2330 | 0 | let dst_start = out_row * embed_dim; |
2331 | 0 |
|
2332 | 0 | // Copy entire row |
2333 | 0 | result.data[dst_start..dst_start + embed_dim] |
2334 | 0 | .copy_from_slice(&self.data[src_start..src_start + embed_dim]); |
2335 | 0 | } |
2336 | | |
2337 | 0 | Ok(result) |
2338 | 0 | } |
2339 | | |
2340 | | /// Lookup embeddings with gradient tracking support (for training) |
2341 | | /// |
2342 | | /// Returns both the embeddings and a sparse gradient accumulator. |
2343 | | /// This is useful for sparse gradient updates in training. |
2344 | | /// |
2345 | | /// # Arguments |
2346 | | /// |
2347 | | /// * `indices` - Slice of indices into the embedding table |
2348 | | /// |
2349 | | /// # Returns |
2350 | | /// |
2351 | | /// Tuple of (embeddings, unique_indices) where unique_indices can be used |
2352 | | /// for sparse gradient updates |
2353 | | /// |
2354 | | /// # Errors |
2355 | | /// |
2356 | | /// Returns `InvalidInput` if any index is out of bounds |
2357 | 0 | pub fn embedding_lookup_sparse( |
2358 | 0 | &self, |
2359 | 0 | indices: &[usize], |
2360 | 0 | ) -> Result<(Matrix<f32>, Vec<usize>), TruenoError> { |
2361 | 0 | let embeddings = self.embedding_lookup(indices)?; |
2362 | | |
2363 | | // Get unique indices for sparse gradient updates |
2364 | 0 | let mut unique: Vec<usize> = indices.to_vec(); |
2365 | 0 | unique.sort_unstable(); |
2366 | 0 | unique.dedup(); |
2367 | | |
2368 | 0 | Ok((embeddings, unique)) |
2369 | 0 | } |
2370 | | |
2371 | | /// 2D Max Pooling operation for CNN downsampling |
2372 | | /// |
2373 | | /// Applies max pooling over a 2D input tensor with specified kernel size and stride. |
2374 | | /// Input shape: (height, width), Output shape: ((height - kh) / sh + 1, (width - kw) / sw + 1) |
2375 | | /// |
2376 | | /// # Arguments |
2377 | | /// * `kernel` - (kernel_height, kernel_width) pooling window size |
2378 | | /// * `stride` - (stride_height, stride_width) step size |
2379 | | /// |
2380 | | /// # Examples |
2381 | | /// ``` |
2382 | | /// use trueno::matrix::Matrix; |
2383 | | /// let input = Matrix::from_vec(4, 4, vec![ |
2384 | | /// 1.0, 2.0, 3.0, 4.0, |
2385 | | /// 5.0, 6.0, 7.0, 8.0, |
2386 | | /// 9.0, 10.0, 11.0, 12.0, |
2387 | | /// 13.0, 14.0, 15.0, 16.0, |
2388 | | /// ]).unwrap(); |
2389 | | /// let pooled = input.max_pool2d((2, 2), (2, 2)).unwrap(); |
2390 | | /// assert_eq!(pooled.shape(), (2, 2)); |
2391 | | /// assert_eq!(pooled.get(0, 0), Some(&6.0)); // max of [1,2,5,6] |
2392 | | /// assert_eq!(pooled.get(1, 1), Some(&16.0)); // max of [11,12,15,16] |
2393 | | /// ``` |
2394 | 0 | pub fn max_pool2d( |
2395 | 0 | &self, |
2396 | 0 | kernel: (usize, usize), |
2397 | 0 | stride: (usize, usize), |
2398 | 0 | ) -> Result<Matrix<f32>, TruenoError> { |
2399 | 0 | let (kh, kw) = kernel; |
2400 | 0 | let (sh, sw) = stride; |
2401 | | |
2402 | 0 | if kh == 0 || kw == 0 || sh == 0 || sw == 0 { |
2403 | 0 | return Err(TruenoError::InvalidInput( |
2404 | 0 | "Kernel and stride dimensions must be positive".into(), |
2405 | 0 | )); |
2406 | 0 | } |
2407 | | |
2408 | 0 | if kh > self.rows || kw > self.cols { |
2409 | 0 | return Err(TruenoError::InvalidInput(format!( |
2410 | 0 | "Kernel size ({}, {}) larger than input ({}, {})", |
2411 | 0 | kh, kw, self.rows, self.cols |
2412 | 0 | ))); |
2413 | 0 | } |
2414 | | |
2415 | 0 | let out_h = (self.rows - kh) / sh + 1; |
2416 | 0 | let out_w = (self.cols - kw) / sw + 1; |
2417 | 0 | let mut result = Matrix::new(out_h, out_w); |
2418 | | |
2419 | 0 | for i in 0..out_h { |
2420 | 0 | for j in 0..out_w { |
2421 | 0 | let mut max_val = f32::NEG_INFINITY; |
2422 | 0 | for ki in 0..kh { |
2423 | 0 | for kj in 0..kw { |
2424 | 0 | let val = self.data[(i * sh + ki) * self.cols + (j * sw + kj)]; |
2425 | 0 | max_val = max_val.max(val); |
2426 | 0 | } |
2427 | | } |
2428 | 0 | result.data[i * out_w + j] = max_val; |
2429 | | } |
2430 | | } |
2431 | | |
2432 | 0 | Ok(result) |
2433 | 0 | } |
2434 | | |
2435 | | /// 2D Average Pooling operation for CNN downsampling |
2436 | | /// |
2437 | | /// Applies average pooling over a 2D input tensor with specified kernel size and stride. |
2438 | | /// Input shape: (height, width), Output shape: ((height - kh) / sh + 1, (width - kw) / sw + 1) |
2439 | | /// |
2440 | | /// # Arguments |
2441 | | /// * `kernel` - (kernel_height, kernel_width) pooling window size |
2442 | | /// * `stride` - (stride_height, stride_width) step size |
2443 | | /// |
2444 | | /// # Examples |
2445 | | /// ``` |
2446 | | /// use trueno::matrix::Matrix; |
2447 | | /// let input = Matrix::from_vec(4, 4, vec![ |
2448 | | /// 1.0, 2.0, 3.0, 4.0, |
2449 | | /// 5.0, 6.0, 7.0, 8.0, |
2450 | | /// 9.0, 10.0, 11.0, 12.0, |
2451 | | /// 13.0, 14.0, 15.0, 16.0, |
2452 | | /// ]).unwrap(); |
2453 | | /// let pooled = input.avg_pool2d((2, 2), (2, 2)).unwrap(); |
2454 | | /// assert_eq!(pooled.shape(), (2, 2)); |
2455 | | /// assert!((pooled.get(0, 0).unwrap() - 3.5).abs() < 1e-5); // avg of [1,2,5,6] |
2456 | | /// ``` |
2457 | 0 | pub fn avg_pool2d( |
2458 | 0 | &self, |
2459 | 0 | kernel: (usize, usize), |
2460 | 0 | stride: (usize, usize), |
2461 | 0 | ) -> Result<Matrix<f32>, TruenoError> { |
2462 | 0 | let (kh, kw) = kernel; |
2463 | 0 | let (sh, sw) = stride; |
2464 | | |
2465 | 0 | if kh == 0 || kw == 0 || sh == 0 || sw == 0 { |
2466 | 0 | return Err(TruenoError::InvalidInput( |
2467 | 0 | "Kernel and stride dimensions must be positive".into(), |
2468 | 0 | )); |
2469 | 0 | } |
2470 | | |
2471 | 0 | if kh > self.rows || kw > self.cols { |
2472 | 0 | return Err(TruenoError::InvalidInput(format!( |
2473 | 0 | "Kernel size ({}, {}) larger than input ({}, {})", |
2474 | 0 | kh, kw, self.rows, self.cols |
2475 | 0 | ))); |
2476 | 0 | } |
2477 | | |
2478 | 0 | let out_h = (self.rows - kh) / sh + 1; |
2479 | 0 | let out_w = (self.cols - kw) / sw + 1; |
2480 | 0 | let kernel_size = (kh * kw) as f32; |
2481 | 0 | let mut result = Matrix::new(out_h, out_w); |
2482 | | |
2483 | 0 | for i in 0..out_h { |
2484 | 0 | for j in 0..out_w { |
2485 | 0 | let mut sum = 0.0; |
2486 | 0 | for ki in 0..kh { |
2487 | 0 | for kj in 0..kw { |
2488 | 0 | sum += self.data[(i * sh + ki) * self.cols + (j * sw + kj)]; |
2489 | 0 | } |
2490 | | } |
2491 | 0 | result.data[i * out_w + j] = sum / kernel_size; |
2492 | | } |
2493 | | } |
2494 | | |
2495 | 0 | Ok(result) |
2496 | 0 | } |
2497 | | |
2498 | | /// Top-K selection: returns the k largest elements and their indices |
2499 | | /// |
2500 | | /// Useful for beam search, sampling, and ranking operations. |
2501 | | /// Searches row-major order and returns (values, indices) sorted descending. |
2502 | | /// |
2503 | | /// # Examples |
2504 | | /// ``` |
2505 | | /// use trueno::matrix::Matrix; |
2506 | | /// let m = Matrix::from_vec(2, 3, vec![1.0, 5.0, 3.0, 2.0, 6.0, 4.0]).unwrap(); |
2507 | | /// let (values, indices) = m.topk(2).unwrap(); |
2508 | | /// assert_eq!(values, vec![6.0, 5.0]); |
2509 | | /// assert_eq!(indices, vec![4, 1]); // flat indices |
2510 | | /// ``` |
2511 | 0 | pub fn topk(&self, k: usize) -> Result<(Vec<f32>, Vec<usize>), TruenoError> { |
2512 | 0 | if k == 0 { |
2513 | 0 | return Ok((vec![], vec![])); |
2514 | 0 | } |
2515 | | |
2516 | 0 | let k = k.min(self.data.len()); |
2517 | 0 | let mut indexed: Vec<(usize, f32)> = self.data.iter().copied().enumerate().collect(); |
2518 | | |
2519 | | // Partial sort - only sort k elements |
2520 | 0 | indexed.select_nth_unstable_by(k.saturating_sub(1), |a, b| { |
2521 | 0 | b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal) |
2522 | 0 | }); |
2523 | | |
2524 | 0 | indexed.truncate(k); |
2525 | 0 | indexed.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal)); |
2526 | | |
2527 | 0 | let values: Vec<f32> = indexed.iter().map(|(_, v)| *v).collect(); |
2528 | 0 | let indices: Vec<usize> = indexed.iter().map(|(i, _)| *i).collect(); |
2529 | | |
2530 | 0 | Ok((values, indices)) |
2531 | 0 | } |
2532 | | |
2533 | | /// Gather elements along axis using indices |
2534 | | /// |
2535 | | /// For 2D matrix with axis=0: output[i] = self[indices[i], :] |
2536 | | /// For 2D matrix with axis=1: output[:, i] = self[:, indices[i]] |
2537 | | /// |
2538 | | /// # Examples |
2539 | | /// ``` |
2540 | | /// use trueno::matrix::Matrix; |
2541 | | /// let m = Matrix::from_vec(3, 2, vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).unwrap(); |
2542 | | /// let gathered = m.gather(&[2, 0], 0).unwrap(); // Select rows 2 and 0 |
2543 | | /// assert_eq!(gathered.shape(), (2, 2)); |
2544 | | /// assert_eq!(gathered.get(0, 0), Some(&5.0)); // Row 2 |
2545 | | /// assert_eq!(gathered.get(1, 0), Some(&1.0)); // Row 0 |
2546 | | /// ``` |
2547 | 0 | pub fn gather(&self, indices: &[usize], axis: usize) -> Result<Matrix<f32>, TruenoError> { |
2548 | 0 | match axis { |
2549 | | 0 => { |
2550 | | // Gather rows |
2551 | 0 | let mut result = Matrix::new(indices.len(), self.cols); |
2552 | 0 | for (out_i, &idx) in indices.iter().enumerate() { |
2553 | 0 | if idx >= self.rows { |
2554 | 0 | return Err(TruenoError::InvalidInput(format!( |
2555 | 0 | "Index {} out of bounds for axis 0 with size {}", |
2556 | 0 | idx, self.rows |
2557 | 0 | ))); |
2558 | 0 | } |
2559 | 0 | for j in 0..self.cols { |
2560 | 0 | result.data[out_i * self.cols + j] = self.data[idx * self.cols + j]; |
2561 | 0 | } |
2562 | | } |
2563 | 0 | Ok(result) |
2564 | | } |
2565 | | 1 => { |
2566 | | // Gather columns |
2567 | 0 | let mut result = Matrix::new(self.rows, indices.len()); |
2568 | 0 | for i in 0..self.rows { |
2569 | 0 | for (out_j, &idx) in indices.iter().enumerate() { |
2570 | 0 | if idx >= self.cols { |
2571 | 0 | return Err(TruenoError::InvalidInput(format!( |
2572 | 0 | "Index {} out of bounds for axis 1 with size {}", |
2573 | 0 | idx, self.cols |
2574 | 0 | ))); |
2575 | 0 | } |
2576 | 0 | result.data[i * indices.len() + out_j] = self.data[i * self.cols + idx]; |
2577 | | } |
2578 | | } |
2579 | 0 | Ok(result) |
2580 | | } |
2581 | 0 | _ => Err(TruenoError::InvalidInput(format!( |
2582 | 0 | "Axis {} not supported for 2D matrix (use 0 or 1)", |
2583 | 0 | axis |
2584 | 0 | ))), |
2585 | | } |
2586 | 0 | } |
2587 | | |
2588 | | /// Pad matrix with a constant value |
2589 | | /// |
2590 | | /// # Arguments |
2591 | | /// * `padding` - ((top, bottom), (left, right)) padding amounts |
2592 | | /// * `value` - constant value to pad with (usually 0.0) |
2593 | | /// |
2594 | | /// # Examples |
2595 | | /// ``` |
2596 | | /// use trueno::matrix::Matrix; |
2597 | | /// let m = Matrix::from_vec(2, 2, vec![1.0, 2.0, 3.0, 4.0]).unwrap(); |
2598 | | /// let padded = m.pad(((1, 1), (1, 1)), 0.0).unwrap(); |
2599 | | /// assert_eq!(padded.shape(), (4, 4)); |
2600 | | /// assert_eq!(padded.get(0, 0), Some(&0.0)); // top-left padding |
2601 | | /// assert_eq!(padded.get(1, 1), Some(&1.0)); // original (0,0) |
2602 | | /// ``` |
2603 | 0 | pub fn pad( |
2604 | 0 | &self, |
2605 | 0 | padding: ((usize, usize), (usize, usize)), |
2606 | 0 | value: f32, |
2607 | 0 | ) -> Result<Matrix<f32>, TruenoError> { |
2608 | 0 | let ((top, bottom), (left, right)) = padding; |
2609 | 0 | let new_rows = self.rows + top + bottom; |
2610 | 0 | let new_cols = self.cols + left + right; |
2611 | | |
2612 | 0 | let mut result = Matrix::from_vec(new_rows, new_cols, vec![value; new_rows * new_cols])?; |
2613 | | |
2614 | | // Copy original data |
2615 | 0 | for i in 0..self.rows { |
2616 | 0 | for j in 0..self.cols { |
2617 | 0 | result.data[(i + top) * new_cols + (j + left)] = self.data[i * self.cols + j]; |
2618 | 0 | } |
2619 | | } |
2620 | | |
2621 | 0 | Ok(result) |
2622 | 0 | } |
2623 | | } |
2624 | | |
2625 | | |
2626 | | // Tests (~2.6K lines extracted for TDG compliance) |
2627 | | #[cfg(test)] |
2628 | | mod tests; |