WebAssembly SIMD128 support #824

Closed
opened 2023-03-13 18:22:17 -05:00 by easyaspi314 · 4 comments
easyaspi314 commented 2023-03-13 18:22:17 -05:00 (Migrated from github.com)

WebAssembly's SIMD128 extension can use the double-width NEON approach almost literally (not XXH-ified):

#include <wasm_simd128.h>

void XXH3_accumulate_512_wasm_simd128(void *restrict acc, const void *restrict input, const void *restrict secret)
{
    v128_t *xacc = acc;
    const char *xinput = input;
    const char *xsecret = secret;
    for (int i = 0; i + 1 < 4; i += 2) {
        v128_t data_vec_1 = wasm_v128_load(xinput + 16 * i);
        v128_t data_vec_2 = wasm_v128_load(xinput + 16 * (i + 1));

        v128_t key_vec_1 = wasm_v128_load(xsecret + 16 * i);
        v128_t key_vec_2 = wasm_v128_load(xsecret + 16 * (i + 1));

        v128_t data_key_1 = wasm_v128_xor(data_vec_1, key_vec_1);
        v128_t data_key_2 = wasm_v128_xor(data_vec_2, key_vec_2);

        v128_t data_swap_1 = wasm_i64x2_shuffle(data_vec_1, data_vec_1, 1, 0);
        v128_t data_swap_2 = wasm_i64x2_shuffle(data_vec_2, data_vec_2, 1, 0);

        v128_t mixed_lo = wasm_i32x4_shuffle(data_key_1, data_key_2, 0, 2, 4, 6);
        v128_t mixed_hi = wasm_i32x4_shuffle(data_key_1, data_key_2, 1, 3, 5, 7);

        v128_t prod_1 = wasm_u64x2_extmul_low_u32x4(mixed_lo, mixed_hi);
        v128_t prod_2 = wasm_u64x2_extmul_high_u32x4(mixed_lo, mixed_hi);

        v128_t sum_1 = wasm_i64x2_add(prod_1, data_swap_1);
        v128_t sum_2 = wasm_i64x2_add(prod_2, data_swap_2);

        xacc[i]     = wasm_i64x2_add(sum_1, xacc[i]);
        xacc[i + 1] = wasm_i64x2_add(sum_2, xacc[i + 1]);
    }
}

void XXH3_scrambleAcc_wasm_simd128(void *restrict acc, const void *restrict secret)
{
    v128_t *xacc = acc;
    const char *xsecret = secret;
    const v128_t prime = wasm_u64x2_const_splat(XXH_PRIME32_1);
    for (int i = 0; i < 4; i++) {
        v128_t acc_vec = xacc[i];
        v128_t shifted = wasm_u64x2_shr(acc_vec, 47);
        v128_t data_vec = wasm_v128_xor(acc_vec, shifted);
        v128_t key_vec = wasm_v128_load(xsecret + 16 * i);
        v128_t mixed = wasm_v128_xor(data_vec, key_vec);
        xacc[i] = wasm_i64x2_mul(mixed, prime);
    }
}

Edit: I can count.

Dispatching support at runtime is impossible, as the entire binary is validated before executing, as far as I know. However, it can be dispatched as separate modules in JavaScript.

WebAssembly's SIMD128 extension can use the double-width NEON approach almost literally (not XXH-ified): ```c #include <wasm_simd128.h> void XXH3_accumulate_512_wasm_simd128(void *restrict acc, const void *restrict input, const void *restrict secret) { v128_t *xacc = acc; const char *xinput = input; const char *xsecret = secret; for (int i = 0; i + 1 < 4; i += 2) { v128_t data_vec_1 = wasm_v128_load(xinput + 16 * i); v128_t data_vec_2 = wasm_v128_load(xinput + 16 * (i + 1)); v128_t key_vec_1 = wasm_v128_load(xsecret + 16 * i); v128_t key_vec_2 = wasm_v128_load(xsecret + 16 * (i + 1)); v128_t data_key_1 = wasm_v128_xor(data_vec_1, key_vec_1); v128_t data_key_2 = wasm_v128_xor(data_vec_2, key_vec_2); v128_t data_swap_1 = wasm_i64x2_shuffle(data_vec_1, data_vec_1, 1, 0); v128_t data_swap_2 = wasm_i64x2_shuffle(data_vec_2, data_vec_2, 1, 0); v128_t mixed_lo = wasm_i32x4_shuffle(data_key_1, data_key_2, 0, 2, 4, 6); v128_t mixed_hi = wasm_i32x4_shuffle(data_key_1, data_key_2, 1, 3, 5, 7); v128_t prod_1 = wasm_u64x2_extmul_low_u32x4(mixed_lo, mixed_hi); v128_t prod_2 = wasm_u64x2_extmul_high_u32x4(mixed_lo, mixed_hi); v128_t sum_1 = wasm_i64x2_add(prod_1, data_swap_1); v128_t sum_2 = wasm_i64x2_add(prod_2, data_swap_2); xacc[i] = wasm_i64x2_add(sum_1, xacc[i]); xacc[i + 1] = wasm_i64x2_add(sum_2, xacc[i + 1]); } } void XXH3_scrambleAcc_wasm_simd128(void *restrict acc, const void *restrict secret) { v128_t *xacc = acc; const char *xsecret = secret; const v128_t prime = wasm_u64x2_const_splat(XXH_PRIME32_1); for (int i = 0; i < 4; i++) { v128_t acc_vec = xacc[i]; v128_t shifted = wasm_u64x2_shr(acc_vec, 47); v128_t data_vec = wasm_v128_xor(acc_vec, shifted); v128_t key_vec = wasm_v128_load(xsecret + 16 * i); v128_t mixed = wasm_v128_xor(data_vec, key_vec); xacc[i] = wasm_i64x2_mul(mixed, prime); } } ``` **Edit:** I can count. Dispatching support at runtime is impossible, as the entire binary is validated before executing, as far as I know. However, it can be dispatched as separate modules in JavaScript.
easyaspi314 commented 2023-03-14 09:26:21 -05:00 (Migrated from github.com)

