optimize num_chars for aarch64 more

Author: llogiq

src/simd/aarch64.rs

@@ -1,6 +1,6 @@
 use core::arch::aarch64::{
     uint8x16_t, uint8x16x4_t, vaddlvq_u8, vandq_u8, vceqq_u8, vdupq_n_u8, vld1q_u8, vld1q_u8_x4,
-    vmvnq_u8, vsubq_u8,
+    vsubq_u8,
 };
 
 const MASK: [u8; 32] = [
@@ -35,6 +35,10 @@ unsafe fn sum(u8s: uint8x16_t) -> usize {
     vaddlvq_u8(u8s) as usize
 }
 
+unsafe fn sum4(u1: uint8x16_t, u2: uint8x16_t, u3: uint8x16_t, u4: uint8x16_t) -> usize {
+    ((vaddlvq_u8(u1) + vaddlvq_u8(u2)) + (vaddlvq_u8(u3) + vaddlvq_u8(u4))) as usize
+}
+
 #[target_feature(enable = "neon")]
 pub unsafe fn chunk_count(haystack: &[u8], needle: u8) -> usize {
     assert!(haystack.len() >= 16);
@@ -56,7 +60,7 @@ pub unsafe fn chunk_count(haystack: &[u8], needle: u8) -> usize {
             count4 = vsubq_u8(count4, vceqq_u8(h4, needles));
             offset += 64;
         }
-        count += sum(count1) + sum(count2) + sum(count3) + sum(count4);
+        count += sum4(count1, count2, count3, count4);
     }
 
     // 64
@@ -70,7 +74,7 @@ pub unsafe fn chunk_count(haystack: &[u8], needle: u8) -> usize {
         count4 = vsubq_u8(count4, vceqq_u8(h4, needles));
         offset += 64;
     }
-    count += sum(count1) + sum(count2) + sum(count3) + sum(count4);
+    count += sum4(count1, count2, count3, count4);
 
     let mut counts = vdupq_n_u8(0);
     // 16
@@ -93,11 +97,11 @@ pub unsafe fn chunk_count(haystack: &[u8], needle: u8) -> usize {
 }
 
 #[target_feature(enable = "neon")]
-unsafe fn is_leading_utf8_byte(u8s: uint8x16_t) -> uint8x16_t {
-    vmvnq_u8(vceqq_u8(
+unsafe fn is_following_utf8_byte(u8s: uint8x16_t) -> uint8x16_t {
+    vceqq_u8(
         vandq_u8(u8s, vdupq_n_u8(0b1100_0000)),
         vdupq_n_u8(0b1000_0000),
-    ))
+    )
 }
 
 #[target_feature(enable = "neon")]
@@ -108,50 +112,51 @@ pub unsafe fn chunk_num_chars(utf8_chars: &[u8]) -> usize {
     let mut count = 0;
 
     // 4080
-    while utf8_chars.len() >= offset + 16 * 255 {
-        let mut counts = vdupq_n_u8(0);
+    while utf8_chars.len() >= offset + 64 * 255 {
+        let (mut count1, mut count2, mut count3, mut count4) =
+            (vdupq_n_u8(0), vdupq_n_u8(0), vdupq_n_u8(0), vdupq_n_u8(0));
 
         for _ in 0..255 {
-            counts = vsubq_u8(
-                counts,
-                is_leading_utf8_byte(u8x16_from_offset(utf8_chars, offset)),
-            );
-            offset += 16;
+            let uint8x16x4_t(h1, h2, h3, h4) = u8x16_x4_from_offset(utf8_chars, offset);
+            count1 = vsubq_u8(count1,is_following_utf8_byte(h1));
+            count2 = vsubq_u8(count2,is_following_utf8_byte(h2));
+            count3 = vsubq_u8(count3,is_following_utf8_byte(h3));
+            count4 = vsubq_u8(count4,is_following_utf8_byte(h4));
+            offset += 64;
         }
-        count += sum(counts);
+        count += sum4(count1, count2, count3, count4);
     }
 
-    // 2048
-    if utf8_chars.len() >= offset + 16 * 128 {
-        let mut counts = vdupq_n_u8(0);
-        for _ in 0..128 {
-            counts = vsubq_u8(
-                counts,
-                is_leading_utf8_byte(u8x16_from_offset(utf8_chars, offset)),
-            );
-            offset += 16;
+    // 4080
+    let (mut count1, mut count2, mut count3, mut count4) =
+        (vdupq_n_u8(0), vdupq_n_u8(0), vdupq_n_u8(0), vdupq_n_u8(0));
+        for _ in 0..(utf8_chars.len() - offset) / 64 {
+            let uint8x16x4_t(h1, h2, h3, h4) = u8x16_x4_from_offset(utf8_chars, offset);
+            count1 = vsubq_u8(count1, is_following_utf8_byte(h1));
+            count2 = vsubq_u8(count2, is_following_utf8_byte(h2));
+            count3 = vsubq_u8(count3, is_following_utf8_byte(h3));
+            count4 = vsubq_u8(count4, is_following_utf8_byte(h4));
+            offset += 64;
         }
-        count += sum(counts);
-    }
-
+        count += sum4(count1, count2, count3, count4);
     // 16
     let mut counts = vdupq_n_u8(0);
     for i in 0..(utf8_chars.len() - offset) / 16 {
         counts = vsubq_u8(
             counts,
-            is_leading_utf8_byte(u8x16_from_offset(utf8_chars, offset + i * 16)),
+            is_following_utf8_byte(u8x16_from_offset(utf8_chars, offset + i * 16)),
         );
     }
     if utf8_chars.len() % 16 != 0 {
         counts = vsubq_u8(
             counts,
             vandq_u8(
-                is_leading_utf8_byte(u8x16_from_offset(utf8_chars, utf8_chars.len() - 16)),
+                is_following_utf8_byte(u8x16_from_offset(utf8_chars, utf8_chars.len() - 16)),
                 u8x16_from_offset(&MASK, utf8_chars.len() % 16),
             ),
         );
     }
     count += sum(counts);
 
-    count
+    utf8_chars.len() - count
 }