diff --git a/benchmark/micro/micro_bench.cpp b/benchmark/micro/micro_bench.cpp index 7e4e64702..8a4ca05b4 100644 --- a/benchmark/micro/micro_bench.cpp +++ b/benchmark/micro/micro_bench.cpp @@ -210,6 +210,7 @@ void registerMicroBenchmarks() registerFeatureGenIntegerBench(1 << 20); registerFeatureGenIntegerBench(1 << 20); registerFeatureGenIntegerBench(1 << 20); + registerPivCoHuffmanBenchmarks(); } } // namespace zstrong::bench::micro diff --git a/benchmark/micro/micro_bench.h b/benchmark/micro/micro_bench.h index d06a4edec..e6c6e1bb2 100644 --- a/benchmark/micro/micro_bench.h +++ b/benchmark/micro/micro_bench.h @@ -61,5 +61,6 @@ class MiscMicroBenchmarkTestcase : public BenchmarkTestcase { * Generates and registers all micro benchmark cases. */ void registerMicroBenchmarks(); +void registerPivCoHuffmanBenchmarks(); } // namespace zstrong::bench::micro diff --git a/benchmark/micro/micro_pivco_huffman.cpp b/benchmark/micro/micro_pivco_huffman.cpp new file mode 100644 index 000000000..36a270497 --- /dev/null +++ b/benchmark/micro/micro_pivco_huffman.cpp @@ -0,0 +1,681 @@ +// Copyright (c) Meta Platforms, Inc. and affiliates. + +#include "benchmark/micro/micro_bench.h" + +#include +#include +#include +#include +#include + +#include "benchmark/benchmark_config.h" +#include "openzl/codecs/pivco_huffman/arch/decode_pivco_arch.h" +#include "openzl/codecs/pivco_huffman/arch/encode_pivco_arch.h" +#include "openzl/common/assertion.h" + +namespace zstrong::bench::micro { +namespace { + +struct EncodeArch { + std::string name; + const ZL_PivCoHuffmanEncode* kernels; +}; + +struct DecodeArch { + std::string name; + const ZL_PivCoHuffmanDecode* kernels; +}; + +std::vector supportedEncodeArchs() +{ + ZL_cpuid_t const cpuid = ZL_cpuid(); + std::vector out; + std::vector const archs = { + { "generic", &ZL_PivCoHuffmanEncode_generic }, + }; + for (auto const& arch : archs) { + if (arch.kernels->supported(&cpuid)) { + out.push_back(arch); + } + } + return out; +} + +std::vector supportedDecodeArchs() +{ + ZL_cpuid_t const cpuid = ZL_cpuid(); + std::vector out; + std::vector const archs = { + { "generic", &ZL_PivCoHuffmanDecode_generic }, + }; + for (auto const& arch : archs) { + if (arch.kernels->supported(&cpuid)) { + out.push_back(arch); + } + } + return out; +} + +const EncodeArch& defaultEncodeArch(const std::vector& archs) +{ + // Match the arch the engine would pick at runtime. + const ZL_PivCoHuffmanEncode* const selected = + ZL_PivCoHuffmanEncode_select(nullptr); + for (auto const& arch : archs) { + if (arch.kernels == selected) { + return arch; + } + } + ZL_ABORT(); +} + +const DecodeArch& defaultDecodeArch(const std::vector& archs) +{ + // Match the arch the engine would pick at runtime. + const ZL_PivCoHuffmanDecode* const selected = + ZL_PivCoHuffmanDecode_select(nullptr); + for (auto const& arch : archs) { + if (arch.kernels == selected) { + return arch; + } + } + ZL_ABORT(); +} + +size_t bitmapBytes(size_t bits) +{ + return (bits + 7) / 8; +} + +template +void requireEqualPrefix( + const std::vector& lhs, + const std::vector& rhs, + size_t size) +{ + ZL_REQUIRE_GE(lhs.size(), size); + ZL_REQUIRE_GE(rhs.size(), size); + for (size_t i = 0; i < size; ++i) { + ZL_REQUIRE_EQ(lhs[i], rhs[i]); + } +} + +struct PartitionData { + explicit PartitionData(size_t n) + : ranks(n + ZL_PIVCO_HUFFMAN_SLOP), + bitmap(bitmapBytes(n) + ZL_PIVCO_HUFFMAN_SLOP), + lhs(n + ZL_PIVCO_HUFFMAN_SLOP), + rhs(n + ZL_PIVCO_HUFFMAN_SLOP), + size(n) + { + // Ranks cycle through the whole byte range so roughly half land on each + // side of rightRank. + for (size_t i = 0; i < n; ++i) { + ranks[i] = (uint8_t)(((i * 257) + 0x1234) & 0xFF); + } + } + + std::vector ranks; + std::vector bitmap; + std::vector lhs; + std::vector rhs; + uint8_t rightRank = 0x80; + size_t size; +}; + +struct FlatData { + FlatData(size_t n, size_t d) + : ranks(n + ZL_PIVCO_HUFFMAN_SLOP), + bitmap(bitmapBytes(n * d) + ZL_PIVCO_HUFFMAN_SLOP), + symbols((size_t)1 << d), + out(n + ZL_PIVCO_HUFFMAN_SLOP), + size(n), + depth(d), + // At depth 8 the index spans the whole byte range, so rankBegin + // must be 0 to keep rankBegin + index within a byte. + rankBegin(d == 8 ? 0 : 17) + { + // Each rank's offset from rankBegin must fit in `depth` bits, so keep + // ranks within [rankBegin, rankBegin + 2^depth). + uint8_t const indexMask = (uint8_t)(((size_t)1 << d) - 1); + for (size_t i = 0; i < n; ++i) { + ranks[i] = (uint8_t)(rankBegin + (uint8_t)(i & indexMask)); + } + for (size_t i = 0; i < symbols.size(); ++i) { + symbols[i] = (uint8_t)((i * 37 + d * 11) & 0xFF); + } + ZL_PivCoHuffmanEncode_generic.packFlatDepth( + bitmap.data(), depth, ranks.data(), size, rankBegin); + } + + std::vector ranks; + std::vector bitmap; + std::vector symbols; + std::vector out; + size_t size; + size_t depth; + uint8_t rankBegin; +}; + +struct MergeData { + explicit MergeData(size_t n) + : bitmap(bitmapBytes(n) + ZL_PIVCO_HUFFMAN_SLOP), + out(n + ZL_PIVCO_HUFFMAN_SLOP), + totalSize(n) + { + for (size_t i = 0; i < n; ++i) { + bool const bit = ((i * 11 + n) % 5) < 2; + if (bit) { + bitmap[i / 8] |= (uint8_t)(1u << (i & 7)); + rhs.push_back((uint8_t)(0x80u + ((rhs.size() * 13) & 0x7F))); + } else { + lhs.push_back((uint8_t)(0x10u + ((lhs.size() * 17) & 0x6F))); + } + } + lhsSize = lhs.size(); + rhsSize = rhs.size(); + lhs.resize(lhsSize + ZL_PIVCO_HUFFMAN_SLOP, 0xEE); + rhs.resize(rhsSize + ZL_PIVCO_HUFFMAN_SLOP, 0xDD); + } + + std::vector bitmap; + std::vector lhs; + std::vector rhs; + std::vector out; + size_t lhsSize = 0; + size_t rhsSize = 0; + size_t totalSize; +}; + +void verifyPartitionFull(const EncodeArch& arch, const PartitionData& data) +{ + std::vector expectedBitmap(data.bitmap.size()); + std::vector expectedLhs(data.lhs.size()); + std::vector expectedRhs(data.rhs.size()); + size_t const expectedOnes = ZL_PivCoHuffmanEncode_generic.partitionFull( + expectedBitmap.data(), + expectedLhs.data(), + expectedRhs.data(), + data.ranks.data(), + data.size, + data.rightRank); + + std::vector bitmap(data.bitmap.size()); + std::vector lhs(data.lhs.size()); + std::vector rhs(data.rhs.size()); + size_t const ones = arch.kernels->partitionFull( + bitmap.data(), + lhs.data(), + rhs.data(), + data.ranks.data(), + data.size, + data.rightRank); + ZL_REQUIRE_EQ(ones, expectedOnes); + requireEqualPrefix(bitmap, expectedBitmap, bitmapBytes(data.size)); + requireEqualPrefix(lhs, expectedLhs, data.size - expectedOnes); + requireEqualPrefix(rhs, expectedRhs, expectedOnes); +} + +void verifyPartitionLeft(const EncodeArch& arch, const PartitionData& data) +{ + std::vector expectedBitmap(data.bitmap.size()); + std::vector expectedLhs(data.lhs.size()); + size_t const expectedOnes = ZL_PivCoHuffmanEncode_generic.partitionLeft( + expectedBitmap.data(), + expectedLhs.data(), + data.ranks.data(), + data.size, + data.rightRank); + + std::vector bitmap(data.bitmap.size()); + std::vector lhs(data.lhs.size()); + size_t const ones = arch.kernels->partitionLeft( + bitmap.data(), + lhs.data(), + data.ranks.data(), + data.size, + data.rightRank); + ZL_REQUIRE_EQ(ones, expectedOnes); + requireEqualPrefix(bitmap, expectedBitmap, bitmapBytes(data.size)); + requireEqualPrefix(lhs, expectedLhs, data.size - expectedOnes); +} + +void verifyPartitionRight(const EncodeArch& arch, const PartitionData& data) +{ + std::vector expectedBitmap(data.bitmap.size()); + std::vector expectedRhs(data.rhs.size()); + size_t const expectedOnes = ZL_PivCoHuffmanEncode_generic.partitionRight( + expectedBitmap.data(), + expectedRhs.data(), + data.ranks.data(), + data.size, + data.rightRank); + + std::vector bitmap(data.bitmap.size()); + std::vector rhs(data.rhs.size()); + size_t const ones = arch.kernels->partitionRight( + bitmap.data(), + rhs.data(), + data.ranks.data(), + data.size, + data.rightRank); + ZL_REQUIRE_EQ(ones, expectedOnes); + requireEqualPrefix(bitmap, expectedBitmap, bitmapBytes(data.size)); + requireEqualPrefix(rhs, expectedRhs, expectedOnes); +} + +void verifyPartitionNone(const EncodeArch& arch, const PartitionData& data) +{ + std::vector expectedBitmap(data.bitmap.size()); + ZL_PivCoHuffmanEncode_generic.partitionNone( + expectedBitmap.data(), + data.ranks.data(), + data.size, + data.rightRank); + + std::vector bitmap(data.bitmap.size()); + arch.kernels->partitionNone( + bitmap.data(), data.ranks.data(), data.size, data.rightRank); + requireEqualPrefix(bitmap, expectedBitmap, bitmapBytes(data.size)); +} + +void verifyPackFlatDepth(const EncodeArch& arch, const FlatData& data) +{ + std::vector expectedBitmap(data.bitmap.size()); + ZL_PivCoHuffmanEncode_generic.packFlatDepth( + expectedBitmap.data(), + data.depth, + data.ranks.data(), + data.size, + data.rankBegin); + + std::vector bitmap(data.bitmap.size()); + arch.kernels->packFlatDepth( + bitmap.data(), + data.depth, + data.ranks.data(), + data.size, + data.rankBegin); + requireEqualPrefix( + bitmap, expectedBitmap, bitmapBytes(data.size * data.depth)); +} + +void verifyMergeVectorVector(const DecodeArch& arch, const MergeData& data) +{ + std::vector expectedOut(data.out.size()); + size_t const expectedOnes = ZL_PivCoHuffmanDecode_generic.mergeVectorVector( + expectedOut.data(), + expectedOut.size(), + data.bitmap.data(), + data.bitmap.size(), + data.lhs.data(), + data.lhsSize, + data.rhs.data(), + data.rhsSize); + + std::vector out(data.out.size()); + size_t const ones = arch.kernels->mergeVectorVector( + out.data(), + out.size(), + data.bitmap.data(), + data.bitmap.size(), + data.lhs.data(), + data.lhsSize, + data.rhs.data(), + data.rhsSize); + ZL_REQUIRE_EQ(ones, expectedOnes); + requireEqualPrefix(out, expectedOut, data.totalSize); +} + +void verifyMergeConstantVector(const DecodeArch& arch, const MergeData& data) +{ + std::vector expectedOut(data.out.size()); + uint8_t const lhs = 0x3C; + size_t const expectedOnes = + ZL_PivCoHuffmanDecode_generic.mergeConstantVector( + expectedOut.data(), + expectedOut.size(), + data.bitmap.data(), + data.bitmap.size(), + lhs, + data.lhsSize, + data.rhs.data(), + data.rhsSize); + + std::vector out(data.out.size()); + size_t const ones = arch.kernels->mergeConstantVector( + out.data(), + out.size(), + data.bitmap.data(), + data.bitmap.size(), + lhs, + data.lhsSize, + data.rhs.data(), + data.rhsSize); + ZL_REQUIRE_EQ(ones, expectedOnes); + requireEqualPrefix(out, expectedOut, data.totalSize); +} + +void verifyMergeVectorConstant(const DecodeArch& arch, const MergeData& data) +{ + std::vector expectedOut(data.out.size()); + uint8_t const rhs = 0xC3; + size_t const expectedOnes = + ZL_PivCoHuffmanDecode_generic.mergeVectorConstant( + expectedOut.data(), + expectedOut.size(), + data.bitmap.data(), + data.bitmap.size(), + data.lhs.data(), + data.lhsSize, + rhs, + data.rhsSize); + + std::vector out(data.out.size()); + size_t const ones = arch.kernels->mergeVectorConstant( + out.data(), + out.size(), + data.bitmap.data(), + data.bitmap.size(), + data.lhs.data(), + data.lhsSize, + rhs, + data.rhsSize); + ZL_REQUIRE_EQ(ones, expectedOnes); + requireEqualPrefix(out, expectedOut, data.totalSize); +} + +void verifyMergeFlatDepth(const DecodeArch& arch, const FlatData& data) +{ + std::vector expectedOut(data.out.size()); + ZL_PivCoHuffmanDecode_generic.mergeFlatDepth( + expectedOut.data(), + data.size, + expectedOut.size(), + data.bitmap.data(), + data.bitmap.size(), + data.depth, + data.symbols.data()); + + std::vector out(data.out.size()); + arch.kernels->mergeFlatDepth( + out.data(), + data.size, + out.size(), + data.bitmap.data(), + data.bitmap.size(), + data.depth, + data.symbols.data()); + requireEqualPrefix(out, expectedOut, data.size); +} + +std::string benchName(const std::string& arch, const std::string& kernel) +{ + return "MCR / PivCoHuffman / " + arch + " / " + kernel; +} + +void registerPartitionFullBenchmark(EncodeArch arch) +{ + auto data = std::make_shared(64 * 1024); + verifyPartitionFull(arch, *data); + RegisterBenchmark( + benchName(arch.name, "partitionFull"), + [arch, data](benchmark::State& state) { + size_t ones = 0; + for (auto _ : state) { + ones += arch.kernels->partitionFull( + data->bitmap.data(), + data->lhs.data(), + data->rhs.data(), + data->ranks.data(), + data->size, + data->rightRank); + } + benchmark::DoNotOptimize(ones); + state.SetBytesProcessed( + (int64_t)data->size * (int64_t)state.iterations()); + }); +} + +void registerPartitionLeftBenchmark(EncodeArch arch) +{ + auto data = std::make_shared(64 * 1024); + verifyPartitionLeft(arch, *data); + RegisterBenchmark( + benchName(arch.name, "partitionLeft"), + [arch, data](benchmark::State& state) { + size_t ones = 0; + for (auto _ : state) { + ones += arch.kernels->partitionLeft( + data->bitmap.data(), + data->lhs.data(), + data->ranks.data(), + data->size, + data->rightRank); + } + benchmark::DoNotOptimize(ones); + state.SetBytesProcessed( + (int64_t)data->size * (int64_t)state.iterations()); + }); +} + +void registerPartitionRightBenchmark(EncodeArch arch) +{ + auto data = std::make_shared(64 * 1024); + verifyPartitionRight(arch, *data); + RegisterBenchmark( + benchName(arch.name, "partitionRight"), + [arch, data](benchmark::State& state) { + size_t ones = 0; + for (auto _ : state) { + ones += arch.kernels->partitionRight( + data->bitmap.data(), + data->rhs.data(), + data->ranks.data(), + data->size, + data->rightRank); + } + benchmark::DoNotOptimize(ones); + state.SetBytesProcessed( + (int64_t)data->size * (int64_t)state.iterations()); + }); +} + +void registerPartitionNoneBenchmark(EncodeArch arch) +{ + auto data = std::make_shared(64 * 1024); + verifyPartitionNone(arch, *data); + RegisterBenchmark( + benchName(arch.name, "partitionNone"), + [arch, data](benchmark::State& state) { + for (auto _ : state) { + arch.kernels->partitionNone( + data->bitmap.data(), + data->ranks.data(), + data->size, + data->rightRank); + } + benchmark::DoNotOptimize(data->bitmap.data()); + state.SetBytesProcessed( + (int64_t)data->size * (int64_t)state.iterations()); + }); +} + +void registerPackFlatDepthBenchmark(EncodeArch arch, size_t depth) +{ + auto data = std::make_shared(64 * 1024, depth); + verifyPackFlatDepth(arch, *data); + RegisterBenchmark( + benchName(arch.name, "packFlatDepth(depth=" + std::to_string(depth)) + + ")", + [arch, data](benchmark::State& state) { + for (auto _ : state) { + arch.kernels->packFlatDepth( + data->bitmap.data(), + data->depth, + data->ranks.data(), + data->size, + data->rankBegin); + } + benchmark::DoNotOptimize(data->bitmap.data()); + state.SetBytesProcessed( + (int64_t)data->size * (int64_t)state.iterations()); + }); +} + +void registerMergeVectorVectorBenchmark(DecodeArch arch) +{ + auto data = std::make_shared(64 * 1024); + verifyMergeVectorVector(arch, *data); + RegisterBenchmark( + benchName(arch.name, "mergeVectorVector"), + [arch, data](benchmark::State& state) { + size_t ones = 0; + for (auto _ : state) { + ones += arch.kernels->mergeVectorVector( + data->out.data(), + data->out.size(), + data->bitmap.data(), + data->bitmap.size(), + data->lhs.data(), + data->lhsSize, + data->rhs.data(), + data->rhsSize); + } + benchmark::DoNotOptimize(ones); + state.SetBytesProcessed( + (int64_t)data->totalSize * (int64_t)state.iterations()); + }); +} + +void registerMergeConstantVectorBenchmark(DecodeArch arch) +{ + auto data = std::make_shared(64 * 1024); + verifyMergeConstantVector(arch, *data); + RegisterBenchmark( + benchName(arch.name, "mergeConstantVector"), + [arch, data](benchmark::State& state) { + size_t ones = 0; + uint8_t const lhs = 0x3C; + for (auto _ : state) { + ones += arch.kernels->mergeConstantVector( + data->out.data(), + data->out.size(), + data->bitmap.data(), + data->bitmap.size(), + lhs, + data->lhsSize, + data->rhs.data(), + data->rhsSize); + } + benchmark::DoNotOptimize(ones); + state.SetBytesProcessed( + (int64_t)data->totalSize * (int64_t)state.iterations()); + }); +} + +void registerMergeVectorConstantBenchmark(DecodeArch arch) +{ + auto data = std::make_shared(64 * 1024); + verifyMergeVectorConstant(arch, *data); + RegisterBenchmark( + benchName(arch.name, "mergeVectorConstant"), + [arch, data](benchmark::State& state) { + size_t ones = 0; + uint8_t const rhs = 0xC3; + for (auto _ : state) { + ones += arch.kernels->mergeVectorConstant( + data->out.data(), + data->out.size(), + data->bitmap.data(), + data->bitmap.size(), + data->lhs.data(), + data->lhsSize, + rhs, + data->rhsSize); + } + benchmark::DoNotOptimize(ones); + state.SetBytesProcessed( + (int64_t)data->totalSize * (int64_t)state.iterations()); + }); +} + +void registerMergeFlatDepthBenchmark(DecodeArch arch, size_t depth) +{ + auto data = std::make_shared(64 * 1024, depth); + verifyMergeFlatDepth(arch, *data); + RegisterBenchmark( + benchName( + arch.name, "mergeFlatDepth(depth=" + std::to_string(depth)) + + ")", + [arch, data](benchmark::State& state) { + for (auto _ : state) { + arch.kernels->mergeFlatDepth( + data->out.data(), + data->size, + data->out.size(), + data->bitmap.data(), + data->bitmap.size(), + data->depth, + data->symbols.data()); + } + benchmark::DoNotOptimize(data->out.data()); + state.SetBytesProcessed( + (int64_t)data->size * (int64_t)state.iterations()); + }); +} + +void registerEncodeBenchmarks(EncodeArch arch) +{ + ZL_REQUIRE_NN(arch.kernels->partitionFull); + ZL_REQUIRE_NN(arch.kernels->partitionLeft); + ZL_REQUIRE_NN(arch.kernels->partitionRight); + ZL_REQUIRE_NN(arch.kernels->partitionNone); + ZL_REQUIRE_NN(arch.kernels->packFlatDepth); + registerPartitionFullBenchmark(arch); + registerPartitionLeftBenchmark(arch); + registerPartitionRightBenchmark(arch); + registerPartitionNoneBenchmark(arch); + for (size_t depth = 1; depth <= 8; ++depth) { + registerPackFlatDepthBenchmark(arch, depth); + } +} + +void registerDecodeBenchmarks(DecodeArch arch) +{ + ZL_REQUIRE_NN(arch.kernels->mergeVectorVector); + ZL_REQUIRE_NN(arch.kernels->mergeConstantVector); + ZL_REQUIRE_NN(arch.kernels->mergeVectorConstant); + ZL_REQUIRE_NN(arch.kernels->mergeFlatDepth); + registerMergeVectorVectorBenchmark(arch); + registerMergeConstantVectorBenchmark(arch); + registerMergeVectorConstantBenchmark(arch); + for (size_t depth = 1; depth <= 8; ++depth) { + registerMergeFlatDepthBenchmark(arch, depth); + } +} + +} // namespace + +void registerPivCoHuffmanBenchmarks() +{ + auto const encodeArchs = supportedEncodeArchs(); + auto const decodeArchs = supportedDecodeArchs(); + + for (auto const& arch : encodeArchs) { + registerEncodeBenchmarks(arch); + } + for (auto const& arch : decodeArchs) { + registerDecodeBenchmarks(arch); + } + + EncodeArch defaultEncode = defaultEncodeArch(encodeArchs); + defaultEncode.name = "default"; + registerEncodeBenchmarks(defaultEncode); + + DecodeArch defaultDecode = defaultDecodeArch(decodeArchs); + defaultDecode.name = "default"; + registerDecodeBenchmarks(defaultDecode); +} + +} // namespace zstrong::bench::micro diff --git a/src/openzl/codecs/common/bitstream/ff_bitstream.h b/src/openzl/codecs/common/bitstream/ff_bitstream.h index b1d7202ee..31bfa72bd 100644 --- a/src/openzl/codecs/common/bitstream/ff_bitstream.h +++ b/src/openzl/codecs/common/bitstream/ff_bitstream.h @@ -31,6 +31,10 @@ ZL_INLINE ZL_Report ZS_BitCStreamFF_finish(ZS_BitCStreamFF* bits); ZL_INLINE void ZS_BitCStreamFF_write(ZS_BitCStreamFF* bits, size_t value, size_t nbBits); ZL_INLINE void ZS_BitCStreamFF_flush(ZS_BitCStreamFF* bits); +ZL_INLINE uint8_t* ZS_BitCStreamFF_reserveAlignedBits( + ZS_BitCStreamFF* bits, + size_t nbBits); +ZL_INLINE void ZS_BitCStreamFF_commitReservedBits(ZS_BitCStreamFF* bits); ZL_INLINE void ZS_BitCStreamFF_writeExpGolomb( ZS_BitCStreamFF* bits, @@ -52,6 +56,7 @@ typedef struct { uint8_t const* ptr; uint8_t const* limit; uint8_t const* end; + uint8_t const* begin; } ZS_BitDStreamFF; ZL_INLINE ZS_BitDStreamFF @@ -62,6 +67,9 @@ ZL_INLINE size_t ZS_BitDStreamFF_peek(ZS_BitDStreamFF const* bits, size_t nbBits); ZL_INLINE void ZS_BitDStreamFF_skip(ZS_BitDStreamFF* bits, size_t nbBits); ZL_INLINE void ZS_BitDStreamFF_reload(ZS_BitDStreamFF* bits); +ZL_INLINE uint8_t const* ZS_BitDStreamFF_popAlignedBits( + ZS_BitDStreamFF* bits, + size_t nbBits); ZL_INLINE uint32_t ZS_BitDStreamFF_readExpGolomb(ZS_BitDStreamFF* bits, size_t order) @@ -80,11 +88,13 @@ ZS_BitDStreamFF_readExpGolomb(ZS_BitDStreamFF* bits, size_t order) ZL_INLINE ZS_BitCStreamFF ZS_BitCStreamFF_init(uint8_t* dst, size_t dstCapacity) { + const size_t limit = + dstCapacity < sizeof(size_t) ? 0 : dstCapacity - sizeof(size_t) + 1; return (ZS_BitCStreamFF){ .container = 0, .nbBits = 0, .ptr = dst, - .limit = dst + dstCapacity - sizeof(size_t), + .limit = dst + limit, .end = dst + dstCapacity, .begin = dst, }; @@ -94,7 +104,7 @@ ZL_INLINE ZL_Report ZS_BitCStreamFF_finish(ZS_BitCStreamFF* bits) { ZL_RESULT_DECLARE_SCOPE_REPORT((ZL_OperationContext*)NULL); size_t bytesToWrite = (bits->nbBits + 7) / 8; - if (bits->end < bits->ptr + bytesToWrite) + if ((size_t)(bits->end - bits->ptr) < bytesToWrite) ZL_ERR(internalBuffer_tooSmall); if (bytesToWrite) { ZL_ASSERT_EQ( @@ -116,7 +126,7 @@ ZS_BitCStreamFF_write(ZS_BitCStreamFF* bits, size_t value, size_t nbBits) ZL_INLINE void ZS_BitCStreamFF_flush(ZS_BitCStreamFF* bits) { - if (bits->ptr > bits->limit) { + if (bits->ptr >= bits->limit) { return; } size_t const nbBytes = bits->nbBits >> 3; @@ -126,6 +136,62 @@ ZL_INLINE void ZS_BitCStreamFF_flush(ZS_BitCStreamFF* bits) bits->container >>= (nbBytes << 3); } +ZL_INLINE uint8_t* ZS_BitCStreamFF_reserveAlignedBits( + ZS_BitCStreamFF* bits, + size_t nbBits) +{ + if (nbBits > SIZE_MAX - 7) { + return NULL; + } + + size_t const bufferedBytes = (bits->nbBits + 7) / 8; + if (bufferedBytes > (size_t)(bits->end - bits->ptr)) { + return NULL; + } + + size_t const nbBytes = (nbBits + 7) / 8; + size_t const available = (size_t)(bits->end - bits->ptr) - bufferedBytes; + if (nbBytes > available) { + return NULL; + } + + ZL_ASSERT_LE(bufferedBytes, sizeof(size_t)); + if (bufferedBytes != 0) { + ZL_writeLE64_N(bits->ptr, bits->container, bufferedBytes); + bits->ptr += bufferedBytes; + bits->container = 0; + bits->nbBits = 0; + } + ZL_ASSERT_EQ(bits->nbBits, 0); + + uint8_t* const out = bits->ptr; + bits->ptr = out + (nbBits / 8); + bits->nbBits = nbBits & 7; + bits->container = 0; + return out; +} + +ZL_INLINE void ZS_BitCStreamFF_commitReservedBits(ZS_BitCStreamFF* bits) +{ + if (bits->nbBits == 0) { + bits->container = 0; + return; + } + bits->container = bits->ptr[0] & (((size_t)1 << bits->nbBits) - 1); +} + +// Little-endian load of the first @p nbBytes (< sizeof(size_t)) bytes at @p src +// into a container word. Used for streams shorter than a full word, where every +// byte lives in the container. +ZL_INLINE size_t ZS_BitDStreamFF_loadPartial(uint8_t const* src, size_t nbBytes) +{ + size_t container = 0; + for (size_t i = 0; i < nbBytes; ++i) { + container |= (size_t)src[i] << (i << 3); + } + return container; +} + ZL_INLINE ZS_BitDStreamFF ZS_BitDStreamFF_init(uint8_t const* src, size_t srcSize) { @@ -136,19 +202,19 @@ ZS_BitDStreamFF_init(uint8_t const* src, size_t srcSize) .ptr = src, .limit = src + srcSize - sizeof(size_t) + 1, .end = src + srcSize, + .begin = src, }; return bits; } else { - ZS_BitDStreamFF bits = { - .container = 0, - .nbBitsRead = (sizeof(size_t) - srcSize) * 8, - .ptr = src + srcSize, - .limit = src, - .end = src + srcSize, + uint8_t const* const end = srcSize > 0 ? src + srcSize : src; + ZS_BitDStreamFF bits = { + .container = ZS_BitDStreamFF_loadPartial(src, srcSize), + .nbBitsRead = (sizeof(size_t) - srcSize) * 8, + .ptr = end, + .limit = src, + .end = end, + .begin = src, }; - for (size_t i = 0; i < srcSize; ++i) { - bits.container |= (size_t)src[i] << (i << 3); - } return bits; } } @@ -159,7 +225,13 @@ ZL_INLINE ZL_Report ZS_BitDStreamFF_finish(ZS_BitDStreamFF const* bits) if (bits->nbBitsRead > ZS_BITSTREAM_READ_MAX_BITS) { ZL_ERR(GENERIC); } - return ZL_returnSuccess(); + size_t bytesRead = (size_t)(bits->ptr - bits->begin); + bytesRead += bits->nbBitsRead >> 3; + bytesRead += (bits->nbBitsRead & 7) != 0; + if ((size_t)(bits->end - bits->begin) < sizeof(size_t)) { + bytesRead -= sizeof(size_t); + } + return ZL_returnValue(bytesRead); } ZL_INLINE size_t ZS_BitDStreamFF_read(ZS_BitDStreamFF* bits, size_t nbBits) @@ -189,22 +261,93 @@ ZL_INLINE void ZS_BitDStreamFF_skip(ZS_BitDStreamFF* bits, size_t nbBits) ZL_INLINE void ZS_BitDStreamFF_reload(ZS_BitDStreamFF* bits) { - bits->ptr += bits->nbBitsRead >> 3; - if (ZL_LIKELY(bits->ptr < bits->limit)) { - size_t const next = ZL_readLEST(bits->ptr); + uint8_t const* const ptr = bits->ptr + (bits->nbBitsRead >> 3); + if (ZL_LIKELY(ptr < bits->limit)) { + bits->ptr = ptr; bits->nbBitsRead &= 7; - bits->container = next >> bits->nbBitsRead; + bits->container = ZL_readLEST(ptr) >> bits->nbBitsRead; return; } - if (bits->ptr >= bits->end) + // Past the end: leave ptr and nbBitsRead untouched. This keeps the consumed + // bit count -- (ptr - begin) * 8 + nbBitsRead -- exact for finish() and + // popAlignedBits(), while an over-read still leaves nbBitsRead too large + // for finish() to accept. + if (ptr >= bits->end) return; - uint8_t const* const limit = bits->limit - 1; - size_t const skippedBits = (size_t)((bits->ptr - limit) << 3); - size_t const next = ZL_readLEST(limit); + uint8_t const* const tail = bits->limit - 1; + size_t const skippedBits = (size_t)((ptr - tail) << 3); + bits->ptr = ptr; bits->nbBitsRead &= 7; - bits->container = next >> (bits->nbBitsRead + skippedBits); + bits->container = ZL_readLEST(tail) >> (bits->nbBitsRead + skippedBits); +} + +ZL_INLINE uint8_t const* ZS_BitDStreamFF_popAlignedBits( + ZS_BitDStreamFF* bits, + size_t nbBits) +{ + if (nbBits > SIZE_MAX - 7) { + return NULL; + } + + // check if the stream is already in an invalid state + if (bits->nbBitsRead > ZS_BITSTREAM_READ_MAX_BITS) { + return NULL; + } + + size_t const streamSize = (size_t)(bits->end - bits->begin); + if (streamSize > SIZE_MAX / 8) { + return NULL; + } + + size_t consumedBits = + (size_t)(bits->ptr - bits->begin) * 8 + bits->nbBitsRead; + if (streamSize < sizeof(size_t)) { + consumedBits -= sizeof(size_t) * 8; + } + + size_t const bitSize = streamSize * 8; + if (consumedBits > bitSize || consumedBits > SIZE_MAX - 7) { + return NULL; + } + + size_t const alignedBits = ((consumedBits + 7) / 8) * 8; + if (nbBits > bitSize - alignedBits) { + return NULL; + } + + uint8_t const* const out = bits->begin + alignedBits / 8; + + // Reposition the stream just past the aligned region we handed out. + size_t const bitPos = alignedBits + nbBits; + uint8_t const* const pos = bits->begin + bitPos / 8; + size_t const bitShift = bitPos & 7; + + if (streamSize >= sizeof(size_t)) { + // Refill the 64-bit window at `pos`, clamping the load back from the + // end when fewer than 8 bytes remain; the matching shift exposes the + // same bit either way. A shift of the full word width (the whole + // stream consumed) leaves an empty window. + uint8_t const* const tail = bits->end - sizeof(size_t); + uint8_t const* const load = pos <= tail ? pos : tail; + size_t const shift = bitPos - (size_t)(load - bits->begin) * 8; + bits->container = + shift < sizeof(size_t) * 8 ? ZL_readLEST(load) >> shift : 0; + bits->nbBitsRead = bitShift; + bits->ptr = pos; + } else { + // Short stream: every byte already lives in the container. + size_t const remaining = (size_t)(bits->end - pos); + bits->container = 0; + for (size_t i = 0; i < remaining; ++i) { + bits->container |= (size_t)pos[i] << (i << 3); + } + bits->container >>= bitShift; + bits->nbBitsRead = (sizeof(size_t) - remaining) * 8 + bitShift; + bits->ptr = bits->end; + } + return out; } ZL_END_C_DECLS diff --git a/src/openzl/codecs/pivco_huffman/README.md b/src/openzl/codecs/pivco_huffman/README.md new file mode 100644 index 000000000..5e625781b --- /dev/null +++ b/src/openzl/codecs/pivco_huffman/README.md @@ -0,0 +1,22 @@ +# PivCo-Huffman + +This is an implementation of PivCo-Huffman coding from Marcin Zukowski. +The code is based on ideas from: + +- [The paper](https://arxiv.org/abs/2606.05765) +- [The repo](https://github.com/MarcinZukowski/pivco-huffman) +- [Ryg's blog](https://fgiesen.wordpress.com/2026/06/21/pivco-huffman-merge-operations/) + +OpenZL re-implements PivCo-Huffman instead of using the upstream repo for several reasons: + +1. OpenZL needs full control over the wire-format. +2. OpenZL needs to be hardened against decoding corrupted data. +3. OpenZL needs to be able to dynamically dispatch to kernels based on the CPUID, +e.g. to detect AVX512 support at runtime. + +The goal is to work with upstream to factor out the primitives so that OpenZL can +import and use upstream's primitives, as those are the most complex part of PivCo, +and are independent of the wire format. + +In the meantime, any meaningful improvements will be upstreamed, so there is a +single reference implementation for PivCo Huffman. diff --git a/src/openzl/codecs/pivco_huffman/arch/common_pivco_arch.h b/src/openzl/codecs/pivco_huffman/arch/common_pivco_arch.h new file mode 100644 index 000000000..48c0d5d17 --- /dev/null +++ b/src/openzl/codecs/pivco_huffman/arch/common_pivco_arch.h @@ -0,0 +1,20 @@ +// Copyright (c) Meta Platforms, Inc. and affiliates. +#ifndef OPENZL_CODECS_PIVCO_HUFFMAN_ARCH_COMMON_PIVCO_ARCH_H +#define OPENZL_CODECS_PIVCO_HUFFMAN_ARCH_COMMON_PIVCO_ARCH_H + +#include "openzl/shared/portability.h" + +ZL_BEGIN_C_DECLS + +/** + * Slop bytes that every bitmap and rank buffer must reserve past its logical + * length. The encode and decode kernels may over-read and over-write up to this + * many trailing bytes where noted in the docs as `SLOP`, so callers pad their + * buffers by this amount. It is sized to the widest fixed-width group the SIMD + * kernels process. + */ +#define ZL_PIVCO_HUFFMAN_SLOP 64 + +ZL_END_C_DECLS + +#endif diff --git a/src/openzl/codecs/pivco_huffman/arch/decode_pivco_arch.c b/src/openzl/codecs/pivco_huffman/arch/decode_pivco_arch.c new file mode 100644 index 000000000..f9e984558 --- /dev/null +++ b/src/openzl/codecs/pivco_huffman/arch/decode_pivco_arch.c @@ -0,0 +1,204 @@ +// Copyright (c) Meta Platforms, Inc. and affiliates. + +#include "openzl/codecs/pivco_huffman/arch/decode_pivco_arch.h" + +#include + +#include "openzl/codecs/common/bitstream/ff_bitstream.h" + +// Portable reference implementation of the PivCo-Huffman decode kernels. +// +// This file is written for clarity, not speed. It is used for correctness +// testing and as the fallback on hardware without a specialized kernel, which +// is not expected to run this code in practice. The architecture-specific +// kernels (x86, arm, avx512) are the fast paths and MUST decode the exact same +// bitmaps as the code here. +// +// Every bitmap is packed little-endian, least-significant-bit first -- exactly +// the layout ZS_BitDStreamFF reads -- so we lean on that primitive instead of +// hand-rolling the bit math. To keep the loops trivially correct we reload the +// 64-bit window after every element. That is wasteful but easy to follow, which +// is the point of this file. + +const ZL_PivCoHuffmanDecode* ZL_PivCoHuffmanDecode_select( + const ZL_cpuid_t* cpuid) +{ + (void)cpuid; + return &ZL_PivCoHuffmanDecode_generic; +} + +ZL_INLINE size_t bitmapBytes(size_t bits) +{ + return (bits + 7) / 8; +} + +// One side of a merge: either a single byte repeated (a constant leaf) or a +// vector of decoded bytes consumed in order. `pos` counts how many have been +// taken so far. +typedef struct { + const uint8_t* vec; // NULL for a constant source + uint8_t constant; + size_t size; + size_t pos; +} MergeSource; + +static MergeSource vectorSource(const uint8_t* vec, size_t size) +{ + return (MergeSource){ .vec = vec, .constant = 0, .size = size, .pos = 0 }; +} + +static MergeSource constantSource(uint8_t value, size_t size) +{ + return (MergeSource){ + .vec = NULL, .constant = value, .size = size, .pos = 0 + }; +} + +static uint8_t mergeSourceNext(MergeSource* source) +{ + uint8_t value; + if (source->vec == NULL) { + value = source->constant; + } else if (source->pos < source->size) { + value = source->vec[source->pos]; + } else { + // More bits selected this source than it has bytes: the bitstream is + // corrupt. Return a placeholder; the caller detects the corruption + // because the returned one-count will not match the expected size. + value = 0; + } + ++source->pos; + return value; +} + +// Reconstructs lhs.size + rhs.size bytes into @p out from a partition bitmap: +// bit i selects the right source (1) or the left source (0) for output position +// i, taking one byte from the chosen source. Returns the number of one bits. +static size_t mergeGeneric( + uint8_t* out, + size_t outCapacity, + const uint8_t* bitmap, + size_t bitmapCapacity, + MergeSource lhs, + MergeSource rhs) +{ + const size_t outSize = lhs.size + rhs.size; + assert(outCapacity >= outSize); + assert(bitmapCapacity >= bitmapBytes(outSize)); + (void)outCapacity; // only read by the assert above + + ZS_BitDStreamFF bits = ZS_BitDStreamFF_init(bitmap, bitmapCapacity); + + size_t ones = 0; + for (size_t i = 0; i < outSize; ++i) { + const size_t bit = ZS_BitDStreamFF_read(&bits, 1); + // Reload every bit so the 64-bit window never underflows. + ZS_BitDStreamFF_reload(&bits); + + if (bit != 0) { + out[i] = mergeSourceNext(&rhs); + ++ones; + } else { + out[i] = mergeSourceNext(&lhs); + } + } + return ones; +} + +static size_t mergeVectorVector( + uint8_t* out, + size_t outCapacity, + const uint8_t* bitmap, + size_t bitmapCapacity, + const uint8_t* lhs, + size_t lhsSize, + const uint8_t* rhs, + size_t rhsSize) +{ + return mergeGeneric( + out, + outCapacity, + bitmap, + bitmapCapacity, + vectorSource(lhs, lhsSize), + vectorSource(rhs, rhsSize)); +} + +static size_t mergeConstantVector( + uint8_t* out, + size_t outCapacity, + const uint8_t* bitmap, + size_t bitmapCapacity, + uint8_t lhs, + size_t lhsSize, + const uint8_t* rhs, + size_t rhsSize) +{ + return mergeGeneric( + out, + outCapacity, + bitmap, + bitmapCapacity, + constantSource(lhs, lhsSize), + vectorSource(rhs, rhsSize)); +} + +static size_t mergeVectorConstant( + uint8_t* out, + size_t outCapacity, + const uint8_t* bitmap, + size_t bitmapCapacity, + const uint8_t* lhs, + size_t lhsSize, + uint8_t rhs, + size_t rhsSize) +{ + return mergeGeneric( + out, + outCapacity, + bitmap, + bitmapCapacity, + vectorSource(lhs, lhsSize), + constantSource(rhs, rhsSize)); +} + +// Expands a flat leaf: reads one @p depth-bit index per output and looks it up +// in @p symbols (which has 2^depth entries). +static void mergeFlatDepth( + uint8_t* out, + size_t outSize, + size_t outCapacity, + const uint8_t* bitmap, + size_t bitmapCapacity, + size_t depth, + const uint8_t* symbols) +{ + assert(depth >= 1); + assert(depth <= 8); + assert(outCapacity >= outSize); + assert(bitmapCapacity >= bitmapBytes(outSize * depth)); + (void)outCapacity; // only read by the assert above + + ZS_BitDStreamFF bits = ZS_BitDStreamFF_init(bitmap, bitmapCapacity); + + for (size_t i = 0; i < outSize; ++i) { + const size_t index = ZS_BitDStreamFF_read(&bits, depth); + // Reload every index so the 64-bit window never underflows. + ZS_BitDStreamFF_reload(&bits); + out[i] = symbols[index]; + } +} + +static bool supported(const ZL_cpuid_t* cpuid) +{ + (void)cpuid; + return true; +} + +const ZL_PivCoHuffmanDecode ZL_PivCoHuffmanDecode_generic = { + .supported = supported, + .mergeVectorVector = mergeVectorVector, + .mergeConstantVector = mergeConstantVector, + .mergeVectorConstant = mergeVectorConstant, + .mergeFlatDepth = mergeFlatDepth, +}; diff --git a/src/openzl/codecs/pivco_huffman/arch/decode_pivco_arch.h b/src/openzl/codecs/pivco_huffman/arch/decode_pivco_arch.h new file mode 100644 index 000000000..e90d6bb9d --- /dev/null +++ b/src/openzl/codecs/pivco_huffman/arch/decode_pivco_arch.h @@ -0,0 +1,131 @@ +// Copyright (c) Meta Platforms, Inc. and affiliates. +#ifndef OPENZL_CODECS_PIVCO_HUFFMAN_ARCH_DECODE_PIVCO_ARCH_H +#define OPENZL_CODECS_PIVCO_HUFFMAN_ARCH_DECODE_PIVCO_ARCH_H + +#include +#include + +#include "openzl/codecs/pivco_huffman/arch/common_pivco_arch.h" +#include "openzl/shared/cpu.h" +#include "openzl/shared/portability.h" + +ZL_BEGIN_C_DECLS + +typedef struct { + /** + * @returns true iff the CPU supports this implementation. + * @note The rest of the function pointers may be NULL if not supported. + */ + bool (*supported)(const ZL_cpuid_t* cpuid); + + /** + * Merges @p lhs and @p rhs into @p out based on the bits in @p bitmap. + * + * @param outCapacity Must be at least `lhsSize + rhsSize` bytes, but may + * be larger to indicate that over-writes are okay. + * @param bitmapCapacity Must be at least `(lhsSize + rhsSize + 7) / 8` + * bytes, but may be larger to indicate that + * over-reads are okay. + * @param lhs Must be at least `lhsSize + SLOP` bytes, + * where the first @p lhsSize bytes are valid. + * @param rhs Must be at least `rhsSize + SLOP` bytes, + * where the first @p rhsSize bytes are valid. + * + * @returns The number of ones in the bitmap. If this is not equal to + * @p rhsSize then the data was corrupt, and the output is unspecified. + */ + size_t (*mergeVectorVector)( + uint8_t* out, + size_t outCapacity, + const uint8_t* bitmap, + size_t bitmapCapacity, + const uint8_t* lhs, + size_t lhsSize, + const uint8_t* rhs, + size_t rhsSize); + + /** + * Merges @p lhs and @p rhs into @p out based on the bits in @p bitmap, + * where @p lhs is a constant. + * + * @param outCapacity Must be at least `lhsSize + rhsSize` bytes, but may + * be larger to indicate that over-writes are okay. + * @param bitmapCapacity Must be at least `(lhsSize + rhsSize + 7) / 8` + * bytes, but may be larger to indicate that + * over-reads are okay. + * @param rhs Must be at least `rhsSize + SLOP` bytes, + * where the first @p rhsSize bytes are valid. + * + * @returns The number of ones in the bitmap. If this is not equal to + * @p rhsSize then the data was corrupt, and the output is unspecified. + */ + size_t (*mergeConstantVector)( + uint8_t* out, + size_t outCapacity, + const uint8_t* bitmap, + size_t bitmapCapacity, + uint8_t lhs, + size_t lhsSize, + const uint8_t* rhs, + size_t rhsSize); + + /** + * Merges @p lhs and @p rhs into @p out based on the bits in @p bitmap, + * where @p rhs is a constant. + * + * @param outCapacity Must be at least `lhsSize + rhsSize` bytes, but may + * be larger to indicate that over-writes are okay. + * @param bitmapCapacity Must be at least `(lhsSize + rhsSize + 7) / 8` + * bytes, but may be larger to indicate that + * over-reads are okay. + * @param lhs Must be at least `lhsSize + SLOP` bytes, + * where the first @p lhsSize bytes are valid. + * + * @returns The number of ones in the bitmap. If this is not equal to + * @p rhsSize then the data was corrupt, and the output is unspecified. + */ + size_t (*mergeVectorConstant)( + uint8_t* out, + size_t outCapacity, + const uint8_t* bitmap, + size_t bitmapCapacity, + const uint8_t* lhs, + size_t lhsSize, + uint8_t rhs, + size_t rhsSize); + + /** + * Reads @p outSize packed indices of @p depth bits each from @p bitmap and + * fills @p out with `symbols[idx]` for each packed index. + * + * @param out Buffer of size @p outCapacity + * @param outSize Number of output symbols + * @param outCapacity Must be at least @p outSize. It can be larger to + * indicate that over-writing is okay + * @param bitmapCapacity Must be at least `(outSize * depth + 7) / 8` bytes, + * but may be larger to indicate that over-reads are + * okay + * @param depth The number of bits each index in @p bitmap takes + * @param symbols Must be @p 2^depth bytes large + */ + void (*mergeFlatDepth)( + uint8_t* out, + size_t outSize, + size_t outCapacity, + const uint8_t* bitmap, + size_t bitmapCapacity, + size_t depth, + const uint8_t* symbols); +} ZL_PivCoHuffmanDecode; + +/** + * @returns The best kernel for the CPU. + */ +const ZL_PivCoHuffmanDecode* ZL_PivCoHuffmanDecode_select( + const ZL_cpuid_t* cpuid); + +extern const ZL_PivCoHuffmanDecode ZL_PivCoHuffmanDecode_generic; + +ZL_END_C_DECLS + +#endif diff --git a/src/openzl/codecs/pivco_huffman/arch/encode_pivco_arch.c b/src/openzl/codecs/pivco_huffman/arch/encode_pivco_arch.c new file mode 100644 index 000000000..46457bef0 --- /dev/null +++ b/src/openzl/codecs/pivco_huffman/arch/encode_pivco_arch.c @@ -0,0 +1,147 @@ +// Copyright (c) Meta Platforms, Inc. and affiliates. + +#include "openzl/codecs/pivco_huffman/arch/encode_pivco_arch.h" + +#include + +#include "openzl/codecs/common/bitstream/ff_bitstream.h" + +// Portable reference implementation of the PivCo-Huffman encode kernels. +// +// This file is written for clarity, not speed. It is used for correctness +// testing and as the fallback on hardware without a specialized kernel, which +// is not expected to run this code in practice. The architecture-specific +// kernels (x86, arm, avx512) are the fast paths and MUST produce byte-identical +// bitmaps to the code here. +// +// Every bitmap is packed little-endian, least-significant-bit first -- exactly +// the layout ZS_BitCStreamFF writes -- so we lean on that primitive instead of +// hand-rolling the bit math. To keep the loops trivially correct we flush the +// 64-bit window after every element. That is wasteful but easy to follow, which +// is the point of this file. + +const ZL_PivCoHuffmanEncode* ZL_PivCoHuffmanEncode_select( + const ZL_cpuid_t* cpuid) +{ + (void)cpuid; + return &ZL_PivCoHuffmanEncode_generic; +} + +static size_t bitmapBytes(size_t bits) +{ + return (bits + 7) / 8; +} + +// Partitions @p ranks into a left group (rank < rightRank) and a right group +// (rank >= rightRank). Emits one bit per rank into @p bitmap (0 = left, +// 1 = right) and, when the corresponding output is non-NULL, copies each rank +// into @p lhs / @p rhs in order. Returns the number of right (1) bits. +static size_t partitionGeneric( + uint8_t* bitmap, + uint8_t* lhs, + uint8_t* rhs, + const uint8_t* ranks, + size_t numRanks, + uint8_t rightRank) +{ + // The bitmap holds numRanks bits; the +SLOP matches the over-write budget + // the kernel contract guarantees, leaving room for the bitstream's 8-byte + // stores. + ZS_BitCStreamFF out = ZS_BitCStreamFF_init( + bitmap, bitmapBytes(numRanks) + ZL_PIVCO_HUFFMAN_SLOP); + + size_t ones = 0; + size_t zeros = 0; + for (size_t i = 0; i < numRanks; ++i) { + const uint8_t rank = ranks[i]; + const bool isRight = rank >= rightRank; + + ZS_BitCStreamFF_write(&out, isRight ? 1 : 0, 1); + // Flush every bit so the accumulator never holds more than 8 bits. + ZS_BitCStreamFF_flush(&out); + + if (isRight) { + if (rhs != NULL) { + rhs[ones] = rank; + } + ++ones; + } else { + if (lhs != NULL) { + lhs[zeros] = rank; + } + ++zeros; + } + } + (void)ZS_BitCStreamFF_finish(&out); + return ones; +} + +static size_t partitionLeft( + uint8_t* bitmap, + uint8_t* lhs, + const uint8_t* ranks, + size_t numRanks, + uint8_t rightRank) +{ + return partitionGeneric(bitmap, lhs, NULL, ranks, numRanks, rightRank); +} + +static size_t partitionRight( + uint8_t* bitmap, + uint8_t* rhs, + const uint8_t* ranks, + size_t numRanks, + uint8_t rightRank) +{ + return partitionGeneric(bitmap, NULL, rhs, ranks, numRanks, rightRank); +} + +static void partitionNone( + uint8_t* bitmap, + const uint8_t* ranks, + size_t numRanks, + uint8_t rightRank) +{ + (void)partitionGeneric(bitmap, NULL, NULL, ranks, numRanks, rightRank); +} + +// Packs one @p depth-bit index per rank into @p bitmap. The index is the rank's +// offset from @p rankBegin (its position within a flat leaf) and must fit in +// @p depth bits. +static void packFlatDepth( + uint8_t* bitmap, + size_t depth, + const uint8_t* ranks, + size_t numRanks, + uint8_t rankBegin) +{ + assert(depth >= 1); + assert(depth <= 8); + + ZS_BitCStreamFF out = ZS_BitCStreamFF_init( + bitmap, bitmapBytes(numRanks * depth) + ZL_PIVCO_HUFFMAN_SLOP); + + for (size_t i = 0; i < numRanks; ++i) { + const size_t index = (size_t)(ranks[i] - rankBegin); + assert(index < ((size_t)1 << depth)); + + ZS_BitCStreamFF_write(&out, index, depth); + ZS_BitCStreamFF_flush(&out); + } + (void)ZS_BitCStreamFF_finish(&out); +} + +static bool supported(const ZL_cpuid_t* cpuid) +{ + (void)cpuid; + return true; +} + +const ZL_PivCoHuffmanEncode ZL_PivCoHuffmanEncode_generic = { + .supported = supported, + .partitionFull = partitionGeneric, + .partitionLeft = partitionLeft, + .partitionRight = partitionRight, + .partitionNone = partitionNone, + .packFlatDepth = packFlatDepth, +}; diff --git a/src/openzl/codecs/pivco_huffman/arch/encode_pivco_arch.h b/src/openzl/codecs/pivco_huffman/arch/encode_pivco_arch.h new file mode 100644 index 000000000..728ebbfd5 --- /dev/null +++ b/src/openzl/codecs/pivco_huffman/arch/encode_pivco_arch.h @@ -0,0 +1,127 @@ +// Copyright (c) Meta Platforms, Inc. and affiliates. +#ifndef OPENZL_CODECS_PIVCO_HUFFMAN_ARCH_ENCODE_PIVCO_ARCH_H +#define OPENZL_CODECS_PIVCO_HUFFMAN_ARCH_ENCODE_PIVCO_ARCH_H + +#include +#include + +#include "openzl/codecs/pivco_huffman/arch/common_pivco_arch.h" +#include "openzl/shared/cpu.h" +#include "openzl/shared/portability.h" + +ZL_BEGIN_C_DECLS + +typedef struct { + /** + * @returns true iff the CPU supports this implementation. + * @note The rest of the function pointers may be NULL if not supported. + */ + bool (*supported)(const ZL_cpuid_t* cpuid); + + /** + * Partitions @p ranks into @p lhs and @p rhs based on @p rightRank. + * If the rank is below @p rightRank, then it goes to @p lhs, otherwise it + * goes to @p rhs. The partition decision is for each rank is written + * as a bit in @p bitmap, 0 for left and 1 for right. + * + * @param bitmap Must be `(numRanks + 7) / 8 + SLOP` bytes + * @param lhs Must be `numRanks + SLOP` elements large + * @param rhs Must be `numRanks + SLOP` elements large + * @param ranks Must be `numRanks + SLOP` elements large, + * where the first @p numRanks elements are valid. + * + * @returns the number of ones in the bitmap + * + * @note @p lhs or @p rhs may alias @p ranks + */ + size_t (*partitionFull)( + uint8_t* bitmap, + uint8_t* lhs, + uint8_t* rhs, + const uint8_t* ranks, + size_t numRanks, + uint8_t rightRank); + + /** + * Partitions @p ranks into @p lhs based on @p rightRank. If the rank is + * below @p rightRank, then it goes to @p lhs. The partition decision is for + * each rank is written as a bit in @p bitmap, 0 for left and 1 for + * right. + * + * @param bitmap Must be `(numRanks + 7) / 8 + SLOP` bytes + * @param lhs Must be `numRanks + SLOP` elements large + * @param ranks Must be `numRanks + SLOP` elements large, + * where the first @p numRanks elements are valid. + * + * @returns the number of ones in the bitmap + */ + size_t (*partitionLeft)( + uint8_t* bitmap, + uint8_t* lhs, + const uint8_t* ranks, + size_t numRanks, + uint8_t rightRank); + + /** + * Partitions @p ranks into @p rhs based on @p rightRank. If the rank is + * at least @p rightRank, then it goes to @p rhs. The partition decision is + * for each rank is written as a bit in @p bitmap, 0 for left and 1 for + * right. + * + * @param bitmap Must be `(numRanks + 7) / 8 + SLOP` bytes + * @param rhs Must be `numRanks + SLOP` elements large + * @param ranks Must be `numRanks + SLOP` elements large, + * where the first @p numRanks elements are valid. + * + * @returns the number of ones in the bitmap + */ + size_t (*partitionRight)( + uint8_t* bitmap, + uint8_t* rhs, + const uint8_t* ranks, + size_t numRanks, + uint8_t rightRank); + + /** + * Partitions @p ranks based on @p rightRank. If the rank is below + * @p rightRank, then it goes to the left, otherwise it goes to the right. + * That decision is written as a bit in @p bitmap, 0 for left and 1 for + * right. + * + * @param bitmap Must be `(numRanks + 7) / 8 + SLOP` bytes + * @param ranks Must be `numRanks + SLOP` elements large, + * where the first @p numRanks elements are valid. + */ + void (*partitionNone)( + uint8_t* bitmap, + const uint8_t* ranks, + size_t numRanks, + uint8_t rightRank); + + /** + * Subtracts @p rankBegin from each rank in @p ranks, the result of which + * must be at most @p depth bits, and stores it into @p bitmap. + * + * @param bitmap Must be `(numRanks * depth + 7) / 8 + SLOP` bytes + * @param ranks Must be `numRanks + SLOP` elements large, + * where the first @p numRanks elements are valid. + */ + void (*packFlatDepth)( + uint8_t* bitmap, + size_t depth, + const uint8_t* ranks, + size_t numRanks, + uint8_t rankBegin); +} ZL_PivCoHuffmanEncode; + +/** + * @returns The best kernel for the CPU. + */ +const ZL_PivCoHuffmanEncode* ZL_PivCoHuffmanEncode_select( + const ZL_cpuid_t* cpuid); + +extern const ZL_PivCoHuffmanEncode ZL_PivCoHuffmanEncode_generic; + +ZL_END_C_DECLS + +#endif diff --git a/src/openzl/codecs/pivco_huffman/common_pivco_kernel.c b/src/openzl/codecs/pivco_huffman/common_pivco_kernel.c new file mode 100644 index 000000000..ccc0ff69c --- /dev/null +++ b/src/openzl/codecs/pivco_huffman/common_pivco_kernel.c @@ -0,0 +1,375 @@ +// Copyright (c) Meta Platforms, Inc. and affiliates. + +#include "openzl/codecs/pivco_huffman/common_pivco_kernel.h" + +#include +#include + +#include "openzl/shared/bits.h" +#include "openzl/shared/utils.h" + +int ZL_PivCoHuffman_computeTableLog(const uint8_t* weights, size_t weightsSize) +{ + if (weightsSize == 0 || weightsSize > ZL_PIVCO_MAX_SYMBOLS) { + return -1; + } + assert(weights != NULL); + + // A complete prefix code satisfies sum(2^(weight-1)) == 2^tableLog over the + // non-zero weights. Accumulate that sum, masking the shift to stay defined + // for the (rejected) out-of-range weights, and branchlessly track validity. + bool invalid = false; + uint32_t sum = 0; + for (size_t i = 0; i < weightsSize; ++i) { + const uint8_t w = weights[i]; + invalid |= w > ZL_PIVCO_MAX_TABLE_LOG; + sum += (1u << (w & 31)) >> 1; + } + + if (invalid) { + return -1; + } + + if (sum == 0 || !ZL_isPow2(sum)) { + return -1; + } + + const int tableLog = ZL_highbit32(sum); + if (tableLog > ZL_PIVCO_MAX_TABLE_LOG) { + return -1; + } + + return tableLog; +} + +void ZL_PivCoHuffman_countWeights( + uint16_t weightCounts[16], + const uint8_t* weights, + size_t numWeights) +{ + assert(numWeights <= 256); + ZL_STATIC_ASSERT(ZL_PIVCO_MAX_TABLE_LOG < 16, "Assumption"); +#if ZL_HAS_SSSE3 + // A weight histogram has at most 16 buckets and typically heavy collisions, + // so a scalar histogram serializes on a few hot counters. Instead keep all + // 16 counts in one SIMD register: for each weight, compare it against the + // lane indices 0..15 and subtract the (0/-1) match mask, incrementing the + // matching lane. + __m128i const iota = + _mm_setr_epi8(0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15); + __m128i count = _mm_setzero_si128(); + for (size_t i = 0; i < numWeights; ++i) { + __m128i const inc = + _mm_cmpeq_epi8(_mm_set1_epi8((char)weights[i]), iota); + count = _mm_sub_epi8(count, inc); + } + // Lanes are 8-bit, so a count of exactly 256 wraps to 0. The only way every + // lane reads 0 after 256 weights is a single bucket that wrapped: recover + // it explicitly. + bool const everyCountIsZero = + _mm_movemask_epi8(_mm_cmpeq_epi8(count, _mm_setzero_si128())) + == 0xFFFF; + if (numWeights == 256 && everyCountIsZero) { + memset(weightCounts, 0, sizeof(*weightCounts) * 16); + weightCounts[weights[0]] = 256; + return; + } + + // Widen the 8-bit lane counts to the 16-bit output. + const __m128i lo = _mm_unpacklo_epi8(count, _mm_setzero_si128()); + const __m128i hi = _mm_unpackhi_epi8(count, _mm_setzero_si128()); + _mm_storeu_si128((__m128i_u*)&weightCounts[0], lo); + _mm_storeu_si128((__m128i_u*)&weightCounts[8], hi); +#else + // Scalar fallback: four independent histograms unrolled over the input to + // hide the load-update-store latency, then summed. uint8 counters cannot + // overflow because each accumulates at most numWeights/4 < 256 increments. + uint8_t weightCounts0[16] = { 0 }; + uint8_t weightCounts1[16] = { 0 }; + uint8_t weightCounts2[16] = { 0 }; + uint8_t weightCounts3[16] = { 0 }; + + size_t const prefix = numWeights % 4; + for (size_t i = 0; i < prefix; ++i) { + weightCounts0[weights[i]] = (uint8_t)(weightCounts0[weights[i]] + 1); + } + for (size_t i = prefix; i < numWeights; i += 4) { + weightCounts0[weights[i + 0]] = + (uint8_t)(weightCounts0[weights[i + 0]] + 1); + weightCounts1[weights[i + 1]] = + (uint8_t)(weightCounts1[weights[i + 1]] + 1); + weightCounts2[weights[i + 2]] = + (uint8_t)(weightCounts2[weights[i + 2]] + 1); + weightCounts3[weights[i + 3]] = + (uint8_t)(weightCounts3[weights[i + 3]] + 1); + } + for (size_t i = 0; i < 16; ++i) { + weightCounts[i] = (uint16_t)(weightCounts0[i] + weightCounts1[i] + + weightCounts2[i] + weightCounts3[i]); + } +#endif +} + +/** + * Worst-case leaf count: one single-symbol leaf per symbol, plus the flat + * leaves. Flattening a weight bucket can add at most one flat leaf to each of + * the weight levels above it, bounding the extra leaves by the triangular + * number 1 + 2 + ... + ZL_PIVCO_MAX_WEIGHT. + */ +#define PIVCO_LEAF_STORAGE_SIZE \ + (ZL_PIVCO_MAX_SYMBOLS \ + + ((ZL_PIVCO_MAX_WEIGHT * (ZL_PIVCO_MAX_WEIGHT + 1)) / 2)) +/** + * Each symbol is stored once for its single-symbol leaf and at most once more + * inside a flat leaf. + */ +#define PIVCO_SYMBOL_STORAGE_SIZE (2 * ZL_PIVCO_MAX_SYMBOLS) + +/** + * A leaf under construction. It owns `numSymbols` symbols stored contiguously + * at `symbolStorage + symbolOffset`. A single-symbol leaf has numSymbols == 1; + * a flat leaf has a power-of-two numSymbols. Leaves exist only while building; + * the finished tree is purely rank-indexed. + * + * A flat leaf copies its symbols into symbolStorage instead of referencing its + * constituent single-symbol leaves in place. That is required: flattening a + * lower weight can append a flat leaf onto those leaves' slots (see appendLeaf + * reusing positions vacated by peeling), so the in-place symbols are not stable + * once a leaf has been flattened. + */ +typedef struct { + uint16_t symbolOffset; + uint16_t numSymbols; +} ZL_PivCoLeaf; + +/** + * Mutable scratch for ZL_PivCoHuffmanTree_build. `leaves` is grouped by + * weight: weight `w`'s leaves occupy [weightOffsets[w], weightOffsets[w + 1]), + * and weightCounts[w] tracks how many of those slots are used. weightCounts + * needs 16 lanes so countWeights can fill it with a single SIMD store. + * `symbolStorage` backs the leaves' symbols and is bump-allocated by + * `nextSymbol`. + */ +typedef struct { + ZL_PivCoLeaf leaves[PIVCO_LEAF_STORAGE_SIZE]; + uint8_t symbolStorage[PIVCO_SYMBOL_STORAGE_SIZE]; + uint16_t weightCounts[16]; + uint16_t weightOffsets[ZL_PIVCO_MAX_WEIGHT + 2]; + uint16_t nextSymbol; +} ZL_PivCoBuilder; + +/** @returns log2(numSymbols): the leaf's flat depth. */ +static size_t leafFlatDepth(const ZL_PivCoLeaf* leaf) +{ + return (size_t)ZL_highbit32((uint32_t)leaf->numSymbols); +} + +/** + * Reserves @p count contiguous bytes in the builder's symbolStorage and + * advances past them. @returns A pointer to the reserved bytes. + */ +static uint8_t* appendSymbols(ZL_PivCoBuilder* builder, size_t count) +{ + assert((int)count <= PIVCO_SYMBOL_STORAGE_SIZE - builder->nextSymbol); + uint8_t* const symbols = builder->symbolStorage + builder->nextSymbol; + builder->nextSymbol += (uint16_t)count; + return symbols; +} + +/** + * Appends an (uninitialized) leaf to weight bucket @p weight, consuming one of + * the slots reserved for that bucket by weightOffsets. @returns The new leaf, + * for the caller to fill in. + */ +static ZL_PivCoLeaf* appendLeaf(ZL_PivCoBuilder* builder, int weight) +{ + assert(weight >= 0); + assert(weight <= ZL_PIVCO_MAX_WEIGHT); + + size_t const pos = (size_t)builder->weightOffsets[weight] + + builder->weightCounts[weight]; + assert(pos < builder->weightOffsets[weight + 1]); + + ZL_PivCoLeaf* const leaf = &builder->leaves[pos]; + ++builder->weightCounts[weight]; + return leaf; +} + +/** + * Collapses runs of equal-weight single-symbol leaves into flat leaves. + * + * `2^k` leaves of the same weight `w` (same code length) form a complete + * subtree rooted `k` levels up, i.e. a single leaf of weight `w + k` holding + * those `2^k` symbols. Decoding that leaf is a flat `k`-bit lookup instead of a + * chain of binary pivots, so it is both smaller and faster. + * + * For each weight from shallowest (largest weight) to deepest (smallest + * weight), repeatedly peel the largest power-of-two block (taken from the end + * of the bucket) while more than two leaves remain; a residual of one or two is + * left for the ordinary binary pivot, which a flat node could not improve on. + * New flat leaves land in higher weight (shallower) buckets that have already + * been processed, so they are never re-flattened. + */ +static void optimizeLeaves(ZL_PivCoBuilder* builder, int tableLog) +{ + for (int weight = tableLog; weight > 0; --weight) { + while (builder->weightCounts[weight] > 2) { + uint16_t const numWeights = builder->weightCounts[weight]; + int const flatBits = ZL_highbit32(numWeights); + size_t const flatNum = (size_t)1 << flatBits; + size_t const flatOff = (size_t)numWeights - flatNum; + int const flatWeight = weight + flatBits; + + ZL_PivCoLeaf* const leaf = appendLeaf(builder, flatWeight); + uint8_t* const symbols = appendSymbols(builder, flatNum); + + // Copy the peeled suffix's symbols into the flat leaf now: a later, + // lower-weight flattening can append over these single leaves' + // slots, so they must be captured before that can happen. + for (size_t i = 0; i < flatNum; ++i) { + const ZL_PivCoLeaf* child = + &builder->leaves + [builder->weightOffsets[weight] + flatOff + i]; + assert(child->numSymbols == 1); + symbols[i] = builder->symbolStorage[child->symbolOffset]; + } + + leaf->symbolOffset = (uint16_t)(symbols - builder->symbolStorage); + leaf->numSymbols = (uint16_t)flatNum; + builder->weightCounts[weight] -= (uint16_t)flatNum; + } + } +} + +/** + * Walks the builder's leaves in canonical order and fills the tree's + * rank-indexed arrays (symbolToRank, rankToSymbol, rankToFlatDepth, + * rankToCodeword) and numRanks. + * + * Leaves are visited shallowest-first (level 0 == the largest weight). Within a + * level, canonical Huffman codewords are consecutive integers, so we keep a + * running `codeword` counter that increments per leaf and shifts left by one + * when descending a level. A flat leaf of depth `flatBits` occupies + * `2^flatBits` consecutive ranks/codewords, one per contained symbol. + */ +static void assignRanksAndCodewords( + ZL_PivCoHuffmanTree* tree, + const ZL_PivCoBuilder* builder) +{ + uint32_t codeword = 0; + uint16_t rank = 0; + for (size_t level = 0; level < tree->numLevels; ++level) { + size_t const weight = (size_t)tree->tableLog + 1 - level; + + for (size_t idx = 0; idx < builder->weightCounts[weight]; ++idx) { + size_t const leafIndex = + (size_t)builder->weightOffsets[weight] + idx; + const ZL_PivCoLeaf* leaf = &builder->leaves[leafIndex]; + size_t const flatBits = leafFlatDepth(leaf); + size_t const totalBits = level + flatBits; + assert(totalBits <= 16); + assert(rank <= ZL_PIVCO_MAX_SYMBOLS - leaf->numSymbols); + + const uint8_t* symbols = + builder->symbolStorage + leaf->symbolOffset; + + for (size_t flatIdx = 0; flatIdx < leaf->numSymbols; ++flatIdx) { + uint8_t const symbol = symbols[flatIdx]; + uint16_t const symbolRank = (uint16_t)(rank + flatIdx); + // The symbol's codeword is the leaf's prefix followed by its + // flat index, left-justified into 16 bits. totalBits == 0 only + // for the single-symbol (constant) tree, which has no codeword + // bits; special-cased to avoid a 16-bit shift by 16. + uint32_t const bits = + (uint32_t)((codeword << flatBits) + flatIdx); + uint16_t const symbolCodeword = totalBits == 0 + ? 0 + : (uint16_t)(bits << (16 - totalBits)); + tree->symbolToRank[symbol] = (uint8_t)symbolRank; + tree->rankToSymbol[symbolRank] = symbol; + tree->rankToFlatDepth[symbolRank] = (uint8_t)flatBits; + tree->rankToCodeword[symbolRank] = symbolCodeword; + } + rank += leaf->numSymbols; + + ++codeword; + } + + if (level + 1 < tree->numLevels) { + codeword <<= 1; + } + } + + // Sanity check: a complete canonical code leaves the codeword counter at + // exactly 2^(numLevels - 1) (one full code space at the deepest level). + assert(codeword == ((uint32_t)1 << (tree->numLevels - 1))); + assert(rank <= ZL_PIVCO_MAX_SYMBOLS); + tree->numRanks = rank; +} + +void ZL_PivCoHuffmanTree_build( + ZL_PivCoHuffmanTree* tree, + const uint8_t* weights, + size_t weightsSize, + int tableLog) +{ + assert(tree != NULL); + assert(tableLog >= 0); + assert(tableLog <= ZL_PIVCO_MAX_TABLE_LOG); + assert(ZL_PivCoHuffman_computeTableLog(weights, weightsSize) == tableLog); + + memset(tree, 0, sizeof(*tree)); + tree->tableLog = tableLog; + + ZL_PivCoBuilder builder; + builder.nextSymbol = 0; + ZL_PivCoHuffman_countWeights(builder.weightCounts, weights, weightsSize); + + // Lay out `leaves` grouped by weight (a prefix sum over weightCounts). + // Beyond its single-symbol leaves, a weight-w bucket reserves a gap of w + // slots for flat leaves that flattening lower weights deposits here: each + // lower weight w' < w contributes at most one (its peels have strictly + // decreasing depth), so w slots always suffice. + builder.weightOffsets[0] = 0; + builder.weightOffsets[1] = 0; + for (size_t weight = 2; weight < ZL_PIVCO_MAX_WEIGHT + 2; ++weight) { + builder.weightOffsets[weight] = + (uint16_t)(builder.weightOffsets[weight - 1] + + builder.weightCounts[weight - 1] + (weight - 1)); + } + + // Seed each present symbol as a single-symbol leaf in its weight bucket. + // weightPos tracks the next free slot per bucket as we fill it. + { + uint16_t weightPos[ZL_PIVCO_MAX_WEIGHT + 2]; + memcpy(weightPos, builder.weightOffsets, sizeof(weightPos)); + for (size_t symbol = 0; symbol < weightsSize; ++symbol) { + uint8_t const weight = weights[symbol]; + if (weight == 0) { + continue; + } + + ZL_PivCoLeaf* leaf = &builder.leaves[weightPos[weight]++]; + uint8_t* const symbols = appendSymbols(&builder, 1); + symbols[0] = (uint8_t)symbol; + leaf->symbolOffset = (uint16_t)(symbols - builder.symbolStorage); + leaf->numSymbols = 1; + } + } + + optimizeLeaves(&builder, tableLog); + + // The number of levels is fixed by the deepest non-empty level, which is + // the smallest weight bucket that still has leaves (flattening can empty + // out the deepest buckets). level == tableLog + 1 - weight, so the deepest + // level + 1 == tableLog + 2 - minWeight. + size_t minWeight = 1; + while (builder.weightCounts[minWeight] == 0) { + ++minWeight; + } + assert(minWeight <= (size_t)tableLog + 1); + tree->numLevels = (uint16_t)((size_t)tableLog + 2 - minWeight); + + assignRanksAndCodewords(tree, &builder); + assert(tree->numRanks != 0); +} diff --git a/src/openzl/codecs/pivco_huffman/common_pivco_kernel.h b/src/openzl/codecs/pivco_huffman/common_pivco_kernel.h new file mode 100644 index 000000000..1489cad74 --- /dev/null +++ b/src/openzl/codecs/pivco_huffman/common_pivco_kernel.h @@ -0,0 +1,229 @@ +// Copyright (c) Meta Platforms, Inc. and affiliates. + +#ifndef OPENZL_CODECS_PIVCO_COMMON_PIVCO_KERNEL_H +#define OPENZL_CODECS_PIVCO_COMMON_PIVCO_KERNEL_H + +#include +#include +#include + +#include "openzl/shared/portability.h" + +ZL_BEGIN_C_DECLS + +/** + * The PivCo-Huffman coding tree. + * + * The tree is the structure shared by the encoder and decoder. It is derived + * deterministically from the zstd-style Huffman `weights` alone (it is never + * serialized), so both sides build an identical tree from the same weights. + * + * Concepts used throughout this codec: + * + * - weight: a zstd Huffman weight. For a symbol with code length `L` under a + * table of `tableLog` bits, weight = `tableLog + 1 - L`. A weight of 0 means + * the symbol is absent. Larger weight == shorter code == closer to the root. + * + * - rank: symbols are ordered canonically (shortest codes first), and a + * symbol's rank is its index in that order. PivCo coding operates on + * contiguous rank ranges rather than on symbols directly, and the tree is + * stored entirely as rank-indexed arrays. + * + * - leaf: each present symbol starts as a single-symbol leaf. A run of leaves + * that share the same code length may be collapsed into one multi-symbol + * "flat leaf" of `2^flatDepth` symbols, which the codec processes with a + * flat fixed-width bitmap instead of a chain of binary splits. A leaf + * occupies a contiguous rank range; flattening happens only while building + * (see the .c) and leaves no leaf objects in the tree. + * + * - level: the depth at which pivco coding splits. At each level a rank + * range is partitioned into the codewords whose bit at that level is 0 vs 1; + * ZL_PivCoHuffmanTree_splitRank finds that partition point from the + * per-rank codewords (`rankToCodeword`). + */ + +/** + * Default bytes of output per pivco block. The block size is a parameter to + * the encoder/decoder kernels; this is the value the binding uses and records + * in the codec header when the input spans more than one block. It can be + * changed at any time without breaking the wire format, and is just a + * recommended default value. + */ +#define ZL_PIVCO_DEFAULT_BLOCK_SIZE ((size_t)(32 * 1024)) +/** + * Blocks larger than this are disallowed by the encoder and decoder. + * The encoder asserts the contract is respected, and the decoder validates the + * block size does not exceed this bound. + * + * The bound is chosen so that manipulating bits based on this number of bytes + * is still comfortably below overlowing a U32. + */ +#define ZL_PIVCO_MAX_BLOCK_SIZE ((size_t)1 << 28) +#define ZL_PIVCO_MAX_SYMBOLS 256 +#define ZL_PIVCO_MAX_TABLE_LOG 12 +/** + * The shallowest leaf (the flat root) sits one level above the longest code, + * so weights run 1..tableLog+1. + */ +#define ZL_PIVCO_MAX_WEIGHT (ZL_PIVCO_MAX_TABLE_LOG + 1) +/** + * Max nodes in the pivot tree: a binary tree with up to ZL_PIVCO_MAX_SYMBOLS + * leaves has at most 2 * ZL_PIVCO_MAX_SYMBOLS - 1 nodes. The encoder uses + * this to bound per-block overhead. + */ +#define ZL_PIVCO_MAX_TREE_NODES (2 * ZL_PIVCO_MAX_SYMBOLS - 1) + +typedef struct ZL_PivCoHuffmanTree_s { + /** symbol -> rank. Used by the encoder to map source bytes to ranks. */ + uint8_t symbolToRank[ZL_PIVCO_MAX_SYMBOLS]; + /** + * rank -> the rank's symbol. Within a leaf, ranks are in flat order, so a + * leaf's symbols are a contiguous slice of this array. + */ + uint8_t rankToSymbol[ZL_PIVCO_MAX_SYMBOLS]; + /** + * rank -> the flat depth (log2 of the symbol count) of the leaf containing + * that rank. A leaf starting at `r` therefore spans `1 << + * rankToFlatDepth[r]` ranks, which is how a leaf is distinguished from an + * internal node. + */ + uint8_t rankToFlatDepth[ZL_PIVCO_MAX_SYMBOLS]; + /** + * rank -> the rank's canonical codeword, left-justified (MSB-aligned) into + * 16 bits. Sorted by rank, so splitRank can scan it for a level's 0/1 bit + * boundary. + */ + uint16_t rankToCodeword[ZL_PIVCO_MAX_SYMBOLS]; + /** Number of levels with at least one leaf (tree depth). */ + uint16_t numLevels; + /** Number of ranks == number of present symbols. */ + uint16_t numRanks; + /** Huffman table log: the longest code is `tableLog` bits (see weight). */ + int tableLog; +} ZL_PivCoHuffmanTree; + +/** + * @returns Whether the rank range [firstRank, rankEnd) is exactly one leaf + * (a base case for the recursive encode/decode), as opposed to an internal node + * spanning multiple leaves. + * + * @pre The range is a valid tree-node range: 0 <= firstRank < rankEnd <= + * numRanks, and firstRank is a leaf boundary. Every caller derives the range + * from the tree itself (`splitRank` / `numRanks`), which guarantees this; the + * range never comes from untrusted bitstream data. + */ +ZL_INLINE bool ZL_PivCoHuffmanTree_rangeIsLeaf( + const ZL_PivCoHuffmanTree* tree, + size_t firstRank, + size_t rankEnd) +{ + // The leaf starting at firstRank spans 2^flatDepth ranks; the range is that + // leaf iff it has exactly that length. A longer range spans into the next + // leaf, and (per the precondition) sub-leaf ranges never occur. + return ((size_t)1 << tree->rankToFlatDepth[firstRank]) + == rankEnd - firstRank; +} + +/** + * @returns The flat depth (log2 of the symbol count) of the leaf starting at + * @p firstRank; 0 for a constant (single-symbol) leaf. + * @pre firstRank is a leaf boundary. + */ +ZL_INLINE size_t ZL_PivCoHuffmanTree_leafFlatDepth( + const ZL_PivCoHuffmanTree* tree, + size_t firstRank) +{ + return tree->rankToFlatDepth[firstRank]; +} + +/** + * @returns Whether the rank range [firstRank, rankEnd) is a single constant + * (single-symbol) leaf, which contributes nothing to the bitstream. + * @pre The range is a valid tree-node range (see rangeIsLeaf). + */ +ZL_INLINE bool ZL_PivCoHuffmanTree_rangeIsConstantLeaf( + const ZL_PivCoHuffmanTree* tree, + size_t firstRank, + size_t rankEnd) +{ + return ZL_PivCoHuffmanTree_rangeIsLeaf(tree, firstRank, rankEnd) + && ZL_PivCoHuffmanTree_leafFlatDepth(tree, firstRank) == 0; +} + +/** + * @returns The symbols of the leaf starting at @p firstRank, in flat order + * (`1 << leafFlatDepth` of them, laid out contiguously in rank order). + * @pre firstRank is a leaf boundary. + */ +ZL_INLINE const uint8_t* ZL_PivCoHuffmanTree_leafSymbols( + const ZL_PivCoHuffmanTree* tree, + size_t firstRank) +{ + return &tree->rankToSymbol[firstRank]; +} + +/** + * @returns The split rank of the internal node covering [firstRank, rankEnd) at + * @p level: the first rank whose codeword has bit @p level set, i.e. the + * boundary between the node's 0-bit (left) and 1-bit (right) children. + * @pre The node is an internal node (not a single leaf), so + * firstRank + 1 < rankEnd <= numRanks. + */ +ZL_INLINE uint16_t ZL_PivCoHuffmanTree_splitRank( + const ZL_PivCoHuffmanTree* tree, + size_t level, + size_t firstRank, + size_t rankEnd) +{ + (void)rankEnd; + assert(level < tree->numLevels); + assert(firstRank + 1 < rankEnd); + assert(rankEnd <= tree->numRanks); + + // Codewords are MSB-aligned, so level L's bit is the (L+1)-th from the top. + // Ranks are in canonical (codeword) order, so within the range bit `level` + // reads 0 over a prefix then 1 over the rest; scan for the first 1. The + // scan is short -- shorter (higher-weight) codewords are visited first -- + // and the range is tiny, so a linear scan beats a binary search. + uint16_t const mask = (uint16_t)(0x8000u >> level); + assert((tree->rankToCodeword[firstRank] & mask) == 0); + + size_t splitRank = firstRank + 1; + while ((tree->rankToCodeword[splitRank] & mask) == 0) { + assert(splitRank < rankEnd); + ++splitRank; + } + assert(firstRank < splitRank); + assert(splitRank < rankEnd); + return (uint16_t)splitRank; +} + +/** + * Builds a tree from validated weights. + * @pre tableLog == ZL_PivCoHuffman_computeTableLog(weights, weightsSize) + */ +void ZL_PivCoHuffmanTree_build( + ZL_PivCoHuffmanTree* tree, + const uint8_t* weights, + size_t weightsSize, + int tableLog); + +/** + * @returns The table log of the Huffman tree described by @p weights or -1 if + * the weights are invalid. + */ +int ZL_PivCoHuffman_computeTableLog(const uint8_t* weights, size_t weightsSize); + +/** + * Efficiently counts weight frequency and outputs to @p weightCounts. + * + * @pre numWeights <= 256 + */ +void ZL_PivCoHuffman_countWeights( + uint16_t weightCounts[16], + const uint8_t* weights, + size_t numWeights); + +ZL_END_C_DECLS + +#endif diff --git a/src/openzl/codecs/pivco_huffman/decode_pivco_kernel.c b/src/openzl/codecs/pivco_huffman/decode_pivco_kernel.c new file mode 100644 index 000000000..7f0573789 --- /dev/null +++ b/src/openzl/codecs/pivco_huffman/decode_pivco_kernel.c @@ -0,0 +1,380 @@ +// Copyright (c) Meta Platforms, Inc. and affiliates. + +#include "openzl/codecs/pivco_huffman/decode_pivco_kernel.h" + +#include +#include + +#include "openzl/codecs/common/bitstream/ff_bitstream.h" +#include "openzl/codecs/pivco_huffman/common_pivco_kernel.h" +#include "openzl/shared/bits.h" +#include "openzl/shared/utils.h" + +/** + * The decoded result of a tree node. A constant node covers a single-symbol + * leaf: the whole rank range is `symbol` and `output` is unused (the caller + * fills the run). Otherwise the node's decoded bytes have been written to + * `output`. + */ +typedef struct { + bool isConstant; + uint8_t symbol; + uint8_t* output; +} ZL_PivCoHuffmanDecodeResult; + +size_t ZL_PivCoHuffmanDecode_scratchBytes(size_t dstSize, size_t blockSize) +{ + // The recursion ping-pongs between two slop-padded buffers (see decodeNode) + // so dst is never used as scratch -- only each block's root merge writes + // into dst, honoring its exact (slop-free) capacity. The largest block is + // min(dstSize, blockSize) bytes long, and each buffer holds one block plus + // the kernels' over-write slop. + blockSize = ZL_MIN(dstSize, blockSize); + return 2 * (blockSize + ZL_PIVCO_HUFFMAN_SLOP); +} + +/** + * Reads a value that the encoder wrote with the fixed width needed to represent + * @p cap (the width is derived from @p cap, not stored). @returns the value; + * the caller must validate it against @p cap, since a corrupt stream can yield + * a larger value. + */ +static size_t bitReaderDecodeCappedValue(ZS_BitDStreamFF* reader, size_t cap) +{ + size_t const numBits = (size_t)ZL_nextPow2(cap + 1); + return ZS_BitDStreamFF_read(reader, numBits); +} + +/** + * Interleaves the two decoded children @p lhs and @p rhs into @p output (of + * @p outputSize bytes, @p outputCapacity available) per @p bitmap: bit i picks + * rhs (1) or lhs (0) for output position i. Each child is either a constant + * symbol or a decoded vector, dispatched to the matching merge kernel; when + * both are constant the bitmap alone reconstructs the output (no count was + * sent). + * + * @returns true on success; false when a vector merge's one-count disagrees + * with + * @p rhsSize, which indicates a corrupt bitstream. + */ +static bool mergeDecodeResults( + const ZL_PivCoHuffmanDecode* kernels, + uint8_t* output, + size_t outputSize, + size_t outputCapacity, + const uint8_t* bitmap, + size_t bitmapBytes, + const ZL_PivCoHuffmanDecodeResult* lhs, + size_t lhsSize, + const ZL_PivCoHuffmanDecodeResult* rhs, + size_t rhsSize) +{ + if (lhs->isConstant && rhs->isConstant) { + uint8_t symbols[2] = { lhs->symbol, rhs->symbol }; + kernels->mergeFlatDepth( + output, + outputSize, + outputCapacity, + bitmap, + bitmapBytes, + 1, + symbols); + return true; + } + + size_t ones; + if (lhs->isConstant) { + ones = kernels->mergeConstantVector( + output, + outputCapacity, + bitmap, + bitmapBytes, + lhs->symbol, + lhsSize, + rhs->output, + rhsSize); + } else if (rhs->isConstant) { + ones = kernels->mergeVectorConstant( + output, + outputCapacity, + bitmap, + bitmapBytes, + lhs->output, + lhsSize, + rhs->symbol, + rhsSize); + } else { + ones = kernels->mergeVectorVector( + output, + outputCapacity, + bitmap, + bitmapBytes, + lhs->output, + lhsSize, + rhs->output, + rhsSize); + } + + return ones == rhsSize; +} + +/** + * Decodes the tree node covering ranks [firstRank, rankEnd) at @p level, + * producing @p numCodewords output symbols, and writes the node's result to + * @p out (capacity @p outCapacity). + * + * @p out is write-only: only this node's final merge touches it, so it may be + * the exactly-sized destination buffer. The merge honors @p outCapacity, only + * over-writing past the logical size when the capacity leaves room -- which it + * never does for the last block written to dst. + * + * The subtree recurses entirely within the two slop-padded ping-pong buffers + * @p scratch1 and @p scratch2 (capacities @p scratch1Capacity and + * @p scratch2Capacity, each >= numCodewords + ZL_PIVCO_HUFFMAN_SLOP). This + * node decodes its two children packed into @p scratch1 (lhs then rhs) using + * @p scratch2 as their recursion buffer, then merges @p scratch1 into @p out. + * The children recurse with scratch1/scratch2 rotated, so @p out is never + * handed down as scratch -- the destination never doubles as recursion scratch, + * and every recursive kernel always has slop in scratch1/scratch2 to over-write + * into. + * + * @note that @p out and @p scratch2 may alias. + */ +static bool decodeNode( + const ZL_PivCoHuffmanTree* tree, + const ZL_PivCoHuffmanDecode* kernels, + ZS_BitDStreamFF* reader, + size_t level, + size_t firstRank, + size_t rankEnd, + size_t numCodewords, + uint8_t* out, + size_t outCapacity, + uint8_t* scratch1, + size_t scratch1Capacity, + uint8_t* scratch2, + size_t scratch2Capacity, + ZL_PivCoHuffmanDecodeResult* result) +{ + if (ZL_PivCoHuffmanTree_rangeIsLeaf(tree, firstRank, rankEnd)) { + const uint8_t* const symbols = + ZL_PivCoHuffmanTree_leafSymbols(tree, firstRank); + const size_t depth = ZL_PivCoHuffmanTree_leafFlatDepth(tree, firstRank); + // A depth-0 leaf is a single symbol: the whole range is constant and + // nothing is read from the bitstream. + if (depth == 0) { + result->isConstant = true; + result->symbol = symbols[0]; + result->output = NULL; + return true; + } + + // A flat leaf packs `depth` bits per output, indexing into `symbols`. + // numCodewords is bounded by the block length (<= dstSize), so the + // bit-count product never overflows for any realizable output. + assert(numCodewords <= SIZE_MAX / depth); + const size_t bitmapBits = numCodewords * depth; + const uint8_t* const bitmap = + ZS_BitDStreamFF_popAlignedBits(reader, bitmapBits); + if (bitmap == NULL) { + return false; + } + const size_t bitmapBytes = (bitmapBits + 7) / 8; + kernels->mergeFlatDepth( + out, + numCodewords, + outCapacity, + bitmap, + bitmapBytes, + depth, + symbols); + result->isConstant = false; + result->symbol = 0; + result->output = out; + return true; + } + // A non-leaf range deeper than the tree is corrupt (no split exists). + if (level >= tree->numLevels) { + return false; + } + + // Internal node: it splits [firstRank, rankEnd) at splitRank into a left + // (0-bit) and right (1-bit) child, recombined via the partition bitmap. + const size_t splitRank = + ZL_PivCoHuffmanTree_splitRank(tree, level, firstRank, rankEnd); + bool const lhsIsConstant = + ZL_PivCoHuffmanTree_rangeIsConstantLeaf(tree, firstRank, splitRank); + bool const rhsIsConstant = + ZL_PivCoHuffmanTree_rangeIsConstantLeaf(tree, splitRank, rankEnd); + + const uint8_t* const bitmap = + ZS_BitDStreamFF_popAlignedBits(reader, numCodewords); + if (bitmap == NULL) { + return false; + } + const size_t bitmapBytes = (numCodewords + 7) / 8; + + // The encoder omits numOnes when both children are constant (the bitmap + // alone reconstructs the output); otherwise it is read from the stream. + size_t numOnes = 0; + if (!(lhsIsConstant && rhsIsConstant)) { + numOnes = bitReaderDecodeCappedValue(reader, numCodewords); + if (numOnes > numCodewords) { + return false; + } + } + + assert(numCodewords + ZL_PIVCO_HUFFMAN_SLOP <= scratch1Capacity); + const size_t numZeros = numCodewords - numOnes; + ZL_PivCoHuffmanDecodeResult lhs; + ZL_PivCoHuffmanDecodeResult rhs; + + // The children are packed into `scratch1` (lhs at [0, numZeros), rhs at + // [numZeros, numCodewords)). Each child recurses with scratch1/scratch2 + // rotated: it packs its own children into `scratch2` and uses its slice of + // `scratch1` as its recursion buffer. `out` is never passed down, so the + // destination is never used as recursion scratch. + if (!decodeNode( + tree, + kernels, + reader, + level + 1, + firstRank, + splitRank, + numZeros, + scratch1, + scratch1Capacity, + scratch2, + scratch2Capacity, + scratch1, + scratch1Capacity, + &lhs)) { + return false; + } + if (!decodeNode( + tree, + kernels, + reader, + level + 1, + splitRank, + rankEnd, + numOnes, + scratch1 + numZeros, + scratch1Capacity - numZeros, + scratch2, + scratch2Capacity, + scratch1 + numZeros, + scratch1Capacity - numZeros, + &rhs)) { + return false; + } + + if (!mergeDecodeResults( + kernels, + out, + numCodewords, + outCapacity, + bitmap, + bitmapBytes, + &lhs, + numZeros, + &rhs, + numOnes)) { + return false; + } + + result->isConstant = false; + result->symbol = 0; + result->output = out; + return true; +} + +bool ZL_PivCoHuffman_decode( + uint8_t* dst, + size_t dstSize, + uint8_t* scratch, + size_t scratchBytes, + const uint8_t* weights, + size_t weightsSize, + const uint8_t* bitstream, + size_t bitstreamSize, + size_t blockSize, + const ZL_PivCoHuffmanDecode* kernels) +{ + if (dstSize == 0) { + return bitstreamSize == 0; + } + + // For a non-empty output the block size bounds the decode loop (a zero + // block size would never make progress) and the scratch sizing, so its + // bounds must be validated -- it comes from the untrusted codec header. + if (blockSize == 0 || blockSize > ZL_PIVCO_MAX_BLOCK_SIZE) { + return false; + } + + if (scratchBytes < ZL_PivCoHuffmanDecode_scratchBytes(dstSize, blockSize)) { + return false; + } + + if (kernels == NULL) { + kernels = ZL_PivCoHuffmanDecode_select(NULL); + } + + ZL_PivCoHuffmanTree tree; + int const tableLog = ZL_PivCoHuffman_computeTableLog(weights, weightsSize); + if (tableLog < 0) { + return false; + } + ZL_PivCoHuffmanTree_build(&tree, weights, weightsSize, tableLog); + assert(tree.numRanks != 0); + + ZS_BitDStreamFF reader = ZS_BitDStreamFF_init(bitstream, bitstreamSize); + + // Two slop-padded ping-pong buffers carved from scratch for writing the + // intermediate outputs, while guaranteeing SLOP bytes are available for + // over-read when merging them. + const size_t bufCapacity = + ZL_MIN(dstSize, blockSize) + ZL_PIVCO_HUFFMAN_SLOP; + uint8_t* const scratch1 = scratch; + uint8_t* const scratch2 = scratch + bufCapacity; + + for (size_t off = 0; off < dstSize; off += blockSize) { + const size_t blockLen = ZL_MIN(blockSize, dstSize - off); + ZL_PivCoHuffmanDecodeResult result; + if (!decodeNode( + &tree, + kernels, + &reader, + 0, + 0, + tree.numRanks, + blockLen, + dst + off, + dstSize - off, + scratch1, + bufCapacity, + scratch2, + bufCapacity, + &result)) { + return false; + } + + // An internal or flat-leaf root merged its result straight into dst; a + // constant root writes nothing, so fill it here. + if (result.isConstant) { + memset(dst + off, result.symbol, blockLen); + } + } + + // A well-formed bitstream is consumed exactly: any leftover bytes mean the + // input was corrupt. + const ZL_Report ret = ZS_BitDStreamFF_finish(&reader); + if (ZL_isError(ret)) { + return false; + } + if (ZL_validResult(ret) != bitstreamSize) { + return false; + } + + return true; +} diff --git a/src/openzl/codecs/pivco_huffman/decode_pivco_kernel.h b/src/openzl/codecs/pivco_huffman/decode_pivco_kernel.h new file mode 100644 index 000000000..15ee008ce --- /dev/null +++ b/src/openzl/codecs/pivco_huffman/decode_pivco_kernel.h @@ -0,0 +1,47 @@ +// Copyright (c) Meta Platforms, Inc. and affiliates. + +#ifndef OPENZL_CODECS_PIVCO_HUFFMAN_DECODE_PIVCO_KERNEL_H +#define OPENZL_CODECS_PIVCO_HUFFMAN_DECODE_PIVCO_KERNEL_H + +#include +#include + +#include "openzl/codecs/pivco_huffman/arch/decode_pivco_arch.h" +#include "openzl/shared/portability.h" + +ZL_BEGIN_C_DECLS + +/** + * @returns The number of bytes required by the decoder scratch. + */ +size_t ZL_PivCoHuffmanDecode_scratchBytes(size_t dstSize, size_t blockSize); + +/** + * Decodes @p bitstream into @p dst using the same zstd-style Huffman weights + * that were provided to the encoder. The weights are not carried in + * @p bitstream; they must be supplied separately. + * + * @param scratch Working buffer of at least + * ZL_PivCoHuffmanDecode_scratchBytes(dstSize, blockSize) bytes. + * @param blockSize The number of decoded bytes per pivco block; must equal + * the block size the encoder used. + * @param kernels The kernel to use for decoding or NULL to use the default. + * + * @returns true on success; false if @p scratchBytes is too small, @p blockSize + * is out of range, or the bitstream is corrupt. + */ +bool ZL_PivCoHuffman_decode( + uint8_t* dst, + size_t dstSize, + uint8_t* scratch, + size_t scratchBytes, + const uint8_t* weights, + size_t weightsSize, + const uint8_t* bitstream, + size_t bitstreamSize, + size_t blockSize, + const ZL_PivCoHuffmanDecode* kernels); + +ZL_END_C_DECLS + +#endif diff --git a/src/openzl/codecs/pivco_huffman/encode_pivco_kernel.c b/src/openzl/codecs/pivco_huffman/encode_pivco_kernel.c new file mode 100644 index 000000000..59add3267 --- /dev/null +++ b/src/openzl/codecs/pivco_huffman/encode_pivco_kernel.c @@ -0,0 +1,305 @@ +// Copyright (c) Meta Platforms, Inc. and affiliates. + +#include "openzl/codecs/pivco_huffman/encode_pivco_kernel.h" + +#include +#include + +#include "openzl/codecs/common/bitstream/ff_bitstream.h" +#include "openzl/codecs/pivco_huffman/common_pivco_kernel.h" +#include "openzl/shared/bits.h" +#include "openzl/shared/utils.h" + +size_t ZL_PivCoHuffmanEncode_scratchElements(size_t srcSize, size_t blockSize) +{ + blockSize = ZL_MIN(srcSize, blockSize); + return 2 * (blockSize + ZL_PIVCO_HUFFMAN_SLOP); +} + +size_t ZL_PivCoHuffmanEncode_bound(size_t srcSize, size_t blockSize) +{ + if (srcSize == 0) { + return 0; + } + + assert(blockSize > 0); + assert(blockSize <= ZL_PIVCO_MAX_BLOCK_SIZE); + + // Worst-case size: the code data (at most ZL_PIVCO_MAX_TABLE_LOG bits per + // symbol) plus fixed per-node overhead for every node in every block, plus + // the kernels' over-write slop. + blockSize = ZL_MIN(srcSize, blockSize); + const size_t blocks = srcSize / blockSize + (srcSize % blockSize != 0); + const size_t countBits = (size_t)ZL_nextPow2(blockSize + 1); + const size_t overheadBitsPerNode = 7 + countBits; + const size_t overheadBitsPerBlock = + overheadBitsPerNode * ZL_PIVCO_MAX_TREE_NODES; + const size_t overheadBits = blocks * overheadBitsPerBlock; + const size_t overheadBytes = (overheadBits + 7) / 8; + + // Avoid a multiply by 8, so that if the bound overflows a size_t, we know + // the true-bound is >= SIZE_MAX. + assert(ZL_PIVCO_MAX_TABLE_LOG == 8 + 4); + const size_t dataBytes = srcSize + srcSize / 2; + + const size_t bound = dataBytes + overheadBytes + ZL_PIVCO_HUFFMAN_SLOP; + + if (bound < srcSize) { + return SIZE_MAX; + } else { + return bound; + } +} + +/** + * Writes @p value (which must be in [0, cap]) to @p writer using the fixed + * number of bits needed to represent @p cap, then flushes. The width depends + * only on @p cap, so the decoder recovers it without it being stored. + */ +static void +bitWriterEncodeCappedValue(ZS_BitCStreamFF* writer, size_t value, size_t cap) +{ + assert(value <= cap); + const size_t numBits = (size_t)ZL_nextPow2(cap + 1); + assert(numBits <= ZS_BITSTREAM_WRITE_MAX_BITS); + ZS_BitCStreamFF_write(writer, value, numBits); + ZS_BitCStreamFF_flush(writer); +} + +/** + * Reserves a byte-aligned region of @p numBits bits in @p writer for a bitmap + * the caller writes directly. @returns A pointer to the region, or NULL if it + * would not leave at least ZL_PIVCO_HUFFMAN_SLOP trailing bytes (which the + * kernels may over-write) -- i.e. the output buffer is exhausted. + */ +static uint8_t* bitWriterReserveBitmap(ZS_BitCStreamFF* writer, size_t numBits) +{ + uint8_t* const bitmap = ZS_BitCStreamFF_reserveAlignedBits(writer, numBits); + if (bitmap == NULL) { + return NULL; + } + + size_t const numBytes = (numBits + 7) / 8; + size_t const available = (size_t)(writer->end - bitmap); + if (available - numBytes < ZL_PIVCO_HUFFMAN_SLOP) { + return NULL; + } + + return bitmap; +} + +/** + * Translates @p numSymbols source @p symbols into ranks via @p symbolToRank, + * writing the result to @p ranks. This is the per-symbol mapping the rest of + * the encoder operates on. + * @pre Every source symbol is present in the tree (has a non-zero weight). + */ +static void buildRankStream( + uint8_t* restrict ranks, + const uint8_t* restrict symbols, + size_t numSymbols, + const uint8_t* restrict symbolToRank) +{ + for (size_t i = 0; i < numSymbols; ++i) { + ranks[i] = symbolToRank[symbols[i]]; + } +} + +/** + * Recursively encodes the pivco-tree node covering ranks [firstRank, rankEnd) + * at @p level into @p writer. @p nodeRanks holds the node's @p numRanks ranks + * (in input order); @p nodeScratch is working space for partitioning. + * + * A leaf node emits a flat bitmap (or nothing, for a constant leaf). An + * internal node partitions its ranks at the split rank into a 0-bit (left) and + * 1-bit (right) group, writes the partition bitmap and -- unless both children + * are constant -- the 1-bit count, then recurses into each non-constant child. + * The children reuse @p nodeRanks and @p nodeScratch as their own rank/scratch + * buffers without overlapping. + * + * @returns true on success, false if the output buffer is exhausted. + */ +static bool encodeNode( + const ZL_PivCoHuffmanTree* tree, + const ZL_PivCoHuffmanEncode* kernels, + ZS_BitCStreamFF* writer, + uint8_t* nodeRanks, + uint8_t* nodeScratch, + size_t numRanks, + size_t level, + size_t firstRank, + size_t rankEnd) +{ + if (ZL_PivCoHuffmanTree_rangeIsLeaf(tree, firstRank, rankEnd)) { + const size_t depth = ZL_PivCoHuffmanTree_leafFlatDepth(tree, firstRank); + // A constant (depth-0) leaf emits nothing -- the weights alone identify + // the symbol. A flat leaf packs `depth` bits per rank into a bitmap. + if (depth != 0) { + uint8_t* const bitmap = + bitWriterReserveBitmap(writer, numRanks * depth); + if (bitmap == NULL) { + return false; + } + kernels->packFlatDepth( + bitmap, depth, nodeRanks, numRanks, (uint8_t)firstRank); + ZS_BitCStreamFF_commitReservedBits(writer); + } + return true; + } + + // Internal node: split [firstRank, rankEnd) at splitRank into a left + // (0-bit) and right (1-bit) group, partition the ranks accordingly, and + // emit the partition bitmap (a constant side needs no separate rank + // buffer). + size_t const splitRank = + ZL_PivCoHuffmanTree_splitRank(tree, level, firstRank, rankEnd); + bool const lhsIsConstant = + ZL_PivCoHuffmanTree_rangeIsConstantLeaf(tree, firstRank, splitRank); + bool const rhsIsConstant = + ZL_PivCoHuffmanTree_rangeIsConstantLeaf(tree, splitRank, rankEnd); + + uint8_t* const lhsRanks = nodeScratch; + uint8_t* const rhsRanks = lhsIsConstant ? nodeScratch : nodeRanks; + + uint8_t* const bitmap = bitWriterReserveBitmap(writer, numRanks); + if (bitmap == NULL) { + return false; + } + + size_t numOnes = 0; + if (lhsIsConstant && rhsIsConstant) { + // Both children are constant leaves; the decoder recovers the partition + // straight from the bitmap, so numOnes is neither computed nor written. + kernels->partitionNone(bitmap, nodeRanks, numRanks, (uint8_t)splitRank); + } else if (lhsIsConstant) { + numOnes = kernels->partitionRight( + bitmap, rhsRanks, nodeRanks, numRanks, (uint8_t)splitRank); + } else if (rhsIsConstant) { + numOnes = kernels->partitionLeft( + bitmap, lhsRanks, nodeRanks, numRanks, (uint8_t)splitRank); + } else { + numOnes = kernels->partitionFull( + bitmap, + lhsRanks, + rhsRanks, + nodeRanks, + numRanks, + (uint8_t)splitRank); + } + ZS_BitCStreamFF_commitReservedBits(writer); + + if (!(lhsIsConstant && rhsIsConstant)) { + bitWriterEncodeCappedValue(writer, numOnes, numRanks); + } + + // Hand each child a working buffer carved from this node's two buffers: the + // partitioned ranks live in one, so each child's scratch is the unused tail + // of the other (or all of nodeRanks when its sibling is constant). + const size_t numZeros = numRanks - numOnes; + uint8_t* const lhsScratch = rhsIsConstant ? nodeRanks : rhsRanks + numOnes; + uint8_t* const rhsScratch = lhsIsConstant ? nodeRanks : lhsRanks + numZeros; + + bool success = true; + if (!lhsIsConstant) { + success &= encodeNode( + tree, + kernels, + writer, + lhsRanks, + lhsScratch, + numZeros, + level + 1, + firstRank, + splitRank); + } + if (!rhsIsConstant) { + success &= encodeNode( + tree, + kernels, + writer, + rhsRanks, + rhsScratch, + numOnes, + level + 1, + splitRank, + rankEnd); + } + return success; +} + +size_t ZL_PivCoHuffman_encode( + uint8_t* dst, + size_t dstCapacity, + uint8_t* scratch, + size_t scratchElements, + const uint8_t* weights, + size_t weightsSize, + int tableLog, + const uint8_t* src, + size_t srcSize, + size_t blockSize, + const ZL_PivCoHuffmanEncode* kernels) +{ + if (srcSize == 0 || tableLog == 0) { + // Empty or constant input ==> empty output + return 0; + } + + assert(srcSize != 0); + assert(blockSize != 0); + assert(blockSize <= ZL_PIVCO_MAX_BLOCK_SIZE); + assert(weightsSize != 0); + assert(weightsSize <= ZL_PIVCO_MAX_SYMBOLS); + assert(tableLog > 0 && tableLog <= ZL_PIVCO_MAX_TABLE_LOG); + assert(tableLog == ZL_PivCoHuffman_computeTableLog(weights, weightsSize)); + + if (scratchElements + < ZL_PivCoHuffmanEncode_scratchElements(srcSize, blockSize)) { + return SIZE_MAX; + } + + if (kernels == NULL) { + kernels = ZL_PivCoHuffmanEncode_select(NULL); + } + + ZL_PivCoHuffmanTree tree; + ZL_PivCoHuffmanTree_build(&tree, weights, weightsSize, tableLog); + assert(tree.numRanks != 0); + + // The largest block is min(srcSize, blockSize) bytes long. + blockSize = ZL_MIN(srcSize, blockSize); + + // Scratch is split into the per-block rank stream and the recursion's + // partitioning workspace, each cappedBlock + SLOP bytes. + uint8_t* const ranks = scratch; + uint8_t* const nodeScratch = scratch + blockSize + ZL_PIVCO_HUFFMAN_SLOP; + + ZS_BitCStreamFF writer = ZS_BitCStreamFF_init(dst, dstCapacity); + + for (size_t off = 0; off < srcSize; off += blockSize) { + const size_t blockLen = ZL_MIN(blockSize, srcSize - off); + buildRankStream(ranks, src + off, blockLen, tree.symbolToRank); + const bool success = encodeNode( + &tree, + kernels, + &writer, + ranks, + nodeScratch, + blockLen, + 0, + 0, + tree.numRanks); + if (!success) { + return SIZE_MAX; + } + } + + const ZL_Report ret = ZS_BitCStreamFF_finish(&writer); + if (ZL_isError(ret)) { + return SIZE_MAX; + } + const size_t dstSize = ZL_validResult(ret); + assert(dstSize != 0); + assert(dstSize + ZL_PIVCO_HUFFMAN_SLOP <= dstCapacity); + return dstSize; +} diff --git a/src/openzl/codecs/pivco_huffman/encode_pivco_kernel.h b/src/openzl/codecs/pivco_huffman/encode_pivco_kernel.h new file mode 100644 index 000000000..954351fe2 --- /dev/null +++ b/src/openzl/codecs/pivco_huffman/encode_pivco_kernel.h @@ -0,0 +1,65 @@ +// Copyright (c) Meta Platforms, Inc. and affiliates. + +#ifndef OPENZL_CODECS_PIVCO_HUFFMAN_ENCODE_PIVCO_KERNEL_H +#define OPENZL_CODECS_PIVCO_HUFFMAN_ENCODE_PIVCO_KERNEL_H + +#include +#include + +#include "openzl/codecs/pivco_huffman/arch/encode_pivco_arch.h" +#include "openzl/shared/portability.h" + +ZL_BEGIN_C_DECLS + +/** + * @returns the number of bytes required by the encoder scratch. + */ +size_t ZL_PivCoHuffmanEncode_scratchElements(size_t srcSize, size_t blockSize); + +/** + * @returns the output capacity including slop bytes to guarantee success + * (except for integer overflow cases, where SIZE_MAX is returned). + * + * Smaller destination sizes can be passed to the encoder safely, but then it + * may fail to compress and return an error. + * + * @pre blockSize > 0 && blockSize <= ZL_PIVCO_MAX_BLOCK_SIZE. + */ +size_t ZL_PivCoHuffmanEncode_bound(size_t srcSize, size_t blockSize); + +/** + * Encodes @p src using zstd-style Huffman weights and returns the number of + * bytes written to @p dst. The weights are not emitted into @p dst; the decoder + * must receive the same weights separately. + * + * @param kernels The kernel to use for encoding or NULL to use the default. + * + * @param blockSize The number of input bytes coded per pivco block. Must + * match the block size the decoder uses. srcSize > blockSize produces multiple + * blocks. + * + * @pre blockSize > 0 && blockSize <= ZL_PIVCO_MAX_BLOCK_SIZE. + * @pre scratchElements >= + * ZL_PivCoHuffmanEncode_scratchElements(srcSize, blockSize). + * @pre weights describes a valid zstd-style Huffman weight table. + * @pre tableLog == ZL_PivCoHuffman_computeTableLog(weights, weightsSize). + * @pre Every symbol in src is less than weightsSize and has non-zero weight. + * + * @returns The encoded size in bytes or SIZE_MAX upon failure. + */ +size_t ZL_PivCoHuffman_encode( + uint8_t* dst, + size_t dstCapacity, + uint8_t* scratch, + size_t scratchElements, + const uint8_t* weights, + size_t weightsSize, + int tableLog, + const uint8_t* src, + size_t srcSize, + size_t blockSize, + const ZL_PivCoHuffmanEncode* kernels); + +ZL_END_C_DECLS + +#endif diff --git a/src/openzl/shared/cpu.h b/src/openzl/shared/cpu.h index 9f84e58c4..fe4289c2f 100644 --- a/src/openzl/shared/cpu.h +++ b/src/openzl/shared/cpu.h @@ -201,6 +201,13 @@ B(avx512vl, 31) #define C(name, bit) X(name, f7c, bit) C(prefetchwt1, 0) C(avx512vbmi, 1) +C(avx512vbmi2, 6) +C(vaes, 9) +C(vpclmulqdq, 10) +C(avx512vnni, 11) +C(avx512bitalg, 12) +C(avx512vpopcntdq, 14) +C(rdpid, 22) #undef C #undef X diff --git a/tests/unittest/common/test_bitstream.cpp b/tests/unittest/common/test_bitstream.cpp index b6f4d9e20..93bf3c096 100644 --- a/tests/unittest/common/test_bitstream.cpp +++ b/tests/unittest/common/test_bitstream.cpp @@ -1,9 +1,13 @@ // Copyright (c) Meta Platforms, Inc. and affiliates. #include +#include +#include #include #include #include +#include +#include #include @@ -257,7 +261,13 @@ class RoundTripTest { ZS_BitDStreamFF_reload(&bits); } - return !ZL_isError(ZS_BitDStreamFF_finish(&bits)); + ZL_Report const report = ZS_BitDStreamFF_finish(&bits); + if (ZL_isError(report)) { + return false; + } + // finish() reports the exact number of bytes consumed, which must equal + // the encoded size regardless of any extra trailing capacity. + return ZL_validResult(report) == bitSize; } template @@ -440,6 +450,165 @@ INSTANTIATE_TEST_SUITE_P( BitstreamImpl::ZS, BitstreamImpl::ZS_BF)); +// Mask of the low @p nbBits bits. +size_t lowMask(size_t nbBits) +{ + if (nbBits == 0) + return 0; + if (nbBits >= sizeof(size_t) * 8) + return ~(size_t)0; + return ((size_t)1 << nbBits) - 1; +} + +// Round-trips `leading` bits, a byte-aligned raw region holding `region` bits, +// then `trailing` bits. The region is written with the reserve/commit aligned +// API and read back with ZS_BitDStreamFF_popAlignedBits, mirroring how the +// PivCo-Huffman kernels embed raw bitmaps in the stream. Verifies every bit +// and the raw region survive. When `extra` is set the decoder gets slack +// capacity past the encoded size (the large-stream path); otherwise it sees the +// exact end, exercising the near-end and short-stream paths in popAlignedBits. +// +// `reload` reloads the window after every field read. The kernels do not do +// this (popAlignedBits refills the window itself), but it must still be a valid +// thing to do: popAlignedBits has to recover the exact bit position from any +// reloaded stream state, including a short stream whose reload ran past the +// end. +void roundTripAlignedField( + const std::vector>& leading, + const std::vector& region, + const std::vector>& trailing, + bool extra, + bool reload) +{ + size_t const regionBytes = (region.size() + 7) / 8; + size_t totalBits = 0; + for (auto const& w : leading) + totalBits += w.second; + for (auto const& w : trailing) + totalBits += w.second; + std::vector buf(totalBits / 8 + regionBytes + 64, 0); + + ZS_BitCStreamFF writer = ZS_BitCStreamFF_init(buf.data(), buf.size()); + for (auto const& [value, nbBits] : leading) { + ZS_BitCStreamFF_write(&writer, value, nbBits); + ZS_BitCStreamFF_flush(&writer); + } + uint8_t* const slot = + ZS_BitCStreamFF_reserveAlignedBits(&writer, region.size()); + ASSERT_NE(slot, nullptr); + std::memset(slot, 0, regionBytes); + for (size_t i = 0; i < region.size(); ++i) { + if (region[i]) + slot[i / 8] |= (uint8_t)(1u << (i & 7)); + } + ZS_BitCStreamFF_commitReservedBits(&writer); + for (auto const& [value, nbBits] : trailing) { + ZS_BitCStreamFF_write(&writer, value, nbBits); + ZS_BitCStreamFF_flush(&writer); + } + ZL_Report const report = ZS_BitCStreamFF_finish(&writer); + ASSERT_ZS_VALID(report); + size_t const encodedSize = ZL_validResult(report); + + ZS_BitDStreamFF reader = + ZS_BitDStreamFF_init(buf.data(), extra ? buf.size() : encodedSize); + for (auto const& [value, nbBits] : leading) { + EXPECT_EQ( + ZS_BitDStreamFF_read(&reader, nbBits), value & lowMask(nbBits)); + if (reload) + ZS_BitDStreamFF_reload(&reader); + } + uint8_t const* const popped = + ZS_BitDStreamFF_popAlignedBits(&reader, region.size()); + ASSERT_NE(popped, nullptr); + for (size_t i = 0; i < region.size(); ++i) { + EXPECT_EQ((popped[i / 8] >> (i & 7)) & 1u, region[i] ? 1u : 0u) + << "region bit " << i; + } + for (auto const& [value, nbBits] : trailing) { + EXPECT_EQ( + ZS_BitDStreamFF_read(&reader, nbBits), value & lowMask(nbBits)); + if (reload) + ZS_BitDStreamFF_reload(&reader); + } + ASSERT_ZS_VALID(ZS_BitDStreamFF_finish(&reader)); +} + +std::vector makeRegionBits(size_t nbBits) +{ + std::vector bits(nbBits); + for (size_t i = 0; i < nbBits; ++i) + bits[i] = (((i * 7) + 3) % 5) < 2; + return bits; +} + +TEST(BitstreamAlignedBitsTest, ReserveCommitPopRoundTrip) +{ + using BitWrites = std::vector>; + BitWrites const someBits = { + { 0x1, 1 }, { 0x2, 3 }, { 0xABCD, 13 }, { 0x7F, 7 } + }; + + struct Scenario { + const char* name; + BitWrites leading; + size_t regionBits; + BitWrites trailing; + }; + std::vector const scenarios = { + // Region alone, with nothing around it. + { "empty_region", {}, 0, {} }, + { "byte_region_only", {}, 8, {} }, + { "aligned_region_only", {}, 64, {} }, + // Region followed by bits (window stays mid-stream). + { "region_then_bits", {}, 24, someBits }, + // Bits then a region that ends the stream (window lands at the end). + { "bits_then_region", someBits, 40, {} }, + // Region surrounded on both sides. + { "bits_region_bits", someBits, 32, someBits }, + // Unaligned leading bits force the region past padding. + { "unaligned_leading", { { 0x5, 3 } }, 16, { { 0x9, 5 } } }, + // Region whose bit count is not a multiple of 8 (partial trailing + // byte). + { "partial_region", someBits, 19, someBits }, + { "partial_region_at_end", someBits, 13, {} }, + // Tiny total so the exact-capacity decode hits the short-stream path. + { "tiny_small_stream", { { 0x1, 2 } }, 8, {} }, + // Large region exercises the steady-state window refill. + { "large_region", someBits, 1000, someBits }, + }; + + for (auto const& s : scenarios) { + for (bool extra : { true, false }) { + for (bool reload : { false, true }) { + SCOPED_TRACE( + std::string(s.name) + (extra ? " extra" : " exact") + + (reload ? " reload" : " noreload")); + roundTripAlignedField( + s.leading, + makeRegionBits(s.regionBits), + s.trailing, + extra, + reload); + } + } + } +} + +TEST(BitstreamAlignedBitsTest, ReserveReturnsNullWhenBufferTooSmall) +{ + std::vector buf(4, 0); + ZS_BitCStreamFF writer = ZS_BitCStreamFF_init(buf.data(), buf.size()); + EXPECT_EQ(ZS_BitCStreamFF_reserveAlignedBits(&writer, 8 * 64), nullptr); +} + +TEST(BitstreamAlignedBitsTest, PopReturnsNullPastEndOfStream) +{ + std::vector buf(8, 0xAB); + ZS_BitDStreamFF reader = ZS_BitDStreamFF_init(buf.data(), buf.size()); + EXPECT_EQ(ZS_BitDStreamFF_popAlignedBits(&reader, 8 * 64), nullptr); +} + struct BenchmarkResult { int64_t encodeNs{ 0 }; int64_t decodeNs{ 0 }; diff --git a/tests/unittest/transforms/test_pivco_huffman.cpp b/tests/unittest/transforms/test_pivco_huffman.cpp new file mode 100644 index 000000000..e0f0525bd --- /dev/null +++ b/tests/unittest/transforms/test_pivco_huffman.cpp @@ -0,0 +1,969 @@ +// Copyright (c) Meta Platforms, Inc. and affiliates. + +#include +#include +#include +#include +#include +#include + +#include + +#include "tests/utils.h" + +#include "openzl/codecs/pivco_huffman/arch/common_pivco_arch.h" +#include "openzl/codecs/pivco_huffman/arch/decode_pivco_arch.h" +#include "openzl/codecs/pivco_huffman/arch/encode_pivco_arch.h" +#include "openzl/codecs/pivco_huffman/common_pivco_kernel.h" +#include "openzl/codecs/pivco_huffman/decode_pivco_kernel.h" +#include "openzl/codecs/pivco_huffman/encode_pivco_kernel.h" + +namespace { + +struct EncodeArch { + const char* name; + const ZL_PivCoHuffmanEncode* kernels; +}; + +struct DecodeArch { + const char* name; + const ZL_PivCoHuffmanDecode* kernels; +}; + +std::vector supportedEncodeArchs() +{ + ZL_cpuid_t const cpuid = ZL_cpuid(); + std::vector out; + std::vector const archs = { + { "generic", &ZL_PivCoHuffmanEncode_generic }, + }; + for (auto const& arch : archs) { + if (arch.kernels->supported(&cpuid)) { + out.push_back(arch); + } + } + return out; +} + +std::vector supportedDecodeArchs() +{ + ZL_cpuid_t const cpuid = ZL_cpuid(); + std::vector out; + std::vector const archs = { + { "generic", &ZL_PivCoHuffmanDecode_generic }, + }; + for (auto const& arch : archs) { + if (arch.kernels->supported(&cpuid)) { + out.push_back(arch); + } + } + return out; +} + +template +std::vector firstN(const std::vector& in, size_t n) +{ + return std::vector(in.begin(), in.begin() + n); +} + +size_t bitmapBytes(size_t bits) +{ + return (bits + 7) / 8; +} + +void setBitmapBit(std::vector& bitmap, size_t bit) +{ + bitmap[bit / 8] |= (uint8_t)(1u << (bit & 7)); +} + +enum class TopBitPattern { + AllZero, + AllOne, + Mixed, +}; + +uint8_t makeRank(size_t i, size_t size, TopBitPattern pattern) +{ + switch (pattern) { + case TopBitPattern::AllZero: + return 0x10; + case TopBitPattern::AllOne: + return 0x90; + case TopBitPattern::Mixed: + return (((i * 37) + size) % 7) < 3 ? 0x90 : 0x10; + } + return 0; +} + +std::vector makeRanks(size_t size, TopBitPattern pattern) +{ + std::vector ranks(size + ZL_PIVCO_HUFFMAN_SLOP, 0xA5); + for (size_t i = 0; i < size; ++i) { + ranks[i] = makeRank(i, size, pattern); + } + return ranks; +} + +std::vector +makeFlatDepthRanks(size_t size, size_t depth, uint8_t rankBegin) +{ + std::vector ranks(size + ZL_PIVCO_HUFFMAN_SLOP, 0xA5); + uint8_t const symbolMask = (uint8_t)((1u << depth) - 1u); + for (size_t i = 0; i < size; ++i) { + ranks[i] = (uint8_t)(rankBegin + (uint8_t)(i & symbolMask)); + } + return ranks; +} + +struct PartitionReference { + std::vector bitmap; + std::vector lhs; + std::vector rhs; +}; + +PartitionReference referencePartition( + const std::vector& ranks, + size_t size, + uint8_t rightRank) +{ + PartitionReference out; + out.bitmap.assign(bitmapBytes(size), 0); + for (size_t i = 0; i < size; ++i) { + uint8_t const rank = ranks[i]; + if (rank >= rightRank) { + setBitmapBit(out.bitmap, i); + out.rhs.push_back(rank); + } else { + out.lhs.push_back(rank); + } + } + return out; +} + +std::vector referencePackFlatDepth( + const std::vector& ranks, + size_t size, + size_t depth, + uint8_t rankBegin) +{ + std::vector bitmap(bitmapBytes(size * depth), 0); + for (size_t i = 0; i < size; ++i) { + uint8_t const idx = (uint8_t)(ranks[i] - rankBegin); + for (size_t bit = 0; bit < depth; ++bit) { + if ((idx & (1u << bit)) != 0) { + setBitmapBit(bitmap, i * depth + bit); + } + } + } + return bitmap; +} + +std::vector referenceMergeFlatDepth( + const std::vector& bitmap, + size_t outSize, + size_t depth, + const std::vector& symbols) +{ + std::vector out(outSize); + uint8_t const mask = (uint8_t)((1u << depth) - 1); + for (size_t i = 0; i < outSize; ++i) { + size_t const bitOffset = i * depth; + size_t const byte = bitOffset / 8; + size_t const shift = bitOffset & 7; + uint16_t bits = bitmap[byte]; + if (byte + 1 < bitmap.size()) { + bits |= (uint16_t)(bitmap[byte + 1] << 8); + } + out[i] = symbols[(bits >> shift) & mask]; + } + return out; +} + +std::vector boundarySizes() +{ + return { + 0, 1, 2, 3, 4, 5, 7, 8, 9, 15, 16, 17, 31, 32, + 33, 63, 64, 65, 127, 128, 129, 255, 256, 257, 511, 512, 513, + }; +} + +std::vector makeMergeBits(size_t size, TopBitPattern pattern) +{ + std::vector bits(size); + for (size_t i = 0; i < size; ++i) { + switch (pattern) { + case TopBitPattern::AllZero: + bits[i] = false; + break; + case TopBitPattern::AllOne: + bits[i] = true; + break; + case TopBitPattern::Mixed: + bits[i] = (((i * 11) + size) % 5) < 2; + break; + } + } + return bits; +} + +std::vector bitmapFromBits( + const std::vector& bits, + size_t extraCapacity) +{ + size_t const exactBytes = bitmapBytes(bits.size()); + std::vector bitmap(exactBytes + extraCapacity, 0); + for (size_t i = 0; i < bits.size(); ++i) { + if (bits[i]) { + setBitmapBit(bitmap, i); + } + } + for (size_t i = exactBytes; i < bitmap.size(); ++i) { + bitmap[i] = 0xA5; + } + return bitmap; +} + +struct MergeInputs { + std::vector expected; + std::vector lhs; + std::vector rhs; + size_t ones = 0; +}; + +MergeInputs makeMergeInputs(const std::vector& bits) +{ + MergeInputs out; + out.expected.reserve(bits.size()); + for (size_t i = 0; i < bits.size(); ++i) { + if (bits[i]) { + uint8_t const value = + (uint8_t)(0x80u + ((out.rhs.size() * 13) & 0x7F)); + out.rhs.push_back(value); + out.expected.push_back(value); + ++out.ones; + } else { + uint8_t const value = + (uint8_t)(0x10u + ((out.lhs.size() * 17) & 0x6F)); + out.lhs.push_back(value); + out.expected.push_back(value); + } + } + out.lhs.resize(out.lhs.size() + ZL_PIVCO_HUFFMAN_SLOP, 0xEE); + out.rhs.resize(out.rhs.size() + ZL_PIVCO_HUFFMAN_SLOP, 0xDD); + return out; +} + +struct HuffmanScenario { + const char* name; + std::vector weights; + std::vector alphabet; +}; + +std::vector makeWeights( + size_t size, + std::initializer_list> entries) +{ + std::vector weights(size); + for (auto const& entry : entries) { + weights[entry.first] = entry.second; + } + return weights; +} + +std::vector huffmanScenarios() +{ + return { + { + "single", + makeWeights(43, { { 42, 1 } }), + { 42 }, + }, + { + "two_equal", + makeWeights(2, { { 0, 1 }, { 1, 1 } }), + { 0, 1 }, + }, + { + "short_plus_pair", + makeWeights(3, { { 0, 2 }, { 1, 1 }, { 2, 1 } }), + { 0, 1, 2 }, + }, + { + "short_plus_flat4", + makeWeights( + 5, + { { 0, 3 }, { 1, 1 }, { 2, 1 }, { 3, 1 }, { 4, 1 } }), + { 0, 1, 2, 3, 4 }, + }, + { + "two_flat_children", + makeWeights( + 6, + { { 0, 3 }, + { 1, 3 }, + { 2, 2 }, + { 3, 2 }, + { 4, 2 }, + { 5, 2 } }), + { 0, 1, 2, 3, 4, 5 }, + }, + }; +} + +std::vector makeHuffmanData( + const std::vector& alphabet, + size_t size) +{ + std::vector data(size); + for (size_t i = 0; i < size; ++i) { + data[i] = alphabet[((i * 37) + size) % alphabet.size()]; + } + return data; +} + +void expectPivCoRoundTrip( + const HuffmanScenario& scenario, + const std::vector& data, + const ZL_PivCoHuffmanEncode* encodeKernels, + const ZL_PivCoHuffmanDecode* decodeKernels, + size_t blockSize = ZL_PIVCO_DEFAULT_BLOCK_SIZE) +{ + size_t const encodeScratchElements = + ZL_PivCoHuffmanEncode_scratchElements(data.size(), blockSize); + size_t const decodeScratchBytes = + ZL_PivCoHuffmanDecode_scratchBytes(data.size(), blockSize); + std::vector encodeScratch(encodeScratchElements); + std::vector encoded( + ZL_PivCoHuffmanEncode_bound(data.size(), blockSize)); + int const tableLog = ZL_PivCoHuffman_computeTableLog( + scenario.weights.data(), scenario.weights.size()); + ASSERT_GE(tableLog, 0); + + // A constant input encodes to an empty bitstream (encodedSize == 0); only + // SIZE_MAX signals failure. + size_t const encodedSize = ZL_PivCoHuffman_encode( + encoded.data(), + encoded.size(), + encodeScratch.data(), + encodeScratch.size(), + scenario.weights.data(), + scenario.weights.size(), + tableLog, + data.data(), + data.size(), + blockSize, + encodeKernels); + ASSERT_NE(encodedSize, SIZE_MAX); + ASSERT_LE(encodedSize, encoded.size()); + encoded.resize(encodedSize); + + std::vector decodeScratch(decodeScratchBytes); + std::vector decoded(data.size(), 0xCC); + ASSERT_TRUE(ZL_PivCoHuffman_decode( + decoded.data(), + decoded.size(), + decodeScratch.data(), + decodeScratch.size(), + scenario.weights.data(), + scenario.weights.size(), + encoded.data(), + encoded.size(), + blockSize, + decodeKernels)); + EXPECT_EQ(decoded, data); +} + +// Encodes @p data with the generic kernel and returns the encoded bitstream +// (resized to its exact length). Used by the decode-rejection tests, which need +// the raw encoded bytes to corrupt before decoding. +std::vector encodePivCo( + const HuffmanScenario& scenario, + const std::vector& data, + size_t blockSize = ZL_PIVCO_DEFAULT_BLOCK_SIZE) +{ + std::vector encodeScratch( + ZL_PivCoHuffmanEncode_scratchElements(data.size(), blockSize)); + std::vector encoded( + ZL_PivCoHuffmanEncode_bound(data.size(), blockSize)); + int const tableLog = ZL_PivCoHuffman_computeTableLog( + scenario.weights.data(), scenario.weights.size()); + EXPECT_GE(tableLog, 0); + size_t const encodedSize = ZL_PivCoHuffman_encode( + encoded.data(), + encoded.size(), + encodeScratch.data(), + encodeScratch.size(), + scenario.weights.data(), + scenario.weights.size(), + tableLog, + data.data(), + data.size(), + blockSize, + &ZL_PivCoHuffmanEncode_generic); + EXPECT_NE(encodedSize, SIZE_MAX); + encoded.resize(encodedSize == SIZE_MAX ? 0 : encodedSize); + return encoded; +} + +void expectBitmapPrefix( + const std::vector& actual, + const std::vector& expected) +{ + ASSERT_LE(expected.size(), actual.size()); + EXPECT_EQ(firstN(actual, expected.size()), expected); +} + +void expectBytesAfterPrefix( + const std::vector& actual, + size_t prefixSize, + uint8_t sentinel) +{ + ASSERT_LE(prefixSize, actual.size()); + std::vector const actualTail( + actual.begin() + prefixSize, actual.end()); + std::vector const expectedTail(actualTail.size(), sentinel); + EXPECT_EQ(actualTail, expectedTail); +} + +} // namespace + +TEST(PivCoHuffmanArchTest, PartitionKernelsMatchReference) +{ + for (auto const& arch : supportedEncodeArchs()) { + ASSERT_NE(arch.kernels->partitionFull, nullptr) << arch.name; + ASSERT_NE(arch.kernels->partitionLeft, nullptr) << arch.name; + ASSERT_NE(arch.kernels->partitionRight, nullptr) << arch.name; + ASSERT_NE(arch.kernels->partitionNone, nullptr) << arch.name; + + for (TopBitPattern pattern : { TopBitPattern::AllZero, + TopBitPattern::AllOne, + TopBitPattern::Mixed }) { + for (size_t size : boundarySizes()) { + SCOPED_TRACE( + std::string(arch.name) + + " size=" + std::to_string(size)); + uint8_t const rightRank = 0x80; + auto const ranks = makeRanks(size, pattern); + auto const ref = referencePartition(ranks, size, rightRank); + size_t const bytes = bitmapBytes(size); + + std::vector bitmap( + bytes + ZL_PIVCO_HUFFMAN_SLOP, 0xCC); + std::vector lhs(size + ZL_PIVCO_HUFFMAN_SLOP, 0xEE); + std::vector rhs(size + ZL_PIVCO_HUFFMAN_SLOP, 0xDD); + size_t const ones = arch.kernels->partitionFull( + bitmap.data(), + lhs.data(), + rhs.data(), + ranks.data(), + size, + rightRank); + EXPECT_EQ(ones, ref.rhs.size()); + expectBitmapPrefix(bitmap, ref.bitmap); + EXPECT_EQ(firstN(lhs, ref.lhs.size()), ref.lhs); + EXPECT_EQ(firstN(rhs, ref.rhs.size()), ref.rhs); + + std::fill(bitmap.begin(), bitmap.end(), 0xCC); + std::fill(lhs.begin(), lhs.end(), 0xEE); + EXPECT_EQ( + arch.kernels->partitionLeft( + bitmap.data(), + lhs.data(), + ranks.data(), + size, + rightRank), + ref.rhs.size()); + expectBitmapPrefix(bitmap, ref.bitmap); + EXPECT_EQ(firstN(lhs, ref.lhs.size()), ref.lhs); + + std::fill(bitmap.begin(), bitmap.end(), 0xCC); + std::fill(rhs.begin(), rhs.end(), 0xDD); + EXPECT_EQ( + arch.kernels->partitionRight( + bitmap.data(), + rhs.data(), + ranks.data(), + size, + rightRank), + ref.rhs.size()); + expectBitmapPrefix(bitmap, ref.bitmap); + EXPECT_EQ(firstN(rhs, ref.rhs.size()), ref.rhs); + + std::fill(bitmap.begin(), bitmap.end(), 0xCC); + arch.kernels->partitionNone( + bitmap.data(), ranks.data(), size, rightRank); + expectBitmapPrefix(bitmap, ref.bitmap); + } + } + } +} + +TEST(PivCoHuffmanArchTest, PartitionFullSupportsDocumentedInputAliasing) +{ + for (auto const& arch : supportedEncodeArchs()) { + for (TopBitPattern pattern : { TopBitPattern::AllZero, + TopBitPattern::AllOne, + TopBitPattern::Mixed }) { + for (size_t size : boundarySizes()) { + SCOPED_TRACE( + std::string(arch.name) + + " size=" + std::to_string(size)); + uint8_t const rightRank = 0x80; + auto const ranks = makeRanks(size, pattern); + auto const ref = referencePartition(ranks, size, rightRank); + size_t const bytes = bitmapBytes(size); + + std::vector bitmap( + bytes + ZL_PIVCO_HUFFMAN_SLOP, 0xCC); + std::vector lhsAlias = ranks; + std::vector rhs(size + ZL_PIVCO_HUFFMAN_SLOP, 0xDD); + EXPECT_EQ( + arch.kernels->partitionFull( + bitmap.data(), + lhsAlias.data(), + rhs.data(), + lhsAlias.data(), + size, + rightRank), + ref.rhs.size()); + expectBitmapPrefix(bitmap, ref.bitmap); + EXPECT_EQ(firstN(lhsAlias, ref.lhs.size()), ref.lhs); + EXPECT_EQ(firstN(rhs, ref.rhs.size()), ref.rhs); + + std::fill(bitmap.begin(), bitmap.end(), 0xCC); + std::vector lhs(size + ZL_PIVCO_HUFFMAN_SLOP, 0xEE); + std::vector rhsAlias = ranks; + EXPECT_EQ( + arch.kernels->partitionFull( + bitmap.data(), + lhs.data(), + rhsAlias.data(), + rhsAlias.data(), + size, + rightRank), + ref.rhs.size()); + expectBitmapPrefix(bitmap, ref.bitmap); + EXPECT_EQ(firstN(lhs, ref.lhs.size()), ref.lhs); + EXPECT_EQ(firstN(rhsAlias, ref.rhs.size()), ref.rhs); + } + } + } +} + +TEST(PivCoHuffmanArchTest, FlatDepthPackAndMergeMatchReference) +{ + auto const sizes = boundarySizes(); + for (size_t depth = 1; depth <= 8; ++depth) { + std::vector symbols((size_t)1 << depth); + for (size_t i = 0; i < symbols.size(); ++i) { + symbols[i] = (uint8_t)((i * 37 + depth * 11) & 0xFF); + } + + for (size_t size : sizes) { + uint8_t const rankBegin = depth == 8 ? 0 : 17; + auto const ranks = makeFlatDepthRanks(size, depth, rankBegin); + auto const packed = + referencePackFlatDepth(ranks, size, depth, rankBegin); + auto const merged = + referenceMergeFlatDepth(packed, size, depth, symbols); + + for (auto const& arch : supportedEncodeArchs()) { + ASSERT_NE(arch.kernels->packFlatDepth, nullptr) << arch.name; + SCOPED_TRACE( + std::string("pack ") + arch.name + + " depth=" + std::to_string(depth) + + " size=" + std::to_string(size)); + size_t const guardOffset = + packed.size() + ZL_PIVCO_HUFFMAN_SLOP; + std::vector bitmap( + guardOffset + ZL_PIVCO_HUFFMAN_SLOP, 0xCC); + arch.kernels->packFlatDepth( + bitmap.data(), depth, ranks.data(), size, rankBegin); + expectBitmapPrefix(bitmap, packed); + expectBytesAfterPrefix(bitmap, guardOffset, 0xCC); + } + + for (auto const& arch : supportedDecodeArchs()) { + ASSERT_NE(arch.kernels->mergeFlatDepth, nullptr) << arch.name; + SCOPED_TRACE( + std::string("merge ") + arch.name + + " depth=" + std::to_string(depth) + + " size=" + std::to_string(size)); + std::vector out(size + ZL_PIVCO_HUFFMAN_SLOP, 0xCC); + arch.kernels->mergeFlatDepth( + out.data(), + size, + out.size(), + packed.data(), + packed.size(), + depth, + symbols.data()); + EXPECT_EQ(firstN(out, size), merged); + } + } + } +} + +TEST(PivCoHuffmanArchTest, MergeKernelsMatchReference) +{ + for (auto const& arch : supportedDecodeArchs()) { + ASSERT_NE(arch.kernels->mergeVectorVector, nullptr) << arch.name; + ASSERT_NE(arch.kernels->mergeConstantVector, nullptr) << arch.name; + ASSERT_NE(arch.kernels->mergeVectorConstant, nullptr) << arch.name; + + for (TopBitPattern pattern : { TopBitPattern::AllZero, + TopBitPattern::AllOne, + TopBitPattern::Mixed }) { + for (size_t size : boundarySizes()) { + auto const bits = makeMergeBits(size, pattern); + auto const inputs = makeMergeInputs(bits); + + // The expected outputs depend only on `bits`/`inputs`, so build + // them once here rather than inside the capacity loops below. + uint8_t const lhsConstant = 0x3C; + uint8_t const rhsConstant = 0xC3; + std::vector expectedConstantVector(size); + std::vector expectedVectorConstant(size); + for (size_t i = 0, lhsPos = 0, rhsPos = 0; i < size; ++i) { + expectedConstantVector[i] = + bits[i] ? inputs.rhs[rhsPos++] : lhsConstant; + expectedVectorConstant[i] = + bits[i] ? rhsConstant : inputs.lhs[lhsPos++]; + } + + for (size_t outExtra : + { (size_t)0, (size_t)ZL_PIVCO_HUFFMAN_SLOP }) { + for (size_t bitmapExtra : + { (size_t)0, (size_t)ZL_PIVCO_HUFFMAN_SLOP }) { + SCOPED_TRACE( + std::string(arch.name) + + " size=" + std::to_string(size) + " outExtra=" + + std::to_string(outExtra) + " bitmapExtra=" + + std::to_string(bitmapExtra)); + auto const bitmap = bitmapFromBits(bits, bitmapExtra); + + std::vector out(size + outExtra, 0xCC); + EXPECT_EQ( + arch.kernels->mergeVectorVector( + out.data(), + out.size(), + bitmap.data(), + bitmap.size(), + inputs.lhs.data(), + inputs.lhs.size() + - ZL_PIVCO_HUFFMAN_SLOP, + inputs.rhs.data(), + inputs.rhs.size() + - ZL_PIVCO_HUFFMAN_SLOP), + inputs.ones); + EXPECT_EQ(firstN(out, size), inputs.expected); + + std::fill(out.begin(), out.end(), 0xCC); + EXPECT_EQ( + arch.kernels->mergeConstantVector( + out.data(), + out.size(), + bitmap.data(), + bitmap.size(), + lhsConstant, + inputs.lhs.size() + - ZL_PIVCO_HUFFMAN_SLOP, + inputs.rhs.data(), + inputs.rhs.size() + - ZL_PIVCO_HUFFMAN_SLOP), + inputs.ones); + EXPECT_EQ(firstN(out, size), expectedConstantVector); + + std::fill(out.begin(), out.end(), 0xCC); + EXPECT_EQ( + arch.kernels->mergeVectorConstant( + out.data(), + out.size(), + bitmap.data(), + bitmap.size(), + inputs.lhs.data(), + inputs.lhs.size() + - ZL_PIVCO_HUFFMAN_SLOP, + rhsConstant, + inputs.rhs.size() + - ZL_PIVCO_HUFFMAN_SLOP), + inputs.ones); + EXPECT_EQ(firstN(out, size), expectedVectorConstant); + } + } + } + } + } +} + +TEST(PivCoHuffmanKernelTest, RoundTripsAcrossKernelImplementations) +{ + auto const encodeArchs = supportedEncodeArchs(); + auto const decodeArchs = supportedDecodeArchs(); + for (auto const& scenario : huffmanScenarios()) { + for (size_t size : { (size_t)1, + (size_t)2, + (size_t)3, + (size_t)7, + (size_t)63, + (size_t)128, + (size_t)257, + (size_t)4096 }) { + auto const data = makeHuffmanData(scenario.alphabet, size); + for (auto const& encodeArch : encodeArchs) { + for (auto const& decodeArch : decodeArchs) { + SCOPED_TRACE( + std::string(scenario.name) + + " size=" + std::to_string(size) + " encode=" + + encodeArch.name + " decode=" + decodeArch.name); + expectPivCoRoundTrip( + scenario, + data, + encodeArch.kernels, + decodeArch.kernels); + } + } + } + } +} + +TEST(PivCoHuffmanKernelTest, RoundTripsLoremRepro) +{ + std::vector const data( + openzl::tests::kLoremTestInput.begin(), + openzl::tests::kLoremTestInput.end()); + ASSERT_FALSE(data.empty()); + + // Exact HUF weights the binding computes for kLoremTestInput (tableLog 10). + HuffmanScenario scenario; + scenario.name = "lorem"; + scenario.weights = { + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 3, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 8, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 5, 0, 5, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 2, 0, 1, 1, 1, 1, 0, 0, 2, 0, 0, 1, 3, 3, 0, 2, 1, 0, 1, + 0, 1, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 7, 4, 6, 5, 8, 4, 4, 4, + 7, 2, 0, 7, 6, 7, 6, 5, 5, 7, 7, 7, 7, 5, 0, 2, + }; + ASSERT_EQ( + ZL_PivCoHuffman_computeTableLog( + scenario.weights.data(), scenario.weights.size()), + 10); + + for (auto const& enc : supportedEncodeArchs()) { + for (auto const& dec : supportedDecodeArchs()) { + expectPivCoRoundTrip(scenario, data, enc.kernels, dec.kernels); + } + } +} + +TEST(PivCoHuffmanKernelTest, RoundTripsLargeAlphabet) +{ + // Valid complete code: single leaves at lengths 4,5,6,7 plus 226 leaves at + // length 8, giving tableLog == 8 and a large, deep tree. + HuffmanScenario scenario; + scenario.name = "large_alphabet"; + size_t const numSymbols = 230; + scenario.weights.assign(numSymbols, 1); + scenario.weights[0] = 5; + scenario.weights[1] = 4; + scenario.weights[2] = 3; + scenario.weights[3] = 2; + for (size_t s = 0; s < numSymbols; ++s) { + scenario.alphabet.push_back((uint8_t)s); + } + ASSERT_EQ( + ZL_PivCoHuffman_computeTableLog( + scenario.weights.data(), scenario.weights.size()), + 8); + + for (size_t size : { (size_t)1000, (size_t)2758, (size_t)5000 }) { + auto const data = makeHuffmanData(scenario.alphabet, size); + for (auto const& enc : supportedEncodeArchs()) { + for (auto const& dec : supportedDecodeArchs()) { + expectPivCoRoundTrip(scenario, data, enc.kernels, dec.kernels); + } + } + } +} + +TEST(PivCoHuffmanKernelTest, RoundTripsMultipleBlocks) +{ + HuffmanScenario const scenario = huffmanScenarios().back(); + auto const data = makeHuffmanData( + scenario.alphabet, ZL_PIVCO_DEFAULT_BLOCK_SIZE + 123); + + expectPivCoRoundTrip( + scenario, + data, + &ZL_PivCoHuffmanEncode_generic, + &ZL_PivCoHuffmanDecode_generic); +} + +TEST(PivCoHuffmanKernelTest, RoundTripsCustomBlockSize) +{ + // A small block size forces many blocks at a non-default size; the encoder + // and decoder must agree on the block boundaries. + HuffmanScenario const scenario = huffmanScenarios().back(); + for (size_t blockSize : + { (size_t)1, (size_t)7, (size_t)64, (size_t)1000 }) { + auto const data = makeHuffmanData(scenario.alphabet, 4096); + for (auto const& enc : supportedEncodeArchs()) { + for (auto const& dec : supportedDecodeArchs()) { + SCOPED_TRACE( + std::string("blockSize=") + std::to_string(blockSize) + + " encode=" + enc.name + " decode=" + dec.name); + expectPivCoRoundTrip( + scenario, data, enc.kernels, dec.kernels, blockSize); + } + } + } +} + +TEST(PivCoHuffmanKernelTest, BuildsSymbolRanksAndSplitRanks) +{ + { + auto const weights = makeWeights(3, { { 0, 2 }, { 1, 1 }, { 2, 1 } }); + int const tableLog = + ZL_PivCoHuffman_computeTableLog(weights.data(), weights.size()); + ASSERT_GE(tableLog, 0); + + ZL_PivCoHuffmanTree tree; + ZL_PivCoHuffmanTree_build( + &tree, weights.data(), weights.size(), tableLog); + + EXPECT_EQ((int)tree.symbolToRank[0], 0); + EXPECT_EQ((int)tree.symbolToRank[1], 1); + EXPECT_EQ((int)tree.symbolToRank[2], 2); + EXPECT_EQ(tree.numRanks, 3); + EXPECT_EQ(ZL_PivCoHuffmanTree_splitRank(&tree, 0, 0, tree.numRanks), 1); + EXPECT_EQ(ZL_PivCoHuffmanTree_splitRank(&tree, 1, 1, tree.numRanks), 2); + } + + { + auto const weights = makeWeights( + 5, { { 0, 3 }, { 1, 1 }, { 2, 1 }, { 3, 1 }, { 4, 1 } }); + int const tableLog = + ZL_PivCoHuffman_computeTableLog(weights.data(), weights.size()); + ASSERT_GE(tableLog, 0); + + ZL_PivCoHuffmanTree tree; + ZL_PivCoHuffmanTree_build( + &tree, weights.data(), weights.size(), tableLog); + + for (uint8_t symbol = 0; symbol < weights.size(); ++symbol) { + EXPECT_EQ((int)tree.symbolToRank[symbol], (int)symbol); + } + EXPECT_TRUE(ZL_PivCoHuffmanTree_rangeIsLeaf(&tree, 1, 5)); + EXPECT_EQ(ZL_PivCoHuffmanTree_leafFlatDepth(&tree, 1), 2); + EXPECT_EQ(tree.numRanks, 5); + EXPECT_EQ(ZL_PivCoHuffmanTree_splitRank(&tree, 0, 0, tree.numRanks), 1); + } +} + +TEST(PivCoHuffmanKernelTest, DecodeEmptyOutputRequiresEmptyPayload) +{ + HuffmanScenario const scenario = huffmanScenarios()[1]; + uint8_t const payload = 0xA5; + + EXPECT_TRUE(ZL_PivCoHuffman_decode( + nullptr, + 0, + nullptr, + 0, + scenario.weights.data(), + scenario.weights.size(), + nullptr, + 0, + ZL_PIVCO_DEFAULT_BLOCK_SIZE, + &ZL_PivCoHuffmanDecode_generic)); + EXPECT_FALSE(ZL_PivCoHuffman_decode( + nullptr, + 0, + nullptr, + 0, + scenario.weights.data(), + scenario.weights.size(), + &payload, + 1, + ZL_PIVCO_DEFAULT_BLOCK_SIZE, + &ZL_PivCoHuffmanDecode_generic)); + + // An empty output has no blocks, so the block size is irrelevant: the + // binding decodes an empty input with the block size defaulted to + // decodedSize == 0 (the encoder omits it for a single block), and an + // out-of-range block size is equally harmless when there is nothing to + // decode. Both must still succeed on an empty payload. + EXPECT_TRUE(ZL_PivCoHuffman_decode( + nullptr, + 0, + nullptr, + 0, + scenario.weights.data(), + scenario.weights.size(), + nullptr, + 0, + 0, + &ZL_PivCoHuffmanDecode_generic)); + EXPECT_TRUE(ZL_PivCoHuffman_decode( + nullptr, + 0, + nullptr, + 0, + scenario.weights.data(), + scenario.weights.size(), + nullptr, + 0, + ZL_PIVCO_MAX_BLOCK_SIZE + 1, + &ZL_PivCoHuffmanDecode_generic)); +} + +TEST(PivCoHuffmanKernelTest, DecodeRejectsTruncatedBitstream) +{ + HuffmanScenario const scenario = huffmanScenarios()[2]; + auto const data = makeHuffmanData(scenario.alphabet, 32); + size_t const blockSize = ZL_PIVCO_DEFAULT_BLOCK_SIZE; + std::vector const encoded = encodePivCo(scenario, data, blockSize); + ASSERT_FALSE(encoded.empty()); + + std::vector decodeScratch( + ZL_PivCoHuffmanDecode_scratchBytes(data.size(), blockSize)); + std::vector decoded(data.size()); + EXPECT_FALSE(ZL_PivCoHuffman_decode( + decoded.data(), + decoded.size(), + decodeScratch.data(), + decodeScratch.size(), + scenario.weights.data(), + scenario.weights.size(), + encoded.data(), + encoded.size() - 1, + blockSize, + &ZL_PivCoHuffmanDecode_generic)); +} + +TEST(PivCoHuffmanKernelTest, DecodeRejectsCorruptNodeCount) +{ + HuffmanScenario const scenario = huffmanScenarios()[2]; + std::vector const data = { 0, 1 }; + size_t const blockSize = ZL_PIVCO_DEFAULT_BLOCK_SIZE; + std::vector encoded = encodePivCo(scenario, data, blockSize); + ASSERT_FALSE(encoded.empty()); + + encoded[0] = (uint8_t)((encoded[0] & 0x03u) | 0x0Cu); + + std::vector decodeScratch( + ZL_PivCoHuffmanDecode_scratchBytes(data.size(), blockSize)); + std::vector decoded(data.size()); + EXPECT_FALSE(ZL_PivCoHuffman_decode( + decoded.data(), + decoded.size(), + decodeScratch.data(), + decodeScratch.size(), + scenario.weights.data(), + scenario.weights.size(), + encoded.data(), + encoded.size(), + blockSize, + &ZL_PivCoHuffmanDecode_generic)); +}