diff --git a/kitty/simd-string-512.c b/kitty/simd-string-512.c index 198f44e02..29cf97e46 100644 --- a/kitty/simd-string-512.c +++ b/kitty/simd-string-512.c @@ -4,12 +4,14 @@ * * Distributed under terms of the GPL3 license. * - * AVX-512 (F, BW, VL, VBMI2) implementation of utf8_decode_to_esc. Unlike the - * 128/256 bit implementations in simd-string-impl.h this is x86-64 only and - * uses native intrinsics, since the algorithm is built around mask registers - * and vpcompressb, which have no efficient equivalents on other platforms. - * The overall algorithm is the same as the one in simd-string-impl.h, see - * there for detailed comments, only the differences are commented here. + * AVX-512 (F, BW, VL, VBMI2) implementations of utf8_decode_to_esc, + * find_either_of_two_bytes, xor_data64 and printable_ascii_run_length. Unlike + * the 128/256 bit implementations in simd-string-impl.h these are x86-64 only + * and use native intrinsics, since the algorithms are built around mask + * registers and masked loads/stores (and vpcompressb for UTF-8 decoding), + * which have no efficient equivalents on other platforms. The overall UTF-8 + * decoding algorithm is the same as the one in simd-string-impl.h, see there + * for detailed comments, only the differences are commented here. */ #include "data-types.h" @@ -20,6 +22,63 @@ #include "charsets.h" #include +// Masked loads and stores fault-suppress the masked out bytes, so unlike the +// 128/256 bit implementations no alignment gymnastics are needed to avoid +// reading or writing beyond the ends of buffers. + +const uint8_t * +find_either_of_two_bytes_512(const uint8_t *haystack, const size_t sz, const uint8_t a, const uint8_t b) { + const __m512i a_vec = _mm512_set1_epi8((char)a), b_vec = _mm512_set1_epi8((char)b); + const uint8_t *ans = NULL; + for (size_t i = 0; i < sz; i += 64) { + const uint64_t chunk_bits = sz - i >= 64 ? ~0ull : (1ull << (sz - i)) - 1; + // the masked load zeroes the bytes beyond the end of the haystack, exclude + // them from the matches as a or b could be zero + const __m512i chunk = _mm512_maskz_loadu_epi8(chunk_bits, haystack + i); + const uint64_t matches = chunk_bits & (_mm512_cmpeq_epi8_mask(chunk, a_vec) | _mm512_cmpeq_epi8_mask(chunk, b_vec)); + if (matches) { + ans = haystack + i + __builtin_ctzll(matches); + break; + } + } + _mm256_zeroupper(); + return ans; +} + +void +xor_data64_512(const uint8_t key[64], uint8_t *data, const size_t data_sz) { + // The key size equals the register width, so unlike the 128/256 bit + // implementations no rotation of the key is ever needed. + const __m512i key_vec = _mm512_loadu_si512(key); + size_t i = 0; + for (; i + 64 <= data_sz; i += 64) _mm512_storeu_si512(data + i, _mm512_xor_si512(_mm512_loadu_si512(data + i), key_vec)); + if (i < data_sz) { + const uint64_t tail_bits = (1ull << (data_sz - i)) - 1; + _mm512_mask_storeu_epi8(data + i, tail_bits, _mm512_xor_si512(_mm512_maskz_loadu_epi8(tail_bits, data + i), key_vec)); + } + _mm256_zeroupper(); +} + +size_t +printable_ascii_run_length_512(const uint32_t *chars, const size_t sz) { + // Length of the prefix of chars that contains only printable ASCII codepoints, 32 <= ch <= 126 + const __m512i lower = _mm512_set1_epi32(32), upper = _mm512_set1_epi32(126); + size_t ans = sz; + for (size_t i = 0; i < sz; i += 16) { + const uint16_t chunk_bits = sz - i >= 16 ? 0xffff : (uint16_t)((1u << (sz - i)) - 1); + // the masked load zeroes the chars beyond the end of the buffer, exclude them + // from the non printable chars as zero is itself non printable + const __m512i chunk = _mm512_maskz_loadu_epi32(chunk_bits, chars + i); + const uint16_t non_printable = chunk_bits & (_mm512_cmplt_epu32_mask(chunk, lower) | _mm512_cmpgt_epu32_mask(chunk, upper)); + if (non_printable) { + ans = i + __builtin_ctz(non_printable); + break; + } + } + _mm256_zeroupper(); + return ans; +} + #define do_one_byte \ const uint8_t ch = src[pos++]; \ switch (decode_utf8(&d->state.cur, &d->state.codep, ch)) { \ @@ -302,4 +361,19 @@ utf8_decode_to_esc_512(UTF8Decoder *d UNUSED, const uint8_t *src UNUSED, size_t fatal("No AVX-512 implementation for this platform"); } +const uint8_t * +find_either_of_two_bytes_512(const uint8_t *haystack UNUSED, const size_t sz UNUSED, const uint8_t a UNUSED, const uint8_t b UNUSED) { + fatal("No AVX-512 implementation for this platform"); +} + +void +xor_data64_512(const uint8_t key[64] UNUSED, uint8_t *data UNUSED, const size_t data_sz UNUSED) { + fatal("No AVX-512 implementation for this platform"); +} + +size_t +printable_ascii_run_length_512(const uint32_t *chars UNUSED, const size_t sz UNUSED) { + fatal("No AVX-512 implementation for this platform"); +} + #endif // x86-64 diff --git a/kitty/simd-string.c b/kitty/simd-string.c index 5651ea217..3a418ba9b 100644 --- a/kitty/simd-string.c +++ b/kitty/simd-string.c @@ -140,6 +140,7 @@ test_find_either_of_two_bytes(PyObject *self UNUSED, PyObject *args) { case 1: func = find_either_of_two_bytes_scalar; break; case 2: func = find_either_of_two_bytes_128; break; case 3: func = find_either_of_two_bytes_256; break; + case 4: func = find_either_of_two_bytes_512; break; case 0: break; default: PyErr_SetString(PyExc_ValueError, "Unknown which_function"); return NULL; } @@ -167,6 +168,7 @@ test_printable_ascii_run_length(PyObject *self UNUSED, PyObject *args) { case 1: func = printable_ascii_run_length_scalar; break; case 2: func = printable_ascii_run_length_128; break; case 3: func = printable_ascii_run_length_256; break; + case 4: func = printable_ascii_run_length_512; break; case 0: break; default: PyErr_SetString(PyExc_ValueError, "Unknown which_function"); return NULL; } @@ -188,6 +190,7 @@ test_xor64(PyObject *self UNUSED, PyObject *args) { case 1: func = xor_data64_scalar; break; case 2: func = xor_data64_128; break; case 3: func = xor_data64_256; break; + case 4: func = xor_data64_512; break; case 0: break; default: PyErr_SetString(PyExc_ValueError, "Unknown which_function"); return NULL; } @@ -274,15 +277,18 @@ init_simd(void *x) { if (has_avx512) { A(has_avx512, True); utf8_decode_to_esc_impl = utf8_decode_to_esc_512; + find_either_of_two_bytes_impl = find_either_of_two_bytes_512; + xor_data64_impl = xor_data64_512; + printable_ascii_run_length_impl = printable_ascii_run_length_512; } else { A(has_avx512, False); } if (has_avx2) { A(has_avx2, True); - find_either_of_two_bytes_impl = find_either_of_two_bytes_256; + if (find_either_of_two_bytes_impl == find_either_of_two_bytes_scalar) find_either_of_two_bytes_impl = find_either_of_two_bytes_256; if (utf8_decode_to_esc_impl == utf8_decode_to_esc_scalar) utf8_decode_to_esc_impl = utf8_decode_to_esc_256; - xor_data64_impl = xor_data64_256; - printable_ascii_run_length_impl = printable_ascii_run_length_256; + if (xor_data64_impl == xor_data64_scalar) xor_data64_impl = xor_data64_256; + if (printable_ascii_run_length_impl == printable_ascii_run_length_scalar) printable_ascii_run_length_impl = printable_ascii_run_length_256; } else { A(has_avx2, False); } diff --git a/kitty/simd-string.h b/kitty/simd-string.h index 083bb11cb..90fa3543e 100644 --- a/kitty/simd-string.h +++ b/kitty/simd-string.h @@ -68,7 +68,10 @@ bool utf8_decode_to_esc_256(UTF8Decoder *d, const uint8_t *src, size_t src_sz); bool utf8_decode_to_esc_512(UTF8Decoder *d, const uint8_t *src, size_t src_sz); const uint8_t *find_either_of_two_bytes_128(const uint8_t *haystack, const size_t sz, const uint8_t a, const uint8_t b); const uint8_t *find_either_of_two_bytes_256(const uint8_t *haystack, const size_t sz, const uint8_t a, const uint8_t b); +const uint8_t *find_either_of_two_bytes_512(const uint8_t *haystack, const size_t sz, const uint8_t a, const uint8_t b); void xor_data64_128(const uint8_t key[64], uint8_t *data, const size_t data_sz); void xor_data64_256(const uint8_t key[64], uint8_t *data, const size_t data_sz); +void xor_data64_512(const uint8_t key[64], uint8_t *data, const size_t data_sz); size_t printable_ascii_run_length_128(const uint32_t *chars, const size_t sz); size_t printable_ascii_run_length_256(const uint32_t *chars, const size_t sz); +size_t printable_ascii_run_length_512(const uint32_t *chars, const size_t sz); diff --git a/kitty_tests/graphics.py b/kitty_tests/graphics.py index 12d8330a1..142e48cf9 100644 --- a/kitty_tests/graphics.py +++ b/kitty_tests/graphics.py @@ -11,7 +11,7 @@ from contextlib import suppress from dataclasses import dataclass from io import BytesIO -from kitty.fast_data_types import base64_decode, base64_encode, has_avx2, has_sse4_2, load_png_data, shm_unlink, shm_write, test_xor64 +from kitty.fast_data_types import base64_decode, base64_encode, load_png_data, shm_unlink, shm_write from .base import BaseTest, parse_bytes @@ -197,30 +197,6 @@ def make_send_command(screen): class TestGraphics(BaseTest): - def test_xor_data(self): - base_data = b'\x01' * 64 - key = b'\x02' * 64 - sizes = [] - if has_sse4_2: - sizes.append(2) - if has_avx2: - sizes.append(3) - sizes.append(0) - - def t(key, data, align_offset=0): - expected = test_xor64(key, data, 1, 0) - for which_function in sizes: - actual = test_xor64(key, data, which_function, align_offset) - self.ae(expected, actual, f'{align_offset=} {len(data)=}') - - t(key, b'') - - for base in (b'abc', base_data): - for extra in range(len(base_data)): - for align_offset in range(64): - data = base + base_data[:extra] - t(key, data, align_offset) - def test_disk_cache(self): s = self.create_screen() dc = s.grman.disk_cache diff --git a/kitty_tests/parser.py b/kitty_tests/parser.py index 332471a84..bb3d9b54e 100644 --- a/kitty_tests/parser.py +++ b/kitty_tests/parser.py @@ -16,6 +16,7 @@ from kitty.fast_data_types import ( test_find_either_of_two_bytes, test_printable_ascii_run_length, test_utf8_decode_to_sentinel, + test_xor64, ) from .base import BaseTest, parse_bytes @@ -743,6 +744,8 @@ class TestParser(BaseTest): sizes.append(2) if has_avx2: sizes.append(3) + if has_avx512: + sizes.append(4) sizes.append(0) def test(buf, a, b, align_offset=0): @@ -753,7 +756,7 @@ class TestParser(BaseTest): self.ae(expected, actual, f'Failed for: {buf!r} {a=} {b=} at {sz=} and {align_offset=}') q = 'abc' - for off in range(32): + for off in range(64): test(q, '<', '>', off) test(q, ' ', 'b', off) test(q, '<', 'a', off) @@ -761,9 +764,9 @@ class TestParser(BaseTest): test(q, 'c', '>', off) def tests(buf, a, b): - for sz in (0, 16, 32, 64, 79): + for sz in (0, 16, 32, 64, 79, 128, 133): buf = (' ' * sz) + buf - for align_offset in range(32): + for align_offset in range(64): test(buf, a, b, align_offset) tests('', '<', '>') @@ -781,6 +784,8 @@ class TestParser(BaseTest): impls.append(2) if has_avx2: impls.append(3) + if has_avx512: + impls.append(4) impls.append(0) def test(text): @@ -799,6 +804,33 @@ class TestParser(BaseTest): test(prefix + bad + 'xyz') test(prefix + bad * 3 + prefix) + def test_xor_data64(self): + # varying bytes in the key and data so that alignment/key rotation bugs are caught + base_data = bytes(range(64)) + key = bytes(range(101, 165)) + sizes = [] + if has_sse4_2: + sizes.append(2) + if has_avx2: + sizes.append(3) + if has_avx512: + sizes.append(4) + sizes.append(0) + + def t(key, data, align_offset=0): + expected = test_xor64(key, data, 1, 0) + for which_function in sizes: + actual = test_xor64(key, data, which_function, align_offset) + self.ae(expected, actual, f'{align_offset=} {len(data)=}') + + t(key, b'') + + for base in (b'abc', base_data, base_data * 2): + for extra in range(len(base_data)): + for align_offset in range(64): + data = base + base_data[:extra] + t(key, data, align_offset) + def test_esc_codes(self): s = self.create_screen() pb = partial(self.parse_bytes_dump, s)