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..46ddcc253 --- /dev/null +++ b/benchmark/micro/micro_pivco_huffman.cpp @@ -0,0 +1,683 @@ +// 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 }, + { "avx512", &ZL_PivCoHuffmanEncode_avx512 } + }; + 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 }, + { "avx512", &ZL_PivCoHuffmanDecode_avx512 }, + }; + 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/include/openzl/zl_version.h b/include/openzl/zl_version.h index e915ac3ad..ffd9789c4 100644 --- a/include/openzl/zl_version.h +++ b/include/openzl/zl_version.h @@ -54,7 +54,7 @@ extern "C" { /// format changes. But note that once a library with /// max format version X is released, we must support X /// through our support window. -#define ZL_MAX_FORMAT_VERSION (26) +#define ZL_MAX_FORMAT_VERSION (27) /// Minimum wire format version required to support chunking. #define ZL_CHUNK_VERSION_MIN (21) diff --git a/scripts/gen_pivco_huffman_tables.py b/scripts/gen_pivco_huffman_tables.py new file mode 100644 index 000000000..f053ceeb7 --- /dev/null +++ b/scripts/gen_pivco_huffman_tables.py @@ -0,0 +1,403 @@ +#!/usr/bin/env python3 +# Copyright (c) Meta Platforms, Inc. and affiliates. + +from __future__ import annotations + +from pathlib import Path + + +ROOT = Path(__file__).resolve().parents[1] +SCRIPT = "scripts/gen_pivco_huffman_tables.py" +GENERATED = "generated" + + +def header_guard(path: str) -> str: + return path.replace("src/", "").replace("/", "_").replace(".", "_").upper() + + +def hex8(value: int) -> str: + assert 0 <= value <= 0xFF + return f"0x{value:02x}" + + +def hex64(value: int) -> str: + assert 0 <= value <= 0xFFFFFFFFFFFFFFFF + return f"0x{value:016x}ULL" + + +def pack_le64(values: list[int]) -> int: + assert len(values) == 8 + out = 0 + for i, value in enumerate(values): + out |= value << (8 * i) + return out + + +def repeat_byte(value: int) -> int: + assert 0 <= value <= 0xFF + return value * 0x0101010101010101 + + +def popcounts() -> list[int]: + return [mask.bit_count() for mask in range(256)] + + +def zeros_counts() -> list[int]: + return [8 - mask.bit_count() for mask in range(256)] + + +def compress_u16(mask: int, pad: int) -> list[int]: + out = [] + for lane in range(8): + if (mask & (1 << lane)) != 0: + out += [2 * lane, 2 * lane + 1] + out += [pad] * (16 - len(out)) + return out + + +def compress_u16_pair(pad: int) -> list[list[int]]: + return [ + compress_u16(mask, pad) + compress_u16((~mask) & 0xFF, pad) + for mask in range(256) + ] + + +def partition_lo16(mask: int) -> list[int]: + # Low-half gather control for the 16-wide partition: right ranks (set lanes) + # packed forward at the front, left ranks (unset lanes) packed reversed at the + # far back; the gap in between is filled by the high-half control. + out = [0x80] * 16 + pos = 0 + for lane in range(8): + if (mask >> lane) & 1: + out[pos] = lane + pos += 1 + pos = 15 + for lane in range(8): + if not ((mask >> lane) & 1): + out[pos] = lane + pos -= 1 + return out + + +def partition_hi_base16(mask: int) -> list[int]: + # High-half gather control before zero-padding: right ranks packed forward at + # the front, left ranks packed reversed ending at lane 7. Values are source + # byte indices (8 + lane), since the high half lives in bytes 8..15. + out = [0x80] * 16 + pos = 0 + for lane in range(8): + if (mask >> lane) & 1: + out[pos] = 8 + lane + pos += 1 + pos = 7 + for lane in range(8): + if not ((mask >> lane) & 1): + out[pos] = 8 + lane + pos -= 1 + return out + + +def partition_lo() -> list[list[int]]: + return [partition_lo16(mask) for mask in range(256)] + + +def partition_hi_padded() -> list[list[int]]: + # Pad each high-half control with 8 zero (0x80) bytes on either end so a + # misaligned load at offset (8 - popcount(loMask)) slides the high contribution + # into the correct output lanes. + pad = [0x80] * 8 + return [pad + partition_hi_base16(mask) + pad for mask in range(256)] + + +def rhs_select8(mask: int) -> list[int]: + rhs = 0 + out = [] + for bit in range(8): + if (mask & (1 << bit)) != 0: + out.append(rhs) + rhs += 1 + else: + out.append(0x80) + return out + + +def lhs_select8(mask: int) -> list[int]: + lhs = 0 + out = [] + for bit in range(8): + if (mask & (1 << bit)) == 0: + out.append(lhs) + lhs += 1 + else: + out.append(0x80) + return out + + +def select8x2() -> list[list[int]]: + return [rhs_select8(mask) + lhs_select8(mask) for mask in range(256)] + + +def byte_splatx2() -> list[list[int]]: + return [[mask.bit_count()] * 8 + [8 - mask.bit_count()] * 8 for mask in range(256)] + + +def select8_u64() -> list[int]: + return [pack_le64(rhs_select8(mask)) for mask in range(256)] + + +def byte_splat_u64() -> list[int]: + return [repeat_byte(mask.bit_count()) for mask in range(256)] + + +def neon_expand8(mask: int, lhs_base: int, rhs_base: int) -> list[int]: + lhs = lhs_base + rhs = rhs_base + out = [] + for bit in range(8): + if (mask & (1 << bit)) != 0: + out.append(rhs) + rhs += 1 + else: + out.append(lhs) + lhs += 1 + return out + + +def neon_expand8_table(rhs_base: int) -> list[list[int]]: + return [neon_expand8(mask, lhs_base=0, rhs_base=rhs_base) for mask in range(256)] + + +def neon_expand8_pre() -> list[list[list[int]]]: + return [ + [ + neon_expand8(mask, lhs_base=8 - prior_ones, rhs_base=16 + prior_ones) + for mask in range(256) + ] + for prior_ones in range(9) + ] + + +def avx512_unpack_permute(depth: int) -> list[int]: + out = [] + for group in range(8): + for i in range(8): + out.append(group * depth + i if i < depth else 0) + return out + + +def avx512_shift_control(depth: int) -> list[int]: + return [i * depth for _group in range(8) for i in range(8)] + + +def avx512_unpack_permute_tables() -> list[list[int]]: + return [avx512_unpack_permute(depth) for depth in range(2, 8)] + + +def avx512_shift_control_tables() -> list[list[int]]: + return [avx512_shift_control(depth) for depth in range(2, 8)] + + +def format_u8_values(values: list[int], per_line: int = 16) -> str: + lines = [] + for i in range(0, len(values), per_line): + lines.append( + " " + ", ".join(str(value) for value in values[i : i + per_line]) + "," + ) + return "\n".join(lines) + + +def format_u8_table_1d(name: str, values: list[int], align: int) -> str: + return ( + f"static const ZL_ALIGNED({align}) uint8_t {name}[{len(values)}] = {{\n" + f"{format_u8_values(values)}\n" + "};" + ) + + +def format_u8_row(row: list[int]) -> str: + return "{ " + ", ".join(hex8(value) for value in row) + " }" + + +def format_u8_table_2d(name: str, rows: list[list[int]], align: int) -> str: + row_len = len(rows[0]) + body = ",\n".join(" " + format_u8_row(row) for row in rows) + return ( + f"static const ZL_ALIGNED({align}) uint8_t {name}[{len(rows)}][{row_len}] = {{\n" + f"{body}\n" + "};" + ) + + +def format_u8_table_3d(name: str, planes: list[list[list[int]]], align: int) -> str: + plane_len = len(planes[0]) + row_len = len(planes[0][0]) + formatted_planes = [] + for plane in planes: + rows = ",\n".join(" " + format_u8_row(row) for row in plane) + formatted_planes.append(f" {{\n{rows}\n }}") + body = ",\n".join(formatted_planes) + return ( + f"static const ZL_ALIGNED({align}) uint8_t " + f"{name}[{len(planes)}][{plane_len}][{row_len}] = {{\n" + f"{body},\n" + "};" + ) + + +def format_u64_table(name: str, values: list[int], align: int) -> str: + lines = [] + for i in range(0, len(values), 4): + line = ", ".join(hex64(value) for value in values[i : i + 4]) + lines.append(f" {line},") + body = "\n".join(lines) + return ( + f"static const ZL_ALIGNED({align}) uint64_t {name}[{len(values)}] = {{\n" + f"{body}\n" + "};" + ) + + +def render_header(path: str, body: str) -> str: + guard = header_guard(path) + return f"""// Copyright (c) Meta Platforms, Inc. and affiliates. +#ifndef {guard} +#define {guard} + +/** + * @{GENERATED} by {SCRIPT} + * Please don't modify this file directly! + */ + +#include "openzl/shared/portability.h" + +ZL_BEGIN_C_DECLS + +// clang-format off +{body} +// clang-format on + +ZL_END_C_DECLS + +#endif // {guard} +""" + + +def render_x86_tables() -> str: + body = "\n\n".join( + [ + format_u8_table_2d("ZL_kPivCoHuffmanSelect8x2", select8x2(), align=16), + format_u8_table_2d("ZL_kPivCoHuffmanByteSplatx2", byte_splatx2(), align=16), + format_u8_table_1d("ZL_kPivCoHuffmanPopcount8", popcounts(), align=16), + ] + ) + return render_header( + "src/openzl/codecs/pivco_huffman/arch/common_pivco_x86_tables.h", + body, + ) + + +def render_x86_merge_tables() -> str: + body = "\n\n".join( + [ + format_u64_table("ZL_kPivCoHuffmanSelect8", select8_u64(), align=16), + format_u64_table("ZL_kPivCoHuffmanByteSplat", byte_splat_u64(), align=16), + ] + ) + return render_header( + "src/openzl/codecs/pivco_huffman/arch/common_pivco_x86_merge_tables.h", + body, + ) + + +def render_x86_partition_tables() -> str: + body = "\n\n".join( + [ + format_u8_table_2d( + "ZL_kPivCoHuffmanCompressU16Pair", compress_u16_pair(0x80), align=32 + ), + format_u8_table_2d("ZL_kPivCoHuffmanPartitionLo", partition_lo(), align=16), + format_u8_table_2d( + "ZL_kPivCoHuffmanPartitionHiPadded", + partition_hi_padded(), + align=32, + ), + ] + ) + return render_header( + "src/openzl/codecs/pivco_huffman/arch/common_pivco_x86_partition_tables.h", + body, + ) + + +def render_neon_tables() -> str: + body = "\n\n".join( + [ + format_u8_table_1d("ZL_kPivCoHuffmanNeonPopcount", popcounts(), align=64), + format_u8_table_1d( + "ZL_kPivCoHuffmanNeonZeroCount", zeros_counts(), align=64 + ), + format_u8_table_2d( + "ZL_kPivCoHuffmanNeonCompressU16", + compress_u16_pair(0xFF), + align=32, + ), + format_u8_table_2d( + "ZL_kPivCoHuffmanNeonExpand8", + neon_expand8_table(rhs_base=8), + align=32, + ), + format_u8_table_2d( + "ZL_kPivCoHuffmanNeonExpand8Q", + neon_expand8_table(rhs_base=16), + align=32, + ), + format_u8_table_3d( + "ZL_kPivCoHuffmanNeonExpand8Pre", + neon_expand8_pre(), + align=64, + ), + ] + ) + return render_header( + "src/openzl/codecs/pivco_huffman/arch/common_pivco_neon_tables.h", + body, + ) + + +def render_avx512_tables() -> str: + body = "\n\n".join( + [ + format_u8_table_2d( + "ZL_kPivCoHuffmanAvx512UnpackPermute", + avx512_unpack_permute_tables(), + align=64, + ), + format_u8_table_2d( + "ZL_kPivCoHuffmanAvx512ShiftControl", + avx512_shift_control_tables(), + align=64, + ), + ] + ) + return render_header( + "src/openzl/codecs/pivco_huffman/arch/common_pivco_avx512_tables.h", + body, + ) + + +def main() -> None: + outputs = { + "src/openzl/codecs/pivco_huffman/arch/common_pivco_x86_tables.h": render_x86_tables(), + "src/openzl/codecs/pivco_huffman/arch/common_pivco_x86_merge_tables.h": render_x86_merge_tables(), + "src/openzl/codecs/pivco_huffman/arch/common_pivco_x86_partition_tables.h": render_x86_partition_tables(), + "src/openzl/codecs/pivco_huffman/arch/common_pivco_neon_tables.h": render_neon_tables(), + "src/openzl/codecs/pivco_huffman/arch/common_pivco_avx512_tables.h": render_avx512_tables(), + } + for path, contents in outputs.items(): + (ROOT / path).write_text(contents) + print(f"generated {path}") + + +if __name__ == "__main__": + main() diff --git a/src/openzl/codecs/common/bitstream/ff_bitstream.h b/src/openzl/codecs/common/bitstream/ff_bitstream.h index b1d7202ee..63cf95be5 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,90 @@ 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 = ZS_BitDStreamFF_loadPartial(pos, remaining); + 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/decoder_registry.c b/src/openzl/codecs/decoder_registry.c index 18acbaf26..dbf575f4d 100644 --- a/src/openzl/codecs/decoder_registry.c +++ b/src/openzl/codecs/decoder_registry.c @@ -27,6 +27,7 @@ #include "openzl/codecs/parse_int/graph_parse_int.h" #include "openzl/codecs/partition/decode_partition_binding.h" #include "openzl/codecs/partition/decode_partition_bitpack_fusion.h" +#include "openzl/codecs/pivco_huffman/decode_pivco_binding.h" #include "openzl/codecs/prefix/decode_prefix_binding.h" #include "openzl/codecs/quantize/decode_quantize_binding.h" #include "openzl/codecs/range_pack/decode_range_pack_binding.h" @@ -138,6 +139,7 @@ const StandardDTransform SDecoders_array[ZL_StandardTransformID_end] = { REGISTER_TTRANSFORM_G(ZL_StandardTransformID_lz, 24, DI_LZ, LZ_GRAPH), REGISTER_TTRANSFORM_G(ZL_StandardTransformID_mux_lengths, 24, DI_MUX_LENGTHS, MUX_LENGTHS_GRAPH), REGISTER_TTRANSFORM_G(ZL_StandardTransformID_sparse_num, 26, DI_SPARSE_NUM, SPARSE_NUM_GRAPH), + REGISTER_TTRANSFORM_G(ZL_StandardTransformID_pivco_huffman, 27, DI_PIVCO_HUFFMAN, PIVCO_HUFFMAN_GRAPH), REGISTER_VOTRANSFORM_G(ZL_StandardTransformID_splitn, 9, DI_SPLITN, GRAPH_VO_SERIAL), REGISTER_VOTRANSFORM_G(ZL_StandardTransformID_splitn_struct, 14, DI_SPLITN_STRUCT, GRAPH_VO_STRUCT), diff --git a/src/openzl/codecs/encoder_registry.c b/src/openzl/codecs/encoder_registry.c index 70dadca8a..d53e09d37 100644 --- a/src/openzl/codecs/encoder_registry.c +++ b/src/openzl/codecs/encoder_registry.c @@ -26,6 +26,7 @@ #include "openzl/codecs/mux_lengths/encode_mux_lengths_binding.h" #include "openzl/codecs/parse_int/encode_parse_int_binding.h" #include "openzl/codecs/partition/encode_partition_binding.h" +#include "openzl/codecs/pivco_huffman/encode_pivco_binding.h" #include "openzl/codecs/prefix/encode_prefix_binding.h" #include "openzl/codecs/quantize/encode_quantize_binding.h" #include "openzl/codecs/range_pack/encode_range_pack_binding.h" @@ -161,6 +162,7 @@ const CNode ER_standardNodes[STANDARD_ENCODERS_NB] = { REGISTER_TRANSFORM(ZL_PrivateStandardNodeID_dedup_num_trusted, ZL_StandardTransformID_dedup_num, 16, 200, EI_DEDUP_NUM_TRUSTED), REGISTER_TRANSFORM(ZL_PrivateStandardNodeID_lz4, ZL_StandardTransformID_lz4, 23, 200, EI_LZ4), REGISTER_TRANSFORM(ZL_PrivateStandardNodeID_bitSplit, ZL_StandardTransformID_bitSplit, 24, 200, EI_BITSPLIT), + REGISTER_TRANSFORM(ZL_PrivateStandardNodeID_pivco_huffman, ZL_StandardTransformID_pivco_huffman, 27, 203, EI_PIVCO_HUFFMAN), // Deprecated Nodes REGISTER_DEPRECATED_TRANSFORM(ZL_PrivateStandardNodeID_rolz_deprecated, ZL_StandardTransformID_rolz, 3, 12, 200, EI_ROLZ), 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/common_pivco_avx512_tables.h b/src/openzl/codecs/pivco_huffman/arch/common_pivco_avx512_tables.h new file mode 100644 index 000000000..21fb585b2 --- /dev/null +++ b/src/openzl/codecs/pivco_huffman/arch/common_pivco_avx512_tables.h @@ -0,0 +1,36 @@ +// Copyright (c) Meta Platforms, Inc. and affiliates. +#ifndef OPENZL_CODECS_PIVCO_HUFFMAN_ARCH_COMMON_PIVCO_AVX512_TABLES_H +#define OPENZL_CODECS_PIVCO_HUFFMAN_ARCH_COMMON_PIVCO_AVX512_TABLES_H + +/** + * @generated by scripts/gen_pivco_huffman_tables.py + * Please don't modify this file directly! + */ + +#include "openzl/shared/portability.h" + +ZL_BEGIN_C_DECLS + +// clang-format off +static const ZL_ALIGNED(64) uint8_t ZL_kPivCoHuffmanAvx512UnpackPermute[6][64] = { + { 0x00, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x02, 0x03, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x04, 0x05, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x06, 0x07, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x08, 0x09, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x0a, 0x0b, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x0c, 0x0d, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x0e, 0x0f, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00 }, + { 0x00, 0x01, 0x02, 0x00, 0x00, 0x00, 0x00, 0x00, 0x03, 0x04, 0x05, 0x00, 0x00, 0x00, 0x00, 0x00, 0x06, 0x07, 0x08, 0x00, 0x00, 0x00, 0x00, 0x00, 0x09, 0x0a, 0x0b, 0x00, 0x00, 0x00, 0x00, 0x00, 0x0c, 0x0d, 0x0e, 0x00, 0x00, 0x00, 0x00, 0x00, 0x0f, 0x10, 0x11, 0x00, 0x00, 0x00, 0x00, 0x00, 0x12, 0x13, 0x14, 0x00, 0x00, 0x00, 0x00, 0x00, 0x15, 0x16, 0x17, 0x00, 0x00, 0x00, 0x00, 0x00 }, + { 0x00, 0x01, 0x02, 0x03, 0x00, 0x00, 0x00, 0x00, 0x04, 0x05, 0x06, 0x07, 0x00, 0x00, 0x00, 0x00, 0x08, 0x09, 0x0a, 0x0b, 0x00, 0x00, 0x00, 0x00, 0x0c, 0x0d, 0x0e, 0x0f, 0x00, 0x00, 0x00, 0x00, 0x10, 0x11, 0x12, 0x13, 0x00, 0x00, 0x00, 0x00, 0x14, 0x15, 0x16, 0x17, 0x00, 0x00, 0x00, 0x00, 0x18, 0x19, 0x1a, 0x1b, 0x00, 0x00, 0x00, 0x00, 0x1c, 0x1d, 0x1e, 0x1f, 0x00, 0x00, 0x00, 0x00 }, + { 0x00, 0x01, 0x02, 0x03, 0x04, 0x00, 0x00, 0x00, 0x05, 0x06, 0x07, 0x08, 0x09, 0x00, 0x00, 0x00, 0x0a, 0x0b, 0x0c, 0x0d, 0x0e, 0x00, 0x00, 0x00, 0x0f, 0x10, 0x11, 0x12, 0x13, 0x00, 0x00, 0x00, 0x14, 0x15, 0x16, 0x17, 0x18, 0x00, 0x00, 0x00, 0x19, 0x1a, 0x1b, 0x1c, 0x1d, 0x00, 0x00, 0x00, 0x1e, 0x1f, 0x20, 0x21, 0x22, 0x00, 0x00, 0x00, 0x23, 0x24, 0x25, 0x26, 0x27, 0x00, 0x00, 0x00 }, + { 0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x00, 0x00, 0x06, 0x07, 0x08, 0x09, 0x0a, 0x0b, 0x00, 0x00, 0x0c, 0x0d, 0x0e, 0x0f, 0x10, 0x11, 0x00, 0x00, 0x12, 0x13, 0x14, 0x15, 0x16, 0x17, 0x00, 0x00, 0x18, 0x19, 0x1a, 0x1b, 0x1c, 0x1d, 0x00, 0x00, 0x1e, 0x1f, 0x20, 0x21, 0x22, 0x23, 0x00, 0x00, 0x24, 0x25, 0x26, 0x27, 0x28, 0x29, 0x00, 0x00, 0x2a, 0x2b, 0x2c, 0x2d, 0x2e, 0x2f, 0x00, 0x00 }, + { 0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x00, 0x07, 0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d, 0x00, 0x0e, 0x0f, 0x10, 0x11, 0x12, 0x13, 0x14, 0x00, 0x15, 0x16, 0x17, 0x18, 0x19, 0x1a, 0x1b, 0x00, 0x1c, 0x1d, 0x1e, 0x1f, 0x20, 0x21, 0x22, 0x00, 0x23, 0x24, 0x25, 0x26, 0x27, 0x28, 0x29, 0x00, 0x2a, 0x2b, 0x2c, 0x2d, 0x2e, 0x2f, 0x30, 0x00, 0x31, 0x32, 0x33, 0x34, 0x35, 0x36, 0x37, 0x00 } +}; + +static const ZL_ALIGNED(64) uint8_t ZL_kPivCoHuffmanAvx512ShiftControl[6][64] = { + { 0x00, 0x02, 0x04, 0x06, 0x08, 0x0a, 0x0c, 0x0e, 0x00, 0x02, 0x04, 0x06, 0x08, 0x0a, 0x0c, 0x0e, 0x00, 0x02, 0x04, 0x06, 0x08, 0x0a, 0x0c, 0x0e, 0x00, 0x02, 0x04, 0x06, 0x08, 0x0a, 0x0c, 0x0e, 0x00, 0x02, 0x04, 0x06, 0x08, 0x0a, 0x0c, 0x0e, 0x00, 0x02, 0x04, 0x06, 0x08, 0x0a, 0x0c, 0x0e, 0x00, 0x02, 0x04, 0x06, 0x08, 0x0a, 0x0c, 0x0e, 0x00, 0x02, 0x04, 0x06, 0x08, 0x0a, 0x0c, 0x0e }, + { 0x00, 0x03, 0x06, 0x09, 0x0c, 0x0f, 0x12, 0x15, 0x00, 0x03, 0x06, 0x09, 0x0c, 0x0f, 0x12, 0x15, 0x00, 0x03, 0x06, 0x09, 0x0c, 0x0f, 0x12, 0x15, 0x00, 0x03, 0x06, 0x09, 0x0c, 0x0f, 0x12, 0x15, 0x00, 0x03, 0x06, 0x09, 0x0c, 0x0f, 0x12, 0x15, 0x00, 0x03, 0x06, 0x09, 0x0c, 0x0f, 0x12, 0x15, 0x00, 0x03, 0x06, 0x09, 0x0c, 0x0f, 0x12, 0x15, 0x00, 0x03, 0x06, 0x09, 0x0c, 0x0f, 0x12, 0x15 }, + { 0x00, 0x04, 0x08, 0x0c, 0x10, 0x14, 0x18, 0x1c, 0x00, 0x04, 0x08, 0x0c, 0x10, 0x14, 0x18, 0x1c, 0x00, 0x04, 0x08, 0x0c, 0x10, 0x14, 0x18, 0x1c, 0x00, 0x04, 0x08, 0x0c, 0x10, 0x14, 0x18, 0x1c, 0x00, 0x04, 0x08, 0x0c, 0x10, 0x14, 0x18, 0x1c, 0x00, 0x04, 0x08, 0x0c, 0x10, 0x14, 0x18, 0x1c, 0x00, 0x04, 0x08, 0x0c, 0x10, 0x14, 0x18, 0x1c, 0x00, 0x04, 0x08, 0x0c, 0x10, 0x14, 0x18, 0x1c }, + { 0x00, 0x05, 0x0a, 0x0f, 0x14, 0x19, 0x1e, 0x23, 0x00, 0x05, 0x0a, 0x0f, 0x14, 0x19, 0x1e, 0x23, 0x00, 0x05, 0x0a, 0x0f, 0x14, 0x19, 0x1e, 0x23, 0x00, 0x05, 0x0a, 0x0f, 0x14, 0x19, 0x1e, 0x23, 0x00, 0x05, 0x0a, 0x0f, 0x14, 0x19, 0x1e, 0x23, 0x00, 0x05, 0x0a, 0x0f, 0x14, 0x19, 0x1e, 0x23, 0x00, 0x05, 0x0a, 0x0f, 0x14, 0x19, 0x1e, 0x23, 0x00, 0x05, 0x0a, 0x0f, 0x14, 0x19, 0x1e, 0x23 }, + { 0x00, 0x06, 0x0c, 0x12, 0x18, 0x1e, 0x24, 0x2a, 0x00, 0x06, 0x0c, 0x12, 0x18, 0x1e, 0x24, 0x2a, 0x00, 0x06, 0x0c, 0x12, 0x18, 0x1e, 0x24, 0x2a, 0x00, 0x06, 0x0c, 0x12, 0x18, 0x1e, 0x24, 0x2a, 0x00, 0x06, 0x0c, 0x12, 0x18, 0x1e, 0x24, 0x2a, 0x00, 0x06, 0x0c, 0x12, 0x18, 0x1e, 0x24, 0x2a, 0x00, 0x06, 0x0c, 0x12, 0x18, 0x1e, 0x24, 0x2a, 0x00, 0x06, 0x0c, 0x12, 0x18, 0x1e, 0x24, 0x2a }, + { 0x00, 0x07, 0x0e, 0x15, 0x1c, 0x23, 0x2a, 0x31, 0x00, 0x07, 0x0e, 0x15, 0x1c, 0x23, 0x2a, 0x31, 0x00, 0x07, 0x0e, 0x15, 0x1c, 0x23, 0x2a, 0x31, 0x00, 0x07, 0x0e, 0x15, 0x1c, 0x23, 0x2a, 0x31, 0x00, 0x07, 0x0e, 0x15, 0x1c, 0x23, 0x2a, 0x31, 0x00, 0x07, 0x0e, 0x15, 0x1c, 0x23, 0x2a, 0x31, 0x00, 0x07, 0x0e, 0x15, 0x1c, 0x23, 0x2a, 0x31, 0x00, 0x07, 0x0e, 0x15, 0x1c, 0x23, 0x2a, 0x31 } +}; +// clang-format on + +ZL_END_C_DECLS + +#endif // OPENZL_CODECS_PIVCO_HUFFMAN_ARCH_COMMON_PIVCO_AVX512_TABLES_H 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..d45968f13 --- /dev/null +++ b/src/openzl/codecs/pivco_huffman/arch/decode_pivco_arch.c @@ -0,0 +1,212 @@ +// 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) +{ + ZL_cpuid_t localCpuid; + if (cpuid == NULL) { + localCpuid = ZL_cpuid(); + cpuid = &localCpuid; + } + + if (ZL_PivCoHuffmanDecode_avx512.supported(cpuid)) { + return &ZL_PivCoHuffmanDecode_avx512; + } + 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..2ee3d9258 --- /dev/null +++ b/src/openzl/codecs/pivco_huffman/arch/decode_pivco_arch.h @@ -0,0 +1,132 @@ +// 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; +extern const ZL_PivCoHuffmanDecode ZL_PivCoHuffmanDecode_avx512; + +ZL_END_C_DECLS + +#endif diff --git a/src/openzl/codecs/pivco_huffman/arch/decode_pivco_avx512.c b/src/openzl/codecs/pivco_huffman/arch/decode_pivco_avx512.c new file mode 100644 index 000000000..8d52114e1 --- /dev/null +++ b/src/openzl/codecs/pivco_huffman/arch/decode_pivco_avx512.c @@ -0,0 +1,630 @@ +// Copyright (c) Meta Platforms, Inc. and affiliates. + +#include "openzl/codecs/pivco_huffman/arch/decode_pivco_arch.h" + +#if ZL_ARCH_X86_64 && ZL_HAS_ATTRIBUTE(__target__) + +# include +# include + +# include "common_pivco_avx512_tables.h" +# include "openzl/shared/mem.h" +# include "openzl/shared/utils.h" + +# define ZL_AVX512_ATTR \ + ZL_TARGET_ATTRIBUTE( \ + "avx512vbmi,avx512vbmi2,avx512f,avx512vl,avx512bw,bmi2") +# define ZL_AVX512_INLINE ZL_FORCE_INLINE ZL_AVX512_ATTR + +static bool supported(const ZL_cpuid_t* cpuid) +{ + return cpuid != NULL && ZL_cpuid_avx512vbmi(*cpuid) + && ZL_cpuid_avx512vbmi2(*cpuid) && ZL_cpuid_avx512f(*cpuid) + && ZL_cpuid_avx512vl(*cpuid) && ZL_cpuid_avx512bw(*cpuid) + && ZL_cpuid_bmi2(*cpuid); +} + +/************************************************************** + * mergeVector* + **************************************************************/ + +/// @returns a 64-bit lane mask with the low @p bits bits set. Selects the valid +/// lanes/bytes of a partial 64-wide block. +static __mmask64 tailMask(size_t bits) +{ + return bits >= 64 ? ~(__mmask64)0 : (((__mmask64)1 << bits) - 1); +} + +/** + * Merges up to 64 output bytes for one 64-bit slice of the partition @p bits: + * bit i picks the right child (1) or left child (0) for output lane i. A vector + * child supplies bytes consecutively from its stream; a constant child supplies + * a splatted value. + * + * The merge uses `expandloadu`, the inverse of the encoder's `compressstoreu`: + * it reads contiguous bytes from a child stream and scatters them into the + * lanes selected by a mask, so each child's bytes land back in their original + * positions. + * + * @param lhsPos/rhsPos In/out cursors into the (contiguous) child streams; + * advanced by the number of bytes this block consumed from each. + * @param lhsVec/rhsVec Splatted constant values, used when that side is + * constant (kLhsVector/kRhsVector false). + * @param numBits Valid lanes in this block (64 for a full block; the tail mask + * keeps the loop from reading past either stream). + * @param kLhsVector/kRhsVector Compile-time flags selecting vector vs constant + * for each side (at least one is a vector). + */ +ZL_AVX512_INLINE __m512i mergeVectorBlock( + const uint8_t* lhs, + size_t* lhsPos, + __m512i lhsVec, + const uint8_t* rhs, + size_t* rhsPos, + __m512i rhsVec, + uint64_t bits, + size_t numBits, + const bool kLhsVector, + const bool kRhsVector) +{ + // Restrict the partition to valid lanes; lhsMask/rhsMask are complementary + // within those lanes. blockOnes counts how many bytes the right child + // feeds. + __mmask64 valid = tailMask(numBits); + __mmask64 rhsMask = (__mmask64)bits & valid; + __mmask64 lhsMask = (__mmask64)(~bits) & valid; + const size_t blockOnes = (size_t)ZL_popcount64((uint64_t)rhsMask); + + __m512i merged; + if (kLhsVector && kRhsVector) { + // Expand each contiguous stream into its selected lanes (zero + // elsewhere); the two are disjoint, so OR fuses them into the output. + const __m512i expandedRhs = + _mm512_maskz_expandloadu_epi8(rhsMask, rhs + *rhsPos); + const __m512i expandedLhs = + _mm512_maskz_expandloadu_epi8(lhsMask, lhs + *lhsPos); + merged = _mm512_or_si512(expandedLhs, expandedRhs); + } else if (kLhsVector) { + // rhs is constant: start from the rhs splat and overwrite the 0-bit + // lanes with expanded lhs bytes. + merged = _mm512_mask_expandloadu_epi8(rhsVec, lhsMask, lhs + *lhsPos); + } else { + assert(kRhsVector); + // lhs is constant: start from the lhs splat and overwrite the 1-bit + // lanes with expanded rhs bytes. + merged = _mm512_mask_expandloadu_epi8(lhsVec, rhsMask, rhs + *rhsPos); + } + + // Advance each vector cursor by the bytes it actually supplied. + if (kLhsVector) { + *lhsPos += numBits - blockOnes; + } + if (kRhsVector) { + *rhsPos += blockOnes; + } + return merged; +} + +/** + * Shared body of the three merge kernels. Reconstructs `lhsSize + rhsSize` + * output bytes by walking the partition bitmap 64 bits at a time and merging + * each block with mergeVectorBlock. A constant side passes its value via + * @p lhsValue / @p rhsValue (splatted once up front); a vector side reads from + * @p lhs / @p rhs. At least one side is a vector. + * + * @returns the number of 1-bits consumed (the right-child size as recomputed + * from the bitmap). The caller compares it to @p rhsSize: a mismatch means a + * corrupt bitstream. The early returns inside the loops fire as soon as a + * vector cursor overruns its stream -- the count returned there is guaranteed + * to differ from @p rhsSize, so corruption is still reported (and we stop + * before reading further out of bounds). + * + * @note bitmapCapacity/outCapacity are asserted, then ignored: the loop relies + * on the documented SLOP so the full-width 64-byte loads/stores are safe. + */ +ZL_AVX512_INLINE size_t mergeVectorImpl( + uint8_t* out, + size_t outCapacity, + const uint8_t* bitmap, + size_t bitmapCapacity, + const uint8_t* lhs, + uint8_t lhsValue, + size_t lhsSize, + const uint8_t* rhs, + uint8_t rhsValue, + size_t rhsSize, + const bool kLhsVector, + const bool kRhsVector) +{ + const size_t outSize = (size_t)(lhsSize + rhsSize); + assert(outCapacity >= outSize); + assert(bitmapCapacity >= (outSize + 7) / 8); + assert(kLhsVector || kRhsVector); + (void)outCapacity; + (void)bitmapCapacity; + + size_t lhsPos = 0; + size_t rhsPos = 0; + + const __m512i lhsVec = _mm512_set1_epi8((char)lhsValue); + const __m512i rhsVec = _mm512_set1_epi8((char)rhsValue); + + // Main loop: process kUnroll outputs at a time in fully-unrolled 64-wide + // steps. Every block here is full-width (numBits == 64) and uses unmasked + // loads/stores, which is the fast path. + size_t bitOffset = 0; + const size_t kUnroll = 256; + const size_t outLimit = outSize & ~(kUnroll - 1); + for (; bitOffset < outLimit; bitOffset += kUnroll) { +# ifdef __clang__ +# pragma clang loop unroll(full) +# endif + for (size_t u = 0; u < kUnroll; u += 64) { + const uint64_t bits = ZL_readLE64(bitmap + bitOffset / 8 + u / 8); + const __m512i merged = mergeVectorBlock( + lhs, + &lhsPos, + lhsVec, + rhs, + &rhsPos, + rhsVec, + bits, + 64, + kLhsVector, + kRhsVector); + _mm512_storeu_si512(out + bitOffset + u, merged); + // Corruption guard: a cursor past its stream means the bitmap had + // more 1s (or 0s) than the child has bytes. Bail with a count that + // can't equal rhsSize so the caller flags it. + if (kRhsVector && ZL_UNLIKELY(rhsPos > rhsSize)) { + return rhsPos; + } + if (kLhsVector && ZL_UNLIKELY(lhsPos > lhsSize)) { + assert(outSize - lhsPos < rhsSize); + return outSize - lhsPos; + } + } + } + + // Tail: the remaining < kUnroll outputs, handled 64 at a time with the last + // block masked to numBits valid lanes. + for (; bitOffset < outSize; bitOffset += 64) { + const size_t numBits = ZL_MIN(outSize - bitOffset, 64); + const __mmask64 valid = tailMask(numBits); + const uint64_t bits = + ZL_readLE64_N(bitmap + bitOffset / 8, (numBits + 7) / 8); + const __m512i merged = mergeVectorBlock( + lhs, + &lhsPos, + lhsVec, + rhs, + &rhsPos, + rhsVec, + bits, + numBits, + kLhsVector, + kRhsVector); + _mm512_mask_storeu_epi8(out + bitOffset, valid, merged); + if (kRhsVector && ZL_UNLIKELY(rhsPos > rhsSize)) { + return rhsPos; + } + if (kLhsVector && ZL_UNLIKELY(lhsPos > lhsSize)) { + assert(outSize - lhsPos < rhsSize); + return outSize - lhsPos; + } + } + + // ones count: rhsPos directly, or (for a constant rhs) the lanes not taken + // from the lhs vector. + return kRhsVector ? rhsPos : outSize - lhsPos; +} + +/// Merge where both children are decoded vectors. @see +/// ZL_PivCoHuffmanDecode::mergeVectorVector. +static ZL_AVX512_ATTR 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 mergeVectorImpl( + out, + outCapacity, + bitmap, + bitmapCapacity, + lhs, + 0, + lhsSize, + rhs, + 0, + rhsSize, + true, + true); +} + +/// Merge where the left child is a constant symbol and the right is a vector. +/// @see ZL_PivCoHuffmanDecode::mergeConstantVector. +static ZL_AVX512_ATTR 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 mergeVectorImpl( + out, + outCapacity, + bitmap, + bitmapCapacity, + NULL, + lhs, + lhsSize, + rhs, + 0, + rhsSize, + false, + true); +} + +/// Merge where the left child is a vector and the right is a constant symbol. +/// @see ZL_PivCoHuffmanDecode::mergeVectorConstant. +static ZL_AVX512_ATTR 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 mergeVectorImpl( + out, + outCapacity, + bitmap, + bitmapCapacity, + lhs, + 0, + lhsSize, + NULL, + rhs, + rhsSize, + true, + false); +} + +/************************************************************** + * mergeFlatDepth* + **************************************************************/ + +/// @returns the permute write-mask for unpacking: within each of the 8 u64 +/// lanes (one per group of 8 indices) keep the low @p kDepth bytes -- the +/// packed bytes that group occupies -- and zero the rest. Matches the layout +/// UnpackPermute scatters into. +ZL_AVX512_INLINE __mmask64 unpackGroupByteMask(size_t kDepth) +{ + __mmask64 mask = 0; + for (size_t group = 0; group < 8; ++group) { + mask |= (__mmask64)(tailMask(kDepth) << (group * 8)); + } + return mask; +} + +/** + * Unpacks one 64-symbol block of depth-`d` indices. Produces a __m512i with one + * depth-bit index per byte lane. + * + * 1. Load the packed bytes (loadMask covers the 8*d bytes for 64 indices). + * 2. permutexvar via @p unpackPermute regathers each group's `d` packed bytes + * into the low bytes of one u64 lane, so each lane holds 8*d contiguous + * bits == 8 indices. @p unpackMask zeroes the unused high bytes of each + * lane. + * 3. multishift via @p shiftControl ({0,d,2d,...,7d}) copies the byte starting + * at each index's bit position into a byte lane, so lane j holds index j in + * its low d bits (with high bits spilled from the neighbour). + * 4. AND with @p indexMask (low d bits) clears that spill. + */ +ZL_AVX512_INLINE __m512i unpackFlatDepth( + const uint8_t* bitmap, + __mmask64 loadMask, + __m512i unpackPermute, + __mmask64 unpackMask, + __m512i shiftControl, + __m512i indexMask) +{ + const __m512i packed = _mm512_maskz_loadu_epi8(loadMask, bitmap); + // Group each packed byte run into one u64 lane for multishift unpacking + const __m512i groups = + _mm512_maskz_permutexvar_epi8(unpackMask, unpackPermute, packed); + // Extract one depth-bit index into each byte lane + return _mm512_and_si512( + _mm512_multishift_epi64_epi8(shiftControl, groups), indexMask); +} + +/** + * Builds the symbol lookup table(s) for depths 1..7, laid out for the shuffle + * primitive each tier uses (see shuffleSymbolsDepth): + * - depth <= 4: the 2^depth symbols (zero-padded to 16) replicated into all + * four 128-bit lanes, because `_mm512_shuffle_epi8` looks up within each + * lane. + * - depth 5..6: the 2^depth symbols in one 64-byte table for `permutexvar`. + * - depth 7: 128 symbols split across two 64-byte pages (lut0 = 0..63, + * lut1 = 64..127), selected later by index bit 6. + */ +ZL_AVX512_INLINE void buildShuffleTables( + uint8_t* lut0, + uint8_t* lut1, + size_t kDepth, + const uint8_t* symbols) +{ + for (size_t i = 0; i < 64; ++i) { + lut0[i] = 0; + lut1[i] = 0; + } + + const size_t numSymbols = (size_t)1 << kDepth; + if (kDepth <= 4) { + for (size_t lane = 0; lane < 4; ++lane) { + for (size_t idx = 0; idx < 16; ++idx) { + lut0[lane * 16 + idx] = idx < numSymbols ? symbols[idx] : 0; + } + } + } else if (kDepth <= 6) { + for (size_t idx = 0; idx < numSymbols; ++idx) { + lut0[idx] = symbols[idx]; + } + } else { + for (size_t idx = 0; idx < 64; ++idx) { + lut0[idx] = symbols[idx]; + lut1[idx] = symbols[idx + 64]; + } + } +} + +/// Maps 64 depth-7 indices (0..127) to symbols. permutexvar handles a 64-entry +/// page, so look up both pages with the low 6 bits and pick per lane by bit 6. +ZL_AVX512_INLINE __m512i +shuffleSymbols7(__m512i indices, __m512i lut0, __m512i lut1) +{ + // Depth 7 needs two 64-byte LUT pages: bit 6 selects the page. + const __m512i localIndexMask = _mm512_set1_epi8(63); + const __m512i pageBit = _mm512_set1_epi8(64); + const __m512i localIndices = _mm512_and_si512(indices, localIndexMask); + const __m512i lowSymbols = _mm512_permutexvar_epi8(localIndices, lut0); + const __m512i highSymbols = _mm512_permutexvar_epi8(localIndices, lut1); + const __mmask64 highPageMask = _mm512_test_epi8_mask(indices, pageBit); + return _mm512_mask_mov_epi8(lowSymbols, highPageMask, highSymbols); +} + +/// Maps per-byte indices to symbols for depths 1..7, picking the right shuffle +/// primitive for the table layout buildShuffleTables produced: +/// - depth <= 4: in-lane `shuffle_epi8` (the table is replicated per 128 +/// bits). +/// - depth 5..6: full-width `permutexvar` over a single 64-byte page. +/// - depth 7: two-page lookup (shuffleSymbols7). Depth 8 is handled +/// separately. +ZL_AVX512_INLINE __m512i shuffleSymbolsDepth( + const size_t kDepth, + __m512i indices, + __m512i lut0, + __m512i lut1) +{ + if (kDepth <= 4) { + return _mm512_shuffle_epi8(lut0, indices); + } + if (kDepth <= 6) { + return _mm512_permutexvar_epi8(indices, lut0); + } + assert(kDepth == 7); + return shuffleSymbols7(indices, lut0, lut1); +} + +/** + * Expands a flat leaf at depth 2..7: reads @p outSize packed @p kDepth-bit + * indices from @p bitmap and writes `symbols[index]` for each. Per 64-symbol + * block: unpackFlatDepth recovers the indices, shuffleSymbolsDepth maps them to + * symbols. Loop-invariant control vectors (the generated permute/shift tables, + * the unpack mask, and the index mask) and the symbol LUTs are built once up + * front. The tail handles the final < 64 symbols with masked load/store. + */ +ZL_AVX512_INLINE void mergeFlatDepthImpl( + uint8_t* out, + size_t outSize, + const uint8_t* bitmap, + size_t bitmapCapacity, + const uint8_t* symbols, + const size_t kDepth) +{ + size_t outIdx = 0; + const __m512i unpackPermute = _mm512_load_si512( + (const void*)ZL_kPivCoHuffmanAvx512UnpackPermute[kDepth - 2]); + const __m512i shiftControl = _mm512_load_si512( + (const void*)ZL_kPivCoHuffmanAvx512ShiftControl[kDepth - 2]); + const __mmask64 unpackMask = unpackGroupByteMask(kDepth); + const size_t inputBytes = 8 * kDepth; + const __mmask64 fullLoadMask = tailMask(inputBytes); + const __m512i indexMask = _mm512_set1_epi8((char)((1u << kDepth) - 1u)); + + ZL_ALIGNED(64) uint8_t lut0Bytes[64]; + ZL_ALIGNED(64) uint8_t lut1Bytes[64]; + buildShuffleTables(lut0Bytes, lut1Bytes, kDepth, symbols); + const __m512i lut0 = _mm512_load_si512((const void*)lut0Bytes); + const __m512i lut1 = _mm512_load_si512((const void*)lut1Bytes); + + (void)bitmapCapacity; + for (; outIdx + 64 <= outSize; outIdx += 64) { + const size_t byteOffset = (outIdx / 8) * kDepth; + // Load only the bytes needed for this 64-symbol block + const __m512i indices = unpackFlatDepth( + bitmap + byteOffset, + fullLoadMask, + unpackPermute, + unpackMask, + shiftControl, + indexMask); + _mm512_storeu_si512( + out + outIdx, shuffleSymbolsDepth(kDepth, indices, lut0, lut1)); + } + + if (outIdx < outSize) { + const size_t remaining = outSize - outIdx; + const size_t byteOffset = (outIdx / 8) * kDepth; + const size_t tailBytes = (remaining * kDepth + 7) / 8; + // Mask both the packed input bytes and output symbols + const __m512i indices = unpackFlatDepth( + bitmap + byteOffset, + tailMask(tailBytes), + unpackPermute, + unpackMask, + shiftControl, + indexMask); + _mm512_mask_storeu_epi8( + out + outIdx, + tailMask(remaining), + shuffleSymbolsDepth(kDepth, indices, lut0, lut1)); + } +} + +// Generates mergeFlatDepth2..7, each pinning the depth so the unpack shifts and +// table indexing fold to constants. +# define DEFINE_MERGE_FLAT_DEPTH(DEPTH) \ + static ZL_AVX512_ATTR void mergeFlatDepth##DEPTH( \ + uint8_t* out, \ + size_t outSize, \ + const uint8_t* bitmap, \ + size_t bitmapCapacity, \ + const uint8_t* symbols) \ + { \ + mergeFlatDepthImpl( \ + out, outSize, bitmap, bitmapCapacity, symbols, DEPTH); \ + } + +DEFINE_MERGE_FLAT_DEPTH(2) +DEFINE_MERGE_FLAT_DEPTH(3) +DEFINE_MERGE_FLAT_DEPTH(4) +DEFINE_MERGE_FLAT_DEPTH(5) +DEFINE_MERGE_FLAT_DEPTH(6) +DEFINE_MERGE_FLAT_DEPTH(7) + +/// Expands a depth-1 flat leaf: one bit per symbol selecting symbols[0] (0) or +/// symbols[1] (1). Each 64-bit packed word becomes a 64-lane blend, so no +/// unpacking or table is needed. +static ZL_AVX512_ATTR void mergeFlatDepth1( + uint8_t* out, + size_t outSize, + const uint8_t* bitmap, + size_t bitmapCapacity, + const uint8_t* symbols) +{ + (void)bitmapCapacity; + + size_t outIdx = 0; + const __m512i lhs = _mm512_set1_epi8((char)symbols[0]); + const __m512i rhs = _mm512_set1_epi8((char)symbols[1]); + for (; outIdx + 64 <= outSize; outIdx += 64) { + const __mmask64 mask = (__mmask64)ZL_readLE64(bitmap + outIdx / 8); + _mm512_storeu_si512(out + outIdx, _mm512_mask_mov_epi8(lhs, mask, rhs)); + } + if (outIdx < outSize) { + const size_t remaining = outSize - outIdx; + const size_t byteOffset = outIdx / 8; + const size_t tailBytes = (remaining + 7) / 8; + const __mmask64 valid = tailMask(remaining); + const __m128i packed = _mm_maskz_loadu_epi8( + (__mmask16)tailMask(tailBytes), bitmap + byteOffset); + // Move the packed tail bytes into a mask and ignore invalid lanes. + const __mmask64 mask = (__mmask64)_mm_cvtsi128_si64(packed) & valid; + _mm512_mask_storeu_epi8( + out + outIdx, valid, _mm512_mask_mov_epi8(lhs, mask, rhs)); + } +} + +static ZL_AVX512_ATTR void mergeFlatDepth8( + uint8_t* out, + size_t outSize, + const uint8_t* bitmap, + size_t bitmapCapacity, + const uint8_t* symbols) +{ + // Just auto-vectorize + assert(bitmapCapacity >= outSize); + (void)bitmapCapacity; + for (size_t i = 0; i < outSize; ++i) { + out[i] = symbols[bitmap[i]]; + } +} + +/// Entry point for flat-leaf expansion: dispatches to the depth-specialized +/// unpacker (depths 1..8). @see ZL_PivCoHuffmanDecode::mergeFlatDepth. +static ZL_AVX512_ATTR 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(outCapacity >= outSize); + assert(depth >= 1 && depth <= 8); + assert(bitmapCapacity >= (outSize * depth + 7) / 8); + (void)outCapacity; + + // The contract guarantees the exact packed bytes are readable. Extra + // bitmap capacity only determines when 64-byte overreads are allowed. + switch (depth) { + case 1: + mergeFlatDepth1(out, outSize, bitmap, bitmapCapacity, symbols); + return; + case 2: + mergeFlatDepth2(out, outSize, bitmap, bitmapCapacity, symbols); + return; + case 3: + mergeFlatDepth3(out, outSize, bitmap, bitmapCapacity, symbols); + return; + case 4: + mergeFlatDepth4(out, outSize, bitmap, bitmapCapacity, symbols); + return; + case 5: + mergeFlatDepth5(out, outSize, bitmap, bitmapCapacity, symbols); + return; + case 6: + mergeFlatDepth6(out, outSize, bitmap, bitmapCapacity, symbols); + return; + case 7: + mergeFlatDepth7(out, outSize, bitmap, bitmapCapacity, symbols); + return; + default: + mergeFlatDepth8(out, outSize, bitmap, bitmapCapacity, symbols); + return; + } +} + +const ZL_PivCoHuffmanDecode ZL_PivCoHuffmanDecode_avx512 = { + .supported = supported, + .mergeVectorVector = mergeVectorVector, + .mergeConstantVector = mergeConstantVector, + .mergeVectorConstant = mergeVectorConstant, + .mergeFlatDepth = mergeFlatDepth, +}; +#else + +/// Non-x86-64 build: the AVX-512 kernels don't exist, so report unsupported and +/// leave the rest of the function table NULL. +static bool supported(const ZL_cpuid_t* cpuid) +{ + (void)cpuid; + return false; +} + +const ZL_PivCoHuffmanDecode ZL_PivCoHuffmanDecode_avx512 = { + .supported = supported, +}; +#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..ead519dfa --- /dev/null +++ b/src/openzl/codecs/pivco_huffman/arch/encode_pivco_arch.c @@ -0,0 +1,155 @@ +// 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) +{ + ZL_cpuid_t localCpuid; + if (cpuid == NULL) { + localCpuid = ZL_cpuid(); + cpuid = &localCpuid; + } + + if (ZL_PivCoHuffmanEncode_avx512.supported(cpuid)) { + return &ZL_PivCoHuffmanEncode_avx512; + } + 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..f8bb3aa67 --- /dev/null +++ b/src/openzl/codecs/pivco_huffman/arch/encode_pivco_arch.h @@ -0,0 +1,128 @@ +// 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; +extern const ZL_PivCoHuffmanEncode ZL_PivCoHuffmanEncode_avx512; + +ZL_END_C_DECLS + +#endif diff --git a/src/openzl/codecs/pivco_huffman/arch/encode_pivco_avx512.c b/src/openzl/codecs/pivco_huffman/arch/encode_pivco_avx512.c new file mode 100644 index 000000000..b2b4d7c68 --- /dev/null +++ b/src/openzl/codecs/pivco_huffman/arch/encode_pivco_avx512.c @@ -0,0 +1,507 @@ +// Copyright (c) Meta Platforms, Inc. and affiliates. + +#include "openzl/codecs/pivco_huffman/arch/encode_pivco_arch.h" + +#if ZL_ARCH_X86_64 && ZL_HAS_ATTRIBUTE(__target__) + +# include +# include + +# include "openzl/shared/bits.h" +# include "openzl/shared/mem.h" + +# define ZL_AVX512_ATTR \ + ZL_TARGET_ATTRIBUTE( \ + "avx512vbmi,avx512vbmi2,avx512f,avx512vl,avx512bw,bmi2") +# define ZL_AVX512_INLINE ZL_FORCE_INLINE ZL_AVX512_ATTR + +static bool supported(const ZL_cpuid_t* cpuid) +{ + return cpuid != NULL && ZL_cpuid_avx512vbmi(*cpuid) + && ZL_cpuid_avx512vbmi2(*cpuid) && ZL_cpuid_avx512f(*cpuid) + && ZL_cpuid_avx512vl(*cpuid) && ZL_cpuid_avx512bw(*cpuid) + && ZL_cpuid_bmi2(*cpuid); +} + +/// @returns a 64-bit lane mask with the low @p lanes bits set (all 64 when +/// @p lanes >= 64). Used to select the valid byte lanes of a partial 64-element +/// tail block. +ZL_AVX512_INLINE __mmask64 tailMask64(size_t lanes) +{ + return lanes >= 64 ? ~(__mmask64)0 : (((__mmask64)1 << lanes) - 1); +} + +/// @returns the number of whole bytes needed to hold @p lanes packed +/// @p depth-bit indices (ceil(lanes*depth / 8)). +ZL_AVX512_INLINE size_t packedBytes(size_t lanes, size_t depth) +{ + return (lanes * depth + 7) / 8; +} + +/** + * @param kPartitionLhs/kPartitionRhs Compile-time flags (always passed as + * literals so the branches fold away) selecting which child streams to produce. + * They are used instead of testing `lhs == NULL` / `rhs == NULL` because + * partitionFull may legitimately be called with a NULL child, and a runtime + * NULL test would add a branch to the hot loop. + * + * @note Writes whole 64-bit words to @p bitmap, so the final word may spill up + * to 7 bytes past `(numRanks + 7) / 8`; covered by SLOP. + */ +ZL_AVX512_INLINE size_t partitionImpl( + uint8_t* bitmap, + uint8_t* lhs, + uint8_t* rhs, + const uint8_t* ranks, + size_t numRanks, + uint8_t rightRank, + const bool kPartitionLhs, + const bool kPartitionRhs) +{ + const __m512i threshold = _mm512_set1_epi8((char)rightRank); + size_t zeros = 0; + size_t ones = 0; + size_t i = 0; + + for (; i + 64 <= numRanks; i += 64) { + const __m512i rankVec = _mm512_loadu_si512((const void*)(ranks + i)); + // One mask bit per lane: set where rank >= rightRank (the right child). + // The 64-bit mask is the partition bitmap for these 64 ranks. + const __mmask64 bits = + _mm512_cmp_epu8_mask(rankVec, threshold, _MM_CMPINT_GE); + ZL_writeLE64(bitmap + i / 8, (uint64_t)bits); + + const size_t blockOnes = (size_t)ZL_popcount64((uint64_t)bits); + if (kPartitionRhs) { + // Compress-store gathers the masked lanes into a contiguous run, + // appending this block's right-child ranks after the previous ones. + _mm512_mask_compressstoreu_epi8(rhs + ones, bits, rankVec); + ones += blockOnes; + } + if (kPartitionLhs) { + // ~bits selects the left child; in a full 64-lane block every + // inverted bit is a valid lane, so no extra masking is needed. + _mm512_mask_compressstoreu_epi8(lhs + zeros, ~bits, rankVec); + zeros += 64 - blockOnes; + } + } + + if (i < numRanks) { + const size_t lanes = numRanks - i; + const __mmask64 valid = tailMask64(lanes); + const __m512i rankVec = + _mm512_maskz_loadu_epi8(valid, (const void*)(ranks + i)); + // Force out-of-range tail lanes to 0 so they never look like a 1-bit. + const __mmask64 bits = + _mm512_cmp_epu8_mask(rankVec, threshold, _MM_CMPINT_GE) & valid; + // Writes up to 7 bytes beyond the end of bitmap. Ok because of SLOP. + ZL_writeLE64(bitmap + i / 8, (uint64_t)bits); + + const size_t blockOnes = (size_t)ZL_popcount64((uint64_t)bits); + if (kPartitionRhs) { + _mm512_mask_compressstoreu_epi8(rhs + ones, bits, rankVec); + ones += blockOnes; + } + if (kPartitionLhs) { + // For the tail, restrict the left child to valid lanes as well, + // otherwise the padding lanes (which are 0-bits) would be stored. + _mm512_mask_compressstoreu_epi8( + lhs + zeros, ~bits & valid, rankVec); + zeros += lanes - blockOnes; + } + } + + if (kPartitionRhs) { + return ones; + } else if (kPartitionLhs) { + return numRanks - zeros; + } else { + return 0; + } +} + +static ZL_AVX512_ATTR size_t partitionFull( + uint8_t* bitmap, + uint8_t* lhs, + uint8_t* rhs, + const uint8_t* ranks, + size_t numRanks, + uint8_t rightRank) +{ + return partitionImpl( + bitmap, lhs, rhs, ranks, numRanks, rightRank, true, true); +} + +static ZL_AVX512_ATTR size_t partitionLeft( + uint8_t* bitmap, + uint8_t* lhs, + const uint8_t* ranks, + size_t numRanks, + uint8_t rightRank) +{ + return partitionImpl( + bitmap, lhs, NULL, ranks, numRanks, rightRank, true, false); +} + +static ZL_AVX512_ATTR size_t partitionRight( + uint8_t* bitmap, + uint8_t* rhs, + const uint8_t* ranks, + size_t numRanks, + uint8_t rightRank) +{ + return partitionImpl( + bitmap, NULL, rhs, ranks, numRanks, rightRank, false, true); +} + +static ZL_AVX512_ATTR void partitionNone( + uint8_t* bitmap, + const uint8_t* ranks, + size_t numRanks, + uint8_t rightRank) +{ + (void)partitionImpl( + bitmap, NULL, NULL, ranks, numRanks, rightRank, false, false); +} + +/// Loads the @p valid lanes of 64 ranks and subtracts @p rankBegin from each, +/// yielding one flat-leaf index per byte lane (the rank's offset within its +/// flat leaf, < 2^depth). Lanes outside @p valid are forced to 0 so padding +/// never contributes spurious set bits to the packed output. Pass @p valid == +/// -1 for a full block: the all-ones mask folds away to a plain load with no +/// zeroing. +ZL_AVX512_INLINE __m512i +loadRankIndexBytes64(const uint8_t* ranks, uint8_t rankBegin, __mmask64 valid) +{ + const __m512i rankVec = _mm512_maskz_loadu_epi8(valid, (const void*)ranks); + const __m512i indices = + _mm512_sub_epi8(rankVec, _mm512_set1_epi8((char)rankBegin)); + return _mm512_maskz_mov_epi8(valid, indices); +} + +/** + * Fuses each adjacent (even, odd) pair of byte indices into one 16-bit lane, + * concatenating their @p depth-bit fields: result lane = even | (odd << depth). + * This is the first packing step for depths 2..7 (a 64-byte vector becomes 32 + * 16-bit lanes each holding two indices). + * + * Each index is < 2^depth, so even and odd never overlap and the "sum" the + * intrinsics compute is exactly a bitwise concatenation. + * + * For depth < 7 this is a single `maddubs`: it multiplies each even byte by 1 + * and each odd byte by `1 << depth`, summing the adjacent pair into a 16-bit + * lane. depth == 7 cannot use that path: the odd-lane multiplier would be + * `1 << 7 == 0x80`, which `maddubs` treats as a *signed* -128. So depth 7 falls + * back to mask-and-shift: keep the even byte's low 7 bits, shift the odd byte's + * 7 bits down by one (from bit 8 to bit 7), and OR them together. + */ +ZL_AVX512_INLINE __m512i packBytePairs16(__m512i indices, const size_t kDepth) +{ + if (kDepth == 7) { + const __m512i lo = + _mm512_and_si512(indices, _mm512_set1_epi16((short)0x007f)); + const __m512i hi = _mm512_srli_epi16( + _mm512_and_si512(indices, _mm512_set1_epi16((short)0x7f00)), 1); + return _mm512_or_si512(lo, hi); + } + + const int pairMultiplierValue = (int)(((1u << kDepth) << 8) | 1u); + return _mm512_maddubs_epi16( + indices, _mm512_set1_epi16((short)pairMultiplierValue)); +} + +/** + * Packs one 64-index block at depth 2 (output: 16 bytes, 4 indices/byte). + * `packBytePairs16` first fuses pairs into 4-bit fields (two 2-bit indices per + * 16-bit lane); then `madd_epi16` with {1, 1<<4} fuses two adjacent lanes into + * an 8-bit field (four 2-bit indices) in the low byte of each 32-bit lane; + * finally `cvtepi32_storeu_epi8` narrows each 32-bit lane to that low byte. + */ +ZL_AVX512_INLINE void +packFlatDepth2Block(uint8_t* bitmap, __m512i indices, __mmask16 storeMask) +{ + const __m512i pairs16 = packBytePairs16(indices, 2); + const __m512i quads32 = + _mm512_madd_epi16(pairs16, _mm512_set1_epi32(0x00100001)); + _mm512_mask_cvtepi32_storeu_epi8(bitmap, storeMask, quads32); +} + +/// Packs depth-2 indices (2 bits each, 16 bytes per 64-index block). +static ZL_AVX512_ATTR void packFlatDepth2( + uint8_t* bitmap, + const uint8_t* ranks, + size_t numRanks, + uint8_t rankBegin) +{ + size_t idx = 0; + size_t outIdx = 0; + for (; idx + 64 <= numRanks; idx += 64) { + packFlatDepth2Block( + bitmap + outIdx, + loadRankIndexBytes64(ranks + idx, rankBegin, (__mmask64)-1), + (__mmask16)0xffff); + outIdx += 16; + } + if (idx < numRanks) { + const size_t lanes = numRanks - idx; + const __mmask64 valid = tailMask64(lanes); + packFlatDepth2Block( + bitmap + outIdx, + loadRankIndexBytes64(ranks + idx, rankBegin, valid), + (__mmask16)tailMask64(packedBytes(lanes, 2))); + } +} + +/** + * Packs one 64-index block at depth 4 (output: 32 bytes, 2 indices/byte). + * `packBytePairs16` already fuses each pair into an 8-bit field (two 4-bit + * indices) in the low byte of every 16-bit lane, so we just narrow each 16-bit + * lane to its low byte. + */ +ZL_AVX512_INLINE void +packFlatDepth4Block(uint8_t* bitmap, __m512i indices, __mmask32 storeMask) +{ + const __m512i pairs16 = packBytePairs16(indices, 4); + _mm512_mask_cvtepi16_storeu_epi8(bitmap, storeMask, pairs16); +} + +/// Packs depth-4 indices (4 bits each, 32 bytes per 64-index block). +static ZL_AVX512_ATTR void packFlatDepth4( + uint8_t* bitmap, + const uint8_t* ranks, + size_t numRanks, + uint8_t rankBegin) +{ + size_t idx = 0; + size_t outIdx = 0; + for (; idx + 64 <= numRanks; idx += 64) { + packFlatDepth4Block( + bitmap + outIdx, + loadRankIndexBytes64(ranks + idx, rankBegin, (__mmask64)-1), + (__mmask32)0xffffffff); + outIdx += 32; + } + if (idx < numRanks) { + const size_t lanes = numRanks - idx; + const __mmask64 valid = tailMask64(lanes); + packFlatDepth4Block( + bitmap + outIdx, + loadRankIndexBytes64(ranks + idx, rankBegin, valid), + (__mmask32)tailMask64(packedBytes(lanes, 4))); + } +} + +/** + * Builds the compress-store mask for the bitpack packer. The packer produces 8 + * bytes per group of 8 indices but only the low `packedBytes(8, kDepth)` of + * each are meaningful; this mask selects, for each of the 8 groups, exactly the + * meaningful bytes so `compressstoreu` writes them contiguously with no gaps. + * + * @param lanes Number of valid indices (64 for a full block, fewer for a tail); + * groups beyond @p lanes contribute nothing. + */ +ZL_AVX512_INLINE __mmask64 packBitpackStoreMask(size_t kDepth, size_t lanes) +{ + __mmask64 storeMask = 0; + for (size_t group = 0; group < 8; ++group) { + const size_t firstLane = group * 8; + if (lanes <= firstLane) { + break; + } + + size_t lanesInGroup = lanes - firstLane; + if (lanesInGroup > 8) { + lanesInGroup = 8; + } + + // Each group occupies one byte-octet of the mask; set the low + // bytesInGroup bits of that octet. + const size_t bytesInGroup = packedBytes(lanesInGroup, kDepth); + storeMask |= (__mmask64)(tailMask64(bytesInGroup) << (group * 8)); + } + return storeMask; +} + +/** + * Packs one 64-index block for the depths (3, 5, 6, 7) whose fields don't land + * on byte boundaries. Strategy: build 8*kDepth contiguous bits per group of 8 + * indices inside one 64-bit lane, then compress out the padding bytes. + * + * Steps: + * 1. packBytePairs16 -> 2 indices (2*kDepth bits) per 16-bit lane. + * 2. madd_epi16 with {1, 1<<(2*kDepth)} -> 4 indices (4*kDepth bits) in the + * low bits of each 32-bit lane (quads32). + * 3. Fuse the two 32-bit halves of each 64-bit lane into 8 contiguous indices: + * keep the low half as-is, and shift the high half down so its 4*kDepth + * bits begin right where the low half's bits end (at bit 4*kDepth) -- i.e. + * right by (32 - 4*kDepth). OR them into octets64 (8*kDepth contiguous + * bits/lane). Those bits already sit in the lane's low `kDepth` bytes in + * little-endian order, so no further byte extraction is needed. + * 4. compressstoreu drops the padding bytes (the high 8-kDepth bytes of each + * group) via @p storeMask, writing the packed bytes contiguously. + */ +ZL_AVX512_INLINE void packFlatDepthBitpackBlock( + uint8_t* bitmap, + __m512i indices, + __mmask64 storeMask, + const size_t kDepth) +{ + const __m512i kQuadPackMultiplier = + _mm512_set1_epi32((int)(((1u << (2 * kDepth)) << 16) | 1u)); + const __m512i kLow32LaneMask = _mm512_set1_epi64((long long)0xffffffffULL); + const __m512i pairs16 = packBytePairs16(indices, kDepth); + const __m512i quads32 = _mm512_madd_epi16(pairs16, kQuadPackMultiplier); + const __m512i lowQuads32 = _mm512_and_si512(quads32, kLow32LaneMask); + const __m512i highQuads32 = _mm512_andnot_si512(kLow32LaneMask, quads32); + const __m512i octets64 = _mm512_or_si512( + lowQuads32, + _mm512_srli_epi64(highQuads32, (unsigned int)(32 - 4 * kDepth))); + _mm512_mask_compressstoreu_epi8(bitmap, storeMask, octets64); +} + +/** + * Drives packFlatDepthBitpackBlock over the whole input, precomputing the + * loop-invariant control vector: + * - quadPackMultiplier {1, 1<<(2*depth)} fuses two pairs into four indices. + */ +ZL_AVX512_INLINE void packFlatDepthBitpackImpl( + uint8_t* bitmap, + const uint8_t* ranks, + size_t numRanks, + uint8_t rankBegin, + const size_t kDepth) +{ + const __mmask64 fullStoreMask = packBitpackStoreMask(kDepth, 64); + + size_t idx = 0; + size_t outIdx = 0; + for (; idx + 64 <= numRanks; idx += 64) { + packFlatDepthBitpackBlock( + bitmap + outIdx, + loadRankIndexBytes64(ranks + idx, rankBegin, (__mmask64)-1), + fullStoreMask, + kDepth); + // Each full block emits depth bytes per 8 indices == 8*depth bytes. + outIdx += 8 * kDepth; + } + if (idx < numRanks) { + const size_t lanes = numRanks - idx; + const __mmask64 valid = tailMask64(lanes); + packFlatDepthBitpackBlock( + bitmap + outIdx, + loadRankIndexBytes64(ranks + idx, rankBegin, valid), + packBitpackStoreMask(kDepth, lanes), + kDepth); + } +} + +// Generates packFlatDepth3/5/6/7, each a thin wrapper that pins the depth so +// the shifts/multipliers in the bitpack impl fold to constants. +# define DEFINE_PACK_FLAT_DEPTH_BITPACK(DEPTH) \ + static ZL_AVX512_ATTR void packFlatDepth##DEPTH( \ + uint8_t* bitmap, \ + const uint8_t* ranks, \ + size_t numRanks, \ + uint8_t rankBegin) \ + { \ + packFlatDepthBitpackImpl( \ + bitmap, ranks, numRanks, rankBegin, DEPTH); \ + } + +DEFINE_PACK_FLAT_DEPTH_BITPACK(3) +DEFINE_PACK_FLAT_DEPTH_BITPACK(5) +DEFINE_PACK_FLAT_DEPTH_BITPACK(6) +DEFINE_PACK_FLAT_DEPTH_BITPACK(7) + +# undef DEFINE_PACK_FLAT_DEPTH_BITPACK + +/// Packs depth-8 indices: each index already fills a whole byte, so packing is +/// just `rank - rankBegin` stored one byte per index (64 bytes per block). +static ZL_AVX512_ATTR void packFlatDepth8( + uint8_t* bitmap, + const uint8_t* ranks, + size_t numRanks, + uint8_t rankBegin) +{ + size_t idx = 0; + size_t outIdx = 0; + for (; idx + 64 <= numRanks; idx += 64) { + _mm512_storeu_si512( + (void*)(bitmap + outIdx), + loadRankIndexBytes64(ranks + idx, rankBegin, (__mmask64)-1)); + outIdx += 64; + } + if (idx < numRanks) { + const size_t lanes = numRanks - idx; + const __mmask64 valid = tailMask64(lanes); + _mm512_mask_storeu_epi8( + bitmap + outIdx, + valid, + loadRankIndexBytes64(ranks + idx, rankBegin, valid)); + } +} + +/// Entry point for flat-leaf packing: dispatches to the depth-specialized +/// packer (depths 1..8). @see ZL_PivCoHuffmanEncode::packFlatDepth. +static ZL_AVX512_ATTR void packFlatDepth( + uint8_t* bitmap, + size_t depth, + const uint8_t* ranks, + size_t numRanks, + uint8_t rankBegin) +{ + switch (depth) { + case 1: + // Exactly equivalent to partitionNone, but need to pass in + // rightRank (split point) instead of rankBegin. + partitionNone(bitmap, ranks, numRanks, rankBegin + 1); + return; + case 2: + packFlatDepth2(bitmap, ranks, numRanks, rankBegin); + return; + case 3: + packFlatDepth3(bitmap, ranks, numRanks, rankBegin); + return; + case 4: + packFlatDepth4(bitmap, ranks, numRanks, rankBegin); + return; + case 5: + packFlatDepth5(bitmap, ranks, numRanks, rankBegin); + return; + case 6: + packFlatDepth6(bitmap, ranks, numRanks, rankBegin); + return; + case 7: + packFlatDepth7(bitmap, ranks, numRanks, rankBegin); + return; + default: + assert(depth == 8); + packFlatDepth8(bitmap, ranks, numRanks, rankBegin); + return; + } +} + +const ZL_PivCoHuffmanEncode ZL_PivCoHuffmanEncode_avx512 = { + .supported = supported, + .partitionFull = partitionFull, + .partitionLeft = partitionLeft, + .partitionRight = partitionRight, + .partitionNone = partitionNone, + .packFlatDepth = packFlatDepth, +}; + +#else + +/// Non-x86-64 build: the AVX-512 kernels don't exist, so report unsupported and +/// leave the rest of the function table NULL. +static bool supported(const ZL_cpuid_t* cpuid) +{ + (void)cpuid; + return false; +} + +const ZL_PivCoHuffmanEncode ZL_PivCoHuffmanEncode_avx512 = { + .supported = supported, +}; + +#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_binding.c b/src/openzl/codecs/pivco_huffman/decode_pivco_binding.c new file mode 100644 index 000000000..828eee276 --- /dev/null +++ b/src/openzl/codecs/pivco_huffman/decode_pivco_binding.c @@ -0,0 +1,101 @@ +// Copyright (c) Meta Platforms, Inc. and affiliates. + +#include "openzl/codecs/pivco_huffman/decode_pivco_binding.h" + +#include + +#include "openzl/codecs/pivco_huffman/common_pivco_kernel.h" +#include "openzl/codecs/pivco_huffman/decode_pivco_kernel.h" +#include "openzl/shared/varint.h" +#include "openzl/zl_data.h" +#include "openzl/zl_errors.h" + +typedef struct { + size_t decodedSize; + size_t blockSize; +} PivCoHuffmanHeader; + +static ZL_Report decodeHeader(ZL_Decoder* dictx, PivCoHuffmanHeader* parsed) +{ + ZL_RESULT_DECLARE_SCOPE_REPORT(dictx); + ZL_ERR_IF_NULL(parsed, GENERIC); + + const ZL_RBuffer header = ZL_Decoder_getCodecHeader(dictx); + ZL_ERR_IF_EQ(header.size, 0, corruption); + ZL_ERR_IF_NULL(header.start, corruption); + const uint8_t* ptr = (const uint8_t*)header.start; + const uint8_t* const end = ptr + header.size; + + ZL_TRY_LET_CONST(uint64_t, decodedSize64, ZL_varintDecode(&ptr, end)); + ZL_ERR_IF_GT(decodedSize64, (uint64_t)SIZE_MAX, corruption); + parsed->decodedSize = (size_t)decodedSize64; + + // The block size is optional: when present it follows the decoded size, + // otherwise it defaults to the decoded size (a single block). + if (ptr != end) { + ZL_TRY_LET_CONST(uint64_t, blockSize64, ZL_varintDecode(&ptr, end)); + ZL_ERR_IF_GT(blockSize64, (uint64_t)SIZE_MAX, corruption); + ZL_ERR_IF_EQ(blockSize64, 0, corruption); + parsed->blockSize = (size_t)blockSize64; + } else { + parsed->blockSize = parsed->decodedSize; + } + // Explicitly ignore unconsumed bytes in the header. + // This allows the encoder to add extra information in the future without + // breaking backwards compatibility. + // For example: The encoder could add a jump table to the offset of each + // encoded block for parallel decoding. + + return ZL_returnSuccess(); +} + +ZL_Report DI_pivco_huffman(ZL_Decoder* dictx, const ZL_Input* ins[]) +{ + ZL_RESULT_DECLARE_SCOPE_REPORT(dictx); + + const ZL_Input* const weightsStream = ins[0]; + const ZL_Input* const bitstream = ins[1]; + + ZL_ASSERT_EQ(ZL_Input_type(weightsStream), ZL_Type_numeric); + ZL_ERR_IF_NE(ZL_Input_eltWidth(weightsStream), 1, corruption); + + ZL_ASSERT_EQ(ZL_Input_type(bitstream), ZL_Type_serial); + ZL_ASSERT_EQ(ZL_Input_eltWidth(bitstream), 1); + + PivCoHuffmanHeader header = { 0, 0 }; + ZL_ERR_IF_ERR(decodeHeader(dictx, &header)); + + const size_t weightsSize = ZL_Input_numElts(weightsStream); + const size_t bitstreamSize = ZL_Input_numElts(bitstream); + ZL_ERR_IF_GT(weightsSize, ZL_PIVCO_MAX_SYMBOLS, corruption); + ZL_ERR_IF_NE(header.decodedSize == 0, weightsSize == 0, corruption); + + ZL_Output* const out = + ZL_Decoder_create1OutStream(dictx, header.decodedSize, 1); + ZL_ERR_IF_NULL(out, allocation); + + const size_t scratchBytes = ZL_PivCoHuffmanDecode_scratchBytes( + header.decodedSize, header.blockSize); + uint8_t* const scratch = ZL_Decoder_getScratchSpace(dictx, scratchBytes); + ZL_ERR_IF_NULL(scratch, allocation); + + const uint8_t* const weights = ZL_Input_ptr(weightsStream); + ZL_ERR_IF_NULL(weights, corruption); + ZL_ERR_IF_NOT( + ZL_PivCoHuffman_decode( + ZL_Output_ptr(out), + header.decodedSize, + scratch, + scratchBytes, + weights, + weightsSize, + ZL_Input_ptr(bitstream), + bitstreamSize, + header.blockSize, + NULL), + corruption, + "PivCo-Huffman decode failed"); + ZL_ERR_IF_ERR(ZL_Output_commit(out, header.decodedSize)); + + return ZL_returnSuccess(); +} diff --git a/src/openzl/codecs/pivco_huffman/decode_pivco_binding.h b/src/openzl/codecs/pivco_huffman/decode_pivco_binding.h new file mode 100644 index 000000000..1b40e25bf --- /dev/null +++ b/src/openzl/codecs/pivco_huffman/decode_pivco_binding.h @@ -0,0 +1,22 @@ +// Copyright (c) Meta Platforms, Inc. and affiliates. + +#ifndef OPENZL_CODECS_PIVCO_HUFFMAN_DECODE_PIVCO_BINDING_H +#define OPENZL_CODECS_PIVCO_HUFFMAN_DECODE_PIVCO_BINDING_H + +#include "openzl/codecs/pivco_huffman/graph_pivco_huffman.h" +#include "openzl/shared/portability.h" +#include "openzl/zl_dtransform.h" + +ZL_BEGIN_C_DECLS + +ZL_Report DI_pivco_huffman(ZL_Decoder* dictx, const ZL_Input* ins[]); + +#define DI_PIVCO_HUFFMAN(id) \ + { \ + .transform_f = DI_pivco_huffman, \ + .name = "!zl.pivco_huffman", \ + } + +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..45925cc73 --- /dev/null +++ b/src/openzl/codecs/pivco_huffman/decode_pivco_kernel.h @@ -0,0 +1,51 @@ +// 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 weights The Huffman weights to use for decoding, which will be + * validated during decoding. + * @param bitstream The PivCo Huffman encoded bitstream, which will be + * validated during decoding. + * @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_binding.c b/src/openzl/codecs/pivco_huffman/encode_pivco_binding.c new file mode 100644 index 000000000..dd503336a --- /dev/null +++ b/src/openzl/codecs/pivco_huffman/encode_pivco_binding.c @@ -0,0 +1,156 @@ +// Copyright (c) Meta Platforms, Inc. and affiliates. + +#include "openzl/codecs/pivco_huffman/encode_pivco_binding.h" + +#define HUF_STATIC_LINKING_ONLY + +#include +#include + +#include "openzl/codecs/pivco_huffman/common_pivco_kernel.h" +#include "openzl/codecs/pivco_huffman/encode_pivco_kernel.h" +#include "openzl/fse/huf.h" +#include "openzl/shared/histogram.h" +#include "openzl/shared/varint.h" +#include "openzl/zl_data.h" +#include "openzl/zl_errors.h" + +static ZL_Report buildWeights( + ZL_Encoder* eictx, + uint8_t weights[ZL_PIVCO_MAX_SYMBOLS], + size_t* weightsSize, + const uint8_t* src, + size_t srcSize) +{ + ZL_RESULT_DECLARE_SCOPE_REPORT(eictx); + ZL_ERR_IF_NULL(weights, GENERIC); + ZL_ERR_IF_NULL(weightsSize, GENERIC); + *weightsSize = 0; + memset(weights, 0, ZL_PIVCO_MAX_SYMBOLS); + + if (srcSize == 0) { + return ZL_returnValue(0); + } + ZL_ERR_IF_NULL(src, node_invalid_input); + ZL_ERR_IF_GT( + srcSize, + UINT_MAX, + node_invalid_input, + "PivCo-Huffman input is too large to histogram"); + + ZL_Histogram8 hist; + ZL_Histogram_init(&hist.base, 255); + ZL_Histogram_build(&hist.base, src, srcSize, 1); + + const uint32_t maxSymbol = hist.base.maxSymbol; + const size_t cardinality = hist.base.cardinality; + const uint32_t* const count = hist.base.count; + + *weightsSize = (size_t)maxSymbol + 1; + if (cardinality == 1) { + weights[maxSymbol] = 1; + return ZL_returnValue(0); + } + + HUF_CREATE_STATIC_CTABLE(ctable, HUF_SYMBOLVALUE_MAX); + unsigned tableLog = + HUF_optimalTableLog(ZL_PIVCO_MAX_TABLE_LOG, srcSize, maxSymbol); + const size_t hufRet = HUF_buildCTable(ctable, count, maxSymbol, tableLog); + ZL_ERR_IF( + HUF_isError(hufRet), + GENERIC, + "HUF_buildCTable failed: %s", + HUF_getErrorName(hufRet)); + tableLog = (unsigned)hufRet; + ZL_ERR_IF_GT(tableLog, ZL_PIVCO_MAX_TABLE_LOG, GENERIC); + + for (unsigned symbol = 0; symbol <= maxSymbol; ++symbol) { + const unsigned numBits = HUF_getNbBitsFromCTable(ctable, symbol); + const uint8_t weight = + numBits == 0 ? 0 : (uint8_t)(tableLog + 1 - numBits); + weights[symbol] = weight; + } + + return ZL_returnValue(tableLog); +} + +/** + * Writes the codec header: the decoded size, followed -- only when the input + * spans more than one block -- by the block size. A single-block input omits + * the block size; the decoder then defaults it to the decoded size. + */ +static size_t +encodeHeader(uint8_t* header, size_t decodedSize, size_t blockSize) +{ + uint8_t* ptr = header; + ptr += ZL_varintEncode((uint64_t)decodedSize, ptr); + if (blockSize < decodedSize) { + ptr += ZL_varintEncode((uint64_t)blockSize, ptr); + } + return (size_t)(ptr - header); +} + +ZL_Report +EI_pivco_huffman(ZL_Encoder* eictx, const ZL_Input* ins[], size_t nbIns) +{ + ZL_RESULT_DECLARE_SCOPE_REPORT(eictx); + + ZL_ASSERT_EQ(nbIns, 1); + const ZL_Input* const in = ins[0]; + ZL_ASSERT_EQ(ZL_Input_type(in), ZL_Type_serial); + + const uint8_t* const src = ZL_Input_ptr(in); + size_t const srcSize = ZL_Input_numElts(in); + + uint8_t weights[ZL_PIVCO_MAX_SYMBOLS]; + size_t weightsSize; + ZL_TRY_LET_CONST( + size_t, + tableLog, + buildWeights(eictx, weights, &weightsSize, src, srcSize)); + + // In the future the block size could be configurable. + size_t const blockSize = ZL_PIVCO_DEFAULT_BLOCK_SIZE; + + // Header is the decoded size plus an optional block size (see + // encodeHeader). + enum { kHeaderCapacity = 2 * ZL_VARINT_LENGTH_64 }; + uint8_t header[kHeaderCapacity]; + size_t const headerSize = encodeHeader(header, srcSize, blockSize); + ZL_Encoder_sendCodecHeader(eictx, header, headerSize); + + ZL_Output* const weightsStream = + ZL_Encoder_createTypedStream(eictx, 0, weightsSize, 1); + ZL_ERR_IF_NULL(weightsStream, allocation); + if (weightsSize != 0) { + memcpy(ZL_Output_ptr(weightsStream), weights, weightsSize); + } + ZL_ERR_IF_ERR(ZL_Output_commit(weightsStream, weightsSize)); + + const size_t dstCapacity = ZL_PivCoHuffmanEncode_bound(srcSize, blockSize); + ZL_Output* const bitstream = + ZL_Encoder_createTypedStream(eictx, 1, dstCapacity, 1); + ZL_ERR_IF_NULL(bitstream, allocation); + + const size_t scratchElements = + ZL_PivCoHuffmanEncode_scratchElements(srcSize, blockSize); + uint8_t* const scratch = ZL_Encoder_getScratchSpace(eictx, scratchElements); + ZL_ERR_IF_NULL(scratch, allocation); + + const size_t dstSize = ZL_PivCoHuffman_encode( + ZL_Output_ptr(bitstream), + dstCapacity, + scratch, + scratchElements, + weights, + weightsSize, + (int)tableLog, + src, + srcSize, + blockSize, + NULL); + ZL_ERR_IF_EQ(dstSize, SIZE_MAX, GENERIC, "PivCo-Huffman encoding failed"); + ZL_ERR_IF_ERR(ZL_Output_commit(bitstream, dstSize)); + + return ZL_returnSuccess(); +} diff --git a/src/openzl/codecs/pivco_huffman/encode_pivco_binding.h b/src/openzl/codecs/pivco_huffman/encode_pivco_binding.h new file mode 100644 index 000000000..c8aee4bcd --- /dev/null +++ b/src/openzl/codecs/pivco_huffman/encode_pivco_binding.h @@ -0,0 +1,24 @@ +// Copyright (c) Meta Platforms, Inc. and affiliates. + +#ifndef OPENZL_CODECS_PIVCO_HUFFMAN_ENCODE_PIVCO_BINDING_H +#define OPENZL_CODECS_PIVCO_HUFFMAN_ENCODE_PIVCO_BINDING_H + +#include "openzl/codecs/pivco_huffman/graph_pivco_huffman.h" +#include "openzl/shared/portability.h" +#include "openzl/zl_ctransform.h" + +ZL_BEGIN_C_DECLS + +ZL_Report +EI_pivco_huffman(ZL_Encoder* eictx, const ZL_Input* ins[], size_t nbIns); + +#define EI_PIVCO_HUFFMAN(id) \ + { \ + .gd = PIVCO_HUFFMAN_GRAPH(id), \ + .transform_f = EI_pivco_huffman, \ + .name = "!zl.pivco_huffman", \ + } + +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..300d24bbf --- /dev/null +++ b/src/openzl/codecs/pivco_huffman/encode_pivco_kernel.h @@ -0,0 +1,68 @@ +// 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 weights The Huffman weights to use for encoding. + * @param tableLog The Huffman table log to use for encoding. + * @param scratch Working buffer of at least + * ZL_PivCoHuffmanEncode_scratchElements(srcSize, blockSize) bytes. + * @param blockSize The number of input bytes coded per pivco block. Must + * match the block size the decoder uses. srcSize > blockSize produces multiple + * blocks. + * @param kernels The kernel to use for encoding or NULL to use the default. + * + * @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/codecs/pivco_huffman/graph_pivco_huffman.h b/src/openzl/codecs/pivco_huffman/graph_pivco_huffman.h new file mode 100644 index 000000000..c745f5cc6 --- /dev/null +++ b/src/openzl/codecs/pivco_huffman/graph_pivco_huffman.h @@ -0,0 +1,15 @@ +// Copyright (c) Meta Platforms, Inc. and affiliates. + +#ifndef OPENZL_CODECS_PIVCO_HUFFMAN_GRAPH_PIVCO_HUFFMAN_H +#define OPENZL_CODECS_PIVCO_HUFFMAN_GRAPH_PIVCO_HUFFMAN_H + +#include "openzl/zl_data.h" + +#define PIVCO_HUFFMAN_GRAPH(id) \ + { \ + .CTid = id, \ + .inputTypes = ZL_STREAMTYPELIST(ZL_Type_serial), \ + .soTypes = ZL_STREAMTYPELIST(ZL_Type_numeric, ZL_Type_serial), \ + } + +#endif diff --git a/src/openzl/codecs/pivco_huffman/spec.md b/src/openzl/codecs/pivco_huffman/spec.md new file mode 100644 index 000000000..349cd96d7 --- /dev/null +++ b/src/openzl/codecs/pivco_huffman/spec.md @@ -0,0 +1,210 @@ +## PivCo-Huffman Decoder Specification + +### Overview + +`pivco_huffman` is a static, byte-oriented Huffman codec. Rather than decoding +one codeword at a time, *pivco* coding decodes a whole block of symbols at +once: it walks the Huffman tree breadth-first and, at each internal node, splits +the block's symbols into the two children according to a single codeword bit. The +split decisions are stored as bitmaps, which the decoder replays to scatter +symbols into place. + +The Huffman tree is never serialized. Both sides rebuild an identical tree +deterministically from the transmitted weights (see *Tree construction*), so the +bitstream carries only the per-node partition bitmaps and a small count per node. + +The codec is endian-independent: every multi-byte quantity on the wire is either +an varint or a little-endian bit packing, both defined below. + +### Inputs + +The decoder takes two streams: + +1. `weights` — a numeric stream of 1-byte zstd-style Huffman weights. +2. `bitstream` — the pivco-coded payload (a serial byte stream). + +It produces one serial stream of `decoded_size` bytes. + +### Codec header + +The codec header is carried out-of-band from `bitstream`. Multi-byte integers are +varints. + +``` +decoded_size : varint +block_size : varint (optional) +... trailing bytes ignored +``` + +- `decoded_size` is the number of output bytes. +- `block_size` is the number of output bytes per pivco block. It is **optional**: + it is present exactly when the encoder split the input into more than one block. + Its presence is detected by whether any header bytes remain after `decoded_size` + has been parsed: + - If absent, `block_size` defaults to `decoded_size` (the whole output is one + block). In which case `decoded_size` must be no greater than `ZL_PIVCO_MAX_BLOCK_SIZE` (2^28). + - If present, it must be non-zero, and the decoder rejects it if it exceeds + `ZL_PIVCO_MAX_BLOCK_SIZE` (2^28). A `block_size` that does not match the one + the encoder used yields corrupt output (caught as a parse failure or by the + frame's integrity check). +- Any bytes after the fields the decoder understands are **ignored**. This is + reserved for future, backward-compatible extensions (for example, a jump table + of per-block bitstream offsets to enable parallel decoding). + +### weights stream + +The `weights` stream length is the number of weight entries and must be no greater +than 256. Weight entry `i` describes byte symbol `i`; symbols beyond the last +entry are implicitly absent. + +Each weight must be in the range 0..12. A zero weight means the symbol is absent +from the alphabet. A non-zero weight follows the zstd convention: a symbol with +code length `L` under a table of `tableLog` bits has weight `tableLog + 1 - L`, so +a larger weight means a shorter code. + +The non-zero weights must describe a complete prefix code. Equivalently, + +``` +sum over non-zero weights of (1 << (weight - 1)) +``` + +must be a non-zero power of two; `tableLog` is its base-2 logarithm and must be no +greater than 12. + +`decoded_size == 0` is encoded with an empty `weights` stream and an empty +`bitstream`; the two emptiness conditions must agree (`decoded_size == 0` iff +`weights` is empty), otherwise the frame is corrupt. + +### Tree construction + +The tree is a function of `weights` alone. An independent decoder must reproduce +the construction below exactly, because the resulting structure (in particular the +flat-leaf grouping) determines the bitstream layout. Encoder and decoder run the +identical deterministic builder. + +1. **Validate and derive `tableLog`** from the weights, as in *weights stream* + above. + +2. **Ranks (canonical order).** Order all present (non-zero weight) symbols by + *descending weight* (shortest code first); break ties by *ascending symbol + value*. A symbol's **rank** is its index in this order. `numRanks` is the count + of present symbols. PivCo coding operates on contiguous rank ranges. + +3. **Canonical codewords.** Visiting ranks in order, assign canonical Huffman + codewords: start at 0, add 1 per leaf, and shift left by 1 each time the code + length increases by one level. Each codeword is considered left-justified + (most-significant bit first). "Level `L`" denotes the `(L+1)`-th bit from the + most-significant end of the codeword; level 0 is the root. + +4. **Flat leaves.** A run of `2^k` symbols that share one code length forms a + complete depth-`k` subtree. The codec collapses such a run into a single + **flat leaf** holding `2^k` symbols, decoded by a `k`-bit table lookup instead + of `k` binary splits. The grouping rule (which both sides must follow + identically): + + Process weight buckets from shallowest (largest weight) to deepest (smallest + weight). Within a bucket, while more than two single-symbol leaves remain, + peel the largest power-of-two block (`2^k`, `k = floor(log2(remaining))`) off + the **end** of the bucket — i.e. the highest symbol values — into a flat leaf + whose weight becomes `weight + k` (a shallower level). A residual of one or two + leaves is left as ordinary single-symbol leaves, since a flat node could not + improve on a single binary split. The shallowest-first order matters: a peeled + flat leaf lands in a shallower (larger-weight) bucket, which has already been + processed, so it is never itself re-flattened. + + Rank and codeword assignment then visits levels shallowest-first; within a + level it takes the remaining single-symbol leaves (ascending symbol value) + followed by the flat leaves collapsed into that level. A flat leaf occupies + `2^k` consecutive ranks, one per contained symbol, and a flat leaf's symbols + occupy a contiguous slice of the rank order. + +A single-symbol leaf has flat depth 0 and is also called a **constant** leaf: it +contributes no bits to the bitstream. The whole-alphabet edge case of one present +symbol produces a tree of one rank, one level, and a constant root. + +### Bitstream model + +`bitstream` is read as bits packed **little-endian, least-significant-bit first**: +bit index `j` lives in byte `j / 8` at bit position `j % 8` (bit 0 is the LSB). +Integers wider than one bit are packed low bits first. The decoder maintains a bit +cursor and uses two read operations: + +- **alignedBitmap(n):** first advance the cursor to the next byte boundary, + discarding 0–7 padding bits; then take the next `ceil(n / 8)` bytes as an + `n`-bit bitmap and advance the cursor by exactly `n` bits. Bitmaps therefore + always start on a byte boundary, but the cursor may end mid-byte. +- **readInt(n):** read the next `n` bits (low bits first) as an unsigned integer, + with no alignment. `n == 0` reads nothing and yields 0. + +When decoding finishes, the number of bytes consumed — the final partial byte +counted as whole — must equal the `bitstream` length exactly. Any leftover or +missing byte is corruption. + +### Decoding algorithm + +``` +for off = 0; off < decoded_size; off += block_size: + blockLen = min(block_size, decoded_size - off) + decodeNode(level = 0, range = [0, numRanks), count = blockLen) + -> writes blockLen bytes to output[off ..] +``` + +`decodeNode(level, [firstRank, rankEnd), count)` produces `count` output bytes: + +**Leaf node** — `[firstRank, rankEnd)` is exactly one leaf, i.e. its flat depth +`d` satisfies `2^d == rankEnd - firstRank`: + +- `d == 0` (constant leaf): emit `count` copies of the leaf's single symbol. Read + nothing. +- `d > 0` (flat leaf): `alignedBitmap(count * d)`. For each output position + `i` in `[0, count)`, the `d` bits at `[i*d, (i+1)*d)` (low bit first) form an + index `idx` in `[0, 2^d)`; emit the leaf's `idx`-th symbol, which is + `rankToSymbol[firstRank + idx]`. + +**Internal node** — otherwise: + +1. If `level >= numLevels`, the range cannot split: corruption. +2. Compute `splitRank` from the tree: the boundary in `[firstRank, rankEnd)` + between codewords whose level-`L` bit is 0 (the lower ranks, the **left** + child) and 1 (the higher ranks, the **right** child). This comes entirely from + the tree (the weights), never from the bitstream. +3. From the tree, determine whether the left child `[firstRank, splitRank)` and + the right child `[splitRank, rankEnd)` are each a constant leaf. +4. `alignedBitmap(count)` — the **partition bitmap**. Bit `i` (low bit first) + selects, for output position `i`, the right child (1) or the left child (0). +5. If **not** both children are constant, read the right child's size: + `numOnes = readInt(nextPow2(count + 1))`, where `nextPow2(x)` is the number of + bits needed to represent values `0..x-1` (i.e. `ceil(log2(x))`). Reject if + `numOnes > count`. `numZeros = count - numOnes`. `numOnes` must equal the + number of set bits in the partition bitmap, otherwise the frame is corrupt. + If both children are constant, no integer is read; the partition bitmap alone + reconstructs the output. +6. Recurse, **left child first**: `decodeNode(level+1, [firstRank, splitRank), + numZeros)` then `decodeNode(level+1, [splitRank, rankEnd), numOnes)`. The left + subtree is fully decoded before the right (a pre-order traversal of the tree). +7. **Merge** into this node's output: for each position `i`, take the next byte + from the right child if bit `i` is set, else from the left child, consuming + each child's bytes in order. A constant child supplies its single symbol for + each of its positions; a non-constant child supplies its decoded bytes. + +Because a node emits its own bitmap and count before recursing, and recurses left +before right, the bitstream is exactly a pre-order serialization of the tree: +`[node bitmap][node count?] [left subtree...] [right subtree...]`, with each +bitmap re-aligned to a byte boundary and each count packed immediately after its +bitmap. + +### Outputs + +The decoder produces one serial stream of `decoded_size` byte elements. + +### Encoder guidance + +Encoders should emit the shortest `weights` stream that includes every non-zero +symbol weight; symbols above `weights.size() - 1` are implicitly absent. The block +size written into the header must match the one used to encode, and is omitted +(defaulting to `decoded_size`) when the input is a single block. + +For an empty input, encoders must emit `decoded_size == 0`, an empty `weights` +stream, and an empty `bitstream`. For a constant input, encoders should emit a +single non-zero weight of 1 for the repeated symbol and an empty `bitstream` (the +constant root emits no bits). diff --git a/src/openzl/common/wire_format.h b/src/openzl/common/wire_format.h index f41658034..0580e203e 100644 --- a/src/openzl/common/wire_format.h +++ b/src/openzl/common/wire_format.h @@ -318,6 +318,8 @@ typedef enum { ZL_StandardTransformID_sparse_num = 66, + ZL_StandardTransformID_pivco_huffman = 67, + ZL_StandardTransformID_end = 128 // last id, used to detect end of ID range (impacts // header encoding) give some room to be able to add new diff --git a/src/openzl/compress/private_nodes.h b/src/openzl/compress/private_nodes.h index 30764099f..5d2857a4d 100644 --- a/src/openzl/compress/private_nodes.h +++ b/src/openzl/compress/private_nodes.h @@ -96,6 +96,12 @@ extern "C" { #define ZL_NODE_SPLIT_BY_STRUCT (ZL_NodeID){ZL_PrivateStandardNodeID_split_by_struct} +// pivco_huffman +// Input : 1 serial stream +// Outputs: 2 streams -- the numeric zstd-style Huffman weights and the serial +// pivco-coded bitstream (see codecs/pivco_huffman/spec.md) +#define ZL_NODE_PIVCO_HUFFMAN ZL_MAKE_NODE_ID(ZL_PrivateStandardNodeID_pivco_huffman) + #define ZL_GRAPH_FIELD_LZ_LITERALS (ZL_GraphID){ ZL_PrivateStandardGraphID_field_lz_literals } #define ZL_GRAPH_FIELD_LZ_LITERALS_CHANNEL (ZL_GraphID){ ZL_PrivateStandardGraphID_field_lz_literals_channel } @@ -155,6 +161,7 @@ typedef enum { ZL_PrivateStandardNodeID_transpose_split8_deprecated, ZL_PrivateStandardNodeID_lz4, + ZL_PrivateStandardNodeID_pivco_huffman, ZL_PrivateStandardNodeID_end // last id, used to detect out-of-bound enum // values 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/src/openzl/shared/portability.h b/src/openzl/shared/portability.h index 36d66c3f8..60e5d082a 100644 --- a/src/openzl/shared/portability.h +++ b/src/openzl/shared/portability.h @@ -92,7 +92,7 @@ ZL_BEGIN_C_DECLS # define ZL_PREFETCH_L2(ptr) ((void)(ptr)) #endif -#if defined(__GNUC__) || defined(__ICCARM__) +#if defined(__GNUC__) || defined(__ICCARM__) || ZL_HAS_ATTRIBUTE(__target__) # define ZL_TARGET_ATTRIBUTE(target) __attribute__((__target__(target))) #else # define ZL_TARGET_ATTRIBUTE(target) @@ -104,6 +104,12 @@ ZL_BEGIN_C_DECLS # define ZL_HAS_BUILTIN(x) 0 #endif +#ifdef __has_attribute +# define ZL_HAS_ATTRIBUTE(x) __has_attribute(x) +#else +# define ZL_HAS_ATTRIBUTE(x) 0 +#endif + #ifdef __has_feature # define ZL_HAS_FEATURE(x) __has_feature(x) #else diff --git a/tests/registry/OpenZLComponents.h b/tests/registry/OpenZLComponents.h index 84676eadc..fafd5383f 100644 --- a/tests/registry/OpenZLComponents.h +++ b/tests/registry/OpenZLComponents.h @@ -90,6 +90,7 @@ enum class OpenZLComponentID { Lz, MuxLengths, SparseNum, + PivCoHuffman, // Must be last enum value NumComponents, }; @@ -166,6 +167,7 @@ std::unique_ptr makeCompressSmallLengthsComponent(); std::unique_ptr makeLzComponent(); std::unique_ptr makeMuxLengthsComponent(); std::unique_ptr makeSparseNumComponent(); +std::unique_ptr makePivCoHuffmanComponent(); } // namespace components @@ -303,6 +305,8 @@ inline std::unique_ptr makeOpenZLComponent( return components::makeMuxLengthsComponent(); case OpenZLComponentID::SparseNum: return components::makeSparseNumComponent(); + case OpenZLComponentID::PivCoHuffman: + return components::makePivCoHuffmanComponent(); case OpenZLComponentID::NumComponents: default: throw std::runtime_error("Invalid component"); diff --git a/tests/registry/components/PivCoHuffman.cpp b/tests/registry/components/PivCoHuffman.cpp new file mode 100644 index 000000000..557d248de --- /dev/null +++ b/tests/registry/components/PivCoHuffman.cpp @@ -0,0 +1,135 @@ +// Copyright (c) Meta Platforms, Inc. and affiliates. + +#include "openzl/compress/private_nodes.h" +#include "tests/datagen/structures/CompressibleStringProducer.h" +#include "tests/registry/OpenZLComponents.h" +#include "tests/registry/OpenZLInput.h" +#include "tests/utils.h" + +namespace openzl::tests::components { +namespace { + +class PivCoHuffmanComponent : public OpenZLComponent { + public: + std::string name() const override + { + return "PivCoHuffman"; + } + + int minFormatVersion() const override + { + return 27; + } + + std::vector predefinedNodes(Compressor&) const override + { + return { ZL_NODE_PIVCO_HUFFMAN }; + } + + std::vector> predefinedInputs() const override + { + std::vector> inputs; + inputs.push_back(SerialOpenZLInput::make("")); + inputs.push_back(SerialOpenZLInput::make("x")); + inputs.push_back(SerialOpenZLInput::make("xy")); + inputs.push_back(SerialOpenZLInput::make("abc")); + inputs.push_back(SerialOpenZLInput::make(std::string(1000, 'x'))); + inputs.push_back(SerialOpenZLInput::make(kLoremTestInput)); + inputs.push_back(SerialOpenZLInput::make(kAudioPCMS32LETestInput)); + + std::string allBytes; + allBytes.reserve(256); + for (int b = 0; b < 256; ++b) { + allBytes.push_back(static_cast(b)); + } + inputs.push_back(SerialOpenZLInput::make(std::move(allBytes))); + + // Larger than ZL_PIVCO_DEFAULT_BLOCK_SIZE (32 KiB), so the encoder + // spans multiple blocks and records the optional block size in the + // codec header. + std::string multiBlock; + multiBlock.reserve(40000); + for (size_t i = 0; i < 40000; ++i) { + multiBlock.push_back( + static_cast('a' + (i * 7 + i / 13) % 11)); + } + inputs.push_back(SerialOpenZLInput::make(std::move(multiBlock))); + + return inputs; + } + + std::vector> generateInputs( + datagen::DataGen& gen, + size_t num, + size_t maxInputSize, + const Compressor&, + GraphID) const override + { + std::vector> inputs; + inputs.reserve(num); + for (size_t i = 0; i < num; ++i) { + datagen::CompressibleStringProducer producer( + gen.getRandWrapper(), + gen.usize_range("input_size", 0, maxInputSize), + gen.u32_range("match_prob", 0, 100) / 100.0); + inputs.push_back(SerialOpenZLInput::make(producer("input"))); + } + return inputs; + } + + std::vector benchmarks( + Compressor& compressor, + datagen::DataGen& gen) const override + { + auto graph = buildTrivialGraph(compressor.get(), ZL_NODE_PIVCO_HUFFMAN); + + std::vector benchmarks; + + struct Params { + size_t inputSize; + size_t numInputs; + size_t compressibility; + }; + // Sizes tuned to provide approximately equal weight to each benchmark + // in compression speed. + constexpr std::array kParams = { + Params{ 100, 600, 10 }, Params{ 1000, 500, 10 }, + Params{ 10000, 150, 10 }, Params{ 100000, 20, 10 }, + Params{ 100000, 4, 20 }, Params{ 100000, 4, 40 }, + Params{ 100000, 4, 60 }, Params{ 100000, 4, 80 }, + Params{ 100000, 4, 90 }, + }; + + benchmarks.reserve(kParams.size()); + for (const auto& params : kParams) { + std::vector> inputs; + inputs.reserve(params.numInputs); + datagen::CompressibleStringProducer producer( + gen.getRandWrapper(), + params.inputSize, + params.compressibility / 100.0); + for (size_t i = 0; i < params.numInputs; ++i) { + inputs.push_back(SerialOpenZLInput::make(producer("input"))); + } + OpenZLComponent::Benchmark benchmark = { + .name = "InputSize:" + std::to_string(params.inputSize) + + "/Compressibility:" + + std::to_string(params.compressibility), + .graph = graph, + .inputs = std::move(inputs), + }; + benchmarks.push_back(std::move(benchmark)); + } + + return benchmarks; + } +}; + +} // namespace + +std::unique_ptr makePivCoHuffmanComponent() +{ + return std::make_unique(); +} + +} // namespace openzl::tests::components 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..97c1fbe65 --- /dev/null +++ b/tests/unittest/transforms/test_pivco_huffman.cpp @@ -0,0 +1,971 @@ +// 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 }, + { "avx512", &ZL_PivCoHuffmanEncode_avx512 }, + }; + 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 }, + { "avx512", &ZL_PivCoHuffmanDecode_avx512 }, + }; + 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)); +}