avx512: use masked load for in complete last block

Author: hkratz

src/implementation/algorithm.rs

@@ -240,13 +240,7 @@ macro_rules! algorithm_simd {
             }
 
             if idx < len {
-                let mut tmpbuf = TempSimdChunk::new();
-                crate::implementation::helpers::memcpy_unaligned_nonoverlapping_inline_opt_lt_64(
-                    input.as_ptr().add(idx),
-                    tmpbuf.0.as_mut_ptr(),
-                    len - idx,
-                );
-                let simd_input = SimdInput::new(tmpbuf.0.as_ptr());
+                let simd_input = SimdInput::new_partial(input.as_ptr().add(idx), len-idx);
                 algorithm.check_utf8(simd_input);
             }
             algorithm.check_incomplete_pending();
@@ -332,14 +326,7 @@ macro_rules! algorithm_simd {
                 break;
             }
             if idx < len {
-                let mut tmpbuf = TempSimdChunk::new();
-                crate::implementation::helpers::memcpy_unaligned_nonoverlapping_inline_opt_lt_64(
-                    input.as_ptr().add(idx),
-                    tmpbuf.0.as_mut_ptr(),
-                    len - idx,
-                );
-                let simd_input = SimdInput::new(tmpbuf.0.as_ptr());
-
+                let simd_input = SimdInput::new_partial(input.as_ptr().add(idx), len-idx);
                 algorithm.check_utf8(simd_input);
             }
             algorithm.check_incomplete_pending();
@@ -537,6 +524,18 @@ macro_rules! simd_input_128_bit {
                 }
             }
 
+            $(#[$feat])*
+            #[inline]
+            unsafe fn new_partial(ptr: *const u8, len: usize) -> Self {
+                let mut tmpbuf = TempSimdChunk::new();
+                crate::implementation::helpers::memcpy_unaligned_nonoverlapping_inline_opt_lt_64(
+                    ptr,
+                    tmpbuf.0.as_mut_ptr(),
+                    len,
+                );
+                Self::new(tmpbuf.0.as_ptr())
+            }
+
             $(#[$feat])*
             #[inline]
             unsafe fn is_ascii(&self) -> bool {
@@ -568,6 +567,18 @@ macro_rules! simd_input_256_bit {
                 }
             }
 
+            $(#[$feat])*
+            #[inline]
+            unsafe fn new_partial(ptr: *const u8, len: usize) -> Self {
+                let mut tmpbuf = TempSimdChunk::new();
+                crate::implementation::helpers::memcpy_unaligned_nonoverlapping_inline_opt_lt_64(
+                    ptr,
+                    tmpbuf.0.as_mut_ptr(),
+                    len,
+                );
+                Self::new(tmpbuf.0.as_ptr())
+            }
+
             $(#[$feat])*
             #[inline]
             unsafe fn is_ascii(&self) -> bool {
@@ -595,6 +606,17 @@ macro_rules! simd_input_512_bit {
                 }
             }
 
+
+            $(#[$feat])*
+            #[inline]
+            unsafe fn new_partial(ptr: *const u8, len: usize) -> Self {
+                Self {
+                    vals: [
+                        SimdU8Value::load_from_partial(ptr, len),
+                    ],
+                }
+            }
+
             $(#[$feat])*
             #[inline]
             unsafe fn is_ascii(&self) -> bool {

src/implementation/x86/avx512.rs

@@ -2,9 +2,10 @@
 
 use core::arch::x86_64::{
     __m512i, _mm512_alignr_epi8, _mm512_and_si512, _mm512_cmpgt_epi8_mask, _mm512_loadu_si512,
-    _mm512_maskz_abs_epi8, _mm512_or_si512, _mm512_permutex2var_epi64, _mm512_set1_epi8,
-    _mm512_set_epi64, _mm512_setzero_si512, _mm512_shuffle_epi8, _mm512_srli_epi16,
-    _mm512_subs_epu8, _mm512_test_epi8_mask, _mm512_xor_si512, _mm_prefetch, _MM_HINT_T0,
+    _mm512_maskz_abs_epi8, _mm512_maskz_loadu_epi8, _mm512_or_si512, _mm512_permutex2var_epi64,
+    _mm512_set1_epi8, _mm512_set_epi64, _mm512_setzero_si512, _mm512_shuffle_epi8,
+    _mm512_srli_epi16, _mm512_subs_epu8, _mm512_test_epi8_mask, _mm512_xor_si512, _mm_prefetch,
+    _MM_HINT_T0,
 };
 use core::arch::x86_64::{_mm512_movepi8_mask, _mm512_set_epi8};
 
@@ -106,6 +107,13 @@ impl SimdU8Value {
         Self::from(_mm512_loadu_si512(ptr.cast::<__m512i>()))
     }
 
+    #[target_feature(enable = "avx512f,avx512bw,avx512vbmi")]
+    #[inline]
+    unsafe fn load_from_partial(ptr: *const u8, len: usize) -> Self {
+        let res = _mm512_maskz_loadu_epi8(u64::MAX >> (64 - len), ptr.cast::<i8>());
+        Self::from(res)
+    }
+
     #[flexpect::e(clippy::too_many_arguments)]
     #[target_feature(enable = "avx512f,avx512bw,avx512vbmi")]
     #[inline]
@@ -266,6 +274,5 @@ unsafe fn simd_prefetch(ptr: *const u8) {
 }
 
 const PREFETCH: bool = true;
-use crate::implementation::helpers::TempSimdChunkA64 as TempSimdChunk;
 simd_input_512_bit!(#[target_feature(enable = "avx512f,avx512bw,avx512vbmi")]);
 algorithm_simd!(#[target_feature(enable = "avx512f,avx512bw,avx512vbmi")]);