CERBERUS/kitty
1
0
Fork 0
mirror of https://github.com/kovidgoyal/kitty.git synced 2026-09-04 00:46:17 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 3

Add AVX 512 implementations for remaining SIMD functions

Browse source
This commit is contained in:
Kovid Goyal 2026-08-18 09:03:29 +05:30
parent acadec4a70
commit 0ecb10d158
No known key found for this signature in database
GPG key ID: 06BC317B515ACE7C
5 changed files with 128 additions and 37 deletions
Download patch file Download diff file Expand all files Collapse all files

86
kitty/simd-string-512.c
View file

@ -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 <immintrin.h>
// 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

12
kitty/simd-string.c
View file

@ -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);
}

3
kitty/simd-string.h
View file

@ -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);

26
kitty_tests/graphics.py
View file

@ -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

38
kitty_tests/parser.py
View file

@ -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)