More work on SIMD utf-8 decode

kitty/simd-string-impl.h

@@ -22,25 +22,34 @@ _Pragma("clang diagnostic pop")
 #define CONCAT(A, B) A##B
 #define CONCAT_EXPAND(A, B) CONCAT(A,B)
 #define FUNC(name) CONCAT_EXPAND(name##_, BITS)
-#define integer_t CONCAT_EXPAND(CONCAT_EXPAND(__m, BITS), i)
+#define integer_t CONCAT_EXPAND(CONCAT_EXPAND(simde__m, BITS), i)
+#define shift_right_by_bytes128 simde_mm_srli_si128
 
 #if BITS == 128
 #define set1_epi8 simde_mm_set1_epi8
 #define load_unaligned simde_mm_loadu_si128
+#define store_aligned simde_mm_store_si128
 #define cmpeq_epi8 simde_mm_cmpeq_epi8
 #define or_si simde_mm_or_si128
 #define movemask_epi8 simde_mm_movemask_epi8
+#define extract_lower_quarter_as_chars simde_mm_cvtepu8_epi32
+#define shift_left_by_chars simde_mm_slli_epi32
+#define shift_left_by_bytes simde_mm_slli_si128
 #else
 #define set1_epi8 simde_mm256_set1_epi8
 #define load_unaligned simde_mm256_loadu_si256
+#define store_aligned simde_mm256_store_si256
 #define cmpeq_epi8 simde_mm256_cmpeq_epi8
 #define or_si simde_mm256_or_si256
 #define movemask_epi8 simde_mm256_movemask_epi8
