(a: &[f32], b: &[f32])
| 11 | |
| 12 | #[target_feature(enable = "avx512f")] |
| 13 | unsafe fn l2_impl(a: &[f32], b: &[f32]) -> f32 { |
| 14 | assert_eq!(a.len(), b.len(), "avx512 l2_impl: length mismatch"); |
| 15 | unsafe { |
| 16 | use std::arch::x86_64::*; |
| 17 | let n = a.len(); |
| 18 | let mut sum = _mm512_setzero_ps(); |
| 19 | let chunks = n / 16; |
| 20 | for i in 0..chunks { |
| 21 | let off = i * 16; |
| 22 | let va = _mm512_loadu_ps(a.as_ptr().add(off)); |
| 23 | let vb = _mm512_loadu_ps(b.as_ptr().add(off)); |
| 24 | let diff = _mm512_sub_ps(va, vb); |
| 25 | sum = _mm512_fmadd_ps(diff, diff, sum); |
| 26 | } |
| 27 | let mut result = _mm512_reduce_add_ps(sum); |
| 28 | for i in (chunks * 16)..n { |
| 29 | let d = a[i] - b[i]; |
| 30 | result += d * d; |
| 31 | } |
| 32 | result |
| 33 | } |
| 34 | } |
| 35 | |
| 36 | pub fn cosine_distance(a: &[f32], b: &[f32]) -> f32 { |
| 37 | assert_eq!(a.len(), b.len(), "avx512 cosine: length mismatch"); |
no test coverage detected