From f48e4ffd5e69656de28c9cf0203bdc9019cd4b0a Mon Sep 17 00:00:00 2001 From: Kovid Goyal Date: Sun, 11 Feb 2024 11:07:35 +0530 Subject: [PATCH] Port aligned load based find algorithm to C --- kitty/simd-string-impl.h | 84 +++++++++++++++++++++++----------------- kitty/simd-string.h | 5 +-- 2 files changed, 51 insertions(+), 38 deletions(-) diff --git a/kitty/simd-string-impl.h b/kitty/simd-string-impl.h index b7d115e31..462f6c414 100644 --- a/kitty/simd-string-impl.h +++ b/kitty/simd-string-impl.h @@ -51,7 +51,7 @@ END_IGNORE_DIAGNOSTIC #define set_epi8 simde_mm_set_epi8 #define add_epi8 simde_mm_add_epi8 #define load_unaligned simde_mm_loadu_si128 -#define load_aligned simde_mm_load_si128 +#define load_aligned(x) simde_mm_load_si128((const integer_t*)(x)) #define store_unaligned simde_mm_storeu_si128 #define cmpeq_epi8 simde_mm_cmpeq_epi8 #define cmplt_epi8 simde_mm_cmplt_epi8 @@ -98,7 +98,7 @@ FUNC(is_zero)(const integer_t a) { return simde_mm_testz_si128(a, a); } #define set_epi8 simde_mm256_set_epi8 #define add_epi8 simde_mm256_add_epi8 #define load_unaligned simde_mm256_loadu_si256 -#define load_aligned simde_mm256_load_si256 +#define load_aligned(x) simde_mm256_load_si256((const integer_t*)(x)) #define store_unaligned simde_mm256_storeu_si256 #define cmpeq_epi8 simde_mm256_cmpeq_epi8 #define cmpgt_epi8 simde_mm256_cmpgt_epi8 @@ -229,7 +229,7 @@ bytes_to_first_match(const integer_t vec) { const uint64_t m = movemask_arm128(v static inline int bytes_to_first_match_ignoring_leading_n(const integer_t vec, uintptr_t num_ignored) { uint64_t m = movemask_arm128(vec); - m >>= num_ignored >> 2; + m >>= num_ignored << 2; return m ? (__builtin_ctzll(m) >> 2) : -1; } @@ -246,15 +246,19 @@ bytes_to_first_match(const integer_t vec) { static inline int bytes_to_first_match_ignoring_leading_n(const integer_t vec, uintptr_t num_ignored) { - num_ignored >>= 2; - simde__m128i v = simde_mm256_extracti128_si256(vec, 0); - uint64_t m = movemask_arm128(vec); - m >>= num_ignored; - if (m) return __builtin_ctzll(m) >> 2; - v = simde_mm256_extracti128_si256(vec, 1); - m = movemask_arm128(vec); - m >>= num_ignored; - return m ? (16 + (__builtin_ctzll(m) >> 2)) : -1; + uint64_t m; + int offset; + if (num_ignored < 16) { + m = ((uint64_t)movemask_arm128(simde_mm256_extracti128_si256(vec, 0))) >> (num_ignored << 2); + if (m) return (__builtin_ctzll(m) >> 2); + offset = 16 - num_ignored; + num_ignored = 0; + } else { + num_ignored -= 16; + offset = 0; + } + m = ((uint64_t)movemask_arm128(simde_mm256_extracti128_si256(vec, 1))) >> (num_ignored << 2); + return m ? (offset + (__builtin_ctzll(m) >> 2)) : -1; } #endif @@ -285,37 +289,47 @@ FUNC(zero_last_n_bytes)(integer_t vec, const integer_t index, char n) { return andnot_si(mask, vec); } -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) { - zero_upper(); - const integer_t a_vec = set1_epi8(a), b_vec = set1_epi8(b); - const uint8_t* limit = haystack + sz; - integer_t chunk; - #define check_chunk() { \ - const int n = bytes_to_first_match(or_si(cmpeq_epi8(chunk, a_vec), cmpeq_epi8(chunk, b_vec))); \ if (n > -1) { \ - const uint8_t *ans = haystack + n; \ + const uint8_t *ans = haystack + n + unaligned_bytes; \ zero_upper(); \ return ans < limit ? ans : NULL; \ }} - // check the first possibly unaligned chunk - chunk = load_unaligned(haystack); - check_chunk(); - const uintptr_t unaligned_leading_count = sizeof(integer_t) - (((uintptr_t)haystack) & (sizeof(integer_t) - 1)); - haystack += unaligned_leading_count; // advance to the first aligned chunk - - // Iterate over aligned chunks, this repeats checking of - // (sizeof(integer_t) - unaligned_leading_count) bytes, but better than a branch - for (; haystack < limit; haystack += sizeof(integer_t)) { - chunk = load_aligned((integer_t*)haystack); - check_chunk(); - } - zero_upper(); - return NULL; +#define find_match(haystack, sz, get_test_vec) { \ + zero_upper(); \ + const uint8_t* limit = haystack + sz; \ + integer_t chunk; int n; \ +\ + const uintptr_t addr = (uintptr_t)haystack; \ + uintptr_t unaligned_bytes = addr & (sizeof(integer_t) - 1); \ + haystack -= unaligned_bytes; /* align haystack to the first sizeof(integer_t) boundary <= original position */ \ + chunk = load_aligned(haystack); \ + n = bytes_to_first_match_ignoring_leading_n(get_test_vec(chunk), unaligned_bytes); \ + check_chunk(); \ + haystack += sizeof(integer_t); \ + unaligned_bytes = 0; \ + /* Iterate over aligned chunks */ \ + for (; haystack < limit; haystack += sizeof(integer_t)) { \ + chunk = load_aligned(haystack); \ + n = bytes_to_first_match(get_test_vec(chunk)); \ + check_chunk(); \ + } \ + zero_upper(); \ + return NULL;\ } +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) { + if (!sz) return NULL; + const integer_t a_vec = set1_epi8(a), b_vec = set1_epi8(b); +#define get_test_from_chunk(chunk) (or_si(cmpeq_epi8(chunk, a_vec), cmpeq_epi8(chunk, b_vec))) + find_match(haystack, sz, get_test_from_chunk); +#undef get_test_from_chunk +} + +#undef check_chunk + #define output_increment sizeof(integer_t)/sizeof(uint32_t) static inline void diff --git a/kitty/simd-string.h b/kitty/simd-string.h index 606846ac4..2dc8b4b64 100644 --- a/kitty/simd-string.h +++ b/kitty/simd-string.h @@ -42,9 +42,8 @@ static inline void utf8_decoder_free(UTF8Decoder *d) { // Pass a PyModule PyObject* as the argument. Must be called once at application startup bool init_simd(void* module); -// Requires 31 bytes to the right of haystack to be readable. Returns pointer to -// first position in haystack that contains either of the two chars or NULL if -// not found. +// Returns pointer to first position in haystack that contains either of the +// two chars or NULL if not found. const uint8_t* find_either_of_two_bytes(const uint8_t *haystack, const size_t sz, const uint8_t a, const uint8_t b); // SIMD implementations, internal use