Deduplicate scroll and SIMD paths while preserving performance

1 files changed, 61 insertions, 90 deletions

diff --git a/src/kernels.h b/src/kernels.h
index 57066c0..5cf0d93 100644
--- a/src/kernels.h
+++ b/src/kernels.h
@@ -197,23 +197,14 @@ static inline TERM_ALWAYS_INLINE void bytes_copy_small(uint8_t *restrict dst,
         if (k < n)
             memcpy(dst + n - 16, src + n - 16, 16);
     } else if (n >= 8) {
-        uint64_t a, b;
-        memcpy(&a, src, 8);
-        memcpy(&b, src + n - 8, 8);
-        memcpy(dst, &a, 8);
-        memcpy(dst + n - 8, &b, 8);
+        memcpy(dst, src, 8);
+        memcpy(dst + n - 8, src + n - 8, 8);
     } else if (n >= 4) {
-        uint32_t a, b;
-        memcpy(&a, src, 4);
-        memcpy(&b, src + n - 4, 4);
-        memcpy(dst, &a, 4);
-        memcpy(dst + n - 4, &b, 4);
+        memcpy(dst, src, 4);
+        memcpy(dst + n - 4, src + n - 4, 4);
     } else if (n >= 2) {
-        uint16_t a, b;
-        memcpy(&a, src, 2);
-        memcpy(&b, src + n - 2, 2);
-        memcpy(dst, &a, 2);
-        memcpy(dst + n - 2, &b, 2);
+        memcpy(dst, src, 2);
+        memcpy(dst + n - 2, src + n - 2, 2);
     } else if (n == 1) {
         dst[0] = (uint8_t)src[0];
     }
@@ -365,6 +356,26 @@ static inline int scan_ascii_lf_tail(const char *buf, int i, int len, uint16_t *
 }
 
 #if defined(TERM_KERNEL_AVX512)
+TERM_TARGET_AVX512
+static inline int store_lf_offsets_avx512(uint16_t *ends, uint64_t nl,
+                                         __m512i lanes, __m128i block)
+{
+    unsigned n = (unsigned)term_popcount64(nl);
+    __m512i packed = _mm512_maskz_compress_epi8(nl, lanes);
+    __m128i offsets = _mm_add_epi16(_mm_cvtepu8_epi16(_mm512_castsi512_si128(packed)), block);
+    _mm_storeu_si128((__m128i *)ends, offsets);
+    if (TERM_EXPECT(n > 8, 0)) {
+        __m512i base = _mm512_broadcast_i32x4(block);
+        __m512i wide = _mm512_add_epi16(_mm512_cvtepu8_epi16(_mm512_castsi512_si256(packed)), base);
+        _mm512_storeu_si512(ends, wide);
+        if (n > 32) {
+            wide = _mm512_add_epi16(_mm512_cvtepu8_epi16(_mm512_extracti64x4_epi64(packed, 1)), base);
+            _mm512_storeu_si512(ends + 32, wide);
+        }
+    }
+    return (int)n;
+}
+
 TERM_TARGET_AVX512 TERM_MAYBE_UNUSED
 static int scan_ascii_lf_avx512(const char *buf, int len, uint16_t *ends)
 {
@@ -388,20 +399,7 @@ static int scan_ascii_lf_avx512(const char *buf, int len, uint16_t *ends)
         }
         if (bad) break;
         for (int q = 0; q < 4; q++) {
-            unsigned n = (unsigned)term_popcount64(nl[q]);
-            __m512i packed = _mm512_maskz_compress_epi8(nl[q], lanes);
-            __m128i offsets = _mm_add_epi16(_mm_cvtepu8_epi16(_mm512_castsi512_si128(packed)), block);
-            _mm_storeu_si128((__m128i *)(ends + count + 1), offsets);
-            if (TERM_EXPECT(n > 8, 0)) {
-                __m512i base = _mm512_broadcast_i32x4(block);
-                __m512i wide = _mm512_add_epi16(_mm512_cvtepu8_epi16(_mm512_castsi512_si256(packed)), base);
-                _mm512_storeu_si512(ends + count + 1, wide);
-                if (n > 32) {
-                    wide = _mm512_add_epi16(_mm512_cvtepu8_epi16(_mm512_extracti64x4_epi64(packed, 1)), base);
-                    _mm512_storeu_si512(ends + count + 33, wide);
-                }
-            }
-            count += (int)n;
+            count += store_lf_offsets_avx512(ends + count + 1, nl[q], lanes, block);
             block = _mm_add_epi16(block, _mm_set1_epi16(64));
         }
     }
