xn_gemm/
lib.rs

1#![cfg_attr(not(feature = "std"), no_std)]
2#![warn(rust_2018_idioms)]
3
4mod gemm;
5
6#[cfg(feature = "f16")]
7pub use crate::gemm::f16;
8pub use crate::gemm::{c32, c64, gemm};
9pub use gemm_common::packed_cache;
10pub use gemm_common::Parallelism;
11
12pub use gemm_common::gemm::{
13    get_lhs_packing_threshold_multi_thread, get_lhs_packing_threshold_single_thread,
14    get_rhs_packing_threshold, get_threading_threshold, set_lhs_packing_threshold_multi_thread,
15    set_lhs_packing_threshold_single_thread, set_rhs_packing_threshold, set_threading_threshold,
16    DEFAULT_LHS_PACKING_THRESHOLD_MULTI_THREAD, DEFAULT_LHS_PACKING_THRESHOLD_SINGLE_THREAD,
17    DEFAULT_RHS_PACKING_THRESHOLD, DEFAULT_THREADING_THRESHOLD,
18};
19pub use gemm_common::{get_wasm_simd128, set_wasm_simd128, DEFAULT_WASM_SIMD128};
20
21#[cfg(test)]
22mod tests {
23    use super::*;
24    extern crate alloc;
25    use alloc::{vec, vec::Vec};
26    use num_traits::Float;
27
28    /// A cached lhs panel must give the same answer as packing per call. One test, since
29    /// enabling the cache is a process-wide switch.
30    #[test]
31    fn test_packed_lhs_cache_hit_matches_miss() {
32        // Row-major cases cover the widened path: uncached, a row-contiguous lhs is read in
33        // place rather than packed. 1536 rows is past 2 * mc, never prepacked before.
34        let cases = [
35            (64, 16, 128, false),
36            (96, 33, 100, false),
37            (512, 16, 384, true),
38            (1536, 16, 512, true),
39            (512, 128, 300, true),
40        ];
41        packed_cache::set_enabled(true);
42        packed_cache::set_budget_mb(1024);
43        packed_cache::clear();
44
45        for (m, n, k, row_major) in cases {
46            let (lhs_rs, lhs_cs) = if row_major { (k as isize, 1) } else { (1, m as isize) };
47            let a: Vec<f32> = (0..m * k).map(|_| rand::random()).collect();
48            let b: Vec<f32> = (0..k * n).map(|_| rand::random()).collect();
49            let mut want = vec![0f32; m * n];
50            let mut got = vec![0f32; m * n];
51            // SAFETY: shapes and strides match the buffers, which are sized to suit.
52            unsafe {
53                gemm::gemm_fallback(
54                    m, n, k, want.as_mut_ptr(), m as isize, 1, false, a.as_ptr(), lhs_cs,
55                    lhs_rs, b.as_ptr(), k as isize, 1, 0.0, 1.0,
56                );
57                // Three rounds: the probe, the call that packs, then the calls that hit.
58                for _ in 0..3 {
59                    for parallelism in [
60                        Parallelism::None,
61                        #[cfg(feature = "rayon")]
62                        Parallelism::Rayon(0),
63                    ] {
64                        gemm::gemm(
65                            m, n, k, got.as_mut_ptr(), m as isize, 1, false, a.as_ptr(),
66                            lhs_cs, lhs_rs, b.as_ptr(), k as isize, 1, 0.0, 1.0, false,
67                            false, false, parallelism,
68                        );
69                        for (c, d) in got.iter().zip(want.iter()) {
70                            assert_approx_eq::assert_approx_eq!(c, d, 1e-3);
71                        }
72                    }
73                }
74            }
75            assert!(packed_cache::stats().1 > 0, "the cached path was never taken");
76        }
77        packed_cache::clear();
78        packed_cache::set_enabled(false);
79
80        // With the switch off, only what the caller declares is cached. A weight is the rhs of
81        // the row-major product a linear layer issues, `x @ w^T`, and `gemm` transposes that
82        // into a product whose lhs is the weight: the declaration follows it there.
83        let (m, n, k) = (16, 512, 384);
84        let x: Vec<f32> = (0..m * k).map(|_| rand::random()).collect();
85        let w: Vec<f32> = (0..n * k).map(|_| rand::random()).collect();
86        let linear = |dst: &mut [f32]| unsafe {
87            // SAFETY: x is m-by-k and w is n-by-k, both row-major, and dst is m-by-n.
88            gemm::gemm(
89                m, n, k, dst.as_mut_ptr(), 1, n as isize, false, x.as_ptr(), 1, k as isize,
90                w.as_ptr(), k as isize, 1, 0.0, 1.0, false, false, false, Parallelism::None,
91            )
92        };
93        let mut want = vec![0f32; m * n];
94        linear(&mut want);
95        assert_eq!(packed_cache::stats().1, 0, "nothing is cached undeclared");
96        let in_w = packed_cache::Constant::new(w.as_ptr(), w.len());
97        for _ in 0..3 {
98            let mut got = vec![0f32; m * n];
99            packed_cache::with_constant(in_w, || linear(&mut got));
100            assert_eq!(got, want, "a cached panel is a copy: bit for bit the same");
101        }
102        assert_eq!(packed_cache::stats().1, 1, "the declared weight is cached");
103
104        // An operand written in place where the fingerprint does not look, as a key/value cache
105        // is: never cached while undeclared, even next to a declared one, so every call sees
106        // the write. `keys` is the rhs of `q @ keys^T`, so the lhs once transposed.
107        let (t, len, d) = (4, 256, 64);
108        let q: Vec<f32> = (0..t * d).map(|_| rand::random()).collect();
109        let mut keys: Vec<f32> = (0..len * d).map(|_| rand::random()).collect();
110        let scores = |keys: &[f32]| {
111            let mut dst = vec![0f32; t * len];
112            // SAFETY: q is t-by-d and keys is len-by-d, both row-major, and dst is t-by-len.
113            unsafe {
114                gemm::gemm(
115                    t, len, d, dst.as_mut_ptr(), 1, len as isize, false, q.as_ptr(), 1,
116                    d as isize, keys.as_ptr(), d as isize, 1, 0.0, 1.0, false, false, false,
117                    Parallelism::None,
118                )
119            };
120            dst
121        };
122        let in_q = packed_cache::Constant::new(q.as_ptr(), q.len());
123        for step in 0..4 {
124            // Row 1 is never sampled: the fingerprint reads rows `i * len / 8`.
125            keys[d + step] += 1.0;
126            let got = packed_cache::with_constant(in_q, || scores(&keys));
127            assert_eq!(got, scores(&keys.clone()), "step {step} saw the write");
128        }
129        assert_eq!(packed_cache::stats().1, 1, "only the declared weight holds a panel");
130        packed_cache::clear();
131    }
132
133    #[test]
134    fn test_gemm_f16() {
135        let mut mnks = vec![];
136        mnks.push((4, 4, 4));
137        mnks.push((63, 2, 10));
138        mnks.push((16, 2, 1));
139        mnks.push((0, 0, 4));
140        mnks.push((16, 1, 1));
141        mnks.push((16, 3, 1));
142        mnks.push((16, 4, 1));
143        mnks.push((16, 1, 2));
144        mnks.push((16, 2, 2));
145        mnks.push((16, 3, 2));
146        mnks.push((16, 4, 2));
147        mnks.push((16, 16, 1));
148        mnks.push((64, 64, 0));
149        mnks.push((256, 256, 256));
150        mnks.push((4096, 4096, 4));
151        mnks.push((64, 64, 4));
152        mnks.push((0, 64, 4));
153        mnks.push((64, 0, 4));
154        mnks.push((8, 16, 1));
155        mnks.push((16, 8, 1));
156        mnks.push((1, 1, 2));
157        mnks.push((1024, 1024, 1));
158        mnks.push((1024, 1024, 4));
159        mnks.push((63, 1, 10));
160        mnks.push((63, 3, 10));
161        mnks.push((63, 4, 10));
162        mnks.push((1, 63, 10));
163        mnks.push((2, 63, 10));
164        mnks.push((3, 63, 10));
165        mnks.push((4, 63, 10));
166
167        // gemv shapes above the threading threshold. k is kept modest because the
168        // f16 fallback accumulates in f16.
169        mnks.push((1, 2048, 384));
170        mnks.push((2048, 1, 384));
171        for (m, n, k) in mnks {
172            #[cfg(feature = "std")]
173            dbg!(m, n, k);
174            for parallelism in [
175                Parallelism::None,
176                #[cfg(feature = "rayon")]
177                Parallelism::Rayon(0),
178                #[cfg(feature = "rayon")]
179                Parallelism::Rayon(128),
180            ] {
181                for alpha in [0.0, 1.0, 2.3] {
182                    for beta in [0.0, 1.0, 2.3] {
183                        #[cfg(feature = "std")]
184                        dbg!(alpha, beta, parallelism);
185
186                        for colmajor in [true, false] {
187                            let alpha = f16::from_f32(alpha);
188                            let beta = f16::from_f32(beta);
189                            let a_vec: Vec<f16> = (0..(m * k))
190                                .map(|_| f16::from_f32(rand::random()))
191                                .collect();
192                            let b_vec: Vec<f16> = (0..(k * n))
193                                .map(|_| f16::from_f32(rand::random()))
194                                .collect();
195                            let mut c_vec: Vec<f16> = (0..(m * n))
196                                .map(|_| f16::from_f32(rand::random()))
197                                .collect();
198                            let mut d_vec = c_vec.clone();
199
200                            unsafe {
201                                gemm::gemm(
202                                    m,
203                                    n,
204                                    k,
205                                    c_vec.as_mut_ptr(),
206                                    if colmajor { m } else { 1 } as isize,
207                                    if colmajor { 1 } else { n } as isize,
208                                    true,
209                                    a_vec.as_ptr(),
210                                    m as isize,
211                                    1,
212                                    b_vec.as_ptr(),
213                                    k as isize,
214                                    1,
215                                    alpha,
216                                    beta,
217                                    false,
218                                    false,
219                                    false,
220                                    parallelism,
221                                );
222
223                                gemm::gemm_fallback(
224                                    m,
225                                    n,
226                                    k,
227                                    d_vec.as_mut_ptr(),
228                                    if colmajor { m } else { 1 } as isize,
229                                    if colmajor { 1 } else { n } as isize,
230                                    true,
231                                    a_vec.as_ptr(),
232                                    m as isize,
233                                    1,
234                                    b_vec.as_ptr(),
235                                    k as isize,
236                                    1,
237                                    alpha,
238                                    beta,
239                                );
240                            }
241                            let eps = f16::from_f32(1e-1);
242                            for (c, d) in c_vec.iter().zip(d_vec.iter()) {
243                                let eps_rel = c.abs() * eps;
244                                let eps_abs = eps;
245                                let eps = if eps_rel > eps_abs { eps_rel } else { eps_abs };
246                                assert_approx_eq::assert_approx_eq!(c, d, eps);
247                            }
248                        }
249                    }
250                }
251            }
252        }
253    }
254
255    #[test]
256    fn test_gemm_f32() {
257        set_wasm_simd128(true);
258
259        let mut mnks = vec![];
260        mnks.push((63, 2, 10));
261        mnks.push((1, 2, 10));
262        mnks.push((1, 63, 10));
263
264        // large m to trigger parallelized rhs packing with big number of threads and small n
265        mnks.push((2048, 255, 255));
266
267        mnks.push((256, 256, 256));
268        mnks.push((4096, 4096, 4));
269        mnks.push((64, 64, 4));
270        mnks.push((0, 64, 4));
271        mnks.push((64, 0, 4));
272        mnks.push((0, 0, 4));
273        mnks.push((64, 64, 0));
274        mnks.push((16, 1, 1));
275        mnks.push((16, 2, 1));
276        mnks.push((16, 3, 1));
277        mnks.push((16, 4, 1));
278        mnks.push((16, 1, 2));
279        mnks.push((16, 2, 2));
280        mnks.push((16, 3, 2));
281        mnks.push((16, 4, 2));
282        mnks.push((16, 16, 1));
283        mnks.push((8, 16, 1));
284        mnks.push((16, 8, 1));
285        mnks.push((1, 1, 2));
286        mnks.push((4, 4, 4));
287        mnks.push((1024, 1024, 1));
288        mnks.push((1024, 1024, 4));
289        mnks.push((63, 1, 10));
290        mnks.push((63, 3, 10));
291        mnks.push((63, 4, 10));
292        mnks.push((2, 63, 10));
293        mnks.push((3, 63, 10));
294        mnks.push((4, 63, 10));
295
296        // gemv shapes above the threading threshold. depth is kept short so the f32
297        // accumulation stays within the fixed absolute tolerance used below.
298        mnks.push((1, 8192, 128));
299        mnks.push((8192, 1, 128));
300        for (m, n, k) in mnks {
301            #[cfg(feature = "std")]
302            dbg!(m, n, k);
303            for parallelism in [
304                Parallelism::None,
305                #[cfg(feature = "rayon")]
306                Parallelism::Rayon(0),
307                #[cfg(feature = "rayon")]
308                Parallelism::Rayon(128),
309            ] {
310                for alpha in [0.0, 1.0, 2.3] {
311                    for beta in [0.0, 1.0, 2.3] {
312                        #[cfg(feature = "std")]
313                        dbg!(alpha, beta, parallelism);
314                        for colmajor in [true, false] {
315                            let a_vec: Vec<f32> = (0..(m * k)).map(|_| rand::random()).collect();
316                            let b_vec: Vec<f32> = (0..(k * n)).map(|_| rand::random()).collect();
317                            let mut c_vec: Vec<f32> =
318                                (0..(m * n)).map(|_| rand::random()).collect();
319                            let mut d_vec = c_vec.clone();
320
321                            unsafe {
322                                gemm::gemm(
323                                    m,
324                                    n,
325                                    k,
326                                    c_vec.as_mut_ptr(),
327                                    if colmajor { m } else { 1 } as isize,
328                                    if colmajor { 1 } else { n } as isize,
329                                    true,
330                                    a_vec.as_ptr(),
331                                    m as isize,
332                                    1,
333                                    b_vec.as_ptr(),
334                                    k as isize,
335                                    1,
336                                    alpha,
337                                    beta,
338                                    false,
339                                    false,
340                                    false,
341                                    parallelism,
342                                );
343
344                                gemm::gemm_fallback(
345                                    m,
346                                    n,
347                                    k,
348                                    d_vec.as_mut_ptr(),
349                                    if colmajor { m } else { 1 } as isize,
350                                    if colmajor { 1 } else { n } as isize,
351                                    true,
352                                    a_vec.as_ptr(),
353                                    m as isize,
354                                    1,
355                                    b_vec.as_ptr(),
356                                    k as isize,
357                                    1,
358                                    alpha,
359                                    beta,
360                                );
361                            }
362                            for (c, d) in c_vec.iter().zip(d_vec.iter()) {
363                                assert_approx_eq::assert_approx_eq!(c, d, 1e-3);
364                            }
365                        }
366                    }
367                }
368            }
369        }
370    }
371
372    #[test]
373    fn test_gemm_f64() {
374        set_wasm_simd128(true);
375
376        let mut mnks = vec![];
377        mnks.push((63, 2, 10));
378        mnks.push((1, 2, 10));
379        mnks.push((1, 63, 10));
380
381        // large m to trigger parallelized rhs packing with big number of threads and small n
382        mnks.push((2048, 255, 255));
383
384        mnks.push((256, 256, 256));
385        mnks.push((4096, 4096, 4));
386        mnks.push((64, 64, 4));
387        mnks.push((0, 64, 4));
388        mnks.push((64, 0, 4));
389        mnks.push((0, 0, 4));
390        mnks.push((64, 64, 0));
391        mnks.push((16, 1, 1));
392        mnks.push((16, 2, 1));
393        mnks.push((16, 3, 1));
394        mnks.push((16, 4, 1));
395        mnks.push((16, 1, 2));
396        mnks.push((16, 2, 2));
397        mnks.push((16, 3, 2));
398        mnks.push((16, 4, 2));
399        mnks.push((16, 16, 1));
400        mnks.push((8, 16, 1));
401        mnks.push((16, 8, 1));
402        mnks.push((1, 1, 2));
403        mnks.push((4, 4, 4));
404        mnks.push((1024, 1024, 1));
405        mnks.push((1024, 1024, 4));
406        mnks.push((63, 1, 10));
407        mnks.push((63, 3, 10));
408        mnks.push((63, 4, 10));
409        mnks.push((2, 63, 10));
410        mnks.push((3, 63, 10));
411        mnks.push((4, 63, 10));
412
413        // gemv shapes above the threading threshold: the first two split the output
414        // dimension, the last two are short-output/long-depth and take the k-split.
415        mnks.push((1, 1024, 1024));
416        mnks.push((1024, 1, 1024));
417        mnks.push((1, 128, 8192));
418        mnks.push((128, 1, 8192));
419        for (m, n, k) in mnks {
420            #[cfg(feature = "std")]
421            dbg!(m, n, k);
422            for parallelism in [
423                Parallelism::None,
424                #[cfg(feature = "rayon")]
425                Parallelism::Rayon(0),
426                #[cfg(feature = "rayon")]
427                Parallelism::Rayon(128),
428            ] {
429                for alpha in [0.0, 1.0, 2.3] {
430                    for beta in [0.0, 1.0, 2.3] {
431                        #[cfg(feature = "std")]
432                        dbg!(alpha, beta, parallelism);
433                        for colmajor in [true, false] {
434                            let a_vec: Vec<f64> = (0..(m * k)).map(|_| rand::random()).collect();
435                            let b_vec: Vec<f64> = (0..(k * n)).map(|_| rand::random()).collect();
436                            let mut c_vec: Vec<f64> =
437                                (0..(m * n)).map(|_| rand::random()).collect();
438                            let mut d_vec = c_vec.clone();
439
440                            unsafe {
441                                gemm::gemm(
442                                    m,
443                                    n,
444                                    k,
445                                    c_vec.as_mut_ptr(),
446                                    if colmajor { m } else { 1 } as isize,
447                                    if colmajor { 1 } else { n } as isize,
448                                    true,
449                                    a_vec.as_ptr(),
450                                    m as isize,
451                                    1,
452                                    b_vec.as_ptr(),
453                                    k as isize,
454                                    1,
455                                    alpha,
456                                    beta,
457                                    false,
458                                    false,
459                                    false,
460                                    parallelism,
461                                );
462
463                                gemm::gemm_fallback(
464                                    m,
465                                    n,
466                                    k,
467                                    d_vec.as_mut_ptr(),
468                                    if colmajor { m } else { 1 } as isize,
469                                    if colmajor { 1 } else { n } as isize,
470                                    true,
471                                    a_vec.as_ptr(),
472                                    m as isize,
473                                    1,
474                                    b_vec.as_ptr(),
475                                    k as isize,
476                                    1,
477                                    alpha,
478                                    beta,
479                                );
480                            }
481                            for (c, d) in c_vec.iter().zip(d_vec.iter()) {
482                                assert_approx_eq::assert_approx_eq!(c, d);
483                            }
484                        }
485                    }
486                }
487            }
488        }
489    }
490
491    #[test]
492    fn test_gemm_cplx32() {
493        let mut mnks = vec![];
494        mnks.push((4, 4, 4));
495        mnks.push((0, 64, 4));
496        mnks.push((64, 0, 4));
497        mnks.push((0, 0, 4));
498        mnks.push((64, 64, 4));
499        mnks.push((64, 64, 0));
500        mnks.push((6, 3, 1));
501        mnks.push((1, 1, 2));
502        mnks.push((128, 128, 128));
503        mnks.push((16, 1, 1));
504        mnks.push((16, 2, 1));
505        mnks.push((16, 3, 1));
506        mnks.push((16, 4, 1));
507        mnks.push((16, 1, 2));
508        mnks.push((16, 2, 2));
509        mnks.push((16, 3, 2));
510        mnks.push((16, 4, 2));
511        mnks.push((16, 16, 1));
512        mnks.push((8, 16, 1));
513        mnks.push((16, 8, 1));
514        mnks.push((1024, 1024, 4));
515        mnks.push((1024, 1024, 1));
516        mnks.push((63, 1, 10));
517        mnks.push((63, 2, 10));
518        mnks.push((63, 3, 10));
519        mnks.push((63, 4, 10));
520        mnks.push((1, 63, 10));
521        mnks.push((2, 63, 10));
522        mnks.push((3, 63, 10));
523        mnks.push((4, 63, 10));
524
525        // gemv shapes above the threading threshold. depth is kept short so the f32
526        // accumulation stays within the fixed absolute tolerance used below.
527        mnks.push((1, 8192, 128));
528        mnks.push((8192, 1, 128));
529        for (m, n, k) in mnks {
530            #[cfg(feature = "std")]
531            dbg!(m, n, k);
532
533            let zero = c32::new(0.0, 0.0);
534            let one = c32::new(1.0, 0.0);
535            let arbitrary = c32::new(2.3, 4.1);
536            for alpha in [zero, one, arbitrary] {
537                for beta in [zero, one, arbitrary] {
538                    #[cfg(feature = "std")]
539                    dbg!(alpha, beta);
540                    for conj_dst in [false, true] {
541                        for conj_lhs in [false, true] {
542                            for conj_rhs in [false, true] {
543                                #[cfg(feature = "std")]
544                                dbg!(conj_dst);
545                                #[cfg(feature = "std")]
546                                dbg!(conj_lhs);
547                                #[cfg(feature = "std")]
548                                dbg!(conj_rhs);
549                                for colmajor in [true, false] {
550                                    let a_vec: Vec<f32> =
551                                        (0..(2 * m * k)).map(|_| rand::random()).collect();
552                                    let b_vec: Vec<f32> =
553                                        (0..(2 * k * n)).map(|_| rand::random()).collect();
554                                    let mut c_vec: Vec<f32> =
555                                        (0..(2 * m * n)).map(|_| rand::random()).collect();
556                                    let mut d_vec = c_vec.clone();
557
558                                    unsafe {
559                                        gemm::gemm(
560                                            m,
561                                            n,
562                                            k,
563                                            c_vec.as_mut_ptr() as *mut c32,
564                                            if colmajor { m } else { 1 } as isize,
565                                            if colmajor { 1 } else { n } as isize,
566                                            true,
567                                            a_vec.as_ptr() as *const c32,
568                                            m as isize,
569                                            1,
570                                            b_vec.as_ptr() as *const c32,
571                                            k as isize,
572                                            1,
573                                            alpha,
574                                            beta,
575                                            conj_dst,
576                                            conj_lhs,
577                                            conj_rhs,
578                                            #[cfg(feature = "rayon")]
579                                            Parallelism::Rayon(0),
580                                            #[cfg(not(feature = "rayon"))]
581                                            Parallelism::None,
582                                        );
583
584                                        gemm::gemm_cplx_fallback(
585                                            m,
586                                            n,
587                                            k,
588                                            d_vec.as_mut_ptr() as *mut c32,
589                                            if colmajor { m } else { 1 } as isize,
590                                            if colmajor { 1 } else { n } as isize,
591                                            true,
592                                            a_vec.as_ptr() as *const c32,
593                                            m as isize,
594                                            1,
595                                            b_vec.as_ptr() as *const c32,
596                                            k as isize,
597                                            1,
598                                            alpha,
599                                            beta,
600                                            conj_dst,
601                                            conj_lhs,
602                                            conj_rhs,
603                                        );
604                                    }
605                                    for (c, d) in c_vec.iter().zip(d_vec.iter()) {
606                                        assert_approx_eq::assert_approx_eq!(c, d, 1e-3);
607                                    }
608                                }
609                            }
610                        }
611                    }
612                }
613            }
614        }
615    }
616
617    #[test]
618    fn test_gemm_cplx64() {
619        let mut mnks = vec![];
620        mnks.push((4, 4, 4));
621        mnks.push((0, 64, 4));
622        mnks.push((64, 0, 4));
623        mnks.push((0, 0, 4));
624        mnks.push((64, 64, 4));
625        mnks.push((64, 64, 0));
626        mnks.push((6, 3, 1));
627        mnks.push((1, 1, 2));
628        mnks.push((128, 128, 128));
629        mnks.push((16, 1, 1));
630        mnks.push((16, 2, 1));
631        mnks.push((16, 3, 1));
632        mnks.push((16, 4, 1));
633        mnks.push((16, 1, 2));
634        mnks.push((16, 2, 2));
635        mnks.push((16, 3, 2));
636        mnks.push((16, 4, 2));
637        mnks.push((16, 16, 1));
638        mnks.push((8, 16, 1));
639        mnks.push((16, 8, 1));
640        mnks.push((1024, 1024, 4));
641        mnks.push((1024, 1024, 1));
642        mnks.push((63, 1, 10));
643        mnks.push((63, 2, 10));
644        mnks.push((63, 3, 10));
645        mnks.push((63, 4, 10));
646        mnks.push((1, 63, 10));
647        mnks.push((2, 63, 10));
648        mnks.push((3, 63, 10));
649        mnks.push((4, 63, 10));
650
651        // gemv shapes above the threading threshold: the first two split the output
652        // dimension, the last two are short-output/long-depth and take the k-split.
653        mnks.push((1, 1024, 1024));
654        mnks.push((1024, 1, 1024));
655        mnks.push((1, 128, 8192));
656        mnks.push((128, 1, 8192));
657        for (m, n, k) in mnks {
658            #[cfg(feature = "std")]
659            dbg!(m, n, k);
660
661            let zero = c64::new(0.0, 0.0);
662            let one = c64::new(1.0, 0.0);
663            let arbitrary = c64::new(2.3, 4.1);
664            for alpha in [zero, one, arbitrary] {
665                for beta in [zero, one, arbitrary] {
666                    #[cfg(feature = "std")]
667                    dbg!(alpha, beta);
668                    for conj_dst in [false, true] {
669                        for conj_lhs in [false, true] {
670                            for conj_rhs in [false, true] {
671                                #[cfg(feature = "std")]
672                                dbg!(conj_dst);
673                                #[cfg(feature = "std")]
674                                dbg!(conj_lhs);
675                                #[cfg(feature = "std")]
676                                dbg!(conj_rhs);
677                                for colmajor in [true, false] {
678                                    let a_vec: Vec<f64> =
679                                        (0..(2 * m * k)).map(|_| rand::random()).collect();
680                                    let b_vec: Vec<f64> =
681                                        (0..(2 * k * n)).map(|_| rand::random()).collect();
682                                    let mut c_vec: Vec<f64> =
683                                        (0..(2 * m * n)).map(|_| rand::random()).collect();
684                                    let mut d_vec = c_vec.clone();
685
686                                    unsafe {
687                                        gemm::gemm(
688                                            m,
689                                            n,
690                                            k,
691                                            c_vec.as_mut_ptr() as *mut c64,
692                                            if colmajor { m } else { 1 } as isize,
693                                            if colmajor { 1 } else { n } as isize,
694                                            true,
695                                            a_vec.as_ptr() as *const c64,
696                                            m as isize,
697                                            1,
698                                            b_vec.as_ptr() as *const c64,
699                                            k as isize,
700                                            1,
701                                            alpha,
702                                            beta,
703                                            conj_dst,
704                                            conj_lhs,
705                                            conj_rhs,
706                                            #[cfg(feature = "rayon")]
707                                            Parallelism::Rayon(0),
708                                            #[cfg(not(feature = "rayon"))]
709                                            Parallelism::None,
710                                        );
711
712                                        gemm::gemm_cplx_fallback(
713                                            m,
714                                            n,
715                                            k,
716                                            d_vec.as_mut_ptr() as *mut c64,
717                                            if colmajor { m } else { 1 } as isize,
718                                            if colmajor { 1 } else { n } as isize,
719                                            true,
720                                            a_vec.as_ptr() as *const c64,
721                                            m as isize,
722                                            1,
723                                            b_vec.as_ptr() as *const c64,
724                                            k as isize,
725                                            1,
726                                            alpha,
727                                            beta,
728                                            conj_dst,
729                                            conj_lhs,
730                                            conj_rhs,
731                                        );
732                                    }
733                                    for (c, d) in c_vec.iter().zip(d_vec.iter()) {
734                                        assert_approx_eq::assert_approx_eq!(c, d);
735                                    }
736                                }
737                            }
738                        }
739                    }
740                }
741            }
742        }
743    }
744}
xn_gemm/lib.rs

xn_gemm/
lib.rs
