Files
Jungfraujoch/tests/ZSTDCompressorTest.cpp
T
leonarski_fandClaude Opus 5 fbc5078399 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) <noreply@anthropic.com>
2026-08-03 21:27:27 +02:00

440 lines
18 KiB
C++

// SPDX-FileCopyrightText: 2024 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
// SPDX-License-Identifier: GPL-3.0-only
#include <catch2/catch_all.hpp>
#include <bitshuffle/bitshuffle.h>
#include <bitshuffle_hperf/bitshuffle.h>
#include <zstd.h>
#include <iostream>
#include <random>
#include "../compression/JFJochCompressor.h"
#include "../compression/JFJochDecompress.h"
#include "../common/DiffractionExperiment.h"
#include "../writer/HDF5Objects.h"
TEST_CASE("JFjochZstdCompressor_Raw_Block","[ZSTD]") {
uint8_t src[128*8];
uint8_t dst[128*8+3];
for (int i = 0; i < 128*8; i++)
src[i] = i % 256;
size_t dst_size = JFJochZstdCompressor::RawBlock(dst, src, 128 * 8, false);
REQUIRE(dst_size == 128*8+3);
int diff = 0;
for (int i = 0; i < 128*8; i++) {
if (dst[i+3] != i % 256) diff++;
}
REQUIRE(diff == 0);
}
TEST_CASE("JFjochZstdCompressor_RLE_Block","[ZSTD]") {
uint8_t dst[4];
size_t dst_size = JFJochZstdCompressor::RLEBlock(dst, 0xAF, 128 * 8, false);
REQUIRE(dst_size == 4);
REQUIRE(dst[3] == 0xAF);
}
TEST_CASE("JFjochZstdCompressor_Frame_onlyRLE","[ZSTD]") {
JFJochZstdCompressor compressor;
uint8_t src[128 * 8];
uint8_t dst[128 * 8 * 2];
uint8_t decomp_buffer[128 * 8 * 2];
for (int i = 0; i < 128 * 8; i++) {
src[i] = 0;
dst[i] = 0xFF;
}
size_t dst_size = compressor.Compress(dst, (uint64_t *) src, 128*8, 128*8);
REQUIRE(ZSTD_decompress(decomp_buffer, 128 * 8 * 2, dst, dst_size) == 8 * 128);
size_t diff = 0;
for (int i = 0; i < 128 * 8; i++) if (src[i] != decomp_buffer[i]) diff++;
REQUIRE(diff == 0);
REQUIRE(dst_size == 4 + 9);
REQUIRE(dst[9 + 3] == 0);
}
TEST_CASE("JFjochZstdCompressor_Frame_onlyRLE_1block","[ZSTD]") {
JFJochZstdCompressor compressor;
uint8_t src[ZSTD_BLOCKSIZE_MAX * 2];
uint8_t dst[ZSTD_BLOCKSIZE_MAX * 4];
uint8_t decomp_buffer[ZSTD_BLOCKSIZE_MAX * 4];
for (int i = 0; i < ZSTD_BLOCKSIZE_MAX * 2; i++) {
src[i] = 0;
dst[i] = 0xFF;
}
size_t dst_size = compressor.Compress(dst, (uint64_t *) src, ZSTD_BLOCKSIZE_MAX, ZSTD_BLOCKSIZE_MAX);
REQUIRE(dst_size == 4 + 9);
REQUIRE(dst[9 + 3] == 0);
REQUIRE(ZSTD_decompress(decomp_buffer, ZSTD_BLOCKSIZE_MAX * 4, dst, dst_size) == ZSTD_BLOCKSIZE_MAX);
size_t diff = 0;
for (int i = 0; i < ZSTD_BLOCKSIZE_MAX; i++) if (src[i] != decomp_buffer[i]) diff++;
REQUIRE(diff == 0);
}
TEST_CASE("JFjochZstdCompressor_Frame_onlyRLE_2blocks","[ZSTD]") {
JFJochZstdCompressor compressor;
uint8_t src[ZSTD_BLOCKSIZE_MAX * 2];
uint8_t dst[ZSTD_BLOCKSIZE_MAX * 4];
uint8_t decomp_buffer[ZSTD_BLOCKSIZE_MAX * 4];
for (int i = 0; i < ZSTD_BLOCKSIZE_MAX * 2; i++) {
src[i] = 0;
dst[i] = 0xFF;
}
size_t dst_size = compressor.Compress(dst, (uint64_t *) src, ZSTD_BLOCKSIZE_MAX * 2, ZSTD_BLOCKSIZE_MAX * 2);
REQUIRE(dst_size == 8 + 9);
REQUIRE(dst[9 + 3] == 0);
REQUIRE(dst[9 + 7] == 0);
REQUIRE(ZSTD_decompress(decomp_buffer, ZSTD_BLOCKSIZE_MAX * 4, dst, dst_size) == ZSTD_BLOCKSIZE_MAX * 2);
size_t diff = 0;
for (int i = 0; i < ZSTD_BLOCKSIZE_MAX * 2; i++) if (src[i] != decomp_buffer[i]) diff++;
REQUIRE(diff == 0);
}
TEST_CASE("JFjochZstdCompressor_Frame_onlyRLE_0xFF_2blocks","[ZSTD]") {
JFJochZstdCompressor compressor;
uint8_t src[ZSTD_BLOCKSIZE_MAX * 2];
uint8_t dst[ZSTD_BLOCKSIZE_MAX * 4];
uint8_t decomp_buffer[ZSTD_BLOCKSIZE_MAX * 4];
for (int i = 0; i < ZSTD_BLOCKSIZE_MAX * 2; i++) {
src[i] = 0xFF;
dst[i] = 0x0;
}
size_t dst_size = compressor.Compress(dst, (uint64_t *) src, ZSTD_BLOCKSIZE_MAX * 2, ZSTD_BLOCKSIZE_MAX * 2);
REQUIRE(dst_size == 8 + 9);
REQUIRE(dst[9 + 3] == 0xFF);
REQUIRE(dst[9 + 7] == 0xFF);
REQUIRE(ZSTD_decompress(decomp_buffer, ZSTD_BLOCKSIZE_MAX * 4, dst, dst_size) == ZSTD_BLOCKSIZE_MAX * 2);
size_t diff = 0;
for (int i = 0; i < ZSTD_BLOCKSIZE_MAX * 2; i++) if (src[i] != decomp_buffer[i]) diff++;
REQUIRE(diff == 0);
}
TEST_CASE("JFjochZstdCompressor_Frame_onlyRAW","[ZSTD]") {
JFJochZstdCompressor compressor;
uint8_t src[128 * 8];
uint8_t dst[128 * 8 * 2];
uint8_t decomp_buffer[128 * 8 * 2];
for (int i = 0; i < 128 * 8; i++)
src[i] = 1 + i % 128;
size_t dst_size = compressor.Compress(dst, (uint64_t *) src, 128 * 8, 128*8);
REQUIRE(ZSTD_decompress(decomp_buffer, 128 * 8 * 2, dst, dst_size) == 8 * 128);
size_t diff = 0;
for (int i = 0; i < 128 * 8; i++) if (src[i] != decomp_buffer[i]) diff++;
REQUIRE(diff == 0);
REQUIRE(dst_size == 3 + 9 + 128 * 8);
diff = 0;
for (int i = 0; i < 128 * 8; i++) {
if (dst[i + 3 + 9] != 1 + i % 128)
diff++;
}
REQUIRE(diff == 0);
}
TEST_CASE("JFjochZstdCompressor_Frame_mixed","[ZSTD]") {
JFJochZstdCompressor compressor;
uint8_t src[128 * 8];
uint8_t dst[128 * 8 * 2];
uint8_t decomp_buffer[128 * 8 * 2];
for (int i = 0; i < 128*8; i++)
src[i] = (i / 128 + 1) % 2 ;
size_t dst_size = compressor.Compress(dst, (uint64_t *) src, 128 * 8, 128 * 8);
REQUIRE(dst_size == 9 + 4 * 4 + 4 * (3+128));
REQUIRE(ZSTD_decompress(decomp_buffer, 128*8*2, dst, dst_size) == 8 * 128);
size_t diff = 0;
for (int i = 0; i < 128*8; i++) if (src[i] != decomp_buffer[i]) diff++;
REQUIRE(diff == 0);
}
TEST_CASE("JFjochZstdCompressor_Frame_simulated","[ZSTD]") {
for(int seed = 0; seed < 200; seed++) {
JFJochZstdCompressor compressor;
// Predictable random number generator
std::mt19937 g1(seed);
std::uniform_int_distribution<int16_t> distribution;
std::vector<int16_t> input(RAW_MODULE_SIZE);
for (auto &i: input) i = distribution(g1);
std::vector<char> input_shuffled(RAW_MODULE_SIZE * sizeof(uint16_t));
bshuf_bitshuffle(input.data(), input_shuffled.data(), RAW_MODULE_SIZE, 2, 4096);
std::vector<uint8_t> input_compressed(ZSTD_compressBound(RAW_MODULE_SIZE * sizeof(uint16_t)));
std::vector<uint8_t> output(RAW_MODULE_SIZE * sizeof(uint16_t) * 2);
size_t dst_size = compressor.Compress(input_compressed.data(), (uint64_t *) input_shuffled.data(),
RAW_MODULE_SIZE * sizeof(uint16_t), RAW_MODULE_SIZE * sizeof(uint16_t));
REQUIRE(ZSTD_decompress(output.data(), RAW_MODULE_SIZE * sizeof(uint16_t) * 2,
input_compressed.data(), dst_size) == RAW_MODULE_SIZE * sizeof(uint16_t));
REQUIRE(memcmp(input_shuffled.data(), output.data(), RAW_MODULE_SIZE * sizeof(uint16_t)) == 0);
}
}
TEST_CASE("JFjochZstdCompressor_Frame_zeroes","[ZSTD]") {
JFJochZstdCompressor compressor;
std::vector<char> input_shuffled(RAW_MODULE_SIZE * sizeof(uint16_t));
for (auto &i: input_shuffled) i = 0x0u;
std::vector<uint8_t> input_compressed(ZSTD_compressBound(RAW_MODULE_SIZE * sizeof(uint16_t)));
std::vector<uint8_t> output(RAW_MODULE_SIZE * sizeof(uint16_t) * 2);
size_t dst_size = compressor.Compress(input_compressed.data(), (uint64_t *) input_shuffled.data(),
RAW_MODULE_SIZE * sizeof(uint16_t), RAW_MODULE_SIZE * sizeof(uint16_t));
REQUIRE(dst_size == 9 + 4 * 8);
REQUIRE(ZSTD_decompress(output.data(), RAW_MODULE_SIZE * sizeof(uint16_t) * 2,
input_compressed.data(), dst_size) == RAW_MODULE_SIZE * sizeof(uint16_t));
REQUIRE(memcmp(input_shuffled.data(), output.data(), RAW_MODULE_SIZE * sizeof(uint16_t)) == 0);
}
TEST_CASE("JFjochZstdCompressor_Frame_ones","[ZSTD]") {
JFJochZstdCompressor compressor;
std::vector<char> input_shuffled(RAW_MODULE_SIZE * sizeof(uint16_t));
for (auto &i: input_shuffled) i = 0xFFu;
std::vector<uint8_t> input_compressed(ZSTD_compressBound(RAW_MODULE_SIZE * sizeof(uint16_t)));
std::vector<uint8_t> output(RAW_MODULE_SIZE * sizeof(uint16_t) * 2);
size_t dst_size = compressor.Compress(input_compressed.data(), (uint64_t *) input_shuffled.data(),
RAW_MODULE_SIZE * sizeof(uint16_t), RAW_MODULE_SIZE * sizeof(uint16_t));
REQUIRE(dst_size == 9 + 4 * 8);
REQUIRE(ZSTD_decompress(output.data(), RAW_MODULE_SIZE * sizeof(uint16_t) * 2,
input_compressed.data(), dst_size) == RAW_MODULE_SIZE * sizeof(uint16_t));
REQUIRE(memcmp(input_shuffled.data(), output.data(), RAW_MODULE_SIZE * sizeof(uint16_t)) == 0);
}
TEST_CASE("JFJochCompressor_JFJochDecompressor_ZSTD","[ZSTD]") {
DiffractionExperiment x(DetJF4M());
x.Compression(CompressionAlgorithm::BSHUF_ZSTD).BitDepthImage(32).PixelSigned(true);
std::vector<int32_t> image(x.GetPixelsNum());
for (auto &i: image)
i = 345;
JFJochBitShuffleCompressor compressor(x.GetCompressionAlgorithm());
std::vector<char> tmp(x.GetPixelsNum() * sizeof(int32_t) * 4 + 12);
auto tmp_size = compressor.Compress(tmp.data(), tmp.size(), image);
tmp.resize(tmp_size);
std::vector<int32_t> output;
REQUIRE_NOTHROW(JFJochDecompress(output, x.GetCompressionAlgorithm(), tmp, x.GetPixelsNum()));
REQUIRE(output.size() == x.GetPixelsNum());
REQUIRE(memcmp(image.data(), output.data(), x.GetPixelsNum() * sizeof(int32_t)) == 0);
}
TEST_CASE("JFJochCompressor_DestTooSmall_Throws","[ZSTD]") {
DiffractionExperiment x(DetJF4M());
x.Compression(CompressionAlgorithm::BSHUF_ZSTD).BitDepthImage(32).PixelSigned(true);
std::vector<int32_t> image(x.GetPixelsNum(), 345);
JFJochBitShuffleCompressor compressor(x.GetCompressionAlgorithm());
// A destination far too small for the compressed output must throw, not overflow it.
std::vector<char> tiny(64);
REQUIRE_THROWS_AS(compressor.Compress(tiny.data(), tiny.size(), image),
CompressionBufferTooSmallException);
// A buffer sized to the worst case never throws.
std::vector<char> big(MaxCompressedSize(x.GetCompressionAlgorithm(), x.GetPixelsNum(), sizeof(int32_t)));
REQUIRE_NOTHROW(compressor.Compress(big.data(), big.size(), image));
}
TEST_CASE("JFJochCompressor_JFJochDecompressor_LZ4","[ZSTD]") {
DiffractionExperiment x(DetJF4M());
x.Compression(CompressionAlgorithm::BSHUF_LZ4).BitDepthImage(32).PixelSigned(true);
std::vector<int32_t> image(x.GetPixelsNum());
for (auto &i: image)
i = 5678;
JFJochBitShuffleCompressor compressor(x.GetCompressionAlgorithm());
std::vector<char> tmp(x.GetPixelsNum() * sizeof(int32_t) * 4 + 12);
auto tmp_size = compressor.Compress(tmp.data(), tmp.size(), image);
tmp.resize(tmp_size);
std::vector<int32_t> output;
REQUIRE_NOTHROW(JFJochDecompress(output, x.GetCompressionAlgorithm(), tmp, x.GetPixelsNum()));
REQUIRE(output.size() == x.GetPixelsNum());
REQUIRE(memcmp(image.data(), output.data(), x.GetPixelsNum() * sizeof(int32_t)) == 0);
}
TEST_CASE("JFJochDecompressor_None","[ZSTD]") {
DiffractionExperiment x(DetJF4M());
x.Compression(CompressionAlgorithm::NO_COMPRESSION).BitDepthImage(32).PixelSigned(true);
std::vector<int32_t> image(x.GetPixelsNum());
for (auto &i: image)
i = 578;
std::vector<int32_t> output;
REQUIRE_NOTHROW(JFJochDecompress(output, x.GetCompressionAlgorithm(), image, x.GetPixelsNum()));
REQUIRE(output.size() == x.GetPixelsNum());
REQUIRE(memcmp(image.data(), output.data(), x.GetPixelsNum() * sizeof(int32_t)) == 0);
}
TEST_CASE("Bitshuffle_ZSTD","[ZSTD]") {
std::vector<int32_t> image(RAW_MODULE_SIZE * sizeof(int32_t));
std::vector<char> compressed(bshuf_compress_zstd_bound(RAW_MODULE_SIZE, 4, 0));
std::vector<int32_t> decompressed(RAW_MODULE_SIZE * sizeof(int32_t));
for (int i = 0; i < RAW_MODULE_SIZE; i++)
image[i] = i;
auto out_size = bshuf_compress_zstd(image.data(), compressed.data(), RAW_MODULE_SIZE, 4, 0,0);
REQUIRE(out_size > 0);
REQUIRE(bshuf_decompress_zstd(compressed.data(), decompressed.data(), RAW_MODULE_SIZE, 4, 0) == out_size);
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<char> 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<char>(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
// confirm the algorithm always produces valid output, even when the ratio is poor.
template<class T>
static void RequireHuffRoundTrip(const std::vector<T> &image) {
JFJochBitShuffleCompressor compressor(CompressionAlgorithm::BSHUF_ZSTD_RLE_HUFF);
auto compressed = compressor.Compress(image);
REQUIRE(compressed.size() <= (size_t) MaxCompressedSize(CompressionAlgorithm::BSHUF_ZSTD_RLE_HUFF,
image.size(), sizeof(T)));
std::vector<T> output;
REQUIRE_NOTHROW(JFJochDecompress(output, CompressionAlgorithm::BSHUF_ZSTD_RLE_HUFF,
compressed, image.size()));
REQUIRE(output.size() == image.size());
REQUIRE(memcmp(image.data(), output.data(), image.size() * sizeof(T)) == 0);
}
TEST_CASE("ZstdHuff_PureZeros", "[ZSTD]") {
std::vector<int16_t> image(200000, 0); // all runs, no literals
RequireHuffRoundTrip(image);
}
TEST_CASE("ZstdHuff_Poisson10", "[ZSTD]") {
std::vector<uint16_t> image(200000);
std::mt19937 gen(12345);
std::poisson_distribution<int> dist(10);
for (auto &v : image) v = (uint16_t) dist(gen);
RequireHuffRoundTrip(image);
}
TEST_CASE("ZstdHuff_RandomIncompressible", "[ZSTD]") {
std::vector<uint16_t> image(200000); // Mersenne-Twister noise: ~zero compression
std::mt19937 gen(98765);
for (auto &v : image) v = (uint16_t) gen();
RequireHuffRoundTrip(image); // must still be valid and within bound
}
TEST_CASE("ZstdHuff_MaskLike", "[ZSTD]") {
std::vector<uint32_t> image(200000, 0); // extreme sparsity like a pixel_mask
std::mt19937 gen(555);
for (auto &v : image) {
uint32_t r = gen() % 100;
if (r < 5) v = 1;
else if (r == 5) v = 0x80000000u; // module-gap style flag
}
RequireHuffRoundTrip(image);
}
TEST_CASE("ZstdHuff_BenchmarkImage", "[ZSTD]") {
RegisterHDF5Filter(); // bitshuffle filter, needed to read the compressed benchmark dataset
HDF5ReadOnlyFile data("../../tests/test_data/compression_benchmark.h5");
HDF5DataSet dataset(data, "/entry/data/data");
HDF5DataSpace file_space(dataset);
auto dims = file_space.GetDimensions();
std::vector<int16_t> image(dims[1] * dims[2]);
std::vector<hsize_t> start = {0, 0, 0}, size = {1, dims[1], dims[2]};
dataset.ReadVector(image, start, size);
RequireHuffRoundTrip(image);
}
// On-the-record proof that a Compressed_Block with Huffman literals and Number_of_Sequences = 0 is
// a valid Zstandard block that outputs its literals. Per zstd_compression_format.md: "if
// (Number_of_Sequences == 0) ... Block's decompressed content is defined solely by the Literals
// Section content." We build exactly such a block and confirm stock ZSTD_decompress returns it.
TEST_CASE("ZstdHuff_LiteralsOnly_NbSeqZero", "[ZSTD]") {
std::vector<uint8_t> lits(4096);
std::mt19937 gen(7);
for (auto &b : lits) b = (uint8_t)(gen() % 6 + 1); // 6-symbol alphabet: Huffman-compressible, no 0x00/0xFF runs
JFJochZstdHuffCompressor huff;
std::vector<uint8_t> frame(lits.size() * 2 + 4096);
size_t fsz = huff.Compress(frame.data(), (const uint64_t *) lits.data(), lits.size());
// Frame = Magic(4) + Frame_Header_Descriptor(1) + Frame_Content_Size(4) = 9 bytes, then one Block.
// Block_Header (3 bytes, little-endian): [ Last_Block:1 | Block_Type:2 | Block_Size:21 ].
uint32_t bh = frame[9] | (frame[10] << 8) | (frame[11] << 16);
bool last_block = bh & 1;
uint32_t block_type = (bh >> 1) & 3;
uint32_t block_size = bh >> 3;
REQUIRE(last_block); // a single block for this all-literal input
REQUIRE(block_type == 2); // Compressed_Block (Huffman literals path, not Raw/RLE)
REQUIRE(frame[12 + block_size - 1] == 0x00); // its Sequences_Section is Number_of_Sequences = 0
std::vector<uint8_t> out(lits.size());
size_t r = ZSTD_decompress(out.data(), out.size(), frame.data(), fsz);
REQUIRE_FALSE(ZSTD_isError(r));
REQUIRE(r == lits.size());
REQUIRE(memcmp(lits.data(), out.data(), lits.size()) == 0);
}