@@ -411,20 +409,7 @@ static int scan_ascii_lf_avx512(const char *buf, int len, uint16_t *ends)
         uint64_t bad = _mm512_cmpge_epu8_mask(
             _mm512_sub_epi8(v, _mm512_set1_epi8(0x20)), _mm512_set1_epi8(0x5f)) & ~nl;
         if (bad) break;
-        unsigned n = (unsigned)term_popcount64(nl);
-        __m512i packed = _mm512_maskz_compress_epi8(nl, lanes);
-        __m128i offsets = _mm_add_epi16(_mm_cvtepu8_epi16(_mm512_castsi512_si128(packed)), block);
-        _mm_storeu_si128((__m128i *)(ends + count + 1), offsets);
-        if (TERM_EXPECT(n > 8, 0)) {
-            __m512i base = _mm512_broadcast_i32x4(block);
-            __m512i wide = _mm512_add_epi16(_mm512_cvtepu8_epi16(_mm512_castsi512_si256(packed)), base);
-            _mm512_storeu_si512(ends + count + 1, wide);
-            if (n > 32) {
-                wide = _mm512_add_epi16(_mm512_cvtepu8_epi16(_mm512_extracti64x4_epi64(packed, 1)), base);
-                _mm512_storeu_si512(ends + count + 33, wide);
-            }
-        }
-        count += (int)n;
+        count += store_lf_offsets_avx512(ends + count + 1, nl, lanes, block);
         block = _mm_add_epi16(block, _mm_set1_epi16(64));
     }
     return scan_ascii_lf_tail(buf, i, len, ends, count);
@@ -672,6 +657,32 @@ static inline void put_ascii_chunk_avx2(struct cell * restrict dst, const char *
 }
 #endif
 
+#if defined(__aarch64__)
+static inline void store_ascii_quad_neon(struct cell *dst, uint32x4_t cps,
+                                         uint32x4_t fg, uint32x4_t bg_flags)
+{
+    uint64x2_t lo = vreinterpretq_u64_u32(vzip1q_u32(cps, fg));
+    uint64x2_t hi = vreinterpretq_u64_u32(vzip2q_u32(cps, fg));
+    uint64x2_t tail = vreinterpretq_u64_u32(bg_flags);
+    uint32x4x4_t cells = {{
+        vreinterpretq_u32_u64(vzip1q_u64(lo, tail)),
+        vreinterpretq_u32_u64(vzip2q_u64(lo, tail)),
+        vreinterpretq_u32_u64(vzip1q_u64(hi, tail)),
+        vreinterpretq_u32_u64(vzip2q_u64(hi, tail)),
+    }};
+    vst1q_u32_x4((uint32_t *)(void *)dst, cells);
+}
+
+static inline void put_ascii_quad_neon(struct cell *dst, const char *src,
+                                       uint32x4_t fg, uint32x4_t bg_flags)
+{
+    uint32_t bytes;
+    memcpy(&bytes, src, sizeof(bytes));
+    uint16x4_t cps = vget_low_u16(vmovl_u8(vreinterpret_u8_u32(vdup_n_u32(bytes))));
+    store_ascii_quad_neon(dst, vmovl_u16(cps), fg, bg_flags);
+}
+#endif
+
 TERM_ALWAYS_INLINE
 static inline void put_ascii_chunk(struct cell * restrict dst, const char * restrict src, int n,
                                   uint32_t fg, uint32_t bg, uint32_t flags)
