Introduce ARM Neon and SSE2 SIMD. #743
Conversation
|
The gain seem to be 7% on real word benchmarks: Also note that I did one more refactoring to make the introduction of SIMD easier, so you still have a conflict. |
ext/json/ext/generator/simd.h
Outdated
| uint8x16x4_t load_uint8x16_4(const unsigned char *table, int offset) { | ||
| uint8x16x4_t tab; | ||
| for(int i=0; i<4; i++) { | ||
| tab.val[i] = vld1q_u8(table+offset+(i*16)); | ||
| } | ||
| return tab; | ||
| } |
There was a problem hiding this comment.
Isn't that just vld4q_u8?
https://developer.arm.com/architectures/instruction-sets/intrinsics/vld4q_u8
There was a problem hiding this comment.
Unfortunately it's not. vld4q_u8 interleaves the data among the 4 vector registers.
% cat load-test.c
#include <stdio.h>
#include <stdint.h>
#include <arm_neon.h>
void print_vec(char *msg, uint8x16_t vec) {
printf("%s\n[ ", msg);
uint8_t store[16] = {0};
vst1q_u8(store, vec);
for(int i=0; i<16; i++) {
printf("%3d ", store[i]);
}
printf("]\n");
}
uint8x16x4_t load_table(uint8_t *table, int offset) {
uint8x16x4_t tab;
for(int i=0; i<4; i++) {
tab.val[i] = vld1q_u8(table+offset+(i*16));
}
return tab;
}
int main(void) {
uint8_t table[256];
for(int i=0; i<256; i++) {
table[i] = i;
}
uint8x16x4_t tab1 = load_table(table, 0);
print_vec("tab1.val[0]", tab1.val[0]);
print_vec("tab1.val[1]", tab1.val[1]);
print_vec("tab1.val[2]", tab1.val[2]);
print_vec("tab1.val[3]", tab1.val[3]);
printf("\n");
uint8x16x4_t tab1_2 = vld4q_u8(table);
print_vec("tab1_2.val[0]", tab1_2.val[0]);
print_vec("tab1_2.val[1]", tab1_2.val[1]);
print_vec("tab1_2.val[2]", tab1_2.val[2]);
print_vec("tab1_2.val[3]", tab1_2.val[3]);
return 0;
}
% ./load-test
tab1.val[0]
[ 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 ]
tab1.val[1]
[ 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 ]
tab1.val[2]
[ 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 ]
tab1.val[3]
[ 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 ]
tab1_2.val[0]
[ 0 4 8 12 16 20 24 28 32 36 40 44 48 52 56 60 ]
tab1_2.val[1]
[ 1 5 9 13 17 21 25 29 33 37 41 45 49 53 57 61 ]
tab1_2.val[2]
[ 2 6 10 14 18 22 26 30 34 38 42 46 50 54 58 62 ]
tab1_2.val[3]
[ 3 7 11 15 19 23 27 31 35 39 43 47 51 55 59 63 ]
There was a problem hiding this comment.
Wow, that's so weird.
Well, maybe that loop should be unrolled then, I suspect the compiler does it, but might as well be explicit.
|
Can you just include the implementation for the regular escaping? I'm not sure the script safe version is quite worth it. |
…tion. Also store the potential matches directly rather than looking up values in the escape table.
ext/json/ext/generator/generator.c
Outdated
| if ((ch_len = search_escape_basic_neon_advance_lut(search)) != 0) { | ||
| return ch_len; | ||
| } | ||
|
|
||
| // if ((ch_len = search_escape_basic_neon_advance_rules(search)) != 0) { | ||
| // return ch_len; | ||
| // } |
There was a problem hiding this comment.
Seems like it's a toss up which one is the best. It might be an artifact that my M1 Macbook Air is passively cooled and it gets warm after I run it over and over.
|
Comparison between Running it a second time: |
…e only need 128 bytes for the lookup table as the top 128 bytes are all zeros.
|
Not sure why but it's way more modest on my machine (Air M3): |
|
Apologies for going dark for a while. I've been trying to make incremental improvements on a different branch (found here). My hope was using a move mask would be faster than Feel free to try it out though. |
That's no worries at all. I want to release a After that I think I can start merging some SIMD stuff. I'd like to go with the smaller possible useful SIMD acceleration to ensure it doesn't cause issues with people. If it works well, we can then go farther. So yeah, no rush. |
|
@byroot if you have a few minutes, would you be able to checkout this branch and benchmark it against master. You'll have to tweak your compare script a bit to compile this branch with This branch uses the bit twiddling sort of platform agnostic SIMD code if the SIMD code is disabled via a The results on my M1: |
|
With that compilation flag and compared to |
|
From a co-worker with an M4 Pro: |
|
From another co-worker with an M1 Pro: |
|
I just pushed a PR #769 to this repo which also employs SIMD to speed up string escapes. I am really sorry that we both worked in that area at the same time; after I started my work I didn't check back with this repo for a while (and I should have done that.) I believe the main difference between my PR and yours seem that mine supports x86 as well. It is doing this by using a cross-platform shim I want to suggest to collaborate on getting SIMD support in one way or another. 👋 |
|
Hi @radiospiel, I'll take a look at #769. I originally started working on #730 which supports Neon, SSE 4.2 and AVX2 with runtime detection support. The PR got a bit big so I decided to close it and implement each instruction set individually. Additionally, @byroot refactored the code quite a bit to make the SIMD implementation quite a bit easier. There are two implementations in this PR, one uses a lookup table and the other is rule-based. Both seem to have similar performance on my machine. On my machine I see a 11%-48% improvement depending on the benchmark. A few of my co-workers saw various speedups depending on their machine. I should probably mark this PR as "Ready for Review". However, I'm happy to collaborate either on this or your PR. Edit: oh yeah, there is an old-school bit-twiddling SIMD approach in pure C: #738 |
|
Thank you, @samyron . I became painfully aware of the work you did when I tried to merge master into my branch, because the interface's of the escape functions had been changed; my implementation relies on a "escape me a The main difference between your approach and mine is that you switch out the search functionality, depending on the availability of SIMD, while I switch out the SIMD primitives instead. This allows me to have working implementations for X86, ARM, and bit-twiddling; but only a handful of primitives are available because NEON and AVX are different, so your approach should allow for per-hardware type optimal implementations. I have a busy week ahead of me, but I will definitively take a look end of the week. I will also benchmark on Graviton instances; most ARM server workloads are probably not on a Apple Silicon CPU after all :) Happy to benchmark this PR as well. Can you share a benchmark script that produces the most useful output for you? I would be especially interested in understanding how you get the "before" and "after" entries in the benchmark output :) Speaking of benchmarks:
This is magnitudes more than the numbers posted here. I have seen a 48% posted above (on the |
Apologies, yes, that was a typo. I'll fix it in the comment above |
|
@samyron I reran benchmarks (link). Both our PRs show a substantial improvement over the baseline, the only significant difference is on short strings.
strings.short is a test on a 13-byte string I believe such short strings are relevant, because JSON object keys are probably quite often shorter than 16 byte; my PR applies SIMD for strings of 8 byte and more (link). (The value of 8 seemed beneficial and looked nice, but I should probably retest this with smaller values.) Maybe you could be able to support that as well? |
|
@byroot we have two competing implementations of the same approach. While mine is probably more beneficial in the short term (because it also supports x86), I believe that @samyron 's approach has more future potential, because it allows handcrafted SIMD implementations that are fundamentally different between NEON and SSE2. (and it certainly can be extended to also support shorter strings, see comment above.) Also, transplanting a x86 implementation from my PR into @samyron 's shouldn't be too hard to achieve. I see the following alternatives:
What do you all think about that? ☝️ |
| #ifdef ENABLE_SIMD | ||
|
|
||
| #if defined(__ARM_NEON) || defined(__ARM_NEON__) || defined(__aarch64__) || defined(_M_ARM64) | ||
| #include <arm_neon.h> |
This comment was marked as resolved.
This comment was marked as resolved.
Sorry, something went wrong.
…htly slower than the rules based implementation.
|
I removed the lookup-table based implementation as it was slightly slower than the direct compare/rules based approach. At this point, subject to PR feedback or any CI issues identified, I'm going to call this PR finished. I'm going to continue to look for additional optimizations but I do not want to hold up progress any longer. |
|
Sounds good. I'm back so I'll find some time to review this deeper soon (I might directly push to your branch for simple stuff). The main concern will be compatibility with |
ext/json/ext/generator/simd.h
Outdated
| return tab; | ||
| } | ||
|
|
||
| void print_uint8x16(char *msg, uint8x16_t vec) { |
There was a problem hiding this comment.
I am not sure we should keep these print_XXX methods.
ext/json/ext/generator/generator.c
Outdated
| if (popcnt >= sizeof(uint8x16_t)/2) { | ||
| return sizeof(uint8x16_t); | ||
| } |
There was a problem hiding this comment.
This is a heuristic. I didn't run exhaustive tests to determine at which point it makes sense to fall back to searching over the chunk of input.
There was a problem hiding this comment.
I'm currently wondering if it wouldn't be worth to not do the popcount at all, and directly build the mask?
There was a problem hiding this comment.
I pushed af859c2. I'd like your feedback on it.
|
The one area I'm currently focused on, possibly for a follow up PR is how the code branches to the current On |
|
I found (and fixed) a bug. Please do not merge this yet. This fails because the This was introduced within the last few commits as an optimization to avoid calling I added two new tests as well as a fix. I'd like to do more performance analysis. This mostly affects the synthetic benchmarks as they are purposely escape-heavy. However, I'd like to give this a bit more analysis. I'd also like to go through the test suite a bit more to ensure all edge cases are handled. My apologies. Another edit: Of course after closing my computer for the night I realized there was a much simpler fix. https://github.com/samyron/json/pull/new/arm-neon-simd-v2-fix2. I'm not sure I like it and I need to compare the both fixes with respect to performance. Note: I do have an additional commit on this branch I haven't pushed yet. |
|
I did some quick performance analysis this morning and the https://github.com/samyron/json/pull/new/arm-neon-simd-v2-fix2 fix was significantly better than the fix in this branch. Adding a conditional to the I'd still like to look over the tests a bit more to ensure everything is covered. |
|
It's now faster on all macro-benchs on my machine: |
| * | ||
| * If the sum is zero, no bytes need to be escaped and we can skip 16 bytes. | ||
| * | ||
| * If the sum is greater than or equal to 8, then we can assume that at least half of the bytes in chunk. |
There was a problem hiding this comment.
It's unclear to me what this sentence means. Is it unfinished?
On my machine is perform as well if not sligthly better,
less code and less conditional seem attractive.
```
== Encoding activitypub.json (52595 bytes)
ruby 3.4.2 (2025-02-15 revision d2930f8e7a) +YJIT +PRISM [arm64-darwin24]
Warming up --------------------------------------
after 3.026k i/100ms
Calculating -------------------------------------
after 30.476k (± 1.6%) i/s (32.81 μs/i) - 154.326k in 5.065214s
Comparison:
before: 29732.0 i/s
after: 30476.0 i/s - 1.03x faster
== Encoding citm_catalog.json (500298 bytes)
ruby 3.4.2 (2025-02-15 revision d2930f8e7a) +YJIT +PRISM [arm64-darwin24]
Warming up --------------------------------------
after 153.000 i/100ms
Calculating -------------------------------------
after 1.547k (± 0.7%) i/s (646.38 μs/i) - 7.803k in 5.043956s
Comparison:
before: 1561.9 i/s
after: 1547.1 i/s - same-ish: difference falls within error
== Encoding twitter.json (466906 bytes)
ruby 3.4.2 (2025-02-15 revision d2930f8e7a) +YJIT +PRISM [arm64-darwin24]
Warming up --------------------------------------
after 309.000 i/100ms
Calculating -------------------------------------
after 3.108k (± 1.1%) i/s (321.74 μs/i) - 15.759k in 5.070855s
Comparison:
before: 3012.7 i/s
after: 3108.1 i/s - 1.03x faster
```
byroot
force-pushed
the
arm-neon-simd-v2
branch
from
42744f6 to
51635ad
Compare
byroot
force-pushed
the
arm-neon-simd-v2
branch
from
51635ad to
c999baf
Compare
|
Alright. I think it looks good to me. I've pushed some small simplification for NEON which I'd like your opinion on. If you think it's OK then Ineed to do the same change for SSE2, otherwise we revert back to checking Other than that I fixed a few typos and added a CI job that disable SIMD. Once we're settled on that pocount thing, I'll cleanup the git history and merge. |
Version 2 of the introduction of ARM Neon SIMD.
There are currently two implementations:
Benchmarks (Lookup table)
Benchmarks (Rules based)
I am still working on this but I wanted to share progress.
Edit: Looks like I missed one commit so I'll have to resolve some merge conflicts.