Improve performance of SIMD UTF-8 decoding

Optimise the AVX 2 implementation reducing number of operations per block
and also add an AVX-512 implementation used when the CPU supports it at
runtime.
This commit is contained in:
Kovid Goyal 2026-08-16 20:25:03 +05:30
parent e62ec6d540
commit 07ea7583ac
No known key found for this signature in database
GPG key ID: 06BC317B515ACE7C
8 changed files with 606 additions and 179 deletions

View file

@ -188,6 +188,8 @@ Detailed list of changes
- Improve throughput when processing large amounts of text by ~50% for ASCII and ~15% for Unicode by writing runs of plain ASCII chars to the screen in batches and skipping unnecessary bookkeeping in the scrolling and tab handling hot paths
- Speed up SIMD UTF-8 decoding: ~20% faster for ASCII and ~50% faster for multi-byte text on both x86 and ARM, by processing pairs of vectors of plain ASCII text at a time and compacting decoded codepoints within 128-bit lanes, avoiding expensive cross-lane operations. Also add an AVX-512 based decoder, another ~75% faster for multi-byte text, used automatically on CPUs that support it (Intel Ice Lake+, AMD Zen 4+). Additionally fix the SIMD decoders failing to emit U+FFFD for an incomplete UTF-8 sequence cut off by an escape code
- The :opt:`scrollbar` option now takes a new value ``scrolled-or-hovered`` to also show the scrollbar when the mouse moves over the scrollbar region (:pull:`10345`)
- A new :code:`kitten @ screenshot` remote control command to take a pixel perfect PNG screenshot of an OS Window, tab or window

305
kitty/simd-string-512.c Normal file
View file

