Optimize oj_dump_cstr using SSE4.2 and SSSE3. (#973)

* Refactored and simplified the ARM Neon implementation. This optimizes the worst case synthetic benchmarks.

* Move the definition of FORCE_INLINE out of the HAVE_SIMD_NEON block.

* Formatting.

* Renamed the expected real-world test data files.

* Initial optimization of oj_dump_cstr using SSE4.2 on x86-64 platforms.

* Refactor the code to reduce duplication.

* SSE42 bugfixes.

* Formatting.

* Added SIMD_MINIMUM_THRESHOLD after accidentally removing it when merging the develop branch.

* More fixes after merging develop into this branch.

* clang-format

* Refactor oj_dump_cstr to use runtime SIMD detection of SSE4.

* Wrap SSE4.2 initialization code in an ifdef.

* formatting

Author: samyron

ext/oj/dump.c

@@ -201,6 +201,45 @@ void initialize_neon(void) {
 }
 #endif
 
+#ifdef HAVE_SIMD_SSE4_2
+
+static __m128i hibit_friendly_chars_sse42[8];
+
+// From: https://stackoverflow.com/questions/36998538/fastest-way-to-horizontally-sum-sse-unsigned-byte-vector
+inline uint32_t _mm_sum_epu8(const __m128i v) {
+    __m128i vsum = _mm_sad_epu8(v, _mm_setzero_si128());
+    return _mm_cvtsi128_si32(vsum) + _mm_extract_epi16(vsum, 4);
+}
+
+inline static OJ_TARGET_SSE42 size_t hibit_friendly_size_sse42(const uint8_t *str, size_t len) {
+    size_t size = 0;
+    size_t i    = 0;
+
+    for (; i + sizeof(__m128i) <= len; i += sizeof(__m128i), str += sizeof(__m128i)) {
+        size += sizeof(__m128i);
+
+        __m128i chunk = _mm_loadu_si128((__m128i *)str);
+        __m128i tmp   = vector_lookup_sse42(chunk, hibit_friendly_chars_sse42, 8);
+        size += _mm_sum_epu8(tmp);
+    }
+    size_t total = size + calculate_string_size(str, len - i, hibit_friendly_chars);
+    return total;
+}
+
+void OJ_TARGET_SSE42 initialize_sse42(void) {
+    for (int i = 0; i < 8; i++) {
+        hibit_friendly_chars_sse42[i] = _mm_sub_epi8(
+            _mm_loadu_si128((__m128i *)(hibit_friendly_chars + i * sizeof(__m128i))),
+            _mm_set1_epi8('1'));
+    }
+}
+
+#else
+
+#define SIMD_TARGET
+
+#endif /* HAVE_SIMD_SSE4_2 */
+
 inline static size_t hibit_friendly_size(const uint8_t *str, size_t len) {
 #ifdef HAVE_SIMD_NEON
     size_t size = 0;
@@ -220,6 +259,13 @@ inline static size_t hibit_friendly_size(const uint8_t *str, size_t len) {
 
     size_t total = size + calculate_string_size(str, len - i, hibit_friendly_chars);
     return total;
+#elif defined(HAVE_SIMD_SSE4_2)
+    if (SIMD_Impl == SIMD_SSE42) {
+        if (len >= sizeof(__m128i)) {
+            return hibit_friendly_size_sse42(str, len);
+        }
+    }
+    return calculate_string_size(str, len, hibit_friendly_chars);
 #else
     return calculate_string_size(str, len, hibit_friendly_chars);
 #endif
@@ -944,6 +990,34 @@ neon_update(const char *str, uint8x16x4_t *cmap_neon, int neon_table_size, bool
     return result;
 }
 
+#elif defined(HAVE_SIMD_SSE4_2)
+typedef struct _sse42_match_result {
+    __m128i actions;
+    bool    needs_escape;
+    int     escape_mask;
+    bool    has_some_hibit;
+    bool    do_unicode_validation;
+} sse42_match_result;
+
+static inline OJ_TARGET_SSE42 sse42_match_result
+sse42_update(const char *str, __m128i *cmap_sse42, int sse42_tab_size, bool do_unicode_validation, bool has_hi) {
+    sse42_match_result result = {.has_some_hibit = false, .do_unicode_validation = false};
+
+    __m128i chunk        = _mm_loadu_si128((__m128i *)str);
+    __m128i actions      = vector_lookup_sse42(chunk, cmap_sse42, sse42_tab_size);
+    __m128i needs_escape = _mm_xor_si128(_mm_cmpeq_epi8(actions, _mm_setzero_si128()), _mm_set1_epi8(0xFF));
+    result.actions       = _mm_add_epi8(actions, _mm_set1_epi8('1'));
+
+    result.escape_mask  = _mm_movemask_epi8(needs_escape);
+    result.needs_escape = result.escape_mask != 0;
+    if (has_hi && do_unicode_validation) {
+        __m128i has_some_hibit       = _mm_and_si128(chunk, _mm_set1_epi8(0x80));
+        result.has_some_hibit        = _mm_movemask_epi8(has_some_hibit) != 0;
+        result.do_unicode_validation = has_hi && do_unicode_validation && result.has_some_hibit;
+    }
+    return result;
+}
+
 #endif /* HAVE_SIMD_NEON */
 
 static inline FORCE_INLINE const char *process_character(char         action,
@@ -1023,6 +1097,9 @@ void oj_dump_cstr(const char *str, size_t cnt, bool is_sym, bool escape1, Out ou
 #ifdef HAVE_SIMD_NEON
     uint8x16x4_t *cmap_neon       = NULL;
     int           neon_table_size = 0;
+#elif defined(HAVE_SIMD_SSE4_2)
+    __m128i *cmap_sse42 = NULL;
+    int      sse42_tab_size;
 #endif /* HAVE_SIMD_NEON */
     const char *orig                  = str;
     bool        has_hi                = false;
@@ -1091,6 +1168,9 @@ void oj_dump_cstr(const char *str, size_t cnt, bool is_sym, bool escape1, Out ou
 #ifdef HAVE_SIMD_NEON
         cmap_neon       = hibit_friendly_chars_neon;
         neon_table_size = 2;
+#elif defined(HAVE_SIMD_SSE4_2)
+        cmap_sse42     = hibit_friendly_chars_sse42;
+        sse42_tab_size = 8;
 #endif /* HAVE_NEON_SIMD */
         size = hibit_friendly_size((uint8_t *)str, cnt);
     }
@@ -1118,21 +1198,32 @@ void oj_dump_cstr(const char *str, size_t cnt, bool is_sym, bool escape1, Out ou
         if (is_sym) {
             *out->cur++ = ':';
         }
-#ifdef HAVE_SIMD_NEON
-        const char *chunk_start;
-        const char *chunk_end;
-        const char *cursor   = str;
-        bool        use_neon = (cmap_neon != NULL && cnt >= (sizeof(uint8x16_t))) ? true : false;
-        char        matches[16];
+
+#if defined(HAVE_SIMD_NEON) || defined(HAVE_SIMD_SSE4_2)
+
 #define SEARCH_FLUSH                                  \
     if (str > cursor) {                               \
         APPEND_CHARS(out->cur, cursor, str - cursor); \
         cursor = str;                                 \
     }
 
-#endif /* HAVE_SIMD_NEON */
+        const char *chunk_start;
+        const char *chunk_end;
+        const char *cursor = str;
+        char        matches[16];
+#endif /* HAVE_SIMD_NEON || HAVE_SIMD_SSE4_2 */
+
+#if defined(HAVE_SIMD_NEON)
+        bool use_simd = (cmap_neon != NULL && cnt >= (sizeof(uint8x16_t))) ? true : false;
+#elif defined(HAVE_SIMD_SSE4_2)
+        bool use_simd = false;
+        if (SIMD_Impl == SIMD_SSE42) {
+            use_simd = (cmap_sse42 != NULL && cnt >= (sizeof(__m128i))) ? true : false;
+        }
+#endif
+
 #ifdef HAVE_SIMD_NEON
-        if (use_neon) {
+        if (use_simd) {
             while (str < end) {
                 const char *chunk_ptr = NULL;
                 if (str + sizeof(uint8x16_t) <= end) {
@@ -1195,7 +1286,55 @@ void oj_dump_cstr(const char *str, size_t cnt, bool is_sym, bool escape1, Out ou
             }
             SEARCH_FLUSH;
         }
-#endif /* HAVE_SIMD_NEON */
+#endif
+
+#ifdef HAVE_SIMD_SSE4_2
+        if (SIMD_Impl == SIMD_SSE42) {
+            if (use_simd) {
+                while (str < end) {
+                    const char *chunk_ptr = NULL;
+                    if (str + sizeof(__m128i) <= end) {
+                        chunk_ptr   = str;
+                        chunk_start = str;
+                        chunk_end   = str + sizeof(__m128i);
+                    } else if ((end - str) >= SIMD_MINIMUM_THRESHOLD) {
+                        memset(out->cur, 'A', sizeof(__m128i));
+                        memcpy(out->cur, str, (end - str));
+                        chunk_ptr   = out->cur;
+                        chunk_start = str;
+                        chunk_end   = end;
+                    } else {
+                        break;
+                    }
+                    sse42_match_result result = sse42_update(chunk_ptr,
+                                                             cmap_sse42,
+                                                             sse42_tab_size,
+                                                             do_unicode_validation,
+                                                             has_hi);
+                    if ((result.do_unicode_validation) || result.needs_escape) {
+                        SEARCH_FLUSH;
+                        _mm_storeu_si128((__m128i *)matches, result.actions);
+                        while (str < chunk_end) {
+                            long match_index = str - chunk_start;
+                            str              = process_character(matches[match_index],
+                                                    str,
+                                                    end,
+                                                    out,
+                                                    orig,
+                                                    do_unicode_validation,
+                                                    &check_start);
+                            str++;
+                        }
+                        cursor = str;
+                        continue;
+                    }
+                    str = chunk_end;
+                }
+                SEARCH_FLUSH;
+            }
+        }
+#endif /* HAVE_SIMD_SSE4_2 */
+
         for (; str < end; str++) {
             str = process_character(cmap[(uint8_t)*str], str, end, out, orig, do_unicode_validation, &check_start);
         }

ext/oj/oj.c

@@ -167,6 +167,8 @@ pthread_mutex_t oj_cache_mutex;
 VALUE oj_cache_mutex = Qnil;
 #endif
 
+SIMD_Implementation SIMD_Impl = SIMD_NONE;
+
 extern void oj_parser_init();
 
 const char oj_json_class[] = "json_class";
@@ -2152,10 +2154,18 @@ void Init_oj(void) {
 #endif
     oj_init_doc();
 
+    SIMD_Impl = oj_get_simd_implementation();
+
     oj_parser_init();
     oj_scanner_init();
 
 #ifdef HAVE_SIMD_NEON
     initialize_neon();
 #endif /* HAVE_SIMD_NEON */
+
+#ifdef HAVE_SIMD_SSE4_2
+    if (SIMD_Impl == SIMD_SSE42) {
+        initialize_sse42();
+    }
+#endif /* HAVE_SIMD_SSE4_2 */
 }

ext/oj/simd.h

@@ -8,6 +8,9 @@
 // SIMD implementation enum - used for runtime selection
 typedef enum _simd_implementation { SIMD_NONE, SIMD_NEON, SIMD_SSE2, SIMD_SSE42 } SIMD_Implementation;
 
+// Define in oj.c.
+extern SIMD_Implementation SIMD_Impl;
+
 // Runtime CPU detection function (implemented in oj.c)
 SIMD_Implementation oj_get_simd_implementation(void);
 
@@ -133,4 +136,37 @@ static inline int oj_ctz_fallback(unsigned int x) {
 #define SIMD_TYPE "none"
 #endif
 
+#if defined(HAVE_SIMD_SSE4_2)
+
+#define SIMD_MINIMUM_THRESHOLD 6
+
+extern void initialize_sse42(void);
+
+static inline OJ_TARGET_SSE42 __m128i vector_lookup_sse42(__m128i input, __m128i *lookup_table, int tab_size) {
+    // Extract high 4 bits to determine which 16-byte chunk (0-15)
+    __m128i hi_index = _mm_and_si128(_mm_srli_epi32(input, 4), _mm_set1_epi8(0x0F));
+
+    // Extract low 4 bits for index within the chunk (0-15)
+    __m128i low_index = _mm_and_si128(input, _mm_set1_epi8(0x0F));
+
+    // Perform lookups in all 16 tables
+    __m128i results[16];
+    for (int i = 0; i < tab_size; i++) {
+        results[i] = _mm_shuffle_epi8(lookup_table[i], low_index);
+    }
+
+    // Create masks for each chunk and blend results
+    __m128i final_result = _mm_setzero_si128();
+
+    for (int i = 0; i < tab_size; i++) {
+        __m128i mask          = _mm_cmpeq_epi8(hi_index, _mm_set1_epi8(i));
+        __m128i masked_result = _mm_and_si128(mask, results[i]);
+        final_result          = _mm_or_si128(final_result, masked_result);
+    }
+
+    return final_result;
+}
+
+#endif
+
 #endif /* OJ_SIMD_H */