Optimize 'json_parse_string' using SIMD.

Author: samyron

ext/json/ext/generator/depend

@@ -1,4 +1,4 @@
 generator.o: generator.c $(srcdir)/../fbuffer/fbuffer.h
 generator.o: generator.c $(srcdir)/../vendor/fpconv.c
 generator.o: generator.c $(srcdir)/../vendor/jeaiii-ltoa.h
-generator.o: generator.c $(srcdir)/simd.h
+generator.o: generator.c $(srcdir)/../simd/simd.h

ext/json/ext/generator/generator.c

@@ -5,7 +5,7 @@
 #include <math.h>
 #include <ctype.h>
 
-#include "simd.h"
+#include "../simd/simd.h"
 
 /* ruby api and some helpers */
 
@@ -304,28 +304,6 @@ static inline FORCE_INLINE unsigned char neon_next_match(search_state *search)
     return 1;
 }
 
-// See: https://community.arm.com/arm-community-blogs/b/servers-and-cloud-computing-blog/posts/porting-x86-vector-bitmask-optimizations-to-arm-neon
-static inline FORCE_INLINE uint64_t neon_match_mask(uint8x16_t matches)
-{
-    const uint8x8_t res = vshrn_n_u16(vreinterpretq_u16_u8(matches), 4);
-    const uint64_t mask = vget_lane_u64(vreinterpret_u64_u8(res), 0);
-    return mask & 0x8888888888888888ull;
-}
-
-static inline FORCE_INLINE uint64_t neon_rules_update(const char *ptr)
-{
-    uint8x16_t chunk = vld1q_u8((const unsigned char *)ptr);
-
-    // Trick: c < 32 || c == 34 can be factored as c ^ 2 < 33
-    // https://lemire.me/blog/2025/04/13/detect-control-characters-quotes-and-backslashes-efficiently-using-swar/
-    const uint8x16_t too_low_or_dbl_quote = vcltq_u8(veorq_u8(chunk, vdupq_n_u8(2)), vdupq_n_u8(33));
-
-    uint8x16_t has_backslash = vceqq_u8(chunk, vdupq_n_u8('\\'));
-    uint8x16_t needs_escape  = vorrq_u8(too_low_or_dbl_quote, has_backslash);
-
-    return neon_match_mask(needs_escape);
-}
-
 static inline unsigned char search_escape_basic_neon(search_state *search)
 {
     if (RB_UNLIKELY(search->has_matches)) {
@@ -380,14 +358,8 @@ static inline unsigned char search_escape_basic_neon(search_state *search)
     * no bytes need to be escaped and we can continue to the next chunk. If the mask is not 0 then we
     * have at least one byte that needs to be escaped.
     */
-    while (search->ptr + sizeof(uint8x16_t) <= search->end) {
-        uint64_t mask = neon_rules_update(search->ptr);
 
-        if (!mask) {
-            search->ptr += sizeof(uint8x16_t);
-            continue;
-        }
-        search->matches_mask = mask;
+    if (string_scan_simd_neon(&search->ptr, search->end, &search->matches_mask)) {
         search->has_matches = true;
         search->chunk_base = search->ptr;
         search->chunk_end = search->ptr + sizeof(uint8x16_t);
@@ -399,7 +371,7 @@ static inline unsigned char search_escape_basic_neon(search_state *search)
     if (remaining >= SIMD_MINIMUM_THRESHOLD) {
         char *s = copy_remaining_bytes(search, sizeof(uint8x16_t), remaining);
 
-        uint64_t mask = neon_rules_update(s);
+        uint64_t mask = compute_chunk_mask_neon(s);
 
         if (!mask) {
             // Nothing to escape, ensure search_flush doesn't do anything by setting
@@ -428,11 +400,6 @@ static inline unsigned char search_escape_basic_neon(search_state *search)
 
 #ifdef HAVE_SIMD_SSE2
 
-#define _mm_cmpge_epu8(a, b) _mm_cmpeq_epi8(_mm_max_epu8(a, b), a)
-#define _mm_cmple_epu8(a, b) _mm_cmpge_epu8(b, a)
-#define _mm_cmpgt_epu8(a, b) _mm_xor_si128(_mm_cmple_epu8(a, b), _mm_set1_epi8(-1))
-#define _mm_cmplt_epu8(a, b) _mm_cmpgt_epu8(b, a)
-
 static inline FORCE_INLINE unsigned char sse2_next_match(search_state *search)
 {
     int mask = search->matches_mask;
@@ -457,18 +424,6 @@ static inline FORCE_INLINE unsigned char sse2_next_match(search_state *search)
 #define TARGET_SSE2
 #endif
 
-static inline TARGET_SSE2 FORCE_INLINE int sse2_update(const char *ptr)
-{
-    __m128i chunk         = _mm_loadu_si128((__m128i const*)ptr);
-
-    // Trick: c < 32 || c == 34 can be factored as c ^ 2 < 33
-    // https://lemire.me/blog/2025/04/13/detect-control-characters-quotes-and-backslashes-efficiently-using-swar/
-    __m128i too_low_or_dbl_quote = _mm_cmplt_epu8(_mm_xor_si128(chunk, _mm_set1_epi8(2)), _mm_set1_epi8(33));
-    __m128i has_backslash = _mm_cmpeq_epi8(chunk, _mm_set1_epi8('\\'));
-    __m128i needs_escape  = _mm_or_si128(too_low_or_dbl_quote, has_backslash);
-    return _mm_movemask_epi8(needs_escape);
-}
-
 static inline TARGET_SSE2 FORCE_INLINE unsigned char search_escape_basic_sse2(search_state *search)
 {
     if (RB_UNLIKELY(search->has_matches)) {
@@ -487,17 +442,10 @@ static inline TARGET_SSE2 FORCE_INLINE unsigned char search_escape_basic_sse2(se
         }
     }
 
-    while (search->ptr + sizeof(__m128i) <= search->end) {
-        int needs_escape_mask = sse2_update(search->ptr);
-
-        if (needs_escape_mask == 0) {
-            search->ptr += sizeof(__m128i);
-            continue;
-        }
-
+    if (string_scan_simd_sse2(&search->ptr, search->end, &search->matches_mask)) {
         search->has_matches = true;
-        search->matches_mask = needs_escape_mask;
         search->chunk_base = search->ptr;
+        search->chunk_end = search->ptr + sizeof(__m128i);
         return sse2_next_match(search);
     }
 
@@ -506,7 +454,7 @@ static inline TARGET_SSE2 FORCE_INLINE unsigned char search_escape_basic_sse2(se
     if (remaining >= SIMD_MINIMUM_THRESHOLD) {
         char *s = copy_remaining_bytes(search, sizeof(__m128i), remaining);
 
-        int needs_escape_mask = sse2_update(s);
+        int needs_escape_mask = compute_chunk_mask_sse2(s);
 
         if (needs_escape_mask == 0) {
             // Nothing to escape, ensure search_flush doesn't do anything by setting

ext/json/ext/generator/simd.h

@@ -1 +1,2 @@
 parser.o: parser.c $(srcdir)/../fbuffer/fbuffer.h
+parser.o: parser.c $(srcdir)/../simd/simd.h

ext/json/ext/parser/depend

@@ -8,4 +8,32 @@
 
 append_cflags("-std=c99")
 
+if enable_config('parser-use-simd', default=!ENV["JSON_DISABLE_SIMD"])
+    if RbConfig::CONFIG['host_cpu'] =~ /^(arm.*|aarch64.*)/
+      # Try to compile a small program using NEON instructions
+      if have_header('arm_neon.h')
+        have_type('uint8x16_t', headers=['arm_neon.h']) && try_compile(<<~'SRC')
+          #include <arm_neon.h>
+          int main() {
+              uint8x16_t test = vdupq_n_u8(32);
+              return 0;
+          }
+        SRC
+          $defs.push("-DJSON_ENABLE_SIMD")
+      end
+    end
+
+    if have_header('x86intrin.h') && have_type('__m128i', headers=['x86intrin.h']) && try_compile(<<~'SRC')
+      #include <x86intrin.h>
+      int main() {
+          __m128i test = _mm_set1_epi8(32);
+          return 0;
+      }
+      SRC
+        $defs.push("-DJSON_ENABLE_SIMD")
+    end
+
+    have_header('cpuid.h')
+  end
+
 create_makefile 'json/ext/parser'

ext/json/ext/parser/extconf.rb

@@ -20,6 +20,8 @@ typedef unsigned char _Bool;
 #endif
 #endif
 
+#include "../simd/simd.h"
+
 #ifndef RB_UNLIKELY
 #define RB_UNLIKELY(expr) expr
 #endif
@@ -879,7 +881,7 @@ static inline VALUE json_push_value(JSON_ParserState *state, JSON_ParserConfig *
     return value;
 }
 
-static const bool string_scan[256] = {
+static const bool string_scan_table[256] = {
     // ASCII Control Characters
      1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
      1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
@@ -892,32 +894,71 @@ static const bool string_scan[256] = {
      0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
 };
 
+#if (defined(__GNUC__ ) || defined(__clang__))
+#define FORCE_INLINE __attribute__((always_inline))
+#else
+#define FORCE_INLINE
+#endif
+
+#ifdef HAVE_SIMD
+static SIMD_Implementation simd_impl = SIMD_NONE;
+#endif /* HAVE_SIMD */
+
+static inline bool FORCE_INLINE string_scan(JSON_ParserState *state)
+{
+#ifdef HAVE_SIMD
+#if defined(HAVE_SIMD_NEON)
+    
+    uint64_t mask = 0;
+    if (string_scan_simd_neon(&state->cursor, state->end, &mask)) {
+        state->cursor += trailing_zeros64(mask) >> 2;
+        return 1;
+    }
+
+#elif defined(HAVE_SIMD_SSE2)
+    if (simd_impl == SIMD_SSE2) {
+        int mask = 0;
+        if (string_scan_simd_sse2(&state->cursor, state->end, &mask)) {
+            state->cursor += trailing_zeros(mask);
+            return 1;
+        }
+    }
+#endif /* HAVE_SIMD_NEON or HAVE_SIMD_SSE2 */
+#endif /* HAVE_SIMD */
+
+    while (state->cursor < state->end) {
+        if (RB_UNLIKELY(string_scan_table[(unsigned char)*state->cursor])) {
+            return 1;
+        }
+        *state->cursor++;
+    }
+    return 0;
+}
+
 static inline VALUE json_parse_string(JSON_ParserState *state, JSON_ParserConfig *config, bool is_name)
 {
     state->cursor++;
     const char *start = state->cursor;
     bool escaped = false;
 
-    while (state->cursor < state->end) {
-        if (RB_UNLIKELY(string_scan[(unsigned char)*state->cursor])) {
-            switch (*state->cursor) {
-                case '"': {
-                    VALUE string = json_decode_string(state, config, start, state->cursor, escaped, is_name);
-                    state->cursor++;
-                    return json_push_value(state, config, string);
-                }
-                case '\\': {
-                    state->cursor++;
-                    escaped = true;
-                    if ((unsigned char)*state->cursor < 0x20) {
-                        raise_parse_error("invalid ASCII control character in string: %s", state);
-                    }
-                    break;
-                }
-                default:
+    while (RB_UNLIKELY(string_scan(state))) {
+        switch (*state->cursor) {
+            case '"': {
+                VALUE string = json_decode_string(state, config, start, state->cursor, escaped, is_name);
+                state->cursor++;
+                return json_push_value(state, config, string);
+            }
+            case '\\': {
+                state->cursor++;
+                escaped = true;
+                if ((unsigned char)*state->cursor < 0x20) {
                     raise_parse_error("invalid ASCII control character in string: %s", state);
-                    break;
+                }
+                break;
             }
+            default:
+                raise_parse_error("invalid ASCII control character in string: %s", state);
+                break;
         }
 
         state->cursor++;
@@ -1459,4 +1500,8 @@ void Init_parser(void)
     binary_encindex = rb_ascii8bit_encindex();
     utf8_encindex = rb_utf8_encindex();
     enc_utf8 = rb_utf8_encoding();
+
+#ifdef HAVE_SIMD
+    simd_impl = find_simd_implementation();
+#endif
 }