ayo/kitty
1
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Only use SIMD if CPU supports it at runtime

Browse source
This commit is contained in:
Kovid Goyal 2023-11-10 18:39:04 +05:30
parent 19a41b4d9a
commit b032313c45
No known key found for this signature in database
GPG key ID: 06BC317B515ACE7C
4 changed files with 47 additions and 20 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

58
kitty/simd-string.c
View file

@ -9,6 +9,7 @@
#include "simd-string.h"
#include <immintrin.h>
// ByteLoader {{{
uint8_t
byte_loader_peek(const ByteLoader *self) {
#if __BYTE_ORDER__ == __ORDER_LITTLE_ENDIAN__
@ -56,13 +57,15 @@ byte_loader_skip(ByteLoader *self) {
self->num_left = 0;
}
}
// }}}
// find_either_of_two_bytes {{{
#define haszero(v) (((v) - 0x0101010101010101ULL) & ~(v) & 0x8080808080808080ULL)
#define prepare_for_hasvalue(n) (~0ULL/255 * (n))
#define hasvalue(x,n) (haszero((x) ^ (n)))
static const uint8_t*
find_either_of_two_bytes_simple(const uint8_t *haystack, const size_t sz, const uint8_t x, const uint8_t y) {
static uint8_t*
find_either_of_two_bytes_simple(uint8_t *haystack, const size_t sz, const uint8_t x, const uint8_t y) {
ByteLoader it; byte_loader_init(&it, (uint8_t*)haystack, sz);
// first align by testing the first few bytes one at a time
@ -74,7 +77,7 @@ find_either_of_two_bytes_simple(const uint8_t *haystack, const size_t sz, const
const BYTE_LOADER_T a = prepare_for_hasvalue(x), b = prepare_for_hasvalue(y);
while (it.num_left) {
if (hasvalue(it.m, a) || hasvalue(it.m, b)) {
const uint8_t *ans = haystack + sz - it.num_left, q = hasvalue(it.m, a) ? x : y;
uint8_t *ans = haystack + sz - it.num_left, q = hasvalue(it.m, a) ? x : y;
while (it.num_left) {
if (byte_loader_next(&it) == q) return ans;
ans++;
@ -87,11 +90,11 @@ find_either_of_two_bytes_simple(const uint8_t *haystack, const size_t sz, const
}
#undef SHIFT_OP
static const uint8_t*
find_either_of_two_bytes_simd_impl(const uint8_t *haystack, const uint8_t* needle_, size_t sz) {
static uint8_t*
find_either_of_two_bytes_sse4_2_impl(uint8_t *haystack, const uint8_t* needle_, size_t sz) {
size_t extra = (uintptr_t)haystack % sizeof(__m128i);
if (extra) { // need aligned loads for performance so search first few bytes by hand
const uint8_t *ans = find_either_of_two_bytes_simple(haystack, MIN(sz, extra), needle_[0], needle_[1]);
uint8_t *ans = find_either_of_two_bytes_simple(haystack, MIN(sz, extra), needle_[0], needle_[1]);
if (ans) return ans;
extra = MIN(extra, sz);
sz -= extra;
@ -109,23 +112,27 @@ find_either_of_two_bytes_simd_impl(const uint8_t *haystack, const uint8_t* needl
return NULL;
}
static const uint8_t*
find_either_of_two_bytes_simd(uint8_t *haystack, const size_t sz, const uint8_t x, const uint8_t y) {
static uint8_t*
find_either_of_two_bytes_sse4_2(uint8_t *haystack, const size_t sz, const uint8_t x, const uint8_t y) {
uint8_t before = haystack[sz];
haystack[sz] = 0;
uint8_t needle[16] = {x, y, 0,0,0,0,0,0,0,0,0,0,0,0,0,0};
const uint8_t *ans = find_either_of_two_bytes_simd_impl(haystack, needle, sz);
uint8_t *ans = find_either_of_two_bytes_sse4_2_impl(haystack, needle, sz);
haystack[sz] = before;
return ans;
}
static uint8_t* (*find_either_of_two_bytes_impl)(uint8_t*, const size_t, const uint8_t, const uint8_t) = find_either_of_two_bytes_simple;
uint8_t*
find_either_of_two_bytes(uint8_t *haystack, const size_t sz, const uint8_t a, const uint8_t b) {
return (uint8_t*)find_either_of_two_bytes_simd(haystack, sz, a, b);
return (uint8_t*)find_either_of_two_bytes_impl(haystack, sz, a, b);
}
// }}}
static const uint8_t*
find_byte_not_in_range_simple(const uint8_t *haystack, const size_t sz, const uint8_t a, const uint8_t b) {
// find_byte_not_in_range {{{
static uint8_t*
find_byte_not_in_range_simple(uint8_t *haystack, const size_t sz, const uint8_t a, const uint8_t b) {
ByteLoader it; byte_loader_init(&it, haystack, sz);
while (it.num_left) {
const uint8_t ch = byte_loader_next(&it);
@ -134,12 +141,12 @@ find_byte_not_in_range_simple(const uint8_t *haystack, const size_t sz, const ui
return NULL;
}
static const uint8_t*
find_byte_not_in_range_simd_impl(const uint8_t *haystack, const uint8_t* needle_, size_t sz) {
static uint8_t*
find_byte_not_in_range_sse4_2_impl(uint8_t *haystack, const uint8_t* needle_, size_t sz) {
size_t extra = (uintptr_t)haystack % sizeof(__m128i);
if (extra) { // need aligned loads for performance so search first few bytes by hand
const uint8_t *ans = find_byte_not_in_range_simple((const uint8_t*)haystack, MIN(sz, extra), needle_[0], needle_[1]);
if (ans) return (const uint8_t*)ans;
uint8_t *ans = find_byte_not_in_range_simple(haystack, MIN(sz, extra), needle_[0], needle_[1]);
if (ans) return ans;
extra = MIN(extra, sz);
sz -= extra;
haystack += extra;
@ -162,14 +169,29 @@ find_byte_not_in_range_sse4_2(uint8_t *haystack, const size_t sz, const uint8_t
uint8_t before = haystack[sz];
haystack[sz] = 0;
uint8_t needle[16] = {a, b, 0, 0, 0,0,0,0,0,0,0,0,0,0,0,0};
uint8_t *ans = (uint8_t*)find_byte_not_in_range_simd_impl((uint8_t*)haystack, needle, sz);
uint8_t *ans = (uint8_t*)find_byte_not_in_range_sse4_2_impl((uint8_t*)haystack, needle, sz);
haystack[sz] = before;
return ans;
}
static uint8_t* (*find_byte_not_in_range_impl)(uint8_t *haystack, const size_t sz, const uint8_t a, const uint8_t b) = find_byte_not_in_range_simple;
uint8_t*
find_byte_not_in_range(uint8_t *haystack, const size_t sz, const uint8_t a, const uint8_t b) {
return (uint8_t*)find_byte_not_in_range_sse4_2(haystack, sz, a, b);
return (uint8_t*)find_byte_not_in_range_impl(haystack, sz, a, b);
}
// }}}
bool
init_simd(void *x) {
PyObject *module = (PyObject*)x;
bool has_sse4_2 = __builtin_cpu_supports("sse4.2") != 0; bool has_avx2 = __builtin_cpu_supports("avx2");
if (0 != PyModule_AddObjectRef(module, "has_sse4_2", has_sse4_2 ? Py_True : Py_False)) return false;
if (0 != PyModule_AddObjectRef(module, "has_avx2", has_avx2 ? Py_True : Py_False)) return false;
if (has_sse4_2) {
find_byte_not_in_range_impl = find_byte_not_in_range_sse4_2;
find_either_of_two_bytes_impl = find_either_of_two_bytes_sse4_2;
}
return true;
}

3
kitty/simd-string.h
View file

@ -21,6 +21,9 @@ uint8_t byte_loader_peek(const ByteLoader *self);
void byte_loader_init(ByteLoader *self, const uint8_t *buf, unsigned int sz);
uint8_t byte_loader_next(ByteLoader *self);
// Pass a PyModule PyObject* as the argument. Must be called once at application startup
bool init_simd(void* module);
// Requires haystack[sz] to be writable and 7 bytes to the left of haystack to
// be readable. Returns pointer to first position in haystack that contains
// either of the two chars or NULL if not found.

1
kitty/vt-parser.c
View file

@ -1747,6 +1747,7 @@ void vt_parser_set_pending_activated_at(Parser*p, monotonic_t n) {
#define EXTRA_INIT \
if (0 != PyModule_AddIntConstant(module, "VT_PARSER_BUFFER_SIZE", BUF_SZ)) return 0; \
if (0 != PyModule_AddIntConstant(module, "VT_PARSER_MAX_ESCAPE_CODE_SIZE", MAX_ESCAPE_CODE_LENGTH)) return 0; \
if (!init_simd(module)) return 0; \
INIT_TYPE(Parser)

5
setup.py
View file

@ -491,8 +491,9 @@ def init_env(
cflags.append('-g3')
ldflags.append('-lprofiler')
if not native_optimizations:
cflags.append('-msse4.2')
# SIMD instructions
cflags.append('-msse4.2')
cflags.append('-mavx2')
library_paths: Dict[str, List[str]] = {}
def add_lpath(which: str, name: str, val: Optional[str]) -> None: