From fbc507839977d2cdcc0e9fbb8a4f66f3a739e154 Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Mon, 3 Aug 2026 21:27:27 +0200 Subject: [PATCH] compression: use the NEON bitshuffle on aarch64 bitshuffle_hperf is an x86-only implementation - its entire vector body sits behind __i386__/__x86_64__, so on aarch64 every entry point compiles down to the scalar fallback. Measured against its own SIMD path that costs 8.3x on encode and 3.8x on decode, and it is the transform behind every compressed image the writer produces and every one the reader, preview and XDS plugin take apart again. The classic bitshuffle vendored beside it does have an aarch64 NEON path, and is already compiled into the same target, so this costs nothing new. BitShuffleBlock.h picks bshuf_trans_bit_elem/bshuf_untrans_bit_elem there and keeps bitshuf_encode_block / bitshuf_decode_block everywhere else, where hperf is about twice classic SSE2 and remains the better choice. The expected aarch64 gain is ~2.5x encode / ~1.7x decode: classic NEON is 128-bit and carries an extra pass, so it recovers part of the gap rather than all of it. The condition mirrors USEARMNEON in bitshuffle_core.c exactly, because with NEON off the classic scalar path is slower than hperf's and must not be selected. Swapping implementations is only safe while the two agree bit for bit - otherwise an ARM build would write files an x86 build could not read. They do: verified byte-identical output and mutual cross-decoding for elem_size 1/2/4/8 over block sizes from 8 to 65536 elements. Both are always compiled in, so the new test holds them to it on every architecture, not just the one that would notice. Co-Authored-By: Claude Opus 5 (1M context) --- compression/BitShuffleBlock.h | 42 ++++++++++++++++++++++++++++++++ compression/CMakeLists.txt | 1 + compression/JFJochCompressor.cpp | 4 +-- compression/JFJochDecompress.h | 14 +++++------ tests/ZSTDCompressorTest.cpp | 30 +++++++++++++++++++++++ 5 files changed, 82 insertions(+), 9 deletions(-) create mode 100644 compression/BitShuffleBlock.h diff --git a/compression/BitShuffleBlock.h b/compression/BitShuffleBlock.h new file mode 100644 index 00000000..b740a4f6 --- /dev/null +++ b/compression/BitShuffleBlock.h @@ -0,0 +1,42 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#pragma once + +#include +#include + +// One bitshuffle block, transformed by whichever of the two vendored implementations is SIMD on +// this architecture. The two write byte-identical output and each decodes the other's, so the file +// format does not depend on the build host. +// +// bitshuffle_hperf is x86-only: outside SSE2 its whole vector body is compiled out and what remains +// is a scalar fallback. The classic bitshuffle has an aarch64 NEON path (bitshuffle_core.c, +// USEARMNEON), so aarch64 uses that instead - measured against the hperf scalar fallback it is +// ~2.5x on encode and ~1.7x on decode. Everywhere else hperf wins outright (~2x over classic SSE2), +// so it stays the default. +// +// The condition mirrors USEARMNEON in bitshuffle_core.c exactly. With NEON off the classic path +// falls back to a scalar of its own that is slower than hperf's, so it must not be selected then. +#if (defined(__ARM_NEON__) || (__ARM_NEON)) && defined(__aarch64__) + +// The classic entry points allocate their 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); +} + +#else + +inline int64_t JFJochBitShuffleBlock(char *out, const char *in, char *scratch, size_t size, size_t elem_size) { + return bitshuf_encode_block(out, in, scratch, size, elem_size); +} + +inline int64_t JFJochBitUnshuffleBlock(char *out, const char *in, char *scratch, size_t size, size_t elem_size) { + return bitshuf_decode_block(out, in, scratch, size, elem_size); +} + +#endif diff --git a/compression/CMakeLists.txt b/compression/CMakeLists.txt index 73a875d3..935ee6df 100644 --- a/compression/CMakeLists.txt +++ b/compression/CMakeLists.txt @@ -13,6 +13,7 @@ ADD_LIBRARY(Compression STATIC JFJochCompressor.cpp JFJochCompressor.h JFJochDecompress.h + BitShuffleBlock.h MaxCompressedSize.cpp MaxCompressedSize.h) set_target_properties(Compression PROPERTIES POSITION_INDEPENDENT_CODE ON) diff --git a/compression/JFJochCompressor.cpp b/compression/JFJochCompressor.cpp index b837fa5b..600d55ee 100644 --- a/compression/JFJochCompressor.cpp +++ b/compression/JFJochCompressor.cpp @@ -6,10 +6,10 @@ #include #include #include -#include #include #include +#include "BitShuffleBlock.h" #include "../common/JFJochException.h" extern "C" { @@ -45,7 +45,7 @@ JFJochBitShuffleCompressor::JFJochBitShuffleCompressor(CompressionAlgorithm in_a size_t JFJochBitShuffleCompressor::CompressBlock(char *dest, const char *source, size_t nelements, size_t elem_size) { // Assert nelements < block_size const char *src_ptr; - int64_t bshuf_ret = bitshuf_encode_block(tmp_space.data(), source, scratch.data(), nelements, elem_size); + int64_t bshuf_ret = JFJochBitShuffleBlock(tmp_space.data(), source, scratch.data(), nelements, elem_size); if (bshuf_ret < 0) throw JFJochException(JFJochExceptionCategory::Compression, "bshuf_trans_bit_elem error"); src_ptr = tmp_space.data(); diff --git a/compression/JFJochDecompress.h b/compression/JFJochDecompress.h index 6f07e5a4..bb6d8d55 100644 --- a/compression/JFJochDecompress.h +++ b/compression/JFJochDecompress.h @@ -8,10 +8,10 @@ #include #include -#include #include #include +#include "BitShuffleBlock.h" #include "../compression/CompressionAlgorithmEnum.h" #include "../common/JFJochException.h" @@ -90,12 +90,12 @@ inline size_t JFJochDecompressHperfPtr(uint8_t *output, throw JFJochException(JFJochExceptionCategory::Compression, "Algorithm not supported"); } - if (bitshuf_decode_block(reinterpret_cast(dst_ptr), - decompressed_block.data(), - scratch.data(), - current_nelements, - elem_size) < 0) - throw JFJochException(JFJochExceptionCategory::Compression, "bitshuffle_hperf decode error"); + if (JFJochBitUnshuffleBlock(reinterpret_cast(dst_ptr), + decompressed_block.data(), + scratch.data(), + current_nelements, + elem_size) < 0) + throw JFJochException(JFJochExceptionCategory::Compression, "bitshuffle block decode error"); src_ptr += compressed_size; dst_ptr += decompressed_size; diff --git a/tests/ZSTDCompressorTest.cpp b/tests/ZSTDCompressorTest.cpp index 2f40d06a..5a497601 100644 --- a/tests/ZSTDCompressorTest.cpp +++ b/tests/ZSTDCompressorTest.cpp @@ -3,6 +3,7 @@ #include #include +#include #include #include #include @@ -317,6 +318,35 @@ TEST_CASE("Bitshuffle_ZSTD","[ZSTD]") { REQUIRE(memcmp(image.data(), decompressed.data(), RAW_MODULE_SIZE*sizeof(uint32_t)) == 0); } + +// BitShuffleBlock.h picks between the two vendored bitshuffle implementations by architecture: +// bitshuffle_hperf on x86, the classic one on aarch64, where only the latter has SIMD. That is only +// safe as long as the two agree bit for bit - otherwise an ARM build would write files an x86 build +// cannot read. Both are always compiled in, so this holds them to it on every architecture. +TEST_CASE("Bitshuffle_hperf_matches_classic","[ZSTD]") { + for (size_t elem_size : {1, 2, 4, 8}) { + // Element counts the compressor actually produces (128 kB / 16 kB blocks), plus the + // smallest legal block and a short trailing one. + const size_t sizes[] = {8, 64, 16384 / elem_size, 131072 / elem_size}; + for (size_t size : sizes) { + std::vector image(size * elem_size), hperf(size * elem_size), + classic(size * elem_size), scratch(size * elem_size), decoded(size * elem_size); + std::mt19937 gen(1234); + for (auto &c : image) + c = static_cast(gen() & 0xff); + + REQUIRE(bitshuf_encode_block(hperf.data(), image.data(), scratch.data(), size, elem_size) >= 0); + REQUIRE(bshuf_trans_bit_elem(image.data(), classic.data(), size, elem_size) >= 0); + REQUIRE(hperf == classic); + + // Each implementation must also decode what the other wrote. + REQUIRE(bitshuf_decode_block(decoded.data(), classic.data(), scratch.data(), size, elem_size) >= 0); + REQUIRE(decoded == image); + REQUIRE(bshuf_untrans_bit_elem(hperf.data(), decoded.data(), size, elem_size) >= 0); + REQUIRE(decoded == image); + } + } +} // ---- BSHUF_ZSTD_RLE_HUFF (RLE runs + adaptive Huffman literals, standard zstd frame) ---- // Round-trips through JFJochDecompress (which calls stock ZSTD_decompress per block) and checks // the output never exceeds the advertised worst-case size. Run on good and adversarial data to