@@ -696,55 +707,15 @@ static inline void put_ascii_chunk(struct cell * restrict dst, const char * rest
                                        vcreate_u32(((uint64_t)flags << 32) | bg));
     for (; k + 8 <= n; k += 8) {
         uint16x8_t cp16 = vmovl_u8(vld1_u8((const uint8_t *)src + k));
-        uint32x4_t cp_lo = vmovl_u16(vget_low_u16(cp16));
-        uint32x4_t cp_hi = vmovl_u16(vget_high_u16(cp16));
-
-        uint32x4_t c_fg_01 = vzip1q_u32(cp_lo, fg4);
-        uint32x4_t c_fg_23 = vzip2q_u32(cp_lo, fg4);
-        uint32x4_t c_fg_45 = vzip1q_u32(cp_hi, fg4);
-        uint32x4_t c_fg_67 = vzip2q_u32(cp_hi, fg4);
-
-        uint32x4_t cell0 = (uint32x4_t)vzip1q_u64((uint64x2_t)c_fg_01, (uint64x2_t)bg_flags);
-        uint32x4_t cell1 = (uint32x4_t)vzip2q_u64((uint64x2_t)c_fg_01, (uint64x2_t)bg_flags);
-        uint32x4_t cell2 = (uint32x4_t)vzip1q_u64((uint64x2_t)c_fg_23, (uint64x2_t)bg_flags);
-        uint32x4_t cell3 = (uint32x4_t)vzip2q_u64((uint64x2_t)c_fg_23, (uint64x2_t)bg_flags);
-        uint32x4_t cell4 = (uint32x4_t)vzip1q_u64((uint64x2_t)c_fg_45, (uint64x2_t)bg_flags);
-        uint32x4_t cell5 = (uint32x4_t)vzip2q_u64((uint64x2_t)c_fg_45, (uint64x2_t)bg_flags);
-        uint32x4_t cell6 = (uint32x4_t)vzip1q_u64((uint64x2_t)c_fg_67, (uint64x2_t)bg_flags);
-        uint32x4_t cell7 = (uint32x4_t)vzip2q_u64((uint64x2_t)c_fg_67, (uint64x2_t)bg_flags);
-
-        uint32x4x4_t q0 = {{ cell0, cell1, cell2, cell3 }};
-        uint32x4x4_t q1 = {{ cell4, cell5, cell6, cell7 }};
-        vst1q_u32_x4((uint32_t *)(void *)(dst + k + 0), q0);
-        vst1q_u32_x4((uint32_t *)(void *)(dst + k + 4), q1);
-    }
-    for (; k + 4 <= n; k += 4) {
-        uint32_t bytes;
-        memcpy(&bytes, src + k, sizeof(bytes));
-        uint16x4_t cp16 = vget_low_u16(vmovl_u8(vreinterpret_u8_u32(vdup_n_u32(bytes))));
-        uint32x4_t cp_lo = vmovl_u16(cp16);
-        uint32x4_t c_fg_01 = vzip1q_u32(cp_lo, fg4);
-        uint32x4_t c_fg_23 = vzip2q_u32(cp_lo, fg4);
-        uint32x4_t cell0 = (uint32x4_t)vzip1q_u64((uint64x2_t)c_fg_01, (uint64x2_t)bg_flags);
-        uint32x4_t cell1 = (uint32x4_t)vzip2q_u64((uint64x2_t)c_fg_01, (uint64x2_t)bg_flags);
-        uint32x4_t cell2 = (uint32x4_t)vzip1q_u64((uint64x2_t)c_fg_23, (uint64x2_t)bg_flags);
-        uint32x4_t cell3 = (uint32x4_t)vzip2q_u64((uint64x2_t)c_fg_23, (uint64x2_t)bg_flags);
-        uint32x4x4_t q0 = {{ cell0, cell1, cell2, cell3 }};
-        vst1q_u32_x4((uint32_t *)(void *)(dst + k + 0), q0);
+        store_ascii_quad_neon(dst + k, vmovl_u16(vget_low_u16(cp16)), fg4, bg_flags);
+        store_ascii_quad_neon(dst + k + 4, vmovl_u16(vget_high_u16(cp16)), fg4, bg_flags);
+    }
+    if (k + 4 <= n) {
+        put_ascii_quad_neon(dst + k, src + k, fg4, bg_flags);
+        k += 4;
     }
     if (k < n && n >= 4) {
-        uint32_t bytes;
-        memcpy(&bytes, src + n - 4, sizeof(bytes));
-        uint16x4_t cp16 = vget_low_u16(vmovl_u8(vreinterpret_u8_u32(vdup_n_u32(bytes))));
-        uint32x4_t cp_lo = vmovl_u16(cp16);
-        uint32x4_t c_fg_01 = vzip1q_u32(cp_lo, fg4);
-        uint32x4_t c_fg_23 = vzip2q_u32(cp_lo, fg4);
-        uint32x4_t cell0 = (uint32x4_t)vzip1q_u64((uint64x2_t)c_fg_01, (uint64x2_t)bg_flags);
-        uint32x4_t cell1 = (uint32x4_t)vzip2q_u64((uint64x2_t)c_fg_01, (uint64x2_t)bg_flags);
-        uint32x4_t cell2 = (uint32x4_t)vzip1q_u64((uint64x2_t)c_fg_23, (uint64x2_t)bg_flags);
-        uint32x4_t cell3 = (uint32x4_t)vzip2q_u64((uint64x2_t)c_fg_23, (uint64x2_t)bg_flags);
-        uint32x4x4_t q0 = {{ cell0, cell1, cell2, cell3 }};
-        vst1q_u32_x4((uint32_t *)(void *)(dst + n - 4), q0);
+        put_ascii_quad_neon(dst + n - 4, src + n - 4, fg4, bg_flags);
         k = n;
     }
 #elif defined(__x86_64__)