Clang 15, Node 18.13.0 (V8 10.2.154.23-node.21, a.k.a. Chromium 102), Cortex X1 gets 13.1 GB/s vs 6.6 GB/s from an initial implementation, so this seems like it is worth it.

I haven't tried Firefox or on x86 yet.

I expect x86 to not get as much because the shuffle and extmul would not work out as cleanly in SSE or AVX — WASM was definitely designed for NEON and not SSE.

Ideally it would be this, but I don't know if the JIT is smart enough to optimize the shuffle to unpacklo and unpackhi, and it might do two shuffles and a blend for each one.

    __m128i mixed_lo = _mm_unpacklo_epi32(data_key_1, data_key_2);
    __m128i mixed_hi = _mm_unpackhi_epi32(data_key_1, data_key_2);

    __m128i prod_1_lhs = _mm_shuffle_epi32(mixed_lo, _MM_SHUFFLE(0, 1, 0, 0);
    __m128i prod_1_rhs = _mm_shuffle_epi32(mixed_hi, _MM_SHUFFLE(0, 1, 0, 0);
    __m128i prod_1 = _mm_mul_epu32(prod_1_lhs, prod_1_rhs);
    __m128i prod_2_lhs = _mm_shuffle_epi32(mixed_lo, _MM_SHUFFLE(0, 3, 0, 2);
    __m128i prod_2_rhs = _mm_shuffle_epi32(mixed_hi, _MM_SHUFFLE(0, 3, 0, 2);
    __m128i prod_2 = _mm_mul_epu32(prod_2_lhs, prod_2_rhs);

Edit again: checked v8's source code, it would emit two _mm_shuffle_epi8s and a _mm_or_si128 per shuffle. No special cases except for when the two arguments are identical.

It still would probably be faster than scalar despite the extra shuffling.

Update: 15.3 GB/s if I convert the data_swap's shuffle to a swizzle so V8 doesn't generate inline shuffle masks each time. It also should be smaller because the shuffle is an 18 byte opcode lol

Clang 15, Node 18.13.0 (V8 10.2.154.23-node.21, a.k.a. Chromium 102), Cortex X1 gets 13.1 GB/s vs 6.6 GB/s from an initial implementation, so this seems like it is worth it. I haven't tried Firefox or on x86 yet. I expect x86 to not get as much because the shuffle and extmul would not work out as cleanly in SSE or AVX — WASM was definitely designed for NEON and not SSE. Ideally it would be this, but I don't know if the JIT is smart enough to optimize the shuffle to unpacklo and unpackhi, and it might do two shuffles and a blend for each one. ```c __m128i mixed_lo = _mm_unpacklo_epi32(data_key_1, data_key_2); __m128i mixed_hi = _mm_unpackhi_epi32(data_key_1, data_key_2); __m128i prod_1_lhs = _mm_shuffle_epi32(mixed_lo, _MM_SHUFFLE(0, 1, 0, 0); __m128i prod_1_rhs = _mm_shuffle_epi32(mixed_hi, _MM_SHUFFLE(0, 1, 0, 0); __m128i prod_1 = _mm_mul_epu32(prod_1_lhs, prod_1_rhs); __m128i prod_2_lhs = _mm_shuffle_epi32(mixed_lo, _MM_SHUFFLE(0, 3, 0, 2); __m128i prod_2_rhs = _mm_shuffle_epi32(mixed_hi, _MM_SHUFFLE(0, 3, 0, 2); __m128i prod_2 = _mm_mul_epu32(prod_2_lhs, prod_2_rhs); ``` Edit again: checked v8's source code, it would emit two `_mm_shuffle_epi8`s and a `_mm_or_si128` per shuffle. No special cases except for when the two arguments are identical. It still would probably be faster than scalar despite the extra shuffling. Update: 15.3 GB/s if I convert the data_swap's shuffle to a swizzle so V8 doesn't generate inline shuffle masks each time. It also should be smaller because the shuffle is an 18 byte opcode lol
easyaspi314 commented 2023-03-14 10:44:30 -05:00 (Migrated from github.com)

Same version (with the updated swizzle), Clang 15, Node 18.14.0 (V8 10.2.154.23-node.22, a.k.a. Chromium 102), Core i7-8650u, it is also really good, 10.1 GB/s vs 5.5 GB/s.

Same version (with the updated swizzle), Clang 15, Node 18.14.0 (V8 10.2.154.23-node.22, a.k.a. Chromium 102), Core i7-8650u, it is also really good, 10.1 GB/s vs 5.5 GB/s.
easyaspi314 commented 2023-03-14 11:26:05 -05:00 (Migrated from github.com)

WIP implementation is at #825

WIP implementation is at #825
Cyan4973 commented 2023-07-15 07:58:16 -05:00 (Migrated from github.com)

merged

merged
Sign in to join this conversation.
No milestone
No assignees
1 participant
Notifications
Due date
The due date is invalid or out of range. Please use the format "yyyy-mm-dd".

No due date set.

Dependencies

No dependencies set

Reference
gh-Cyan4973/xxHash#824
No description provided.