@ -0,0 +1,305 @@
/*
* simd-string-512.c
* Copyright (C) 2026 Kovid Goyal <kovid at kovidgoyal.net>
*
* 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.
*/
#include "data-types.h"
#include "simd-string.h"
#if (defined(__x86_64__) || defined(__amd64__)) && !defined(KITTY_NO_SIMD)
#include "charsets.h"
#include <immintrin.h>
#define do_one_byte \
const uint8_t ch = src[pos++]; \
switch (decode_utf8(&d->state.cur, &d->state.codep, ch)) { \
case UTF8_ACCEPT: d->output.storage[d->output.pos++] = d->state.codep; break; \
case UTF8_REJECT: { \
const bool prev_was_accept = d->state.prev == UTF8_ACCEPT; \
zero_at_ptr(&d->state); \
d->output.storage[d->output.pos++] = 0xfffd; \
if (!prev_was_accept) { \
pos--; \
continue; /* so that prev is correct */ \
} \
} break; \
} \
d->state.prev = d->state.cur;
static size_t
scalar_decode_to_accept(UTF8Decoder *d, const uint8_t *src, size_t src_sz) {
size_t pos = 0;
utf8_decoder_ensure_capacity(d, src_sz);
while (pos < src_sz && d->state.cur != UTF8_ACCEPT) { do_one_byte }
return pos;
}
static size_t
scalar_decode_all(UTF8Decoder *d, const uint8_t *src, size_t src_sz) {
size_t pos = 0;
utf8_decoder_ensure_capacity(d, src_sz);
while (pos < src_sz) { do_one_byte }
return pos;
}
#undef do_one_byte
// Shift bytes towards higher memory addresses, filling with zeroes. Done by
// first constructing a copy of the register with the 128-bit lanes shifted up
// by one and then using it to fill in the bytes shifted across lane boundaries.
#define shift_right_by_bytes512(A, n) _mm512_alignr_epi8(A, _mm512_alignr_epi64(A, _mm512_setzero_si512(), 6), 16 - (n))
// The caller must ensure output has sufficient capacity and that src_sz bytes
// are readable from src. Widens directly from memory, which is faster than
// extracting the 128-bit quarters from an already loaded register.
static inline void
output_plain_ascii_from_memory_512(UTF8Decoder *d, const uint8_t *src, const size_t src_sz) {
uint32_t *out = d->output.storage + d->output.pos;
for (size_t i = 0; i < src_sz; i += 16) _mm512_storeu_si512(out + i, _mm512_cvtepu8_epi32(_mm_loadu_si128((const __m128i *)(src + i))));
d->output.pos += src_sz;
}
// The caller must ensure output has sufficient capacity
static inline void
output_plain_ascii_512(UTF8Decoder *d, const __m512i vec, const size_t src_sz) {
uint32_t *out = d->output.storage + d->output.pos;
_mm512_storeu_si512(out, _mm512_cvtepu8_epi32(_mm512_castsi512_si128(vec)));
if (src_sz > 16) {
_mm512_storeu_si512(out + 16, _mm512_cvtepu8_epi32(_mm512_extracti32x4_epi32(vec, 1)));
if (src_sz > 32) {
_mm512_storeu_si512(out + 32, _mm512_cvtepu8_epi32(_mm512_extracti32x4_epi32(vec, 2)));
if (src_sz > 48) { _mm512_storeu_si512(out + 48, _mm512_cvtepu8_epi32(_mm512_extracti32x4_epi32(vec, 3))); }
}
}
d->output.pos += src_sz;
}
// The caller must ensure output has sufficient capacity. The outputN vectors
// contain the three little-endian bytes of each codepoint, already compacted.
static inline void
output_unicode_512(UTF8Decoder *d, __m512i output1, __m512i output2, __m512i output3, const unsigned num_codepoints) {
uint32_t *out = d->output.storage + d->output.pos;
for (unsigned i = 0; i < num_codepoints; i += 16) {
const __m512i unpacked1 = _mm512_cvtepu8_epi32(_mm512_castsi512_si128(output1));
const __m512i unpacked2 = _mm512_slli_epi32(_mm512_cvtepu8_epi32(_mm512_castsi512_si128(output2)), 8);
const __m512i unpacked3 = _mm512_slli_epi32(_mm512_cvtepu8_epi32(_mm512_castsi512_si128(output3)), 16);
_mm512_storeu_si512(out + i, _mm512_or_si512(_mm512_or_si512(unpacked1, unpacked2), unpacked3));
// shift the registers down by 16 bytes for the next group of codepoints
output1 = _mm512_alignr_epi64(_mm512_setzero_si512(), output1, 2);
output2 = _mm512_alignr_epi64(_mm512_setzero_si512(), output2, 2);
output3 = _mm512_alignr_epi64(_mm512_setzero_si512(), output3, 2);
}
d->output.pos += num_codepoints;
}
bool
utf8_decode_to_esc_512(UTF8Decoder *d, const uint8_t *src_data, size_t src_len) {
d->output.pos = 0;
d->num_consumed = 0;
if (d->state.cur != UTF8_ACCEPT) {
// Finish the trailing sequence only
d->num_consumed = scalar_decode_to_accept(d, src_data, src_len);
src_data += d->num_consumed;
src_len -= d->num_consumed;
}
utf8_decoder_ensure_capacity(d, src_len + 64);
const __m512i esc_vec = _mm512_set1_epi8(0x1b), zero = _mm512_setzero_si512(), one = _mm512_set1_epi8(1), two = _mm512_set1_epi8(2),
three = _mm512_set1_epi8(3);
const uint8_t *limit = src_data + src_len, *p = src_data, *start_of_current_chunk = src_data;
bool sentinel_found = false;
unsigned chunk_src_sz = 0;
unsigned num_of_trailing_bytes = 0;
bool prev_chunk_was_all_ascii = true;
#define abort_with_invalid_utf8() \
{ \
scalar_decode_all(d, start_of_current_chunk, chunk_src_sz + num_of_trailing_bytes); \
d->num_consumed += num_of_trailing_bytes; \
break; \
}
#define handle_trailing_bytes() \
if (num_of_trailing_bytes) { \
if (p >= limit) { \
scalar_decode_all(d, p - num_of_trailing_bytes, num_of_trailing_bytes); \
d->num_consumed += num_of_trailing_bytes; \
break; \
} \
p -= num_of_trailing_bytes; \
}
while (p < limit && !sentinel_found) {
// Fast path: process pairs of full chunks that contain only ASCII and no ESC
if (prev_chunk_was_all_ascii) {
while ((size_t)(limit - p) >= 128) {
const __m512i v1 = _mm512_loadu_si512(p), v2 = _mm512_loadu_si512(p + 64);
const __mmask64 bad = _mm512_cmpeq_epi8_mask(v1, esc_vec) | _mm512_cmpeq_epi8_mask(v2, esc_vec) | _mm512_movepi8_mask(_mm512_or_si512(v1, v2));
if (bad) break;
output_plain_ascii_from_memory_512(d, p, 128);
d->num_consumed += 128;
p += 128;
}
if (p >= limit) break;
}
chunk_src_sz = MIN((size_t)(limit - p), 64u);
// The masked load both avoids reading beyond the end of the input and zeroes the tail bytes
const uint64_t chunk_bits = chunk_src_sz == 64 ? ~0ull : (1ull << chunk_src_sz) - 1;
__m512i vec = _mm512_maskz_loadu_epi8(chunk_bits, p);
start_of_current_chunk = p;
p += chunk_src_sz;
uint64_t ascii_mask = _mm512_movepi8_mask(vec); // bit set for every non-ASCII byte
const uint64_t esc_mask = _mm512_cmpeq_epi8_mask(vec, esc_vec);
if (esc_mask) {
sentinel_found = true;
chunk_src_sz = __builtin_ctzll(esc_mask);
const uint64_t kept_bits = (1ull << chunk_src_sz) - 1;
ascii_mask &= kept_bits;
vec = _mm512_maskz_mov_epi8(kept_bits, vec);
d->num_consumed += chunk_src_sz + 1; // esc is also consumed
if (!chunk_src_sz) continue;
} else d->num_consumed += chunk_src_sz;
num_of_trailing_bytes = 0;
bool check_for_trailing_bytes = !sentinel_found;
start_classification:
if (!ascii_mask) { // no bytes with high bit (0x80) set, so just plain ASCII
output_plain_ascii_512(d, vec, chunk_src_sz);
prev_chunk_was_all_ascii = true;
handle_trailing_bytes();
continue;
}
prev_chunk_was_all_ascii = false;
// Classify the bytes by the length of the UTF-8 sequence they start. As in the
// 128/256-bit implementations, this is an initial, potential classification:
// 0xC0, 0xC1 and 0xF5..0xFF are classified as starter bytes here and rejected
// in the validation checks below.
const __mmask64 starts_two = _mm512_cmpge_epu8_mask(vec, _mm512_set1_epi8((char)0xc0));
const __mmask64 starts_three = _mm512_cmpge_epu8_mask(vec, _mm512_set1_epi8((char)0xe0));
const __mmask64 starts_four = _mm512_cmpge_epu8_mask(vec, _mm512_set1_epi8((char)0xf0));
// count contains the number of bytes in the sequence for the start byte of every sequence and zero elsewhere
__m512i count = _mm512_maskz_mov_epi8(starts_two, two);
count = _mm512_mask_add_epi8(count, starts_three, count, one);
count = _mm512_mask_add_epi8(count, starts_four, count, one);
// the UTF-8 marker bits to strip from each byte: 0x80 for ASCII and continuation bytes, the sequence length markers for starters
__m512i strip = _mm512_set1_epi8((char)0x80);
strip = _mm512_mask_mov_epi8(strip, starts_two, _mm512_set1_epi8((char)0xc0));
strip = _mm512_mask_mov_epi8(strip, starts_three, _mm512_set1_epi8((char)0xe0));
strip = _mm512_mask_mov_epi8(strip, starts_four, _mm512_set1_epi8((char)0xf0));
// counts contains the number of bytes remaining in each UTF-8 sequence of 2 or more bytes
const __m512i count_subs1 = _mm512_subs_epu8(count, one);
__m512i counts = _mm512_add_epi8(count, shift_right_by_bytes512(count_subs1, 1));
counts = _mm512_add_epi8(counts, shift_right_by_bytes512(_mm512_subs_epu8(counts, two), 2));
// bytes that are neither ASCII nor the last byte of a multi-byte sequence
const uint64_t discarded_locations = _mm512_cmpgt_epu8_mask(counts, one);
// check for an incomplete trailing utf8 sequence
if (check_for_trailing_bytes && ((discarded_locations >> (chunk_src_sz - 1)) & 1)) {
check_for_trailing_bytes = false;
if (start_of_current_chunk[chunk_src_sz - 1] >= 0xc0) num_of_trailing_bytes = 1; // 2-, 3- and 4-byte characters with only 1 byte left
else if (chunk_src_sz > 1 && start_of_current_chunk[chunk_src_sz - 2] >= 0xe0)
num_of_trailing_bytes = 2; // 3- and 4-byte characters with only 1 byte left
else if (chunk_src_sz > 2 && start_of_current_chunk[chunk_src_sz - 3] >= 0xf0)
num_of_trailing_bytes = 3; // 4-byte characters with only 3 bytes left
// num_of_trailing_bytes can be zero, when overlapping sequences near the end
// of the chunk make counts[last] > 1 without an actual incomplete trailing
// sequence. Reclassification then detects the overlap as invalid.
if (num_of_trailing_bytes) {
chunk_src_sz -= num_of_trailing_bytes;
d->num_consumed -= num_of_trailing_bytes;
if (!chunk_src_sz) { abort_with_invalid_utf8(); }
const uint64_t kept_bits = (1ull << chunk_src_sz) - 1;
ascii_mask &= kept_bits;
vec = _mm512_maskz_mov_epi8(kept_bits, vec);
}
goto start_classification;
}
// Validation, see simd-string-impl.h for detailed comments. Note that unlike
// there, ill-formed chunks are detected with mask registers directly.
uint64_t chunk_is_invalid = ascii_mask ^ _mm512_cmpgt_epu8_mask(counts, zero);
// 2-byte sequence starters must be >= 0xC2 to avoid overlong encodings
chunk_is_invalid |= _mm512_mask_cmplt_epu8_mask(starts_two, vec, _mm512_set1_epi8((char)0xc2));
// 4-byte sequence starters must be <= 0xF4 to stay within the Unicode codespace
chunk_is_invalid |= _mm512_cmpgt_epu8_mask(vec, _mm512_set1_epi8((char)0xf4));
// continuation bytes' positions must not have starter bytes
chunk_is_invalid |= _mm512_mask_cmpgt_epu8_mask(starts_two, counts, count);
// second byte restrictions for E0, ED, F0 and F4 starters. These are stricter than
// the checks in simd-string-impl.h, also matching ASCII bytes at the follower
// positions, which is fine as those chunks are ill-formed anyway and the scalar
// fallback decodes them correctly.
const __m512i prev_bytes = shift_right_by_bytes512(vec, 1);
chunk_is_invalid |= _mm512_mask_cmplt_epu8_mask(_mm512_cmpeq_epi8_mask(prev_bytes, _mm512_set1_epi8((char)0xe0)), vec, _mm512_set1_epi8((char)0xa0));
chunk_is_invalid |= _mm512_mask_cmpgt_epu8_mask(_mm512_cmpeq_epi8_mask(prev_bytes, _mm512_set1_epi8((char)0xed)), vec, _mm512_set1_epi8((char)0x9f));
chunk_is_invalid |= _mm512_mask_cmplt_epu8_mask(_mm512_cmpeq_epi8_mask(prev_bytes, _mm512_set1_epi8((char)0xf0)), vec, _mm512_set1_epi8((char)0x90));
chunk_is_invalid |= _mm512_mask_cmpgt_epu8_mask(_mm512_cmpeq_epi8_mask(prev_bytes, _mm512_set1_epi8((char)0xf4)), vec, _mm512_set1_epi8((char)0x8f));
if (chunk_is_invalid) { abort_with_invalid_utf8(); }
// Assemble the three output bytes for each codepoint, at the position of the
// last byte of its sequence, see simd-string-impl.h for detailed comments.
vec = _mm512_andnot_si512(strip, vec);
const __mmask64 count1_locations = _mm512_cmpeq_epu8_mask(counts, one);
const __mmask64 count2_locations = _mm512_cmpeq_epu8_mask(counts, two);
const __mmask64 count3_locations = _mm512_cmpeq_epu8_mask(counts, three);
const __mmask64 count4_locations = _mm512_cmpeq_epu8_mask(counts, _mm512_set1_epi8(4));
const __m512i vec_non_ascii = _mm512_maskz_mov_epi8(_mm512_cmpgt_epu8_mask(counts, zero), vec);
const __m512i prev_non_ascii = shift_right_by_bytes512(vec_non_ascii, 1);
const __m512i low_two_bits_of_starter = _mm512_and_si512(_mm512_slli_epi16(prev_non_ascii, 6), _mm512_set1_epi8((char)0xc0));
__m512i output1 = _mm512_mask_mov_epi8(vec, count1_locations, _mm512_or_si512(vec, low_two_bits_of_starter));
__m512i output2 = _mm512_srli_epi32(_mm512_maskz_mov_epi8(count2_locations, vec), 2);
output2 = _mm512_or_si512(
output2, _mm512_and_si512(_mm512_set1_epi8((char)0xf0), _mm512_slli_epi16(_mm512_maskz_mov_epi8(count3_locations << 1, prev_non_ascii), 4)));
output2 = shift_right_by_bytes512(_mm512_maskz_mov_epi8(count2_locations, output2), 1);
__m512i output3 = _mm512_and_si512(three, _mm512_srli_epi32(vec, 4));
output3 = _mm512_or_si512(
output3, _mm512_and_si512(_mm512_set1_epi8((char)0xfc), _mm512_slli_epi16(_mm512_maskz_mov_epi8(count4_locations << 1, prev_non_ascii), 2)));
output3 = shift_right_by_bytes512(_mm512_maskz_mov_epi8(count3_locations, output3), 2);
// Discard the continuation bytes, compacting the codepoints. This replaces the
// whole prefix sum, move and shuffle machinery of the 128/256 implementations.
const uint64_t kept_locations = ~discarded_locations & (chunk_src_sz == 64 ? ~0ull : (1ull << chunk_src_sz) - 1);
output1 = _mm512_maskz_compress_epi8(kept_locations, output1);
output2 = _mm512_maskz_compress_epi8(kept_locations, output2);
output3 = _mm512_maskz_compress_epi8(kept_locations, output3);
const unsigned num_codepoints = __builtin_popcountll(kept_locations);
output_unicode_512(d, output1, output2, output3, num_codepoints);
handle_trailing_bytes();
}
if (sentinel_found && d->state.cur != UTF8_ACCEPT) {
// an incomplete UTF-8 sequence was cut off by the sentinel, matching
// the scalar implementation, emit a replacement char for it
utf8_decoder_ensure_capacity(d, 1);
d->output.storage[d->output.pos++] = 0xfffd;
zero_at_ptr(&d->state);
}
_mm256_zeroupper();
return sentinel_found;
#undef abort_with_invalid_utf8
#undef handle_trailing_bytes
}
#else
bool
utf8_decode_to_esc_512(UTF8Decoder *d UNUSED, const uint8_t *src UNUSED, size_t src_sz UNUSED) {
fatal("No AVX-512 implementation for this platform");
}
#endif // x86-64

