create_str_impl_avx512vl() remove support for AVX10.2/256

ijl · ijl · commit 96b270bd17c6 · 2025-07-10T14:56:23.000Z
diff --git a/src/str/avx512.rs b/src/str/avx512.rs
@@ -3,14 +3,14 @@
 use crate::str::pyunicode_new::*;
 
 use core::arch::x86_64::{
-    _mm256_and_si256, _mm256_cmpgt_epu8_mask, _mm256_cmpneq_epi8_mask, _mm256_loadu_epi8,
-    _mm256_mask_cmpneq_epi8_mask, _mm256_maskz_loadu_epi8, _mm256_max_epu8, _mm256_set1_epi8,
+    _mm512_and_si512, _mm512_cmpgt_epu8_mask, _mm512_cmpneq_epi8_mask, _mm512_loadu_epi8,
+    _mm512_mask_cmpneq_epi8_mask, _mm512_maskz_loadu_epi8, _mm512_max_epu8, _mm512_set1_epi8,
 };
 
 #[inline(never)]
 #[target_feature(enable = "avx512f,avx512bw,avx512vl,bmi2")]
 pub(crate) unsafe fn create_str_impl_avx512vl(buf: &str) -> *mut pyo3_ffi::PyObject {
-    const STRIDE: usize = 32;
+    const STRIDE: usize = 64;
 
     let buf_ptr = buf.as_bytes().as_ptr().cast::<i8>();
     let buf_len = buf.len();
@@ -20,40 +20,40 @@ pub(crate) unsafe fn create_str_impl_avx512vl(buf: &str) -> *mut pyo3_ffi::PyObj
     let num_loops = buf_len / STRIDE;
     let remainder = buf_len % STRIDE;
 
-    let remainder_mask: u32 = !(u32::MAX << remainder);
-    let mut str_vec = _mm256_maskz_loadu_epi8(remainder_mask, buf_ptr);
+    let remainder_mask: u64 = !(u64::MAX << remainder);
+    let mut str_vec = _mm512_maskz_loadu_epi8(remainder_mask, buf_ptr);
     let sptr = buf_ptr.add(remainder);
 
     for i in 0..num_loops {
-        str_vec = _mm256_max_epu8(
+        str_vec = _mm512_max_epu8(
             str_vec,
-            _mm256_loadu_epi8(sptr.add(STRIDE * i).cast::<i8>()),
+            _mm512_loadu_epi8(sptr.add(STRIDE * i).cast::<i8>()),
         );
     }
 
     #[allow(overflowing_literals)]
-    let vec_128 = _mm256_set1_epi8(0b10000000i8);
-    if _mm256_cmpgt_epu8_mask(str_vec, vec_128) == 0 {
+    let vec_128 = _mm512_set1_epi8(0b10000000i8);
+    if _mm512_cmpgt_epu8_mask(str_vec, vec_128) == 0 {
         pyunicode_ascii(buf.as_bytes().as_ptr(), buf_len)
     } else {
         #[allow(overflowing_literals)]
-        let is_four = _mm256_cmpgt_epu8_mask(str_vec, _mm256_set1_epi8(239i8)) != 0;
+        let is_four = _mm512_cmpgt_epu8_mask(str_vec, _mm512_set1_epi8(239i8)) != 0;
         #[allow(overflowing_literals)]
-        let is_not_latin = _mm256_cmpgt_epu8_mask(str_vec, _mm256_set1_epi8(195i8)) != 0;
+        let is_not_latin = _mm512_cmpgt_epu8_mask(str_vec, _mm512_set1_epi8(195i8)) != 0;
         #[allow(overflowing_literals)]
-        let multibyte = _mm256_set1_epi8(0b11000000i8);
+        let multibyte = _mm512_set1_epi8(0b11000000i8);
 
-        let mut num_chars = _mm256_mask_cmpneq_epi8_mask(
+        let mut num_chars = _mm512_mask_cmpneq_epi8_mask(
             remainder_mask,
-            _mm256_and_si256(_mm256_maskz_loadu_epi8(remainder_mask, buf_ptr), multibyte),
+            _mm512_and_si512(_mm512_maskz_loadu_epi8(remainder_mask, buf_ptr), multibyte),
             vec_128,
         )
         .count_ones() as usize;
 
         for i in 0..num_loops {
-            num_chars += _mm256_cmpneq_epi8_mask(
-                _mm256_and_si256(
-                    _mm256_loadu_epi8(sptr.add(STRIDE * i).cast::<i8>()),
+            num_chars += _mm512_cmpneq_epi8_mask(
+                _mm512_and_si512(
+                    _mm512_loadu_epi8(sptr.add(STRIDE * i).cast::<i8>()),
                     multibyte,
                 ),
                 vec_128,
