diff --git a/docs/release.md b/docs/release.md index 1f439a9c..fc1d3b59 100644 --- a/docs/release.md +++ b/docs/release.md @@ -14,6 +14,16 @@ ## Unreleased +### Enhancements + +* Blosc byte-shuffle and unshuffle now use NEON SIMD on aarch64 (Apple Silicon, AWS + Graviton, etc). c-blosc 1.x only ships SSE2/AVX2 kernels, so on ARM it previously fell + back to a scalar loop that dominated Blosc runtime. The NEON kernels are taken from + c-blosc2 and produce byte-identical output. On an Apple M5, decoding 1 MiB lz4+shuffle + chunks is ~5-6x faster and encoding ~2x faster. Controlled by the new `neon` meson + option (default `auto`). + By {user}`Davis Bennett `. + (release_0.17.0)= * Allow building against a system Zlib. Note, this is only required if blosc is still diff --git a/meson.options b/meson.options index f055af87..63378bb0 100644 --- a/meson.options +++ b/meson.options @@ -10,3 +10,5 @@ option('sse2', type: 'feature', value: 'auto', description: 'Enable SSE2 SIMD optimizations') option('avx2', type: 'feature', value: 'auto', description: 'Enable AVX2 SIMD optimizations') +option('neon', type: 'feature', value: 'auto', + description: 'Enable NEON SIMD shuffle on aarch64') diff --git a/src/numcodecs/_blosc_neon/shuffle-generic-neon.c b/src/numcodecs/_blosc_neon/shuffle-generic-neon.c new file mode 100644 index 00000000..a7e46306 --- /dev/null +++ b/src/numcodecs/_blosc_neon/shuffle-generic-neon.c @@ -0,0 +1,372 @@ +/********************************************************************* + NEON-accelerated drop-in replacement for c-blosc's shuffle-generic.c. + + c-blosc 1.x has no ARM SIMD shuffle, so on aarch64 its dispatcher always + selects blosc_internal_(un)shuffle_generic, a scalar byte-by-byte loop. + This file defines those two symbols with NEON kernels for type sizes + 2, 4, 8 and 16, falling back to the scalar loop for other type sizes and + for the tail of each block. The byte layout is identical to the generic + routines, so the Blosc format is unchanged. + + The kernels below (shuffleN_neon / unshuffleN_neon) are copied verbatim + from c-blosc2 (blosc/shuffle-neon.c): + + Copyright (c) 2021 Lucian Marc + https://blosc.org + License: BSD 3-Clause (see c-blosc/LICENSE.txt) +**********************************************************************/ + +#include "shuffle-generic.h" + +#if !defined(__ARM_NEON) && !defined(_M_ARM64) +#error "shuffle-generic-neon.c requires a NEON-capable target" +#endif + +#include +#include + +/* Routine optimized for shuffling a buffer for a type size of 2 bytes. */ +static void +shuffle2_neon(uint8_t *const dest, const uint8_t *const src, + const size_t vectorizable_elements, const size_t total_elements) { + size_t i, k; + static const size_t bytesoftype = 2; + uint8x16x2_t r0; + + for (i = 0, k = 0; i < vectorizable_elements * bytesoftype; i += 32, k++) { + /* Load (and permute) 32 bytes to the structure r0 */ + r0 = vld2q_u8(src + i); + /* Store the results in the destination vector */ + vst1q_u8(dest + total_elements * 0 + k * 16, r0.val[0]); + vst1q_u8(dest + total_elements * 1 + k * 16, r0.val[1]); + } +} + +/* Routine optimized for shuffling a buffer for a type size of 4 bytes. */ +static void +shuffle4_neon(uint8_t *const dest, const uint8_t *const src, + const size_t vectorizable_elements, const size_t total_elements) { + size_t i, k; + static const size_t bytesoftype = 4; + uint8x16x4_t r0; + + for (i = 0, k = 0; i < vectorizable_elements * bytesoftype; i += 64, k++) { + /* Load (and permute) 64 bytes to the structure r0 */ + r0 = vld4q_u8(src + i); + /* Store the results in the destination vector */ + vst1q_u8(dest + total_elements * 0 + k * 16, r0.val[0]); + vst1q_u8(dest + total_elements * 1 + k * 16, r0.val[1]); + vst1q_u8(dest + total_elements * 2 + k * 16, r0.val[2]); + vst1q_u8(dest + total_elements * 3 + k * 16, r0.val[3]); + } +} + +/* Routine optimized for shuffling a buffer for a type size of 8 bytes. */ +static void +shuffle8_neon(uint8_t *const dest, const uint8_t *const src, + const size_t vectorizable_elements, const size_t total_elements) { + size_t i, k; + static const size_t bytesoftype = 8; + uint8x8x2_t r0[4]; + uint16x4x2_t r1[4]; + uint32x2x2_t r2[4]; + + for (i = 0, k = 0; i < vectorizable_elements * bytesoftype; i += 64, k++) { + /* Load and interleave groups of 8 bytes (64 bytes) to the structure r0 */ + r0[0] = vzip_u8(vld1_u8(src + i + 0 * 8), vld1_u8(src + i + 1 * 8)); + r0[1] = vzip_u8(vld1_u8(src + i + 2 * 8), vld1_u8(src + i + 3 * 8)); + r0[2] = vzip_u8(vld1_u8(src + i + 4 * 8), vld1_u8(src + i + 5 * 8)); + r0[3] = vzip_u8(vld1_u8(src + i + 6 * 8), vld1_u8(src + i + 7 * 8)); + /* Interleave 16 bytes */ + r1[0] = vzip_u16(vreinterpret_u16_u8(r0[0].val[0]), vreinterpret_u16_u8(r0[1].val[0])); + r1[1] = vzip_u16(vreinterpret_u16_u8(r0[0].val[1]), vreinterpret_u16_u8(r0[1].val[1])); + r1[2] = vzip_u16(vreinterpret_u16_u8(r0[2].val[0]), vreinterpret_u16_u8(r0[3].val[0])); + r1[3] = vzip_u16(vreinterpret_u16_u8(r0[2].val[1]), vreinterpret_u16_u8(r0[3].val[1])); + /* Interleave 32 bytes */ + r2[0] = vzip_u32(vreinterpret_u32_u16(r1[0].val[0]), vreinterpret_u32_u16(r1[2].val[0])); + r2[1] = vzip_u32(vreinterpret_u32_u16(r1[0].val[1]), vreinterpret_u32_u16(r1[2].val[1])); + r2[2] = vzip_u32(vreinterpret_u32_u16(r1[1].val[0]), vreinterpret_u32_u16(r1[3].val[0])); + r2[3] = vzip_u32(vreinterpret_u32_u16(r1[1].val[1]), vreinterpret_u32_u16(r1[3].val[1])); + /* Store the results in the destination vector */ + vst1_u8(dest + k * 8 + 0 * total_elements, vreinterpret_u8_u32(r2[0].val[0])); + vst1_u8(dest + k * 8 + 1 * total_elements, vreinterpret_u8_u32(r2[0].val[1])); + vst1_u8(dest + k * 8 + 2 * total_elements, vreinterpret_u8_u32(r2[1].val[0])); + vst1_u8(dest + k * 8 + 3 * total_elements, vreinterpret_u8_u32(r2[1].val[1])); + vst1_u8(dest + k * 8 + 4 * total_elements, vreinterpret_u8_u32(r2[2].val[0])); + vst1_u8(dest + k * 8 + 5 * total_elements, vreinterpret_u8_u32(r2[2].val[1])); + vst1_u8(dest + k * 8 + 6 * total_elements, vreinterpret_u8_u32(r2[3].val[0])); + vst1_u8(dest + k * 8 + 7 * total_elements, vreinterpret_u8_u32(r2[3].val[1])); + } +} + +/* Routine optimized for shuffling a buffer for a type size of 16 bytes. */ +static void +shuffle16_neon(uint8_t *const dest, const uint8_t *const src, + const size_t vectorizable_elements, const size_t total_elements) { + size_t i, k; + static const size_t bytesoftype = 16; + uint8x8x2_t r0[8]; + uint16x4x2_t r1[8]; + uint32x2x2_t r2[8]; + + for (i = 0, k = 0; i < vectorizable_elements * bytesoftype; i += 128, k++) { + /* Load and interleave groups of 16 bytes (128 bytes) to the structure r0 */ + r0[0] = vzip_u8(vld1_u8(src + i + 0 * 8), vld1_u8(src + i + 2 * 8)); + r0[1] = vzip_u8(vld1_u8(src + i + 1 * 8), vld1_u8(src + i + 3 * 8)); + r0[2] = vzip_u8(vld1_u8(src + i + 4 * 8), vld1_u8(src + i + 6 * 8)); + r0[3] = vzip_u8(vld1_u8(src + i + 5 * 8), vld1_u8(src + i + 7 * 8)); + r0[4] = vzip_u8(vld1_u8(src + i + 8 * 8), vld1_u8(src + i + 10 * 8)); + r0[5] = vzip_u8(vld1_u8(src + i + 9 * 8), vld1_u8(src + i + 11 * 8)); + r0[6] = vzip_u8(vld1_u8(src + i + 12 * 8), vld1_u8(src + i + 14 * 8)); + r0[7] = vzip_u8(vld1_u8(src + i + 13 * 8), vld1_u8(src + i + 15 * 8)); + /* Interleave 16 bytes */ + r1[0] = vzip_u16(vreinterpret_u16_u8(r0[0].val[0]), vreinterpret_u16_u8(r0[2].val[0])); + r1[1] = vzip_u16(vreinterpret_u16_u8(r0[0].val[1]), vreinterpret_u16_u8(r0[2].val[1])); + r1[2] = vzip_u16(vreinterpret_u16_u8(r0[1].val[0]), vreinterpret_u16_u8(r0[3].val[0])); + r1[3] = vzip_u16(vreinterpret_u16_u8(r0[1].val[1]), vreinterpret_u16_u8(r0[3].val[1])); + r1[4] = vzip_u16(vreinterpret_u16_u8(r0[4].val[0]), vreinterpret_u16_u8(r0[6].val[0])); + r1[5] = vzip_u16(vreinterpret_u16_u8(r0[4].val[1]), vreinterpret_u16_u8(r0[6].val[1])); + r1[6] = vzip_u16(vreinterpret_u16_u8(r0[5].val[0]), vreinterpret_u16_u8(r0[7].val[0])); + r1[7] = vzip_u16(vreinterpret_u16_u8(r0[5].val[1]), vreinterpret_u16_u8(r0[7].val[1])); + /* Interleave 32 bytes */ + r2[0] = vzip_u32(vreinterpret_u32_u16(r1[0].val[0]), vreinterpret_u32_u16(r1[4].val[0])); + r2[1] = vzip_u32(vreinterpret_u32_u16(r1[0].val[1]), vreinterpret_u32_u16(r1[4].val[1])); + r2[2] = vzip_u32(vreinterpret_u32_u16(r1[1].val[0]), vreinterpret_u32_u16(r1[5].val[0])); + r2[3] = vzip_u32(vreinterpret_u32_u16(r1[1].val[1]), vreinterpret_u32_u16(r1[5].val[1])); + r2[4] = vzip_u32(vreinterpret_u32_u16(r1[2].val[0]), vreinterpret_u32_u16(r1[6].val[0])); + r2[5] = vzip_u32(vreinterpret_u32_u16(r1[2].val[1]), vreinterpret_u32_u16(r1[6].val[1])); + r2[6] = vzip_u32(vreinterpret_u32_u16(r1[3].val[0]), vreinterpret_u32_u16(r1[7].val[0])); + r2[7] = vzip_u32(vreinterpret_u32_u16(r1[3].val[1]), vreinterpret_u32_u16(r1[7].val[1])); + /* Store the results to the destination vector */ + vst1_u8(dest + k * 8 + 0 * total_elements, vreinterpret_u8_u32(r2[0].val[0])); + vst1_u8(dest + k * 8 + 1 * total_elements, vreinterpret_u8_u32(r2[0].val[1])); + vst1_u8(dest + k * 8 + 2 * total_elements, vreinterpret_u8_u32(r2[1].val[0])); + vst1_u8(dest + k * 8 + 3 * total_elements, vreinterpret_u8_u32(r2[1].val[1])); + vst1_u8(dest + k * 8 + 4 * total_elements, vreinterpret_u8_u32(r2[2].val[0])); + vst1_u8(dest + k * 8 + 5 * total_elements, vreinterpret_u8_u32(r2[2].val[1])); + vst1_u8(dest + k * 8 + 6 * total_elements, vreinterpret_u8_u32(r2[3].val[0])); + vst1_u8(dest + k * 8 + 7 * total_elements, vreinterpret_u8_u32(r2[3].val[1])); + vst1_u8(dest + k * 8 + 8 * total_elements, vreinterpret_u8_u32(r2[4].val[0])); + vst1_u8(dest + k * 8 + 9 * total_elements, vreinterpret_u8_u32(r2[4].val[1])); + vst1_u8(dest + k * 8 + 10 * total_elements, vreinterpret_u8_u32(r2[5].val[0])); + vst1_u8(dest + k * 8 + 11 * total_elements, vreinterpret_u8_u32(r2[5].val[1])); + vst1_u8(dest + k * 8 + 12 * total_elements, vreinterpret_u8_u32(r2[6].val[0])); + vst1_u8(dest + k * 8 + 13 * total_elements, vreinterpret_u8_u32(r2[6].val[1])); + vst1_u8(dest + k * 8 + 14 * total_elements, vreinterpret_u8_u32(r2[7].val[0])); + vst1_u8(dest + k * 8 + 15 * total_elements, vreinterpret_u8_u32(r2[7].val[1])); + } +} + +/* Routine optimized for unshuffling a buffer for a type size of 2 bytes. */ +static void +unshuffle2_neon(uint8_t *const dest, const uint8_t *const src, + const size_t vectorizable_elements, const size_t total_elements) { + size_t i, k; + static const size_t bytesoftype = 2; + uint8x16x2_t r0; + + for (i = 0, k = 0; i < vectorizable_elements * bytesoftype; i += 32, k++) { + /* Load 32 bytes to the structure r0 */ + r0.val[0] = vld1q_u8(src + total_elements * 0 + k * 16); + r0.val[1] = vld1q_u8(src + total_elements * 1 + k * 16); + /* Store (with permutation) the results in the destination vector */ + vst2q_u8(dest + k * 32, r0); + } +} + +/* Routine optimized for unshuffling a buffer for a type size of 4 bytes. */ +static void +unshuffle4_neon(uint8_t *const dest, const uint8_t *const src, + const size_t vectorizable_elements, const size_t total_elements) { + size_t i, k; + static const size_t bytesoftype = 4; + uint8x16x4_t r0; + + for (i = 0, k = 0; i < vectorizable_elements * bytesoftype; i += 64, k++) { + /* load 64 bytes to the structure r0 */ + r0.val[0] = vld1q_u8(src + total_elements * 0 + k * 16); + r0.val[1] = vld1q_u8(src + total_elements * 1 + k * 16); + r0.val[2] = vld1q_u8(src + total_elements * 2 + k * 16); + r0.val[3] = vld1q_u8(src + total_elements * 3 + k * 16); + /* Store (with permutation) the results in the destination vector */ + vst4q_u8(dest + k * 64, r0); + } +} + +/* Routine optimized for unshuffling a buffer for a type size of 8 bytes. */ +static void +unshuffle8_neon(uint8_t *const dest, const uint8_t *const src, + const size_t vectorizable_elements, const size_t total_elements) { + size_t i, k; + static const size_t bytesoftype = 8; + uint8x8x2_t r0[4]; + uint16x4x2_t r1[4]; + uint32x2x2_t r2[4]; + + for (i = 0, k = 0; i < vectorizable_elements * bytesoftype; i += 64, k++) { + /* Load and interleave groups of 8 bytes (64 bytes) to the structure r0 */ + r0[0] = vzip_u8(vld1_u8(src + 0 * total_elements + k * 8), vld1_u8(src + 1 * total_elements + k * 8)); + r0[1] = vzip_u8(vld1_u8(src + 2 * total_elements + k * 8), vld1_u8(src + 3 * total_elements + k * 8)); + r0[2] = vzip_u8(vld1_u8(src + 4 * total_elements + k * 8), vld1_u8(src + 5 * total_elements + k * 8)); + r0[3] = vzip_u8(vld1_u8(src + 6 * total_elements + k * 8), vld1_u8(src + 7 * total_elements + k * 8)); + /* Interleave 16 bytes */ + r1[0] = vzip_u16(vreinterpret_u16_u8(r0[0].val[0]), vreinterpret_u16_u8(r0[1].val[0])); + r1[1] = vzip_u16(vreinterpret_u16_u8(r0[0].val[1]), vreinterpret_u16_u8(r0[1].val[1])); + r1[2] = vzip_u16(vreinterpret_u16_u8(r0[2].val[0]), vreinterpret_u16_u8(r0[3].val[0])); + r1[3] = vzip_u16(vreinterpret_u16_u8(r0[2].val[1]), vreinterpret_u16_u8(r0[3].val[1])); + /* Interleave 32 bytes */ + r2[0] = vzip_u32(vreinterpret_u32_u16(r1[0].val[0]), vreinterpret_u32_u16(r1[2].val[0])); + r2[1] = vzip_u32(vreinterpret_u32_u16(r1[0].val[1]), vreinterpret_u32_u16(r1[2].val[1])); + r2[2] = vzip_u32(vreinterpret_u32_u16(r1[1].val[0]), vreinterpret_u32_u16(r1[3].val[0])); + r2[3] = vzip_u32(vreinterpret_u32_u16(r1[1].val[1]), vreinterpret_u32_u16(r1[3].val[1])); + /* Store the results in the destination vector */ + vst1_u8(dest + i + 0 * 8, vreinterpret_u8_u32(r2[0].val[0])); + vst1_u8(dest + i + 1 * 8, vreinterpret_u8_u32(r2[0].val[1])); + vst1_u8(dest + i + 2 * 8, vreinterpret_u8_u32(r2[1].val[0])); + vst1_u8(dest + i + 3 * 8, vreinterpret_u8_u32(r2[1].val[1])); + vst1_u8(dest + i + 4 * 8, vreinterpret_u8_u32(r2[2].val[0])); + vst1_u8(dest + i + 5 * 8, vreinterpret_u8_u32(r2[2].val[1])); + vst1_u8(dest + i + 6 * 8, vreinterpret_u8_u32(r2[3].val[0])); + vst1_u8(dest + i + 7 * 8, vreinterpret_u8_u32(r2[3].val[1])); + + } +} + +/* Routine optimized for unshuffling a buffer for a type size of 16 bytes. */ +static void +unshuffle16_neon(uint8_t *const dest, const uint8_t *const src, + const size_t vectorizable_elements, const size_t total_elements) { + size_t i, k; + static const size_t bytesoftype = 16; + uint8x8x2_t r0[8]; + uint16x4x2_t r1[8]; + uint32x2x2_t r2[8]; + + for (i = 0, k = 0; i < vectorizable_elements * bytesoftype; i += 128, k++) { + /* Load and interleave groups of 16 bytes (128 bytes) to the structure r0*/ + r0[0] = vzip_u8(vld1_u8(src + k * 8 + 0 * total_elements), vld1_u8(src + k * 8 + 1 * total_elements)); + r0[1] = vzip_u8(vld1_u8(src + k * 8 + 2 * total_elements), vld1_u8(src + k * 8 + 3 * total_elements)); + r0[2] = vzip_u8(vld1_u8(src + k * 8 + 4 * total_elements), vld1_u8(src + k * 8 + 5 * total_elements)); + r0[3] = vzip_u8(vld1_u8(src + k * 8 + 6 * total_elements), vld1_u8(src + k * 8 + 7 * total_elements)); + r0[4] = vzip_u8(vld1_u8(src + k * 8 + 8 * total_elements), vld1_u8(src + k * 8 + 9 * total_elements)); + r0[5] = vzip_u8(vld1_u8(src + k * 8 + 10 * total_elements), vld1_u8(src + k * 8 + 11 * total_elements)); + r0[6] = vzip_u8(vld1_u8(src + k * 8 + 12 * total_elements), vld1_u8(src + k * 8 + 13 * total_elements)); + r0[7] = vzip_u8(vld1_u8(src + k * 8 + 14 * total_elements), vld1_u8(src + k * 8 + 15 * total_elements)); + /* Interleave 16 bytes */ + r1[0] = vzip_u16(vreinterpret_u16_u8(r0[0].val[0]), vreinterpret_u16_u8(r0[1].val[0])); + r1[1] = vzip_u16(vreinterpret_u16_u8(r0[0].val[1]), vreinterpret_u16_u8(r0[1].val[1])); + r1[2] = vzip_u16(vreinterpret_u16_u8(r0[2].val[0]), vreinterpret_u16_u8(r0[3].val[0])); + r1[3] = vzip_u16(vreinterpret_u16_u8(r0[2].val[1]), vreinterpret_u16_u8(r0[3].val[1])); + r1[4] = vzip_u16(vreinterpret_u16_u8(r0[4].val[0]), vreinterpret_u16_u8(r0[5].val[0])); + r1[5] = vzip_u16(vreinterpret_u16_u8(r0[4].val[1]), vreinterpret_u16_u8(r0[5].val[1])); + r1[6] = vzip_u16(vreinterpret_u16_u8(r0[6].val[0]), vreinterpret_u16_u8(r0[7].val[0])); + r1[7] = vzip_u16(vreinterpret_u16_u8(r0[6].val[1]), vreinterpret_u16_u8(r0[7].val[1])); + /* Interleave 32 bytes */ + r2[0] = vzip_u32(vreinterpret_u32_u16(r1[0].val[0]), vreinterpret_u32_u16(r1[2].val[0])); + r2[1] = vzip_u32(vreinterpret_u32_u16(r1[0].val[1]), vreinterpret_u32_u16(r1[2].val[1])); + r2[2] = vzip_u32(vreinterpret_u32_u16(r1[1].val[0]), vreinterpret_u32_u16(r1[3].val[0])); + r2[3] = vzip_u32(vreinterpret_u32_u16(r1[1].val[1]), vreinterpret_u32_u16(r1[3].val[1])); + r2[4] = vzip_u32(vreinterpret_u32_u16(r1[4].val[0]), vreinterpret_u32_u16(r1[6].val[0])); + r2[5] = vzip_u32(vreinterpret_u32_u16(r1[4].val[1]), vreinterpret_u32_u16(r1[6].val[1])); + r2[6] = vzip_u32(vreinterpret_u32_u16(r1[5].val[0]), vreinterpret_u32_u16(r1[7].val[0])); + r2[7] = vzip_u32(vreinterpret_u32_u16(r1[5].val[1]), vreinterpret_u32_u16(r1[7].val[1])); + /* Store the results in the destination vector */ + vst1_u8(dest + i + 0 * 8, vreinterpret_u8_u32(r2[0].val[0])); + vst1_u8(dest + i + 1 * 8, vreinterpret_u8_u32(r2[4].val[0])); + vst1_u8(dest + i + 2 * 8, vreinterpret_u8_u32(r2[0].val[1])); + vst1_u8(dest + i + 3 * 8, vreinterpret_u8_u32(r2[4].val[1])); + vst1_u8(dest + i + 4 * 8, vreinterpret_u8_u32(r2[1].val[0])); + vst1_u8(dest + i + 5 * 8, vreinterpret_u8_u32(r2[5].val[0])); + vst1_u8(dest + i + 6 * 8, vreinterpret_u8_u32(r2[1].val[1])); + vst1_u8(dest + i + 7 * 8, vreinterpret_u8_u32(r2[5].val[1])); + vst1_u8(dest + i + 8 * 8, vreinterpret_u8_u32(r2[2].val[0])); + vst1_u8(dest + i + 9 * 8, vreinterpret_u8_u32(r2[6].val[0])); + vst1_u8(dest + i + 10 * 8, vreinterpret_u8_u32(r2[2].val[1])); + vst1_u8(dest + i + 11 * 8, vreinterpret_u8_u32(r2[6].val[1])); + vst1_u8(dest + i + 12 * 8, vreinterpret_u8_u32(r2[3].val[0])); + vst1_u8(dest + i + 13 * 8, vreinterpret_u8_u32(r2[7].val[0])); + vst1_u8(dest + i + 14 * 8, vreinterpret_u8_u32(r2[3].val[1])); + vst1_u8(dest + i + 15 * 8, vreinterpret_u8_u32(r2[7].val[1])); + } +} + +/* Size in bytes of the chunk processed by one iteration of a NEON kernel, + or 0 if there is no NEON kernel for this type size. */ +static size_t neon_chunk_size(const size_t bytesoftype) { + switch (bytesoftype) { + case 2: + case 4: + return bytesoftype * 16; + case 8: + case 16: + return bytesoftype * 8; + default: + return 0; + } +} + +/* Shuffle a block. This can never fail. */ +void blosc_internal_shuffle_generic(const size_t bytesoftype, const size_t blocksize, + const uint8_t* const _src, uint8_t* const _dest) +{ + const size_t chunk = neon_chunk_size(bytesoftype); + size_t vectorizable_bytes, vectorizable_elements, total_elements; + + if (chunk == 0 || blocksize < chunk) { + shuffle_generic_inline(bytesoftype, 0, blocksize, _src, _dest); + return; + } + vectorizable_bytes = blocksize - (blocksize % chunk); + vectorizable_elements = vectorizable_bytes / bytesoftype; + total_elements = blocksize / bytesoftype; + + switch (bytesoftype) { + case 2: + shuffle2_neon(_dest, _src, vectorizable_elements, total_elements); + break; + case 4: + shuffle4_neon(_dest, _src, vectorizable_elements, total_elements); + break; + case 8: + shuffle8_neon(_dest, _src, vectorizable_elements, total_elements); + break; + case 16: + shuffle16_neon(_dest, _src, vectorizable_elements, total_elements); + break; + } + /* Handles the non-vectorizable tail and copies any leftover bytes. */ + shuffle_generic_inline(bytesoftype, vectorizable_bytes, blocksize, _src, _dest); +} + +/* Unshuffle a block. This can never fail. */ +void blosc_internal_unshuffle_generic(const size_t bytesoftype, const size_t blocksize, + const uint8_t* const _src, uint8_t* const _dest) +{ + const size_t chunk = neon_chunk_size(bytesoftype); + size_t vectorizable_bytes, vectorizable_elements, total_elements; + + if (chunk == 0 || blocksize < chunk) { + unshuffle_generic_inline(bytesoftype, 0, blocksize, _src, _dest); + return; + } + vectorizable_bytes = blocksize - (blocksize % chunk); + vectorizable_elements = vectorizable_bytes / bytesoftype; + total_elements = blocksize / bytesoftype; + + switch (bytesoftype) { + case 2: + unshuffle2_neon(_dest, _src, vectorizable_elements, total_elements); + break; + case 4: + unshuffle4_neon(_dest, _src, vectorizable_elements, total_elements); + break; + case 8: + unshuffle8_neon(_dest, _src, vectorizable_elements, total_elements); + break; + case 16: + unshuffle16_neon(_dest, _src, vectorizable_elements, total_elements); + break; + } + /* Handles the non-vectorizable tail and copies any leftover bytes. */ + unshuffle_generic_inline(bytesoftype, vectorizable_bytes, blocksize, _src, _dest); +} diff --git a/src/numcodecs/meson.build b/src/numcodecs/meson.build index 41a3a7a2..ed813f01 100644 --- a/src/numcodecs/meson.build +++ b/src/numcodecs/meson.build @@ -12,6 +12,7 @@ endif # Correct for cross-compilation: checks what the *target* supports. have_sse2 = false have_avx2 = false +have_neon = false if host_machine.cpu_family() == 'x86_64' # Honor legacy environment variables for backwards compatibility with @@ -38,6 +39,17 @@ if host_machine.cpu_family() == 'x86_64' endif endif +# NEON is mandatory on aarch64, so no runtime detection is needed. +if host_machine.cpu_family() == 'aarch64' and not get_option('neon').disabled() + have_neon = cc.compiles(''' + #if !defined(__ARM_NEON) && !defined(_M_ARM64) + #error "no NEON" + #endif + #include + int main(void) { uint8x16_t v = vdupq_n_u8(0); return vgetq_lane_u8(v, 0); } + ''', name: 'NEON intrinsics') +endif + # --- Vendored zstd --- zstd_dep = dependency('libzstd', required: get_option('system_zstd')) @@ -165,10 +177,17 @@ if not blosc_dep.found() '../../c-blosc/blosc/blosclz.c', '../../c-blosc/blosc/fastcopy.c', '../../c-blosc/blosc/shuffle.c', - '../../c-blosc/blosc/shuffle-generic.c', '../../c-blosc/blosc/bitshuffle-generic.c', ) + # c-blosc 1.x has no ARM SIMD shuffle; on aarch64 swap in a NEON + # implementation of the generic (un)shuffle routines. + if have_neon + blosc_sources += files('_blosc_neon/shuffle-generic-neon.c') + else + blosc_sources += files('../../c-blosc/blosc/shuffle-generic.c') + endif + blosc_c_args = c_args + ['-DHAVE_LZ4', '-DHAVE_ZLIB', '-DHAVE_ZSTD'] if have_sse2 diff --git a/tests/test_blosc.py b/tests/test_blosc.py index b920f03e..b54efcf6 100644 --- a/tests/test_blosc.py +++ b/tests/test_blosc.py @@ -330,3 +330,15 @@ def test_config_no_typesize(): codec = Blosc(shuffle=Blosc.SHUFFLE, typesize=5) config = codec.get_config() assert "typesize" not in config + + +@pytest.mark.parametrize('shuffle', [Blosc.SHUFFLE, Blosc.BITSHUFFLE]) +@pytest.mark.parametrize('blocksize', [0, 2**8, 4097]) +@pytest.mark.parametrize('nbytes', [1, 31, 33, 127, 1000, 65537]) +@pytest.mark.parametrize('typesize', [1, 2, 3, 4, 7, 8, 12, 16, 17, 32]) +def test_shuffle_roundtrip_typesizes(typesize, nbytes, blocksize, shuffle): + # Exercises the SIMD (un)shuffle kernels, including the scalar tail used when a + # block is not a multiple of the vector width or the typesize. + src = np.random.default_rng(0).integers(0, 64, nbytes, dtype='u1').tobytes() + codec = Blosc(cname='lz4', clevel=5, shuffle=shuffle, blocksize=blocksize, typesize=typesize) + assert bytes(codec.decode(codec.encode(src))) == src