(a_hot: &[u8], b_hot: &[u8], dim: usize)
| 21 | #[cfg(target_arch = "x86_64")] |
| 22 | #[target_feature(enable = "avx512f,avx512bw")] |
| 23 | unsafe fn ternary_dot_avx512(a_hot: &[u8], b_hot: &[u8], dim: usize) -> i32 { |
| 24 | use std::arch::x86_64::*; |
| 25 | |
| 26 | let a_trits = unpack_hot(a_hot, dim); |
| 27 | let b_trits = unpack_hot(b_hot, dim); |
| 28 | |
| 29 | let len = a_trits.len(); |
| 30 | let chunks = len / 64; |
| 31 | let mut sum = _mm512_setzero_si512(); |
| 32 | |
| 33 | for i in 0..chunks { |
| 34 | let a_ptr = a_trits.as_ptr().add(i * 64) as *const __m512i; |
| 35 | let b_ptr = b_trits.as_ptr().add(i * 64) as *const __m512i; |
| 36 | let a_vec = _mm512_loadu_si512(a_ptr); |
| 37 | let b_vec = _mm512_loadu_si512(b_ptr); |
| 38 | let a_lo = _mm512_cvtepi8_epi16(_mm512_castsi512_si256(a_vec)); |
| 39 | let a_hi = _mm512_cvtepi8_epi16(_mm512_extracti64x4_epi64(a_vec, 1)); |
| 40 | let b_lo = _mm512_cvtepi8_epi16(_mm512_castsi512_si256(b_vec)); |
| 41 | let b_hi = _mm512_cvtepi8_epi16(_mm512_extracti64x4_epi64(b_vec, 1)); |
| 42 | let prod_lo = _mm512_mullo_epi16(a_lo, b_lo); |
| 43 | let prod_hi = _mm512_mullo_epi16(a_hi, b_hi); |
| 44 | let ones = _mm512_set1_epi16(1); |
| 45 | let sum_lo = _mm512_madd_epi16(prod_lo, ones); |
| 46 | let sum_hi = _mm512_madd_epi16(prod_hi, ones); |
| 47 | sum = _mm512_add_epi32(sum, _mm512_add_epi32(sum_lo, sum_hi)); |
| 48 | } |
| 49 | |
| 50 | let mut acc = _mm512_reduce_add_epi32(sum); |
| 51 | for i in (chunks * 64)..len { |
| 52 | acc += a_trits[i] as i32 * b_trits[i] as i32; |
| 53 | } |
| 54 | acc |
| 55 | } |
| 56 | |
| 57 | #[cfg(target_arch = "aarch64")] |
| 58 | fn ternary_dot_neon(a_hot: &[u8], b_hot: &[u8], dim: usize) -> i32 { |
no test coverage detected