From 9fa6fddf93f5f61686b27851275a88a4a8285215 Mon Sep 17 00:00:00 2001 From: Sanya239 Date: Mon, 20 Jul 2026 02:00:23 +0300 Subject: [PATCH 01/14] Naive huffman implementation --- CMakeLists.txt | 21 +- include/pixie/huffman.h | 89 +++++ include/pixie/huffman/implementations.h | 12 + include/pixie/huffman/pivco_huffman.h | 343 ++++++++++++++++++++ src/benchmarks/pivco_huffman_benchmarks.cpp | 189 +++++++++++ src/tests/pivco_huffman_tests.cpp | 108 ++++++ 6 files changed, 761 insertions(+), 1 deletion(-) create mode 100644 include/pixie/huffman.h create mode 100644 include/pixie/huffman/implementations.h create mode 100644 include/pixie/huffman/pivco_huffman.h create mode 100644 src/benchmarks/pivco_huffman_benchmarks.cpp create mode 100644 src/tests/pivco_huffman_tests.cpp diff --git a/CMakeLists.txt b/CMakeLists.txt index 6101a98..6828f92 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -256,6 +256,15 @@ if (PIXIE_TESTS) PRIVATE ${sdsl_lite_SOURCE_DIR}/include) endif () + add_executable(pivco_huffman_tests + src/tests/pivco_huffman_tests.cpp) + target_include_directories(pivco_huffman_tests + PUBLIC include) + target_link_libraries(pivco_huffman_tests + gtest + gtest_main + ${PIXIE_DIAGNOSTICS_LIBS}) + set(PIXIE_TEST_TARGETS bit_algorithms_unittests rank_select_unittests @@ -267,7 +276,8 @@ if (PIXIE_TESTS) storage_tests excess_positions_tests excess_record_lows_tests - rmq_tests) + rmq_tests + pivco_huffman_tests) foreach (test_target IN LISTS PIXIE_TEST_TARGETS) gtest_discover_tests(${test_target} DISCOVERY_MODE PRE_TEST @@ -399,6 +409,15 @@ if (PIXIE_BENCHMARKS) benchmark benchmark_main ${PIXIE_DIAGNOSTICS_LIBS}) + + add_executable(pivco_huffman_benchmarks + src/benchmarks/pivco_huffman_benchmarks.cpp) + target_include_directories(pivco_huffman_benchmarks + PUBLIC include) + target_link_libraries(pivco_huffman_benchmarks + benchmark + benchmark_main + ${PIXIE_DIAGNOSTICS_LIBS}) endif () # --------------------------------------------------------------------------- diff --git a/include/pixie/huffman.h b/include/pixie/huffman.h new file mode 100644 index 0000000..68acab5 --- /dev/null +++ b/include/pixie/huffman.h @@ -0,0 +1,89 @@ +#pragma once + +/** + * @file huffman.h + * @brief Common CRTP contract for Huffman entropy codecs. + * + * PivCo-Huffman ("Pivot-Coded Huffman") reuses the wavelet-tree "tree of + * bitmaps" layout to turn sequential, bit-by-bit Huffman-tree traversals into + * vectorizable operations. A codec encodes a byte sequence into a compressed + * stream by building a Huffman-shaped tree of per-node bitmaps, and decodes it + * back by traversing that tree. Concrete implementations live under + * ``. + * + * @see Marcin Zukowski, "PivCo-Huffman", v1.0 (2026). + */ + +#include +#include +#include +#include + +namespace pixie { + +/** + * @brief CRTP facade for Huffman entropy codecs. + * + * The contract is the source of truth for observable codec semantics. Each + * concrete implementation inherits `HuffmanBase` and supplies the + * required `*_impl()` extension points. There is no virtual dispatch: the + * facade delegates statically through CRTP, mirroring the other Pixie families. + * + * Range and ownership conventions: + * - Symbol sequences are zero-based byte streams (`symbol_type`). + * - Compressed streams are byte-oriented views (`std::byte`). + * - The caller keeps any non-owning view alive for the codec lifetime. + * + * @tparam Impl Concrete codec type implementing the `*_impl()` contract. + * + * @see `` for the available concrete + * implementations. + */ +template +class HuffmanBase { + public: + /** @brief Symbol type handled by the codec: one byte per symbol. */ + using symbol_type = std::uint8_t; + + /** + * @brief Number of symbols in the original (uncompressed) stream. + * @return Logical uncompressed symbol count. + */ + std::size_t uncompressed_size() const { + return impl().uncompressed_size_impl(); + } + + /** + * @brief Number of bytes in the compressed representation. + * @return Compressed stream size in bytes. + */ + std::size_t compressed_size() const { return impl().compressed_size_impl(); } + + /** + * @brief Check whether the codec holds no data. + * @return `true` when `uncompressed_size() == 0`. + */ + bool empty() const { return uncompressed_size() == 0; } + + /** + * @brief Read-only view of the compressed byte stream. + * @return Span over the serialized representation. + * + * @note The returned view is invalidated by codec destruction. + */ + std::span compressed_data() const { + return impl().compressed_data_impl(); + } + + /** + * @brief Reconstruct the original symbol sequence. + * @return Decoded symbols of length `uncompressed_size()`. + */ + std::vector decode() const { return impl().decode_impl(); } + + private: + /** @brief Return this facade as its concrete CRTP implementation. */ + const Impl& impl() const { return static_cast(*this); } +}; + +} // namespace pixie diff --git a/include/pixie/huffman/implementations.h b/include/pixie/huffman/implementations.h new file mode 100644 index 0000000..e5455cf --- /dev/null +++ b/include/pixie/huffman/implementations.h @@ -0,0 +1,12 @@ +#pragma once + +/** + * @file implementations.h + * @brief All Huffman codec implementations provided by Pixie. + * + * - `PivCoHuffman`: simple scalar reference codec storing a Huffman-shaped + * tree of per-node routing bitmaps as packed 64-bit words. + */ + +#include +#include diff --git a/include/pixie/huffman/pivco_huffman.h b/include/pixie/huffman/pivco_huffman.h new file mode 100644 index 0000000..c675be7 --- /dev/null +++ b/include/pixie/huffman/pivco_huffman.h @@ -0,0 +1,343 @@ +#pragma once + +/** + * @file pivco_huffman.h + * @brief Simple scalar PivCo-Huffman codec. + * + * Stores a Huffman-shaped tree of per-node routing bitmaps (the "tree of + * bitmaps" layout shared with wavelet trees). Encoding walks each input symbol + * from the root, appending one direction bit to every internal node on its + * path. Decoding walks each output position from the root down to a leaf, + * consuming one bit per internal node. + * + * This is a deliberately simple, unoptimized reference implementation: node + * bitmaps are stored as packed `std::uint64_t` words, traversal is scalar, + * and there are no flat-subtree, non-canonical, SIMD, or selective-ANS + * optimizations. It is intended as a correct baseline for the PivCo-Huffman + * layout and for future optimized variants. + * + * Range and ownership conventions follow `HuffmanBase`: symbols are bytes, + * the compressed stream is a byte view, and the codec owns its serialized form. + */ + +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace pixie { + +/** + * @brief Simple scalar PivCo-Huffman codec. + * @details Implements `HuffmanBase` by building a Huffman tree over the + * byte alphabet, storing one routing bitmap per internal node as + * packed 64-bit words, and decoding by top-down per-position + * traversal. The serialized form is a header followed by the tree + * structure and packed node bitmaps. + */ +class PivCoHuffman : public HuffmanBase { + public: + /** @brief Symbol type handled by the codec: one byte per symbol. */ + using symbol_type = std::uint8_t; + + /** + * @brief Build a codec by encoding @p input. + * @param input Symbol sequence to compress. + */ + explicit PivCoHuffman(std::span input) { + build(input); + serialize(); + } + + /** + * @brief Load a codec from a previously serialized compressed stream. + * @param compressed Compressed byte view produced by another instance. + */ + explicit PivCoHuffman(std::span compressed) { + deserialize(compressed); + } + + // --- CRTP extension points ---------------------------------------------- + + /** @brief Size of the original symbol stream. */ + std::size_t uncompressed_size_impl() const { return uncompressed_size_; } + + /** @brief Size of the compressed byte stream. */ + std::size_t compressed_size_impl() const { return compressed_.size(); } + + /** @brief Read-only view of the compressed byte stream. */ + std::span compressed_data_impl() const { + return compressed_; + } + + /** @brief Reconstruct the original symbol sequence. */ + std::vector decode_impl() const { return decode_from_tree(); } + + private: + /// @brief Sentinel node index meaning "no node". + static constexpr std::size_t kNpos = static_cast(-1); + + /// @brief Byte-alphabet size (256 symbols). + static constexpr std::size_t kAlphabet = 256; + + /// @brief Bits per word used by the packed bitmap. + static constexpr std::size_t kWordBits = 64; + + /** + * @brief Packed routing bitmap stored as 64-bit words. + * @details Bit `i` lives in word `i / 64` at position `i % 64`. This avoids + * the slow bit-proxy access of `std::vector` and allows + * bulk serialization via direct word `memcpy`. + */ + struct Bitmap { + std::vector words; + std::size_t count = 0; + + /** @brief Append one bit to the end of the bitmap. */ + void push_back(bool bit) { + if (count % kWordBits == 0) { + words.push_back(0); + } + if (bit) { + words.back() |= (1ull << (count % kWordBits)); + } + ++count; + } + + /** @brief Read the bit at position @p i. */ + bool get(std::size_t i) const { + return ((words[i / kWordBits] >> (i % kWordBits)) & 1ull) != 0; + } + + /** @brief Number of stored bits. */ + std::size_t size() const { return count; } + + /** @brief Number of 64-bit words backing the bitmap. */ + std::size_t word_count() const { + return (count + kWordBits - 1) / kWordBits; + } + }; + + /** + * @brief A node in the Huffman tree of bitmaps. + * @details Internal nodes carry a routing bitmap with one bit per symbol + * that passes through the node (0 = left, 1 = right), in input + * order. Leaves carry only their symbol. + */ + struct Node { + std::size_t left = kNpos; + std::size_t right = kNpos; + std::uint8_t symbol = 0; + bool is_leaf = false; + Bitmap bits; + }; + + std::size_t uncompressed_size_ = 0; + std::size_t root_ = kNpos; + std::vector nodes_; + std::vector compressed_; + + // --- construction -------------------------------------------------------- + + /** @brief Build the Huffman tree and per-node bitmaps from @p input. */ + void build(std::span input) { + uncompressed_size_ = input.size(); + if (uncompressed_size_ == 0) { + root_ = kNpos; + return; + } + + std::array freq{}; + for (auto s : input) { + freq[s]++; + } + + struct HeapItem { + std::size_t weight; + std::size_t idx; + }; + auto cmp = [](const HeapItem& a, const HeapItem& b) { + return a.weight > b.weight; + }; + std::priority_queue, decltype(cmp)> heap( + cmp); + + for (std::size_t s = 0; s < kAlphabet; s++) { + if (freq[s] > 0) { + Node n; + n.is_leaf = true; + n.symbol = static_cast(s); + nodes_.push_back(std::move(n)); + heap.push({freq[s], nodes_.size() - 1}); + } + } + + while (heap.size() > 1) { + HeapItem a = heap.top(); + heap.pop(); + HeapItem b = heap.top(); + heap.pop(); + Node n; + n.is_leaf = false; + n.left = a.idx; + n.right = b.idx; + nodes_.push_back(std::move(n)); + heap.push({a.weight + b.weight, nodes_.size() - 1}); + } + root_ = heap.top().idx; + + // Record, for each symbol, the path of (node, direction-bit) from root. + std::array>, kAlphabet> paths; + std::vector> path; + assign_paths(root_, path, paths); + + // Append one direction bit per visited node, in input order. + for (auto s : input) { + for (const auto& [idx, bit] : paths[s]) { + nodes_[idx].bits.push_back(bit); + } + } + } + + /** @brief Depth-first assignment of root-to-leaf paths. */ + void assign_paths( + std::size_t idx, + std::vector>& path, + std::array>, kAlphabet>& paths) { + if (nodes_[idx].is_leaf) { + paths[nodes_[idx].symbol] = path; + return; + } + path.emplace_back(idx, false); + assign_paths(nodes_[idx].left, path, paths); + path.pop_back(); + path.emplace_back(idx, true); + assign_paths(nodes_[idx].right, path, paths); + path.pop_back(); + } + + // --- decode -------------------------------------------------------------- + + /** @brief Reconstruct the original sequence by top-down traversal. */ + std::vector decode_from_tree() const { + std::vector out(uncompressed_size_); + if (root_ == kNpos) { + return out; + } + if (nodes_[root_].is_leaf) { + const std::uint8_t sym = nodes_[root_].symbol; + for (auto& s : out) { + s = sym; + } + return out; + } + std::vector cursor(nodes_.size(), 0); + for (std::size_t i = 0; i < uncompressed_size_; i++) { + std::size_t idx = root_; + while (!nodes_[idx].is_leaf) { + const bool bit = nodes_[idx].bits.get(cursor[idx]++); + idx = bit ? nodes_[idx].right : nodes_[idx].left; + } + out[i] = nodes_[idx].symbol; + } + return out; + } + + // --- serialization ------------------------------------------------------- + + /** @brief Write the in-memory tree into the serialized byte buffer. */ + void serialize() { + compressed_.clear(); + write(uncompressed_size_); + if (uncompressed_size_ == 0) { + return; + } + write(root_); + write(nodes_.size()); + for (const auto& n : nodes_) { + write(static_cast(n.is_leaf ? 1 : 0)); + write(n.symbol); + write(n.left); + write(n.right); + write(n.bits.count); + write_words(n.bits.words); + } + } + + /** @brief Rebuild the in-memory tree from a serialized byte buffer. */ + void deserialize(std::span data) { + compressed_.assign(data.begin(), data.end()); + nodes_.clear(); + if (data.empty()) { + uncompressed_size_ = 0; + root_ = kNpos; + return; + } + std::size_t pos = 0; + uncompressed_size_ = read(data, pos); + if (uncompressed_size_ == 0) { + root_ = kNpos; + return; + } + root_ = read(data, pos); + const std::size_t count = read(data, pos); + nodes_.resize(count); + for (std::size_t i = 0; i < count; i++) { + const std::uint8_t flags = read(data, pos); + nodes_[i].is_leaf = (flags & 1) != 0; + nodes_[i].symbol = read(data, pos); + nodes_[i].left = read(data, pos); + nodes_[i].right = read(data, pos); + nodes_[i].bits.count = read(data, pos); + nodes_[i].bits.words = read_words(data, pos, nodes_[i].bits.word_count()); + } + } + + /** @brief Append a little-endian, fixed-width value to the byte buffer. */ + template + void write(const T& value) { + const std::size_t old = compressed_.size(); + compressed_.resize(old + sizeof(T)); + std::memcpy(compressed_.data() + old, &value, sizeof(T)); + } + + /** @brief Append a packed word array directly to the byte buffer. */ + void write_words(const std::vector& words) { + const std::size_t bytes = words.size() * sizeof(std::uint64_t); + const std::size_t old = compressed_.size(); + compressed_.resize(old + bytes); + if (bytes > 0) { + std::memcpy(compressed_.data() + old, words.data(), bytes); + } + } + + /** @brief Read a fixed-width value from @p data at @p pos. */ + template + static T read(std::span data, std::size_t& pos) { + T value; + std::memcpy(&value, data.data() + pos, sizeof(T)); + pos += sizeof(T); + return value; + } + + /** @brief Read a packed word array of @p word_count words from @p data. */ + static std::vector read_words(std::span data, + std::size_t& pos, + std::size_t word_count) { + std::vector words(word_count); + const std::size_t bytes = word_count * sizeof(std::uint64_t); + if (bytes > 0) { + std::memcpy(words.data(), data.data() + pos, bytes); + pos += bytes; + } + return words; + } +}; + +} // namespace pixie diff --git a/src/benchmarks/pivco_huffman_benchmarks.cpp b/src/benchmarks/pivco_huffman_benchmarks.cpp new file mode 100644 index 0000000..5812bf4 --- /dev/null +++ b/src/benchmarks/pivco_huffman_benchmarks.cpp @@ -0,0 +1,189 @@ +#include +#include + +#include +#include +#include +#include +#include + +using pixie::PivCoHuffman; + +// --------------------------------------------------------------------------- +// Configuration via environment variables (follows the Pixie convention used +// by EXCESS_POS_CASES / RECORD_LOWS_CASES etc.): +// PIVCO_BENCH_SIZE - input size in bytes (default 1<<16 = 64 KiB) +// PIVCO_BENCH_SEED - random seed (default 42) +// +// To run a subset of datasets, use Google Benchmark's built-in filter flag, +// e.g. `--benchmark_filter=".*Skewed.*|.*Text.*"`. No code change needed. +// --------------------------------------------------------------------------- + +namespace { + +/// @brief Input size in bytes, overridable via `PIVCO_BENCH_SIZE`. +std::size_t bench_size() { + if (const char* env = std::getenv("PIVCO_BENCH_SIZE")) { + try { + const unsigned long long v = std::stoull(env); + if (v != 0) { + return static_cast(v); + } + } catch (...) { + } + } + return 1ull << 16; +} + +/// @brief RNG seed, overridable via `PIVCO_BENCH_SEED`. +std::uint64_t bench_seed() { + if (const char* env = std::getenv("PIVCO_BENCH_SEED")) { + try { + return std::stoull(env); + } catch (...) { + } + } + return 42; +} + +/// @brief Signature of a dataset generator: produces @p n bytes from @p rng. +using DatasetMaker = std::vector (*)(std::size_t, + std::mt19937_64&); + +// --- built-in datasets ----------------------------------------------------- + +/// @brief Uniformly distributed bytes over the full 256-symbol alphabet. +std::vector make_uniform256(std::size_t n, std::mt19937_64& rng) { + std::vector data(n); + for (auto& b : data) { + b = static_cast(rng()); + } + return data; +} + +/// @brief Binary alphabet {0, 1}: shortest possible codes (depth 1). +std::vector make_binary(std::size_t n, std::mt19937_64& rng) { + std::vector data(n); + for (auto& b : data) { + b = static_cast(rng() & 1u); + } + return data; +} + +/// @brief Four-symbol alphabet {0,1,2,3}: code depth 2. +std::vector make_low_entropy_4(std::size_t n, + std::mt19937_64& rng) { + std::vector data(n); + for (auto& b : data) { + b = static_cast(rng() & 3u); + } + return data; +} + +/// @brief English-text-like distribution: ~17% space, ~80% lowercase letters, +/// ~2% digits, ~1% punctuation. +std::vector make_text_like(std::size_t n, std::mt19937_64& rng) { + std::vector data(n); + std::uniform_int_distribution pick(0, 99); + std::uniform_int_distribution letter('a', 'z'); + std::uniform_int_distribution digit('0', '9'); + std::uniform_int_distribution punct('!', '/'); + for (auto& b : data) { + const int r = pick(rng); + if (r < 17) { + b = ' '; + } else if (r < 97) { + b = static_cast(letter(rng)); + } else if (r < 99) { + b = static_cast(digit(rng)); + } else { + b = static_cast(punct(rng)); + } + } + return data; +} + +/// @brief Highly skewed: 99% one symbol, 1% random over 256. Exercises a deep, +/// lopsided Huffman tree. +std::vector make_skewed99(std::size_t n, std::mt19937_64& rng) { + std::vector data(n, 0); + std::uniform_int_distribution coin(0, 99); + for (auto& b : data) { + if (coin(rng) < 1) { + b = static_cast(rng()); + } + } + return data; +} + +/// @brief Degenerate single-symbol input (root is a leaf): best-case path. +std::vector make_single_symbol(std::size_t n, + std::mt19937_64& rng) { + (void)rng; + return std::vector(n, 123); +} + +/// @brief Compression efficiency in bits per input byte (8 = no compression). +void report_ratio(benchmark::State& state, + std::size_t compressed_bytes, + std::size_t input_bytes) { + state.counters["bpb"] = static_cast(compressed_bytes) * 8.0 / + static_cast(input_bytes); +} + +// --- benchmarks ------------------------------------------------------------ + +/// @brief Measure full encoding: build Huffman tree, fill node bitmaps, and +/// serialize the compressed stream. +void BM_Encode(benchmark::State& state, DatasetMaker make) { + const std::size_t size = bench_size(); + std::mt19937_64 rng(bench_seed()); + const std::vector data = make(size, rng); + + // Encode once outside timing to report compression ratio. + const PivCoHuffman probe(data); + report_ratio(state, probe.compressed_size(), size); + + for (auto _ : state) { + PivCoHuffman codec(data); + benchmark::DoNotOptimize(codec.compressed_data().data()); + benchmark::ClobberMemory(); + } + state.SetBytesProcessed(static_cast(state.iterations() * size)); +} + +/// @brief Measure decoding: reconstruct the original byte stream from the +/// in-memory tree. Encoding is done once, outside the timed loop. +void BM_Decode(benchmark::State& state, DatasetMaker make) { + const std::size_t size = bench_size(); + std::mt19937_64 rng(bench_seed()); + const std::vector data = make(size, rng); + + const PivCoHuffman codec(data); + report_ratio(state, codec.compressed_size(), size); + + for (auto _ : state) { + std::vector out = codec.decode(); + benchmark::DoNotOptimize(out.data()); + benchmark::ClobberMemory(); + } + state.SetBytesProcessed(static_cast(state.iterations() * size)); +} + +} // namespace + +// Register encode + decode for every built-in dataset. Each registration is +// named `BM_/`, so `--benchmark_filter` selects subsets freely. +BENCHMARK_CAPTURE(BM_Encode, Uniform256, make_uniform256); +BENCHMARK_CAPTURE(BM_Encode, Binary, make_binary); +BENCHMARK_CAPTURE(BM_Encode, LowEntropy4, make_low_entropy_4); +BENCHMARK_CAPTURE(BM_Encode, TextLike, make_text_like); +BENCHMARK_CAPTURE(BM_Encode, Skewed99, make_skewed99); +BENCHMARK_CAPTURE(BM_Encode, SingleSymbol, make_single_symbol); + +BENCHMARK_CAPTURE(BM_Decode, Uniform256, make_uniform256); +BENCHMARK_CAPTURE(BM_Decode, Binary, make_binary); +BENCHMARK_CAPTURE(BM_Decode, LowEntropy4, make_low_entropy_4); +BENCHMARK_CAPTURE(BM_Decode, TextLike, make_text_like); +BENCHMARK_CAPTURE(BM_Decode, Skewed99, make_skewed99); +BENCHMARK_CAPTURE(BM_Decode, SingleSymbol, make_single_symbol); diff --git a/src/tests/pivco_huffman_tests.cpp b/src/tests/pivco_huffman_tests.cpp new file mode 100644 index 0000000..fcc656b --- /dev/null +++ b/src/tests/pivco_huffman_tests.cpp @@ -0,0 +1,108 @@ +#include +#include + +#include +#include +#include +#include +#include + +using pixie::PivCoHuffman; + +static_assert( + std::is_base_of_v, PivCoHuffman>); +static_assert(std::is_same_v); + +namespace { +// Round-trip through the in-memory tree of the same codec instance. +std::vector decode_same(const PivCoHuffman& codec) { + return codec.decode(); +} + +// Round-trip through the serialized byte stream: simulate writing the +// compressed form to disk and loading it into a fresh codec instance. +std::vector decode_via_bytes(const PivCoHuffman& codec) { + const auto bytes = codec.compressed_data(); + PivCoHuffman copy(bytes); + return copy.decode(); +} +} // namespace + +TEST(PivCoHuffmanSmoke, EmptyInputRoundTrips) { + const std::vector data; + const PivCoHuffman codec(data); + EXPECT_TRUE(codec.empty()); + EXPECT_EQ(codec.uncompressed_size(), 0u); + EXPECT_EQ(decode_same(codec), data); + EXPECT_EQ(decode_via_bytes(codec), data); +} + +TEST(PivCoHuffmanSmoke, SingleSymbolRoundTrips) { + const std::vector data(1000, 42); + const PivCoHuffman codec(data); + EXPECT_EQ(codec.uncompressed_size(), data.size()); + EXPECT_EQ(decode_same(codec), data); + EXPECT_EQ(decode_via_bytes(codec), data); +} + +TEST(PivCoHuffmanSmoke, TwoSymbolRoundTrips) { + std::vector data; + for (std::size_t i = 0; i < 500; i++) { + data.push_back(static_cast(i % 2)); + } + const PivCoHuffman codec(data); + EXPECT_EQ(decode_same(codec), data); + EXPECT_EQ(decode_via_bytes(codec), data); +} + +TEST(PivCoHuffmanSmoke, KnownInputRoundTrips) { + const std::vector data{3, 2, 0, 3, 1, 1, 2, 0, 3, 1}; + const PivCoHuffman codec(data); + EXPECT_EQ(decode_same(codec), data); + EXPECT_EQ(decode_via_bytes(codec), data); +} + +TEST(PivCoHuffmanSmoke, CompressedDataIsNonEmpty) { + const std::vector data{3, 2, 0, 3, 1, 1, 2, 0, 3, 1}; + const PivCoHuffman codec(data); + EXPECT_GT(codec.compressed_size(), 0u); + EXPECT_EQ(codec.compressed_data().size(), codec.compressed_size()); +} + +TEST(PivCoHuffmanSmoke, FullAlphabetRoundTrips) { + std::vector data(256); + for (std::size_t i = 0; i < 256; i++) { + data[i] = static_cast(i); + } + const PivCoHuffman codec(data); + EXPECT_EQ(decode_same(codec), data); + EXPECT_EQ(decode_via_bytes(codec), data); +} + +TEST(PivCoHuffmanSmoke, RandomUniformRoundTrips) { + std::mt19937 rng(239); + std::uniform_int_distribution dist(0, 255); + std::vector data(10000); + for (auto& b : data) { + b = static_cast(dist(rng)); + } + const PivCoHuffman codec(data); + EXPECT_EQ(decode_same(codec), data); + EXPECT_EQ(decode_via_bytes(codec), data); +} + +TEST(PivCoHuffmanSmoke, SkewedInputRoundTrips) { + // Most symbols are one value, a few are outliers: exercises a deep, lopsided + // Huffman tree without any single-symbol shortcut. + std::mt19937 rng(7); + std::uniform_int_distribution coin(0, 99); + std::vector data(10000, 7); + for (std::size_t i = 0; i < data.size(); i++) { + if (coin(rng) < 5) { + data[i] = static_cast(rng() % 256); + } + } + const PivCoHuffman codec(data); + EXPECT_EQ(decode_same(codec), data); + EXPECT_EQ(decode_via_bytes(codec), data); +} From 9f80561440c38b3a089198db6ad9e97ffec7d9d2 Mon Sep 17 00:00:00 2001 From: Sanya239 Date: Tue, 21 Jul 2026 18:20:00 +0300 Subject: [PATCH 02/14] Bottom-up implementation with naive primitives --- include/pixie/huffman/pivco_huffman.h | 159 ++++++++++++++------ src/benchmarks/pivco_huffman_benchmarks.cpp | 92 +++++------ 2 files changed, 162 insertions(+), 89 deletions(-) diff --git a/include/pixie/huffman/pivco_huffman.h b/include/pixie/huffman/pivco_huffman.h index c675be7..48f40be 100644 --- a/include/pixie/huffman/pivco_huffman.h +++ b/include/pixie/huffman/pivco_huffman.h @@ -7,8 +7,8 @@ * Stores a Huffman-shaped tree of per-node routing bitmaps (the "tree of * bitmaps" layout shared with wavelet trees). Encoding walks each input symbol * from the root, appending one direction bit to every internal node on its - * path. Decoding walks each output position from the root down to a leaf, - * consuming one bit per internal node. + * path. Decoding merges child symbol streams bottom-up from leaves to root, + * using the node bitmap as a selector. * * This is a deliberately simple, unoptimized reference implementation: node * bitmaps are stored as packed `std::uint64_t` words, traversal is scalar, @@ -23,6 +23,7 @@ #include #include +#include #include #include #include @@ -37,9 +38,9 @@ namespace pixie { * @brief Simple scalar PivCo-Huffman codec. * @details Implements `HuffmanBase` by building a Huffman tree over the * byte alphabet, storing one routing bitmap per internal node as - * packed 64-bit words, and decoding by top-down per-position - * traversal. The serialized form is a header followed by the tree - * structure and packed node bitmaps. + * packed 64-bit words, and decoding by bottom-up merging of child + * symbol streams. The serialized form is a header followed by the + * tree structure and packed node bitmaps. */ class PivCoHuffman : public HuffmanBase { public: @@ -99,29 +100,19 @@ class PivCoHuffman : public HuffmanBase { std::vector words; std::size_t count = 0; - /** @brief Append one bit to the end of the bitmap. */ - void push_back(bool bit) { - if (count % kWordBits == 0) { - words.push_back(0); - } - if (bit) { - words.back() |= (1ull << (count % kWordBits)); - } - ++count; - } - - /** @brief Read the bit at position @p i. */ - bool get(std::size_t i) const { - return ((words[i / kWordBits] >> (i % kWordBits)) & 1ull) != 0; - } - - /** @brief Number of stored bits. */ - std::size_t size() const { return count; } - /** @brief Number of 64-bit words backing the bitmap. */ std::size_t word_count() const { return (count + kWordBits - 1) / kWordBits; } + + /** @brief Number of set (1) bits in the bitmap. */ + std::size_t popcount() const { + std::size_t total = 0; + for (std::uint64_t w : words) { + total += std::popcount(w); + } + return total; + } }; /** @@ -168,6 +159,10 @@ class PivCoHuffman : public HuffmanBase { std::priority_queue, decltype(cmp)> heap( cmp); + // At most 2*kAlphabet-1 nodes (leaves + internal). Pre-allocate to avoid + // reallocation during tree construction. + nodes_.reserve(2 * kAlphabet - 1); + for (std::size_t s = 0; s < kAlphabet; s++) { if (freq[s] > 0) { Node n; @@ -187,8 +182,11 @@ class PivCoHuffman : public HuffmanBase { n.is_leaf = false; n.left = a.idx; n.right = b.idx; + std::size_t w = a.weight + b.weight; + n.bits.count = w; + n.bits.words.assign((w + kWordBits - 1) / kWordBits, 0); nodes_.push_back(std::move(n)); - heap.push({a.weight + b.weight, nodes_.size() - 1}); + heap.push({w, nodes_.size() - 1}); } root_ = heap.top().idx; @@ -197,12 +195,9 @@ class PivCoHuffman : public HuffmanBase { std::vector> path; assign_paths(root_, path, paths); - // Append one direction bit per visited node, in input order. - for (auto s : input) { - for (const auto& [idx, bit] : paths[s]) { - nodes_[idx].bits.push_back(bit); - } - } + // Top-down partitioning: recursively split the symbol stream and fill + // pre-allocated bitmaps. Leaves are skipped (zero overhead). + encode_node(root_, 0, input, paths); } /** @brief Depth-first assignment of root-to-leaf paths. */ @@ -222,29 +217,99 @@ class PivCoHuffman : public HuffmanBase { path.pop_back(); } + /** @brief Top-down encoding: recursively partition symbols and fill bitmaps. + * @details At each internal node, fill the pre-allocated bitmap with the + * direction bit of each symbol's code at this depth, and split the + * symbol stream into left (0) and right (1) child vectors. Leaves + * return immediately — zero overhead. + * @param idx Current node index. + * @param depth Current tree depth (indexes into each symbol's code path). + * @param symbols Symbol stream arriving at this node, in order. + * @param paths Per-symbol code paths from `assign_paths`. */ + void encode_node(std::size_t idx, + std::size_t depth, + std::span symbols, + const std::array>, + kAlphabet>& paths) { + if (nodes_[idx].is_leaf) { + return; + } + + auto& n = nodes_[idx]; + std::vector left_symbols, right_symbols; + std::size_t left_size = nodes_[n.left].bits.count; + std::size_t right_size = nodes_[n.right].bits.count; + left_symbols.reserve(left_size); + right_symbols.reserve(right_size); + + // Fill bitmap via incremental word/bit tracking — avoids per-symbol + // division and modulo. + std::size_t word_idx = 0; + std::size_t bit_pos = 0; + for (std::size_t i = 0; i < symbols.size(); i++) { + bool bit = paths[symbols[i]][depth].second; + if (bit) { + n.bits.words[word_idx] |= (1ull << bit_pos); + right_symbols.push_back(symbols[i]); + } else { + left_symbols.push_back(symbols[i]); + } + if (++bit_pos == kWordBits) { + bit_pos = 0; + ++word_idx; + } + } + + encode_node(n.left, depth + 1, left_symbols, paths); + encode_node(n.right, depth + 1, right_symbols, paths); + } + // --- decode -------------------------------------------------------------- - /** @brief Reconstruct the original sequence by top-down traversal. */ + /** @brief Reconstruct the original sequence by bottom-up merging. */ std::vector decode_from_tree() const { - std::vector out(uncompressed_size_); if (root_ == kNpos) { - return out; + return std::vector(); } - if (nodes_[root_].is_leaf) { - const std::uint8_t sym = nodes_[root_].symbol; - for (auto& s : out) { - s = sym; - } - return out; + return decode_node(root_, uncompressed_size_); + } + + /** @brief Recursively decode a node's output stream bottom-up. + * @param weight Number of symbols this node must produce. For internal + * nodes this equals `bits.count`; for leaves it is passed down from + * the parent (no bitmap to read it from). */ + std::vector decode_node(std::size_t idx, + std::size_t weight) const { + const auto& n = nodes_[idx]; + if (n.is_leaf) { + return std::vector(weight, n.symbol); } - std::vector cursor(nodes_.size(), 0); - for (std::size_t i = 0; i < uncompressed_size_; i++) { - std::size_t idx = root_; - while (!nodes_[idx].is_leaf) { - const bool bit = nodes_[idx].bits.get(cursor[idx]++); - idx = bit ? nodes_[idx].right : nodes_[idx].left; + + std::size_t right_weight = n.bits.popcount(); + std::size_t left_weight = n.bits.count - right_weight; + std::vector left_symbols = decode_node(n.left, left_weight); + std::vector right_symbols = decode_node(n.right, right_weight); + + std::vector out(n.bits.count); + std::size_t c_left = 0; + std::size_t c_right = 0; + // Word-wise merge: read each 64-bit word once and shift through its bits, + // avoiding per-bit division/modulo. + for (std::size_t w = 0; w < n.bits.words.size(); w++) { + std::uint64_t word = n.bits.words[w]; + std::size_t base = w * kWordBits; + std::size_t limit = base + kWordBits; + if (limit > n.bits.count) { + limit = n.bits.count; + } + for (std::size_t i = base; i < limit; i++) { + if (word & 1ull) { + out[i] = right_symbols[c_right++]; + } else { + out[i] = left_symbols[c_left++]; + } + word >>= 1; } - out[i] = nodes_[idx].symbol; } return out; } diff --git a/src/benchmarks/pivco_huffman_benchmarks.cpp b/src/benchmarks/pivco_huffman_benchmarks.cpp index 5812bf4..cebb0ed 100644 --- a/src/benchmarks/pivco_huffman_benchmarks.cpp +++ b/src/benchmarks/pivco_huffman_benchmarks.cpp @@ -4,37 +4,24 @@ #include #include #include -#include #include using pixie::PivCoHuffman; // --------------------------------------------------------------------------- -// Configuration via environment variables (follows the Pixie convention used -// by EXCESS_POS_CASES / RECORD_LOWS_CASES etc.): -// PIVCO_BENCH_SIZE - input size in bytes (default 1<<16 = 64 KiB) +// Configuration: // PIVCO_BENCH_SEED - random seed (default 42) // -// To run a subset of datasets, use Google Benchmark's built-in filter flag, -// e.g. `--benchmark_filter=".*Skewed.*|.*Text.*"`. No code change needed. +// Input sizes are swept by Google Benchmark's range mechanism, not an env +// var: 1 KiB -> 1 MiB -> 1 GiB (multiplier 1024). This exercises the codec +// from cache-resident to gigabyte scale in a single run. +// +// To run a subset of datasets or sizes, use Google Benchmark's filter flag, +// e.g. `--benchmark_filter="BM_Encode/Uniform256"`. // --------------------------------------------------------------------------- namespace { -/// @brief Input size in bytes, overridable via `PIVCO_BENCH_SIZE`. -std::size_t bench_size() { - if (const char* env = std::getenv("PIVCO_BENCH_SIZE")) { - try { - const unsigned long long v = std::stoull(env); - if (v != 0) { - return static_cast(v); - } - } catch (...) { - } - } - return 1ull << 16; -} - /// @brief RNG seed, overridable via `PIVCO_BENCH_SEED`. std::uint64_t bench_seed() { if (const char* env = std::getenv("PIVCO_BENCH_SEED")) { @@ -131,31 +118,42 @@ void report_ratio(benchmark::State& state, static_cast(input_bytes); } +/// @brief Common range of input sizes: 1 KiB, 1 MiB, 1 GiB. +/// @details Applied to every registered benchmark. The 1 GiB tier exercises +/// gigabyte-scale inputs; at that size Google Benchmark runs a single +/// iteration per case. +constexpr std::int64_t kSizeLo = 1ull << 10; // 1 KiB +constexpr std::int64_t kSizeHi = 1ull << 30; // 1 GiB +constexpr std::int64_t kSizeMult = 1024; + // --- benchmarks ------------------------------------------------------------ /// @brief Measure full encoding: build Huffman tree, fill node bitmaps, and -/// serialize the compressed stream. +/// serialize the compressed stream. The compressed size is read from +/// the first timed iteration (no separate probe build), which keeps +/// gigabyte-scale setup cost to a single encode. void BM_Encode(benchmark::State& state, DatasetMaker make) { - const std::size_t size = bench_size(); + const std::size_t size = static_cast(state.range(0)); std::mt19937_64 rng(bench_seed()); const std::vector data = make(size, rng); - // Encode once outside timing to report compression ratio. - const PivCoHuffman probe(data); - report_ratio(state, probe.compressed_size(), size); - + std::size_t compressed = 0; for (auto _ : state) { PivCoHuffman codec(data); + compressed = codec.compressed_size(); benchmark::DoNotOptimize(codec.compressed_data().data()); benchmark::ClobberMemory(); } + if (compressed > 0) { + report_ratio(state, compressed, size); + } state.SetBytesProcessed(static_cast(state.iterations() * size)); } /// @brief Measure decoding: reconstruct the original byte stream from the /// in-memory tree. Encoding is done once, outside the timed loop. void BM_Decode(benchmark::State& state, DatasetMaker make) { - const std::size_t size = bench_size(); + const std::size_t size = static_cast(state.range(0)); std::mt19937_64 rng(bench_seed()); const std::vector data = make(size, rng); @@ -172,18 +170,28 @@ void BM_Decode(benchmark::State& state, DatasetMaker make) { } // namespace -// Register encode + decode for every built-in dataset. Each registration is -// named `BM_/`, so `--benchmark_filter` selects subsets freely. -BENCHMARK_CAPTURE(BM_Encode, Uniform256, make_uniform256); -BENCHMARK_CAPTURE(BM_Encode, Binary, make_binary); -BENCHMARK_CAPTURE(BM_Encode, LowEntropy4, make_low_entropy_4); -BENCHMARK_CAPTURE(BM_Encode, TextLike, make_text_like); -BENCHMARK_CAPTURE(BM_Encode, Skewed99, make_skewed99); -BENCHMARK_CAPTURE(BM_Encode, SingleSymbol, make_single_symbol); - -BENCHMARK_CAPTURE(BM_Decode, Uniform256, make_uniform256); -BENCHMARK_CAPTURE(BM_Decode, Binary, make_binary); -BENCHMARK_CAPTURE(BM_Decode, LowEntropy4, make_low_entropy_4); -BENCHMARK_CAPTURE(BM_Decode, TextLike, make_text_like); -BENCHMARK_CAPTURE(BM_Decode, Skewed99, make_skewed99); -BENCHMARK_CAPTURE(BM_Decode, SingleSymbol, make_single_symbol); +// Register encode + decode for every built-in dataset, each swept over the +// full size range. Each registration is named `BM_//`, so +// `--benchmark_filter` selects subsets freely. +#define PIXIE_PIVCO_REGISTER(Op, Dataset, Maker) \ + BENCHMARK_CAPTURE(BM_##Op, Dataset, Maker) \ + ->Unit(benchmark::kMillisecond) \ + ->UseRealTime() \ + ->Repetitions(5) \ + ->ReportAggregatesOnly(true) \ + ->RangeMultiplier(kSizeMult) \ + ->Range(kSizeLo, kSizeHi); + +PIXIE_PIVCO_REGISTER(Encode, Uniform256, make_uniform256) +PIXIE_PIVCO_REGISTER(Encode, Binary, make_binary) +PIXIE_PIVCO_REGISTER(Encode, LowEntropy4, make_low_entropy_4) +PIXIE_PIVCO_REGISTER(Encode, TextLike, make_text_like) +PIXIE_PIVCO_REGISTER(Encode, Skewed99, make_skewed99) +PIXIE_PIVCO_REGISTER(Encode, SingleSymbol, make_single_symbol) + +PIXIE_PIVCO_REGISTER(Decode, Uniform256, make_uniform256) +PIXIE_PIVCO_REGISTER(Decode, Binary, make_binary) +PIXIE_PIVCO_REGISTER(Decode, LowEntropy4, make_low_entropy_4) +PIXIE_PIVCO_REGISTER(Decode, TextLike, make_text_like) +PIXIE_PIVCO_REGISTER(Decode, Skewed99, make_skewed99) +PIXIE_PIVCO_REGISTER(Decode, SingleSymbol, make_single_symbol) From 96293e91e3960100659e0d5116d9abdab4899cac Mon Sep 17 00:00:00 2001 From: Sanya239 Date: Tue, 21 Jul 2026 20:18:19 +0300 Subject: [PATCH 03/14] Add benchmark for file huffman encoding --- CMakeLists.txt | 9 ++ .../pivco_huffman_file_benchmarks.cpp | 135 ++++++++++++++++++ 2 files changed, 144 insertions(+) create mode 100644 src/benchmarks/pivco_huffman_file_benchmarks.cpp diff --git a/CMakeLists.txt b/CMakeLists.txt index 6828f92..72cf718 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -418,6 +418,15 @@ if (PIXIE_BENCHMARKS) benchmark benchmark_main ${PIXIE_DIAGNOSTICS_LIBS}) + + add_executable(pivco_huffman_file_benchmarks + src/benchmarks/pivco_huffman_file_benchmarks.cpp) + target_include_directories(pivco_huffman_file_benchmarks + PUBLIC include) + target_link_libraries(pivco_huffman_file_benchmarks + benchmark + benchmark_main + ${PIXIE_DIAGNOSTICS_LIBS}) endif () # --------------------------------------------------------------------------- diff --git a/src/benchmarks/pivco_huffman_file_benchmarks.cpp b/src/benchmarks/pivco_huffman_file_benchmarks.cpp new file mode 100644 index 0000000..1795103 --- /dev/null +++ b/src/benchmarks/pivco_huffman_file_benchmarks.cpp @@ -0,0 +1,135 @@ +#include +#include + +#include +#include +#include +#include +#include +#include + +using pixie::PivCoHuffman; + +// --------------------------------------------------------------------------- +// File-based PivCo-Huffman benchmark. +// +// Reads a real-world text file (default: prose_pride.txt, "Pride and +// Prejudice" from Project Gutenberg) and measures encode/decode throughput +// on it. This matches the article's evaluation methodology, which uses real +// datasets rather than synthetic random data. +// +// Configuration: +// PIVCO_BENCH_FILE - path to the input file (default: prose_pride.txt) +// --------------------------------------------------------------------------- + +namespace { + +/// @brief Read the entire contents of @p path into a byte vector. +std::vector read_file(const std::string& path) { + std::ifstream file(path, std::ios::binary | std::ios::ate); + if (!file) { + std::fprintf(stderr, "Cannot open '%s'\n", path.c_str()); + std::exit(1); + } + const std::streamsize size = file.tellg(); + file.seekg(0); + std::vector data(static_cast(size)); + if (size > 0) { + file.read(reinterpret_cast(data.data()), size); + } + return data; +} + +/// @brief Input file path, overridable via `PIVCO_BENCH_FILE`. +const std::string& bench_file() { + static const std::string path = [] { + if (const char* env = std::getenv("PIVCO_BENCH_FILE")) { + return std::string(env); + } + return std::string("prose_pride.txt"); + }(); + return path; +} + +/// @brief Compression efficiency in bits per input byte. +void report_ratio(benchmark::State& state, + std::size_t compressed_bytes, + std::size_t input_bytes) { + state.counters["bpb"] = static_cast(compressed_bytes) * 8.0 / + static_cast(input_bytes); +} + +/// @brief Compute Shannon entropy of the data in bits per byte, for reference. +double entropy(std::span data) { + if (data.empty()) { + return 0.0; + } + std::array freq{}; + for (auto b : data) { + freq[b]++; + } + double h = 0.0; + const double n = static_cast(data.size()); + for (std::size_t i = 0; i < 256; i++) { + if (freq[i] > 0) { + const double p = static_cast(freq[i]) / n; + h -= p * std::log2(p); + } + } + return h; +} + +// --- benchmarks ------------------------------------------------------------ + +/// @brief Measure full encoding of the file. +void BM_EncodeFile(benchmark::State& state) { + const std::vector data = read_file(bench_file()); + const std::size_t size = data.size(); + + // Report entropy once as a reference counter. + state.counters["entropy_bpb"] = entropy(data); + + std::size_t compressed = 0; + for (auto _ : state) { + PivCoHuffman codec(data); + compressed = codec.compressed_size(); + benchmark::DoNotOptimize(codec.compressed_data().data()); + benchmark::ClobberMemory(); + } + if (compressed > 0) { + report_ratio(state, compressed, size); + } + state.SetBytesProcessed(static_cast(state.iterations() * size)); +} + +/// @brief Measure decoding of the file. Encoding is done once, outside timing. +void BM_DecodeFile(benchmark::State& state) { + const std::vector data = read_file(bench_file()); + const std::size_t size = data.size(); + + state.counters["entropy_bpb"] = entropy(data); + + const PivCoHuffman codec(data); + report_ratio(state, codec.compressed_size(), size); + + for (auto _ : state) { + std::vector out = codec.decode(); + benchmark::DoNotOptimize(out.data()); + benchmark::ClobberMemory(); + } + state.SetBytesProcessed(static_cast(state.iterations() * size)); +} + +} // namespace + +BENCHMARK(BM_EncodeFile) + ->Unit(benchmark::kMillisecond) + ->UseRealTime() + ->Repetitions(10) + ->ReportAggregatesOnly(true); + +BENCHMARK(BM_DecodeFile) + ->Unit(benchmark::kMillisecond) + ->UseRealTime() + ->Repetitions(10) + ->ReportAggregatesOnly(true); From 09c7e6485d6bb31e6548c09616f0298c1fdbb618 Mon Sep 17 00:00:00 2001 From: Sanya239 Date: Wed, 22 Jul 2026 02:37:18 +0300 Subject: [PATCH 04/14] Speed and memory optimization --- include/pixie/huffman/pivco_huffman.h | 696 ++++++++++++++++++++------ 1 file changed, 535 insertions(+), 161 deletions(-) diff --git a/include/pixie/huffman/pivco_huffman.h b/include/pixie/huffman/pivco_huffman.h index 48f40be..6948e9b 100644 --- a/include/pixie/huffman/pivco_huffman.h +++ b/include/pixie/huffman/pivco_huffman.h @@ -2,7 +2,7 @@ /** * @file pivco_huffman.h - * @brief Simple scalar PivCo-Huffman codec. + * @brief Simple scalar PivCo-Huffman codec with canonical Huffman codes. * * Stores a Huffman-shaped tree of per-node routing bitmaps (the "tree of * bitmaps" layout shared with wavelet trees). Encoding walks each input symbol @@ -10,11 +10,16 @@ * path. Decoding merges child symbol streams bottom-up from leaves to root, * using the node bitmap as a selector. * + * The tree uses canonical Huffman code lengths (<= 15 bits, enforced via the + * package-merge length-limited algorithm): the wire format stores only the 256 + * per-symbol lengths (128 bytes as 4-bit nibbles), and the tree structure is + * reconstructed deterministically at load time. This groups same-length codes + * together, which is a prerequisite for future flat-subtree optimizations and + * reduces serialized header size. + * * This is a deliberately simple, unoptimized reference implementation: node * bitmaps are stored as packed `std::uint64_t` words, traversal is scalar, - * and there are no flat-subtree, non-canonical, SIMD, or selective-ANS - * optimizations. It is intended as a correct baseline for the PivCo-Huffman - * layout and for future optimized variants. + * and there are no flat-subtree, SIMD, or selective-ANS optimizations. * * Range and ownership conventions follow `HuffmanBase`: symbols are bytes, * the compressed stream is a byte view, and the codec owns its serialized form. @@ -22,12 +27,12 @@ #include +#include #include #include #include #include #include -#include #include #include #include @@ -35,12 +40,13 @@ namespace pixie { /** - * @brief Simple scalar PivCo-Huffman codec. + * @brief Simple scalar PivCo-Huffman codec with canonical Huffman codes. * @details Implements `HuffmanBase` by building a Huffman tree over the - * byte alphabet, storing one routing bitmap per internal node as - * packed 64-bit words, and decoding by bottom-up merging of child - * symbol streams. The serialized form is a header followed by the - * tree structure and packed node bitmaps. + * byte alphabet, canonicalizing it (grouping same-length codes), + * storing one routing bitmap per internal node as packed 64-bit + * words, and decoding by bottom-up merging of child symbol streams. + * The serialized form is 256 code lengths followed by packed node + * bitmaps; the tree is reconstructed from the lengths at load time. */ class PivCoHuffman : public HuffmanBase { public: @@ -90,29 +96,26 @@ class PivCoHuffman : public HuffmanBase { /// @brief Bits per word used by the packed bitmap. static constexpr std::size_t kWordBits = 64; + /// @brief Maximum supported Huffman code length (matches common decoders). + static constexpr std::size_t kMaxCodeLength = 15; + /** - * @brief Packed routing bitmap stored as 64-bit words. - * @details Bit `i` lives in word `i / 64` at position `i % 64`. This avoids - * the slow bit-proxy access of `std::vector` and allows - * bulk serialization via direct word `memcpy`. + * @brief Lightweight descriptor for a node's routing bitmap. + * @details The packed 64-bit words live in a single shared arena + * (`arena_`); this struct only records where the node's words begin + * (`offset`, in words) and how many symbols they describe + * (`count`). Splitting descriptor from storage lets one contiguous + * allocation back every internal node, eliminating per-node heap + * traffic during decode and improving cache locality. */ struct Bitmap { - std::vector words; + std::size_t offset = 0; std::size_t count = 0; /** @brief Number of 64-bit words backing the bitmap. */ std::size_t word_count() const { return (count + kWordBits - 1) / kWordBits; } - - /** @brief Number of set (1) bits in the bitmap. */ - std::size_t popcount() const { - std::size_t total = 0; - for (std::uint64_t w : words) { - total += std::popcount(w); - } - return total; - } }; /** @@ -133,10 +136,21 @@ class PivCoHuffman : public HuffmanBase { std::size_t root_ = kNpos; std::vector nodes_; std::vector compressed_; + /// @brief One contiguous allocation backing every internal node's bitmap + /// words. Node `i` owns `nodes_[i].bits.word_count()` words starting + /// at `arena_[nodes_[i].bits.offset]`. + std::vector arena_; + /// @brief Reusable decode scratch: two ping-pong halves of size + /// `uncompressed_size_`. Allocated once and reused across decode + /// calls; `mutable` because `decode_impl()` is const. + mutable std::vector workspace_; + std::array code_lengths_{}; // --- construction -------------------------------------------------------- - /** @brief Build the Huffman tree and per-node bitmaps from @p input. */ + /** @brief Build the canonical Huffman tree and per-node bitmaps from @p + * input. + */ void build(std::span input) { uncompressed_size_ = input.size(); if (uncompressed_size_ == 0) { @@ -144,60 +158,302 @@ class PivCoHuffman : public HuffmanBase { return; } + // 1. Count symbol frequencies. std::array freq{}; for (auto s : input) { freq[s]++; } - struct HeapItem { + // 2. Build a Huffman tree to determine code lengths, then extract lengths. + compute_code_lengths(freq, code_lengths_); + + // 3. Rebuild a canonical tree from the lengths (groups same-length codes). + // Structure only: bitmap weights/offsets are assigned below, not here. + build_canonical_tree(code_lengths_); + + // Single-symbol input: root is a leaf, no internal nodes, no arena. + if (root_ == kNpos || nodes_[root_].is_leaf) { + return; + } + + // 4. Assign exact per-node weights (subtree frequency sums) and arena + // offsets in a single post-order walk, then allocate the arena once. + std::size_t arena_words = 0; + assign_weights_and_offsets(root_, freq, arena_words); + arena_.assign(arena_words, 0); + + // 5. Assign per-symbol code paths, then fill node bitmaps via top-down + // in-place partition into the reusable workspace. This mirrors decode: + // the root reads its input from workspace half 0, writes its bitmap, + // and partitions symbols into half 1 for its children. Each level + // ping-pongs between the two halves — no per-node allocations. + std::array>, kAlphabet> paths; + std::vector> path; + assign_paths(root_, path, paths); + workspace_.resize(uncompressed_size_ * 2); + std::copy(input.begin(), input.end(), workspace_.data()); + encode_partition(root_, 0, 0, paths, freq); + } + + /** @brief Compute length-limited Huffman code lengths via package-merge. + * @details Implements the Larmore-Hirschberg package-merge algorithm to + * produce optimal code lengths subject to `length <= + * kMaxCodeLength`. This guarantees the 4-bit nibble header never truncates a + * length, and keeps all downstream canonical-tree, flat-subtree, and + * reshaping optimizations valid. + * + * Algorithm: maintain a sorted list of "items" per level (1..L). A level-1 + * item is a single coin (one symbol). At each subsequent level, form + * "packages" by pairing consecutive items from the previous level (a + * package commits to selecting both its constituents), then merge with the + * original coins and keep the cheapest `2n-2` items. After L-1 iterations, + * summing each symbol's count across the final `2n-2` items yields its code + * length (each ≤ L, and the Kraft inequality holds exactly). + * + * Each item carries a per-symbol count vector (`Counts`) so that summing at + * the end is a flat array reduction; with `n <= 256` and `L = 15` the total + * work is ~2M additions (< 1 ms), negligible relative to the bitmap fill. */ + void compute_code_lengths(const std::array& freq, + std::array& lengths) { + lengths.fill(0); + + // Collect present symbols, sorted by weight ascending (the coin list S, + // reused at every level of the package-merge). + struct Coin { + std::uint8_t symbol; std::size_t weight; - std::size_t idx; }; - auto cmp = [](const HeapItem& a, const HeapItem& b) { - return a.weight > b.weight; + std::vector coins; + coins.reserve(kAlphabet); + for (std::size_t s = 0; s < kAlphabet; s++) { + if (freq[s] > 0) { + coins.push_back({static_cast(s), freq[s]}); + } + } + std::sort(coins.begin(), coins.end(), + [](const Coin& a, const Coin& b) { return a.weight < b.weight; }); + + const std::size_t n = coins.size(); + if (n == 0) { + return; + } + // Single distinct symbol: must still get a 1-bit code to be decodable. + if (n == 1) { + lengths[coins[0].symbol] = 1; + return; + } + + // Per-item symbol-count vector. uint16 suffices: a symbol's count within a + // single package can grow across levels, but final lengths are <= L = 15. + using Counts = std::array; + struct Item { + std::size_t weight; + Counts counts; }; - std::priority_queue, decltype(cmp)> heap( - cmp); - // At most 2*kAlphabet-1 nodes (leaves + internal). Pre-allocate to avoid - // reallocation during tree construction. - nodes_.reserve(2 * kAlphabet - 1); + // Level-1 list: one coin per symbol. + std::vector cur; + cur.reserve(2 * n); + for (const auto& c : coins) { + Item it{}; + it.weight = c.weight; + it.counts[c.symbol] = 1; + cur.push_back(std::move(it)); + } + + // Keep the cheapest 2n-2 items at each level (the final selection size; + // cheaper items at a level can never displace a more expensive one in an + // optimal solution, so pruning to 2n-2 is safe). + const std::size_t keep = 2 * n - 2; + + // Iterate levels 2..L. Coins are reused at every level via the merge. + for (std::size_t level = 1; level < kMaxCodeLength; level++) { + // Package: pair consecutive items of `cur` (already weight-sorted). + std::vector packages; + packages.reserve(cur.size() / 2); + for (std::size_t i = 0; i + 1 < cur.size(); i += 2) { + Item pkg{}; + pkg.weight = cur[i].weight + cur[i + 1].weight; + for (std::size_t j = 0; j < kAlphabet; j++) { + pkg.counts[j] = cur[i].counts[j] + cur[i + 1].counts[j]; + } + packages.push_back(std::move(pkg)); + } + // Merge coins (sorted) with packages (sorted — pairing preserves order), + // keeping the cheapest `keep` items. Ties prefer coins (deterministic). + // At early levels the combined list may hold fewer than `keep` items; + // keep all of them in that case (the list grows toward `keep` over + // levels as packages proliferate). + const std::size_t available = n + packages.size(); + const std::size_t keep_here = std::min(keep, available); + std::vector next; + next.reserve(keep_here); + std::size_t ci = 0; + std::size_t pi = 0; + while (next.size() < keep_here) { + bool use_coin; + if (ci >= n) { + use_coin = false; + } else if (pi >= packages.size()) { + use_coin = true; + } else { + use_coin = coins[ci].weight <= packages[pi].weight; + } + if (use_coin) { + Item it{}; + it.weight = coins[ci].weight; + it.counts[coins[ci].symbol] = 1; + next.push_back(std::move(it)); + ci++; + } else { + next.push_back(std::move(packages[pi])); + pi++; + } + } + cur = std::move(next); + } + + // Sum symbol counts across the final 2n-2 selected items -> code lengths. + Counts total{}; + for (const auto& it : cur) { + for (std::size_t j = 0; j < kAlphabet; j++) { + total[j] += it.counts[j]; + } + } + for (const auto& c : coins) { + lengths[c.symbol] = static_cast(total[c.symbol]); + } + } + /** @brief Rebuild a canonical Huffman tree from code lengths. + * @details Assigns canonical codes (sorted by length, then symbol), then + * builds the tree top-down by inserting each (symbol, code) pair. + * Same-length codes are grouped together, which is a prerequisite + * for future flat-subtree optimization. */ + void build_canonical_tree( + const std::array& lengths) { + nodes_.clear(); + nodes_.reserve(2 * kAlphabet); + + // Collect symbols present, sorted by (length, symbol). + struct SymbolEntry { + std::uint8_t symbol; + std::uint8_t length; + }; + std::vector entries; + entries.reserve(kAlphabet); for (std::size_t s = 0; s < kAlphabet; s++) { - if (freq[s] > 0) { - Node n; - n.is_leaf = true; - n.symbol = static_cast(s); - nodes_.push_back(std::move(n)); - heap.push({freq[s], nodes_.size() - 1}); + if (lengths[s] > 0) { + entries.push_back({static_cast(s), lengths[s]}); } } + std::sort(entries.begin(), entries.end(), + [](const SymbolEntry& a, const SymbolEntry& b) { + if (a.length != b.length) { + return a.length < b.length; + } + return a.symbol < b.symbol; + }); + + // Special case: single distinct symbol → root is a leaf (no internal nodes, + // no bitmaps). decode_node returns a constant-filled vector. + if (entries.size() == 1) { + nodes_.push_back(Node{}); + root_ = 0; + nodes_[0].is_leaf = true; + nodes_[0].symbol = entries[0].symbol; + return; + } - while (heap.size() > 1) { - HeapItem a = heap.top(); - heap.pop(); - HeapItem b = heap.top(); - heap.pop(); - Node n; - n.is_leaf = false; - n.left = a.idx; - n.right = b.idx; - std::size_t w = a.weight + b.weight; - n.bits.count = w; - n.bits.words.assign((w + kWordBits - 1) / kWordBits, 0); - nodes_.push_back(std::move(n)); - heap.push({w, nodes_.size() - 1}); - } - root_ = heap.top().idx; - - // Record, for each symbol, the path of (node, direction-bit) from root. - std::array>, kAlphabet> paths; - std::vector> path; - assign_paths(root_, path, paths); + // Assign canonical codes: first code at each length is 0, increment + // within a length, shift left when length increases. + std::vector> + codes; // (code, length) + codes.reserve(entries.size()); + std::uint64_t code = 0; + std::uint8_t prev_len = 0; + for (const auto& e : entries) { + if (prev_len == 0) { + prev_len = e.length; + } + while (prev_len < e.length) { + code <<= 1; + prev_len++; + } + codes.push_back({code, e.length}); + code++; + } + + // Build tree by inserting each (symbol, code) pair via bit-by-bit descent. + // Create root. + nodes_.push_back(Node{}); + root_ = 0; + + for (std::size_t i = 0; i < entries.size(); i++) { + std::uint64_t c = codes[i].first; + std::uint8_t len = codes[i].second; + std::size_t cur = root_; + + // Descend len-1 bits, creating internal nodes as needed. + for (std::size_t bit = 0; bit + 1 < len; bit++) { + // Read MSB first (canonical codes are assigned MSB-first). + bool go_right = (c >> (len - 1 - bit)) & 1; + std::size_t& child = go_right ? nodes_[cur].right : nodes_[cur].left; + if (child == kNpos) { + nodes_.push_back(Node{}); + child = nodes_.size() - 1; + } + cur = child; + } + // Last bit → leaf. + bool go_right = (c >> 0) & 1; + std::size_t& child = go_right ? nodes_[cur].right : nodes_[cur].left; + nodes_.push_back(Node{}); + child = nodes_.size() - 1; + nodes_[child].is_leaf = true; + nodes_[child].symbol = entries[i].symbol; + } + + // Bitmap weights and arena offsets are assigned by the caller: `build()` + // derives them from symbol frequencies (exact), `deserialize()` reads + // them from the wire. This keeps the tree builder free of bitmap storage. + } - // Top-down partitioning: recursively split the symbol stream and fill - // pre-allocated bitmaps. Leaves are skipped (zero overhead). - encode_node(root_, 0, input, paths); + /** @brief Post-order assignment of exact per-node weights and arena offsets. + * @details The weight of an internal node is the sum of leaf frequencies in + * its subtree (exact, not a parent-derived upper bound). Assigns + * `bits.count` and `bits.offset` for each internal node and + * accumulates the total arena word count. Returns the subtree + * weight so each parent can sum its two children. */ + std::size_t assign_weights_and_offsets( + std::size_t idx, + const std::array& freq, + std::size_t& arena_words) { + if (idx == kNpos) { + return 0; + } + Node& n = nodes_[idx]; + if (n.is_leaf) { + return freq[n.symbol]; + } + std::size_t left_weight = + assign_weights_and_offsets(n.left, freq, arena_words); + std::size_t right_weight = + assign_weights_and_offsets(n.right, freq, arena_words); + n.bits.count = left_weight + right_weight; + n.bits.offset = arena_words; + arena_words += n.bits.word_count(); + return n.bits.count; + } + + /** @brief Exact weight of a child node, used to place partitions. + * @details Internal nodes carry their precomputed `bits.count`; leaves + * carry their symbol frequency (no bitmap stored). */ + std::size_t child_weight( + std::size_t idx, + const std::array& freq) const { + const Node& n = nodes_[idx]; + return n.is_leaf ? freq[n.symbol] : n.bits.count; } /** @brief Depth-first assignment of root-to-leaf paths. */ @@ -217,42 +473,51 @@ class PivCoHuffman : public HuffmanBase { path.pop_back(); } - /** @brief Top-down encoding: recursively partition symbols and fill bitmaps. - * @details At each internal node, fill the pre-allocated bitmap with the - * direction bit of each symbol's code at this depth, and split the - * symbol stream into left (0) and right (1) child vectors. Leaves - * return immediately — zero overhead. - * @param idx Current node index. - * @param depth Current tree depth (indexes into each symbol's code path). - * @param symbols Symbol stream arriving at this node, in order. - * @param paths Per-symbol code paths from `assign_paths`. */ - void encode_node(std::size_t idx, - std::size_t depth, - std::span symbols, - const std::array>, - kAlphabet>& paths) { + /** @brief Top-down in-place partition encoding into the reusable workspace. + * @details The dual of decode: at each internal node, read the incoming + * symbol stream from workspace half `depth`, fill the node's + * arena bitmap with direction bits, and partition symbols into + * half `depth+1` (left at `out_base`, right at + * `out_base + left_weight`). Leaves return immediately. No + * per-node allocations — the workspace is pre-sized and + * ping-ponged by depth, exactly mirroring decode. + * @param idx Current node index. + * @param depth Current tree depth (selects workspace half + path index). + * @param out_base Offset within the current half where this node's stream + * begins. + * @param paths Per-symbol code paths from `assign_paths`. + * @param freq Symbol frequencies (sizing child partitions). */ + void encode_partition( + std::size_t idx, + std::size_t depth, + std::size_t out_base, + const std::array>, kAlphabet>& + paths, + const std::array& freq) { if (nodes_[idx].is_leaf) { return; } - auto& n = nodes_[idx]; - std::vector left_symbols, right_symbols; - std::size_t left_size = nodes_[n.left].bits.count; - std::size_t right_size = nodes_[n.right].bits.count; - left_symbols.reserve(left_size); - right_symbols.reserve(right_size); + Node& n = nodes_[idx]; + const std::size_t left_weight = child_weight(n.left, freq); + const std::size_t weight = n.bits.count; + + const symbol_type* src = workspace_half(depth) + out_base; + symbol_type* left_dst = workspace_half(depth + 1) + out_base; + symbol_type* right_dst = workspace_half(depth + 1) + out_base + left_weight; - // Fill bitmap via incremental word/bit tracking — avoids per-symbol - // division and modulo. + std::uint64_t* bits = arena_.data() + n.bits.offset; std::size_t word_idx = 0; std::size_t bit_pos = 0; - for (std::size_t i = 0; i < symbols.size(); i++) { - bool bit = paths[symbols[i]][depth].second; - if (bit) { - n.bits.words[word_idx] |= (1ull << bit_pos); - right_symbols.push_back(symbols[i]); + std::size_t c_left = 0; + std::size_t c_right = 0; + for (std::size_t i = 0; i < weight; i++) { + symbol_type s = src[i]; + if (paths[s][depth].second) { + bits[word_idx] |= (1ull << bit_pos); + right_dst[c_right++] = s; } else { - left_symbols.push_back(symbols[i]); + left_dst[c_left++] = s; } if (++bit_pos == kWordBits) { bit_pos = 0; @@ -260,78 +525,137 @@ class PivCoHuffman : public HuffmanBase { } } - encode_node(n.left, depth + 1, left_symbols, paths); - encode_node(n.right, depth + 1, right_symbols, paths); + encode_partition(n.left, depth + 1, out_base, paths, freq); + encode_partition(n.right, depth + 1, out_base + left_weight, paths, freq); } // --- decode -------------------------------------------------------------- - /** @brief Reconstruct the original sequence by bottom-up merging. */ + /** @brief Reconstruct the original sequence by bottom-up merging. + * @details Uses a single reusable workspace buffer (sized to twice the + * uncompressed length, allocated once and reused across calls) + * and ping-pongs between its two halves by tree depth. Each node + * reads its children's outputs from one half and writes its merged + * output to the other, so no per-node allocation occurs after the + * workspace is sized. The final result is copied out of half 0. */ std::vector decode_from_tree() const { if (root_ == kNpos) { return std::vector(); } - return decode_node(root_, uncompressed_size_); + if (nodes_[root_].is_leaf) { + return std::vector(uncompressed_size_, nodes_[root_].symbol); + } + const std::size_t n = uncompressed_size_; + workspace_.resize(n * 2); + decode_node(root_, n, 0, 0); + return std::vector(workspace_.begin(), workspace_.begin() + n); } - /** @brief Recursively decode a node's output stream bottom-up. - * @param weight Number of symbols this node must produce. For internal - * nodes this equals `bits.count`; for leaves it is passed down from - * the parent (no bitmap to read it from). */ - std::vector decode_node(std::size_t idx, - std::size_t weight) const { - const auto& n = nodes_[idx]; + /** @brief Recursively decode a node's output stream bottom-up into the + * reusable workspace. + * @param weight Number of symbols this node must produce. + * @param depth Tree depth, selecting which workspace half receives this + * node's output; children write to the opposite half. + * @param out_base Offset within the selected half where this node's output + * begins. Children are placed at `out_base` (left) and + * `out_base + left_weight` (right) in the opposite half, + * then merged into `[out_base, out_base + weight)` here. */ + void decode_node(std::size_t idx, + std::size_t weight, + std::size_t depth, + std::size_t out_base) const { + const Node& n = nodes_[idx]; + symbol_type* dst = workspace_half(depth) + out_base; if (n.is_leaf) { - return std::vector(weight, n.symbol); + std::fill(dst, dst + weight, n.symbol); + return; } - std::size_t right_weight = n.bits.popcount(); - std::size_t left_weight = n.bits.count - right_weight; - std::vector left_symbols = decode_node(n.left, left_weight); - std::vector right_symbols = decode_node(n.right, right_weight); + const std::size_t right_weight = bitmap_popcount(n.bits); + const std::size_t left_weight = n.bits.count - right_weight; + decode_node(n.left, left_weight, depth + 1, out_base); + decode_node(n.right, right_weight, depth + 1, out_base + left_weight); + + // Children wrote to the opposite half: [out_base, out_base + left_weight) + // holds the left stream, [out_base + left_weight, out_base + weight) + // holds the right stream. Merge both into this node's half by the bitmap. + const symbol_type* src = workspace_half(depth + 1) + out_base; + const symbol_type* left_out = src; + const symbol_type* right_out = src + left_weight; - std::vector out(n.bits.count); + std::size_t out_i = 0; std::size_t c_left = 0; std::size_t c_right = 0; - // Word-wise merge: read each 64-bit word once and shift through its bits, - // avoiding per-bit division/modulo. - for (std::size_t w = 0; w < n.bits.words.size(); w++) { - std::uint64_t word = n.bits.words[w]; - std::size_t base = w * kWordBits; - std::size_t limit = base + kWordBits; - if (limit > n.bits.count) { - limit = n.bits.count; + const std::uint64_t* bits = arena_.data() + n.bits.offset; + const std::size_t words = n.bits.word_count(); + for (std::size_t w = 0; w < words; w++) { + std::uint64_t word = bits[w]; + std::size_t limit = kWordBits; + if (w + 1 == words) { + limit = n.bits.count - w * kWordBits; } - for (std::size_t i = base; i < limit; i++) { - if (word & 1ull) { - out[i] = right_symbols[c_right++]; - } else { - out[i] = left_symbols[c_left++]; - } + for (std::size_t i = 0; i < limit; i++) { + dst[out_i++] = + (word & 1ull) ? right_out[c_right++] : left_out[c_left++]; word >>= 1; } } - return out; + } + + /** @brief Pointer to the workspace half that holds depth-@p depth outputs. */ + symbol_type* workspace_half(std::size_t depth) const { + return workspace_.data() + (depth % 2) * uncompressed_size_; + } + + /** @brief Number of set bits in a node's arena-backed bitmap. */ + std::size_t bitmap_popcount(const Bitmap& b) const { + const std::uint64_t* p = arena_.data() + b.offset; + std::size_t total = 0; + const std::size_t words = b.word_count(); + for (std::size_t i = 0; i < words; i++) { + total += std::popcount(p[i]); + } + return total; } // --- serialization ------------------------------------------------------- - /** @brief Write the in-memory tree into the serialized byte buffer. */ + /** @brief Write the in-memory tree into the serialized byte buffer. + * @details Wire format: + * - uncompressed_size (8 bytes) + * - 256 code lengths as 128 bytes of 4-bit nibbles (0 = absent) + * - For each internal node (pre-order): bits.count (8 bytes) + + * packed words. Leaves are implied by the tree structure + * reconstructed from lengths. */ void serialize() { compressed_.clear(); write(uncompressed_size_); if (uncompressed_size_ == 0) { return; } - write(root_); - write(nodes_.size()); - for (const auto& n : nodes_) { - write(static_cast(n.is_leaf ? 1 : 0)); - write(n.symbol); - write(n.left); - write(n.right); + + // Pack stored code lengths as nibbles: 2 symbols per byte. + for (std::size_t i = 0; i < kAlphabet; i += 2) { + std::uint8_t byte = static_cast( + code_lengths_[i] | (code_lengths_[i + 1] << 4)); + write(byte); + } + + // Write internal node bitmaps in pre-order (matches deserialize order). + serialize_node(root_); + } + + /** @brief Serialize a node's bitmap and recurse (pre-order). */ + void serialize_node(std::size_t idx) { + if (idx == kNpos) { + return; + } + const auto& n = nodes_[idx]; + if (!n.is_leaf) { write(n.bits.count); - write_words(n.bits.words); + write_words(n.bits.offset, n.bits.word_count()); + serialize_node(n.left); + serialize_node(n.right); } } @@ -339,6 +663,7 @@ class PivCoHuffman : public HuffmanBase { void deserialize(std::span data) { compressed_.assign(data.begin(), data.end()); nodes_.clear(); + arena_.clear(); if (data.empty()) { uncompressed_size_ = 0; root_ = kNpos; @@ -350,18 +675,80 @@ class PivCoHuffman : public HuffmanBase { root_ = kNpos; return; } - root_ = read(data, pos); - const std::size_t count = read(data, pos); - nodes_.resize(count); - for (std::size_t i = 0; i < count; i++) { - const std::uint8_t flags = read(data, pos); - nodes_[i].is_leaf = (flags & 1) != 0; - nodes_[i].symbol = read(data, pos); - nodes_[i].left = read(data, pos); - nodes_[i].right = read(data, pos); - nodes_[i].bits.count = read(data, pos); - nodes_[i].bits.words = read_words(data, pos, nodes_[i].bits.word_count()); + + // Read 256 code lengths from nibbles. + for (std::size_t i = 0; i < kAlphabet; i += 2) { + std::uint8_t byte = read(data, pos); + code_lengths_[i] = byte & 0x0F; + code_lengths_[i + 1] = (byte >> 4) & 0x0F; + } + + // Rebuild the canonical tree from lengths. + build_canonical_tree(code_lengths_); + + // Single-symbol input: root is a leaf, no internal nodes, no arena. + if (root_ == kNpos || nodes_[root_].is_leaf) { + return; + } + + // Pass 1: read each internal node's `count` (and skip its word bytes on + // the wire), assign contiguous arena offsets, and size the arena exactly. + // `scan_pos` is discarded — pass 2 re-walks the same region to read words. + const std::size_t bitmap_start = pos; + std::size_t arena_words = 0; + scan_node_counts(root_, data, pos, arena_words); + arena_.assign(arena_words, 0); + + // Pass 2: rewind to the bitmap section and read each node's packed words + // into its arena slot. + pos = bitmap_start; + read_node_words(root_, data, pos); + } + + /** @brief Pass 1: read each internal node's `count` from the wire, assign its + * arena offset, and advance past its word bytes without storing + * them. Sums the total arena word count so the caller allocates once. + * @details Pre-order traversal matches the serialization order. */ + void scan_node_counts(std::size_t idx, + std::span data, + std::size_t& pos, + std::size_t& arena_words) { + if (idx == kNpos) { + return; + } + Node& n = nodes_[idx]; + if (n.is_leaf) { + return; + } + n.bits.count = read(data, pos); + n.bits.offset = arena_words; + arena_words += n.bits.word_count(); + pos += n.bits.word_count() * sizeof(std::uint64_t); + scan_node_counts(n.left, data, pos, arena_words); + scan_node_counts(n.right, data, pos, arena_words); + } + + /** @brief Pass 2: read each internal node's packed words from the wire into + * its arena slot (pre-order, matching serialization). */ + void read_node_words(std::size_t idx, + std::span data, + std::size_t& pos) { + if (idx == kNpos) { + return; + } + Node& n = nodes_[idx]; + if (n.is_leaf) { + return; } + pos += sizeof(std::size_t); // count, already read in pass 1 + const std::size_t words = n.bits.word_count(); + const std::size_t bytes = words * sizeof(std::uint64_t); + if (bytes > 0) { + std::memcpy(arena_.data() + n.bits.offset, data.data() + pos, bytes); + } + pos += bytes; + read_node_words(n.left, data, pos); + read_node_words(n.right, data, pos); } /** @brief Append a little-endian, fixed-width value to the byte buffer. */ @@ -372,13 +759,13 @@ class PivCoHuffman : public HuffmanBase { std::memcpy(compressed_.data() + old, &value, sizeof(T)); } - /** @brief Append a packed word array directly to the byte buffer. */ - void write_words(const std::vector& words) { - const std::size_t bytes = words.size() * sizeof(std::uint64_t); + /** @brief Append a node's arena-backed words directly to the byte buffer. */ + void write_words(std::size_t offset, std::size_t word_count) { + const std::size_t bytes = word_count * sizeof(std::uint64_t); const std::size_t old = compressed_.size(); compressed_.resize(old + bytes); if (bytes > 0) { - std::memcpy(compressed_.data() + old, words.data(), bytes); + std::memcpy(compressed_.data() + old, arena_.data() + offset, bytes); } } @@ -390,19 +777,6 @@ class PivCoHuffman : public HuffmanBase { pos += sizeof(T); return value; } - - /** @brief Read a packed word array of @p word_count words from @p data. */ - static std::vector read_words(std::span data, - std::size_t& pos, - std::size_t word_count) { - std::vector words(word_count); - const std::size_t bytes = word_count * sizeof(std::uint64_t); - if (bytes > 0) { - std::memcpy(words.data(), data.data() + pos, bytes); - pos += bytes; - } - return words; - } }; } // namespace pixie From e18c0438c5e711ed7442c76370dbd676a4d3e29c Mon Sep 17 00:00:00 2001 From: Sanya239 Date: Wed, 22 Jul 2026 02:52:14 +0300 Subject: [PATCH 05/14] Add block-based processing pipeline --- include/pixie/huffman/pivco_huffman.h | 150 +++++++++++++++++++------- 1 file changed, 110 insertions(+), 40 deletions(-) diff --git a/include/pixie/huffman/pivco_huffman.h b/include/pixie/huffman/pivco_huffman.h index 6948e9b..b3b19ba 100644 --- a/include/pixie/huffman/pivco_huffman.h +++ b/include/pixie/huffman/pivco_huffman.h @@ -57,10 +57,7 @@ class PivCoHuffman : public HuffmanBase { * @brief Build a codec by encoding @p input. * @param input Symbol sequence to compress. */ - explicit PivCoHuffman(std::span input) { - build(input); - serialize(); - } + explicit PivCoHuffman(std::span input) { build(input); } /** * @brief Load a codec from a previously serialized compressed stream. @@ -83,8 +80,24 @@ class PivCoHuffman : public HuffmanBase { return compressed_; } - /** @brief Reconstruct the original symbol sequence. */ - std::vector decode_impl() const { return decode_from_tree(); } + /** @brief Reconstruct the original symbol sequence by decoding each block. */ + std::vector decode_impl() const { + if (uncompressed_size_ == 0) { + return std::vector(); + } + std::vector result; + result.reserve(uncompressed_size_); + // Skip the top-level header (total_size + block_size). + std::size_t pos = sizeof(std::size_t) * 2; + std::size_t remaining = uncompressed_size_; + while (remaining > 0) { + deserialize_block(pos); + std::vector block_result = decode_from_tree(); + result.insert(result.end(), block_result.begin(), block_result.end()); + remaining -= block_uncompressed_size_; + } + return result; + } private: /// @brief Sentinel node index meaning "no node". @@ -99,6 +112,13 @@ class PivCoHuffman : public HuffmanBase { /// @brief Maximum supported Huffman code length (matches common decoders). static constexpr std::size_t kMaxCodeLength = 15; + /// @brief Default block size for block-based processing (64 KiB). + /// @details Each block is compressed independently with its own Huffman + /// tree. This keeps the decode workspace (2 × block_size = 128 KiB) + /// within the per-core L2 (256 KiB), eliminating the L2 evictions + /// that hurt decode throughput on whole-input processing. + static constexpr std::size_t kBlockSize = 64 * 1024; + /** * @brief Lightweight descriptor for a node's routing bitmap. * @details The packed 64-bit words live in a single shared arena @@ -132,28 +152,63 @@ class PivCoHuffman : public HuffmanBase { Bitmap bits; }; + /// @brief Total uncompressed size across all blocks (set during + /// build/deserialize, read by `uncompressed_size_impl()`). std::size_t uncompressed_size_ = 0; - std::size_t root_ = kNpos; - std::vector nodes_; + /// @brief Per-block serialized size from the wire (fixed; last block may + /// be smaller). Used by decode to iterate blocks. + std::size_t block_size_ = kBlockSize; + /// @brief Serialized compressed stream (header + per-block data). std::vector compressed_; + + // --- per-block tree state (rebuilt for each block during decode) -------- + // Mutable because `decode_impl()` is const but rebuilds the tree + // per-block from the serialized stream. + mutable std::size_t root_ = kNpos; + mutable std::vector nodes_; /// @brief One contiguous allocation backing every internal node's bitmap /// words. Node `i` owns `nodes_[i].bits.word_count()` words starting /// at `arena_[nodes_[i].bits.offset]`. - std::vector arena_; - /// @brief Reusable decode scratch: two ping-pong halves of size - /// `uncompressed_size_`. Allocated once and reused across decode - /// calls; `mutable` because `decode_impl()` is const. + mutable std::vector arena_; + /// @brief Reusable decode/encode scratch: two ping-pong halves of size + /// `block_uncompressed_size_`. Reused across blocks; `mutable` + /// because `decode_impl()` is const. mutable std::vector workspace_; - std::array code_lengths_{}; + /// @brief Uncompressed size of the current block being processed. + mutable std::size_t block_uncompressed_size_ = 0; + mutable std::array code_lengths_{}; // --- construction -------------------------------------------------------- - /** @brief Build the canonical Huffman tree and per-node bitmaps from @p - * input. - */ + /** @brief Build and serialize all blocks from @p input. + * @details Splits the input into fixed-size blocks (kBlockSize), builds + * an independent Huffman tree per block, and serializes each + * block into `compressed_`. The wire format is: + * - total_uncompressed_size (8 bytes) + * - block_size (8 bytes) + * - Per block: block_uncompressed_size (8 bytes) + 128 bytes + * nibbles + per-internal-node bitmaps. */ void build(std::span input) { uncompressed_size_ = input.size(); - if (uncompressed_size_ == 0) { + compressed_.clear(); + write(uncompressed_size_); + write(block_size_); + + std::size_t offset = 0; + while (offset < input.size()) { + std::size_t remaining = input.size() - offset; + std::size_t this_block = std::min(remaining, block_size_); + build_block(input.subspan(offset, this_block)); + serialize_block(); + offset += this_block; + } + } + + /** @brief Build the canonical Huffman tree and per-node bitmaps for one + * block of @p input. */ + void build_block(std::span input) { + block_uncompressed_size_ = input.size(); + if (block_uncompressed_size_ == 0) { root_ = kNpos; return; } @@ -190,7 +245,7 @@ class PivCoHuffman : public HuffmanBase { std::array>, kAlphabet> paths; std::vector> path; assign_paths(root_, path, paths); - workspace_.resize(uncompressed_size_ * 2); + workspace_.resize(block_uncompressed_size_ * 2); std::copy(input.begin(), input.end(), workspace_.data()); encode_partition(root_, 0, 0, paths, freq); } @@ -331,7 +386,7 @@ class PivCoHuffman : public HuffmanBase { * Same-length codes are grouped together, which is a prerequisite * for future flat-subtree optimization. */ void build_canonical_tree( - const std::array& lengths) { + const std::array& lengths) const { nodes_.clear(); nodes_.reserve(2 * kAlphabet); @@ -543,9 +598,10 @@ class PivCoHuffman : public HuffmanBase { return std::vector(); } if (nodes_[root_].is_leaf) { - return std::vector(uncompressed_size_, nodes_[root_].symbol); + return std::vector(block_uncompressed_size_, + nodes_[root_].symbol); } - const std::size_t n = uncompressed_size_; + const std::size_t n = block_uncompressed_size_; workspace_.resize(n * 2); decode_node(root_, n, 0, 0); return std::vector(workspace_.begin(), workspace_.begin() + n); @@ -604,7 +660,7 @@ class PivCoHuffman : public HuffmanBase { /** @brief Pointer to the workspace half that holds depth-@p depth outputs. */ symbol_type* workspace_half(std::size_t depth) const { - return workspace_.data() + (depth % 2) * uncompressed_size_; + return workspace_.data() + (depth % 2) * block_uncompressed_size_; } /** @brief Number of set bits in a node's arena-backed bitmap. */ @@ -620,17 +676,16 @@ class PivCoHuffman : public HuffmanBase { // --- serialization ------------------------------------------------------- - /** @brief Write the in-memory tree into the serialized byte buffer. - * @details Wire format: - * - uncompressed_size (8 bytes) + /** @brief Append one block's serialized form to `compressed_`. + * @details Per-block wire format: + * - block_uncompressed_size (8 bytes) * - 256 code lengths as 128 bytes of 4-bit nibbles (0 = absent) * - For each internal node (pre-order): bits.count (8 bytes) + * packed words. Leaves are implied by the tree structure * reconstructed from lengths. */ - void serialize() { - compressed_.clear(); - write(uncompressed_size_); - if (uncompressed_size_ == 0) { + void serialize_block() { + write(block_uncompressed_size_); + if (block_uncompressed_size_ == 0) { return; } @@ -659,26 +714,42 @@ class PivCoHuffman : public HuffmanBase { } } - /** @brief Rebuild the in-memory tree from a serialized byte buffer. */ + /** @brief Load the serialized stream and parse the block header. + * @details Copies @p data into `compressed_` and reads the top-level + * header (total_uncompressed_size, block_size). Per-block trees + * are rebuilt on demand during `decode_impl()`. */ void deserialize(std::span data) { compressed_.assign(data.begin(), data.end()); nodes_.clear(); arena_.clear(); + root_ = kNpos; if (data.empty()) { uncompressed_size_ = 0; - root_ = kNpos; return; } std::size_t pos = 0; - uncompressed_size_ = read(data, pos); - if (uncompressed_size_ == 0) { - root_ = kNpos; + uncompressed_size_ = read(compressed_, pos); + block_size_ = read(compressed_, pos); + } + + /** @brief Rebuild the tree and arena for one block from `compressed_` at + * @p pos. Advances @p pos past the block's serialized data. + * @details Reads the per-block header (uncompressed_size + nibbles), + * rebuilds the canonical tree, and reads arena words via a + * two-pass wire walk. */ + void deserialize_block(std::size_t& pos) const { + nodes_.clear(); + arena_.clear(); + root_ = kNpos; + + block_uncompressed_size_ = read(compressed_, pos); + if (block_uncompressed_size_ == 0) { return; } // Read 256 code lengths from nibbles. for (std::size_t i = 0; i < kAlphabet; i += 2) { - std::uint8_t byte = read(data, pos); + std::uint8_t byte = read(compressed_, pos); code_lengths_[i] = byte & 0x0F; code_lengths_[i + 1] = (byte >> 4) & 0x0F; } @@ -693,16 +764,15 @@ class PivCoHuffman : public HuffmanBase { // Pass 1: read each internal node's `count` (and skip its word bytes on // the wire), assign contiguous arena offsets, and size the arena exactly. - // `scan_pos` is discarded — pass 2 re-walks the same region to read words. const std::size_t bitmap_start = pos; std::size_t arena_words = 0; - scan_node_counts(root_, data, pos, arena_words); + scan_node_counts(root_, compressed_, pos, arena_words); arena_.assign(arena_words, 0); // Pass 2: rewind to the bitmap section and read each node's packed words // into its arena slot. pos = bitmap_start; - read_node_words(root_, data, pos); + read_node_words(root_, compressed_, pos); } /** @brief Pass 1: read each internal node's `count` from the wire, assign its @@ -712,7 +782,7 @@ class PivCoHuffman : public HuffmanBase { void scan_node_counts(std::size_t idx, std::span data, std::size_t& pos, - std::size_t& arena_words) { + std::size_t& arena_words) const { if (idx == kNpos) { return; } @@ -732,7 +802,7 @@ class PivCoHuffman : public HuffmanBase { * its arena slot (pre-order, matching serialization). */ void read_node_words(std::size_t idx, std::span data, - std::size_t& pos) { + std::size_t& pos) const { if (idx == kNpos) { return; } From dd11415c7233b76dbceb9b2c6e2ecc1b9ce73db8 Mon Sep 17 00:00:00 2001 From: Sanya239 Date: Wed, 22 Jul 2026 14:37:31 +0300 Subject: [PATCH 06/14] Flat trees optimization --- include/pixie/huffman/pivco_huffman.h | 256 +++++++++++++++++++++++--- 1 file changed, 233 insertions(+), 23 deletions(-) diff --git a/include/pixie/huffman/pivco_huffman.h b/include/pixie/huffman/pivco_huffman.h index b3b19ba..d8bfac0 100644 --- a/include/pixie/huffman/pivco_huffman.h +++ b/include/pixie/huffman/pivco_huffman.h @@ -112,6 +112,13 @@ class PivCoHuffman : public HuffmanBase { /// @brief Maximum supported Huffman code length (matches common decoders). static constexpr std::size_t kMaxCodeLength = 15; + /// @brief Maximum depth of a flat subtree (table size = 2^depth ≤ 256). + /// @details A flat subtree replaces `2^d - 1` internal nodes (each with a + /// 1-bit-per-symbol bitmap and a merge pass) with a single node + /// storing `d` bits per symbol and a lookup table. This turns `d` + /// decode merge passes into a single table-lookup pass. + static constexpr std::uint8_t kMaxFlatDepth = 8; + /// @brief Default block size for block-based processing (64 KiB). /// @details Each block is compressed independently with its own Huffman /// tree. This keeps the decode workspace (2 × block_size = 128 KiB) @@ -130,11 +137,14 @@ class PivCoHuffman : public HuffmanBase { */ struct Bitmap { std::size_t offset = 0; - std::size_t count = 0; - - /** @brief Number of 64-bit words backing the bitmap. */ - std::size_t word_count() const { - return (count + kWordBits - 1) / kWordBits; + std::size_t count = 0; ///< Number of symbols (not bits). + + /** @brief Number of 64-bit words backing the bitmap. + * @param flat_depth 0 for normal nodes (1 bit/symbol); >0 for flat + * nodes (`flat_depth` bits/symbol). */ + std::size_t word_count(std::uint8_t flat_depth = 0) const { + std::size_t bits_per = flat_depth > 0 ? flat_depth : 1; + return (count * bits_per + kWordBits - 1) / kWordBits; } }; @@ -150,6 +160,9 @@ class PivCoHuffman : public HuffmanBase { std::uint8_t symbol = 0; bool is_leaf = false; Bitmap bits; + /// @brief 0 = normal node (1 bit/symbol); >0 = flat subtree of this depth + /// (stores `flat_depth` bits/symbol + a lookup table). + std::uint8_t flat_depth = 0; }; /// @brief Total uncompressed size across all blocks (set during @@ -177,6 +190,10 @@ class PivCoHuffman : public HuffmanBase { /// @brief Uncompressed size of the current block being processed. mutable std::size_t block_uncompressed_size_ = 0; mutable std::array code_lengths_{}; + /// @brief Per-node flat lookup tables. `flat_tables_[i]` is valid when + /// `nodes_[i].flat_depth > 0`; it maps each `d`-bit suffix to its + /// symbol. Rebuilt from the tree structure (not serialized). + mutable std::vector> flat_tables_; // --- construction -------------------------------------------------------- @@ -225,6 +242,7 @@ class PivCoHuffman : public HuffmanBase { // 3. Rebuild a canonical tree from the lengths (groups same-length codes). // Structure only: bitmap weights/offsets are assigned below, not here. build_canonical_tree(code_lengths_); + detect_flat_subtrees(); // Single-symbol input: root is a leaf, no internal nodes, no arena. if (root_ == kNpos || nodes_[root_].is_leaf) { @@ -380,6 +398,142 @@ class PivCoHuffman : public HuffmanBase { } } + // --- bit helpers --------------------------------------------------------- + // Small functions for reading/writing individual bits and multi-bit + // values from/to packed `std::uint64_t` word arrays. Used by both the + // normal (1-bit-per-symbol) and flat (d-bits-per-symbol) code paths. + + /** @brief Read a single bit from a packed word array at bit position @p pos. + */ + static bool read_bit(const std::uint64_t* words, std::size_t pos) { + return (words[pos / kWordBits] >> (pos % kWordBits)) & 1ull; + } + + /** @brief Set a single bit in a packed word array at bit position @p pos. + * @details Assumes the word is pre-zeroed; only sets 1-bits. */ + static void set_bit(std::uint64_t* words, std::size_t pos) { + words[pos / kWordBits] |= (1ull << (pos % kWordBits)); + } + + /** @brief Read @p d bits from a packed word array starting at bit position + * @p pos, MSB-first. Returns the decoded value. */ + static std::uint32_t read_bits(const std::uint64_t* words, + std::size_t pos, + std::uint8_t d) { + std::uint32_t value = 0; + for (std::uint8_t i = 0; i < d; i++) { + value = (value << 1) | read_bit(words, pos + i); + } + return value; + } + + /** @brief Write @p d bits of @p value to a packed word array starting at + * bit position @p pos, MSB-first. Only sets 1-bits (pre-zeroed). */ + static void write_bits(std::uint64_t* words, + std::size_t pos, + std::uint32_t value, + std::uint8_t d) { + for (std::uint8_t i = 0; i < d; i++) { + if ((value >> (d - 1 - i)) & 1u) { + set_bit(words, pos + i); + } + } + } + + // --- flat subtree detection ---------------------------------------------- + + /** @brief Check if all leaves in the subtree rooted at @p idx are at the + * same depth from @p idx. + * @param idx Subtree root. + * @param depth Output: the uniform leaf depth (0 for a leaf itself). + * @returns true if uniform, false if leaves are at mixed depths. */ + bool is_uniform_leaf_depth(std::size_t idx, std::uint8_t& depth) const { + const Node& n = nodes_[idx]; + if (n.is_leaf) { + depth = 0; + return true; + } + std::uint8_t ld = 0, rd = 0; + if (!is_uniform_leaf_depth(n.left, ld) || + !is_uniform_leaf_depth(n.right, rd)) { + return false; + } + if (ld != rd) { + return false; + } + depth = ld + 1; + return true; + } + + /** @brief Walk the tree top-down, marking nodes whose subtrees are uniform + * and shallow enough (≤ kMaxFlatDepth) to flatten. + * @details A flat node of depth @p d stores `d` bits per symbol in its + * bitmap and uses a lookup table of size `2^d` to decode. It + * replaces `2^d - 1` internal nodes, eliminating `d - 1` merge + * passes. Top-down marking ensures the largest eligible subtree + * is flattened; children of a flat node are not visited. */ + void detect_flat_subtrees() const { + flat_tables_.clear(); + flat_tables_.resize(nodes_.size()); + mark_flat(root_); + } + + /** @brief Recursively mark a node as flat if eligible, else recurse. */ + void mark_flat(std::size_t idx) const { + if (idx == kNpos) { + return; + } + Node& n = nodes_[idx]; + if (n.is_leaf) { + return; + } + + std::uint8_t depth = 0; + if (is_uniform_leaf_depth(idx, depth) && depth >= 2 && + depth <= kMaxFlatDepth) { + n.flat_depth = depth; + flat_tables_[idx].assign(std::size_t{1} << depth, 0); + build_flat_table(idx, depth, 0, flat_tables_[idx]); + return; // children are covered by this flat node + } + + mark_flat(n.left); + mark_flat(n.right); + } + + /** @brief Populate the flat lookup table by walking the subtree. + * @param idx Current node. + * @param remaining_depth Bits remaining (counts down to 0 at leaves). + * @param code Suffix accumulated so far (MSB-first: left=0, + * right=1). + * @param table Output table of size `2^total_depth`. */ + void build_flat_table(std::size_t idx, + std::uint8_t remaining_depth, + std::uint32_t code, + std::vector& table) const { + const Node& n = nodes_[idx]; + if (n.is_leaf) { + table[code] = n.symbol; + return; + } + build_flat_table(n.left, remaining_depth - 1, code << 1, table); + build_flat_table(n.right, remaining_depth - 1, (code << 1) | 1, table); + } + + /** @brief Total frequency weight of a subtree (sum of leaf frequencies). */ + std::size_t subtree_weight( + std::size_t idx, + const std::array& freq) const { + if (idx == kNpos) { + return 0; + } + const Node& n = nodes_[idx]; + if (n.is_leaf) { + return freq[n.symbol]; + } + return subtree_weight(n.left, freq) + subtree_weight(n.right, freq); + } + /** @brief Rebuild a canonical Huffman tree from code lengths. * @details Assigns canonical codes (sorted by length, then symbol), then * builds the tree top-down by inserting each (symbol, code) pair. @@ -491,13 +645,25 @@ class PivCoHuffman : public HuffmanBase { if (n.is_leaf) { return freq[n.symbol]; } + + if (n.flat_depth > 0) { + // Flat node: compute weight from subtree, assign arena for its + // d-bits-per-symbol bitmap, but DON'T recurse into children (their + // bitmaps aren't stored — the flat node replaces them). + std::size_t weight = subtree_weight(idx, freq); + n.bits.count = weight; + n.bits.offset = arena_words; + arena_words += n.bits.word_count(n.flat_depth); + return weight; + } + std::size_t left_weight = assign_weights_and_offsets(n.left, freq, arena_words); std::size_t right_weight = assign_weights_and_offsets(n.right, freq, arena_words); n.bits.count = left_weight + right_weight; n.bits.offset = arena_words; - arena_words += n.bits.word_count(); + arena_words += n.bits.word_count(0); return n.bits.count; } @@ -554,6 +720,24 @@ class PivCoHuffman : public HuffmanBase { } Node& n = nodes_[idx]; + + // Flat node: write d-bit suffix per symbol, no partition, no recursion. + if (n.flat_depth > 0) { + const std::uint8_t d = n.flat_depth; + std::uint64_t* bits = arena_.data() + n.bits.offset; + const symbol_type* src = workspace_half(depth) + out_base; + for (std::size_t i = 0; i < n.bits.count; i++) { + // Build the d-bit suffix from the symbol's precomputed code path. + std::uint32_t suffix = 0; + for (std::uint8_t b = 0; b < d; b++) { + suffix = (suffix << 1) | (paths[src[i]][depth + b].second ? 1 : 0); + } + write_bits(bits, i * d, suffix, d); + } + return; + } + + // Normal node: 1-bit-per-symbol bitmap + partition into children. const std::size_t left_weight = child_weight(n.left, freq); const std::size_t weight = n.bits.count; @@ -627,14 +811,23 @@ class PivCoHuffman : public HuffmanBase { return; } + // Flat node: read d-bit suffixes, lookup symbols, no merge. + if (n.flat_depth > 0) { + const std::uint8_t d = n.flat_depth; + const std::uint64_t* bits = arena_.data() + n.bits.offset; + const symbol_type* table = flat_tables_[idx].data(); + for (std::size_t i = 0; i < weight; i++) { + dst[i] = table[read_bits(bits, i * d, d)]; + } + return; + } + + // Normal node: decode children, then merge by 1-bit-per-symbol bitmap. const std::size_t right_weight = bitmap_popcount(n.bits); const std::size_t left_weight = n.bits.count - right_weight; decode_node(n.left, left_weight, depth + 1, out_base); decode_node(n.right, right_weight, depth + 1, out_base + left_weight); - // Children wrote to the opposite half: [out_base, out_base + left_weight) - // holds the left stream, [out_base + left_weight, out_base + weight) - // holds the right stream. Merge both into this node's half by the bitmap. const symbol_type* src = workspace_half(depth + 1) + out_base; const symbol_type* left_out = src; const symbol_type* right_out = src + left_weight; @@ -643,7 +836,7 @@ class PivCoHuffman : public HuffmanBase { std::size_t c_left = 0; std::size_t c_right = 0; const std::uint64_t* bits = arena_.data() + n.bits.offset; - const std::size_t words = n.bits.word_count(); + const std::size_t words = n.bits.word_count(0); for (std::size_t w = 0; w < words; w++) { std::uint64_t word = bits[w]; std::size_t limit = kWordBits; @@ -663,11 +856,13 @@ class PivCoHuffman : public HuffmanBase { return workspace_.data() + (depth % 2) * block_uncompressed_size_; } - /** @brief Number of set bits in a node's arena-backed bitmap. */ + /** @brief Number of set bits in a normal (1-bit-per-symbol) node's bitmap. + * @details Only called for non-flat nodes (flat nodes use table lookup, + * not merge). */ std::size_t bitmap_popcount(const Bitmap& b) const { const std::uint64_t* p = arena_.data() + b.offset; std::size_t total = 0; - const std::size_t words = b.word_count(); + const std::size_t words = b.word_count(0); for (std::size_t i = 0; i < words; i++) { total += std::popcount(p[i]); } @@ -700,18 +895,24 @@ class PivCoHuffman : public HuffmanBase { serialize_node(root_); } - /** @brief Serialize a node's bitmap and recurse (pre-order). */ + /** @brief Serialize a node's bitmap and recurse (pre-order). + * @details Flat nodes are terminal: their children's bitmaps are not + * serialized (the flat node replaces them). */ void serialize_node(std::size_t idx) { if (idx == kNpos) { return; } const auto& n = nodes_[idx]; - if (!n.is_leaf) { - write(n.bits.count); - write_words(n.bits.offset, n.bits.word_count()); - serialize_node(n.left); - serialize_node(n.right); + if (n.is_leaf) { + return; + } + write(n.bits.count); + write_words(n.bits.offset, n.bits.word_count(n.flat_depth)); + if (n.flat_depth > 0) { + return; // flat: children not serialized } + serialize_node(n.left); + serialize_node(n.right); } /** @brief Load the serialized stream and parse the block header. @@ -756,6 +957,7 @@ class PivCoHuffman : public HuffmanBase { // Rebuild the canonical tree from lengths. build_canonical_tree(code_lengths_); + detect_flat_subtrees(); // Single-symbol input: root is a leaf, no internal nodes, no arena. if (root_ == kNpos || nodes_[root_].is_leaf) { @@ -778,7 +980,8 @@ class PivCoHuffman : public HuffmanBase { /** @brief Pass 1: read each internal node's `count` from the wire, assign its * arena offset, and advance past its word bytes without storing * them. Sums the total arena word count so the caller allocates once. - * @details Pre-order traversal matches the serialization order. */ + * @details Pre-order traversal matches the serialization order. Flat nodes + * are terminal (children not visited). */ void scan_node_counts(std::size_t idx, std::span data, std::size_t& pos, @@ -792,14 +995,18 @@ class PivCoHuffman : public HuffmanBase { } n.bits.count = read(data, pos); n.bits.offset = arena_words; - arena_words += n.bits.word_count(); - pos += n.bits.word_count() * sizeof(std::uint64_t); + arena_words += n.bits.word_count(n.flat_depth); + pos += n.bits.word_count(n.flat_depth) * sizeof(std::uint64_t); + if (n.flat_depth > 0) { + return; // flat: children not on the wire + } scan_node_counts(n.left, data, pos, arena_words); scan_node_counts(n.right, data, pos, arena_words); } /** @brief Pass 2: read each internal node's packed words from the wire into - * its arena slot (pre-order, matching serialization). */ + * its arena slot (pre-order, matching serialization). Flat nodes are + * terminal. */ void read_node_words(std::size_t idx, std::span data, std::size_t& pos) const { @@ -811,12 +1018,15 @@ class PivCoHuffman : public HuffmanBase { return; } pos += sizeof(std::size_t); // count, already read in pass 1 - const std::size_t words = n.bits.word_count(); + const std::size_t words = n.bits.word_count(n.flat_depth); const std::size_t bytes = words * sizeof(std::uint64_t); if (bytes > 0) { std::memcpy(arena_.data() + n.bits.offset, data.data() + pos, bytes); } pos += bytes; + if (n.flat_depth > 0) { + return; // flat: children not on the wire + } read_node_words(n.left, data, pos); read_node_words(n.right, data, pos); } From 3c9c33262c467e06717d56d9bbc4ff49098f9fe3 Mon Sep 17 00:00:00 2001 From: Sanya239 Date: Sat, 25 Jul 2026 16:32:22 +0300 Subject: [PATCH 07/14] Slow non-canonical trees --- include/pixie/huffman/pivco_huffman.h | 189 ++++++++++++++++++-------- 1 file changed, 132 insertions(+), 57 deletions(-) diff --git a/include/pixie/huffman/pivco_huffman.h b/include/pixie/huffman/pivco_huffman.h index d8bfac0..cedf123 100644 --- a/include/pixie/huffman/pivco_huffman.h +++ b/include/pixie/huffman/pivco_huffman.h @@ -534,21 +534,27 @@ class PivCoHuffman : public HuffmanBase { return subtree_weight(n.left, freq) + subtree_weight(n.right, freq); } - /** @brief Rebuild a canonical Huffman tree from code lengths. - * @details Assigns canonical codes (sorted by length, then symbol), then - * builds the tree top-down by inserting each (symbol, code) pair. - * Same-length codes are grouped together, which is a prerequisite - * for future flat-subtree optimization. */ + /// @brief A symbol with its Huffman code length, used for tree building. + struct SymbolEntry { + std::uint8_t symbol; + std::uint8_t length; + }; + + /** @brief Rebuild a Huffman tree from code lengths, reshaped to maximize + * flat-subtree opportunities. + * @details Instead of building the canonical tree bit-by-bit (which scatters + * same-length symbols across the tree), this builder groups + * same-length symbols into power-of-two flat subtrees, then + * assembles them into a tree whose root-to-leaf path lengths still + * equal the original code lengths. This is the paper's "non- + * canonical subtree" optimization (Section 2.2.4): same average + * code lengths, different (better) tree shape. */ void build_canonical_tree( const std::array& lengths) const { nodes_.clear(); nodes_.reserve(2 * kAlphabet); // Collect symbols present, sorted by (length, symbol). - struct SymbolEntry { - std::uint8_t symbol; - std::uint8_t length; - }; std::vector entries; entries.reserve(kAlphabet); for (std::size_t s = 0; s < kAlphabet; s++) { @@ -564,8 +570,7 @@ class PivCoHuffman : public HuffmanBase { return a.symbol < b.symbol; }); - // Special case: single distinct symbol → root is a leaf (no internal nodes, - // no bitmaps). decode_node returns a constant-filled vector. + // Special case: single distinct symbol → root is a leaf. if (entries.size() == 1) { nodes_.push_back(Node{}); root_ = 0; @@ -574,58 +579,128 @@ class PivCoHuffman : public HuffmanBase { return; } - // Assign canonical codes: first code at each length is 0, increment - // within a length, shift left when length increases. - std::vector> - codes; // (code, length) - codes.reserve(entries.size()); - std::uint64_t code = 0; - std::uint8_t prev_len = 0; + // Build the reshaped tree using the power-of-two grouping strategy. + build_reshaped_tree(entries); + } + + /** @brief Build a non-canonical tree that groups same-length symbols into + * power-of-two flat subtrees. + * @details Uses a recursive slot-filling approach. At each depth, we have + * `slots` open positions. We fill them with leaves (grouped into + * power-of-2 balanced subtrees for flat-node detection) and + * internal nodes (which recurse deeper). The split is determined + * by the Kraft equality: `leaves_at_depth + 2*internal_at_depth = + * slots`, and `internal_at_depth` creates the slots for the next + * depth. */ + /** @brief Build a balanced binary subtree of depth @p d with @p count + * leaves (count must be 2^d). + * @param symbols Array of symbol values (count entries). + * @param count Number of symbols (must be 2^d). + * @param d Depth of the subtree (0 = single leaf). + * @returns Node index of the subtree root. */ + std::size_t build_balanced_subtree(const std::uint8_t* symbols, + std::size_t count, + std::uint8_t d) const { + if (d == 0) { + nodes_.push_back(Node{}); + std::size_t idx = nodes_.size() - 1; + nodes_[idx].is_leaf = true; + nodes_[idx].symbol = symbols[0]; + return idx; + } + std::size_t half = count / 2; + std::size_t left = build_balanced_subtree(symbols, half, d - 1); + std::size_t right = build_balanced_subtree(symbols + half, half, d - 1); + nodes_.push_back(Node{}); + std::size_t idx = nodes_.size() - 1; + nodes_[idx].left = left; + nodes_[idx].right = right; + return idx; + } + + void build_reshaped_tree(const std::vector& entries) const { + std::array, kMaxCodeLength + 1> by_length; for (const auto& e : entries) { - if (prev_len == 0) { - prev_len = e.length; - } - while (prev_len < e.length) { - code <<= 1; - prev_len++; - } - codes.push_back({code, e.length}); - code++; + by_length[e.length].push_back(e.symbol); } - // Build tree by inserting each (symbol, code) pair via bit-by-bit descent. - // Create root. - nodes_.push_back(Node{}); - root_ = 0; - - for (std::size_t i = 0; i < entries.size(); i++) { - std::uint64_t c = codes[i].first; - std::uint8_t len = codes[i].second; - std::size_t cur = root_; - - // Descend len-1 bits, creating internal nodes as needed. - for (std::size_t bit = 0; bit + 1 < len; bit++) { - // Read MSB first (canonical codes are assigned MSB-first). - bool go_right = (c >> (len - 1 - bit)) & 1; - std::size_t& child = go_right ? nodes_[cur].right : nodes_[cur].left; - if (child == kNpos) { - nodes_.push_back(Node{}); - child = nodes_.size() - 1; - } - cur = child; + // Per-depth consumption cursor. + std::array cursor{}; + root_ = build_slots(by_length, cursor, 1, 2); + } + + /** @brief Fill `slots` open positions at `depth` with leaves and internal + * nodes, returning the root of the resulting balanced subtree. + * @param by_length Symbols grouped by code length. + * @param cursor Per-depth consumption cursor (how many symbols used). + * @param depth Current tree depth (1 = children of root). + * @param slots Number of open positions (must be a power of 2). */ + std::size_t build_slots( + std::array, kMaxCodeLength + 1>& by_length, + std::array& cursor, + std::uint8_t depth, + std::size_t slots) const { + if (slots == 0) { + return kNpos; + } + + // How many leaves to place at this depth. + std::size_t available = by_length[depth].size() - cursor[depth]; + std::size_t leaves_here = std::min(available, slots); + std::size_t internal = slots - leaves_here; + + // Build children left-to-right: first leaves (grouped), then internal. + std::vector children; + + // Group leaves into power-of-2 balanced subtrees. + std::size_t leaf_off = cursor[depth]; + std::size_t remaining = leaves_here; + while (remaining > 0) { + std::uint8_t d = 0; + std::size_t subset = 1; + while (subset * 2 <= remaining && d < kMaxFlatDepth) { + subset *= 2; + d++; } - // Last bit → leaf. - bool go_right = (c >> 0) & 1; - std::size_t& child = go_right ? nodes_[cur].right : nodes_[cur].left; - nodes_.push_back(Node{}); - child = nodes_.size() - 1; - nodes_[child].is_leaf = true; - nodes_[child].symbol = entries[i].symbol; + children.push_back(build_balanced_subtree( + by_length[depth].data() + leaf_off, subset, d)); + leaf_off += subset; + remaining -= subset; } + cursor[depth] += leaves_here; - // Bitmap weights and arena offsets are assigned by the caller: `build()` - // derives them from symbol frequencies (exact), `deserialize()` reads - // them from the wire. This keeps the tree builder free of bitmap storage. + // Internal nodes: each recurses with 2 slots at depth+1. + for (std::size_t i = 0; i < internal; i++) { + children.push_back(build_slots(by_length, cursor, depth + 1, 2)); + } + + // Assemble children into a balanced binary tree. + return assemble_balanced(children); + } + + /** @brief Assemble a list of subtree roots into a balanced binary tree. + * @details Recursively pairs children left-right, creating internal nodes. + */ + std::size_t assemble_balanced(std::vector& children) const { + if (children.empty()) { + return kNpos; + } + if (children.size() == 1) { + return children[0]; + } + // Pair up: left half vs right half. + std::size_t mid = children.size() / 2; + std::vector left_children(children.begin(), + children.begin() + mid); + std::vector right_children(children.begin() + mid, + children.end()); + std::size_t left = assemble_balanced(left_children); + std::size_t right = assemble_balanced(right_children); + nodes_.push_back(Node{}); + std::size_t idx = nodes_.size() - 1; + nodes_[idx].left = left; + nodes_[idx].right = right; + return idx; } /** @brief Post-order assignment of exact per-node weights and arena offsets. From 72bc9fa37c2e54a2220936979c67de2341aae884 Mon Sep 17 00:00:00 2001 From: Sanya239 Date: Sun, 26 Jul 2026 15:10:41 +0300 Subject: [PATCH 08/14] Faster non-canonical trees + a bunch of refactoring --- include/pixie/huffman/pivco_huffman.h | 594 ++++++++++++++------------ 1 file changed, 320 insertions(+), 274 deletions(-) diff --git a/include/pixie/huffman/pivco_huffman.h b/include/pixie/huffman/pivco_huffman.h index cedf123..4ccea14 100644 --- a/include/pixie/huffman/pivco_huffman.h +++ b/include/pixie/huffman/pivco_huffman.h @@ -10,16 +10,15 @@ * path. Decoding merges child symbol streams bottom-up from leaves to root, * using the node bitmap as a selector. * - * The tree uses canonical Huffman code lengths (<= 15 bits, enforced via the - * package-merge length-limited algorithm): the wire format stores only the 256 + * The tree uses canonical Huffman code lengths (<= 15 bits, enforced with a + * zlib-style overflow correction): the wire format stores only the 256 * per-symbol lengths (128 bytes as 4-bit nibbles), and the tree structure is - * reconstructed deterministically at load time. This groups same-length codes - * together, which is a prerequisite for future flat-subtree optimizations and - * reduces serialized header size. + * reconstructed deterministically at load time. Same-length symbols are + * reshaped into flat subtrees as described by the PivCo-Huffman paper. * - * This is a deliberately simple, unoptimized reference implementation: node - * bitmaps are stored as packed `std::uint64_t` words, traversal is scalar, - * and there are no flat-subtree, SIMD, or selective-ANS optimizations. + * This implementation uses packed `std::uint64_t` bitmaps, flat subtrees, + * contiguous arenas, and scalar bottom-up decoding. SIMD and selective ANS + * are intentionally left to specialized future implementations. * * Range and ownership conventions follow `HuffmanBase`: symbols are bytes, * the compressed stream is a byte view, and the codec owns its serialized form. @@ -34,7 +33,6 @@ #include #include #include -#include #include namespace pixie { @@ -119,6 +117,9 @@ class PivCoHuffman : public HuffmanBase { /// decode merge passes into a single table-lookup pass. static constexpr std::uint8_t kMaxFlatDepth = 8; + /// @brief Maximum nodes in a full binary tree over the byte alphabet. + static constexpr std::size_t kMaxTreeNodes = 2 * kAlphabet - 1; + /// @brief Default block size for block-based processing (64 KiB). /// @details Each block is compressed independently with its own Huffman /// tree. This keeps the decode workspace (2 × block_size = 128 KiB) @@ -163,6 +164,23 @@ class PivCoHuffman : public HuffmanBase { /// @brief 0 = normal node (1 bit/symbol); >0 = flat subtree of this depth /// (stores `flat_depth` bits/symbol + a lookup table). std::uint8_t flat_depth = 0; + /// @brief Offset of this flat node's decode table in `flat_tables_`. + std::size_t flat_table_offset = 0; + }; + + /// @brief A symbol with its frequency weight, used for Huffman construction. + struct Leaf { + std::size_t weight; + std::uint8_t symbol; + }; + + /** @brief Compact root-to-leaf code used by the scalar encoder. + * @details Codes are stored MSB-first in the low `length` bits. The + * 15-bit length limit keeps every code in one 16-bit value and + * avoids the per-symbol dynamic path vectors used previously. */ + struct SymbolCode { + std::uint16_t bits = 0; + std::uint8_t length = 0; }; /// @brief Total uncompressed size across all blocks (set during @@ -190,10 +208,8 @@ class PivCoHuffman : public HuffmanBase { /// @brief Uncompressed size of the current block being processed. mutable std::size_t block_uncompressed_size_ = 0; mutable std::array code_lengths_{}; - /// @brief Per-node flat lookup tables. `flat_tables_[i]` is valid when - /// `nodes_[i].flat_depth > 0`; it maps each `d`-bit suffix to its - /// symbol. Rebuilt from the tree structure (not serialized). - mutable std::vector> flat_tables_; + /// @brief Contiguous storage for all flat-subtree decode tables. + mutable std::vector flat_tables_; // --- construction -------------------------------------------------------- @@ -230,18 +246,15 @@ class PivCoHuffman : public HuffmanBase { return; } - // 1. Count symbol frequencies. - std::array freq{}; - for (auto s : input) { - freq[s]++; - } + // 1. Count symbol frequencies using independent partial histograms. + const std::array freq = count_frequencies(input); // 2. Build a Huffman tree to determine code lengths, then extract lengths. compute_code_lengths(freq, code_lengths_); // 3. Rebuild a canonical tree from the lengths (groups same-length codes). // Structure only: bitmap weights/offsets are assigned below, not here. - build_canonical_tree(code_lengths_); + build_noncanonical_tree(code_lengths_); detect_flat_subtrees(); // Single-symbol input: root is a leaf, no internal nodes, no arena. @@ -255,146 +268,187 @@ class PivCoHuffman : public HuffmanBase { assign_weights_and_offsets(root_, freq, arena_words); arena_.assign(arena_words, 0); - // 5. Assign per-symbol code paths, then fill node bitmaps via top-down + // 5. Assign compact per-symbol codes, then fill node bitmaps via top-down // in-place partition into the reusable workspace. This mirrors decode: // the root reads its input from workspace half 0, writes its bitmap, // and partitions symbols into half 1 for its children. Each level // ping-pongs between the two halves — no per-node allocations. - std::array>, kAlphabet> paths; - std::vector> path; - assign_paths(root_, path, paths); + std::array codes{}; + assign_codes(root_, 0, 0, codes); workspace_.resize(block_uncompressed_size_ * 2); std::copy(input.begin(), input.end(), workspace_.data()); - encode_partition(root_, 0, 0, paths, freq); + encode_partition(root_, 0, 0, codes, freq); } - /** @brief Compute length-limited Huffman code lengths via package-merge. - * @details Implements the Larmore-Hirschberg package-merge algorithm to - * produce optimal code lengths subject to `length <= - * kMaxCodeLength`. This guarantees the 4-bit nibble header never truncates a - * length, and keeps all downstream canonical-tree, flat-subtree, and - * reshaping optimizations valid. - * - * Algorithm: maintain a sorted list of "items" per level (1..L). A level-1 - * item is a single coin (one symbol). At each subsequent level, form - * "packages" by pairing consecutive items from the previous level (a - * package commits to selecting both its constituents), then merge with the - * original coins and keep the cheapest `2n-2` items. After L-1 iterations, - * summing each symbol's count across the final `2n-2` items yields its code - * length (each ≤ L, and the Kraft inequality holds exactly). + /** @brief Count a block's byte frequencies with independent accumulators. + * @details Eight partial histograms break the dependency chain caused by + * repeatedly incrementing a hot symbol. The block is at most + * 64 KiB, so 32-bit partial counters cannot overflow. */ + static std::array count_frequencies( + std::span input) { + constexpr std::size_t kLanes = 8; + std::array, kLanes> partial{}; + std::size_t i = 0; + for (; i + kLanes <= input.size(); i += kLanes) { + partial[0][input[i]]++; + partial[1][input[i + 1]]++; + partial[2][input[i + 2]]++; + partial[3][input[i + 3]]++; + partial[4][input[i + 4]]++; + partial[5][input[i + 5]]++; + partial[6][input[i + 6]]++; + partial[7][input[i + 7]]++; + } + for (; i < input.size(); ++i) { + partial[0][input[i]]++; + } + + std::array frequencies{}; + for (std::size_t symbol = 0; symbol < kAlphabet; ++symbol) { + for (std::size_t lane = 0; lane < kLanes; ++lane) { + frequencies[symbol] += partial[lane][symbol]; + } + } + return frequencies; + } + + /** @brief Compute Huffman code lengths, limited to ≤ kMaxCodeLength. + * @details Two-phase approach (same as zlib/zstd): + * 1. Sort present symbols by frequency and build a standard + * Huffman tree with the linear two-queue algorithm, then + * extract per-leaf depths by walking parent pointers. + * 2. If any depth exceeds kMaxCodeLength, apply the zlib-style + * "overflow fixup": collapse all lengths > L into L, then + * rebalance the Kraft sum by borrowing from shorter codes. * - * Each item carries a per-symbol count vector (`Counts`) so that summing at - * the end is a flat array reduction; with `n <= 256` and `L = 15` the total - * work is ~2M additions (< 1 ms), negligible relative to the bitmap fill. */ + * This is O(n log n) for sorting and O(n + L) after that, with fixed-size + * storage for the byte alphabet. */ void compute_code_lengths(const std::array& freq, std::array& lengths) { lengths.fill(0); - // Collect present symbols, sorted by weight ascending (the coin list S, - // reused at every level of the package-merge). - struct Coin { - std::uint8_t symbol; - std::size_t weight; - }; - std::vector coins; - coins.reserve(kAlphabet); + std::array leaves{}; + std::size_t leaf_count = 0; for (std::size_t s = 0; s < kAlphabet; s++) { if (freq[s] > 0) { - coins.push_back({static_cast(s), freq[s]}); + leaves[leaf_count++] = {freq[s], static_cast(s)}; } } - std::sort(coins.begin(), coins.end(), - [](const Coin& a, const Coin& b) { return a.weight < b.weight; }); - - const std::size_t n = coins.size(); - if (n == 0) { + if (leaf_count == 0) { return; } - // Single distinct symbol: must still get a 1-bit code to be decodable. - if (n == 1) { - lengths[coins[0].symbol] = 1; + if (leaf_count == 1) { + lengths[leaves[0].symbol] = 1; return; } - // Per-item symbol-count vector. uint16 suffices: a symbol's count within a - // single package can grow across levels, but final lengths are <= L = 15. - using Counts = std::array; - struct Item { - std::size_t weight; - Counts counts; + std::sort(leaves.begin(), leaves.begin() + leaf_count, + [](const Leaf& a, const Leaf& b) { + if (a.weight != b.weight) { + return a.weight < b.weight; + } + return a.symbol < b.symbol; + }); + + // Build the Huffman tree with the linear two-queue algorithm. Sorted + // leaves form one queue; newly-created internal nodes form the other. + // Internal weights are produced in nondecreasing order, so selecting the + // smaller queue front replaces all binary-heap maintenance. + struct TmpNode { + std::size_t weight = 0; + std::size_t parent = kNpos; + }; + std::array tree{}; + for (std::size_t i = 0; i < leaf_count; ++i) { + tree[i].weight = leaves[i].weight; + } + + std::size_t leaf_head = 0; + std::size_t internal_head = leaf_count; + std::size_t next_internal = leaf_count; + auto take_smallest = [&]() { + const bool take_leaf = + leaf_head < leaf_count && + (internal_head == next_internal || + tree[leaf_head].weight <= tree[internal_head].weight); + return take_leaf ? leaf_head++ : internal_head++; }; - // Level-1 list: one coin per symbol. - std::vector cur; - cur.reserve(2 * n); - for (const auto& c : coins) { - Item it{}; - it.weight = c.weight; - it.counts[c.symbol] = 1; - cur.push_back(std::move(it)); - } - - // Keep the cheapest 2n-2 items at each level (the final selection size; - // cheaper items at a level can never displace a more expensive one in an - // optimal solution, so pruning to 2n-2 is safe). - const std::size_t keep = 2 * n - 2; - - // Iterate levels 2..L. Coins are reused at every level via the merge. - for (std::size_t level = 1; level < kMaxCodeLength; level++) { - // Package: pair consecutive items of `cur` (already weight-sorted). - std::vector packages; - packages.reserve(cur.size() / 2); - for (std::size_t i = 0; i + 1 < cur.size(); i += 2) { - Item pkg{}; - pkg.weight = cur[i].weight + cur[i + 1].weight; - for (std::size_t j = 0; j < kAlphabet; j++) { - pkg.counts[j] = cur[i].counts[j] + cur[i + 1].counts[j]; - } - packages.push_back(std::move(pkg)); + const std::size_t node_count = 2 * leaf_count - 1; + while (next_internal < node_count) { + const std::size_t left = take_smallest(); + const std::size_t right = take_smallest(); + tree[next_internal].weight = tree[left].weight + tree[right].weight; + tree[left].parent = next_internal; + tree[right].parent = next_internal; + ++next_internal; + } + + // Extract code lengths by walking from each leaf to root. + std::uint8_t max_len = 0; + for (std::size_t i = 0; i < leaf_count; ++i) { + std::uint8_t depth = 0; + std::size_t cur = i; + while (tree[cur].parent != kNpos) { + cur = tree[cur].parent; + depth++; } - // Merge coins (sorted) with packages (sorted — pairing preserves order), - // keeping the cheapest `keep` items. Ties prefer coins (deterministic). - // At early levels the combined list may hold fewer than `keep` items; - // keep all of them in that case (the list grows toward `keep` over - // levels as packages proliferate). - const std::size_t available = n + packages.size(); - const std::size_t keep_here = std::min(keep, available); - std::vector next; - next.reserve(keep_here); - std::size_t ci = 0; - std::size_t pi = 0; - while (next.size() < keep_here) { - bool use_coin; - if (ci >= n) { - use_coin = false; - } else if (pi >= packages.size()) { - use_coin = true; - } else { - use_coin = coins[ci].weight <= packages[pi].weight; - } - if (use_coin) { - Item it{}; - it.weight = coins[ci].weight; - it.counts[coins[ci].symbol] = 1; - next.push_back(std::move(it)); - ci++; - } else { - next.push_back(std::move(packages[pi])); - pi++; - } + lengths[leaves[i].symbol] = depth; + if (depth > max_len) { + max_len = depth; + } + } + + if (max_len <= kMaxCodeLength) { + return; + } + + // Length limiting via zlib-style overflow fixup. + limit_code_lengths(lengths, + std::span(leaves.data(), leaf_count)); + } + + /** @brief Limit Huffman code lengths to ≤ kMaxCodeLength. + * @details Collapses all lengths > L into L, rebalances the length + * histogram, then assigns its longest codes to the least frequent + * symbols. */ + void limit_code_lengths(std::array& lengths, + std::span leaves) { + // Count codes at each length and track how many exceeded the limit. This + // is the same overflow definition used by zlib's gen_bitlen(). + std::array bl_count{}; + int overflow = 0; + for (const Leaf& leaf : leaves) { + std::uint8_t len = lengths[leaf.symbol]; + if (len > kMaxCodeLength) { + len = kMaxCodeLength; + ++overflow; } - cur = std::move(next); + ++bl_count[len]; } - // Sum symbol counts across the final 2n-2 selected items -> code lengths. - Counts total{}; - for (const auto& it : cur) { - for (std::size_t j = 0; j < kAlphabet; j++) { - total[j] += it.counts[j]; + // Move one leaf down from the deepest available shorter level, creating + // two children there, and remove one overlong code from level L. + while (overflow > 0) { + int bits = static_cast(kMaxCodeLength) - 1; + while (bits > 0 && bl_count[bits] == 0) { + --bits; } + --bl_count[bits]; + bl_count[bits + 1] += 2; + --bl_count[kMaxCodeLength]; + overflow -= 2; } - for (const auto& c : coins) { - lengths[c.symbol] = static_cast(total[c.symbol]); + + // Leaves arrive sorted by ascending frequency. Assign the longest codes + // first so the fixup preserves the minimum weighted path length for its + // corrected length histogram. + lengths.fill(0); + std::size_t leaf_index = 0; + for (int len = static_cast(kMaxCodeLength); len >= 1; --len) { + for (int count = 0; count < bl_count[len]; ++count) { + lengths[leaves[leaf_index++].symbol] = static_cast(len); + } } } @@ -403,16 +457,14 @@ class PivCoHuffman : public HuffmanBase { // values from/to packed `std::uint64_t` word arrays. Used by both the // normal (1-bit-per-symbol) and flat (d-bits-per-symbol) code paths. - /** @brief Read a single bit from a packed word array at bit position @p pos. - */ - static bool read_bit(const std::uint64_t* words, std::size_t pos) { - return (words[pos / kWordBits] >> (pos % kWordBits)) & 1ull; - } - - /** @brief Set a single bit in a packed word array at bit position @p pos. - * @details Assumes the word is pre-zeroed; only sets 1-bits. */ - static void set_bit(std::uint64_t* words, std::size_t pos) { - words[pos / kWordBits] |= (1ull << (pos % kWordBits)); + /** @brief Reverse the low @p width bits of an at-most-eight-bit value. */ + static std::uint8_t reverse_low_bits(std::uint8_t value, std::uint8_t width) { + value = static_cast(((value & 0x55u) << 1) | + ((value & 0xAAu) >> 1)); + value = static_cast(((value & 0x33u) << 2) | + ((value & 0xCCu) >> 2)); + value = static_cast((value << 4) | (value >> 4)); + return value >> (8 - width); } /** @brief Read @p d bits from a packed word array starting at bit position @@ -420,11 +472,14 @@ class PivCoHuffman : public HuffmanBase { static std::uint32_t read_bits(const std::uint64_t* words, std::size_t pos, std::uint8_t d) { - std::uint32_t value = 0; - for (std::uint8_t i = 0; i < d; i++) { - value = (value << 1) | read_bit(words, pos + i); - } - return value; + const std::size_t word_index = pos / kWordBits; + const std::size_t shift = pos % kWordBits; + std::uint64_t packed = words[word_index] >> shift; + if (shift + d > kWordBits) { + packed |= words[word_index + 1] << (kWordBits - shift); + } + const std::uint8_t mask = static_cast((1u << d) - 1); + return reverse_low_bits(static_cast(packed) & mask, d); } /** @brief Write @p d bits of @p value to a packed word array starting at @@ -433,91 +488,93 @@ class PivCoHuffman : public HuffmanBase { std::size_t pos, std::uint32_t value, std::uint8_t d) { - for (std::uint8_t i = 0; i < d; i++) { - if ((value >> (d - 1 - i)) & 1u) { - set_bit(words, pos + i); - } + const std::uint64_t packed = + reverse_low_bits(static_cast(value), d); + const std::size_t word_index = pos / kWordBits; + const std::size_t shift = pos % kWordBits; + words[word_index] |= packed << shift; + if (shift + d > kWordBits) { + words[word_index + 1] |= packed >> (kWordBits - shift); } } // --- flat subtree detection ---------------------------------------------- - /** @brief Check if all leaves in the subtree rooted at @p idx are at the - * same depth from @p idx. - * @param idx Subtree root. - * @param depth Output: the uniform leaf depth (0 for a leaf itself). - * @returns true if uniform, false if leaves are at mixed depths. */ - bool is_uniform_leaf_depth(std::size_t idx, std::uint8_t& depth) const { - const Node& n = nodes_[idx]; - if (n.is_leaf) { - depth = 0; - return true; - } - std::uint8_t ld = 0, rd = 0; - if (!is_uniform_leaf_depth(n.left, ld) || - !is_uniform_leaf_depth(n.right, rd)) { - return false; - } - if (ld != rd) { - return false; - } - depth = ld + 1; - return true; - } + /** @brief Sentinel value meaning "subtree has non-uniform leaf depths". */ + static constexpr std::uint8_t kNotUniform = 0xFF; - /** @brief Walk the tree top-down, marking nodes whose subtrees are uniform - * and shallow enough (≤ kMaxFlatDepth) to flatten. - * @details A flat node of depth @p d stores `d` bits per symbol in its - * bitmap and uses a lookup table of size `2^d` to decode. It - * replaces `2^d - 1` internal nodes, eliminating `d - 1` merge - * passes. Top-down marking ensures the largest eligible subtree - * is flattened; children of a flat node are not visited. */ + /** @brief Detect maximal flat subtrees in two linear tree passes. + * @details The post-order pass computes uniform leaf depths without + * allocating per node. A top-down pass then selects only maximal + * eligible subtrees and appends their tables to one contiguous + * arena. This follows the paper's flat-subtree optimization while + * avoiding temporary child tables that a bottom-up marker would + * immediately discard when their parent is also flat. */ void detect_flat_subtrees() const { flat_tables_.clear(); - flat_tables_.resize(nodes_.size()); - mark_flat(root_); + if (root_ != kNpos) { + std::array uniform_depths{}; + compute_uniform_depths(root_, uniform_depths); + mark_flat_topdown(root_, uniform_depths); + } } - /** @brief Recursively mark a node as flat if eligible, else recurse. */ - void mark_flat(std::size_t idx) const { - if (idx == kNpos) { - return; + /** @brief Record each subtree's uniform leaf depth in post-order. */ + std::uint8_t compute_uniform_depths( + std::size_t idx, + std::array& depths) const { + const Node& n = nodes_[idx]; + if (n.is_leaf) { + return depths[idx] = 0; + } + + const std::uint8_t left_depth = compute_uniform_depths(n.left, depths); + const std::uint8_t right_depth = compute_uniform_depths(n.right, depths); + + if (left_depth == kNotUniform || right_depth == kNotUniform || + left_depth != right_depth) { + return depths[idx] = kNotUniform; } + return depths[idx] = static_cast(left_depth + 1); + } + + /** @brief Mark maximal flat subtrees and append their lookup tables. */ + void mark_flat_topdown( + std::size_t idx, + const std::array& depths) const { Node& n = nodes_[idx]; if (n.is_leaf) { return; } - std::uint8_t depth = 0; - if (is_uniform_leaf_depth(idx, depth) && depth >= 2 && - depth <= kMaxFlatDepth) { + const std::uint8_t depth = depths[idx]; + if (depth >= 2 && depth <= kMaxFlatDepth) { n.flat_depth = depth; - flat_tables_[idx].assign(std::size_t{1} << depth, 0); - build_flat_table(idx, depth, 0, flat_tables_[idx]); - return; // children are covered by this flat node + n.flat_table_offset = flat_tables_.size(); + flat_tables_.resize(flat_tables_.size() + (std::size_t{1} << depth)); + build_flat_table(idx, 0, flat_tables_.data() + n.flat_table_offset); + return; } - mark_flat(n.left); - mark_flat(n.right); + mark_flat_topdown(n.left, depths); + mark_flat_topdown(n.right, depths); } /** @brief Populate the flat lookup table by walking the subtree. * @param idx Current node. - * @param remaining_depth Bits remaining (counts down to 0 at leaves). * @param code Suffix accumulated so far (MSB-first: left=0, * right=1). * @param table Output table of size `2^total_depth`. */ void build_flat_table(std::size_t idx, - std::uint8_t remaining_depth, std::uint32_t code, - std::vector& table) const { + symbol_type* table) const { const Node& n = nodes_[idx]; if (n.is_leaf) { table[code] = n.symbol; return; } - build_flat_table(n.left, remaining_depth - 1, code << 1, table); - build_flat_table(n.right, remaining_depth - 1, (code << 1) | 1, table); + build_flat_table(n.left, code << 1, table); + build_flat_table(n.right, (code << 1) | 1, table); } /** @brief Total frequency weight of a subtree (sum of leaf frequencies). */ @@ -534,11 +591,9 @@ class PivCoHuffman : public HuffmanBase { return subtree_weight(n.left, freq) + subtree_weight(n.right, freq); } - /// @brief A symbol with its Huffman code length, used for tree building. - struct SymbolEntry { - std::uint8_t symbol; - std::uint8_t length; - }; + using SymbolsByLength = + std::array, kMaxCodeLength + 1>; + using LengthCounts = std::array; /** @brief Rebuild a Huffman tree from code lengths, reshaped to maximize * flat-subtree opportunities. @@ -549,38 +604,37 @@ class PivCoHuffman : public HuffmanBase { * equal the original code lengths. This is the paper's "non- * canonical subtree" optimization (Section 2.2.4): same average * code lengths, different (better) tree shape. */ - void build_canonical_tree( + void build_noncanonical_tree( const std::array& lengths) const { nodes_.clear(); nodes_.reserve(2 * kAlphabet); - // Collect symbols present, sorted by (length, symbol). - std::vector entries; - entries.reserve(kAlphabet); - for (std::size_t s = 0; s < kAlphabet; s++) { - if (lengths[s] > 0) { - entries.push_back({static_cast(s), lengths[s]}); + SymbolsByLength by_length{}; + LengthCounts length_counts{}; + std::size_t symbol_count = 0; + std::uint8_t only_symbol = 0; + for (std::size_t symbol = 0; symbol < kAlphabet; ++symbol) { + const std::uint8_t length = lengths[symbol]; + if (length == 0) { + continue; } + by_length[length][length_counts[length]++] = + static_cast(symbol); + only_symbol = static_cast(symbol); + ++symbol_count; } - std::sort(entries.begin(), entries.end(), - [](const SymbolEntry& a, const SymbolEntry& b) { - if (a.length != b.length) { - return a.length < b.length; - } - return a.symbol < b.symbol; - }); // Special case: single distinct symbol → root is a leaf. - if (entries.size() == 1) { + if (symbol_count == 1) { nodes_.push_back(Node{}); root_ = 0; nodes_[0].is_leaf = true; - nodes_[0].symbol = entries[0].symbol; + nodes_[0].symbol = only_symbol; return; } // Build the reshaped tree using the power-of-two grouping strategy. - build_reshaped_tree(entries); + build_reshaped_tree(by_length, length_counts); } /** @brief Build a non-canonical tree that groups same-length symbols into @@ -618,15 +672,11 @@ class PivCoHuffman : public HuffmanBase { return idx; } - void build_reshaped_tree(const std::vector& entries) const { - std::array, kMaxCodeLength + 1> by_length; - for (const auto& e : entries) { - by_length[e.length].push_back(e.symbol); - } - + void build_reshaped_tree(const SymbolsByLength& by_length, + const LengthCounts& length_counts) const { // Per-depth consumption cursor. - std::array cursor{}; - root_ = build_slots(by_length, cursor, 1, 2); + LengthCounts cursor{}; + root_ = build_slots(by_length, length_counts, cursor, 1, 2); } /** @brief Fill `slots` open positions at `depth` with leaves and internal @@ -635,22 +685,23 @@ class PivCoHuffman : public HuffmanBase { * @param cursor Per-depth consumption cursor (how many symbols used). * @param depth Current tree depth (1 = children of root). * @param slots Number of open positions (must be a power of 2). */ - std::size_t build_slots( - std::array, kMaxCodeLength + 1>& by_length, - std::array& cursor, - std::uint8_t depth, - std::size_t slots) const { + std::size_t build_slots(const SymbolsByLength& by_length, + const LengthCounts& length_counts, + LengthCounts& cursor, + std::uint8_t depth, + std::size_t slots) const { if (slots == 0) { return kNpos; } // How many leaves to place at this depth. - std::size_t available = by_length[depth].size() - cursor[depth]; + std::size_t available = length_counts[depth] - cursor[depth]; std::size_t leaves_here = std::min(available, slots); std::size_t internal = slots - leaves_here; // Build children left-to-right: first leaves (grouped), then internal. - std::vector children; + std::array children{}; + std::size_t child_count = 0; // Group leaves into power-of-2 balanced subtrees. std::size_t leaf_off = cursor[depth]; @@ -662,8 +713,8 @@ class PivCoHuffman : public HuffmanBase { subset *= 2; d++; } - children.push_back(build_balanced_subtree( - by_length[depth].data() + leaf_off, subset, d)); + children[child_count++] = + build_balanced_subtree(by_length[depth].data() + leaf_off, subset, d); leaf_off += subset; remaining -= subset; } @@ -671,17 +722,19 @@ class PivCoHuffman : public HuffmanBase { // Internal nodes: each recurses with 2 slots at depth+1. for (std::size_t i = 0; i < internal; i++) { - children.push_back(build_slots(by_length, cursor, depth + 1, 2)); + children[child_count++] = + build_slots(by_length, length_counts, cursor, depth + 1, 2); } // Assemble children into a balanced binary tree. - return assemble_balanced(children); + return assemble_balanced( + std::span(children.data(), child_count)); } /** @brief Assemble a list of subtree roots into a balanced binary tree. * @details Recursively pairs children left-right, creating internal nodes. */ - std::size_t assemble_balanced(std::vector& children) const { + std::size_t assemble_balanced(std::span children) const { if (children.empty()) { return kNpos; } @@ -690,12 +743,8 @@ class PivCoHuffman : public HuffmanBase { } // Pair up: left half vs right half. std::size_t mid = children.size() / 2; - std::vector left_children(children.begin(), - children.begin() + mid); - std::vector right_children(children.begin() + mid, - children.end()); - std::size_t left = assemble_balanced(left_children); - std::size_t right = assemble_balanced(right_children); + std::size_t left = assemble_balanced(children.first(mid)); + std::size_t right = assemble_balanced(children.subspan(mid)); nodes_.push_back(Node{}); std::size_t idx = nodes_.size() - 1; nodes_[idx].left = left; @@ -752,21 +801,19 @@ class PivCoHuffman : public HuffmanBase { return n.is_leaf ? freq[n.symbol] : n.bits.count; } - /** @brief Depth-first assignment of root-to-leaf paths. */ - void assign_paths( - std::size_t idx, - std::vector>& path, - std::array>, kAlphabet>& paths) { + /** @brief Assign compact MSB-first codes with a depth-first tree walk. */ + void assign_codes(std::size_t idx, + std::uint16_t bits, + std::uint8_t length, + std::array& codes) const { if (nodes_[idx].is_leaf) { - paths[nodes_[idx].symbol] = path; + codes[nodes_[idx].symbol] = {bits, length}; return; } - path.emplace_back(idx, false); - assign_paths(nodes_[idx].left, path, paths); - path.pop_back(); - path.emplace_back(idx, true); - assign_paths(nodes_[idx].right, path, paths); - path.pop_back(); + assign_codes(nodes_[idx].left, static_cast(bits << 1), + length + 1, codes); + assign_codes(nodes_[idx].right, static_cast((bits << 1) | 1), + length + 1, codes); } /** @brief Top-down in-place partition encoding into the reusable workspace. @@ -781,15 +828,13 @@ class PivCoHuffman : public HuffmanBase { * @param depth Current tree depth (selects workspace half + path index). * @param out_base Offset within the current half where this node's stream * begins. - * @param paths Per-symbol code paths from `assign_paths`. + * @param codes Compact per-symbol codes from `assign_codes`. * @param freq Symbol frequencies (sizing child partitions). */ - void encode_partition( - std::size_t idx, - std::size_t depth, - std::size_t out_base, - const std::array>, kAlphabet>& - paths, - const std::array& freq) { + void encode_partition(std::size_t idx, + std::size_t depth, + std::size_t out_base, + const std::array& codes, + const std::array& freq) { if (nodes_[idx].is_leaf) { return; } @@ -801,12 +846,11 @@ class PivCoHuffman : public HuffmanBase { const std::uint8_t d = n.flat_depth; std::uint64_t* bits = arena_.data() + n.bits.offset; const symbol_type* src = workspace_half(depth) + out_base; + const std::uint16_t mask = (std::uint16_t{1} << d) - 1; for (std::size_t i = 0; i < n.bits.count; i++) { - // Build the d-bit suffix from the symbol's precomputed code path. - std::uint32_t suffix = 0; - for (std::uint8_t b = 0; b < d; b++) { - suffix = (suffix << 1) | (paths[src[i]][depth + b].second ? 1 : 0); - } + const SymbolCode code = codes[src[i]]; + const std::uint8_t shift = code.length - depth - d; + const std::uint16_t suffix = (code.bits >> shift) & mask; write_bits(bits, i * d, suffix, d); } return; @@ -827,7 +871,9 @@ class PivCoHuffman : public HuffmanBase { std::size_t c_right = 0; for (std::size_t i = 0; i < weight; i++) { symbol_type s = src[i]; - if (paths[s][depth].second) { + const SymbolCode code = codes[s]; + const bool goes_right = (code.bits >> (code.length - depth - 1)) & 1u; + if (goes_right) { bits[word_idx] |= (1ull << bit_pos); right_dst[c_right++] = s; } else { @@ -839,8 +885,8 @@ class PivCoHuffman : public HuffmanBase { } } - encode_partition(n.left, depth + 1, out_base, paths, freq); - encode_partition(n.right, depth + 1, out_base + left_weight, paths, freq); + encode_partition(n.left, depth + 1, out_base, codes, freq); + encode_partition(n.right, depth + 1, out_base + left_weight, codes, freq); } // --- decode -------------------------------------------------------------- @@ -890,7 +936,7 @@ class PivCoHuffman : public HuffmanBase { if (n.flat_depth > 0) { const std::uint8_t d = n.flat_depth; const std::uint64_t* bits = arena_.data() + n.bits.offset; - const symbol_type* table = flat_tables_[idx].data(); + const symbol_type* table = flat_tables_.data() + n.flat_table_offset; for (std::size_t i = 0; i < weight; i++) { dst[i] = table[read_bits(bits, i * d, d)]; } @@ -1031,7 +1077,7 @@ class PivCoHuffman : public HuffmanBase { } // Rebuild the canonical tree from lengths. - build_canonical_tree(code_lengths_); + build_noncanonical_tree(code_lengths_); detect_flat_subtrees(); // Single-symbol input: root is a leaf, no internal nodes, no arena. From 42f6f44c725558c51dd732ac2e2e72223a987600 Mon Sep 17 00:00:00 2001 From: Sanya239 Date: Sun, 26 Jul 2026 15:11:06 +0300 Subject: [PATCH 09/14] Add more tests --- src/tests/pivco_huffman_tests.cpp | 49 +++++++++++++++++++++++++++++++ 1 file changed, 49 insertions(+) diff --git a/src/tests/pivco_huffman_tests.cpp b/src/tests/pivco_huffman_tests.cpp index fcc656b..52d7fcb 100644 --- a/src/tests/pivco_huffman_tests.cpp +++ b/src/tests/pivco_huffman_tests.cpp @@ -106,3 +106,52 @@ TEST(PivCoHuffmanSmoke, SkewedInputRoundTrips) { EXPECT_EQ(decode_same(codec), data); EXPECT_EQ(decode_via_bytes(codec), data); } + +TEST(PivCoHuffmanSmoke, LengthLimitedTreeRoundTrips) { + // Fibonacci weights create a maximally deep unconstrained Huffman tree and + // exercise the 15-bit length correction. + std::vector data; + std::size_t previous = 1; + std::size_t current = 1; + for (std::uint8_t symbol = 0; symbol < 17; ++symbol) { + const std::size_t frequency = symbol < 2 ? 1 : previous + current; + if (symbol >= 2) { + previous = current; + current = frequency; + } + data.insert(data.end(), frequency, symbol); + } + + const PivCoHuffman codec(data); + EXPECT_EQ(decode_same(codec), data); + EXPECT_EQ(decode_via_bytes(codec), data); +} + +TEST(PivCoHuffmanSmoke, MultipleBlocksRoundTrip) { + std::mt19937 rng(9127); + std::uniform_int_distribution dist(0, 255); + std::vector data(3 * 64 * 1024 + 137); + for (auto& byte : data) { + byte = static_cast(dist(rng)); + } + + const PivCoHuffman codec(data); + EXPECT_EQ(decode_same(codec), data); + EXPECT_EQ(decode_via_bytes(codec), data); +} + +TEST(PivCoHuffmanSmoke, RandomizedAlphabetAndBlockSizesRoundTrip) { + std::mt19937 rng(3319); + for (std::size_t test_case = 0; test_case < 32; ++test_case) { + const std::size_t alphabet_size = 1 + rng() % 256; + const std::size_t size = 1 + rng() % (2 * 64 * 1024); + std::vector data(size); + for (auto& byte : data) { + byte = static_cast(rng() % alphabet_size); + } + + const PivCoHuffman codec(data); + EXPECT_EQ(decode_same(codec), data) << "test case " << test_case; + EXPECT_EQ(decode_via_bytes(codec), data) << "test case " << test_case; + } +} From f1525e846f6ce6852a417a82262e425a2faebf6f Mon Sep 17 00:00:00 2001 From: Sanya239 Date: Sun, 26 Jul 2026 20:12:52 +0300 Subject: [PATCH 10/14] Add SIMD kernel --- include/pixie/huffman/pivco_huffman.h | 65 ++- include/pixie/huffman/pivco_simd.h | 553 ++++++++++++++++++++++++++ 2 files changed, 580 insertions(+), 38 deletions(-) create mode 100644 include/pixie/huffman/pivco_simd.h diff --git a/include/pixie/huffman/pivco_huffman.h b/include/pixie/huffman/pivco_huffman.h index 4ccea14..27d3cb2 100644 --- a/include/pixie/huffman/pivco_huffman.h +++ b/include/pixie/huffman/pivco_huffman.h @@ -17,14 +17,16 @@ * reshaped into flat subtrees as described by the PivCo-Huffman paper. * * This implementation uses packed `std::uint64_t` bitmaps, flat subtrees, - * contiguous arenas, and scalar bottom-up decoding. SIMD and selective ANS - * are intentionally left to specialized future implementations. + * contiguous arenas, and bottom-up decoding. The merge/partition primitives are + * vectorized in `pivco_simd.h` (AVX-512 VBMI2, AVX2, and scalar fallbacks); + * selective ANS remains future work. * * Range and ownership conventions follow `HuffmanBase`: symbols are bytes, * the compressed stream is a byte view, and the codec owns its serialized form. */ #include +#include #include #include @@ -127,6 +129,14 @@ class PivCoHuffman : public HuffmanBase { /// that hurt decode throughput on whole-input processing. static constexpr std::size_t kBlockSize = 64 * 1024; + /// @brief Extra bytes appended to the ping-pong workspace beyond `2 * block`. + /// @details Vectorized merge/partition kernels read and write a full 8- or + /// 64-byte vector on each iteration even when the final group is + /// shorter; this slack makes those over-reads/over-writes stay + /// within the allocated buffer rather than spilling past the end of + /// a child region that abuts the workspace boundary. + static constexpr std::size_t kWorkspaceSlack = 64; + /** * @brief Lightweight descriptor for a node's routing bitmap. * @details The packed 64-bit words live in a single shared arena @@ -275,7 +285,7 @@ class PivCoHuffman : public HuffmanBase { // ping-pongs between the two halves — no per-node allocations. std::array codes{}; assign_codes(root_, 0, 0, codes); - workspace_.resize(block_uncompressed_size_ * 2); + workspace_.resize(block_uncompressed_size_ * 2 + kWorkspaceSlack); std::copy(input.begin(), input.end(), workspace_.data()); encode_partition(root_, 0, 0, codes, freq); } @@ -865,25 +875,19 @@ class PivCoHuffman : public HuffmanBase { symbol_type* right_dst = workspace_half(depth + 1) + out_base + left_weight; std::uint64_t* bits = arena_.data() + n.bits.offset; - std::size_t word_idx = 0; - std::size_t bit_pos = 0; - std::size_t c_left = 0; - std::size_t c_right = 0; - for (std::size_t i = 0; i < weight; i++) { - symbol_type s = src[i]; - const SymbolCode code = codes[s]; - const bool goes_right = (code.bits >> (code.length - depth - 1)) & 1u; - if (goes_right) { - bits[word_idx] |= (1ull << bit_pos); - right_dst[c_right++] = s; - } else { - left_dst[c_left++] = s; - } - if (++bit_pos == kWordBits) { - bit_pos = 0; - ++word_idx; - } + + // Per-symbol direction bit at this node's depth: bit (length-depth-1) of + // each symbol's code. Only symbols passing through (length > depth) are + // routed; absent symbols (length 0) are never read here. + std::array dir{}; + for (std::size_t s = 0; s < kAlphabet; ++s) { + const std::uint8_t len = codes[s].length; + dir[s] = (len > depth) ? static_cast( + (codes[s].bits >> (len - depth - 1)) & 1u) + : 0; } + pivco_partition_encode(left_dst, right_dst, bits, src, dir.data(), weight, + left_weight, weight - left_weight); encode_partition(n.left, depth + 1, out_base, codes, freq); encode_partition(n.right, depth + 1, out_base + left_weight, codes, freq); @@ -907,7 +911,7 @@ class PivCoHuffman : public HuffmanBase { nodes_[root_].symbol); } const std::size_t n = block_uncompressed_size_; - workspace_.resize(n * 2); + workspace_.resize(n * 2 + kWorkspaceSlack); decode_node(root_, n, 0, 0); return std::vector(workspace_.begin(), workspace_.begin() + n); } @@ -953,23 +957,8 @@ class PivCoHuffman : public HuffmanBase { const symbol_type* left_out = src; const symbol_type* right_out = src + left_weight; - std::size_t out_i = 0; - std::size_t c_left = 0; - std::size_t c_right = 0; const std::uint64_t* bits = arena_.data() + n.bits.offset; - const std::size_t words = n.bits.word_count(0); - for (std::size_t w = 0; w < words; w++) { - std::uint64_t word = bits[w]; - std::size_t limit = kWordBits; - if (w + 1 == words) { - limit = n.bits.count - w * kWordBits; - } - for (std::size_t i = 0; i < limit; i++) { - dst[out_i++] = - (word & 1ull) ? right_out[c_right++] : left_out[c_left++]; - word >>= 1; - } - } + pivco_merge_decode(dst, left_out, right_out, bits, n.bits.count); } /** @brief Pointer to the workspace half that holds depth-@p depth outputs. */ diff --git a/include/pixie/huffman/pivco_simd.h b/include/pixie/huffman/pivco_simd.h new file mode 100644 index 0000000..53b89e8 --- /dev/null +++ b/include/pixie/huffman/pivco_simd.h @@ -0,0 +1,553 @@ +#pragma once + +/** + * @file pivco_simd.h + * @brief SIMD kernels for PivCo-Huffman encode/decode primitives. + * + * Implements the two core PivCo-Huffman primitives from the paper with + * architecture-specific kernels and a portable scalar fallback: + * + * - `pivco_merge_decode` (bottom-up decode, `merge_vec_vec`): interleaves two + * dense child symbol streams under a 1-bit-per-symbol routing bitmap into a + * single dense output. A set bitmap bit selects the right child, a clear bit + * the left child. + * - `pivco_partition_encode` (top-down encode partition): the inverse — given + * per-symbol direction bits, splits a dense input into left/right child + * streams and writes the routing bitmap. + * + * Bitmap convention (matches `pivco_huffman.h`): bits are packed LSB-first + * into 64-bit words; bit 0 of word 0 is the first symbol. + * + * Backends, selected by compiler feature macros (`-march`): + * - AVX-512 VBMI2 (`__AVX512VBMI2__`): byte-granular `vpexpandb`/`vpcompressb` + * operating on 64 bytes per iteration. This is the instruction the paper + * (Sect. 6.6) identifies as mapping directly to the partition kernel. + * - AVX2 + SSE4.1: 8-byte-at-a-time `_mm_shuffle_epi8` with 256-entry lookup + * tables, mirroring the paper's NEON `vqtbl1q_u8` technique. + * - Scalar: the original bit-by-bit / symbol-by-symbol loops. + * + * Every SIMD backend produces byte-identical output to the scalar fallback + * (the operation is a deterministic permutation). The existing round-trip tests + * under release and AddressSanitizer serve as the differential oracle for every + * backend the test matrix builds and runs; the AVX-512 VBMI2 backend is + * verified only when built on VBMI2-capable hardware, since there is no runtime + * dispatch. + */ + +#include +#include +#include +#include +#include + +#if defined(__AVX2__) || defined(__AVX512VBMI2__) +#include +#endif + +namespace pixie::pivco_simd_detail { + +/// @brief Bitmap bit-ordering convention used throughout this header. +/// @details Every routing bitmap packs one bit per symbol, LSB-first within +/// each +/// 64-bit word: bit `k` is the direction bit of symbol `k`, where a +/// set bit selects the right child and a clear bit the left child. All +/// mask builders (`build_dir_mask8`, the AVX-512 scalar loop), LUT +/// tables, merge kernels, and partition tails MUST agree on this +/// convention — flipping it in one backend only would make the +/// serialized bitmap backend-dependent and silently break +/// cross-`-march` encode/decode. + +// --------------------------------------------------------------------------- +// Scalar helpers (always compiled; used as tail + fallback). +// --------------------------------------------------------------------------- + +/// @brief Build an 8-bit routing mask for symbols @p s8 (one byte per symbol) +/// under the LSB-first convention (bit k = dir[s8[k]]). +/// @details Shared source of truth for the 8-symbol direction mask; the AVX2 +/// partition kernel and (transitively) the AVX-512 tail both build +/// their per-group mask here. +inline std::uint8_t build_dir_mask8(const std::uint8_t* s8, + const std::uint8_t* dir) noexcept { + return static_cast((dir[s8[0]] << 0) | (dir[s8[1]] << 1) | + (dir[s8[2]] << 2) | (dir[s8[3]] << 3) | + (dir[s8[4]] << 4) | (dir[s8[5]] << 5) | + (dir[s8[6]] << 6) | (dir[s8[7]] << 7)); +} + +/// @brief Read 8 bits starting at bit position @p pos from a packed word array. +/// @pre @p pos + 8 is within the bitmap's bit range; callers stop the SIMD +/// loop once fewer than 8 symbols remain, so the span never exceeds the +/// allocated words. +inline std::uint8_t extract_byte(const std::uint64_t* bits, + std::size_t pos) noexcept { + const std::size_t w = pos / 64; + const std::size_t b = pos % 64; + std::uint64_t v = bits[w] >> b; + if (b + 8 > 64) { + v |= bits[w + 1] << (64 - b); + } + return static_cast(v); +} + +/// @brief Scalar merge tail: process the remaining `< 8` symbols bit-by-bit. +inline void merge_scalar_tail(std::uint8_t* dst, + std::size_t& out_i, + const std::uint8_t* left, + std::size_t& c_left, + const std::uint8_t* right, + std::size_t& c_right, + const std::uint64_t* bits, + std::size_t i, + std::size_t count) noexcept { + for (; i < count; ++i) { + const std::uint8_t bit = + static_cast((bits[i / 64] >> (i % 64)) & 1u); + dst[out_i++] = bit ? right[c_right++] : left[c_left++]; + } +} + +/// @brief Scalar partition tail: route the remaining symbols (fewer than the +/// vector width) one at a time, writing the bitmap bit and appending to +/// the left/right child stream. +/// @details Shared by the AVX2 kernel tail, the AVX-512 kernel tail, and the +/// scalar fallback, so the LSB-first `|=`-into-pre-zeroed-arena +/// convention lives in exactly one place. @p bits is the bitmap base +/// (word array); @p start_bit is the absolute bit offset where the +/// tail begins (must be a multiple of the full-group width so it is +/// already byte-aligned within the word packing). +inline void partition_scalar_tail(std::uint8_t* left_dst, + std::uint8_t* right_dst, + std::uint64_t* bits, + const std::uint8_t* src, + const std::uint8_t* dir, + std::size_t& c_left, + std::size_t& c_right, + std::size_t start_i, + std::size_t weight) noexcept { + std::size_t i = start_i; + std::size_t word_idx = i / 64; + std::size_t bit_pos = i % 64; + for (; i < weight; ++i) { + const std::uint8_t s = src[i]; + if (dir[s]) { + bits[word_idx] |= (1ull << bit_pos); + right_dst[c_right++] = s; + } else { + left_dst[c_left++] = s; + } + if (++bit_pos == 64) { + bit_pos = 0; + ++word_idx; + } + } +} + +// =========================================================================== +// AVX2 backend (SSE4.1 byte-shuffle lookup tables; AVX2 implies SSE4.1). +// =========================================================================== +#if defined(__AVX2__) && defined(__SSE4_1__) + +/// @brief 256-entry × 16-byte shuffle-control table, 16-byte aligned per entry. +/// @details Stored as plain bytes (not `__m128i`) to avoid dropping the +/// `__may_alias__` attribute that `std::array<__m128i>` would warn +/// about; entries are loaded as aligned vectors at use time. +struct ShufTable { + alignas(16) std::uint8_t entry[256][16]; +}; + +/// @brief Load shuffle entry @p m as an aligned 128-bit vector. +inline __m128i load_shuf(const ShufTable& t, std::uint8_t m) noexcept { + return _mm_load_si128( + reinterpret_cast(&t.entry[m])); // NOLINT +} + +/// @brief 256-entry shuffle table for the 8-byte merge. +/// @details For an 8-bit mask `m` (bit i = 1 -> right), entry `m` is a +/// 16-byte shuffle control that, applied to a 128-bit vector holding +/// `left[0..7]` in its low qword and `right[0..7]` in its high qword, +/// produces the merged 8-byte output. Mirrors the paper's NEON +/// `expand_tab`. +inline const ShufTable& merge_lut() { + static const ShufTable lut = [] { + ShufTable t{}; + for (int m = 0; m < 256; ++m) { + int nl = 0; + int nr = 0; + for (int i = 0; i < 8; ++i) { + if ((m >> i) & 1) { + t.entry[m][i] = static_cast(8 + nr); + ++nr; + } else { + t.entry[m][i] = static_cast(nl); + ++nl; + } + } + for (int i = 8; i < 16; ++i) { + t.entry[m][i] = 0x80; // zeroed lanes (unused, stored as full 16 bytes) + } + } + return t; + }(); + return lut; +} + +/// @brief 256-entry shuffle tables for the 8-byte compress (partition). +/// @details `compress_right[m]` gathers the elements at the 1-bit positions of +/// mask `m` (in order) to the front of the 128-bit result; +/// `compress_left[m]` does the same for the 0-bit positions. Trailing +/// lanes are zeroed. Mirrors the paper's NEON `compress_tab`. +inline const ShufTable& compress_right_lut() { + static const ShufTable lut = [] { + ShufTable t{}; + for (int m = 0; m < 256; ++m) { + int p = 0; + for (int i = 0; i < 8; ++i) { + if ((m >> i) & 1) { + t.entry[m][p++] = static_cast(i); + } + } + for (int i = p; i < 16; ++i) { + t.entry[m][i] = 0x80; + } + } + return t; + }(); + return lut; +} + +inline const ShufTable& compress_left_lut() { + static const ShufTable lut = [] { + ShufTable t{}; + for (int m = 0; m < 256; ++m) { + int p = 0; + for (int i = 0; i < 8; ++i) { + if (((m >> i) & 1) == 0) { + t.entry[m][p++] = static_cast(i); + } + } + for (int i = p; i < 16; ++i) { + t.entry[m][i] = 0x80; + } + } + return t; + }(); + return lut; +} + +inline void merge_decode_avx2(std::uint8_t* dst, + const std::uint8_t* left, + const std::uint8_t* right, + const std::uint64_t* bits, + std::size_t count) noexcept { + const auto& mlut = merge_lut(); + std::size_t out_i = 0; + std::size_t c_left = 0; + std::size_t c_right = 0; + std::size_t i = 0; + for (; i + 8 <= count; i += 8) { + const std::uint8_t m8 = extract_byte(bits, i); + std::uint64_t lv = 0; + std::uint64_t rv = 0; + std::memcpy(&lv, left + c_left, 8); + std::memcpy(&rv, right + c_right, 8); + const __m128i combined = + _mm_set_epi64x(static_cast(rv), static_cast(lv)); + const __m128i out = _mm_shuffle_epi8(combined, load_shuf(mlut, m8)); + _mm_storel_epi64(reinterpret_cast<__m128i*>(dst + out_i), out); + out_i += 8; + const int nr = std::popcount(static_cast(m8)); + c_right += static_cast(nr); + c_left += static_cast(8 - nr); + } + merge_scalar_tail(dst, out_i, left, c_left, right, c_right, bits, i, count); +} + +/// @brief Store @p v (low @p n bytes valid) at @p dst without writing past +/// @p dst + @p n. +/// @details The codec partitions into a ping-pong workspace where a node's left +/// and right child regions are contiguous, and a node's right region +/// is immediately followed by its sibling's region (already written by +/// the parent, consumed later in pre-order). A full-width vector store +/// would clobber the start of the adjacent region, so when the valid +/// byte count is less than the vector width the store is masked to +/// exactly +/// @p n bytes. +inline void store_masked_epi64(std::uint8_t* dst, __m128i v, int n) noexcept { + if (n >= 8) { + _mm_storel_epi64(reinterpret_cast<__m128i*>(dst), v); + } else if (n > 0) { + alignas(8) std::uint8_t tmp[8]; + _mm_storel_epi64(reinterpret_cast<__m128i*>(tmp), v); + std::memcpy(dst, tmp, static_cast(n)); + } +} + +inline void partition_encode_avx2(std::uint8_t* left_dst, + std::uint8_t* right_dst, + std::uint64_t* bits, + const std::uint8_t* src, + const std::uint8_t* dir, + std::size_t weight, + std::size_t left_weight, + std::size_t right_weight) noexcept { + const auto& cr = compress_right_lut(); + const auto& cl = compress_left_lut(); + std::uint8_t* bits_bytes = reinterpret_cast(bits); + + const std::size_t full = weight / 8; + std::size_t c_left = 0; + std::size_t c_right = 0; + std::size_t g = 0; + for (; g < full; ++g) { + const std::uint8_t* s8 = src + g * 8; + const __m128i data = _mm_loadl_epi64(reinterpret_cast(s8)); + const std::uint8_t m8 = build_dir_mask8(s8, dir); + // Byte-aligned write: full groups always start on a byte boundary in the + // LSB-first bitmap, so byte g holds exactly symbols [8g, 8g+8). + bits_bytes[g] = m8; + const __m128i rv = _mm_shuffle_epi8(data, load_shuf(cr, m8)); + const __m128i lv = _mm_shuffle_epi8(data, load_shuf(cl, m8)); + const int nr = std::popcount(static_cast(m8)); + const int nl = 8 - nr; + // Both child regions are contiguous with a later-consumed sibling region, + // so neither store may overflow past its own region end. Masked stores + // keep the trailing invalid lanes inside the region when they would + // otherwise spill into the neighbor. + store_masked_epi64(right_dst + c_right, rv, + static_cast(right_weight - c_right)); + store_masked_epi64(left_dst + c_left, lv, + static_cast(left_weight - c_left)); + c_right += static_cast(nr); + c_left += static_cast(nl); + } + + // Scalar tail for the trailing < 8 symbols (byte-unaligned bitmap bits). + partition_scalar_tail(left_dst, right_dst, bits, src, dir, c_left, c_right, + full * 8, weight); +} + +#endif // AVX2 + SSE4.1 + +// =========================================================================== +// AVX-512 VBMI2 backend (byte-granular expand/compress, 64 bytes/iteration). +// =========================================================================== +#if defined(__AVX512VBMI2__) && defined(__AVX512BW__) + +inline void merge_decode_avx512(std::uint8_t* dst, + const std::uint8_t* left, + const std::uint8_t* right, + const std::uint64_t* bits, + std::size_t count) noexcept { + std::size_t out_i = 0; + std::size_t c_left = 0; + std::size_t c_right = 0; + const std::size_t full = count / 64; + for (std::size_t w = 0; w < full; ++w) { + const __mmask64 km = bits[w]; + const __m512i lv = _mm512_loadu_si512(left + c_left); + const __m512i rv = _mm512_loadu_si512(right + c_right); + // Fill clear-bit positions with left symbols, then set-bit positions with + // right symbols. maskz zeroes the right positions first; the second expand + // keeps the left values there. + __m512i d = _mm512_maskz_expand_epi8(~km, lv); + d = _mm512_mask_expand_epi8(d, km, rv); + _mm512_storeu_si512(dst + out_i, d); + const std::size_t nr = std::popcount(bits[w]); + out_i += 64; + c_right += nr; + c_left += 64 - nr; + } + // Tail handled by the AVX2 path if available, otherwise scalar. +#if defined(__AVX2__) && defined(__SSE4_1__) + const std::size_t done = full * 64; + merge_decode_avx2(dst + done, left + c_left, right + c_right, bits + full, + count - done); + (void)out_i; +#else + merge_scalar_tail(dst, out_i, left, c_left, right, c_right, bits, full * 64, + count); +#endif +} + +inline void partition_encode_avx512(std::uint8_t* left_dst, + std::uint8_t* right_dst, + std::uint64_t* bits, + const std::uint8_t* src, + const std::uint8_t* dir, + std::size_t weight, + std::size_t left_weight, + std::size_t right_weight) noexcept { + // Build the routing mask 64 symbols at a time, store the word, then compress + // the 64 input symbols into left/right streams with vpcompressb. + const std::size_t full = weight / 64; + std::size_t c_left = 0; + std::size_t c_right = 0; + std::size_t g = 0; + for (; g < full; ++g) { + const std::uint8_t* s64 = src + g * 64; + const __m512i data = _mm512_loadu_si512(s64); + // Direction mask: bit k = dir[symbol k], LSB-first. x86 has no + // byte-granular gather, so the 64-bit mask is built with a scalar OR loop + // over the dir table. (A qword gather + vpmovb2m would be faster in + // principle but reads all 64 byte sign-bits, corrupting the mask with + // neighboring dir entries.) + __mmask64 km = 0; + for (int k = 0; k < 64; ++k) { + km |= static_cast<__mmask64>(dir[s64[k]]) << k; + } + bits[g] = km; + // Compress: pack right-bound (set-bit) elements to the front of one stream + // and left-bound (clear-bit) elements to the front of another. + const __m512i right_packed = _mm512_maskz_compress_epi8(km, data); + const __m512i left_packed = _mm512_maskz_compress_epi8(~km, data); + const std::size_t nr = std::popcount(km); + const std::size_t nl = 64 - nr; + // Both child regions are contiguous with a later-consumed sibling region, + // so neither store may overflow its own region. Use masked stores that + // write only the valid lanes when the full vector would spill into the + // neighbor. + { + const std::size_t rem = right_weight - c_right; + if (rem >= 64) { + _mm512_storeu_si512(right_dst + c_right, right_packed); + } else { + const __mmask64 rmask = (1ull << rem) - 1; + _mm512_mask_storeu_epi8(right_dst + c_right, rmask, right_packed); + } + } + { + const std::size_t rem = left_weight - c_left; + if (rem >= 64) { + _mm512_storeu_si512(left_dst + c_left, left_packed); + } else { + const __mmask64 lmask = (1ull << rem) - 1; + _mm512_mask_storeu_epi8(left_dst + c_left, lmask, left_packed); + } + } + c_right += nr; + c_left += nl; + } + // Tail for the trailing < 64 symbols. `__AVX512VBMI2__` implies `__AVX2__` + // on every real target, so delegate to the AVX2 kernel (which itself defers + // the final < 8 symbols to the shared `partition_scalar_tail`). There is no + // reachable `#else` branch, so no scalar tail is duplicated here. + const std::size_t done_words = g; + const std::size_t done_syms = g * 64; + partition_encode_avx2(left_dst + c_left, right_dst + c_right, + bits + done_words, src + done_syms, dir, + weight - done_syms, left_weight - c_left, + right_weight - c_right); +} + +#endif // AVX-512 VBMI2 + +// =========================================================================== +// Scalar fallback (no SIMD). +// =========================================================================== + +inline void merge_decode_scalar(std::uint8_t* dst, + const std::uint8_t* left, + const std::uint8_t* right, + const std::uint64_t* bits, + std::size_t count) noexcept { + std::size_t out_i = 0; + std::size_t c_left = 0; + std::size_t c_right = 0; + const std::size_t words = (count + 63) / 64; + for (std::size_t w = 0; w < words; ++w) { + std::uint64_t word = bits[w]; + std::size_t limit = 64; + if (w + 1 == words) { + limit = count - w * 64; + } + for (std::size_t i = 0; i < limit; ++i) { + dst[out_i++] = (word & 1ull) ? right[c_right++] : left[c_left++]; + word >>= 1; + } + } +} + +inline void partition_encode_scalar(std::uint8_t* left_dst, + std::uint8_t* right_dst, + std::uint64_t* bits, + const std::uint8_t* src, + const std::uint8_t* dir, + std::size_t weight, + std::size_t /*left_weight*/, + std::size_t /*right_weight*/) noexcept { + std::size_t c_left = 0; + std::size_t c_right = 0; + partition_scalar_tail(left_dst, right_dst, bits, src, dir, c_left, c_right, 0, + weight); +} + +} // namespace pixie::pivco_simd_detail + +// --------------------------------------------------------------------------- +// Public dispatch entry points. +// --------------------------------------------------------------------------- + +namespace pixie { + +/// @brief Bottom-up decode merge: interleave @p left / @p right child streams +/// under @p bits into @p dst, producing @p count symbols. +/// @param dst Output buffer of (at least) @p count bytes. +/// @param left Left child output stream (consumed left_weight symbols). +/// @param right Right child output stream (consumed right_weight symbols). +/// @param bits Packed 1-bit-per-symbol routing bitmap, LSB-first per word. +/// Bit value 1 selects @p right, 0 selects @p left. +/// @param count Number of symbols (== bitmap bit count). +/// @details The caller must guarantee a few bytes of slack past the end of +/// @p dst, @p left, and @p right for vectorized over-reads/over-writes +/// (the codec's workspace is padded accordingly). +inline void pivco_merge_decode(std::uint8_t* dst, + const std::uint8_t* left, + const std::uint8_t* right, + const std::uint64_t* bits, + std::size_t count) noexcept { +#if defined(__AVX512VBMI2__) && defined(__AVX512BW__) + pivco_simd_detail::merge_decode_avx512(dst, left, right, bits, count); +#elif defined(__AVX2__) && defined(__SSE4_1__) + pivco_simd_detail::merge_decode_avx2(dst, left, right, bits, count); +#else + pivco_simd_detail::merge_decode_scalar(dst, left, right, bits, count); +#endif +} + +/// @brief Top-down encode partition: split @p src into left/right child streams +/// and write the routing bitmap, using per-symbol direction bits. +/// @param left_dst Output for symbols whose direction bit is 0 (exactly +/// @p left_weight symbols). +/// @param right_dst Output for symbols whose direction bit is 1. Must equal +/// @p left_dst + @p left_weight: the two streams are stored +/// contiguously, and each is immediately followed by a +/// later-consumed sibling region, so both stores are masked +/// to never spill past their own region end. +/// @param bits Destination bitmap (pre-zeroed), LSB-first per word. +/// @param src Input symbols to partition. +/// @param dir Per-symbol direction bit table (0 or 1), indexed by +/// symbol. +/// @param weight Number of input symbols (== bitmap bit count). +/// @param left_weight Number of symbols routed left (== bitmap zero-count). +/// @param right_weight Number of symbols routed right (== bitmap one-count). +inline void pivco_partition_encode(std::uint8_t* left_dst, + std::uint8_t* right_dst, + std::uint64_t* bits, + const std::uint8_t* src, + const std::uint8_t* dir, + std::size_t weight, + std::size_t left_weight, + std::size_t right_weight) noexcept { +#if defined(__AVX512VBMI2__) && defined(__AVX512BW__) + pivco_simd_detail::partition_encode_avx512( + left_dst, right_dst, bits, src, dir, weight, left_weight, right_weight); +#elif defined(__AVX2__) && defined(__SSE4_1__) + pivco_simd_detail::partition_encode_avx2(left_dst, right_dst, bits, src, dir, + weight, left_weight, right_weight); +#else + pivco_simd_detail::partition_encode_scalar( + left_dst, right_dst, bits, src, dir, weight, left_weight, right_weight); +#endif +} + +} // namespace pixie From befa6f709dae466057c06076bbdb752d7e825b9c Mon Sep 17 00:00:00 2001 From: Sanya239 Date: Sun, 26 Jul 2026 22:03:45 +0300 Subject: [PATCH 11/14] Rework SIMD kernel --- include/pixie/huffman/pivco_huffman.h | 240 ++++++--- include/pixie/huffman/pivco_simd.h | 742 ++++++++++++++++++++++---- src/tests/pivco_huffman_tests.cpp | 108 ++++ 3 files changed, 902 insertions(+), 188 deletions(-) diff --git a/include/pixie/huffman/pivco_huffman.h b/include/pixie/huffman/pivco_huffman.h index 27d3cb2..15ef355 100644 --- a/include/pixie/huffman/pivco_huffman.h +++ b/include/pixie/huffman/pivco_huffman.h @@ -30,7 +30,6 @@ #include #include -#include #include #include #include @@ -193,6 +192,9 @@ class PivCoHuffman : public HuffmanBase { std::uint8_t length = 0; }; + using DirectionTables = + std::array, kMaxCodeLength>; + /// @brief Total uncompressed size across all blocks (set during /// build/deserialize, read by `uncompressed_size_impl()`). std::size_t uncompressed_size_ = 0; @@ -285,9 +287,10 @@ class PivCoHuffman : public HuffmanBase { // ping-pongs between the two halves — no per-node allocations. std::array codes{}; assign_codes(root_, 0, 0, codes); + const DirectionTables directions = build_direction_tables(codes); workspace_.resize(block_uncompressed_size_ * 2 + kWorkspaceSlack); std::copy(input.begin(), input.end(), workspace_.data()); - encode_partition(root_, 0, 0, codes, freq); + encode_partition(root_, 0, 0, directions, freq); } /** @brief Count a block's byte frequencies with independent accumulators. @@ -477,37 +480,6 @@ class PivCoHuffman : public HuffmanBase { return value >> (8 - width); } - /** @brief Read @p d bits from a packed word array starting at bit position - * @p pos, MSB-first. Returns the decoded value. */ - static std::uint32_t read_bits(const std::uint64_t* words, - std::size_t pos, - std::uint8_t d) { - const std::size_t word_index = pos / kWordBits; - const std::size_t shift = pos % kWordBits; - std::uint64_t packed = words[word_index] >> shift; - if (shift + d > kWordBits) { - packed |= words[word_index + 1] << (kWordBits - shift); - } - const std::uint8_t mask = static_cast((1u << d) - 1); - return reverse_low_bits(static_cast(packed) & mask, d); - } - - /** @brief Write @p d bits of @p value to a packed word array starting at - * bit position @p pos, MSB-first. Only sets 1-bits (pre-zeroed). */ - static void write_bits(std::uint64_t* words, - std::size_t pos, - std::uint32_t value, - std::uint8_t d) { - const std::uint64_t packed = - reverse_low_bits(static_cast(value), d); - const std::size_t word_index = pos / kWordBits; - const std::size_t shift = pos % kWordBits; - words[word_index] |= packed << shift; - if (shift + d > kWordBits) { - words[word_index + 1] |= packed >> (kWordBits - shift); - } - } - // --- flat subtree detection ---------------------------------------------- /** @brief Sentinel value meaning "subtree has non-uniform leaf depths". */ @@ -562,7 +534,8 @@ class PivCoHuffman : public HuffmanBase { n.flat_depth = depth; n.flat_table_offset = flat_tables_.size(); flat_tables_.resize(flat_tables_.size() + (std::size_t{1} << depth)); - build_flat_table(idx, 0, flat_tables_.data() + n.flat_table_offset); + build_flat_table(idx, depth, 0, + flat_tables_.data() + n.flat_table_offset); return; } @@ -572,19 +545,22 @@ class PivCoHuffman : public HuffmanBase { /** @brief Populate the flat lookup table by walking the subtree. * @param idx Current node. + * @param flat_depth Total depth of the flat subtree. * @param code Suffix accumulated so far (MSB-first: left=0, * right=1). * @param table Output table of size `2^total_depth`. */ void build_flat_table(std::size_t idx, + std::uint8_t flat_depth, std::uint32_t code, symbol_type* table) const { const Node& n = nodes_[idx]; if (n.is_leaf) { - table[code] = n.symbol; + table[reverse_low_bits(static_cast(code), flat_depth)] = + n.symbol; return; } - build_flat_table(n.left, code << 1, table); - build_flat_table(n.right, (code << 1) | 1, table); + build_flat_table(n.left, flat_depth, code << 1, table); + build_flat_table(n.right, flat_depth, (code << 1) | 1, table); } /** @brief Total frequency weight of a subtree (sum of leaf frequencies). */ @@ -826,6 +802,91 @@ class PivCoHuffman : public HuffmanBase { length + 1, codes); } + /** @brief Build the per-depth direction tables shared by all tree nodes. + * @details A symbol can occur in only one node at a given depth, so one + * 256-byte table per depth replaces rebuilding the same direction + * information separately for every internal node. */ + static DirectionTables build_direction_tables( + const std::array& codes) { + DirectionTables directions{}; + for (std::size_t symbol = 0; symbol < kAlphabet; ++symbol) { + const SymbolCode code = codes[symbol]; + for (std::uint8_t depth = 0; depth < code.length; ++depth) { + directions[depth][symbol] = static_cast( + (code.bits >> (code.length - depth - 1)) & 1u); + } + } + return directions; + } + + /** @brief Pack one flat node's fixed-width wire codes. + * @details Inverts the node's small decode table once, then packs directly + * from the symbol-to-wire-code lookup. Widths dividing a byte get + * dedicated loops; other widths track word position without a + * division per symbol. */ + static void encode_flat_values(std::uint64_t* bits, + const symbol_type* src, + std::size_t count, + std::uint8_t flat_depth, + const symbol_type* table) { + std::array wire_code; + const std::size_t alphabet_size = std::size_t{1} << flat_depth; + for (std::size_t code = 0; code < alphabet_size; ++code) { + wire_code[table[code]] = static_cast(code); + } + + auto* bytes = reinterpret_cast(bits); + if (flat_depth == 8) { + for (std::size_t i = 0; i < count; ++i) { + bytes[i] = wire_code[src[i]]; + } + return; + } + if (flat_depth == 4) { + std::size_t i = 0; + for (; i + 2 <= count; i += 2) { + bytes[i / 2] = static_cast(wire_code[src[i]] | + (wire_code[src[i + 1]] << 4)); + } + if (i < count) { + bytes[i / 2] = wire_code[src[i]]; + } + return; + } + if (flat_depth == 2) { + std::size_t i = 0; + for (; i + 4 <= count; i += 4) { + bytes[i / 4] = static_cast( + wire_code[src[i]] | (wire_code[src[i + 1]] << 2) | + (wire_code[src[i + 2]] << 4) | (wire_code[src[i + 3]] << 6)); + } + if (i < count) { + std::uint8_t packed = 0; + for (std::size_t lane = 0; i + lane < count; ++lane) { + packed |= + static_cast(wire_code[src[i + lane]] << (2 * lane)); + } + bytes[i / 4] = packed; + } + return; + } + + std::size_t word_index = 0; + std::size_t bit_offset = 0; + for (std::size_t i = 0; i < count; ++i) { + const std::uint64_t packed = wire_code[src[i]]; + bits[word_index] |= packed << bit_offset; + if (bit_offset + flat_depth > kWordBits) { + bits[word_index + 1] |= packed >> (kWordBits - bit_offset); + } + bit_offset += flat_depth; + if (bit_offset >= kWordBits) { + bit_offset -= kWordBits; + ++word_index; + } + } + } + /** @brief Top-down in-place partition encoding into the reusable workspace. * @details The dual of decode: at each internal node, read the incoming * symbol stream from workspace half `depth`, fill the node's @@ -838,12 +899,12 @@ class PivCoHuffman : public HuffmanBase { * @param depth Current tree depth (selects workspace half + path index). * @param out_base Offset within the current half where this node's stream * begins. - * @param codes Compact per-symbol codes from `assign_codes`. + * @param directions Per-depth symbol direction lookup tables. * @param freq Symbol frequencies (sizing child partitions). */ void encode_partition(std::size_t idx, std::size_t depth, std::size_t out_base, - const std::array& codes, + const DirectionTables& directions, const std::array& freq) { if (nodes_[idx].is_leaf) { return; @@ -853,15 +914,14 @@ class PivCoHuffman : public HuffmanBase { // Flat node: write d-bit suffix per symbol, no partition, no recursion. if (n.flat_depth > 0) { - const std::uint8_t d = n.flat_depth; std::uint64_t* bits = arena_.data() + n.bits.offset; const symbol_type* src = workspace_half(depth) + out_base; - const std::uint16_t mask = (std::uint16_t{1} << d) - 1; - for (std::size_t i = 0; i < n.bits.count; i++) { - const SymbolCode code = codes[src[i]]; - const std::uint8_t shift = code.length - depth - d; - const std::uint16_t suffix = (code.bits >> shift) & mask; - write_bits(bits, i * d, suffix, d); + if (n.flat_depth == 1) { + pivco_partition_encode_none(bits, src, directions[depth].data(), + n.bits.count); + } else { + const symbol_type* table = flat_tables_.data() + n.flat_table_offset; + encode_flat_values(bits, src, n.bits.count, n.flat_depth, table); } return; } @@ -875,22 +935,25 @@ class PivCoHuffman : public HuffmanBase { symbol_type* right_dst = workspace_half(depth + 1) + out_base + left_weight; std::uint64_t* bits = arena_.data() + n.bits.offset; - - // Per-symbol direction bit at this node's depth: bit (length-depth-1) of - // each symbol's code. Only symbols passing through (length > depth) are - // routed; absent symbols (length 0) are never read here. - std::array dir{}; - for (std::size_t s = 0; s < kAlphabet; ++s) { - const std::uint8_t len = codes[s].length; - dir[s] = (len > depth) ? static_cast( - (codes[s].bits >> (len - depth - 1)) & 1u) - : 0; - } - pivco_partition_encode(left_dst, right_dst, bits, src, dir.data(), weight, - left_weight, weight - left_weight); - - encode_partition(n.left, depth + 1, out_base, codes, freq); - encode_partition(n.right, depth + 1, out_base + left_weight, codes, freq); + const std::uint8_t* dir = directions[depth].data(); + const bool left_is_leaf = nodes_[n.left].is_leaf; + const bool right_is_leaf = nodes_[n.right].is_leaf; + if (left_is_leaf && right_is_leaf) { + pivco_partition_encode_none(bits, src, dir, weight); + } else if (right_is_leaf) { + pivco_partition_encode_left(left_dst, bits, src, dir, weight, + left_weight); + } else if (left_is_leaf) { + pivco_partition_encode_right(right_dst, bits, src, dir, weight, + weight - left_weight); + } else { + pivco_partition_encode(left_dst, right_dst, bits, src, dir, weight, + left_weight, weight - left_weight); + } + + encode_partition(n.left, depth + 1, out_base, directions, freq); + encode_partition(n.right, depth + 1, out_base + left_weight, directions, + freq); } // --- decode -------------------------------------------------------------- @@ -938,27 +1001,43 @@ class PivCoHuffman : public HuffmanBase { // Flat node: read d-bit suffixes, lookup symbols, no merge. if (n.flat_depth > 0) { - const std::uint8_t d = n.flat_depth; const std::uint64_t* bits = arena_.data() + n.bits.offset; const symbol_type* table = flat_tables_.data() + n.flat_table_offset; - for (std::size_t i = 0; i < weight; i++) { - dst[i] = table[read_bits(bits, i * d, d)]; - } + pivco_flat_decode(dst, bits, table, weight, n.flat_depth); + return; + } + + // Normal node: decode non-leaf children, then select the paper's matching + // MVV/MCV/MVC/MCC merge primitive. Internal children already carry their + // exact stream size, so no bitmap popcount pass is needed. + const Node& left = nodes_[n.left]; + const Node& right = nodes_[n.right]; + const std::uint64_t* bits = arena_.data() + n.bits.offset; + if (left.is_leaf && right.is_leaf) { + pivco_merge_decode_cst_cst(dst, left.symbol, right.symbol, bits, weight); return; } - // Normal node: decode children, then merge by 1-bit-per-symbol bitmap. - const std::size_t right_weight = bitmap_popcount(n.bits); - const std::size_t left_weight = n.bits.count - right_weight; - decode_node(n.left, left_weight, depth + 1, out_base); - decode_node(n.right, right_weight, depth + 1, out_base + left_weight); + const std::size_t left_weight = + left.is_leaf ? weight - right.bits.count : left.bits.count; + const std::size_t right_weight = weight - left_weight; + if (!left.is_leaf) { + decode_node(n.left, left_weight, depth + 1, out_base); + } + if (!right.is_leaf) { + decode_node(n.right, right_weight, depth + 1, out_base + left_weight); + } const symbol_type* src = workspace_half(depth + 1) + out_base; const symbol_type* left_out = src; const symbol_type* right_out = src + left_weight; - - const std::uint64_t* bits = arena_.data() + n.bits.offset; - pivco_merge_decode(dst, left_out, right_out, bits, n.bits.count); + if (left.is_leaf) { + pivco_merge_decode_cst_vec(dst, left.symbol, right_out, bits, weight); + } else if (right.is_leaf) { + pivco_merge_decode_vec_cst(dst, left_out, right.symbol, bits, weight); + } else { + pivco_merge_decode(dst, left_out, right_out, bits, weight); + } } /** @brief Pointer to the workspace half that holds depth-@p depth outputs. */ @@ -966,19 +1045,6 @@ class PivCoHuffman : public HuffmanBase { return workspace_.data() + (depth % 2) * block_uncompressed_size_; } - /** @brief Number of set bits in a normal (1-bit-per-symbol) node's bitmap. - * @details Only called for non-flat nodes (flat nodes use table lookup, - * not merge). */ - std::size_t bitmap_popcount(const Bitmap& b) const { - const std::uint64_t* p = arena_.data() + b.offset; - std::size_t total = 0; - const std::size_t words = b.word_count(0); - for (std::size_t i = 0; i < words; i++) { - total += std::popcount(p[i]); - } - return total; - } - // --- serialization ------------------------------------------------------- /** @brief Append one block's serialized form to `compressed_`. diff --git a/include/pixie/huffman/pivco_simd.h b/include/pixie/huffman/pivco_simd.h index 53b89e8..1c034ca 100644 --- a/include/pixie/huffman/pivco_simd.h +++ b/include/pixie/huffman/pivco_simd.h @@ -14,14 +14,17 @@ * - `pivco_partition_encode` (top-down encode partition): the inverse — given * per-symbol direction bits, splits a dense input into left/right child * streams and writes the routing bitmap. + * - `pivco_flat_decode` (bottom-up `merge_flat_D`): unpacks fixed-width leaf + * indices and translates them through a flat-subtree symbol table. * * Bitmap convention (matches `pivco_huffman.h`): bits are packed LSB-first * into 64-bit words; bit 0 of word 0 is the first symbol. * * Backends, selected by compiler feature macros (`-march`): - * - AVX-512 VBMI2 (`__AVX512VBMI2__`): byte-granular `vpexpandb`/`vpcompressb` - * operating on 64 bytes per iteration. This is the instruction the paper - * (Sect. 6.6) identifies as mapping directly to the partition kernel. + * - AVX-512 VBMI + VBMI2: byte-granular `vpermb`, `vpmultishiftqb`, + * `vpexpandb`, and `vpcompressb` operating on 64 symbols per iteration. + * These extensions implement vector direction lookup, partition, and the + * full `merge_flat_D` family without changing the wire format. * - AVX2 + SSE4.1: 8-byte-at-a-time `_mm_shuffle_epi8` with 256-entry lookup * tables, mirroring the paper's NEON `vqtbl1q_u8` technique. * - Scalar: the original bit-by-bit / symbol-by-symbol loops. @@ -29,8 +32,8 @@ * Every SIMD backend produces byte-identical output to the scalar fallback * (the operation is a deterministic permutation). The existing round-trip tests * under release and AddressSanitizer serve as the differential oracle for every - * backend the test matrix builds and runs; the AVX-512 VBMI2 backend is - * verified only when built on VBMI2-capable hardware, since there is no runtime + * backend the test matrix builds and runs; the AVX-512 backend is verified only + * when built on VBMI/VBMI2-capable hardware, since there is no runtime * dispatch. */ @@ -90,11 +93,14 @@ inline std::uint8_t extract_byte(const std::uint64_t* bits, } /// @brief Scalar merge tail: process the remaining `< 8` symbols bit-by-bit. +template inline void merge_scalar_tail(std::uint8_t* dst, std::size_t& out_i, const std::uint8_t* left, + std::uint8_t left_symbol, std::size_t& c_left, const std::uint8_t* right, + std::uint8_t right_symbol, std::size_t& c_right, const std::uint64_t* bits, std::size_t i, @@ -102,7 +108,40 @@ inline void merge_scalar_tail(std::uint8_t* dst, for (; i < count; ++i) { const std::uint8_t bit = static_cast((bits[i / 64] >> (i % 64)) & 1u); - dst[out_i++] = bit ? right[c_right++] : left[c_left++]; + if (bit) { + if constexpr (kRightConstant) { + dst[out_i++] = right_symbol; + } else { + dst[out_i++] = right[c_right]; + } + ++c_right; + } else { + if constexpr (kLeftConstant) { + dst[out_i++] = left_symbol; + } else { + dst[out_i++] = left[c_left]; + } + ++c_left; + } + } +} + +/// @brief Write only the routing bitmap for a scalar partition tail. +inline void partition_bitmap_tail(std::uint64_t* bits, + const std::uint8_t* src, + const std::uint8_t* dir, + std::size_t start_i, + std::size_t weight) noexcept { + std::size_t word_idx = start_i / 64; + std::size_t bit_pos = start_i % 64; + for (std::size_t i = start_i; i < weight; ++i) { + if (dir[src[i]]) { + bits[word_idx] |= (1ull << bit_pos); + } + if (++bit_pos == 64) { + bit_pos = 0; + ++word_idx; + } } } @@ -115,6 +154,7 @@ inline void merge_scalar_tail(std::uint8_t* dst, /// (word array); @p start_bit is the absolute bit offset where the /// tail begins (must be a multiple of the full-group width so it is /// already byte-aligned within the word packing). +template inline void partition_scalar_tail(std::uint8_t* left_dst, std::uint8_t* right_dst, std::uint64_t* bits, @@ -131,9 +171,15 @@ inline void partition_scalar_tail(std::uint8_t* left_dst, const std::uint8_t s = src[i]; if (dir[s]) { bits[word_idx] |= (1ull << bit_pos); - right_dst[c_right++] = s; + if constexpr (kWriteRight) { + right_dst[c_right] = s; + } + ++c_right; } else { - left_dst[c_left++] = s; + if constexpr (kWriteLeft) { + left_dst[c_left] = s; + } + ++c_left; } if (++bit_pos == 64) { bit_pos = 0; @@ -234,9 +280,12 @@ inline const ShufTable& compress_left_lut() { return lut; } +template inline void merge_decode_avx2(std::uint8_t* dst, const std::uint8_t* left, + std::uint8_t left_symbol, const std::uint8_t* right, + std::uint8_t right_symbol, const std::uint64_t* bits, std::size_t count) noexcept { const auto& mlut = merge_lut(); @@ -246,12 +295,21 @@ inline void merge_decode_avx2(std::uint8_t* dst, std::size_t i = 0; for (; i + 8 <= count; i += 8) { const std::uint8_t m8 = extract_byte(bits, i); - std::uint64_t lv = 0; - std::uint64_t rv = 0; - std::memcpy(&lv, left + c_left, 8); - std::memcpy(&rv, right + c_right, 8); - const __m128i combined = - _mm_set_epi64x(static_cast(rv), static_cast(lv)); + __m128i left_values; + if constexpr (kLeftConstant) { + left_values = _mm_set1_epi8(static_cast(left_symbol)); + } else { + left_values = + _mm_loadl_epi64(reinterpret_cast(left + c_left)); + } + __m128i right_values; + if constexpr (kRightConstant) { + right_values = _mm_set1_epi8(static_cast(right_symbol)); + } else { + right_values = + _mm_loadl_epi64(reinterpret_cast(right + c_right)); + } + const __m128i combined = _mm_unpacklo_epi64(left_values, right_values); const __m128i out = _mm_shuffle_epi8(combined, load_shuf(mlut, m8)); _mm_storel_epi64(reinterpret_cast<__m128i*>(dst + out_i), out); out_i += 8; @@ -259,7 +317,9 @@ inline void merge_decode_avx2(std::uint8_t* dst, c_right += static_cast(nr); c_left += static_cast(8 - nr); } - merge_scalar_tail(dst, out_i, left, c_left, right, c_right, bits, i, count); + merge_scalar_tail( + dst, out_i, left, left_symbol, c_left, right, right_symbol, c_right, bits, + i, count); } /// @brief Store @p v (low @p n bytes valid) at @p dst without writing past @@ -282,6 +342,7 @@ inline void store_masked_epi64(std::uint8_t* dst, __m128i v, int n) noexcept { } } +template inline void partition_encode_avx2(std::uint8_t* left_dst, std::uint8_t* right_dst, std::uint64_t* bits, @@ -290,8 +351,6 @@ inline void partition_encode_avx2(std::uint8_t* left_dst, std::size_t weight, std::size_t left_weight, std::size_t right_weight) noexcept { - const auto& cr = compress_right_lut(); - const auto& cl = compress_left_lut(); std::uint8_t* bits_bytes = reinterpret_cast(bits); const std::size_t full = weight / 8; @@ -300,30 +359,99 @@ inline void partition_encode_avx2(std::uint8_t* left_dst, std::size_t g = 0; for (; g < full; ++g) { const std::uint8_t* s8 = src + g * 8; - const __m128i data = _mm_loadl_epi64(reinterpret_cast(s8)); const std::uint8_t m8 = build_dir_mask8(s8, dir); // Byte-aligned write: full groups always start on a byte boundary in the // LSB-first bitmap, so byte g holds exactly symbols [8g, 8g+8). bits_bytes[g] = m8; - const __m128i rv = _mm_shuffle_epi8(data, load_shuf(cr, m8)); - const __m128i lv = _mm_shuffle_epi8(data, load_shuf(cl, m8)); - const int nr = std::popcount(static_cast(m8)); - const int nl = 8 - nr; - // Both child regions are contiguous with a later-consumed sibling region, - // so neither store may overflow past its own region end. Masked stores - // keep the trailing invalid lanes inside the region when they would - // otherwise spill into the neighbor. - store_masked_epi64(right_dst + c_right, rv, - static_cast(right_weight - c_right)); - store_masked_epi64(left_dst + c_left, lv, - static_cast(left_weight - c_left)); - c_right += static_cast(nr); - c_left += static_cast(nl); + if constexpr (kWriteLeft || kWriteRight) { + const int nr = std::popcount(static_cast(m8)); + const int nl = 8 - nr; + const __m128i data = + _mm_loadl_epi64(reinterpret_cast(s8)); + if constexpr (kWriteRight) { + const __m128i right_values = + _mm_shuffle_epi8(data, load_shuf(compress_right_lut(), m8)); + store_masked_epi64(right_dst + c_right, right_values, + static_cast(right_weight - c_right)); + } + if constexpr (kWriteLeft) { + const __m128i left_values = + _mm_shuffle_epi8(data, load_shuf(compress_left_lut(), m8)); + store_masked_epi64(left_dst + c_left, left_values, + static_cast(left_weight - c_left)); + } + c_right += static_cast(nr); + c_left += static_cast(nl); + } } // Scalar tail for the trailing < 8 symbols (byte-unaligned bitmap bits). - partition_scalar_tail(left_dst, right_dst, bits, src, dir, c_left, c_right, - full * 8, weight); + if constexpr (kWriteLeft || kWriteRight) { + partition_scalar_tail( + left_dst, right_dst, bits, src, dir, c_left, c_right, full * 8, weight); + } else { + partition_bitmap_tail(bits, src, dir, full * 8, weight); + } +} + +/// @brief Decode a 2-bit flat subtree, 32 symbols per iteration. +inline void flat_decode_d2_avx2(std::uint8_t* dst, + const std::uint8_t* packed, + const std::uint8_t* table, + std::size_t count) noexcept { + alignas(16) std::uint8_t table_lanes[16]{}; + std::memcpy(table_lanes, table, 4); + const __m128i lookup = + _mm_load_si128(reinterpret_cast(table_lanes)); + const __m128i mask = _mm_set1_epi8(3); + + std::size_t i = 0; + for (; i + 32 <= count; i += 32) { + const __m128i data = + _mm_loadl_epi64(reinterpret_cast(packed + i / 4)); + const __m128i code0 = _mm_and_si128(data, mask); + const __m128i code1 = _mm_and_si128(_mm_srli_epi16(data, 2), mask); + const __m128i code2 = _mm_and_si128(_mm_srli_epi16(data, 4), mask); + const __m128i code3 = _mm_and_si128(_mm_srli_epi16(data, 6), mask); + const __m128i code01 = _mm_unpacklo_epi8(code0, code1); + const __m128i code23 = _mm_unpacklo_epi8(code2, code3); + const __m128i indices0 = _mm_unpacklo_epi16(code01, code23); + const __m128i indices1 = _mm_unpackhi_epi16(code01, code23); + _mm_storeu_si128(reinterpret_cast<__m128i*>(dst + i), + _mm_shuffle_epi8(lookup, indices0)); + _mm_storeu_si128(reinterpret_cast<__m128i*>(dst + i + 16), + _mm_shuffle_epi8(lookup, indices1)); + } + for (; i < count; ++i) { + dst[i] = table[(packed[i / 4] >> (2 * (i % 4))) & 3u]; + } +} + +/// @brief Decode a 4-bit flat subtree, 32 symbols per iteration. +inline void flat_decode_d4_avx2(std::uint8_t* dst, + const std::uint8_t* packed, + const std::uint8_t* table, + std::size_t count) noexcept { + const __m128i lookup = + _mm_loadu_si128(reinterpret_cast(table)); + const __m128i mask = _mm_set1_epi8(0x0f); + + std::size_t i = 0; + for (; i + 32 <= count; i += 32) { + const __m128i data = + _mm_loadu_si128(reinterpret_cast(packed + i / 2)); + const __m128i low = _mm_and_si128(data, mask); + const __m128i high = _mm_and_si128(_mm_srli_epi16(data, 4), mask); + const __m128i low_symbols = _mm_shuffle_epi8(lookup, low); + const __m128i high_symbols = _mm_shuffle_epi8(lookup, high); + _mm_storeu_si128(reinterpret_cast<__m128i*>(dst + i), + _mm_unpacklo_epi8(low_symbols, high_symbols)); + _mm_storeu_si128(reinterpret_cast<__m128i*>(dst + i + 16), + _mm_unpackhi_epi8(low_symbols, high_symbols)); + } + for (; i < count; ++i) { + dst[i] = table[(packed[i / 2] >> (4 * (i % 2))) & 0x0fu]; + } } #endif // AVX2 + SSE4.1 @@ -333,9 +461,183 @@ inline void partition_encode_avx2(std::uint8_t* left_dst, // =========================================================================== #if defined(__AVX512VBMI2__) && defined(__AVX512BW__) +#if defined(__AVX512VBMI__) + +/// @brief Four-bank byte lookup used for 64-, 128-, and 256-entry tables. +/// @details `vpermb` addresses one 64-byte bank. The high two index bits select +/// the bank after all required bank permutations run in parallel. +template +struct ByteLookupAvx512 { + __m512i bank0; + __m512i bank1; + __m512i bank2; + __m512i bank3; +}; + +/// @brief Load a `2^kDepth` byte table without reading beyond its allocation. +template +inline ByteLookupAvx512 load_byte_lookup_avx512( + const std::uint8_t* table) noexcept { + static_assert(kDepth >= 1 && kDepth <= 8); + constexpr std::size_t kTableSize = std::size_t{1} << kDepth; + constexpr __mmask64 kFirstBankMask = + kTableSize >= 64 ? ~__mmask64{0} + : (__mmask64{1} << kTableSize) - __mmask64{1}; + const __m512i zero = _mm512_setzero_si512(); + ByteLookupAvx512 lookup{ + _mm512_maskz_loadu_epi8(kFirstBankMask, table), zero, zero, zero}; + if constexpr (kDepth >= 7) { + lookup.bank1 = _mm512_loadu_si512(table + 64); + } + if constexpr (kDepth == 8) { + lookup.bank2 = _mm512_loadu_si512(table + 128); + lookup.bank3 = _mm512_loadu_si512(table + 192); + } + return lookup; +} + +/// @brief Translate 64 byte indices through a preloaded `2^kDepth` table. +template +inline __m512i lookup_bytes_avx512( + __m512i indices, + const ByteLookupAvx512& lookup) noexcept { + const __m512i low_indices = _mm512_and_si512(indices, _mm512_set1_epi8(0x3f)); + __m512i result = _mm512_permutexvar_epi8(low_indices, lookup.bank0); + if constexpr (kDepth >= 7) { + const __m512i banks = + _mm512_and_si512(_mm512_srli_epi16(indices, 6), _mm512_set1_epi8(3)); + const __mmask64 bank1 = _mm512_cmpeq_epi8_mask(banks, _mm512_set1_epi8(1)); + result = _mm512_mask_mov_epi8( + result, bank1, _mm512_permutexvar_epi8(low_indices, lookup.bank1)); + if constexpr (kDepth == 8) { + const __mmask64 bank2 = + _mm512_cmpeq_epi8_mask(banks, _mm512_set1_epi8(2)); + const __mmask64 bank3 = + _mm512_cmpeq_epi8_mask(banks, _mm512_set1_epi8(3)); + result = _mm512_mask_mov_epi8( + result, bank2, _mm512_permutexvar_epi8(low_indices, lookup.bank2)); + result = _mm512_mask_mov_epi8( + result, bank3, _mm512_permutexvar_epi8(low_indices, lookup.bank3)); + } + } + return result; +} + +/// @brief Decode one scalar tail after full 64-symbol flat-D groups. +template +inline void flat_decode_tail_avx512(std::uint8_t* dst, + const std::uint64_t* bits, + const std::uint8_t* table, + std::size_t start, + std::size_t count) noexcept { + constexpr std::uint64_t kMask = (std::uint64_t{1} << kDepth) - 1; + for (std::size_t i = start; i < count; ++i) { + const std::size_t bit_position = i * kDepth; + const std::size_t word_index = bit_position / 64; + const std::size_t bit_offset = bit_position % 64; + std::uint64_t packed = bits[word_index] >> bit_offset; + if (bit_offset + kDepth > 64) { + packed |= bits[word_index + 1] << (64 - bit_offset); + } + dst[i] = table[packed & kMask]; + } +} + +/// @brief AVX-512 implementation of the paper's `merge_flat_D` primitive. +/// @details Each 64-symbol group occupies exactly @p kDepth 64-bit words. The +/// words are permuted into eight 8-code windows, aligned with vector +/// shifts, unpacked by `vpmultishiftqb`, and translated by `vpermb`. +template +inline void flat_decode_avx512(std::uint8_t* dst, + const std::uint64_t* bits, + const std::uint8_t* table, + std::size_t count) noexcept { + static_assert(kDepth >= 2 && kDepth <= 8); + const ByteLookupAvx512 lookup = + load_byte_lookup_avx512(table); + + alignas(64) static constexpr auto kWordIndices = [] { + std::array values{}; + for (std::size_t lane = 0; lane < values.size(); ++lane) { + values[lane] = (lane * 8 * kDepth) / 64; + } + return values; + }(); + alignas(64) static constexpr auto kNextWordIndices = [] { + std::array values{}; + for (std::size_t lane = 0; lane < values.size(); ++lane) { + values[lane] = ((lane * 8 * kDepth) / 64 + 1) % 8; + } + return values; + }(); + alignas(64) static constexpr auto kRightShifts = [] { + std::array values{}; + for (std::size_t lane = 0; lane < values.size(); ++lane) { + values[lane] = (lane * 8 * kDepth) % 64; + } + return values; + }(); + alignas(64) static constexpr auto kLeftShifts = [] { + std::array values{}; + for (std::size_t lane = 0; lane < values.size(); ++lane) { + values[lane] = 64 - ((lane * 8 * kDepth) % 64); + } + return values; + }(); + alignas(64) static constexpr auto kMultishiftControl = [] { + std::array values{}; + for (std::size_t lane = 0; lane < 8; ++lane) { + for (std::size_t code = 0; code < 8; ++code) { + values[lane * 8 + code] = static_cast(code * kDepth); + } + } + return values; + }(); + + const __m512i word_indices = _mm512_load_si512(kWordIndices.data()); + const __m512i next_word_indices = _mm512_load_si512(kNextWordIndices.data()); + const __m512i right_shifts = _mm512_load_si512(kRightShifts.data()); + const __m512i left_shifts = _mm512_load_si512(kLeftShifts.data()); + const __m512i multishift_control = + _mm512_load_si512(kMultishiftControl.data()); + const __m512i code_mask = + _mm512_set1_epi8(static_cast((1u << kDepth) - 1u)); + constexpr __mmask8 kPackedWordMask = + static_cast<__mmask8>((1u << kDepth) - 1u); + + std::size_t i = 0; + for (; i + 64 <= count; i += 64) { + __m512i indices; + if constexpr (kDepth == 8) { + // Eight-bit codes are already byte indices; only lookup is required. + indices = + _mm512_loadu_si512(reinterpret_cast(bits) + i); + } else { + const __m512i packed = + _mm512_maskz_loadu_epi64(kPackedWordMask, bits + (i / 64) * kDepth); + const __m512i current_words = + _mm512_permutexvar_epi64(word_indices, packed); + const __m512i next_words = + _mm512_permutexvar_epi64(next_word_indices, packed); + const __m512i windows = + _mm512_or_si512(_mm512_srlv_epi64(current_words, right_shifts), + _mm512_sllv_epi64(next_words, left_shifts)); + indices = _mm512_and_si512( + _mm512_multishift_epi64_epi8(multishift_control, windows), code_mask); + } + _mm512_storeu_si512(dst + i, lookup_bytes_avx512(indices, lookup)); + } + flat_decode_tail_avx512(dst, bits, table, i, count); +} + +#endif // AVX-512 VBMI + +template inline void merge_decode_avx512(std::uint8_t* dst, const std::uint8_t* left, + std::uint8_t left_symbol, const std::uint8_t* right, + std::uint8_t right_symbol, const std::uint64_t* bits, std::size_t count) noexcept { std::size_t out_i = 0; @@ -344,13 +646,32 @@ inline void merge_decode_avx512(std::uint8_t* dst, const std::size_t full = count / 64; for (std::size_t w = 0; w < full; ++w) { const __mmask64 km = bits[w]; - const __m512i lv = _mm512_loadu_si512(left + c_left); - const __m512i rv = _mm512_loadu_si512(right + c_right); + __m512i left_values; + if constexpr (kLeftConstant) { + left_values = _mm512_set1_epi8(static_cast(left_symbol)); + } else { + left_values = _mm512_loadu_si512(left + c_left); + } + __m512i right_values; + if constexpr (kRightConstant) { + right_values = _mm512_set1_epi8(static_cast(right_symbol)); + } else { + right_values = _mm512_loadu_si512(right + c_right); + } // Fill clear-bit positions with left symbols, then set-bit positions with // right symbols. maskz zeroes the right positions first; the second expand // keeps the left values there. - __m512i d = _mm512_maskz_expand_epi8(~km, lv); - d = _mm512_mask_expand_epi8(d, km, rv); + __m512i d; + if constexpr (kLeftConstant) { + d = left_values; + } else { + d = _mm512_maskz_expand_epi8(~km, left_values); + } + if constexpr (kRightConstant) { + d = _mm512_mask_mov_epi8(d, km, right_values); + } else { + d = _mm512_mask_expand_epi8(d, km, right_values); + } _mm512_storeu_si512(dst + out_i, d); const std::size_t nr = std::popcount(bits[w]); out_i += 64; @@ -360,15 +681,20 @@ inline void merge_decode_avx512(std::uint8_t* dst, // Tail handled by the AVX2 path if available, otherwise scalar. #if defined(__AVX2__) && defined(__SSE4_1__) const std::size_t done = full * 64; - merge_decode_avx2(dst + done, left + c_left, right + c_right, bits + full, - count - done); + const std::uint8_t* tail_left = kLeftConstant ? left : left + c_left; + const std::uint8_t* tail_right = kRightConstant ? right : right + c_right; + merge_decode_avx2( + dst + done, tail_left, left_symbol, tail_right, right_symbol, bits + full, + count - done); (void)out_i; #else - merge_scalar_tail(dst, out_i, left, c_left, right, c_right, bits, full * 64, - count); + merge_scalar_tail( + dst, out_i, left, left_symbol, c_left, right, right_symbol, c_right, bits, + full * 64, count); #endif } +template inline void partition_encode_avx512(std::uint8_t* left_dst, std::uint8_t* right_dst, std::uint64_t* bits, @@ -382,50 +708,54 @@ inline void partition_encode_avx512(std::uint8_t* left_dst, const std::size_t full = weight / 64; std::size_t c_left = 0; std::size_t c_right = 0; +#if defined(__AVX512VBMI__) + const ByteLookupAvx512<8> direction_lookup = load_byte_lookup_avx512<8>(dir); +#endif std::size_t g = 0; for (; g < full; ++g) { const std::uint8_t* s64 = src + g * 64; - const __m512i data = _mm512_loadu_si512(s64); - // Direction mask: bit k = dir[symbol k], LSB-first. x86 has no - // byte-granular gather, so the 64-bit mask is built with a scalar OR loop - // over the dir table. (A qword gather + vpmovb2m would be faster in - // principle but reads all 64 byte sign-bits, corrupting the mask with - // neighboring dir entries.) + __m512i data; __mmask64 km = 0; +#if defined(__AVX512VBMI__) + data = _mm512_loadu_si512(s64); + const __m512i directions = lookup_bytes_avx512(data, direction_lookup); + km = _mm512_cmpneq_epi8_mask(directions, _mm512_setzero_si512()); +#else for (int k = 0; k < 64; ++k) { km |= static_cast<__mmask64>(dir[s64[k]]) << k; } +#endif bits[g] = km; - // Compress: pack right-bound (set-bit) elements to the front of one stream - // and left-bound (clear-bit) elements to the front of another. - const __m512i right_packed = _mm512_maskz_compress_epi8(km, data); - const __m512i left_packed = _mm512_maskz_compress_epi8(~km, data); - const std::size_t nr = std::popcount(km); - const std::size_t nl = 64 - nr; - // Both child regions are contiguous with a later-consumed sibling region, - // so neither store may overflow its own region. Use masked stores that - // write only the valid lanes when the full vector would spill into the - // neighbor. - { - const std::size_t rem = right_weight - c_right; - if (rem >= 64) { - _mm512_storeu_si512(right_dst + c_right, right_packed); - } else { - const __mmask64 rmask = (1ull << rem) - 1; - _mm512_mask_storeu_epi8(right_dst + c_right, rmask, right_packed); + if constexpr (kWriteLeft || kWriteRight) { +#if !defined(__AVX512VBMI__) + data = _mm512_loadu_si512(s64); +#endif + const std::size_t nr = std::popcount(km); + const std::size_t nl = 64 - nr; + if constexpr (kWriteRight) { + const __m512i right_packed = _mm512_maskz_compress_epi8(km, data); + const std::size_t remaining = right_weight - c_right; + if (remaining >= 64) { + _mm512_storeu_si512(right_dst + c_right, right_packed); + } else { + const __mmask64 store_mask = (1ull << remaining) - 1; + _mm512_mask_storeu_epi8(right_dst + c_right, store_mask, + right_packed); + } } - } - { - const std::size_t rem = left_weight - c_left; - if (rem >= 64) { - _mm512_storeu_si512(left_dst + c_left, left_packed); - } else { - const __mmask64 lmask = (1ull << rem) - 1; - _mm512_mask_storeu_epi8(left_dst + c_left, lmask, left_packed); + if constexpr (kWriteLeft) { + const __m512i left_packed = _mm512_maskz_compress_epi8(~km, data); + const std::size_t remaining = left_weight - c_left; + if (remaining >= 64) { + _mm512_storeu_si512(left_dst + c_left, left_packed); + } else { + const __mmask64 store_mask = (1ull << remaining) - 1; + _mm512_mask_storeu_epi8(left_dst + c_left, store_mask, left_packed); + } } + c_right += nr; + c_left += nl; } - c_right += nr; - c_left += nl; } // Tail for the trailing < 64 symbols. `__AVX512VBMI2__` implies `__AVX2__` // on every real target, so delegate to the AVX2 kernel (which itself defers @@ -433,10 +763,11 @@ inline void partition_encode_avx512(std::uint8_t* left_dst, // reachable `#else` branch, so no scalar tail is duplicated here. const std::size_t done_words = g; const std::size_t done_syms = g * 64; - partition_encode_avx2(left_dst + c_left, right_dst + c_right, - bits + done_words, src + done_syms, dir, - weight - done_syms, left_weight - c_left, - right_weight - c_right); + std::uint8_t* tail_left = kWriteLeft ? left_dst + c_left : left_dst; + std::uint8_t* tail_right = kWriteRight ? right_dst + c_right : right_dst; + partition_encode_avx2( + tail_left, tail_right, bits + done_words, src + done_syms, dir, + weight - done_syms, left_weight - c_left, right_weight - c_right); } #endif // AVX-512 VBMI2 @@ -445,9 +776,12 @@ inline void partition_encode_avx512(std::uint8_t* left_dst, // Scalar fallback (no SIMD). // =========================================================================== +template inline void merge_decode_scalar(std::uint8_t* dst, const std::uint8_t* left, + std::uint8_t left_symbol, const std::uint8_t* right, + std::uint8_t right_symbol, const std::uint64_t* bits, std::size_t count) noexcept { std::size_t out_i = 0; @@ -461,24 +795,173 @@ inline void merge_decode_scalar(std::uint8_t* dst, limit = count - w * 64; } for (std::size_t i = 0; i < limit; ++i) { - dst[out_i++] = (word & 1ull) ? right[c_right++] : left[c_left++]; + if (word & 1ull) { + if constexpr (kRightConstant) { + dst[out_i++] = right_symbol; + } else { + dst[out_i++] = right[c_right]; + } + ++c_right; + } else { + if constexpr (kLeftConstant) { + dst[out_i++] = left_symbol; + } else { + dst[out_i++] = left[c_left]; + } + ++c_left; + } word >>= 1; } } } +template inline void partition_encode_scalar(std::uint8_t* left_dst, std::uint8_t* right_dst, std::uint64_t* bits, const std::uint8_t* src, const std::uint8_t* dir, - std::size_t weight, - std::size_t /*left_weight*/, - std::size_t /*right_weight*/) noexcept { + std::size_t weight) noexcept { std::size_t c_left = 0; std::size_t c_right = 0; - partition_scalar_tail(left_dst, right_dst, bits, src, dir, c_left, c_right, 0, - weight); + partition_scalar_tail( + left_dst, right_dst, bits, src, dir, c_left, c_right, 0, weight); +} + +/// @brief Portable fixed-width unpack and lookup for flat subtrees. +inline void flat_decode_scalar(std::uint8_t* dst, + const std::uint64_t* bits, + const std::uint8_t* table, + std::size_t count, + std::uint8_t depth) noexcept { + if (depth == 8) { + const auto* bytes = reinterpret_cast(bits); + for (std::size_t i = 0; i < count; ++i) { + dst[i] = table[bytes[i]]; + } + return; + } + + const std::uint64_t mask = (std::uint64_t{1} << depth) - 1; + std::size_t word_index = 0; + std::size_t bit_offset = 0; + for (std::size_t i = 0; i < count; ++i) { + std::uint64_t packed = bits[word_index] >> bit_offset; + if (bit_offset + depth > 64) { + packed |= bits[word_index + 1] << (64 - bit_offset); + } + dst[i] = table[packed & mask]; + bit_offset += depth; + if (bit_offset >= 64) { + bit_offset -= 64; + ++word_index; + } + } +} + +template +inline void merge_decode_dispatch(std::uint8_t* dst, + const std::uint8_t* left, + std::uint8_t left_symbol, + const std::uint8_t* right, + std::uint8_t right_symbol, + const std::uint64_t* bits, + std::size_t count) noexcept { +#if defined(__AVX512VBMI2__) && defined(__AVX512BW__) + merge_decode_avx512( + dst, left, left_symbol, right, right_symbol, bits, count); +#elif defined(__AVX2__) && defined(__SSE4_1__) + merge_decode_avx2( + dst, left, left_symbol, right, right_symbol, bits, count); +#else + merge_decode_scalar( + dst, left, left_symbol, right, right_symbol, bits, count); +#endif +} + +template +inline void partition_encode_dispatch(std::uint8_t* left_dst, + std::uint8_t* right_dst, + std::uint64_t* bits, + const std::uint8_t* src, + const std::uint8_t* dir, + std::size_t weight, + std::size_t left_weight, + std::size_t right_weight) noexcept { +#if defined(__AVX512VBMI2__) && defined(__AVX512BW__) + partition_encode_avx512( + left_dst, right_dst, bits, src, dir, weight, left_weight, right_weight); +#elif defined(__AVX2__) && defined(__SSE4_1__) + partition_encode_avx2( + left_dst, right_dst, bits, src, dir, weight, left_weight, right_weight); +#else + partition_encode_scalar(left_dst, right_dst, bits, + src, dir, weight); + (void)left_weight; + (void)right_weight; +#endif +} + +/// @brief Dispatch the paper's merge-flat-D primitive family. +inline void flat_decode_dispatch(std::uint8_t* dst, + const std::uint64_t* bits, + const std::uint8_t* table, + std::size_t count, + std::uint8_t depth) noexcept { +#if defined(__AVX512VBMI2__) && defined(__AVX512BW__) + if (depth == 1) { + merge_decode_avx512(dst, nullptr, table[0], nullptr, table[1], + bits, count); + return; + } +#elif defined(__AVX2__) && defined(__SSE4_1__) + if (depth == 1) { + merge_decode_avx2(dst, nullptr, table[0], nullptr, table[1], + bits, count); + return; + } +#endif + +#if defined(__AVX512VBMI2__) && defined(__AVX512BW__) && defined(__AVX512VBMI__) + switch (depth) { + case 2: + flat_decode_avx512<2>(dst, bits, table, count); + return; + case 3: + flat_decode_avx512<3>(dst, bits, table, count); + return; + case 4: + flat_decode_avx512<4>(dst, bits, table, count); + return; + case 5: + flat_decode_avx512<5>(dst, bits, table, count); + return; + case 6: + flat_decode_avx512<6>(dst, bits, table, count); + return; + case 7: + flat_decode_avx512<7>(dst, bits, table, count); + return; + case 8: + flat_decode_avx512<8>(dst, bits, table, count); + return; + default: + break; + } +#endif + +#if defined(__AVX2__) && defined(__SSE4_1__) + const auto* packed = reinterpret_cast(bits); + if (depth == 2) { + flat_decode_d2_avx2(dst, packed, table, count); + return; + } + if (depth == 4) { + flat_decode_d4_avx2(dst, packed, table, count); + return; + } +#endif + flat_decode_scalar(dst, bits, table, count, depth); } } // namespace pixie::pivco_simd_detail @@ -505,13 +988,38 @@ inline void pivco_merge_decode(std::uint8_t* dst, const std::uint8_t* right, const std::uint64_t* bits, std::size_t count) noexcept { -#if defined(__AVX512VBMI2__) && defined(__AVX512BW__) - pivco_simd_detail::merge_decode_avx512(dst, left, right, bits, count); -#elif defined(__AVX2__) && defined(__SSE4_1__) - pivco_simd_detail::merge_decode_avx2(dst, left, right, bits, count); -#else - pivco_simd_detail::merge_decode_scalar(dst, left, right, bits, count); -#endif + pivco_simd_detail::merge_decode_dispatch(dst, left, 0, right, 0, + bits, count); +} + +/// @brief Merge two constant-symbol leaves under a routing bitmap. +inline void pivco_merge_decode_cst_cst(std::uint8_t* dst, + std::uint8_t left_symbol, + std::uint8_t right_symbol, + const std::uint64_t* bits, + std::size_t count) noexcept { + pivco_simd_detail::merge_decode_dispatch( + dst, nullptr, left_symbol, nullptr, right_symbol, bits, count); +} + +/// @brief Merge a constant left leaf with a dense right child stream. +inline void pivco_merge_decode_cst_vec(std::uint8_t* dst, + std::uint8_t left_symbol, + const std::uint8_t* right, + const std::uint64_t* bits, + std::size_t count) noexcept { + pivco_simd_detail::merge_decode_dispatch( + dst, nullptr, left_symbol, right, 0, bits, count); +} + +/// @brief Merge a dense left child stream with a constant right leaf. +inline void pivco_merge_decode_vec_cst(std::uint8_t* dst, + const std::uint8_t* left, + std::uint8_t right_symbol, + const std::uint64_t* bits, + std::size_t count) noexcept { + pivco_simd_detail::merge_decode_dispatch( + dst, left, 0, nullptr, right_symbol, bits, count); } /// @brief Top-down encode partition: split @p src into left/right child streams @@ -538,16 +1046,48 @@ inline void pivco_partition_encode(std::uint8_t* left_dst, std::size_t weight, std::size_t left_weight, std::size_t right_weight) noexcept { -#if defined(__AVX512VBMI2__) && defined(__AVX512BW__) - pivco_simd_detail::partition_encode_avx512( - left_dst, right_dst, bits, src, dir, weight, left_weight, right_weight); -#elif defined(__AVX2__) && defined(__SSE4_1__) - pivco_simd_detail::partition_encode_avx2(left_dst, right_dst, bits, src, dir, - weight, left_weight, right_weight); -#else - pivco_simd_detail::partition_encode_scalar( + pivco_simd_detail::partition_encode_dispatch( left_dst, right_dst, bits, src, dir, weight, left_weight, right_weight); -#endif +} + +/// @brief Build a routing bitmap and materialize only the left child stream. +inline void pivco_partition_encode_left(std::uint8_t* left_dst, + std::uint64_t* bits, + const std::uint8_t* src, + const std::uint8_t* dir, + std::size_t weight, + std::size_t left_weight) noexcept { + pivco_simd_detail::partition_encode_dispatch( + left_dst, nullptr, bits, src, dir, weight, left_weight, 0); +} + +/// @brief Build a routing bitmap and materialize only the right child stream. +inline void pivco_partition_encode_right(std::uint8_t* right_dst, + std::uint64_t* bits, + const std::uint8_t* src, + const std::uint8_t* dir, + std::size_t weight, + std::size_t right_weight) noexcept { + pivco_simd_detail::partition_encode_dispatch( + nullptr, right_dst, bits, src, dir, weight, 0, right_weight); +} + +/// @brief Build only the routing bitmap when both children are leaves. +inline void pivco_partition_encode_none(std::uint64_t* bits, + const std::uint8_t* src, + const std::uint8_t* dir, + std::size_t weight) noexcept { + pivco_simd_detail::partition_encode_dispatch( + nullptr, nullptr, bits, src, dir, weight, 0, 0); +} + +/// @brief Unpack and translate a flat subtree's fixed-width wire codes. +inline void pivco_flat_decode(std::uint8_t* dst, + const std::uint64_t* bits, + const std::uint8_t* table, + std::size_t count, + std::uint8_t depth) noexcept { + pivco_simd_detail::flat_decode_dispatch(dst, bits, table, count, depth); } } // namespace pixie diff --git a/src/tests/pivco_huffman_tests.cpp b/src/tests/pivco_huffman_tests.cpp index 52d7fcb..d9883f8 100644 --- a/src/tests/pivco_huffman_tests.cpp +++ b/src/tests/pivco_huffman_tests.cpp @@ -1,6 +1,8 @@ #include #include +#include +#include #include #include #include @@ -79,6 +81,112 @@ TEST(PivCoHuffmanSmoke, FullAlphabetRoundTrips) { EXPECT_EQ(decode_via_bytes(codec), data); } +TEST(PivCoHuffmanSmoke, FlatPowerOfTwoAlphabetsRoundTrip) { + std::mt19937 rng(1871); + for (std::uint8_t depth = 1; depth <= 8; ++depth) { + const std::size_t alphabet_size = std::size_t{1} << depth; + std::vector data; + data.reserve(alphabet_size * 16); + for (std::size_t repetition = 0; repetition < 16; ++repetition) { + for (std::size_t symbol = 0; symbol < alphabet_size; ++symbol) { + data.push_back(static_cast(symbol)); + } + } + std::shuffle(data.begin(), data.end(), rng); + + const PivCoHuffman codec(data); + EXPECT_EQ(decode_same(codec), data) << "flat depth " << +depth; + EXPECT_EQ(decode_via_bytes(codec), data) << "flat depth " << +depth; + } +} + +TEST(PivCoSimd, FlatKernelsDecodeFullGroupsAndTails) { + std::mt19937 rng(4189); + for (std::uint8_t depth = 1; depth <= 8; ++depth) { + const std::size_t alphabet_size = std::size_t{1} << depth; + const std::size_t count = 129 + depth; + std::vector table(alphabet_size); + for (std::size_t i = 0; i < alphabet_size; ++i) { + table[i] = static_cast(i); + } + std::shuffle(table.begin(), table.end(), rng); + + std::vector bits((count * depth + 63) / 64, 0); + std::vector expected(count); + for (std::size_t i = 0; i < count; ++i) { + const std::uint64_t code = rng() % alphabet_size; + const std::size_t bit_position = i * depth; + const std::size_t word_index = bit_position / 64; + const std::size_t bit_offset = bit_position % 64; + bits[word_index] |= code << bit_offset; + if (bit_offset + depth > 64) { + bits[word_index + 1] |= code >> (64 - bit_offset); + } + expected[i] = table[code]; + } + + std::vector decoded(count); + pixie::pivco_flat_decode(decoded.data(), bits.data(), table.data(), count, + depth); + EXPECT_EQ(decoded, expected) << "flat depth " << +depth; + } +} + +TEST(PivCoSimd, PartitionVariantsBuildTheSameDirectionMask) { + constexpr std::size_t kCount = 257; + std::array directions{}; + for (std::size_t symbol = 0; symbol < directions.size(); ++symbol) { + directions[symbol] = + static_cast(((symbol * 73) ^ (symbol >> 2)) & 1u); + } + + std::mt19937 rng(7127); + std::vector input(kCount); + std::vector expected_left; + std::vector expected_right; + std::vector expected_bits((kCount + 63) / 64, 0); + for (std::size_t i = 0; i < input.size(); ++i) { + input[i] = static_cast(rng()); + if (directions[input[i]]) { + expected_bits[i / 64] |= std::uint64_t{1} << (i % 64); + expected_right.push_back(input[i]); + } else { + expected_left.push_back(input[i]); + } + } + + std::vector left(expected_left.size()); + std::vector right(expected_right.size()); + std::vector full_bits(expected_bits.size(), 0); + pixie::pivco_partition_encode(left.data(), right.data(), full_bits.data(), + input.data(), directions.data(), input.size(), + left.size(), right.size()); + EXPECT_EQ(left, expected_left); + EXPECT_EQ(right, expected_right); + EXPECT_EQ(full_bits, expected_bits); + + std::vector left_bits(expected_bits.size(), 0); + std::fill(left.begin(), left.end(), 0); + pixie::pivco_partition_encode_left(left.data(), left_bits.data(), + input.data(), directions.data(), + input.size(), left.size()); + EXPECT_EQ(left, expected_left); + EXPECT_EQ(left_bits, expected_bits); + + std::vector right_bits(expected_bits.size(), 0); + std::fill(right.begin(), right.end(), 0); + pixie::pivco_partition_encode_right(right.data(), right_bits.data(), + input.data(), directions.data(), + input.size(), right.size()); + EXPECT_EQ(right, expected_right); + EXPECT_EQ(right_bits, expected_bits); + + std::vector bitmap_only(expected_bits.size(), 0); + pixie::pivco_partition_encode_none(bitmap_only.data(), input.data(), + directions.data(), input.size()); + EXPECT_EQ(bitmap_only, expected_bits); +} + TEST(PivCoHuffmanSmoke, RandomUniformRoundTrips) { std::mt19937 rng(239); std::uniform_int_distribution dist(0, 255); From 33ce79ee78ea472c539a5a13cb7a4bab4382fc0f Mon Sep 17 00:00:00 2001 From: Sanya239 Date: Mon, 27 Jul 2026 18:42:59 +0300 Subject: [PATCH 12/14] More AVX2 optimizations --- include/pixie/huffman/pivco_huffman.h | 66 +- include/pixie/huffman/pivco_simd.h | 675 +++++++++++++++++++- src/benchmarks/pivco_huffman_benchmarks.cpp | 25 + 3 files changed, 728 insertions(+), 38 deletions(-) diff --git a/include/pixie/huffman/pivco_huffman.h b/include/pixie/huffman/pivco_huffman.h index 15ef355..f3143b4 100644 --- a/include/pixie/huffman/pivco_huffman.h +++ b/include/pixie/huffman/pivco_huffman.h @@ -173,6 +173,8 @@ class PivCoHuffman : public HuffmanBase { /// @brief 0 = normal node (1 bit/symbol); >0 = flat subtree of this depth /// (stores `flat_depth` bits/symbol + a lookup table). std::uint8_t flat_depth = 0; + /// @brief Whether this flat table is the balanced bit-reversal mapping. + bool flat_bitreverse = false; /// @brief Offset of this flat node's decode table in `flat_tables_`. std::size_t flat_table_offset = 0; }; @@ -534,8 +536,8 @@ class PivCoHuffman : public HuffmanBase { n.flat_depth = depth; n.flat_table_offset = flat_tables_.size(); flat_tables_.resize(flat_tables_.size() + (std::size_t{1} << depth)); - build_flat_table(idx, depth, 0, - flat_tables_.data() + n.flat_table_offset); + n.flat_bitreverse = build_flat_table( + idx, depth, 0, flat_tables_.data() + n.flat_table_offset); return; } @@ -548,8 +550,10 @@ class PivCoHuffman : public HuffmanBase { * @param flat_depth Total depth of the flat subtree. * @param code Suffix accumulated so far (MSB-first: left=0, * right=1). - * @param table Output table of size `2^total_depth`. */ - void build_flat_table(std::size_t idx, + * @param table Output table of size `2^total_depth`. + * @return True when the completed table is the balanced bit-reversal + * permutation used by the specialized AVX2 kernels. */ + bool build_flat_table(std::size_t idx, std::uint8_t flat_depth, std::uint32_t code, symbol_type* table) const { @@ -557,10 +561,13 @@ class PivCoHuffman : public HuffmanBase { if (n.is_leaf) { table[reverse_low_bits(static_cast(code), flat_depth)] = n.symbol; - return; + return n.symbol == code; } - build_flat_table(n.left, flat_depth, code << 1, table); - build_flat_table(n.right, flat_depth, (code << 1) | 1, table); + const bool left_bitreverse = + build_flat_table(n.left, flat_depth, code << 1, table); + const bool right_bitreverse = + build_flat_table(n.right, flat_depth, (code << 1) | 1, table); + return left_bitreverse && right_bitreverse; } /** @brief Total frequency weight of a subtree (sum of leaf frequencies). */ @@ -828,13 +835,35 @@ class PivCoHuffman : public HuffmanBase { const symbol_type* src, std::size_t count, std::uint8_t flat_depth, - const symbol_type* table) { - std::array wire_code; + const symbol_type* table, + bool known_bitreverse) { + // Three zero pad bytes make overlapping dword gathers at symbol 255 safe + // in the AVX2 flat encoder. + std::array wire_code{}; const std::size_t alphabet_size = std::size_t{1} << flat_depth; for (std::size_t code = 0; code < alphabet_size; ++code) { wire_code[table[code]] = static_cast(code); } +#if defined(__AVX2__) && defined(__SSE4_1__) && !defined(__AVX512VBMI2__) + if (flat_depth == 8 && + pivco_flat_encode_d8_bitreverse(bits, src, wire_code.data(), count, + known_bitreverse)) { + return; + } +#if defined(__BMI2__) + if (flat_depth >= 2 && flat_depth <= 7 && + pivco_flat_encode_bitreverse(bits, src, wire_code.data(), count, + flat_depth, known_bitreverse)) { + return; + } + if (flat_depth == 3) { + pivco_flat_encode_d3(bits, src, wire_code.data(), count); + return; + } +#endif +#endif + auto* bytes = reinterpret_cast(bits); if (flat_depth == 8) { for (std::size_t i = 0; i < count; ++i) { @@ -921,7 +950,8 @@ class PivCoHuffman : public HuffmanBase { n.bits.count); } else { const symbol_type* table = flat_tables_.data() + n.flat_table_offset; - encode_flat_values(bits, src, n.bits.count, n.flat_depth, table); + encode_flat_values(bits, src, n.bits.count, n.flat_depth, table, + n.flat_bitreverse); } return; } @@ -939,13 +969,24 @@ class PivCoHuffman : public HuffmanBase { const bool left_is_leaf = nodes_[n.left].is_leaf; const bool right_is_leaf = nodes_[n.right].is_leaf; if (left_is_leaf && right_is_leaf) { - pivco_partition_encode_none(bits, src, dir, weight); + pivco_partition_encode_cst_cst(bits, src, nodes_[n.right].symbol, weight); } else if (right_is_leaf) { +#if defined(__AVX2__) && defined(__SSE4_1__) && !defined(__AVX512VBMI2__) + pivco_partition_encode_vec_cst( + left_dst, bits, src, nodes_[n.right].symbol, weight, left_weight); +#else pivco_partition_encode_left(left_dst, bits, src, dir, weight, left_weight); +#endif } else if (left_is_leaf) { +#if defined(__AVX2__) && defined(__SSE4_1__) && !defined(__AVX512VBMI2__) + pivco_partition_encode_cst_vec(right_dst, bits, src, + nodes_[n.left].symbol, weight, + weight - left_weight); +#else pivco_partition_encode_right(right_dst, bits, src, dir, weight, weight - left_weight); +#endif } else { pivco_partition_encode(left_dst, right_dst, bits, src, dir, weight, left_weight, weight - left_weight); @@ -1003,7 +1044,8 @@ class PivCoHuffman : public HuffmanBase { if (n.flat_depth > 0) { const std::uint64_t* bits = arena_.data() + n.bits.offset; const symbol_type* table = flat_tables_.data() + n.flat_table_offset; - pivco_flat_decode(dst, bits, table, weight, n.flat_depth); + pivco_flat_decode(dst, bits, table, weight, n.flat_depth, + n.flat_bitreverse); return; } diff --git a/include/pixie/huffman/pivco_simd.h b/include/pixie/huffman/pivco_simd.h index 1c034ca..165ff67 100644 --- a/include/pixie/huffman/pivco_simd.h +++ b/include/pixie/huffman/pivco_simd.h @@ -77,21 +77,6 @@ inline std::uint8_t build_dir_mask8(const std::uint8_t* s8, (dir[s8[6]] << 6) | (dir[s8[7]] << 7)); } -/// @brief Read 8 bits starting at bit position @p pos from a packed word array. -/// @pre @p pos + 8 is within the bitmap's bit range; callers stop the SIMD -/// loop once fewer than 8 symbols remain, so the span never exceeds the -/// allocated words. -inline std::uint8_t extract_byte(const std::uint64_t* bits, - std::size_t pos) noexcept { - const std::size_t w = pos / 64; - const std::size_t b = pos % 64; - std::uint64_t v = bits[w] >> b; - if (b + 8 > 64) { - v |= bits[w + 1] << (64 - b); - } - return static_cast(v); -} - /// @brief Scalar merge tail: process the remaining `< 8` symbols bit-by-bit. template inline void merge_scalar_tail(std::uint8_t* dst, @@ -280,6 +265,48 @@ inline const ShufTable& compress_left_lut() { return lut; } +/// @brief Merge two constant leaves 32 symbols at a time with AVX2. +/// @details Four routing bytes are widened to four qwords, each byte is +/// broadcast over its eight output lanes, and a fixed bit-position +/// vector expands the bitmap without a shuffle-table lookup. +inline void merge_decode_cst_cst_avx2(std::uint8_t* dst, + std::uint8_t left_symbol, + std::uint8_t right_symbol, + const std::uint64_t* bits, + std::size_t count) noexcept { + const auto* bit_bytes = reinterpret_cast(bits); + const __m256i broadcast_control = + _mm256_setr_epi8(0, 0, 0, 0, 0, 0, 0, 0, 8, 8, 8, 8, 8, 8, 8, 8, 0, 0, 0, + 0, 0, 0, 0, 0, 8, 8, 8, 8, 8, 8, 8, 8); + const __m256i bit_positions = _mm256_setr_epi8( + 1, 2, 4, 8, 16, 32, 64, static_cast(0x80), 1, 2, 4, 8, 16, 32, 64, + static_cast(0x80), 1, 2, 4, 8, 16, 32, 64, static_cast(0x80), + 1, 2, 4, 8, 16, 32, 64, static_cast(0x80)); + const __m256i zero = _mm256_setzero_si256(); + const __m256i left = _mm256_set1_epi8(static_cast(left_symbol)); + const __m256i delta = + _mm256_set1_epi8(static_cast(left_symbol ^ right_symbol)); + + std::size_t i = 0; + for (; i + 32 <= count; i += 32) { + std::uint32_t routing = 0; + std::memcpy(&routing, bit_bytes + i / 8, sizeof(routing)); + const __m128i routing_bytes = _mm_cvtsi32_si128(static_cast(routing)); + const __m256i routing_qwords = _mm256_cvtepu8_epi64(routing_bytes); + const __m256i broadcast = + _mm256_shuffle_epi8(routing_qwords, broadcast_control); + const __m256i clear_bits = + _mm256_cmpeq_epi8(_mm256_and_si256(broadcast, bit_positions), zero); + const __m256i symbols = + _mm256_xor_si256(left, _mm256_andnot_si256(clear_bits, delta)); + _mm256_storeu_si256(reinterpret_cast<__m256i*>(dst + i), symbols); + } + for (; i < count; ++i) { + dst[i] = + ((bits[i / 64] >> (i % 64)) & 1u) != 0 ? right_symbol : left_symbol; + } +} + template inline void merge_decode_avx2(std::uint8_t* dst, const std::uint8_t* left, @@ -288,13 +315,22 @@ inline void merge_decode_avx2(std::uint8_t* dst, std::uint8_t right_symbol, const std::uint64_t* bits, std::size_t count) noexcept { + if constexpr (kLeftConstant && kRightConstant) { + merge_decode_cst_cst_avx2(dst, left_symbol, right_symbol, bits, count); + return; + } + const auto& mlut = merge_lut(); + const auto* bit_bytes = reinterpret_cast(bits); std::size_t out_i = 0; std::size_t c_left = 0; std::size_t c_right = 0; std::size_t i = 0; for (; i + 8 <= count; i += 8) { - const std::uint8_t m8 = extract_byte(bits, i); + // AVX2 is little-endian on every supported x86 target. Eight-symbol + // groups are byte-aligned in the LSB-first bitmap, so loading the routing + // byte directly avoids a word-index/shift sequence in every merge group. + const std::uint8_t m8 = bit_bytes[i / 8]; __m128i left_values; if constexpr (kLeftConstant) { left_values = _mm_set1_epi8(static_cast(left_symbol)); @@ -394,6 +430,60 @@ inline void partition_encode_avx2(std::uint8_t* left_dst, } } +/// @brief AVX2 partition when one child is a constant leaf. +/// @tparam kLeftConstant True when the left child is the constant; the dense +/// right stream is written. False writes the dense left +/// stream and treats the right child as constant. +template +inline void partition_encode_one_constant_avx2( + std::uint8_t* output, + std::uint64_t* bits, + const std::uint8_t* src, + std::uint8_t constant_symbol, + std::size_t weight, + std::size_t output_weight) noexcept { + auto* bit_bytes = reinterpret_cast(bits); + const __m128i constant = _mm_set1_epi8(static_cast(constant_symbol)); + std::size_t output_count = 0; + const std::size_t full = weight / 8; + std::size_t group = 0; + for (; group < full; ++group) { + const std::uint8_t* s8 = src + group * 8; + const __m128i data = _mm_loadl_epi64(reinterpret_cast(s8)); + const std::uint8_t equal_mask = static_cast( + _mm_movemask_epi8(_mm_cmpeq_epi8(data, constant))); + const std::uint8_t routing = + kLeftConstant ? static_cast(~equal_mask) : equal_mask; + bit_bytes[group] = routing; + + const auto& lut = + kLeftConstant ? compress_right_lut() : compress_left_lut(); + const __m128i packed = _mm_shuffle_epi8(data, load_shuf(lut, routing)); + store_masked_epi64(output + output_count, packed, + static_cast(output_weight - output_count)); + const std::size_t right_count = + std::popcount(static_cast(routing)); + output_count += kLeftConstant ? right_count : 8 - right_count; + } + + std::size_t word_index = (group * 8) / 64; + std::size_t bit_position = (group * 8) % 64; + for (std::size_t i = group * 8; i < weight; ++i) { + const bool is_constant = src[i] == constant_symbol; + const bool goes_right = kLeftConstant ? !is_constant : is_constant; + if (goes_right) { + bits[word_index] |= std::uint64_t{1} << bit_position; + } + if (kLeftConstant ? goes_right : !goes_right) { + output[output_count++] = src[i]; + } + if (++bit_position == 64) { + bit_position = 0; + ++word_index; + } + } +} + /// @brief Decode a 2-bit flat subtree, 32 symbols per iteration. inline void flat_decode_d2_avx2(std::uint8_t* dst, const std::uint8_t* packed, @@ -454,6 +544,374 @@ inline void flat_decode_d4_avx2(std::uint8_t* dst, } } +/// @brief Reverse the low @p width bits of one byte. +inline std::uint8_t reverse_low_bits_avx2(std::uint8_t value, + std::uint8_t width) noexcept { + value = static_cast(((value & 0x55u) << 1) | + ((value & 0xaau) >> 1)); + value = static_cast(((value & 0x33u) << 2) | + ((value & 0xccu) >> 2)); + value = static_cast((value << 4) | (value >> 4)); + return value >> (8 - width); +} + +/// @brief Decode a known full-alphabet depth-8 bit-reversal table. +/// @details A balanced tree built from symbols 0..255 stores +/// `table[reverse8(code)] = code`. AVX2 can apply that permutation +/// with two nibble `pshufb` operations, avoiding the arbitrary byte +/// lookup that AVX2 cannot otherwise vectorize efficiently. +inline void flat_decode_d8_bitreverse_unchecked_avx2( + std::uint8_t* dst, + const std::uint8_t* packed, + const std::uint8_t* table, + std::size_t count) noexcept { + alignas(16) static constexpr std::uint8_t kReverseNibble[16] = { + 0, 8, 4, 12, 2, 10, 6, 14, 1, 9, 5, 13, 3, 11, 7, 15}; + const __m128i reverse128 = + _mm_load_si128(reinterpret_cast(kReverseNibble)); + const __m256i reverse = _mm256_broadcastsi128_si256(reverse128); + const __m256i nibble_mask = _mm256_set1_epi8(0x0f); + std::size_t i = 0; + for (; i + 32 <= count; i += 32) { + const __m256i codes = + _mm256_loadu_si256(reinterpret_cast(packed + i)); + const __m256i low = _mm256_and_si256(codes, nibble_mask); + const __m256i high = + _mm256_and_si256(_mm256_srli_epi16(codes, 4), nibble_mask); + const __m256i reversed_low = _mm256_shuffle_epi8(reverse, low); + const __m256i reversed_high = _mm256_shuffle_epi8(reverse, high); + const __m256i symbols = + _mm256_or_si256(_mm256_slli_epi16(reversed_low, 4), reversed_high); + _mm256_storeu_si256(reinterpret_cast<__m256i*>(dst + i), symbols); + } + for (; i < count; ++i) { + dst[i] = table[packed[i]]; + } +} + +/// @brief Decode the common full-alphabet depth-8 bit-reversal table. +/// @returns True when @p table has the reshaped tree's bit-reversal layout and +/// the output was produced; false for an arbitrary 256-byte table. +inline bool flat_decode_d8_bitreverse_avx2(std::uint8_t* dst, + const std::uint8_t* packed, + const std::uint8_t* table, + std::size_t count) noexcept { + for (std::size_t code = 0; code < 256; ++code) { + if (table[code] != + reverse_low_bits_avx2(static_cast(code), 8)) { + return false; + } + } + flat_decode_d8_bitreverse_unchecked_avx2(dst, packed, table, count); + return true; +} + +/// @brief Expand eight tightly packed D-bit codes into eight byte lanes. +/// @details Eight codes occupy exactly D bytes, so every eight-code group is +/// byte-aligned. BMI2 `pdep` inserts the padding bits between codes in +/// one instruction. The constexpr fallback is fully unrolled for +/// AVX2 CPUs that do not also expose BMI2. +template +inline std::uint64_t expand_codes8_avx2(const std::uint8_t* packed) noexcept { + static_assert(kDepth >= 3 && kDepth <= 8); + std::uint64_t dense = 0; + std::memcpy(&dense, packed, kDepth); + if constexpr (kDepth == 8) { + return dense; + } + + constexpr std::uint64_t kCodeMask = (std::uint64_t{1} << kDepth) - 1; +#if defined(__BMI2__) + constexpr std::uint64_t kByteLaneMask = + kCodeMask * UINT64_C(0x0101010101010101); + return _pdep_u64(dense, kByteLaneMask); +#else + std::uint64_t expanded = 0; + for (std::size_t lane = 0; lane < 8; ++lane) { + expanded |= ((dense >> (lane * kDepth)) & kCodeMask) << (lane * 8); + } + return expanded; +#endif +} + +/// @brief Reverse the low D bits independently in 16 byte lanes. +template +inline __m128i reverse_low_bits16_avx2(__m128i values) noexcept { + static_assert(kDepth >= 2 && kDepth <= 7); + alignas(16) static constexpr std::uint8_t kReverseNibble[16] = { + 0, 8, 4, 12, 2, 10, 6, 14, 1, 9, 5, 13, 3, 11, 7, 15}; + const __m128i reverse = + _mm_load_si128(reinterpret_cast(kReverseNibble)); + const __m128i nibble_mask = _mm_set1_epi8(0x0f); + const __m128i low = _mm_and_si128(values, nibble_mask); + const __m128i high = _mm_and_si128(_mm_srli_epi16(values, 4), nibble_mask); + const __m128i reversed = + _mm_or_si128(_mm_slli_epi16(_mm_shuffle_epi8(reverse, low), 4), + _mm_shuffle_epi8(reverse, high)); + constexpr std::uint8_t kCodeMask = (std::uint8_t{1} << kDepth) - 1; + return _mm_and_si128(_mm_srli_epi16(reversed, 8 - kDepth), + _mm_set1_epi8(static_cast(kCodeMask))); +} + +/// @brief Decode a known balanced low-alphabet table as bit reversal. +template +inline void flat_decode_bitreverse_unchecked_avx2(std::uint8_t* dst, + const std::uint8_t* packed, + const std::uint8_t* table, + std::size_t count) noexcept { + static_assert(kDepth >= 6 && kDepth <= 7); + std::size_t i = 0; + for (; i + 16 <= count; i += 16) { + const std::size_t byte_offset = (i * kDepth) / 8; + const std::uint64_t codes0 = + expand_codes8_avx2(packed + byte_offset); + const std::uint64_t codes1 = + expand_codes8_avx2(packed + byte_offset + kDepth); + const __m128i codes = _mm_set_epi64x(static_cast(codes1), + static_cast(codes0)); + _mm_storeu_si128(reinterpret_cast<__m128i*>(dst + i), + reverse_low_bits16_avx2(codes)); + } + constexpr std::uint64_t kCodeMask = (std::uint64_t{1} << kDepth) - 1; + const auto* words = reinterpret_cast(packed); + for (; i < count; ++i) { + const std::size_t bit_position = i * kDepth; + const std::size_t word_index = bit_position / 64; + const std::size_t bit_offset = bit_position % 64; + std::uint64_t code = words[word_index] >> bit_offset; + if (bit_offset + kDepth > 64) { + code |= words[word_index + 1] << (64 - bit_offset); + } + dst[i] = table[code & kCodeMask]; + } +} + +/// @brief Decode a balanced low-alphabet flat table as bit reversal. +/// @returns True when the table has the expected permutation and was decoded. +template +inline bool flat_decode_bitreverse_avx2(std::uint8_t* dst, + const std::uint8_t* packed, + const std::uint8_t* table, + std::size_t count) noexcept { + static_assert(kDepth >= 6 && kDepth <= 7); + constexpr std::size_t kTableSize = std::size_t{1} << kDepth; + for (std::size_t code = 0; code < kTableSize; ++code) { + if (table[code] != + reverse_low_bits_avx2(static_cast(code), kDepth)) { + return false; + } + } + flat_decode_bitreverse_unchecked_avx2(dst, packed, table, count); + return true; +} + +/// @brief Lookup byte indices in a power-of-two collection of 16-byte banks. +/// @details AVX2 has no arbitrary byte permutation across 256 entries. Each +/// leaf performs one `pshufb` using the low nibble; balanced +/// `blendv` levels select the bank with the remaining high bits. +/// The recursion is compile-time and emits straight-line SIMD code. +template +inline __m128i lookup_banked_bytes_avx2(const std::uint8_t* table, + __m128i low_nibbles, + __m128i codes) noexcept { + static_assert(kBankCount > 0 && std::has_single_bit(kBankCount)); + if constexpr (kBankCount == 1) { + const __m128i bank = _mm_loadu_si128( + reinterpret_cast(table + kFirstBank * 16)); + return _mm_shuffle_epi8(bank, low_nibbles); + } else { + constexpr std::size_t kHalf = kBankCount / 2; + constexpr int kSelectBit = 3 + std::countr_zero(kBankCount); + const __m128i low = + lookup_banked_bytes_avx2(table, low_nibbles, codes); + const __m128i high = lookup_banked_bytes_avx2( + table, low_nibbles, codes); + const __m128i select = _mm_slli_epi16(codes, 7 - kSelectBit); + return _mm_blendv_epi8(low, high, select); + } +} + +/// @brief Decode flat depths 3 and 5..7 in 16-symbol AVX2 batches. +/// @details The existing depth-2/depth-4 kernels unpack 32 symbols per batch +/// and remain faster for those widths. Other depths use BMI2 to +/// expand two groups of eight packed codes, followed by a banked +/// `pshufb` lookup. This avoids the former per-symbol shift, branch, +/// and dependent table load loop. +template +inline void flat_decode_banked_avx2(std::uint8_t* dst, + const std::uint8_t* packed, + const std::uint8_t* table, + std::size_t count) noexcept { + static_assert(kDepth == 3 || (kDepth >= 5 && kDepth <= 7)); + constexpr std::size_t kTableSize = std::size_t{1} << kDepth; + constexpr std::size_t kBankCount = kTableSize <= 16 ? 1 : kTableSize / 16; + const __m128i nibble_mask = _mm_set1_epi8(0x0f); + + // A depth-3 table has only eight valid bytes, while the banked lookup loads + // a complete vector. Pad that one small table locally; depths >=5 already + // contain one or more complete 16-byte banks. + alignas(16) std::uint8_t padded_depth3[16]{}; + const std::uint8_t* lookup_table = table; + if constexpr (kDepth == 3) { + std::memcpy(padded_depth3, table, kTableSize); + lookup_table = padded_depth3; + } + + std::size_t i = 0; + for (; i + 16 <= count; i += 16) { + const std::size_t byte_offset = (i * kDepth) / 8; + const std::uint64_t codes0 = + expand_codes8_avx2(packed + byte_offset); + const std::uint64_t codes1 = + expand_codes8_avx2(packed + byte_offset + kDepth); + const __m128i codes = _mm_set_epi64x(static_cast(codes1), + static_cast(codes0)); + const __m128i low_nibbles = _mm_and_si128(codes, nibble_mask); + const __m128i symbols = lookup_banked_bytes_avx2<0, kBankCount>( + lookup_table, low_nibbles, codes); + _mm_storeu_si128(reinterpret_cast<__m128i*>(dst + i), symbols); + } + + const auto* words = reinterpret_cast(packed); + constexpr std::uint64_t kCodeMask = (std::uint64_t{1} << kDepth) - 1; + for (; i < count; ++i) { + const std::size_t bit_position = i * kDepth; + const std::size_t word_index = bit_position / 64; + const std::size_t bit_offset = bit_position % 64; + std::uint64_t code = words[word_index] >> bit_offset; + if (bit_offset + kDepth > 64) { + code |= words[word_index + 1] << (64 - bit_offset); + } + dst[i] = table[code & kCodeMask]; + } +} + +#if defined(__BMI2__) +/// @brief Encode one depth-3 flat subtree with AVX2 gather and BMI2 packing. +/// @details AVX2 has no byte gather, so eight input symbols are widened to +/// dword indices and looked up together. The gathered codes are +/// narrowed back to bytes, then `pext` packs their low three bits. +/// Wider flat depths were benchmarked and deliberately retain their +/// faster scalar encoders. The caller provides three readable pad +/// bytes after the 256-entry code table because dword gathers overlap +/// adjacent byte entries. +inline void flat_encode_d3_avx2(std::uint64_t* bits, + const std::uint8_t* src, + const std::uint8_t* wire_code, + std::size_t count) noexcept { + constexpr std::uint8_t kDepth = 3; + constexpr std::uint64_t kByteLaneMask = UINT64_C(0x0707070707070707); + auto* packed = reinterpret_cast(bits); + const __m256i byte_mask = _mm256_set1_epi32(0xff); + const __m128i zero = _mm_setzero_si128(); + + std::size_t i = 0; + for (; i + 8 <= count; i += 8) { + const __m128i symbols = + _mm_loadl_epi64(reinterpret_cast(src + i)); + const __m256i indices = _mm256_cvtepu8_epi32(symbols); + const __m256i gathered = _mm256_and_si256( + _mm256_i32gather_epi32(reinterpret_cast(wire_code), indices, + 1), + byte_mask); + const __m128i gathered_low = _mm256_castsi256_si128(gathered); + const __m128i gathered_high = _mm256_extracti128_si256(gathered, 1); + const __m128i codes16 = _mm_packus_epi32(gathered_low, gathered_high); + const __m128i codes8 = _mm_packus_epi16(codes16, zero); + const std::uint64_t codes = + static_cast(_mm_cvtsi128_si64(codes8)); + const std::uint64_t dense = _pext_u64(codes, kByteLaneMask); + std::memcpy(packed + (i * kDepth) / 8, &dense, kDepth); + } + + std::size_t word_index = (i * kDepth) / 64; + std::size_t bit_offset = (i * kDepth) % 64; + for (; i < count; ++i) { + const std::uint64_t code = wire_code[src[i]]; + bits[word_index] |= code << bit_offset; + if (bit_offset + kDepth > 64) { + bits[word_index + 1] |= code >> (64 - bit_offset); + } + bit_offset += kDepth; + if (bit_offset >= 64) { + bit_offset -= 64; + ++word_index; + } + } +} + +/// @brief Encode a known balanced low-alphabet table as bit reversal. +template +inline void flat_encode_bitreverse_unchecked_avx2(std::uint64_t* bits, + const std::uint8_t* src, + const std::uint8_t* wire_code, + std::size_t count) noexcept { + static_assert(kDepth >= 2 && kDepth <= 7); + constexpr std::uint64_t kCodeMask = (std::uint64_t{1} << kDepth) - 1; + constexpr std::uint64_t kByteLaneMask = + kCodeMask * UINT64_C(0x0101010101010101); + auto* packed = reinterpret_cast(bits); + std::size_t i = 0; + for (; i + 8 <= count; i += 8) { + const __m128i symbols = + _mm_loadl_epi64(reinterpret_cast(src + i)); + const __m128i codes = reverse_low_bits16_avx2(symbols); + const std::uint64_t byte_codes = + static_cast(_mm_cvtsi128_si64(codes)); + const std::uint64_t dense = _pext_u64(byte_codes, kByteLaneMask); + std::memcpy(packed + (i * kDepth) / 8, &dense, kDepth); + } + + std::size_t word_index = (i * kDepth) / 64; + std::size_t bit_offset = (i * kDepth) % 64; + for (; i < count; ++i) { + const std::uint64_t code = wire_code[src[i]]; + bits[word_index] |= code << bit_offset; + if (bit_offset + kDepth > 64) { + bits[word_index + 1] |= code >> (64 - bit_offset); + } + bit_offset += kDepth; + if (bit_offset >= 64) { + bit_offset -= 64; + ++word_index; + } + } +} + +/// @brief Encode a balanced low-alphabet table as vectorized bit reversal. +/// @returns True when @p wire_code has the expected mapping and was encoded. +template +inline bool flat_encode_bitreverse_avx2(std::uint64_t* bits, + const std::uint8_t* src, + const std::uint8_t* wire_code, + std::size_t count) noexcept { + static_assert(kDepth >= 2 && kDepth <= 7); + constexpr std::size_t kTableSize = std::size_t{1} << kDepth; + for (std::size_t symbol = 0; symbol < kTableSize; ++symbol) { + if (wire_code[symbol] != + reverse_low_bits_avx2(static_cast(symbol), kDepth)) { + return false; + } + } + flat_encode_bitreverse_unchecked_avx2(bits, src, wire_code, count); + return true; +} + +/// @brief Select checked or preclassified balanced-table encoding. +template +inline bool flat_encode_bitreverse_maybe_avx2(std::uint64_t* bits, + const std::uint8_t* src, + const std::uint8_t* wire_code, + std::size_t count, + bool known_bitreverse) noexcept { + if (known_bitreverse) { + flat_encode_bitreverse_unchecked_avx2(bits, src, wire_code, count); + return true; + } + return flat_encode_bitreverse_avx2(bits, src, wire_code, count); +} +#endif + #endif // AVX2 + SSE4.1 // =========================================================================== @@ -907,7 +1365,8 @@ inline void flat_decode_dispatch(std::uint8_t* dst, const std::uint64_t* bits, const std::uint8_t* table, std::size_t count, - std::uint8_t depth) noexcept { + std::uint8_t depth, + bool known_bitreverse) noexcept { #if defined(__AVX512VBMI2__) && defined(__AVX512BW__) if (depth == 1) { merge_decode_avx512(dst, nullptr, table[0], nullptr, table[1], @@ -952,13 +1411,50 @@ inline void flat_decode_dispatch(std::uint8_t* dst, #if defined(__AVX2__) && defined(__SSE4_1__) const auto* packed = reinterpret_cast(bits); - if (depth == 2) { - flat_decode_d2_avx2(dst, packed, table, count); - return; - } - if (depth == 4) { - flat_decode_d4_avx2(dst, packed, table, count); - return; + switch (depth) { + case 2: + flat_decode_d2_avx2(dst, packed, table, count); + return; + case 3: + flat_decode_banked_avx2<3>(dst, packed, table, count); + return; + case 4: + flat_decode_d4_avx2(dst, packed, table, count); + return; + case 5: + flat_decode_banked_avx2<5>(dst, packed, table, count); + return; + case 6: + if (known_bitreverse) { + flat_decode_bitreverse_unchecked_avx2<6>(dst, packed, table, count); + return; + } + if (flat_decode_bitreverse_avx2<6>(dst, packed, table, count)) { + return; + } + flat_decode_banked_avx2<6>(dst, packed, table, count); + return; + case 7: + if (known_bitreverse) { + flat_decode_bitreverse_unchecked_avx2<7>(dst, packed, table, count); + return; + } + if (flat_decode_bitreverse_avx2<7>(dst, packed, table, count)) { + return; + } + flat_decode_banked_avx2<7>(dst, packed, table, count); + return; + case 8: + if (known_bitreverse) { + flat_decode_d8_bitreverse_unchecked_avx2(dst, packed, table, count); + return; + } + if (flat_decode_d8_bitreverse_avx2(dst, packed, table, count)) { + return; + } + break; + default: + break; } #endif flat_decode_scalar(dst, bits, table, count, depth); @@ -1081,13 +1577,140 @@ inline void pivco_partition_encode_none(std::uint64_t* bits, nullptr, nullptr, bits, src, dir, weight, 0, 0); } +/// @brief Build a bitmap for a node whose two children are constant leaves. +/// @details A set bit means @p src equals @p right_symbol. Unlike the generic +/// direction-table path, AVX2 can compare 32 symbols directly and +/// obtain the complete routing word fragment with `vpmovmskb`. +inline void pivco_partition_encode_cst_cst(std::uint64_t* bits, + const std::uint8_t* src, + std::uint8_t right_symbol, + std::size_t weight) noexcept { + std::size_t i = 0; +#if defined(__AVX512VBMI2__) && defined(__AVX512BW__) + const __m512i right = _mm512_set1_epi8(static_cast(right_symbol)); + for (; i + 64 <= weight; i += 64) { + const __m512i symbols = _mm512_loadu_si512(src + i); + bits[i / 64] = _mm512_cmpeq_epi8_mask(symbols, right); + } +#elif defined(__AVX2__) + const __m256i right = _mm256_set1_epi8(static_cast(right_symbol)); + auto* bit_bytes = reinterpret_cast(bits); + for (; i + 32 <= weight; i += 32) { + const __m256i symbols = + _mm256_loadu_si256(reinterpret_cast(src + i)); + const std::uint32_t routing = static_cast( + _mm256_movemask_epi8(_mm256_cmpeq_epi8(symbols, right))); + std::memcpy(bit_bytes + i / 8, &routing, sizeof(routing)); + } +#endif + for (; i < weight; ++i) { + if (src[i] == right_symbol) { + bits[i / 64] |= std::uint64_t{1} << (i % 64); + } + } +} + +#if defined(__AVX2__) && defined(__SSE4_1__) +/// @brief Partition a constant left leaf and materialize the right stream. +inline void pivco_partition_encode_cst_vec(std::uint8_t* right_dst, + std::uint64_t* bits, + const std::uint8_t* src, + std::uint8_t left_symbol, + std::size_t weight, + std::size_t right_weight) noexcept { + pivco_simd_detail::partition_encode_one_constant_avx2( + right_dst, bits, src, left_symbol, weight, right_weight); +} + +/// @brief Partition a dense left stream and a constant right leaf. +inline void pivco_partition_encode_vec_cst(std::uint8_t* left_dst, + std::uint64_t* bits, + const std::uint8_t* src, + std::uint8_t right_symbol, + std::size_t weight, + std::size_t left_weight) noexcept { + pivco_simd_detail::partition_encode_one_constant_avx2( + left_dst, bits, src, right_symbol, weight, left_weight); +} +#endif + +#if defined(__AVX2__) && defined(__SSE4_1__) && defined(__BMI2__) +/// @brief Encode a balanced depth-2..7 table as AVX2 bit reversal. +/// @param known_bitreverse Skip the table scan when construction already +/// classified the mapping. +inline bool pivco_flat_encode_bitreverse( + std::uint64_t* bits, + const std::uint8_t* src, + const std::uint8_t* wire_code, + std::size_t count, + std::uint8_t depth, + bool known_bitreverse = false) noexcept { + switch (depth) { + case 2: + return pivco_simd_detail::flat_encode_bitreverse_maybe_avx2<2>( + bits, src, wire_code, count, known_bitreverse); + case 3: + return pivco_simd_detail::flat_encode_bitreverse_maybe_avx2<3>( + bits, src, wire_code, count, known_bitreverse); + case 4: + return pivco_simd_detail::flat_encode_bitreverse_maybe_avx2<4>( + bits, src, wire_code, count, known_bitreverse); + case 5: + return pivco_simd_detail::flat_encode_bitreverse_maybe_avx2<5>( + bits, src, wire_code, count, known_bitreverse); + case 6: + return pivco_simd_detail::flat_encode_bitreverse_maybe_avx2<6>( + bits, src, wire_code, count, known_bitreverse); + case 7: + return pivco_simd_detail::flat_encode_bitreverse_maybe_avx2<7>( + bits, src, wire_code, count, known_bitreverse); + default: + return false; + } +} + +/// @brief Pack a depth-3 flat subtree through the AVX2/BMI2 kernel. +/// @pre @p wire_code has 259 readable bytes. +inline void pivco_flat_encode_d3(std::uint64_t* bits, + const std::uint8_t* src, + const std::uint8_t* wire_code, + std::size_t count) noexcept { + pivco_simd_detail::flat_encode_d3_avx2(bits, src, wire_code, count); +} +#endif + +#if defined(__AVX2__) && defined(__SSE4_1__) +/// @brief Encode the common full-alphabet bit-reversal table with AVX2. +/// @returns True when @p wire_code is the expected bit-reversal permutation. +/// @param known_bitreverse Skip the table scan when construction already +/// classified the mapping. +inline bool pivco_flat_encode_d8_bitreverse( + std::uint64_t* bits, + const std::uint8_t* src, + const std::uint8_t* wire_code, + std::size_t count, + bool known_bitreverse = false) noexcept { + if (known_bitreverse) { + pivco_simd_detail::flat_decode_d8_bitreverse_unchecked_avx2( + reinterpret_cast(bits), src, wire_code, count); + return true; + } + return pivco_simd_detail::flat_decode_d8_bitreverse_avx2( + reinterpret_cast(bits), src, wire_code, count); +} +#endif + /// @brief Unpack and translate a flat subtree's fixed-width wire codes. +/// @param known_bitreverse Skip table classification when the caller already +/// established the balanced mapping during construction. inline void pivco_flat_decode(std::uint8_t* dst, const std::uint64_t* bits, const std::uint8_t* table, std::size_t count, - std::uint8_t depth) noexcept { - pivco_simd_detail::flat_decode_dispatch(dst, bits, table, count, depth); + std::uint8_t depth, + bool known_bitreverse = false) noexcept { + pivco_simd_detail::flat_decode_dispatch(dst, bits, table, count, depth, + known_bitreverse); } } // namespace pixie diff --git a/src/benchmarks/pivco_huffman_benchmarks.cpp b/src/benchmarks/pivco_huffman_benchmarks.cpp index cebb0ed..297b5c7 100644 --- a/src/benchmarks/pivco_huffman_benchmarks.cpp +++ b/src/benchmarks/pivco_huffman_benchmarks.cpp @@ -67,6 +67,21 @@ std::vector make_low_entropy_4(std::size_t n, return data; } +/// @brief Uniform distribution over a power-of-two alphabet. +/// @details Depths 3--7 exercise every flat width between the four-symbol and +/// full-byte datasets. +template +std::vector make_uniform_pow2(std::size_t n, + std::mt19937_64& rng) { + static_assert(kDepth >= 3 && kDepth <= 7); + constexpr std::uint8_t kMask = (std::uint8_t{1} << kDepth) - 1; + std::vector data(n); + for (auto& b : data) { + b = static_cast(rng()) & kMask; + } + return data; +} + /// @brief English-text-like distribution: ~17% space, ~80% lowercase letters, /// ~2% digits, ~1% punctuation. std::vector make_text_like(std::size_t n, std::mt19937_64& rng) { @@ -185,6 +200,11 @@ void BM_Decode(benchmark::State& state, DatasetMaker make) { PIXIE_PIVCO_REGISTER(Encode, Uniform256, make_uniform256) PIXIE_PIVCO_REGISTER(Encode, Binary, make_binary) PIXIE_PIVCO_REGISTER(Encode, LowEntropy4, make_low_entropy_4) +PIXIE_PIVCO_REGISTER(Encode, Uniform8, make_uniform_pow2<3>) +PIXIE_PIVCO_REGISTER(Encode, Uniform16, make_uniform_pow2<4>) +PIXIE_PIVCO_REGISTER(Encode, Uniform32, make_uniform_pow2<5>) +PIXIE_PIVCO_REGISTER(Encode, Uniform64, make_uniform_pow2<6>) +PIXIE_PIVCO_REGISTER(Encode, Uniform128, make_uniform_pow2<7>) PIXIE_PIVCO_REGISTER(Encode, TextLike, make_text_like) PIXIE_PIVCO_REGISTER(Encode, Skewed99, make_skewed99) PIXIE_PIVCO_REGISTER(Encode, SingleSymbol, make_single_symbol) @@ -192,6 +212,11 @@ PIXIE_PIVCO_REGISTER(Encode, SingleSymbol, make_single_symbol) PIXIE_PIVCO_REGISTER(Decode, Uniform256, make_uniform256) PIXIE_PIVCO_REGISTER(Decode, Binary, make_binary) PIXIE_PIVCO_REGISTER(Decode, LowEntropy4, make_low_entropy_4) +PIXIE_PIVCO_REGISTER(Decode, Uniform8, make_uniform_pow2<3>) +PIXIE_PIVCO_REGISTER(Decode, Uniform16, make_uniform_pow2<4>) +PIXIE_PIVCO_REGISTER(Decode, Uniform32, make_uniform_pow2<5>) +PIXIE_PIVCO_REGISTER(Decode, Uniform64, make_uniform_pow2<6>) +PIXIE_PIVCO_REGISTER(Decode, Uniform128, make_uniform_pow2<7>) PIXIE_PIVCO_REGISTER(Decode, TextLike, make_text_like) PIXIE_PIVCO_REGISTER(Decode, Skewed99, make_skewed99) PIXIE_PIVCO_REGISTER(Decode, SingleSymbol, make_single_symbol) From ce30a28d9794f093a91fc98a599e532c907b63c6 Mon Sep 17 00:00:00 2001 From: Sanya239 Date: Tue, 28 Jul 2026 01:29:44 +0300 Subject: [PATCH 13/14] Add aarch64 NEON optimizations --- include/pixie/huffman/pivco_huffman.h | 23 +- include/pixie/huffman/pivco_simd.h | 854 +++++++++++++++++++++++++- src/tests/pivco_huffman_tests.cpp | 120 ++++ 3 files changed, 992 insertions(+), 5 deletions(-) diff --git a/include/pixie/huffman/pivco_huffman.h b/include/pixie/huffman/pivco_huffman.h index f3143b4..4e823b4 100644 --- a/include/pixie/huffman/pivco_huffman.h +++ b/include/pixie/huffman/pivco_huffman.h @@ -552,7 +552,7 @@ class PivCoHuffman : public HuffmanBase { * right=1). * @param table Output table of size `2^total_depth`. * @return True when the completed table is the balanced bit-reversal - * permutation used by the specialized AVX2 kernels. */ + * permutation used by the specialized SIMD kernels. */ bool build_flat_table(std::size_t idx, std::uint8_t flat_depth, std::uint32_t code, @@ -837,6 +837,13 @@ class PivCoHuffman : public HuffmanBase { std::uint8_t flat_depth, const symbol_type* table, bool known_bitreverse) { +#if defined(PIXIE_PIVCO_NEON_SUPPORT) + if (known_bitreverse) { + pivco_flat_encode_neon(bits, src, nullptr, count, flat_depth, true); + return; + } +#endif + // Three zero pad bytes make overlapping dword gathers at symbol 255 safe // in the AVX2 flat encoder. std::array wire_code{}; @@ -864,6 +871,12 @@ class PivCoHuffman : public HuffmanBase { #endif #endif +#if defined(PIXIE_PIVCO_NEON_SUPPORT) + pivco_flat_encode_neon(bits, src, wire_code.data(), count, flat_depth, + false); + return; +#endif + auto* bytes = reinterpret_cast(bits); if (flat_depth == 8) { for (std::size_t i = 0; i < count; ++i) { @@ -971,7 +984,9 @@ class PivCoHuffman : public HuffmanBase { if (left_is_leaf && right_is_leaf) { pivco_partition_encode_cst_cst(bits, src, nodes_[n.right].symbol, weight); } else if (right_is_leaf) { -#if defined(__AVX2__) && defined(__SSE4_1__) && !defined(__AVX512VBMI2__) +#if ((defined(__AVX2__) && defined(__SSE4_1__)) || \ + defined(PIXIE_PIVCO_NEON_SUPPORT)) && \ + !defined(__AVX512VBMI2__) pivco_partition_encode_vec_cst( left_dst, bits, src, nodes_[n.right].symbol, weight, left_weight); #else @@ -979,7 +994,9 @@ class PivCoHuffman : public HuffmanBase { left_weight); #endif } else if (left_is_leaf) { -#if defined(__AVX2__) && defined(__SSE4_1__) && !defined(__AVX512VBMI2__) +#if ((defined(__AVX2__) && defined(__SSE4_1__)) || \ + defined(PIXIE_PIVCO_NEON_SUPPORT)) && \ + !defined(__AVX512VBMI2__) pivco_partition_encode_cst_vec(right_dst, bits, src, nodes_[n.left].symbol, weight, weight - left_weight); diff --git a/include/pixie/huffman/pivco_simd.h b/include/pixie/huffman/pivco_simd.h index 165ff67..bafe099 100644 --- a/include/pixie/huffman/pivco_simd.h +++ b/include/pixie/huffman/pivco_simd.h @@ -27,6 +27,8 @@ * full `merge_flat_D` family without changing the wire format. * - AVX2 + SSE4.1: 8-byte-at-a-time `_mm_shuffle_epi8` with 256-entry lookup * tables, mirroring the paper's NEON `vqtbl1q_u8` technique. + * - AArch64 NEON: paper-style `vqtbl*`/`vqtbx*` table lookup, compact and + * expand tables, and 16-symbol fixed-width unpacking. * - Scalar: the original bit-by-bit / symbol-by-symbol loops. * * Every SIMD backend produces byte-identical output to the scalar fallback @@ -47,6 +49,11 @@ #include #endif +#if defined(__aarch64__) && (defined(__ARM_NEON) || defined(__ARM_NEON__)) +#include +#define PIXIE_PIVCO_NEON_SUPPORT 1 +#endif + namespace pixie::pivco_simd_detail { /// @brief Bitmap bit-ordering convention used throughout this header. @@ -173,6 +180,761 @@ inline void partition_scalar_tail(std::uint8_t* left_dst, } } +// =========================================================================== +// AArch64 NEON backend (table lookup and compaction, 8/16 symbols per batch). +// =========================================================================== +#if defined(PIXIE_PIVCO_NEON_SUPPORT) + +/// @brief 256 eight-byte table controls for one NEON `vqtbl1_u8` operation. +struct NeonShufTable { + alignas(16) std::uint8_t entry[256][8]; +}; + +/// @brief Paper `expand_tab`: select left/right dense elements for each mask. +inline const NeonShufTable& merge_lut_neon() { + static const NeonShufTable lut = [] { + NeonShufTable table{}; + for (int mask = 0; mask < 256; ++mask) { + int left = 0; + int right = 0; + for (int lane = 0; lane < 8; ++lane) { + if ((mask >> lane) & 1) { + table.entry[mask][lane] = static_cast(8 + right++); + } else { + table.entry[mask][lane] = static_cast(left++); + } + } + } + return table; + }(); + return lut; +} + +/// @brief Paper `compress_tab` for elements routed to the right child. +inline const NeonShufTable& compress_right_lut_neon() { + static const NeonShufTable lut = [] { + NeonShufTable table{}; + for (int mask = 0; mask < 256; ++mask) { + int output = 0; + for (int lane = 0; lane < 8; ++lane) { + if ((mask >> lane) & 1) { + table.entry[mask][output++] = static_cast(lane); + } + } + for (; output < 8; ++output) { + table.entry[mask][output] = 0xff; + } + } + return table; + }(); + return lut; +} + +/// @brief Paper `compress_tab` for elements routed to the left child. +inline const NeonShufTable& compress_left_lut_neon() { + static const NeonShufTable lut = [] { + NeonShufTable table{}; + for (int mask = 0; mask < 256; ++mask) { + int output = 0; + for (int lane = 0; lane < 8; ++lane) { + if (((mask >> lane) & 1) == 0) { + table.entry[mask][output++] = static_cast(lane); + } + } + for (; output < 8; ++output) { + table.entry[mask][output] = 0xff; + } + } + return table; + }(); + return lut; +} + +/// @brief Four 64-byte NEON lookup banks, enough for any byte alphabet. +struct ByteLookupNeon { + uint8x16x4_t bank0; + uint8x16x4_t bank1; + uint8x16x4_t bank2; + uint8x16x4_t bank3; +}; + +/// @brief Return four zero vectors in the layout expected by `vqtbl4*`. +inline uint8x16x4_t zero_lookup_bank_neon() noexcept { + const uint8x16_t zero = vdupq_n_u8(0); + uint8x16x4_t bank; + bank.val[0] = zero; + bank.val[1] = zero; + bank.val[2] = zero; + bank.val[3] = zero; + return bank; +} + +/// @brief Load one complete 64-byte `vqtbl4*` bank. +inline uint8x16x4_t load_lookup_bank_neon(const std::uint8_t* table) noexcept { + uint8x16x4_t bank; + bank.val[0] = vld1q_u8(table); + bank.val[1] = vld1q_u8(table + 16); + bank.val[2] = vld1q_u8(table + 32); + bank.val[3] = vld1q_u8(table + 48); + return bank; +} + +/// @brief Load exactly `2^D` table bytes without reading past small tables. +template +inline ByteLookupNeon load_byte_lookup_neon( + const std::uint8_t* table) noexcept { + static_assert(kDepth >= 1 && kDepth <= 8); + ByteLookupNeon lookup{zero_lookup_bank_neon(), zero_lookup_bank_neon(), + zero_lookup_bank_neon(), zero_lookup_bank_neon()}; + if constexpr (kDepth <= 3) { + std::uint64_t entries = 0; + std::memcpy(&entries, table, std::size_t{1} << kDepth); + lookup.bank0.val[0] = vcombine_u8(vcreate_u8(entries), vdup_n_u8(0)); + } else if constexpr (kDepth == 4) { + lookup.bank0.val[0] = vld1q_u8(table); + } else if constexpr (kDepth == 5) { + lookup.bank0.val[0] = vld1q_u8(table); + lookup.bank0.val[1] = vld1q_u8(table + 16); + } else { + lookup.bank0 = load_lookup_bank_neon(table); + if constexpr (kDepth >= 7) { + lookup.bank1 = load_lookup_bank_neon(table + 64); + } + if constexpr (kDepth == 8) { + lookup.bank2 = load_lookup_bank_neon(table + 128); + lookup.bank3 = load_lookup_bank_neon(table + 192); + } + } + return lookup; +} + +/// @brief Translate sixteen byte indices through a preloaded `2^D` table. +template +inline uint8x16_t lookup_bytes16_neon(const ByteLookupNeon& lookup, + uint8x16_t indices) noexcept { + uint8x16_t result = vqtbl4q_u8(lookup.bank0, indices); + if constexpr (kDepth >= 7) { + result = + vqtbx4q_u8(result, lookup.bank1, vsubq_u8(indices, vdupq_n_u8(64))); + } + if constexpr (kDepth == 8) { + result = + vqtbx4q_u8(result, lookup.bank2, vsubq_u8(indices, vdupq_n_u8(128))); + result = + vqtbx4q_u8(result, lookup.bank3, vsubq_u8(indices, vdupq_n_u8(192))); + } + return result; +} + +/// @brief Translate eight byte indices through a preloaded `2^D` table. +template +inline uint8x8_t lookup_bytes8_neon(const ByteLookupNeon& lookup, + uint8x8_t indices) noexcept { + uint8x8_t result = vqtbl4_u8(lookup.bank0, indices); + if constexpr (kDepth >= 7) { + result = vqtbx4_u8(result, lookup.bank1, vsub_u8(indices, vdup_n_u8(64))); + } + if constexpr (kDepth == 8) { + result = vqtbx4_u8(result, lookup.bank2, vsub_u8(indices, vdup_n_u8(128))); + result = vqtbx4_u8(result, lookup.bank3, vsub_u8(indices, vdup_n_u8(192))); + } + return result; +} + +/// @brief Reverse the low D bits independently in sixteen byte lanes. +template +inline uint8x16_t reverse_low_bits16_neon(uint8x16_t values) noexcept { + static_assert(kDepth >= 2 && kDepth <= 8); + const uint8x16_t reversed = vrbitq_u8(values); + if constexpr (kDepth == 8) { + return reversed; + } else { + return vshrq_n_u8(reversed, 8 - kDepth); + } +} + +/// @brief Scalar bit reverse for the final fewer-than-16 encode symbols. +template +inline std::uint8_t reverse_low_bits_scalar_neon(std::uint8_t value) noexcept { + static_assert(kDepth >= 2 && kDepth <= 8); + value = static_cast((value >> 4) | (value << 4)); + value = static_cast(((value & 0xccu) >> 2) | + ((value & 0x33u) << 2)); + value = static_cast(((value & 0xaau) >> 1) | + ((value & 0x55u) << 1)); + return static_cast(value >> (8 - kDepth)); +} + +/// @brief Translate flat codes through either a table or a known bit reverse. +template +inline uint8x16_t translate_flat_codes16_neon(const ByteLookupNeon* lookup, + uint8x16_t codes) noexcept { + if constexpr (kBitReverse) { + return reverse_low_bits16_neon(codes); + } else { + return lookup_bytes16_neon(*lookup, codes); + } +} + +/// @brief Eight-lane form of `translate_flat_codes16_neon`. +template +inline uint8x8_t translate_flat_codes8_neon(const ByteLookupNeon* lookup, + uint8x8_t codes) noexcept { + if constexpr (kBitReverse) { + return vget_low_u8( + reverse_low_bits16_neon(vcombine_u8(codes, codes))); + } else { + return lookup_bytes8_neon(*lookup, codes); + } +} + +/// @brief Convert eight lanes containing 0/1 into an LSB-first byte mask. +inline std::uint8_t lane_bits_to_mask_neon(uint8x8_t lane_bits) noexcept { + alignas(8) static constexpr std::uint8_t kBitWeights[8] = {1, 2, 4, 8, + 16, 32, 64, 128}; + return vaddv_u8(vmul_u8(lane_bits, vld1_u8(kBitWeights))); +} + +/// @brief Store only the valid prefix of an eight-byte compaction result. +inline void store_masked_neon(std::uint8_t* dst, + uint8x8_t values, + std::size_t count) noexcept { + if (count >= 8) { + vst1_u8(dst, values); + } else if (count != 0) { + alignas(8) std::uint8_t temporary[8]; + vst1_u8(temporary, values); + std::memcpy(dst, temporary, count); + } +} + +/// @brief Compact one eight-symbol partition group into its requested child +/// streams. +template +inline void partition_group8_neon(std::uint8_t* left_dst, + std::uint8_t* right_dst, + uint8x8_t values, + std::uint8_t routing, + std::size_t left_weight, + std::size_t right_weight, + std::size_t& left_count, + std::size_t& right_count) noexcept { + const uint8x16_t table = vcombine_u8(values, vdup_n_u8(0)); + if constexpr (kWriteRight) { + const uint8x8_t controls = + vld1_u8(compress_right_lut_neon().entry[routing]); + store_masked_neon(right_dst + right_count, vqtbl1_u8(table, controls), + right_weight - right_count); + } + if constexpr (kWriteLeft) { + const uint8x8_t controls = vld1_u8(compress_left_lut_neon().entry[routing]); + store_masked_neon(left_dst + left_count, vqtbl1_u8(table, controls), + left_weight - left_count); + } + const std::size_t routed_right = + std::popcount(static_cast(routing)); + right_count += routed_right; + left_count += 8 - routed_right; +} + +/// @brief Compact the non-constant side of one eight-symbol leaf partition. +template +inline void partition_one_constant_group8_neon(std::uint8_t* output, + uint8x8_t values, + uint8x8_t equal_lanes, + std::size_t output_weight, + std::size_t& output_count, + std::uint8_t& routing) noexcept { + const std::uint8_t equal = lane_bits_to_mask_neon(vshr_n_u8(equal_lanes, 7)); + routing = kLeftConstant ? static_cast(~equal) : equal; + const auto& lut = + kLeftConstant ? compress_right_lut_neon() : compress_left_lut_neon(); + const uint8x8_t packed = + vqtbl1_u8(vcombine_u8(values, vdup_n_u8(0)), vld1_u8(lut.entry[routing])); + store_masked_neon(output + output_count, packed, + output_weight - output_count); + const std::size_t routed_right = + std::popcount(static_cast(routing)); + output_count += kLeftConstant ? routed_right : 8 - routed_right; +} + +/// @brief Merge two constant leaves without a shuffle-table load. +inline void merge_decode_cst_cst_neon(std::uint8_t* dst, + std::uint8_t left_symbol, + std::uint8_t right_symbol, + const std::uint64_t* bits, + std::size_t count) noexcept { + alignas(16) static constexpr std::uint8_t kBitPositions[16] = { + 1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128}; + const uint8x16_t positions = vld1q_u8(kBitPositions); + const uint8x16_t left = vdupq_n_u8(left_symbol); + const uint8x16_t right = vdupq_n_u8(right_symbol); + const auto* bit_bytes = reinterpret_cast(bits); + std::size_t i = 0; + for (; i + 16 <= count; i += 16) { + const uint8x16_t routing = vcombine_u8(vdup_n_u8(bit_bytes[i / 8]), + vdup_n_u8(bit_bytes[i / 8 + 1])); + vst1q_u8(dst + i, vbslq_u8(vtstq_u8(routing, positions), right, left)); + } + for (; i < count; ++i) { + dst[i] = + ((bits[i / 64] >> (i % 64)) & 1u) != 0 ? right_symbol : left_symbol; + } +} + +/// @brief Paper-style NEON dense/constant merge in eight-symbol batches. +template +inline void merge_decode_neon(std::uint8_t* dst, + const std::uint8_t* left, + std::uint8_t left_symbol, + const std::uint8_t* right, + std::uint8_t right_symbol, + const std::uint64_t* bits, + std::size_t count) noexcept { + if constexpr (kLeftConstant && kRightConstant) { + merge_decode_cst_cst_neon(dst, left_symbol, right_symbol, bits, count); + return; + } + + const auto& merge_table = merge_lut_neon(); + const auto* bit_bytes = reinterpret_cast(bits); + std::size_t output = 0; + std::size_t left_count = 0; + std::size_t right_count = 0; + std::size_t i = 0; + for (; i + 8 <= count; i += 8) { + const std::uint8_t routing = bit_bytes[i / 8]; + uint8x8_t left_values; + if constexpr (kLeftConstant) { + left_values = vdup_n_u8(left_symbol); + } else { + left_values = vld1_u8(left + left_count); + } + uint8x8_t right_values; + if constexpr (kRightConstant) { + right_values = vdup_n_u8(right_symbol); + } else { + right_values = vld1_u8(right + right_count); + } + const uint8x8_t controls = vld1_u8(merge_table.entry[routing]); + vst1_u8(dst + output, + vqtbl1_u8(vcombine_u8(left_values, right_values), controls)); + output += 8; + const std::size_t routed_right = + std::popcount(static_cast(routing)); + right_count += routed_right; + left_count += 8 - routed_right; + } + merge_scalar_tail( + dst, output, left, left_symbol, left_count, right, right_symbol, + right_count, bits, i, count); +} + +/// @brief Paper-style direction lookup plus NEON `compress_tab` partition. +template +inline void partition_encode_neon(std::uint8_t* left_dst, + std::uint8_t* right_dst, + std::uint64_t* bits, + const std::uint8_t* src, + const std::uint8_t* dir, + std::size_t weight, + std::size_t left_weight, + std::size_t right_weight) noexcept { + const ByteLookupNeon direction_lookup = load_byte_lookup_neon<8>(dir); + auto* bit_bytes = reinterpret_cast(bits); + std::size_t left_count = 0; + std::size_t right_count = 0; + std::size_t i = 0; + for (; i + 16 <= weight; i += 16) { + const uint8x16_t values = vld1q_u8(src + i); + const uint8x16_t directions = + lookup_bytes16_neon<8>(direction_lookup, values); + const std::uint8_t low_routing = + lane_bits_to_mask_neon(vget_low_u8(directions)); + const std::uint8_t high_routing = + lane_bits_to_mask_neon(vget_high_u8(directions)); + bit_bytes[i / 8] = low_routing; + bit_bytes[i / 8 + 1] = high_routing; + partition_group8_neon( + left_dst, right_dst, vget_low_u8(values), low_routing, left_weight, + right_weight, left_count, right_count); + partition_group8_neon( + left_dst, right_dst, vget_high_u8(values), high_routing, left_weight, + right_weight, left_count, right_count); + } + for (; i + 8 <= weight; i += 8) { + const uint8x8_t values = vld1_u8(src + i); + const std::uint8_t routing = + lane_bits_to_mask_neon(lookup_bytes8_neon<8>(direction_lookup, values)); + bit_bytes[i / 8] = routing; + partition_group8_neon( + left_dst, right_dst, values, routing, left_weight, right_weight, + left_count, right_count); + } + + if constexpr (kWriteLeft || kWriteRight) { + partition_scalar_tail(left_dst, right_dst, bits, + src, dir, left_count, + right_count, i, weight); + } else { + partition_bitmap_tail(bits, src, dir, i, weight); + } +} + +/// @brief NEON partition when one child is a constant leaf. +template +inline void partition_encode_one_constant_neon( + std::uint8_t* output, + std::uint64_t* bits, + const std::uint8_t* src, + std::uint8_t constant_symbol, + std::size_t weight, + std::size_t output_weight) noexcept { + const uint8x8_t constant = vdup_n_u8(constant_symbol); + auto* bit_bytes = reinterpret_cast(bits); + std::size_t output_count = 0; + std::size_t i = 0; + for (; i + 16 <= weight; i += 16) { + const uint8x16_t values = vld1q_u8(src + i); + const uint8x16_t equal = vceqq_u8(values, vdupq_n_u8(constant_symbol)); + std::uint8_t low_routing = 0; + std::uint8_t high_routing = 0; + partition_one_constant_group8_neon( + output, vget_low_u8(values), vget_low_u8(equal), output_weight, + output_count, low_routing); + partition_one_constant_group8_neon( + output, vget_high_u8(values), vget_high_u8(equal), output_weight, + output_count, high_routing); + bit_bytes[i / 8] = low_routing; + bit_bytes[i / 8 + 1] = high_routing; + } + for (; i + 8 <= weight; i += 8) { + const uint8x8_t values = vld1_u8(src + i); + std::uint8_t routing = 0; + partition_one_constant_group8_neon( + output, values, vceq_u8(values, constant), output_weight, output_count, + routing); + bit_bytes[i / 8] = routing; + } + + for (; i < weight; ++i) { + const bool is_constant = src[i] == constant_symbol; + const bool goes_right = kLeftConstant ? !is_constant : is_constant; + if (goes_right) { + bits[i / 64] |= std::uint64_t{1} << (i % 64); + } + if (kLeftConstant ? goes_right : !goes_right) { + output[output_count++] = src[i]; + } + } +} + +/// @brief Build a two-leaf routing bitmap sixteen symbols at a time. +inline void partition_encode_cst_cst_neon(std::uint64_t* bits, + const std::uint8_t* src, + std::uint8_t right_symbol, + std::size_t weight) noexcept { + const uint8x16_t right = vdupq_n_u8(right_symbol); + auto* bit_bytes = reinterpret_cast(bits); + std::size_t i = 0; + for (; i + 16 <= weight; i += 16) { + const uint8x16_t equal = vceqq_u8(vld1q_u8(src + i), right); + bit_bytes[i / 8] = lane_bits_to_mask_neon(vshr_n_u8(vget_low_u8(equal), 7)); + bit_bytes[i / 8 + 1] = + lane_bits_to_mask_neon(vshr_n_u8(vget_high_u8(equal), 7)); + } + for (; i < weight; ++i) { + if (src[i] == right_symbol) { + bits[i / 64] |= std::uint64_t{1} << (i % 64); + } + } +} + +/// @brief Load exactly the bytes containing sixteen packed D-bit codes. +template +inline uint8x16_t load_packed_codes16_neon( + const std::uint8_t* packed) noexcept { + static_assert(kDepth >= 3 && kDepth <= 8); + constexpr std::size_t kPackedBytes = 2 * kDepth; + std::uint64_t low = 0; + std::uint64_t high = 0; + constexpr std::size_t kLowBytes = kPackedBytes < 8 ? kPackedBytes : 8; + std::memcpy(&low, packed, kLowBytes); + if constexpr (kPackedBytes > 8) { + std::memcpy(&high, packed + 8, kPackedBytes - 8); + } + return vcombine_u8(vcreate_u8(low), vcreate_u8(high)); +} + +/// @brief Unpack sixteen tightly packed D-bit codes into sixteen byte lanes. +template +inline uint8x16_t unpack_codes16_neon(const std::uint8_t* packed) noexcept { + static_assert(kDepth >= 3 && kDepth <= 7); + alignas(16) static constexpr auto kLowIndices = [] { + std::array controls{}; + for (std::size_t lane = 0; lane < controls.size(); ++lane) { + controls[lane] = static_cast((lane * kDepth) / 8); + } + return controls; + }(); + alignas(16) static constexpr auto kHighIndices = [] { + std::array controls{}; + for (std::size_t lane = 0; lane < controls.size(); ++lane) { + controls[lane] = static_cast((lane * kDepth) / 8 + 1); + } + return controls; + }(); + alignas(16) static constexpr auto kRightShifts = [] { + std::array shifts{}; + for (std::size_t lane = 0; lane < shifts.size(); ++lane) { + shifts[lane] = -static_cast((lane * kDepth) % 8); + } + return shifts; + }(); + alignas(16) static constexpr auto kLeftShifts = [] { + std::array shifts{}; + for (std::size_t lane = 0; lane < shifts.size(); ++lane) { + const std::uint8_t offset = (lane * kDepth) % 8; + shifts[lane] = static_cast(offset == 0 ? 8 : 8 - offset); + } + return shifts; + }(); + + const uint8x16_t bytes = load_packed_codes16_neon(packed); + const uint8x16_t low = vqtbl1q_u8(bytes, vld1q_u8(kLowIndices.data())); + const uint8x16_t high = vqtbl1q_u8(bytes, vld1q_u8(kHighIndices.data())); + const uint8x16_t codes = + vorrq_u8(vshlq_u8(low, vld1q_s8(kRightShifts.data())), + vshlq_u8(high, vld1q_s8(kLeftShifts.data()))); + return vandq_u8(codes, vdupq_n_u8((std::uint8_t{1} << kDepth) - 1)); +} + +/// @brief Decode one flat subtree with NEON table lookup or bit reversal. +template +inline void flat_decode_neon(std::uint8_t* dst, + const std::uint64_t* bits, + const std::uint8_t* table, + std::size_t count) noexcept { + static_assert(kDepth >= 2 && kDepth <= 8); + ByteLookupNeon lookup; + const ByteLookupNeon* lookup_ptr = nullptr; + if constexpr (!kBitReverse) { + lookup = load_byte_lookup_neon(table); + lookup_ptr = &lookup; + } + const auto* packed = reinterpret_cast(bits); + std::size_t i = 0; + if constexpr (kDepth == 2) { + const uint8x8_t mask = vdup_n_u8(0x03); + for (; i + 32 <= count; i += 32) { + const uint8x8_t bytes = vld1_u8(packed + i / 4); + uint8x8x4_t symbols; + symbols.val[0] = translate_flat_codes8_neon<2, kBitReverse>( + lookup_ptr, vand_u8(bytes, mask)); + symbols.val[1] = translate_flat_codes8_neon<2, kBitReverse>( + lookup_ptr, vand_u8(vshr_n_u8(bytes, 2), mask)); + symbols.val[2] = translate_flat_codes8_neon<2, kBitReverse>( + lookup_ptr, vand_u8(vshr_n_u8(bytes, 4), mask)); + symbols.val[3] = translate_flat_codes8_neon<2, kBitReverse>( + lookup_ptr, vshr_n_u8(bytes, 6)); + vst4_u8(dst + i, symbols); + } + } else if constexpr (kDepth == 4) { + const uint8x8_t mask = vdup_n_u8(0x0f); + for (; i + 16 <= count; i += 16) { + const uint8x8_t bytes = vld1_u8(packed + i / 2); + uint8x8x2_t symbols; + symbols.val[0] = translate_flat_codes8_neon<4, kBitReverse>( + lookup_ptr, vand_u8(bytes, mask)); + symbols.val[1] = translate_flat_codes8_neon<4, kBitReverse>( + lookup_ptr, vshr_n_u8(bytes, 4)); + vst2_u8(dst + i, symbols); + } + } else if constexpr (kDepth == 8) { + for (; i + 16 <= count; i += 16) { + const uint8x16_t codes = vld1q_u8(packed + i); + const uint8x16_t symbols = + translate_flat_codes16_neon<8, kBitReverse>(lookup_ptr, codes); + vst1q_u8(dst + i, symbols); + } + } else { + for (; i + 16 <= count; i += 16) { + const auto* group = packed + (i * kDepth) / 8; + const uint8x16_t codes = unpack_codes16_neon(group); + const uint8x16_t symbols = + translate_flat_codes16_neon(lookup_ptr, codes); + vst1q_u8(dst + i, symbols); + } + } + + constexpr std::uint64_t kCodeMask = (std::uint64_t{1} << kDepth) - 1; + for (; i < count; ++i) { + const std::size_t bit_position = i * kDepth; + const std::size_t word_index = bit_position / 64; + const std::size_t bit_offset = bit_position % 64; + std::uint64_t code = bits[word_index] >> bit_offset; + if (bit_offset + kDepth > 64) { + code |= bits[word_index + 1] << (64 - bit_offset); + } + dst[i] = table[code & kCodeMask]; + } +} + +/// @brief Densely pack eight byte codes into the low `8 * D` bits. +template +inline std::uint64_t pack_codes8_neon(std::uint64_t byte_codes) noexcept { + static_assert(kDepth >= 3 && kDepth <= 8); + constexpr std::uint64_t kCodeMask = (std::uint64_t{1} << kDepth) - 1; + std::uint64_t dense = 0; + for (std::size_t lane = 0; lane < 8; ++lane) { + dense |= ((byte_codes >> (lane * 8)) & kCodeMask) << (lane * kDepth); + } + return dense; +} + +/// @brief Encode flat wire codes with NEON lookup/bit-reverse and fixed pack. +template +inline void flat_encode_neon(std::uint64_t* bits, + const std::uint8_t* src, + const std::uint8_t* wire_code, + std::size_t count) noexcept { + static_assert(kDepth >= 2 && kDepth <= 8); + ByteLookupNeon lookup; + const ByteLookupNeon* lookup_ptr = nullptr; + if constexpr (!kBitReverse) { + lookup = load_byte_lookup_neon<8>(wire_code); + lookup_ptr = &lookup; + } + auto* packed = reinterpret_cast(bits); + std::size_t i = 0; + if constexpr (kDepth == 8) { + for (; i + 16 <= count; i += 16) { + const uint8x16_t symbols = vld1q_u8(src + i); + const uint8x16_t codes = + translate_flat_codes16_neon<8, kBitReverse>(lookup_ptr, symbols); + vst1q_u8(packed + i, codes); + } + for (; i < count; ++i) { + if constexpr (kBitReverse) { + packed[i] = reverse_low_bits_scalar_neon<8>(src[i]); + } else { + packed[i] = wire_code[src[i]]; + } + } + return; + } + + for (; i + 16 <= count; i += 16) { + uint8x16_t codes; + if constexpr (kBitReverse) { + codes = reverse_low_bits16_neon(vld1q_u8(src + i)); + } else { + codes = lookup_bytes16_neon<8>(*lookup_ptr, vld1q_u8(src + i)); + } + if constexpr (kDepth == 2) { + const uint8x16_t even = vuzp1q_u8(codes, codes); + const uint8x16_t odd = vuzp2q_u8(codes, codes); + const uint8x8_t pairs = vget_low_u8( + vorrq_u8(even, vshlq_n_u8(odd, static_cast(kDepth)))); + const uint8x8_t packed_codes = + vorr_u8(vuzp1_u8(pairs, pairs), vshl_n_u8(vuzp2_u8(pairs, pairs), 4)); + const std::uint32_t dense = + vget_lane_u32(vreinterpret_u32_u8(packed_codes), 0); + std::memcpy(packed + i / 4, &dense, sizeof(dense)); + } else if constexpr (kDepth == 4) { + const uint8x16_t even = vuzp1q_u8(codes, codes); + const uint8x16_t odd = vuzp2q_u8(codes, codes); + vst1_u8(packed + i / 2, vget_low_u8(vorrq_u8(even, vshlq_n_u8(odd, 4)))); + } else { + const uint64x2_t byte_codes = vreinterpretq_u64_u8(codes); + const std::uint64_t low = + pack_codes8_neon(vgetq_lane_u64(byte_codes, 0)); + const std::uint64_t high = + pack_codes8_neon(vgetq_lane_u64(byte_codes, 1)); + auto* output = packed + (i * kDepth) / 8; + std::memcpy(output, &low, kDepth); + std::memcpy(output + kDepth, &high, kDepth); + } + } + + std::size_t word_index = (i * kDepth) / 64; + std::size_t bit_offset = (i * kDepth) % 64; + for (; i < count; ++i) { + std::uint64_t code; + if constexpr (kBitReverse) { + code = reverse_low_bits_scalar_neon(src[i]); + } else { + code = wire_code[src[i]]; + } + bits[word_index] |= code << bit_offset; + if (bit_offset + kDepth > 64) { + bits[word_index + 1] |= code >> (64 - bit_offset); + } + bit_offset += kDepth; + if (bit_offset >= 64) { + bit_offset -= 64; + ++word_index; + } + } +} + +/// @brief Select table lookup or the preclassified balanced-table fast path. +template +inline void flat_encode_maybe_bitreverse_neon(std::uint64_t* bits, + const std::uint8_t* src, + const std::uint8_t* wire_code, + std::size_t count, + bool known_bitreverse) noexcept { + if (known_bitreverse) { + flat_encode_neon(bits, src, wire_code, count); + } else { + flat_encode_neon(bits, src, wire_code, count); + } +} + +/// @brief Runtime depth dispatch for NEON flat encoding. +inline void flat_encode_dispatch_neon(std::uint64_t* bits, + const std::uint8_t* src, + const std::uint8_t* wire_code, + std::size_t count, + std::uint8_t depth, + bool known_bitreverse) noexcept { + switch (depth) { + case 2: + flat_encode_maybe_bitreverse_neon<2>(bits, src, wire_code, count, + known_bitreverse); + return; + case 3: + flat_encode_maybe_bitreverse_neon<3>(bits, src, wire_code, count, + known_bitreverse); + return; + case 4: + flat_encode_maybe_bitreverse_neon<4>(bits, src, wire_code, count, + known_bitreverse); + return; + case 5: + flat_encode_maybe_bitreverse_neon<5>(bits, src, wire_code, count, + known_bitreverse); + return; + case 6: + flat_encode_maybe_bitreverse_neon<6>(bits, src, wire_code, count, + known_bitreverse); + return; + case 7: + flat_encode_maybe_bitreverse_neon<7>(bits, src, wire_code, count, + known_bitreverse); + return; + case 8: + flat_encode_maybe_bitreverse_neon<8>(bits, src, wire_code, count, + known_bitreverse); + return; + default: + return; + } +} + +#endif // PIXIE_PIVCO_NEON_SUPPORT + // =========================================================================== // AVX2 backend (SSE4.1 byte-shuffle lookup tables; AVX2 implies SSE4.1). // =========================================================================== @@ -1331,6 +2093,9 @@ inline void merge_decode_dispatch(std::uint8_t* dst, #elif defined(__AVX2__) && defined(__SSE4_1__) merge_decode_avx2( dst, left, left_symbol, right, right_symbol, bits, count); +#elif defined(PIXIE_PIVCO_NEON_SUPPORT) + merge_decode_neon( + dst, left, left_symbol, right, right_symbol, bits, count); #else merge_decode_scalar( dst, left, left_symbol, right, right_symbol, bits, count); @@ -1352,6 +2117,9 @@ inline void partition_encode_dispatch(std::uint8_t* left_dst, #elif defined(__AVX2__) && defined(__SSE4_1__) partition_encode_avx2( left_dst, right_dst, bits, src, dir, weight, left_weight, right_weight); +#elif defined(PIXIE_PIVCO_NEON_SUPPORT) + partition_encode_neon( + left_dst, right_dst, bits, src, dir, weight, left_weight, right_weight); #else partition_encode_scalar(left_dst, right_dst, bits, src, dir, weight); @@ -1379,6 +2147,12 @@ inline void flat_decode_dispatch(std::uint8_t* dst, bits, count); return; } +#elif defined(PIXIE_PIVCO_NEON_SUPPORT) + if (depth == 1) { + merge_decode_neon(dst, nullptr, table[0], nullptr, table[1], + bits, count); + return; + } #endif #if defined(__AVX512VBMI2__) && defined(__AVX512BW__) && defined(__AVX512VBMI__) @@ -1457,6 +2231,43 @@ inline void flat_decode_dispatch(std::uint8_t* dst, break; } #endif + +#if defined(PIXIE_PIVCO_NEON_SUPPORT) + switch (depth) { + case 2: + flat_decode_neon<2>(dst, bits, table, count); + return; + case 3: + flat_decode_neon<3>(dst, bits, table, count); + return; + case 4: + flat_decode_neon<4>(dst, bits, table, count); + return; + case 5: + flat_decode_neon<5>(dst, bits, table, count); + return; + case 6: + flat_decode_neon<6>(dst, bits, table, count); + return; + case 7: + if (known_bitreverse) { + flat_decode_neon<7, true>(dst, bits, table, count); + } else { + flat_decode_neon<7, false>(dst, bits, table, count); + } + return; + case 8: + if (known_bitreverse) { + flat_decode_neon<8, true>(dst, bits, table, count); + } else { + flat_decode_neon<8, false>(dst, bits, table, count); + } + return; + default: + break; + } +#endif + (void)known_bitreverse; flat_decode_scalar(dst, bits, table, count, depth); } @@ -1579,8 +2390,8 @@ inline void pivco_partition_encode_none(std::uint64_t* bits, /// @brief Build a bitmap for a node whose two children are constant leaves. /// @details A set bit means @p src equals @p right_symbol. Unlike the generic -/// direction-table path, AVX2 can compare 32 symbols directly and -/// obtain the complete routing word fragment with `vpmovmskb`. +/// direction-table path, SIMD backends compare complete vectors and +/// pack the equality results directly into routing bytes. inline void pivco_partition_encode_cst_cst(std::uint64_t* bits, const std::uint8_t* src, std::uint8_t right_symbol, @@ -1602,6 +2413,10 @@ inline void pivco_partition_encode_cst_cst(std::uint64_t* bits, _mm256_movemask_epi8(_mm256_cmpeq_epi8(symbols, right))); std::memcpy(bit_bytes + i / 8, &routing, sizeof(routing)); } +#elif defined(PIXIE_PIVCO_NEON_SUPPORT) + pivco_simd_detail::partition_encode_cst_cst_neon(bits, src, right_symbol, + weight); + return; #endif for (; i < weight; ++i) { if (src[i] == right_symbol) { @@ -1632,6 +2447,41 @@ inline void pivco_partition_encode_vec_cst(std::uint8_t* left_dst, pivco_simd_detail::partition_encode_one_constant_avx2( left_dst, bits, src, right_symbol, weight, left_weight); } +#elif defined(PIXIE_PIVCO_NEON_SUPPORT) +/// @brief Partition a constant left leaf and materialize the right stream. +inline void pivco_partition_encode_cst_vec(std::uint8_t* right_dst, + std::uint64_t* bits, + const std::uint8_t* src, + std::uint8_t left_symbol, + std::size_t weight, + std::size_t right_weight) noexcept { + pivco_simd_detail::partition_encode_one_constant_neon( + right_dst, bits, src, left_symbol, weight, right_weight); +} + +/// @brief Partition a dense left stream and a constant right leaf. +inline void pivco_partition_encode_vec_cst(std::uint8_t* left_dst, + std::uint64_t* bits, + const std::uint8_t* src, + std::uint8_t right_symbol, + std::size_t weight, + std::size_t left_weight) noexcept { + pivco_simd_detail::partition_encode_one_constant_neon( + left_dst, bits, src, right_symbol, weight, left_weight); +} +#endif + +#if defined(PIXIE_PIVCO_NEON_SUPPORT) +/// @brief Pack a depth-2..8 flat subtree through the AArch64 NEON backend. +inline void pivco_flat_encode_neon(std::uint64_t* bits, + const std::uint8_t* src, + const std::uint8_t* wire_code, + std::size_t count, + std::uint8_t depth, + bool known_bitreverse) noexcept { + pivco_simd_detail::flat_encode_dispatch_neon(bits, src, wire_code, count, + depth, known_bitreverse); +} #endif #if defined(__AVX2__) && defined(__SSE4_1__) && defined(__BMI2__) diff --git a/src/tests/pivco_huffman_tests.cpp b/src/tests/pivco_huffman_tests.cpp index d9883f8..0ab2f89 100644 --- a/src/tests/pivco_huffman_tests.cpp +++ b/src/tests/pivco_huffman_tests.cpp @@ -132,6 +132,64 @@ TEST(PivCoSimd, FlatKernelsDecodeFullGroupsAndTails) { } } +TEST(PivCoSimd, MergeVariantsMatchReferenceAcrossFullGroupsAndTail) { + constexpr std::size_t kCount = 263; + constexpr std::size_t kPadding = 64; + std::vector bits((kCount + 63) / 64, 0); + std::vector left(kCount + kPadding, 0); + std::vector right(kCount + kPadding, 0); + std::vector expected(kCount); + std::size_t left_count = 0; + std::size_t right_count = 0; + for (std::size_t i = 0; i < kCount; ++i) { + const bool goes_right = ((i * 37 + i / 5 + 11) % 13) < 6; + if (goes_right) { + bits[i / 64] |= std::uint64_t{1} << (i % 64); + right[right_count] = static_cast(0x80u + right_count % 97); + expected[i] = right[right_count++]; + } else { + left[left_count] = static_cast(left_count % 113); + expected[i] = left[left_count++]; + } + } + + std::vector decoded(kCount + kPadding, 0); + pixie::pivco_merge_decode(decoded.data(), left.data(), right.data(), + bits.data(), kCount); + EXPECT_TRUE(std::equal(expected.begin(), expected.end(), decoded.begin())); + + constexpr std::uint8_t kLeftSymbol = 17; + constexpr std::uint8_t kRightSymbol = 231; + for (std::size_t i = 0; i < kCount; ++i) { + const bool goes_right = ((bits[i / 64] >> (i % 64)) & 1u) != 0; + expected[i] = goes_right ? kRightSymbol : kLeftSymbol; + } + std::fill(decoded.begin(), decoded.end(), 0); + pixie::pivco_merge_decode_cst_cst(decoded.data(), kLeftSymbol, kRightSymbol, + bits.data(), kCount); + EXPECT_TRUE(std::equal(expected.begin(), expected.end(), decoded.begin())); + + right_count = 0; + for (std::size_t i = 0; i < kCount; ++i) { + const bool goes_right = ((bits[i / 64] >> (i % 64)) & 1u) != 0; + expected[i] = goes_right ? right[right_count++] : kLeftSymbol; + } + std::fill(decoded.begin(), decoded.end(), 0); + pixie::pivco_merge_decode_cst_vec(decoded.data(), kLeftSymbol, right.data(), + bits.data(), kCount); + EXPECT_TRUE(std::equal(expected.begin(), expected.end(), decoded.begin())); + + left_count = 0; + for (std::size_t i = 0; i < kCount; ++i) { + const bool goes_right = ((bits[i / 64] >> (i % 64)) & 1u) != 0; + expected[i] = goes_right ? kRightSymbol : left[left_count++]; + } + std::fill(decoded.begin(), decoded.end(), 0); + pixie::pivco_merge_decode_vec_cst(decoded.data(), left.data(), kRightSymbol, + bits.data(), kCount); + EXPECT_TRUE(std::equal(expected.begin(), expected.end(), decoded.begin())); +} + TEST(PivCoSimd, PartitionVariantsBuildTheSameDirectionMask) { constexpr std::size_t kCount = 257; std::array directions{}; @@ -187,6 +245,68 @@ TEST(PivCoSimd, PartitionVariantsBuildTheSameDirectionMask) { EXPECT_EQ(bitmap_only, expected_bits); } +TEST(PivCoSimd, ConstantLeafBitmapMatchesReferenceAcrossTail) { + constexpr std::size_t kCount = 267; + constexpr std::uint8_t kRightSymbol = 193; + std::vector input(kCount); + std::vector expected_bits((kCount + 63) / 64, 0); + for (std::size_t i = 0; i < kCount; ++i) { + const bool is_right = ((i * 19 + i / 3) % 11) < 5; + input[i] = is_right ? kRightSymbol : static_cast(i % 127); + if (is_right) { + expected_bits[i / 64] |= std::uint64_t{1} << (i % 64); + } + } + + std::vector bits(expected_bits.size(), 0); + pixie::pivco_partition_encode_cst_cst(bits.data(), input.data(), kRightSymbol, + input.size()); + EXPECT_EQ(bits, expected_bits); +} + +#if (defined(__AVX2__) && defined(__SSE4_1__)) || \ + defined(PIXIE_PIVCO_NEON_SUPPORT) +TEST(PivCoSimd, OneConstantPartitionsMatchReferenceAcrossTail) { + constexpr std::size_t kCount = 269; + constexpr std::uint8_t kConstant = 241; + std::vector input(kCount); + std::vector expected_dense; + std::vector expected_bits((kCount + 63) / 64, 0); + for (std::size_t i = 0; i < kCount; ++i) { + const bool is_constant = ((i * 23 + i / 7) % 17) < 8; + input[i] = is_constant ? kConstant : static_cast(i % 191); + if (!is_constant) { + expected_bits[i / 64] |= std::uint64_t{1} << (i % 64); + expected_dense.push_back(input[i]); + } + } + + std::vector dense(expected_dense.size()); + std::vector bits(expected_bits.size(), 0); + pixie::pivco_partition_encode_cst_vec(dense.data(), bits.data(), input.data(), + kConstant, input.size(), dense.size()); + EXPECT_EQ(dense, expected_dense); + EXPECT_EQ(bits, expected_bits); + + std::fill(expected_bits.begin(), expected_bits.end(), 0); + expected_dense.clear(); + for (std::size_t i = 0; i < kCount; ++i) { + const bool is_constant = input[i] == kConstant; + if (is_constant) { + expected_bits[i / 64] |= std::uint64_t{1} << (i % 64); + } else { + expected_dense.push_back(input[i]); + } + } + std::fill(dense.begin(), dense.end(), 0); + std::fill(bits.begin(), bits.end(), 0); + pixie::pivco_partition_encode_vec_cst(dense.data(), bits.data(), input.data(), + kConstant, input.size(), dense.size()); + EXPECT_EQ(dense, expected_dense); + EXPECT_EQ(bits, expected_bits); +} +#endif + TEST(PivCoHuffmanSmoke, RandomUniformRoundTrips) { std::mt19937 rng(239); std::uniform_int_distribution dist(0, 255); From 58937feb6cbe5df20e9fcab954f1807b2bdc1a9a Mon Sep 17 00:00:00 2001 From: Sanya239 Date: Tue, 28 Jul 2026 15:17:21 +0300 Subject: [PATCH 14/14] One huffman-tree instead of one per block --- include/pixie/huffman/pivco_huffman.h | 264 +++++++++++++++----------- src/tests/pivco_huffman_tests.cpp | 17 ++ 2 files changed, 167 insertions(+), 114 deletions(-) diff --git a/include/pixie/huffman/pivco_huffman.h b/include/pixie/huffman/pivco_huffman.h index 4e823b4..27bfb39 100644 --- a/include/pixie/huffman/pivco_huffman.h +++ b/include/pixie/huffman/pivco_huffman.h @@ -11,10 +11,11 @@ * using the node bitmap as a selector. * * The tree uses canonical Huffman code lengths (<= 15 bits, enforced with a - * zlib-style overflow correction): the wire format stores only the 256 - * per-symbol lengths (128 bytes as 4-bit nibbles), and the tree structure is - * reconstructed deterministically at load time. Same-length symbols are - * reshaped into flat subtrees as described by the PivCo-Huffman paper. + * zlib-style overflow correction): the wire format stores one file-level set + * of 256 per-symbol lengths (128 bytes as 4-bit nibbles), and the tree + * structure is reconstructed once at load time and reused by every block. + * Same-length symbols are reshaped into flat subtrees as described by the + * PivCo-Huffman paper. * * This implementation uses packed `std::uint64_t` bitmaps, flat subtrees, * contiguous arenas, and bottom-up decoding. The merge/partition primitives are @@ -44,8 +45,9 @@ namespace pixie { * byte alphabet, canonicalizing it (grouping same-length codes), * storing one routing bitmap per internal node as packed 64-bit * words, and decoding by bottom-up merging of child symbol streams. - * The serialized form is 256 code lengths followed by packed node - * bitmaps; the tree is reconstructed from the lengths at load time. + * The serialized form is one file-level set of 256 code lengths + * followed by block-local packed node bitmaps; the shared tree is + * reconstructed from the lengths at load time. */ class PivCoHuffman : public HuffmanBase { public: @@ -84,17 +86,19 @@ class PivCoHuffman : public HuffmanBase { if (uncompressed_size_ == 0) { return std::vector(); } - std::vector result; - result.reserve(uncompressed_size_); - // Skip the top-level header (total_size + block_size). - std::size_t pos = sizeof(std::size_t) * 2; - std::size_t remaining = uncompressed_size_; - while (remaining > 0) { - deserialize_block(pos); - std::vector block_result = decode_from_tree(); - result.insert(result.end(), block_result.begin(), block_result.end()); - remaining -= block_uncompressed_size_; - } + // SIMD merge kernels may touch one vector past their logical output. Keep + // the same slack as the internal workspace; resize it away without copying + // after the last block has decoded. + std::vector result(uncompressed_size_ + kWorkspaceSlack); + std::size_t pos = blocks_offset_; + std::size_t output_offset = 0; + while (output_offset < uncompressed_size_) { + deserialize_block_payload(pos); + decode_block_into(std::span(result).subspan( + output_offset, block_uncompressed_size_)); + output_offset += block_uncompressed_size_; + } + result.resize(uncompressed_size_); return result; } @@ -122,10 +126,10 @@ class PivCoHuffman : public HuffmanBase { static constexpr std::size_t kMaxTreeNodes = 2 * kAlphabet - 1; /// @brief Default block size for block-based processing (64 KiB). - /// @details Each block is compressed independently with its own Huffman - /// tree. This keeps the decode workspace (2 × block_size = 128 KiB) - /// within the per-core L2 (256 KiB), eliminating the L2 evictions - /// that hurt decode throughput on whole-input processing. + /// @details All blocks share one file-level Huffman tree but carry their own + /// routing bitmaps. This keeps the decode workspace (2 × block_size + /// = 128 KiB) within the per-core L2 while avoiding per-block tree + /// construction and serialization overhead. static constexpr std::size_t kBlockSize = 64 * 1024; /// @brief Extra bytes appended to the ping-pong workspace beyond `2 * block`. @@ -196,6 +200,7 @@ class PivCoHuffman : public HuffmanBase { using DirectionTables = std::array, kMaxCodeLength>; + using FrequencyTable = std::array; /// @brief Total uncompressed size across all blocks (set during /// build/deserialize, read by `uncompressed_size_impl()`). @@ -205,10 +210,12 @@ class PivCoHuffman : public HuffmanBase { std::size_t block_size_ = kBlockSize; /// @brief Serialized compressed stream (header + per-block data). std::vector compressed_; + /// @brief Byte offset of the first block payload in `compressed_`. + std::size_t blocks_offset_ = 0; - // --- per-block tree state (rebuilt for each block during decode) -------- - // Mutable because `decode_impl()` is const but rebuilds the tree - // per-block from the serialized stream. + // --- shared tree and current-block state -------------------------------- + // The topology and flat tables are file-level. Bitmap counts, offsets, + // and words are replaced for each block during const decoding. mutable std::size_t root_ = kNpos; mutable std::vector nodes_; /// @brief One contiguous allocation backing every internal node's bitmap @@ -227,69 +234,95 @@ class PivCoHuffman : public HuffmanBase { // --- construction -------------------------------------------------------- - /** @brief Build and serialize all blocks from @p input. - * @details Splits the input into fixed-size blocks (kBlockSize), builds - * an independent Huffman tree per block, and serializes each - * block into `compressed_`. The wire format is: + /** @brief Build and serialize a shared tree and all block payloads. + * @details Counts each block once, combines those counts to build one + * file-level Huffman tree, then reuses its topology and direction + * tables while encoding every block. The wire format is: * - total_uncompressed_size (8 bytes) * - block_size (8 bytes) - * - Per block: block_uncompressed_size (8 bytes) + 128 bytes - * nibbles + per-internal-node bitmaps. */ + * - 256 file-level code lengths as 128 bytes of nibbles + * - Per block: block_uncompressed_size (8 bytes) followed by + * per-internal-node bitmaps. */ void build(std::span input) { uncompressed_size_ = input.size(); compressed_.clear(); write(uncompressed_size_); write(block_size_); + blocks_offset_ = compressed_.size(); + if (input.empty()) { + return; + } + + const std::size_t block_count = + (input.size() + block_size_ - 1) / block_size_; + std::vector block_frequencies; + block_frequencies.reserve(block_count); + FrequencyTable file_frequencies{}; std::size_t offset = 0; while (offset < input.size()) { - std::size_t remaining = input.size() - offset; - std::size_t this_block = std::min(remaining, block_size_); - build_block(input.subspan(offset, this_block)); - serialize_block(); + const std::size_t this_block = + std::min(input.size() - offset, block_size_); + block_frequencies.push_back( + count_frequencies(input.subspan(offset, this_block))); + const FrequencyTable& frequencies = block_frequencies.back(); + for (std::size_t symbol = 0; symbol < kAlphabet; ++symbol) { + file_frequencies[symbol] += frequencies[symbol]; + } + offset += this_block; + } + + compute_code_lengths(file_frequencies, code_lengths_); + build_noncanonical_tree(code_lengths_); + detect_flat_subtrees(); + serialize_code_lengths(); + blocks_offset_ = compressed_.size(); + + std::array codes{}; + DirectionTables directions{}; + if (root_ != kNpos && !nodes_[root_].is_leaf) { + assign_codes(root_, 0, 0, codes); + directions = build_direction_tables(codes); + } + + offset = 0; + std::size_t block_index = 0; + while (offset < input.size()) { + const std::size_t this_block = + std::min(input.size() - offset, block_size_); + build_block_payload(input.subspan(offset, this_block), + block_frequencies[block_index], directions); + serialize_block_payload(); offset += this_block; + ++block_index; } } - /** @brief Build the canonical Huffman tree and per-node bitmaps for one - * block of @p input. */ - void build_block(std::span input) { + /** @brief Build one block's bitmaps using the file-level tree. */ + void build_block_payload(std::span input, + const FrequencyTable& freq, + const DirectionTables& directions) { block_uncompressed_size_ = input.size(); if (block_uncompressed_size_ == 0) { - root_ = kNpos; return; } - // 1. Count symbol frequencies using independent partial histograms. - const std::array freq = count_frequencies(input); - - // 2. Build a Huffman tree to determine code lengths, then extract lengths. - compute_code_lengths(freq, code_lengths_); - - // 3. Rebuild a canonical tree from the lengths (groups same-length codes). - // Structure only: bitmap weights/offsets are assigned below, not here. - build_noncanonical_tree(code_lengths_); - detect_flat_subtrees(); - // Single-symbol input: root is a leaf, no internal nodes, no arena. if (root_ == kNpos || nodes_[root_].is_leaf) { return; } - // 4. Assign exact per-node weights (subtree frequency sums) and arena + // Assign exact per-node weights (subtree frequency sums) and arena // offsets in a single post-order walk, then allocate the arena once. std::size_t arena_words = 0; assign_weights_and_offsets(root_, freq, arena_words); arena_.assign(arena_words, 0); - // 5. Assign compact per-symbol codes, then fill node bitmaps via top-down - // in-place partition into the reusable workspace. This mirrors decode: + // Fill node bitmaps via top-down in-place partition into the reusable + // workspace. This mirrors decode: // the root reads its input from workspace half 0, writes its bitmap, // and partitions symbols into half 1 for its children. Each level // ping-pongs between the two halves — no per-node allocations. - std::array codes{}; - assign_codes(root_, 0, 0, codes); - const DirectionTables directions = build_direction_tables(codes); workspace_.resize(block_uncompressed_size_ * 2 + kWorkspaceSlack); std::copy(input.begin(), input.end(), workspace_.data()); encode_partition(root_, 0, 0, directions, freq); @@ -299,8 +332,7 @@ class PivCoHuffman : public HuffmanBase { * @details Eight partial histograms break the dependency chain caused by * repeatedly incrementing a hot symbol. The block is at most * 64 KiB, so 32-bit partial counters cannot overflow. */ - static std::array count_frequencies( - std::span input) { + static FrequencyTable count_frequencies(std::span input) { constexpr std::size_t kLanes = 8; std::array, kLanes> partial{}; std::size_t i = 0; @@ -318,7 +350,7 @@ class PivCoHuffman : public HuffmanBase { partial[0][input[i]]++; } - std::array frequencies{}; + FrequencyTable frequencies{}; for (std::size_t symbol = 0; symbol < kAlphabet; ++symbol) { for (std::size_t lane = 0; lane < kLanes; ++lane) { frequencies[symbol] += partial[lane][symbol]; @@ -953,6 +985,9 @@ class PivCoHuffman : public HuffmanBase { } Node& n = nodes_[idx]; + if (n.bits.count == 0) { + return; + } // Flat node: write d-bit suffix per symbol, no partition, no recursion. if (n.flat_depth > 0) { @@ -1016,42 +1051,43 @@ class PivCoHuffman : public HuffmanBase { // --- decode -------------------------------------------------------------- - /** @brief Reconstruct the original sequence by bottom-up merging. - * @details Uses a single reusable workspace buffer (sized to twice the - * uncompressed length, allocated once and reused across calls) - * and ping-pongs between its two halves by tree depth. Each node - * reads its children's outputs from one half and writes its merged - * output to the other, so no per-node allocation occurs after the - * workspace is sized. The final result is copied out of half 0. */ - std::vector decode_from_tree() const { - if (root_ == kNpos) { - return std::vector(); + /** @brief Decode the current block directly into its final output slice. + * @details Internal levels still ping-pong through the reusable two-block + * workspace, but the root writes into @p output. This avoids a + * temporary block vector and the subsequent append/copy. */ + void decode_block_into(std::span output) const { + if (root_ == kNpos || output.empty()) { + return; } if (nodes_[root_].is_leaf) { - return std::vector(block_uncompressed_size_, - nodes_[root_].symbol); + std::fill(output.begin(), output.end(), nodes_[root_].symbol); + return; } - const std::size_t n = block_uncompressed_size_; - workspace_.resize(n * 2 + kWorkspaceSlack); - decode_node(root_, n, 0, 0); - return std::vector(workspace_.begin(), workspace_.begin() + n); + workspace_.resize(output.size() * 2 + kWorkspaceSlack); + decode_node(root_, output.size(), 0, 0, output.data()); } /** @brief Recursively decode a node's output stream bottom-up into the * reusable workspace. * @param weight Number of symbols this node must produce. - * @param depth Tree depth, selecting which workspace half receives this - * node's output; children write to the opposite half. + * @param depth Tree depth. The root writes directly to the caller's + * output; deeper levels select a workspace half. * @param out_base Offset within the selected half where this node's output * begins. Children are placed at `out_base` (left) and * `out_base + left_weight` (right) in the opposite half, - * then merged into `[out_base, out_base + weight)` here. */ + * then merged into `[out_base, out_base + weight)` here. + * @param block_output Final output address used by the root node. */ void decode_node(std::size_t idx, std::size_t weight, std::size_t depth, - std::size_t out_base) const { + std::size_t out_base, + symbol_type* block_output) const { + if (weight == 0) { + return; + } const Node& n = nodes_[idx]; - symbol_type* dst = workspace_half(depth) + out_base; + symbol_type* dst = + depth == 0 ? block_output : workspace_half(depth) + out_base; if (n.is_leaf) { std::fill(dst, dst + weight, n.symbol); return; @@ -1081,10 +1117,11 @@ class PivCoHuffman : public HuffmanBase { left.is_leaf ? weight - right.bits.count : left.bits.count; const std::size_t right_weight = weight - left_weight; if (!left.is_leaf) { - decode_node(n.left, left_weight, depth + 1, out_base); + decode_node(n.left, left_weight, depth + 1, out_base, block_output); } if (!right.is_leaf) { - decode_node(n.right, right_weight, depth + 1, out_base + left_weight); + decode_node(n.right, right_weight, depth + 1, out_base + left_weight, + block_output); } const symbol_type* src = workspace_half(depth + 1) + out_base; @@ -1106,26 +1143,27 @@ class PivCoHuffman : public HuffmanBase { // --- serialization ------------------------------------------------------- - /** @brief Append one block's serialized form to `compressed_`. + /** @brief Append the shared file-level code lengths as packed nibbles. */ + void serialize_code_lengths() { + for (std::size_t i = 0; i < kAlphabet; i += 2) { + std::uint8_t byte = static_cast( + code_lengths_[i] | (code_lengths_[i + 1] << 4)); + write(byte); + } + } + + /** @brief Append one block's serialized payload to `compressed_`. * @details Per-block wire format: * - block_uncompressed_size (8 bytes) - * - 256 code lengths as 128 bytes of 4-bit nibbles (0 = absent) * - For each internal node (pre-order): bits.count (8 bytes) + - * packed words. Leaves are implied by the tree structure - * reconstructed from lengths. */ - void serialize_block() { + * packed words. Leaves and topology are implied by the shared + * file-level code lengths. */ + void serialize_block_payload() { write(block_uncompressed_size_); if (block_uncompressed_size_ == 0) { return; } - // Pack stored code lengths as nibbles: 2 symbols per byte. - for (std::size_t i = 0; i < kAlphabet; i += 2) { - std::uint8_t byte = static_cast( - code_lengths_[i] | (code_lengths_[i + 1] << 4)); - write(byte); - } - // Write internal node bitmaps in pre-order (matches deserialize order). serialize_node(root_); } @@ -1150,10 +1188,10 @@ class PivCoHuffman : public HuffmanBase { serialize_node(n.right); } - /** @brief Load the serialized stream and parse the block header. + /** @brief Load the serialized stream and reconstruct its shared tree. * @details Copies @p data into `compressed_` and reads the top-level - * header (total_uncompressed_size, block_size). Per-block trees - * are rebuilt on demand during `decode_impl()`. */ + * header and file-level code lengths. Per-block bitmap payloads + * are loaded on demand during `decode_impl()`. */ void deserialize(std::span data) { compressed_.assign(data.begin(), data.end()); nodes_.clear(); @@ -1161,38 +1199,36 @@ class PivCoHuffman : public HuffmanBase { root_ = kNpos; if (data.empty()) { uncompressed_size_ = 0; + blocks_offset_ = 0; return; } std::size_t pos = 0; uncompressed_size_ = read(compressed_, pos); block_size_ = read(compressed_, pos); - } - - /** @brief Rebuild the tree and arena for one block from `compressed_` at - * @p pos. Advances @p pos past the block's serialized data. - * @details Reads the per-block header (uncompressed_size + nibbles), - * rebuilds the canonical tree, and reads arena words via a - * two-pass wire walk. */ - void deserialize_block(std::size_t& pos) const { - nodes_.clear(); - arena_.clear(); - root_ = kNpos; - - block_uncompressed_size_ = read(compressed_, pos); - if (block_uncompressed_size_ == 0) { + if (uncompressed_size_ == 0) { + blocks_offset_ = pos; return; } - // Read 256 code lengths from nibbles. for (std::size_t i = 0; i < kAlphabet; i += 2) { - std::uint8_t byte = read(compressed_, pos); + const std::uint8_t byte = read(compressed_, pos); code_lengths_[i] = byte & 0x0F; code_lengths_[i + 1] = (byte >> 4) & 0x0F; } - - // Rebuild the canonical tree from lengths. build_noncanonical_tree(code_lengths_); detect_flat_subtrees(); + blocks_offset_ = pos; + } + + /** @brief Load one block's arena from `compressed_` at @p pos. + * @details Reuses the file-level tree, reads the block's uncompressed size, + * and loads its arena words via a two-pass wire walk. Advances + * @p pos past the block payload. */ + void deserialize_block_payload(std::size_t& pos) const { + block_uncompressed_size_ = read(compressed_, pos); + if (block_uncompressed_size_ == 0) { + return; + } // Single-symbol input: root is a leaf, no internal nodes, no arena. if (root_ == kNpos || nodes_[root_].is_leaf) { diff --git a/src/tests/pivco_huffman_tests.cpp b/src/tests/pivco_huffman_tests.cpp index 0ab2f89..965ac4c 100644 --- a/src/tests/pivco_huffman_tests.cpp +++ b/src/tests/pivco_huffman_tests.cpp @@ -368,6 +368,23 @@ TEST(PivCoHuffmanSmoke, MultipleBlocksRoundTrip) { EXPECT_EQ(decode_via_bytes(codec), data); } +TEST(PivCoHuffmanSmoke, SharedTreeHandlesBlockLocalMissingSymbols) { + constexpr std::size_t kBlockSize = 64 * 1024; + std::vector data; + data.reserve(3 * kBlockSize); + data.insert(data.end(), kBlockSize, 17); + data.insert(data.end(), kBlockSize, 93); + data.insert(data.end(), kBlockSize, 241); + + const PivCoHuffman codec(data); + EXPECT_EQ(codec.decode(), data); + EXPECT_EQ(codec.decode(), data); + + const PivCoHuffman copy(codec.compressed_data()); + EXPECT_EQ(copy.decode(), data); + EXPECT_EQ(copy.decode(), data); +} + TEST(PivCoHuffmanSmoke, RandomizedAlphabetAndBlockSizesRoundTrip) { std::mt19937 rng(3319); for (std::size_t test_case = 0; test_case < 32; ++test_case) {