diff --git a/compression/BitShuffleBlock.h b/compression/BitShuffleBlock.h index a4a73f76a..0f053b16d 100644 --- a/compression/BitShuffleBlock.h +++ b/compression/BitShuffleBlock.h @@ -23,13 +23,24 @@ // which a single configure-time answer could not follow. #if (defined(__ARM_NEON__) || (__ARM_NEON)) && defined(__aarch64__) -// The classic entry points allocate their own block-sized scratch, so the caller's goes unused. +// The two stages of bshuf_untrans_bit_elem_NEON, which bitshuffle_core.c defines but no header +// declares. Calling them directly lets the decode use the caller's scratch instead of the block-sized +// buffer the classic entry point mallocs and frees for every block. +extern "C" { + int64_t bshuf_trans_byte_bitrow_NEON(const void *in, void *out, size_t size, size_t elem_size); + int64_t bshuf_shuffle_bit_eightelem_NEON(const void *in, void *out, size_t size, size_t elem_size); +} + +// The classic encode entry point allocates its own block-sized scratch, so the caller's goes unused. inline int64_t JFJochBitShuffleBlock(char *out, const char *in, char *, size_t size, size_t elem_size) { return bshuf_trans_bit_elem(in, out, size, elem_size); } -inline int64_t JFJochBitUnshuffleBlock(char *out, const char *in, char *, size_t size, size_t elem_size) { - return bshuf_untrans_bit_elem(in, out, size, elem_size); +inline int64_t JFJochBitUnshuffleBlock(char *out, const char *in, char *scratch, size_t size, size_t elem_size) { + const int64_t count = bshuf_trans_byte_bitrow_NEON(in, scratch, size, elem_size); + if (count < 0) + return count; + return bshuf_shuffle_bit_eightelem_NEON(scratch, out, size, elem_size); } #else