+#define extract_lower_half_as_chars simde_mm256_cvtepu8_epi32
+#define shift_left_by_chars simde_mm256_slli_epi32
+#define shift_left_by_bytes simde_mm256_slli_si256
 #endif
 
 static inline const uint8_t*
 FUNC(find_either_of_two_bytes)(const uint8_t *haystack, const size_t sz, const uint8_t a, const uint8_t b) {
-    integer_t a_vec = set1_epi8(a), b_vec = set1_epi8(b);
+    const integer_t a_vec = set1_epi8(a), b_vec = set1_epi8(b);
     for (const uint8_t* limit = haystack + sz; haystack < limit; haystack += sizeof(integer_t)) {
         const integer_t chunk = load_unaligned((integer_t*)haystack);
         const integer_t a_cmp = cmpeq_epi8(chunk, a_vec);
@@ -56,10 +65,57 @@ FUNC(find_either_of_two_bytes)(const uint8_t *haystack, const size_t sz, const u
     return NULL;
 }
 
-static inline unsigned
+static inline bool
 FUNC(utf8_decode_to_esc)(UTF8Decoder *d, const uint8_t *src, size_t src_sz) {
-    (void)d; (void)src; (void)src_sz;
-    return 0;
+    d->output_sz = 0; d->num_consumed = 0;
+    src_sz = MIN(src_sz, sizeof(integer_t));
+    integer_t vec = load_unaligned((integer_t*)src);
+
+    const integer_t esc_vec = set1_epi8(0x1b);
+    const integer_t esc_cmp = cmpeq_epi8(vec, esc_vec);
+    const int esc_test_mask = movemask_epi8(esc_cmp);
+    bool sentinel_found = false;
+    const unsigned num_of_bytes_to_first_esc = __builtin_ctz(esc_test_mask);
+    if (num_of_bytes_to_first_esc < src_sz) {
+        sentinel_found = true;
+        d->num_consumed = num_of_bytes_to_first_esc + 1;  // esc is also consumed
+        src_sz = d->num_consumed - 1;
+    } else d->num_consumed = src_sz;
+
+    const int ascii_test_mask = movemask_epi8(vec);
+    const unsigned num_of_bytes_to_first_non_ascii_byte = __builtin_ctz(ascii_test_mask);
+
+    if (num_of_bytes_to_first_non_ascii_byte >= src_sz) {  // no bytes with high bit (0x80) set, so just plain ASCII
+#if BITS == 128
+        for (const uint32_t *limit = d->output + src_sz, *p = d->output; p < limit; p += sizeof(integer_t)/sizeof(uint32_t)) {
+            const integer_t unpacked = extract_lower_quarter_as_chars(vec);
+            store_aligned((integer_t*)p, unpacked);
+            vec = shift_right_by_bytes128(vec, sizeof(integer_t)/sizeof(uint32_t));
+        }
+#else
+        const uint32_t *limit = d->output + src_sz, *p = d->output;
+        simde__m128i x = simde_mm256_extractf128_si256(vec, 0);
+        integer_t unpacked = extract_lower_half_as_chars(x);
+        store_aligned((integer_t*)p, unpacked); p += sizeof(integer_t)/sizeof(uint32_t);
+        if (p < limit) {
+            x = shift_right_by_bytes128(x, sizeof(integer_t)/sizeof(uint32_t));
+            unpacked = extract_lower_half_as_chars(x);
+            store_aligned((integer_t*)p, unpacked); p += sizeof(integer_t)/sizeof(uint32_t);
+            if (p < limit) {
+                x = simde_mm256_extractf128_si256(vec, 1);
+                unpacked = extract_lower_half_as_chars(x);
+                store_aligned((integer_t*)p, unpacked); p += sizeof(integer_t)/sizeof(uint32_t);
+                if (p < limit) {
+                    x = shift_right_by_bytes128(x, sizeof(integer_t)/sizeof(uint32_t));
+                    unpacked = extract_lower_half_as_chars(x);
+                    store_aligned((integer_t*)p, unpacked); p += sizeof(integer_t)/sizeof(uint32_t);
+                }
+            }
+        }
+#endif
+        d->output_sz = src_sz;
+    }
+    return sentinel_found;
 }
 
 
@@ -67,8 +123,15 @@ FUNC(utf8_decode_to_esc)(UTF8Decoder *d, const uint8_t *src, size_t src_sz) {
 #undef integer_t
 #undef set1_epi8
 #undef load_unaligned
+#undef store_aligned
 #undef cmpeq_epi8
 #undef or_si
 #undef movemask_epi8
 #undef CONCAT
 #undef CONCAT_EXPAND
+#undef BITS
+#undef shift_left_by_chars
+#undef shift_left_by_bytes
+#undef shift_right_by_bytes128
+#undef extract_lower_quarter_as_chars
+#undef extract_lower_half_as_chars

kitty/simd-string.c

@@ -81,19 +81,30 @@ test_utf8_decode_to_sentinel(PyObject *self UNUSED, PyObject *args) {
     static UTF8Decoder d = {0};
     if (!PyArg_ParseTuple(args, "s#|i", &src, &src_sz, &which_function)) return NULL;
     bool found_sentinel = false;
+    bool(*func)(UTF8Decoder*, const uint8_t*, size_t sz) = utf8_decode_to_esc;
     switch(which_function) {
         case -1:
             zero_at_ptr(&d); Py_RETURN_NONE;
         case 1:
-            found_sentinel = utf8_decode_to_esc(&d, src, src_sz); break;
+            func = utf8_decode_to_esc_scalar; break;
         case 2:
-            found_sentinel = utf8_decode_to_esc_128(&d, src, src_sz); break;
+            func = utf8_decode_to_esc_128; break;
         case 3:
-            found_sentinel = utf8_decode_to_esc_256(&d, src, src_sz); break;
-        default:
-            found_sentinel = utf8_decode_to_esc(&d, src, src_sz); break;
+            func = utf8_decode_to_esc_256; break;
     }
-    return Py_BuildValue("ON", found_sentinel ? Py_True : Py_False, PyUnicode_FromKindAndData(PyUnicode_4BYTE_KIND, d.output, d.output_sz));
+    RAII_PyObject(ans, PyUnicode_FromString(""));
+    ssize_t p = 0;
+    while(p < src_sz && !found_sentinel) {
+        found_sentinel = func(&d, src + p, src_sz - p);
+        p += d.num_consumed;
+        if (d.output_sz) {
+            RAII_PyObject(temp, PyUnicode_FromKindAndData(PyUnicode_4BYTE_KIND, d.output, d.output_sz));
+            PyObject *t = PyUnicode_Concat(ans, temp);
+            Py_DECREF(ans);
+            ans = t;
+        }
+    }
+    return Py_BuildValue("OO", found_sentinel ? Py_True : Py_False, ans);
 }
 // }}}
 

kitty_tests/parser.py

@@ -6,7 +6,7 @@
 from binascii import hexlify
 from functools import partial
 
-from kitty.fast_data_types import CURSOR_BLOCK, VT_PARSER_BUFFER_SIZE, base64_decode, base64_encode
+from kitty.fast_data_types import CURSOR_BLOCK, VT_PARSER_BUFFER_SIZE, base64_decode, base64_encode, test_utf8_decode_to_sentinel
 from kitty.notify import NotificationCommand, handle_notification_cmd, notification_activated, reset_registry
 
 from . import BaseTest, parse_bytes
@@ -193,6 +193,20 @@ def test_utf8_parsing(self):
         pb(b'"\xe0\xa0"', '"\ufffd"')
         pb(b'"\xf0\x9f\x98"', '"\ufffd"')
 
+    def test_utf8_simd_decode(self):
+        def t(x, which=2, reset=True):
+            if reset:
+                test_utf8_decode_to_sentinel(b'', -1)
+            expected = test_utf8_decode_to_sentinel(x, 1)
+            actual = test_utf8_decode_to_sentinel(x, which)
+            self.ae(expected, actual)
+        for which in (2, 3):
+            with self.subTest(which=which):
+                x = partial(t, which=which)
+                x('abcd1234efgh5678')
+                x('abc\x1bd1234efgh5678')
+                x('abcd1234efgh5678ijklABCDmnopEFGH')
+
     def test_esc_codes(self):
         s = self.create_screen()
         pb = partial(self.parse_bytes_dump, s)