View file

@ -92,7 +92,6 @@ _Pragma("clang diagnostic pop")
#define subtract_epi8 simde_mm_sub_epi8
#define create_zero_integer simde_mm_setzero_si128
#define create_all_ones_integer() simde_mm_set1_epi64x(-1)
#define sum_bytes sum_bytes_128
#define zero_upper()
static inline int FUNC(is_zero)(const integer_t a) {
@ -225,86 +224,36 @@ w(right, one_byte, 1) w(right, two_bytes, 2) w(right, four_bytes, 4) w(right, ei
#undef shift_right_by_bytes_macro
#undef shift_left_by_bytes_macro
static inline integer_t shuffle_impl256(const integer_t value, const integer_t shuffle) {
#define K0 \
simde_mm256_setr_epi8( \
0x70, \
0x70, \
0x70, \
0x70, \
0x70, \
0x70, \
0x70, \
0x70, \
0x70, \
0x70, \
0x70, \
0x70, \
0x70, \
0x70, \
0x70, \
0x70, \
-16, \
-16, \
-16, \
-16, \
-16, \
-16, \
-16, \
-16, \
-16, \
-16, \
-16, \
-16, \
-16, \
-16, \
-16, \
-16)
// only used with in-lane shuffle indices (0-15 within each 128-bit lane)
#define shuffle_epi8 simde_mm256_shuffle_epi8
#endif
#define K1 \
simde_mm256_setr_epi8( \
-16, \
-16, \
-16, \
-16, \
-16, \
-16, \
-16, \
-16, \
-16, \
-16, \
-16, \
-16, \
-16, \
-16, \
-16, \
-16, \
0x70, \
0x70, \
0x70, \
0x70, \
0x70, \
0x70, \
0x70, \
0x70, \
0x70, \
0x70, \
0x70, \
0x70, \
0x70, \
0x70, \
0x70, \
0x70)
return or_si(
simde_mm256_shuffle_epi8(value, add_epi8(shuffle, K0)),
simde_mm256_shuffle_epi8(simde_mm256_permute4x64_epi64(value, 0x4E), simde_mm256_add_epi8(shuffle, K1)));
#undef K0
#undef K1
}
#define shuffle_epi8 shuffle_impl256
#define sum_bytes(x) (sum_bytes_128(simde_mm256_extracti128_si256(x, 0)) + sum_bytes_128(simde_mm256_extracti128_si256(x, 1)))
// Byte shifts within each 128-bit lane. These are single instructions on x86
// and cheap NEON ext ops on ARM, unlike the cross-lane full-register shifts
// needed at the 256-bit level. Note that in this codebase shift_right means
// moving bytes to higher memory addresses, which corresponds to the intrinsics
// named "slli" and vice versa.
#if KITTY_SIMD_LEVEL == 128
// for 128-bit registers in-lane and full-register byte shifts are the same thing
#define shift_right_in_lane_by_one_byte shift_right_by_one_byte
#define shift_right_in_lane_by_two_bytes shift_right_by_two_bytes
#define shift_right_in_lane_by_four_bytes shift_right_by_four_bytes
#define shift_right_in_lane_by_eight_bytes shift_right_by_eight_bytes
#define shift_left_in_lane_by_one_byte shift_left_by_one_byte
#define shift_left_in_lane_by_two_bytes shift_left_by_two_bytes
#define shift_left_in_lane_by_four_bytes shift_left_by_four_bytes
#define shift_left_in_lane_by_eight_bytes shift_left_by_eight_bytes
#define numbered_bytes_in_lane numbered_bytes
#else
#define shift_right_in_lane_by_one_byte(A) simde_mm256_slli_si256(A, 1)
#define shift_right_in_lane_by_two_bytes(A) simde_mm256_slli_si256(A, 2)
#define shift_right_in_lane_by_four_bytes(A) simde_mm256_slli_si256(A, 4)
#define shift_right_in_lane_by_eight_bytes(A) simde_mm256_slli_si256(A, 8)
#define shift_left_in_lane_by_one_byte(A) simde_mm256_srli_si256(A, 1)
#define shift_left_in_lane_by_two_bytes(A) simde_mm256_srli_si256(A, 2)
#define shift_left_in_lane_by_four_bytes(A) simde_mm256_srli_si256(A, 4)
#define shift_left_in_lane_by_eight_bytes(A) simde_mm256_srli_si256(A, 8)
#define numbered_bytes_in_lane() set_epi8(15, 14, 13, 12, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1, 0, 15, 14, 13, 12, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1, 0)
#endif
#define print_register_as_bytes(r) \
@ -382,8 +331,7 @@ bytes_to_first_match_ignoring_leading_n(const integer_t vec, uintptr_t num_ignor
#else
static inline int
bytes_to_first_match(const integer_t vec) {
static inline int bytes_to_first_match(const integer_t vec) {
return is_zero(vec) ? -1 : __builtin_ctz(movemask_epi8(vec));
}
@ -498,9 +446,9 @@ FUNC(find_either_of_two_bytes)(const uint8_t *haystack, const size_t sz, const u
#define output_increment sizeof(integer_t) / sizeof(uint32_t)
// The caller must ensure output has sufficient capacity
static inline void
FUNC(output_plain_ascii)(UTF8Decoder *d, integer_t vec, size_t src_sz) {
utf8_decoder_ensure_capacity(d, src_sz);
#if KITTY_SIMD_LEVEL == 128
for (const uint32_t *p = d->output.storage + d->output.pos, *limit = p + src_sz; p < limit; p += output_increment) {
const integer_t unpacked = extract_lower_quarter_as_chars(vec);
@ -535,10 +483,13 @@ FUNC(output_plain_ascii)(UTF8Decoder *d, integer_t vec, size_t src_sz) {
d->output.pos += src_sz;
}
static inline void
FUNC(output_unicode)(UTF8Decoder *d, integer_t output1, integer_t output2, integer_t output3, const size_t num_codepoints) {
utf8_decoder_ensure_capacity(d, 64);
// The caller must ensure output has sufficient capacity.
// The outputN vectors must be compacted within each 128-bit lane, the number of
// codepoints in each lane is passed in separately and the lanes are joined by
// overlapping stores, relying on the output overwrite slack.
#if KITTY_SIMD_LEVEL == 128
static inline void
FUNC(output_unicode)(UTF8Decoder *d, integer_t output1, integer_t output2, integer_t output3, const unsigned num_codepoints) {
for (const uint32_t *p = d->output.storage + d->output.pos, *limit = p + num_codepoints; p < limit; p += output_increment) {
const integer_t unpacked1 = extract_lower_quarter_as_chars(output1);
const integer_t unpacked2 = shift_right_by_one_byte(extract_lower_quarter_as_chars(output2));
@ -549,17 +500,21 @@ FUNC(output_unicode)(UTF8Decoder *d, integer_t output1, integer_t output2, integ
output2 = shift_right_by_bytes128(output2, output_increment);
output3 = shift_right_by_bytes128(output3, output_increment);
}
d->output.pos += num_codepoints;
}
#else
static inline void
FUNC(output_unicode)(
UTF8Decoder *d, integer_t output1, integer_t output2, integer_t output3, const unsigned num_codepoints_lane0, const unsigned num_codepoints_lane1) {
uint32_t *p = d->output.storage + d->output.pos;
const uint32_t *limit = p + num_codepoints;
simde__m128i x1, x2, x3;
#define chunk() \
{ \
const integer_t unpacked1 = extract_lower_half_as_chars(x1); \
const integer_t unpacked2 = shift_right_by_one_byte(extract_lower_half_as_chars(x2)); \
const integer_t unpacked3 = shift_right_by_two_bytes(extract_lower_half_as_chars(x3)); \
store_unaligned((integer_t *)p, or_si(or_si(unpacked1, unpacked2), unpacked3)); \
p += output_increment; \
#define chunk() \
{ \
const integer_t unpacked1 = extract_lower_half_as_chars(x1); \
const integer_t unpacked2 = shift_right_in_lane_by_one_byte(extract_lower_half_as_chars(x2)); \
const integer_t unpacked3 = shift_right_in_lane_by_two_bytes(extract_lower_half_as_chars(x3)); \
store_unaligned((integer_t *)p, or_si(or_si(unpacked1, unpacked2), unpacked3)); \
p += output_increment; \
}
#define extract(which) \
x1 = simde_mm256_extracti128_si256(output1, which); \
@ -571,40 +526,25 @@ FUNC(output_unicode)(UTF8Decoder *d, integer_t output1, integer_t output2, integ
x3 = shift_right_by_bytes128(x3, output_increment);
extract(0);
chunk();
if (p < limit) {
if (num_codepoints_lane0 > output_increment) {
shift();
chunk();
}
p = d->output.storage + d->output.pos + num_codepoints_lane0;
extract(1);
chunk();
if (num_codepoints_lane1 > output_increment) {
shift();
chunk();
if (p < limit) {
extract(1);
chunk();
if (p < limit) {
shift();
chunk();
}
}
}
#undef chunk
#undef extract
#undef shift
d->output.pos += num_codepoints_lane0 + num_codepoints_lane1;
}
#endif
d->output.pos += num_codepoints;
}
#undef output_increment
static inline unsigned
sum_bytes_128(simde__m128i v) {
// Use _mm_sad_epu8 to perform a sum of absolute differences against zero
// This sums up all 8-bit integers in the 128-bit vector and packs the result into a 64-bit integer
simde__m128i sum = simde_mm_sad_epu8(v, simde_mm_setzero_si128());
// At this point, the sum of the first half is in the lower 64 bits, and the sum of the second half is in the upper 64 bits.
// Extract the lower and upper 64-bit sums and add them together.
const unsigned lower_sum = simde_mm_cvtsi128_si32(sum); // Extracts the lower 32 bits
const unsigned upper_sum = simde_mm_cvtsi128_si32(simde_mm_srli_si128(sum, 8)); // Extracts the upper 32 bits
return lower_sum + upper_sum; // Final sum of all bytes
}
#define do_one_byte \
const uint8_t ch = src[pos++]; \
switch (decode_utf8(&d->state.cur, &d->state.codep, ch)) { \
@ -655,36 +595,57 @@ FUNC(utf8_decode_to_esc)(UTF8Decoder *d, const uint8_t *src_data, size_t src_len
src_data += d->num_consumed;
src_len -= d->num_consumed;
}
utf8_decoder_ensure_capacity(d, src_len + 64);
const integer_t esc_vec = set1_epi8(0x1b);
const integer_t zero = create_zero_integer(), one = set1_epi8(1), two = set1_epi8(2), three = set1_epi8(3), four = set1_epi8(4),
numbered = numbered_bytes();
numbered = numbered_bytes(), numbered_in_lane = numbered_bytes_in_lane();
const uint8_t *limit = src_data + src_len, *p = src_data, *start_of_current_chunk = src_data;
bool sentinel_found = false;
unsigned chunk_src_sz = 0;
unsigned num_of_trailing_bytes = 0;
bool prev_chunk_was_all_ascii = true; // avoid paying for the fast path attempt in runs of non-ASCII chunks
while (p < limit && !sentinel_found) {
// Fast path: process pairs of full chunks that contain only ASCII and no ESC
if (prev_chunk_was_all_ascii) {
while ((size_t)(limit - p) >= 2 * sizeof(integer_t)) {
const integer_t v1 = load_unaligned((integer_t *)p), v2 = load_unaligned((integer_t *)(p + sizeof(integer_t)));
const integer_t esc_or_non_ascii = or_si(or_si(cmpeq_epi8(v1, esc_vec), cmpeq_epi8(v2, esc_vec)), or_si(v1, v2));
if (movemask_epi8(esc_or_non_ascii)) break;
FUNC(output_plain_ascii)(d, v1, sizeof(integer_t));
FUNC(output_plain_ascii)(d, v2, sizeof(integer_t));
d->num_consumed += 2 * sizeof(integer_t);
p += 2 * sizeof(integer_t);
}
if (p >= limit) break;
}
chunk_src_sz = MIN((size_t)(limit - p), sizeof(integer_t));
integer_t vec = load_unaligned((integer_t *)p);
start_of_current_chunk = p;
p += chunk_src_sz;
const integer_t esc_cmp = cmpeq_epi8(vec, esc_vec);
int num_of_bytes_to_first_esc = bytes_to_first_match(esc_cmp);
if (num_of_bytes_to_first_esc > -1 && (unsigned)num_of_bytes_to_first_esc < chunk_src_sz) {
// bit set for every non-ASCII byte and every ESC byte in the chunk
uint32_t ascii_mask = movemask_epi8(vec);
uint32_t esc_mask = movemask_epi8(cmpeq_epi8(vec, esc_vec));
if (chunk_src_sz < sizeof(integer_t)) {
// exclude the garbage bytes read from beyond limit
const uint32_t chunk_mask = (1u << chunk_src_sz) - 1;
ascii_mask &= chunk_mask;
esc_mask &= chunk_mask;
}
if (esc_mask) {
sentinel_found = true;
chunk_src_sz = num_of_bytes_to_first_esc;
chunk_src_sz = __builtin_ctz(esc_mask);
ascii_mask &= (1u << chunk_src_sz) - 1;
d->num_consumed += chunk_src_sz + 1; // esc is also consumed
if (!chunk_src_sz) continue;
} else d->num_consumed += chunk_src_sz;
if (chunk_src_sz < sizeof(integer_t)) vec = zero_last_n_bytes(vec, sizeof(integer_t) - chunk_src_sz);
num_of_trailing_bytes = 0;
bool check_for_trailing_bytes = !sentinel_found;
debug_register(vec);
int32_t ascii_mask;
#define abort_with_invalid_utf8() \
{ \
@ -704,13 +665,15 @@ FUNC(utf8_decode_to_esc)(UTF8Decoder *d, const uint8_t *src_data, size_t src_len
}
start_classification:
// Check if we have pure ASCII and use fast path
ascii_mask = movemask_epi8(vec);
if (!ascii_mask) { // no bytes with high bit (0x80) set, so just plain ASCII
FUNC(output_plain_ascii)(d, vec, chunk_src_sz);
prev_chunk_was_all_ascii = true;
handle_trailing_bytes();
continue;
}
prev_chunk_was_all_ascii = false;
// zero the bytes not part of the chunk as they would corrupt classification
if (chunk_src_sz < sizeof(integer_t)) vec = zero_last_n_bytes(vec, sizeof(integer_t) - chunk_src_sz);
// Classify the bytes by whether they may be the start of a 2-byte, 3-byte, or 4-byte sequence.
// This is only an initial, potential classification.
// 0xC0 and 0xC1 are initially classified as potential starter bytes of 2-byte sequences.
@ -749,10 +712,16 @@ FUNC(utf8_decode_to_esc)(UTF8Decoder *d, const uint8_t *src_data, size_t src_len
num_of_trailing_bytes = 2; // 3- and 4-byte characters with only 1 byte left
else if (chunk_src_sz > 2 && start_of_current_chunk[chunk_src_sz - 3] >= 0xf0)
num_of_trailing_bytes = 3; // 4-byte characters with only 3 bytes left
chunk_src_sz -= num_of_trailing_bytes;
d->num_consumed -= num_of_trailing_bytes;
if (!chunk_src_sz) { abort_with_invalid_utf8(); }
vec = zero_last_n_bytes(vec, sizeof(integer_t) - chunk_src_sz);
// num_of_trailing_bytes can be zero, when overlapping sequences near the end
// of the chunk make counts[last] > 1 without an actual incomplete trailing
// sequence. Reclassification then detects the overlap as invalid.
if (num_of_trailing_bytes) {
chunk_src_sz -= num_of_trailing_bytes;
d->num_consumed -= num_of_trailing_bytes;
if (!chunk_src_sz) { abort_with_invalid_utf8(); }
ascii_mask &= (1u << chunk_src_sz) - 1;
vec = zero_last_n_bytes(vec, sizeof(integer_t) - chunk_src_sz);
}
goto start_classification;
}
@ -775,7 +744,7 @@ FUNC(utf8_decode_to_esc)(UTF8Decoder *d, const uint8_t *src_data, size_t src_len
// Therefore there is a count mismatch, indicating that the chunk is ill-formed UTF-8.
// (If the following "\x01" were absent, and the "\x7f" were the last byte of the chunk,
// then the `check_for_trailing_bytes` validation above detects the error as a trailing incomplete sequence.)
const int ascii_sequence_count_mismatches = ascii_mask ^ movemask_epi8(cmpgt_epi8(counts, zero));
const uint32_t ascii_sequence_count_mismatches = ascii_mask ^ (uint32_t)movemask_epi8(cmpgt_epi8(counts, zero));
chunk_is_invalid = set1_epi8(ascii_sequence_count_mismatches ? 0xff : 0x00);
// Validate 2-byte sequence starter bytes: 0xC0..0xC1 are invalid (overlong encodings for U+0000..U+007F).
@ -791,32 +760,34 @@ FUNC(utf8_decode_to_esc)(UTF8Decoder *d, const uint8_t *src_data, size_t src_len
// Without this, "\xf0\x90\xc2\x80" would have overlapping sequences, and it would be incorrectly decoded elsewhere as an empty string.
chunk_is_invalid = or_si(chunk_is_invalid, andnot_si(cmplt_epi8(vec, set1_epi8(0xc0)), cmpgt_epi8(counts, count)));
// The second byte checks below all need the previous byte at each position.
// Compute the shifted vector once, comparing it for equality with a starter
// byte gives the same result as shifting the comparison result, since the
// zero byte shifted in at position 0 cannot equal any starter byte.
const integer_t prev_bytes = shift_right_by_one_byte(vec);
// Validate second bytes of E0-starting 3-byte sequences.
// 0xE0 must be followed by 0xA0..0xBF (not 0x80..0x9F) to avoid overlong encodings.
// Without this, "\xe0\x80\x80" would incorrectly be decoded as a "\x00".
const integer_t e0_starter_bytes = cmpeq_epi8(vec, set1_epi8(0xe0));
const integer_t e0_first_follower_bytes = shift_right_by_one_byte(e0_starter_bytes);
const integer_t e0_first_follower_bytes = cmpeq_epi8(prev_bytes, set1_epi8(0xe0));
chunk_is_invalid = or_si(chunk_is_invalid, and_si(e0_first_follower_bytes, cmplt_epi8(and_si(e0_first_follower_bytes, vec), set1_epi8(0xa0))));
// Validate second bytes of ED-starting 3-byte sequences.
// 0xED must be followed by 0x80..0x9F (not 0xA0..0xBF) to avoid UTF-16 surrogates.
// Without this, "\xed\xa0\x80" would incorrectly be decoded as an isolated surrogate "\uD800".
const integer_t ed_starter_bytes = cmpeq_epi8(vec, set1_epi8(0xed));
const integer_t ed_first_follower_bytes = shift_right_by_one_byte(ed_starter_bytes);
const integer_t ed_first_follower_bytes = cmpeq_epi8(prev_bytes, set1_epi8(0xed));
chunk_is_invalid = or_si(chunk_is_invalid, and_si(ed_first_follower_bytes, cmpgt_epi8(and_si(ed_first_follower_bytes, vec), set1_epi8(0x9f))));
// Validate second bytes of F0-starting 4-byte sequences.
// F0 must be followed by 0x90..0xBF (not 0x80..0x8F) to avoid overlong encodings.
// Without this, "\xf0\x80\x80\x80" would incorrectly be decoded as a "\x0000".
const integer_t f0_starter_bytes = cmpeq_epi8(vec, set1_epi8(0xf0));
const integer_t f0_first_follower_bytes = shift_right_by_one_byte(f0_starter_bytes);
const integer_t f0_first_follower_bytes = cmpeq_epi8(prev_bytes, set1_epi8(0xf0));
chunk_is_invalid = or_si(chunk_is_invalid, and_si(f0_first_follower_bytes, cmplt_epi8(and_si(f0_first_follower_bytes, vec), set1_epi8(0x90))));
// Validate second bytes of F4-starting 4-byte sequences.
// F4 must be followed by 0x80..0x8F (not 0x90..0xBF) to stay within the Unicode codespace.
// Without this, "\xf4\x90\x80\x80" would incorrectly be decoded as an ill-formed "\U00110000".
const integer_t f4_starter_bytes = cmpeq_epi8(vec, set1_epi8(0xf4));
const integer_t f4_first_follower_bytes = shift_right_by_one_byte(f4_starter_bytes);
const integer_t f4_first_follower_bytes = cmpeq_epi8(prev_bytes, set1_epi8(0xf4));
chunk_is_invalid = or_si(chunk_is_invalid, and_si(f4_first_follower_bytes, cmpgt_epi8(and_si(f4_first_follower_bytes, vec), set1_epi8(0x8f))));
// Check for any accumulated validation errors and, if found,
@ -863,35 +834,39 @@ FUNC(utf8_decode_to_esc)(UTF8Decoder *d, const uint8_t *src_data, size_t src_len
output3 = shift_right_by_two_bytes(output3);
debug_register(output3);
// Shuffle bytes to remove continuation bytes
integer_t shifts = count_subs1; // number of bytes we need to skip for each UTF-8 sequence
// propagate the shifts to all subsequent bytes by shift and add
shifts = add_epi8(shifts, shift_right_by_one_byte(shifts));
shifts = add_epi8(shifts, shift_right_by_two_bytes(shifts));
shifts = add_epi8(shifts, shift_right_by_four_bytes(shifts));
shifts = add_epi8(shifts, shift_right_by_eight_bytes(shifts));
#if KITTY_SIMD_LEVEL == 256
shifts = add_epi8(shifts, shift_right_by_sixteen_bytes(shifts));
#endif
// Shuffle bytes to remove continuation bytes. The compaction is done
// independently within each 128-bit lane, which needs only cheap in-lane
// shifts and an in-lane shuffle. The lanes are then joined by output_unicode
// using the number of codepoints in each lane. Note that sequences spanning
// a lane boundary are fine: their assembled codepoint lives at the position
// of the last byte of the sequence, all that matters here is how far
// leftwards within its lane that position must move.
const integer_t discarded_locations = cmpgt_epi8(counts, one); // bytes that are not ASCII and not the last byte of a sequence
// number of discarded bytes at or before each position, within its lane,
// computed via prefix sum of ones at discarded locations
integer_t shifts = and_si(one, discarded_locations);
shifts = add_epi8(shifts, shift_right_in_lane_by_one_byte(shifts));
shifts = add_epi8(shifts, shift_right_in_lane_by_two_bytes(shifts));
shifts = add_epi8(shifts, shift_right_in_lane_by_four_bytes(shifts));
shifts = add_epi8(shifts, shift_right_in_lane_by_eight_bytes(shifts));
// zero the shifts for discarded continuation bytes
shifts = and_si(shifts, cmplt_epi8(counts, two));
shifts = andnot_si(discarded_locations, shifts);
// now we need to convert shifts into a mask for the shuffle. The mask has each byte of the
// form 0000xxxx the lower four bits indicating the destination location for the byte. For 256 bit shuffle we use lower 5 bits.
// First we move the numbers in shifts to discard the unwanted UTF-8 sequence bytes. We note that the numbers
// are bounded by sizeof(integer_t) and so we need at most 4 (for 128 bit) or 5 (for 256 bit) moves. The numbers are
// monotonic from left to right and change value only at the end of a UTF-8 sequence. We move them leftwards, accumulating the
// form 0000xxxx the lower four bits indicating the source location within the lane for the byte.
// First we move the numbers in shifts to discard the unwanted UTF-8 sequence bytes.
// Every full lane keeps at least four bytes so the numbers are bounded by 12 and we
// need at most 4 moves. The numbers are monotonic from left to right within a lane and
// change value only at the end of a UTF-8 sequence. We move them leftwards, accumulating the
// moves bit-by-bit.
#define move(shifts, amt, which_bit) blendv_epi8(shifts, shift_left_by_##amt(shifts), shift_left_by_##amt(shift_left_by_bits16(shifts, 8 - which_bit)))
#define move(shifts, amt, which_bit) \
blendv_epi8(shifts, shift_left_in_lane_by_##amt(shifts), shift_left_in_lane_by_##amt(shift_left_by_bits16(shifts, 8 - which_bit)))
shifts = move(shifts, one_byte, 1);
shifts = move(shifts, two_bytes, 2);
shifts = move(shifts, four_bytes, 3);
shifts = move(shifts, eight_bytes, 4);
#if KITTY_SIMD_LEVEL == 256
shifts = move(shifts, sixteen_bytes, 5);
#endif
#undef move
// convert the shifts into a suitable mask for shuffle by adding the byte number to each byte
shifts = add_epi8(shifts, numbered);
// convert the shifts into a suitable mask for the in-lane shuffle by adding the byte number within the lane to each byte
shifts = add_epi8(shifts, numbered_in_lane);
debug_register(shifts);
output1 = shuffle_epi8(output1, shifts);
@ -901,12 +876,26 @@ FUNC(utf8_decode_to_esc)(UTF8Decoder *d, const uint8_t *src_data, size_t src_len
debug_register(output2);
debug_register(output3);
const unsigned num_of_discarded_bytes = sum_bytes(count_subs1);
const unsigned num_codepoints = chunk_src_sz - num_of_discarded_bytes;
debug("num_of_discarded_bytes: %u num_codepoints: %u\n", num_of_discarded_bytes, num_codepoints);
uint32_t kept_mask = ~(uint32_t)movemask_epi8(discarded_locations);
if (chunk_src_sz < 32) kept_mask &= (1u << chunk_src_sz) - 1; // also clears the always zero high bits for the 128-bit case
#if KITTY_SIMD_LEVEL == 128
const unsigned num_codepoints = __builtin_popcount(kept_mask);
debug("num_codepoints: %u\n", num_codepoints);
FUNC(output_unicode)(d, output1, output2, output3, num_codepoints);
#else
const unsigned num_codepoints_lane0 = __builtin_popcount(kept_mask & 0xffff), num_codepoints_lane1 = __builtin_popcount(kept_mask >> 16);
debug("num_codepoints: %u\n", num_codepoints_lane0 + num_codepoints_lane1);
FUNC(output_unicode)(d, output1, output2, output3, num_codepoints_lane0, num_codepoints_lane1);
#endif
handle_trailing_bytes();
}
if (sentinel_found && d->state.cur != UTF8_ACCEPT) {
// an incomplete UTF-8 sequence was cut off by the sentinel, matching
// the scalar implementation, emit a replacement char for it
utf8_decoder_ensure_capacity(d, 1);
d->output.storage[d->output.pos++] = 0xfffd;
zero_at_ptr(&d->state);
}
#ifdef compare_with_scalar
if (debugdec.output.pos != d->output.pos || debugdec.num_consumed != d->num_consumed ||
memcmp(d->output.storage, debugdec.output.storage, d->output.pos * sizeof(d->output.storage[0])) != 0 || sentinel_found != scalar_sentinel_found ||
@ -969,6 +958,15 @@ FUNC(utf8_decode_to_esc)(UTF8Decoder *d, const uint8_t *src_data, size_t src_len
#undef shift_left_by_four_bytes
#undef shift_left_by_eight_bytes
#undef shift_left_by_sixteen_bytes
#undef shift_right_in_lane_by_one_byte
#undef shift_right_in_lane_by_two_bytes
#undef shift_right_in_lane_by_four_bytes
#undef shift_right_in_lane_by_eight_bytes
#undef shift_left_in_lane_by_one_byte
#undef shift_left_in_lane_by_two_bytes
#undef shift_left_in_lane_by_four_bytes
#undef shift_left_in_lane_by_eight_bytes
#undef numbered_bytes_in_lane
#undef shift_left_by_bits16
#undef shift_right_by_bits32
#undef shift_right_by_bytes128
@ -983,7 +981,6 @@ FUNC(utf8_decode_to_esc)(UTF8Decoder *d, const uint8_t *src_data, size_t src_len
#undef shuffle_epi8
#undef numbered_bytes
#undef reverse_numbered_bytes
#undef sum_bytes
#undef is_zero
#undef zero_upper
#undef print_register_as_bytes

View file

@ -8,7 +8,7 @@
#include "data-types.h"
#include "charsets.h"
#include "simd-string.h"
static bool has_sse4_2 = false, has_avx2 = false;
static bool has_sse4_2 = false, has_avx2 = false, has_avx512 = false;
// xor_data64 {{{
static void
@ -95,6 +95,7 @@ test_utf8_decode_to_sentinel(PyObject *self UNUSED, PyObject *args) {
case 1: func = utf8_decode_to_esc_scalar; break;
case 2: func = utf8_decode_to_esc_128; break;
case 3: func = utf8_decode_to_esc_256; break;
case 4: func = utf8_decode_to_esc_512; break;
}
RAII_PyObject(ans, PyUnicode_FromString(""));
ssize_t p = 0;
@ -191,10 +192,12 @@ init_simd(void *x) {
Py_INCREF(Py_##val); \
if (0 != PyModule_AddObject(module, #x, Py_##val)) return false; \
}
#define do_check() \
{ \
has_sse4_2 = __builtin_cpu_supports("sse4.2") != 0; \
has_avx2 = __builtin_cpu_supports("avx2") != 0; \
#define do_check() \
{ \
has_sse4_2 = __builtin_cpu_supports("sse4.2") != 0; \
has_avx2 = __builtin_cpu_supports("avx2") != 0; \
has_avx512 = __builtin_cpu_supports("avx512f") && __builtin_cpu_supports("avx512bw") && __builtin_cpu_supports("avx512vl") && \
__builtin_cpu_supports("avx512vbmi2"); \
}
#ifdef __APPLE__
@ -227,13 +230,20 @@ init_simd(void *x) {
if (simd_env) {
has_sse4_2 = strcmp(simd_env, "128") == 0;
has_avx2 = strcmp(simd_env, "256") == 0;
has_avx512 = strcmp(simd_env, "512") == 0;
}
#undef do_check
if (has_avx512) {
A(has_avx512, True);
utf8_decode_to_esc_impl = utf8_decode_to_esc_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;
utf8_decode_to_esc_impl = utf8_decode_to_esc_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;
} else {
A(has_avx2, False);

View file

@ -62,6 +62,7 @@ void xor_data64(const uint8_t key[64], uint8_t *data, const size_t data_sz);
// SIMD implementations, internal use
bool utf8_decode_to_esc_128(UTF8Decoder *d, const uint8_t *src, size_t src_sz);
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);
void xor_data64_128(const uint8_t key[64], uint8_t *data, const size_t data_sz);

View file

@ -235,9 +235,9 @@ def env_for_python_tests(report_env: bool = False) -> Iterator[None]:
if report_env:
print('Running under CI:', is_ci)
print('Using PATH in test environment:', path)
from kitty.fast_data_types import has_avx2, has_sse4_2
from kitty.fast_data_types import has_avx2, has_avx512, has_sse4_2
print(f'Intrinsics: {has_avx2=} {has_sse4_2=}')
print(f'Intrinsics: {has_avx512=} {has_avx2=} {has_sse4_2=}')
with (
TemporaryDirectory() as tdir,
env_vars(

View file

@ -1,6 +1,7 @@
#!/usr/bin/env python
# License: GPL v3 Copyright: 2016, Kovid Goyal <kovid at kovidgoyal.net>
import random
from binascii import hexlify
from functools import partial
@ -10,6 +11,7 @@ from kitty.fast_data_types import (
base64_decode,
base64_encode,
has_avx2,
has_avx512,
has_sse4_2,
test_find_either_of_two_bytes,
test_utf8_decode_to_sentinel,
@ -185,7 +187,7 @@ class TestParser(BaseTest):
def test_utf8_simd_decode(self):
def unsupported(which):
return (which == 2 and not has_sse4_2) or (which == 3 and not has_avx2)
return (which == 2 and not has_sse4_2) or (which == 3 and not has_avx2) or (which == 4 and not has_avx512)
def reset_state():
test_utf8_decode_to_sentinel(b'', -1)
@ -219,9 +221,10 @@ class TestParser(BaseTest):
return actual
def double_test(x):
for which in (2, 3):
for which in (2, 3, 4):
t(x, which=which)
t(x * 2, which=3)
t(x * 4, which=4)
reset_state()
# incomplete trailer at end of vector
@ -235,7 +238,7 @@ class TestParser(BaseTest):
x('abc\x1bd1234efgh5678')
x('abcd1234efgh5678ijklABCDmnopEFGH')
for which in (2, 3):
for which in (2, 3, 4):
x = partial(t, which=which)
x('abcdef', 'ghijk')
x('2:α3', ':≤4:😸|')
@ -257,7 +260,7 @@ class TestParser(BaseTest):
expected = 'filler' + expected
self.ae(expected, actual, f'Failed for: {src!r} with {which=}')
for which in (1, 2, 3):
for which in (1, 2, 3, 4):
pb = partial(test_expected, which=which)
pb('ニチ', 'ニチ')
pb('\x84\x85', '\x84\x85')
@ -634,6 +637,105 @@ class TestParser(BaseTest):
# Boundary case: too large codepoint (> U+10FFFF)
pb(b'"\xf5\x80\x80\x80"', '"\ufffd\ufffd\ufffd\ufffd"')
# The tests below compare SIMD implementations against the scalar one via t().
# SIMD chunks are 16 (SSE), 32 (AVX2) or 64 (AVX-512) bytes and implementations
# may compact bytes within 16-byte lanes, so slide multi-byte sequences across
# every lane and chunk boundary.
simd_impls = 2, 3, 4
multibyte = '\xb5'.encode(), 'α'.encode(), '≤'.encode(), '\U0001f638'.encode()
for which in simd_impls:
for mb in multibyte:
for filler_len in range(70):
t(b'a' * filler_len + mb + b'b' * 40, which=which)
t(b'a' * filler_len + mb * 3 + b'b' * 40, which=which)
# ESC at and around lane/chunk boundaries, adjacent to multi-byte sequences
# and cutting off incomplete sequences
for which in simd_impls:
for esc_pos in (0, 1, 14, 15, 16, 17, 30, 31, 32, 33, 47, 48, 63, 64, 65):
buf = bytearray(b'x' * 80)
buf[esc_pos] = 0x1B
t(bytes(buf), which=which)
t(b'y' * esc_pos + '≤'.encode() + b'\x1b' + '\u03c6'.encode(), which=which)
t(b'y' * esc_pos + b'\xe2\x89\x1b' + b'z' * 8, which=which)
t(b'y' * esc_pos + b'\xf0\x9f\x98\x1b', which=which)
# Invalid/incomplete sequences sliding across lane/chunk boundaries, both
# followed by more data and truncated at the end of input
bad_seqs = (
b'\x80',
b'\xbf\xbf',
b'\xc0\x80',
b'\xc1\xbf',
b'\xe0\x80\x80',
b'\xe0\x9f\x80',
b'\xed\xa0\x80',
b'\xf0\x8f\x80\x80',
b'\xf4\x90\x80\x80',
b'\xf5\x81\x82\x83',
b'\xfe',
b'\xff',
b'\xc2',
b'\xe2\x89',
b'\xf0\x9f\x98',
)
for which in simd_impls:
for bad in bad_seqs:
for pos in (0, 13, 14, 15, 16, 17, 29, 30, 31, 32, 33, 62, 63, 64):
t(b'k' * pos + bad + b'm' * 9, which=which)
t(b'k' * pos + bad, which=which)
# Overlapping sequences at the end of a chunk. These make counts[last] > 1,
# triggering the incomplete-trailing-sequence check, without there being an
# actual incomplete trailing sequence.
for which in simd_impls:
for sz in (16, 32, 64, 128):
for pattern in (b'\xc0\xefh\x7f', b'\xc2\xefh\x7f', b'\xc0\xefhq', b'\xe0\xa0\xc2\xe2\x89'[:4]):
t(b'k' * (sz - 4) + pattern, which=which)
t(b'k' * (sz - 4) + pattern + b'more data', which=which)
# Dense/alternating sequence widths (maximum compaction churn) and exact
# chunk-multiple sizes
for which in simd_impls:
t(b'\x80' * 32, which=which)
t(b'\xbf' * 64, which=which)
t(('α' * 32).encode(), which=which)
t(('aα' * 24).encode(), which=which)
t(('≤' * 22).encode(), which=which)
t(('\U0001f638' * 17).encode(), which=which)
t(('a\U0001f638' * 13).encode(), which=which)
t(('α≤\U0001f638' * 11).encode(), which=which)
t(b'a' * 32, which=which)
t(b'a' * 33, which=which)
t(b'a' * 64, which=which)
t(b'a' * 65, which=which)
t(b'a' * 128, which=which)
t(b'a' * 129, which=which)
t(('α' * 16).encode() + b'\x1b[m', which=which)
t(('α' * 32).encode() + b'\x1b[m', which=which)
# Seeded random fuzzing against the scalar implementation, including
# split inputs to exercise cross-call state carry-over
rng = random.Random(0x1BADF00D)
interesting = 0x1B, 0x00, 0x7F, 0x80, 0xBF, 0xC0, 0xC1, 0xC2, 0xDF, 0xE0, 0xED, 0xEE, 0xEF, 0xF0, 0xF4, 0xF5, 0xFF
valid_chars = 'a', '\xb5', 'α', '≤', '\U0001f638', '\u07ff', '\u0800', '\ud7ff', '\ue000', '\uffff', '\U00010000', '\U0010ffff'
for which in simd_impls:
for _ in range(300):
parts = []
for _ in range(rng.randint(1, 80)):
r = rng.random()
if r < 0.5:
parts.append(rng.choice(valid_chars).encode())
elif r < 0.8:
parts.append(bytes((rng.choice(interesting),)))
else:
parts.append(bytes((rng.randint(0, 255),)))
data = b''.join(parts)
t(data, which=which)
if len(data) > 2:
cut = rng.randint(1, len(data) - 1)
t(data[:cut], data[cut:], which=which)
def test_find_either_of_two_bytes(self):
sizes = []
if has_sse4_2:

View file

@ -941,6 +941,16 @@ def get_source_specific_cflags(env: Env, src: str) -> List[str]:
ans.append('-x86-use-vzeroupper=0')
else:
ans.append('-mno-vzeroupper')
elif src == 'kitty/simd-string-512.c':
# uses native AVX-512 intrinsics, only compiled on x86-64, selected at runtime only when the CPU supports these features
if env.binary_arch.isa is ISA.AMD64:
ans.extend(('-mavx512f', '-mavx512bw', '-mavx512vl', '-mavx512vbmi2'))
# We have manual vzeroupper so prevent compiler from emitting it causing duplicates
if env.compiler_type is CompilerType.clang:
ans.append('-mllvm')
ans.append('-x86-use-vzeroupper=0')
else:
ans.append('-mno-vzeroupper')
elif src.startswith('3rdparty/base64/lib/arch/'):
if env.binary_arch.isa in (ISA.AMD64, ISA.X86):
q = src.split(os.path.sep)