diff --git a/src/openzl/codecs/common/bitstream/ff_bitstream.h b/src/openzl/codecs/common/bitstream/ff_bitstream.h index b1d7202ee..31bfa72bd 100644 --- a/src/openzl/codecs/common/bitstream/ff_bitstream.h +++ b/src/openzl/codecs/common/bitstream/ff_bitstream.h @@ -31,6 +31,10 @@ ZL_INLINE ZL_Report ZS_BitCStreamFF_finish(ZS_BitCStreamFF* bits); ZL_INLINE void ZS_BitCStreamFF_write(ZS_BitCStreamFF* bits, size_t value, size_t nbBits); ZL_INLINE void ZS_BitCStreamFF_flush(ZS_BitCStreamFF* bits); +ZL_INLINE uint8_t* ZS_BitCStreamFF_reserveAlignedBits( + ZS_BitCStreamFF* bits, + size_t nbBits); +ZL_INLINE void ZS_BitCStreamFF_commitReservedBits(ZS_BitCStreamFF* bits); ZL_INLINE void ZS_BitCStreamFF_writeExpGolomb( ZS_BitCStreamFF* bits, @@ -52,6 +56,7 @@ typedef struct { uint8_t const* ptr; uint8_t const* limit; uint8_t const* end; + uint8_t const* begin; } ZS_BitDStreamFF; ZL_INLINE ZS_BitDStreamFF @@ -62,6 +67,9 @@ ZL_INLINE size_t ZS_BitDStreamFF_peek(ZS_BitDStreamFF const* bits, size_t nbBits); ZL_INLINE void ZS_BitDStreamFF_skip(ZS_BitDStreamFF* bits, size_t nbBits); ZL_INLINE void ZS_BitDStreamFF_reload(ZS_BitDStreamFF* bits); +ZL_INLINE uint8_t const* ZS_BitDStreamFF_popAlignedBits( + ZS_BitDStreamFF* bits, + size_t nbBits); ZL_INLINE uint32_t ZS_BitDStreamFF_readExpGolomb(ZS_BitDStreamFF* bits, size_t order) @@ -80,11 +88,13 @@ ZS_BitDStreamFF_readExpGolomb(ZS_BitDStreamFF* bits, size_t order) ZL_INLINE ZS_BitCStreamFF ZS_BitCStreamFF_init(uint8_t* dst, size_t dstCapacity) { + const size_t limit = + dstCapacity < sizeof(size_t) ? 0 : dstCapacity - sizeof(size_t) + 1; return (ZS_BitCStreamFF){ .container = 0, .nbBits = 0, .ptr = dst, - .limit = dst + dstCapacity - sizeof(size_t), + .limit = dst + limit, .end = dst + dstCapacity, .begin = dst, }; @@ -94,7 +104,7 @@ ZL_INLINE ZL_Report ZS_BitCStreamFF_finish(ZS_BitCStreamFF* bits) { ZL_RESULT_DECLARE_SCOPE_REPORT((ZL_OperationContext*)NULL); size_t bytesToWrite = (bits->nbBits + 7) / 8; - if (bits->end < bits->ptr + bytesToWrite) + if ((size_t)(bits->end - bits->ptr) < bytesToWrite) ZL_ERR(internalBuffer_tooSmall); if (bytesToWrite) { ZL_ASSERT_EQ( @@ -116,7 +126,7 @@ ZS_BitCStreamFF_write(ZS_BitCStreamFF* bits, size_t value, size_t nbBits) ZL_INLINE void ZS_BitCStreamFF_flush(ZS_BitCStreamFF* bits) { - if (bits->ptr > bits->limit) { + if (bits->ptr >= bits->limit) { return; } size_t const nbBytes = bits->nbBits >> 3; @@ -126,6 +136,62 @@ ZL_INLINE void ZS_BitCStreamFF_flush(ZS_BitCStreamFF* bits) bits->container >>= (nbBytes << 3); } +ZL_INLINE uint8_t* ZS_BitCStreamFF_reserveAlignedBits( + ZS_BitCStreamFF* bits, + size_t nbBits) +{ + if (nbBits > SIZE_MAX - 7) { + return NULL; + } + + size_t const bufferedBytes = (bits->nbBits + 7) / 8; + if (bufferedBytes > (size_t)(bits->end - bits->ptr)) { + return NULL; + } + + size_t const nbBytes = (nbBits + 7) / 8; + size_t const available = (size_t)(bits->end - bits->ptr) - bufferedBytes; + if (nbBytes > available) { + return NULL; + } + + ZL_ASSERT_LE(bufferedBytes, sizeof(size_t)); + if (bufferedBytes != 0) { + ZL_writeLE64_N(bits->ptr, bits->container, bufferedBytes); + bits->ptr += bufferedBytes; + bits->container = 0; + bits->nbBits = 0; + } + ZL_ASSERT_EQ(bits->nbBits, 0); + + uint8_t* const out = bits->ptr; + bits->ptr = out + (nbBits / 8); + bits->nbBits = nbBits & 7; + bits->container = 0; + return out; +} + +ZL_INLINE void ZS_BitCStreamFF_commitReservedBits(ZS_BitCStreamFF* bits) +{ + if (bits->nbBits == 0) { + bits->container = 0; + return; + } + bits->container = bits->ptr[0] & (((size_t)1 << bits->nbBits) - 1); +} + +// Little-endian load of the first @p nbBytes (< sizeof(size_t)) bytes at @p src +// into a container word. Used for streams shorter than a full word, where every +// byte lives in the container. +ZL_INLINE size_t ZS_BitDStreamFF_loadPartial(uint8_t const* src, size_t nbBytes) +{ + size_t container = 0; + for (size_t i = 0; i < nbBytes; ++i) { + container |= (size_t)src[i] << (i << 3); + } + return container; +} + ZL_INLINE ZS_BitDStreamFF ZS_BitDStreamFF_init(uint8_t const* src, size_t srcSize) { @@ -136,19 +202,19 @@ ZS_BitDStreamFF_init(uint8_t const* src, size_t srcSize) .ptr = src, .limit = src + srcSize - sizeof(size_t) + 1, .end = src + srcSize, + .begin = src, }; return bits; } else { - ZS_BitDStreamFF bits = { - .container = 0, - .nbBitsRead = (sizeof(size_t) - srcSize) * 8, - .ptr = src + srcSize, - .limit = src, - .end = src + srcSize, + uint8_t const* const end = srcSize > 0 ? src + srcSize : src; + ZS_BitDStreamFF bits = { + .container = ZS_BitDStreamFF_loadPartial(src, srcSize), + .nbBitsRead = (sizeof(size_t) - srcSize) * 8, + .ptr = end, + .limit = src, + .end = end, + .begin = src, }; - for (size_t i = 0; i < srcSize; ++i) { - bits.container |= (size_t)src[i] << (i << 3); - } return bits; } } @@ -159,7 +225,13 @@ ZL_INLINE ZL_Report ZS_BitDStreamFF_finish(ZS_BitDStreamFF const* bits) if (bits->nbBitsRead > ZS_BITSTREAM_READ_MAX_BITS) { ZL_ERR(GENERIC); } - return ZL_returnSuccess(); + size_t bytesRead = (size_t)(bits->ptr - bits->begin); + bytesRead += bits->nbBitsRead >> 3; + bytesRead += (bits->nbBitsRead & 7) != 0; + if ((size_t)(bits->end - bits->begin) < sizeof(size_t)) { + bytesRead -= sizeof(size_t); + } + return ZL_returnValue(bytesRead); } ZL_INLINE size_t ZS_BitDStreamFF_read(ZS_BitDStreamFF* bits, size_t nbBits) @@ -189,22 +261,93 @@ ZL_INLINE void ZS_BitDStreamFF_skip(ZS_BitDStreamFF* bits, size_t nbBits) ZL_INLINE void ZS_BitDStreamFF_reload(ZS_BitDStreamFF* bits) { - bits->ptr += bits->nbBitsRead >> 3; - if (ZL_LIKELY(bits->ptr < bits->limit)) { - size_t const next = ZL_readLEST(bits->ptr); + uint8_t const* const ptr = bits->ptr + (bits->nbBitsRead >> 3); + if (ZL_LIKELY(ptr < bits->limit)) { + bits->ptr = ptr; bits->nbBitsRead &= 7; - bits->container = next >> bits->nbBitsRead; + bits->container = ZL_readLEST(ptr) >> bits->nbBitsRead; return; } - if (bits->ptr >= bits->end) + // Past the end: leave ptr and nbBitsRead untouched. This keeps the consumed + // bit count -- (ptr - begin) * 8 + nbBitsRead -- exact for finish() and + // popAlignedBits(), while an over-read still leaves nbBitsRead too large + // for finish() to accept. + if (ptr >= bits->end) return; - uint8_t const* const limit = bits->limit - 1; - size_t const skippedBits = (size_t)((bits->ptr - limit) << 3); - size_t const next = ZL_readLEST(limit); + uint8_t const* const tail = bits->limit - 1; + size_t const skippedBits = (size_t)((ptr - tail) << 3); + bits->ptr = ptr; bits->nbBitsRead &= 7; - bits->container = next >> (bits->nbBitsRead + skippedBits); + bits->container = ZL_readLEST(tail) >> (bits->nbBitsRead + skippedBits); +} + +ZL_INLINE uint8_t const* ZS_BitDStreamFF_popAlignedBits( + ZS_BitDStreamFF* bits, + size_t nbBits) +{ + if (nbBits > SIZE_MAX - 7) { + return NULL; + } + + // check if the stream is already in an invalid state + if (bits->nbBitsRead > ZS_BITSTREAM_READ_MAX_BITS) { + return NULL; + } + + size_t const streamSize = (size_t)(bits->end - bits->begin); + if (streamSize > SIZE_MAX / 8) { + return NULL; + } + + size_t consumedBits = + (size_t)(bits->ptr - bits->begin) * 8 + bits->nbBitsRead; + if (streamSize < sizeof(size_t)) { + consumedBits -= sizeof(size_t) * 8; + } + + size_t const bitSize = streamSize * 8; + if (consumedBits > bitSize || consumedBits > SIZE_MAX - 7) { + return NULL; + } + + size_t const alignedBits = ((consumedBits + 7) / 8) * 8; + if (nbBits > bitSize - alignedBits) { + return NULL; + } + + uint8_t const* const out = bits->begin + alignedBits / 8; + + // Reposition the stream just past the aligned region we handed out. + size_t const bitPos = alignedBits + nbBits; + uint8_t const* const pos = bits->begin + bitPos / 8; + size_t const bitShift = bitPos & 7; + + if (streamSize >= sizeof(size_t)) { + // Refill the 64-bit window at `pos`, clamping the load back from the + // end when fewer than 8 bytes remain; the matching shift exposes the + // same bit either way. A shift of the full word width (the whole + // stream consumed) leaves an empty window. + uint8_t const* const tail = bits->end - sizeof(size_t); + uint8_t const* const load = pos <= tail ? pos : tail; + size_t const shift = bitPos - (size_t)(load - bits->begin) * 8; + bits->container = + shift < sizeof(size_t) * 8 ? ZL_readLEST(load) >> shift : 0; + bits->nbBitsRead = bitShift; + bits->ptr = pos; + } else { + // Short stream: every byte already lives in the container. + size_t const remaining = (size_t)(bits->end - pos); + bits->container = 0; + for (size_t i = 0; i < remaining; ++i) { + bits->container |= (size_t)pos[i] << (i << 3); + } + bits->container >>= bitShift; + bits->nbBitsRead = (sizeof(size_t) - remaining) * 8 + bitShift; + bits->ptr = bits->end; + } + return out; } ZL_END_C_DECLS diff --git a/src/openzl/codecs/pivco_huffman/README.md b/src/openzl/codecs/pivco_huffman/README.md new file mode 100644 index 000000000..5e625781b --- /dev/null +++ b/src/openzl/codecs/pivco_huffman/README.md @@ -0,0 +1,22 @@ +# PivCo-Huffman + +This is an implementation of PivCo-Huffman coding from Marcin Zukowski. +The code is based on ideas from: + +- [The paper](https://arxiv.org/abs/2606.05765) +- [The repo](https://github.com/MarcinZukowski/pivco-huffman) +- [Ryg's blog](https://fgiesen.wordpress.com/2026/06/21/pivco-huffman-merge-operations/) + +OpenZL re-implements PivCo-Huffman instead of using the upstream repo for several reasons: + +1. OpenZL needs full control over the wire-format. +2. OpenZL needs to be hardened against decoding corrupted data. +3. OpenZL needs to be able to dynamically dispatch to kernels based on the CPUID, +e.g. to detect AVX512 support at runtime. + +The goal is to work with upstream to factor out the primitives so that OpenZL can +import and use upstream's primitives, as those are the most complex part of PivCo, +and are independent of the wire format. + +In the meantime, any meaningful improvements will be upstreamed, so there is a +single reference implementation for PivCo Huffman. diff --git a/src/openzl/codecs/pivco_huffman/arch/common_pivco_arch.h b/src/openzl/codecs/pivco_huffman/arch/common_pivco_arch.h new file mode 100644 index 000000000..48c0d5d17 --- /dev/null +++ b/src/openzl/codecs/pivco_huffman/arch/common_pivco_arch.h @@ -0,0 +1,20 @@ +// Copyright (c) Meta Platforms, Inc. and affiliates. +#ifndef OPENZL_CODECS_PIVCO_HUFFMAN_ARCH_COMMON_PIVCO_ARCH_H +#define OPENZL_CODECS_PIVCO_HUFFMAN_ARCH_COMMON_PIVCO_ARCH_H + +#include "openzl/shared/portability.h" + +ZL_BEGIN_C_DECLS + +/** + * Slop bytes that every bitmap and rank buffer must reserve past its logical + * length. The encode and decode kernels may over-read and over-write up to this + * many trailing bytes where noted in the docs as `SLOP`, so callers pad their + * buffers by this amount. It is sized to the widest fixed-width group the SIMD + * kernels process. + */ +#define ZL_PIVCO_HUFFMAN_SLOP 64 + +ZL_END_C_DECLS + +#endif diff --git a/src/openzl/codecs/pivco_huffman/arch/decode_pivco_arch.c b/src/openzl/codecs/pivco_huffman/arch/decode_pivco_arch.c new file mode 100644 index 000000000..f9e984558 --- /dev/null +++ b/src/openzl/codecs/pivco_huffman/arch/decode_pivco_arch.c @@ -0,0 +1,204 @@ +// Copyright (c) Meta Platforms, Inc. and affiliates. + +#include "openzl/codecs/pivco_huffman/arch/decode_pivco_arch.h" + +#include + +#include "openzl/codecs/common/bitstream/ff_bitstream.h" + +// Portable reference implementation of the PivCo-Huffman decode kernels. +// +// This file is written for clarity, not speed. It is used for correctness +// testing and as the fallback on hardware without a specialized kernel, which +// is not expected to run this code in practice. The architecture-specific +// kernels (x86, arm, avx512) are the fast paths and MUST decode the exact same +// bitmaps as the code here. +// +// Every bitmap is packed little-endian, least-significant-bit first -- exactly +// the layout ZS_BitDStreamFF reads -- so we lean on that primitive instead of +// hand-rolling the bit math. To keep the loops trivially correct we reload the +// 64-bit window after every element. That is wasteful but easy to follow, which +// is the point of this file. + +const ZL_PivCoHuffmanDecode* ZL_PivCoHuffmanDecode_select( + const ZL_cpuid_t* cpuid) +{ + (void)cpuid; + return &ZL_PivCoHuffmanDecode_generic; +} + +ZL_INLINE size_t bitmapBytes(size_t bits) +{ + return (bits + 7) / 8; +} + +// One side of a merge: either a single byte repeated (a constant leaf) or a +// vector of decoded bytes consumed in order. `pos` counts how many have been +// taken so far. +typedef struct { + const uint8_t* vec; // NULL for a constant source + uint8_t constant; + size_t size; + size_t pos; +} MergeSource; + +static MergeSource vectorSource(const uint8_t* vec, size_t size) +{ + return (MergeSource){ .vec = vec, .constant = 0, .size = size, .pos = 0 }; +} + +static MergeSource constantSource(uint8_t value, size_t size) +{ + return (MergeSource){ + .vec = NULL, .constant = value, .size = size, .pos = 0 + }; +} + +static uint8_t mergeSourceNext(MergeSource* source) +{ + uint8_t value; + if (source->vec == NULL) { + value = source->constant; + } else if (source->pos < source->size) { + value = source->vec[source->pos]; + } else { + // More bits selected this source than it has bytes: the bitstream is + // corrupt. Return a placeholder; the caller detects the corruption + // because the returned one-count will not match the expected size. + value = 0; + } + ++source->pos; + return value; +} + +// Reconstructs lhs.size + rhs.size bytes into @p out from a partition bitmap: +// bit i selects the right source (1) or the left source (0) for output position +// i, taking one byte from the chosen source. Returns the number of one bits. +static size_t mergeGeneric( + uint8_t* out, + size_t outCapacity, + const uint8_t* bitmap, + size_t bitmapCapacity, + MergeSource lhs, + MergeSource rhs) +{ + const size_t outSize = lhs.size + rhs.size; + assert(outCapacity >= outSize); + assert(bitmapCapacity >= bitmapBytes(outSize)); + (void)outCapacity; // only read by the assert above + + ZS_BitDStreamFF bits = ZS_BitDStreamFF_init(bitmap, bitmapCapacity); + + size_t ones = 0; + for (size_t i = 0; i < outSize; ++i) { + const size_t bit = ZS_BitDStreamFF_read(&bits, 1); + // Reload every bit so the 64-bit window never underflows. + ZS_BitDStreamFF_reload(&bits); + + if (bit != 0) { + out[i] = mergeSourceNext(&rhs); + ++ones; + } else { + out[i] = mergeSourceNext(&lhs); + } + } + return ones; +} + +static size_t mergeVectorVector( + uint8_t* out, + size_t outCapacity, + const uint8_t* bitmap, + size_t bitmapCapacity, + const uint8_t* lhs, + size_t lhsSize, + const uint8_t* rhs, + size_t rhsSize) +{ + return mergeGeneric( + out, + outCapacity, + bitmap, + bitmapCapacity, + vectorSource(lhs, lhsSize), + vectorSource(rhs, rhsSize)); +} + +static size_t mergeConstantVector( + uint8_t* out, + size_t outCapacity, + const uint8_t* bitmap, + size_t bitmapCapacity, + uint8_t lhs, + size_t lhsSize, + const uint8_t* rhs, + size_t rhsSize) +{ + return mergeGeneric( + out, + outCapacity, + bitmap, + bitmapCapacity, + constantSource(lhs, lhsSize), + vectorSource(rhs, rhsSize)); +} + +static size_t mergeVectorConstant( + uint8_t* out, + size_t outCapacity, + const uint8_t* bitmap, + size_t bitmapCapacity, + const uint8_t* lhs, + size_t lhsSize, + uint8_t rhs, + size_t rhsSize) +{ + return mergeGeneric( + out, + outCapacity, + bitmap, + bitmapCapacity, + vectorSource(lhs, lhsSize), + constantSource(rhs, rhsSize)); +} + +// Expands a flat leaf: reads one @p depth-bit index per output and looks it up +// in @p symbols (which has 2^depth entries). +static void mergeFlatDepth( + uint8_t* out, + size_t outSize, + size_t outCapacity, + const uint8_t* bitmap, + size_t bitmapCapacity, + size_t depth, + const uint8_t* symbols) +{ + assert(depth >= 1); + assert(depth <= 8); + assert(outCapacity >= outSize); + assert(bitmapCapacity >= bitmapBytes(outSize * depth)); + (void)outCapacity; // only read by the assert above + + ZS_BitDStreamFF bits = ZS_BitDStreamFF_init(bitmap, bitmapCapacity); + + for (size_t i = 0; i < outSize; ++i) { + const size_t index = ZS_BitDStreamFF_read(&bits, depth); + // Reload every index so the 64-bit window never underflows. + ZS_BitDStreamFF_reload(&bits); + out[i] = symbols[index]; + } +} + +static bool supported(const ZL_cpuid_t* cpuid) +{ + (void)cpuid; + return true; +} + +const ZL_PivCoHuffmanDecode ZL_PivCoHuffmanDecode_generic = { + .supported = supported, + .mergeVectorVector = mergeVectorVector, + .mergeConstantVector = mergeConstantVector, + .mergeVectorConstant = mergeVectorConstant, + .mergeFlatDepth = mergeFlatDepth, +}; diff --git a/src/openzl/codecs/pivco_huffman/arch/decode_pivco_arch.h b/src/openzl/codecs/pivco_huffman/arch/decode_pivco_arch.h new file mode 100644 index 000000000..e90d6bb9d --- /dev/null +++ b/src/openzl/codecs/pivco_huffman/arch/decode_pivco_arch.h @@ -0,0 +1,131 @@ +// Copyright (c) Meta Platforms, Inc. and affiliates. +#ifndef OPENZL_CODECS_PIVCO_HUFFMAN_ARCH_DECODE_PIVCO_ARCH_H +#define OPENZL_CODECS_PIVCO_HUFFMAN_ARCH_DECODE_PIVCO_ARCH_H + +#include +#include + +#include "openzl/codecs/pivco_huffman/arch/common_pivco_arch.h" +#include "openzl/shared/cpu.h" +#include "openzl/shared/portability.h" + +ZL_BEGIN_C_DECLS + +typedef struct { + /** + * @returns true iff the CPU supports this implementation. + * @note The rest of the function pointers may be NULL if not supported. + */ + bool (*supported)(const ZL_cpuid_t* cpuid); + + /** + * Merges @p lhs and @p rhs into @p out based on the bits in @p bitmap. + * + * @param outCapacity Must be at least `lhsSize + rhsSize` bytes, but may + * be larger to indicate that over-writes are okay. + * @param bitmapCapacity Must be at least `(lhsSize + rhsSize + 7) / 8` + * bytes, but may be larger to indicate that + * over-reads are okay. + * @param lhs Must be at least `lhsSize + SLOP` bytes, + * where the first @p lhsSize bytes are valid. + * @param rhs Must be at least `rhsSize + SLOP` bytes, + * where the first @p rhsSize bytes are valid. + * + * @returns The number of ones in the bitmap. If this is not equal to + * @p rhsSize then the data was corrupt, and the output is unspecified. + */ + size_t (*mergeVectorVector)( + uint8_t* out, + size_t outCapacity, + const uint8_t* bitmap, + size_t bitmapCapacity, + const uint8_t* lhs, + size_t lhsSize, + const uint8_t* rhs, + size_t rhsSize); + + /** + * Merges @p lhs and @p rhs into @p out based on the bits in @p bitmap, + * where @p lhs is a constant. + * + * @param outCapacity Must be at least `lhsSize + rhsSize` bytes, but may + * be larger to indicate that over-writes are okay. + * @param bitmapCapacity Must be at least `(lhsSize + rhsSize + 7) / 8` + * bytes, but may be larger to indicate that + * over-reads are okay. + * @param rhs Must be at least `rhsSize + SLOP` bytes, + * where the first @p rhsSize bytes are valid. + * + * @returns The number of ones in the bitmap. If this is not equal to + * @p rhsSize then the data was corrupt, and the output is unspecified. + */ + size_t (*mergeConstantVector)( + uint8_t* out, + size_t outCapacity, + const uint8_t* bitmap, + size_t bitmapCapacity, + uint8_t lhs, + size_t lhsSize, + const uint8_t* rhs, + size_t rhsSize); + + /** + * Merges @p lhs and @p rhs into @p out based on the bits in @p bitmap, + * where @p rhs is a constant. + * + * @param outCapacity Must be at least `lhsSize + rhsSize` bytes, but may + * be larger to indicate that over-writes are okay. + * @param bitmapCapacity Must be at least `(lhsSize + rhsSize + 7) / 8` + * bytes, but may be larger to indicate that + * over-reads are okay. + * @param lhs Must be at least `lhsSize + SLOP` bytes, + * where the first @p lhsSize bytes are valid. + * + * @returns The number of ones in the bitmap. If this is not equal to + * @p rhsSize then the data was corrupt, and the output is unspecified. + */ + size_t (*mergeVectorConstant)( + uint8_t* out, + size_t outCapacity, + const uint8_t* bitmap, + size_t bitmapCapacity, + const uint8_t* lhs, + size_t lhsSize, + uint8_t rhs, + size_t rhsSize); + + /** + * Reads @p outSize packed indices of @p depth bits each from @p bitmap and + * fills @p out with `symbols[idx]` for each packed index. + * + * @param out Buffer of size @p outCapacity + * @param outSize Number of output symbols + * @param outCapacity Must be at least @p outSize. It can be larger to + * indicate that over-writing is okay + * @param bitmapCapacity Must be at least `(outSize * depth + 7) / 8` bytes, + * but may be larger to indicate that over-reads are + * okay + * @param depth The number of bits each index in @p bitmap takes + * @param symbols Must be @p 2^depth bytes large + */ + void (*mergeFlatDepth)( + uint8_t* out, + size_t outSize, + size_t outCapacity, + const uint8_t* bitmap, + size_t bitmapCapacity, + size_t depth, + const uint8_t* symbols); +} ZL_PivCoHuffmanDecode; + +/** + * @returns The best kernel for the CPU. + */ +const ZL_PivCoHuffmanDecode* ZL_PivCoHuffmanDecode_select( + const ZL_cpuid_t* cpuid); + +extern const ZL_PivCoHuffmanDecode ZL_PivCoHuffmanDecode_generic; + +ZL_END_C_DECLS + +#endif diff --git a/src/openzl/codecs/pivco_huffman/arch/encode_pivco_arch.c b/src/openzl/codecs/pivco_huffman/arch/encode_pivco_arch.c new file mode 100644 index 000000000..46457bef0 --- /dev/null +++ b/src/openzl/codecs/pivco_huffman/arch/encode_pivco_arch.c @@ -0,0 +1,147 @@ +// Copyright (c) Meta Platforms, Inc. and affiliates. + +#include "openzl/codecs/pivco_huffman/arch/encode_pivco_arch.h" + +#include + +#include "openzl/codecs/common/bitstream/ff_bitstream.h" + +// Portable reference implementation of the PivCo-Huffman encode kernels. +// +// This file is written for clarity, not speed. It is used for correctness +// testing and as the fallback on hardware without a specialized kernel, which +// is not expected to run this code in practice. The architecture-specific +// kernels (x86, arm, avx512) are the fast paths and MUST produce byte-identical +// bitmaps to the code here. +// +// Every bitmap is packed little-endian, least-significant-bit first -- exactly +// the layout ZS_BitCStreamFF writes -- so we lean on that primitive instead of +// hand-rolling the bit math. To keep the loops trivially correct we flush the +// 64-bit window after every element. That is wasteful but easy to follow, which +// is the point of this file. + +const ZL_PivCoHuffmanEncode* ZL_PivCoHuffmanEncode_select( + const ZL_cpuid_t* cpuid) +{ + (void)cpuid; + return &ZL_PivCoHuffmanEncode_generic; +} + +static size_t bitmapBytes(size_t bits) +{ + return (bits + 7) / 8; +} + +// Partitions @p ranks into a left group (rank < rightRank) and a right group +// (rank >= rightRank). Emits one bit per rank into @p bitmap (0 = left, +// 1 = right) and, when the corresponding output is non-NULL, copies each rank +// into @p lhs / @p rhs in order. Returns the number of right (1) bits. +static size_t partitionGeneric( + uint8_t* bitmap, + uint8_t* lhs, + uint8_t* rhs, + const uint8_t* ranks, + size_t numRanks, + uint8_t rightRank) +{ + // The bitmap holds numRanks bits; the +SLOP matches the over-write budget + // the kernel contract guarantees, leaving room for the bitstream's 8-byte + // stores. + ZS_BitCStreamFF out = ZS_BitCStreamFF_init( + bitmap, bitmapBytes(numRanks) + ZL_PIVCO_HUFFMAN_SLOP); + + size_t ones = 0; + size_t zeros = 0; + for (size_t i = 0; i < numRanks; ++i) { + const uint8_t rank = ranks[i]; + const bool isRight = rank >= rightRank; + + ZS_BitCStreamFF_write(&out, isRight ? 1 : 0, 1); + // Flush every bit so the accumulator never holds more than 8 bits. + ZS_BitCStreamFF_flush(&out); + + if (isRight) { + if (rhs != NULL) { + rhs[ones] = rank; + } + ++ones; + } else { + if (lhs != NULL) { + lhs[zeros] = rank; + } + ++zeros; + } + } + (void)ZS_BitCStreamFF_finish(&out); + return ones; +} + +static size_t partitionLeft( + uint8_t* bitmap, + uint8_t* lhs, + const uint8_t* ranks, + size_t numRanks, + uint8_t rightRank) +{ + return partitionGeneric(bitmap, lhs, NULL, ranks, numRanks, rightRank); +} + +static size_t partitionRight( + uint8_t* bitmap, + uint8_t* rhs, + const uint8_t* ranks, + size_t numRanks, + uint8_t rightRank) +{ + return partitionGeneric(bitmap, NULL, rhs, ranks, numRanks, rightRank); +} + +static void partitionNone( + uint8_t* bitmap, + const uint8_t* ranks, + size_t numRanks, + uint8_t rightRank) +{ + (void)partitionGeneric(bitmap, NULL, NULL, ranks, numRanks, rightRank); +} + +// Packs one @p depth-bit index per rank into @p bitmap. The index is the rank's +// offset from @p rankBegin (its position within a flat leaf) and must fit in +// @p depth bits. +static void packFlatDepth( + uint8_t* bitmap, + size_t depth, + const uint8_t* ranks, + size_t numRanks, + uint8_t rankBegin) +{ + assert(depth >= 1); + assert(depth <= 8); + + ZS_BitCStreamFF out = ZS_BitCStreamFF_init( + bitmap, bitmapBytes(numRanks * depth) + ZL_PIVCO_HUFFMAN_SLOP); + + for (size_t i = 0; i < numRanks; ++i) { + const size_t index = (size_t)(ranks[i] - rankBegin); + assert(index < ((size_t)1 << depth)); + + ZS_BitCStreamFF_write(&out, index, depth); + ZS_BitCStreamFF_flush(&out); + } + (void)ZS_BitCStreamFF_finish(&out); +} + +static bool supported(const ZL_cpuid_t* cpuid) +{ + (void)cpuid; + return true; +} + +const ZL_PivCoHuffmanEncode ZL_PivCoHuffmanEncode_generic = { + .supported = supported, + .partitionFull = partitionGeneric, + .partitionLeft = partitionLeft, + .partitionRight = partitionRight, + .partitionNone = partitionNone, + .packFlatDepth = packFlatDepth, +}; diff --git a/src/openzl/codecs/pivco_huffman/arch/encode_pivco_arch.h b/src/openzl/codecs/pivco_huffman/arch/encode_pivco_arch.h new file mode 100644 index 000000000..728ebbfd5 --- /dev/null +++ b/src/openzl/codecs/pivco_huffman/arch/encode_pivco_arch.h @@ -0,0 +1,127 @@ +// Copyright (c) Meta Platforms, Inc. and affiliates. +#ifndef OPENZL_CODECS_PIVCO_HUFFMAN_ARCH_ENCODE_PIVCO_ARCH_H +#define OPENZL_CODECS_PIVCO_HUFFMAN_ARCH_ENCODE_PIVCO_ARCH_H + +#include +#include + +#include "openzl/codecs/pivco_huffman/arch/common_pivco_arch.h" +#include "openzl/shared/cpu.h" +#include "openzl/shared/portability.h" + +ZL_BEGIN_C_DECLS + +typedef struct { + /** + * @returns true iff the CPU supports this implementation. + * @note The rest of the function pointers may be NULL if not supported. + */ + bool (*supported)(const ZL_cpuid_t* cpuid); + + /** + * Partitions @p ranks into @p lhs and @p rhs based on @p rightRank. + * If the rank is below @p rightRank, then it goes to @p lhs, otherwise it + * goes to @p rhs. The partition decision is for each rank is written + * as a bit in @p bitmap, 0 for left and 1 for right. + * + * @param bitmap Must be `(numRanks + 7) / 8 + SLOP` bytes + * @param lhs Must be `numRanks + SLOP` elements large + * @param rhs Must be `numRanks + SLOP` elements large + * @param ranks Must be `numRanks + SLOP` elements large, + * where the first @p numRanks elements are valid. + * + * @returns the number of ones in the bitmap + * + * @note @p lhs or @p rhs may alias @p ranks + */ + size_t (*partitionFull)( + uint8_t* bitmap, + uint8_t* lhs, + uint8_t* rhs, + const uint8_t* ranks, + size_t numRanks, + uint8_t rightRank); + + /** + * Partitions @p ranks into @p lhs based on @p rightRank. If the rank is + * below @p rightRank, then it goes to @p lhs. The partition decision is for + * each rank is written as a bit in @p bitmap, 0 for left and 1 for + * right. + * + * @param bitmap Must be `(numRanks + 7) / 8 + SLOP` bytes + * @param lhs Must be `numRanks + SLOP` elements large + * @param ranks Must be `numRanks + SLOP` elements large, + * where the first @p numRanks elements are valid. + * + * @returns the number of ones in the bitmap + */ + size_t (*partitionLeft)( + uint8_t* bitmap, + uint8_t* lhs, + const uint8_t* ranks, + size_t numRanks, + uint8_t rightRank); + + /** + * Partitions @p ranks into @p rhs based on @p rightRank. If the rank is + * at least @p rightRank, then it goes to @p rhs. The partition decision is + * for each rank is written as a bit in @p bitmap, 0 for left and 1 for + * right. + * + * @param bitmap Must be `(numRanks + 7) / 8 + SLOP` bytes + * @param rhs Must be `numRanks + SLOP` elements large + * @param ranks Must be `numRanks + SLOP` elements large, + * where the first @p numRanks elements are valid. + * + * @returns the number of ones in the bitmap + */ + size_t (*partitionRight)( + uint8_t* bitmap, + uint8_t* rhs, + const uint8_t* ranks, + size_t numRanks, + uint8_t rightRank); + + /** + * Partitions @p ranks based on @p rightRank. If the rank is below + * @p rightRank, then it goes to the left, otherwise it goes to the right. + * That decision is written as a bit in @p bitmap, 0 for left and 1 for + * right. + * + * @param bitmap Must be `(numRanks + 7) / 8 + SLOP` bytes + * @param ranks Must be `numRanks + SLOP` elements large, + * where the first @p numRanks elements are valid. + */ + void (*partitionNone)( + uint8_t* bitmap, + const uint8_t* ranks, + size_t numRanks, + uint8_t rightRank); + + /** + * Subtracts @p rankBegin from each rank in @p ranks, the result of which + * must be at most @p depth bits, and stores it into @p bitmap. + * + * @param bitmap Must be `(numRanks * depth + 7) / 8 + SLOP` bytes + * @param ranks Must be `numRanks + SLOP` elements large, + * where the first @p numRanks elements are valid. + */ + void (*packFlatDepth)( + uint8_t* bitmap, + size_t depth, + const uint8_t* ranks, + size_t numRanks, + uint8_t rankBegin); +} ZL_PivCoHuffmanEncode; + +/** + * @returns The best kernel for the CPU. + */ +const ZL_PivCoHuffmanEncode* ZL_PivCoHuffmanEncode_select( + const ZL_cpuid_t* cpuid); + +extern const ZL_PivCoHuffmanEncode ZL_PivCoHuffmanEncode_generic; + +ZL_END_C_DECLS + +#endif diff --git a/src/openzl/shared/cpu.h b/src/openzl/shared/cpu.h index 9f84e58c4..fe4289c2f 100644 --- a/src/openzl/shared/cpu.h +++ b/src/openzl/shared/cpu.h @@ -201,6 +201,13 @@ B(avx512vl, 31) #define C(name, bit) X(name, f7c, bit) C(prefetchwt1, 0) C(avx512vbmi, 1) +C(avx512vbmi2, 6) +C(vaes, 9) +C(vpclmulqdq, 10) +C(avx512vnni, 11) +C(avx512bitalg, 12) +C(avx512vpopcntdq, 14) +C(rdpid, 22) #undef C #undef X diff --git a/tests/unittest/common/test_bitstream.cpp b/tests/unittest/common/test_bitstream.cpp index b6f4d9e20..93bf3c096 100644 --- a/tests/unittest/common/test_bitstream.cpp +++ b/tests/unittest/common/test_bitstream.cpp @@ -1,9 +1,13 @@ // Copyright (c) Meta Platforms, Inc. and affiliates. #include +#include +#include #include #include #include +#include +#include #include @@ -257,7 +261,13 @@ class RoundTripTest { ZS_BitDStreamFF_reload(&bits); } - return !ZL_isError(ZS_BitDStreamFF_finish(&bits)); + ZL_Report const report = ZS_BitDStreamFF_finish(&bits); + if (ZL_isError(report)) { + return false; + } + // finish() reports the exact number of bytes consumed, which must equal + // the encoded size regardless of any extra trailing capacity. + return ZL_validResult(report) == bitSize; } template @@ -440,6 +450,165 @@ INSTANTIATE_TEST_SUITE_P( BitstreamImpl::ZS, BitstreamImpl::ZS_BF)); +// Mask of the low @p nbBits bits. +size_t lowMask(size_t nbBits) +{ + if (nbBits == 0) + return 0; + if (nbBits >= sizeof(size_t) * 8) + return ~(size_t)0; + return ((size_t)1 << nbBits) - 1; +} + +// Round-trips `leading` bits, a byte-aligned raw region holding `region` bits, +// then `trailing` bits. The region is written with the reserve/commit aligned +// API and read back with ZS_BitDStreamFF_popAlignedBits, mirroring how the +// PivCo-Huffman kernels embed raw bitmaps in the stream. Verifies every bit +// and the raw region survive. When `extra` is set the decoder gets slack +// capacity past the encoded size (the large-stream path); otherwise it sees the +// exact end, exercising the near-end and short-stream paths in popAlignedBits. +// +// `reload` reloads the window after every field read. The kernels do not do +// this (popAlignedBits refills the window itself), but it must still be a valid +// thing to do: popAlignedBits has to recover the exact bit position from any +// reloaded stream state, including a short stream whose reload ran past the +// end. +void roundTripAlignedField( + const std::vector>& leading, + const std::vector& region, + const std::vector>& trailing, + bool extra, + bool reload) +{ + size_t const regionBytes = (region.size() + 7) / 8; + size_t totalBits = 0; + for (auto const& w : leading) + totalBits += w.second; + for (auto const& w : trailing) + totalBits += w.second; + std::vector buf(totalBits / 8 + regionBytes + 64, 0); + + ZS_BitCStreamFF writer = ZS_BitCStreamFF_init(buf.data(), buf.size()); + for (auto const& [value, nbBits] : leading) { + ZS_BitCStreamFF_write(&writer, value, nbBits); + ZS_BitCStreamFF_flush(&writer); + } + uint8_t* const slot = + ZS_BitCStreamFF_reserveAlignedBits(&writer, region.size()); + ASSERT_NE(slot, nullptr); + std::memset(slot, 0, regionBytes); + for (size_t i = 0; i < region.size(); ++i) { + if (region[i]) + slot[i / 8] |= (uint8_t)(1u << (i & 7)); + } + ZS_BitCStreamFF_commitReservedBits(&writer); + for (auto const& [value, nbBits] : trailing) { + ZS_BitCStreamFF_write(&writer, value, nbBits); + ZS_BitCStreamFF_flush(&writer); + } + ZL_Report const report = ZS_BitCStreamFF_finish(&writer); + ASSERT_ZS_VALID(report); + size_t const encodedSize = ZL_validResult(report); + + ZS_BitDStreamFF reader = + ZS_BitDStreamFF_init(buf.data(), extra ? buf.size() : encodedSize); + for (auto const& [value, nbBits] : leading) { + EXPECT_EQ( + ZS_BitDStreamFF_read(&reader, nbBits), value & lowMask(nbBits)); + if (reload) + ZS_BitDStreamFF_reload(&reader); + } + uint8_t const* const popped = + ZS_BitDStreamFF_popAlignedBits(&reader, region.size()); + ASSERT_NE(popped, nullptr); + for (size_t i = 0; i < region.size(); ++i) { + EXPECT_EQ((popped[i / 8] >> (i & 7)) & 1u, region[i] ? 1u : 0u) + << "region bit " << i; + } + for (auto const& [value, nbBits] : trailing) { + EXPECT_EQ( + ZS_BitDStreamFF_read(&reader, nbBits), value & lowMask(nbBits)); + if (reload) + ZS_BitDStreamFF_reload(&reader); + } + ASSERT_ZS_VALID(ZS_BitDStreamFF_finish(&reader)); +} + +std::vector makeRegionBits(size_t nbBits) +{ + std::vector bits(nbBits); + for (size_t i = 0; i < nbBits; ++i) + bits[i] = (((i * 7) + 3) % 5) < 2; + return bits; +} + +TEST(BitstreamAlignedBitsTest, ReserveCommitPopRoundTrip) +{ + using BitWrites = std::vector>; + BitWrites const someBits = { + { 0x1, 1 }, { 0x2, 3 }, { 0xABCD, 13 }, { 0x7F, 7 } + }; + + struct Scenario { + const char* name; + BitWrites leading; + size_t regionBits; + BitWrites trailing; + }; + std::vector const scenarios = { + // Region alone, with nothing around it. + { "empty_region", {}, 0, {} }, + { "byte_region_only", {}, 8, {} }, + { "aligned_region_only", {}, 64, {} }, + // Region followed by bits (window stays mid-stream). + { "region_then_bits", {}, 24, someBits }, + // Bits then a region that ends the stream (window lands at the end). + { "bits_then_region", someBits, 40, {} }, + // Region surrounded on both sides. + { "bits_region_bits", someBits, 32, someBits }, + // Unaligned leading bits force the region past padding. + { "unaligned_leading", { { 0x5, 3 } }, 16, { { 0x9, 5 } } }, + // Region whose bit count is not a multiple of 8 (partial trailing + // byte). + { "partial_region", someBits, 19, someBits }, + { "partial_region_at_end", someBits, 13, {} }, + // Tiny total so the exact-capacity decode hits the short-stream path. + { "tiny_small_stream", { { 0x1, 2 } }, 8, {} }, + // Large region exercises the steady-state window refill. + { "large_region", someBits, 1000, someBits }, + }; + + for (auto const& s : scenarios) { + for (bool extra : { true, false }) { + for (bool reload : { false, true }) { + SCOPED_TRACE( + std::string(s.name) + (extra ? " extra" : " exact") + + (reload ? " reload" : " noreload")); + roundTripAlignedField( + s.leading, + makeRegionBits(s.regionBits), + s.trailing, + extra, + reload); + } + } + } +} + +TEST(BitstreamAlignedBitsTest, ReserveReturnsNullWhenBufferTooSmall) +{ + std::vector buf(4, 0); + ZS_BitCStreamFF writer = ZS_BitCStreamFF_init(buf.data(), buf.size()); + EXPECT_EQ(ZS_BitCStreamFF_reserveAlignedBits(&writer, 8 * 64), nullptr); +} + +TEST(BitstreamAlignedBitsTest, PopReturnsNullPastEndOfStream) +{ + std::vector buf(8, 0xAB); + ZS_BitDStreamFF reader = ZS_BitDStreamFF_init(buf.data(), buf.size()); + EXPECT_EQ(ZS_BitDStreamFF_popAlignedBits(&reader, 8 * 64), nullptr); +} + struct BenchmarkResult { int64_t encodeNs{ 0 }; int64_t decodeNs{ 0 }; diff --git a/tests/unittest/transforms/test_pivco_huffman.cpp b/tests/unittest/transforms/test_pivco_huffman.cpp new file mode 100644 index 000000000..196196e93 --- /dev/null +++ b/tests/unittest/transforms/test_pivco_huffman.cpp @@ -0,0 +1,537 @@ +// Copyright (c) Meta Platforms, Inc. and affiliates. + +#include +#include +#include + +#include + +#include "openzl/codecs/pivco_huffman/arch/common_pivco_arch.h" +#include "openzl/codecs/pivco_huffman/arch/decode_pivco_arch.h" +#include "openzl/codecs/pivco_huffman/arch/encode_pivco_arch.h" + +namespace { + +struct EncodeArch { + const char* name; + const ZL_PivCoHuffmanEncode* kernels; +}; + +struct DecodeArch { + const char* name; + const ZL_PivCoHuffmanDecode* kernels; +}; + +std::vector supportedEncodeArchs() +{ + ZL_cpuid_t const cpuid = ZL_cpuid(); + std::vector out; + std::vector const archs = { + { "generic", &ZL_PivCoHuffmanEncode_generic }, + }; + for (auto const& arch : archs) { + if (arch.kernels->supported(&cpuid)) { + out.push_back(arch); + } + } + return out; +} + +std::vector supportedDecodeArchs() +{ + ZL_cpuid_t const cpuid = ZL_cpuid(); + std::vector out; + std::vector const archs = { + { "generic", &ZL_PivCoHuffmanDecode_generic }, + }; + for (auto const& arch : archs) { + if (arch.kernels->supported(&cpuid)) { + out.push_back(arch); + } + } + return out; +} + +template +std::vector firstN(const std::vector& in, size_t n) +{ + return std::vector(in.begin(), in.begin() + n); +} + +size_t bitmapBytes(size_t bits) +{ + return (bits + 7) / 8; +} + +void setBitmapBit(std::vector& bitmap, size_t bit) +{ + bitmap[bit / 8] |= (uint8_t)(1u << (bit & 7)); +} + +enum class TopBitPattern { + AllZero, + AllOne, + Mixed, +}; + +uint8_t makeRank(size_t i, size_t size, TopBitPattern pattern) +{ + switch (pattern) { + case TopBitPattern::AllZero: + return 0x10; + case TopBitPattern::AllOne: + return 0x90; + case TopBitPattern::Mixed: + return (((i * 37) + size) % 7) < 3 ? 0x90 : 0x10; + } + return 0; +} + +std::vector makeRanks(size_t size, TopBitPattern pattern) +{ + std::vector ranks(size + ZL_PIVCO_HUFFMAN_SLOP, 0xA5); + for (size_t i = 0; i < size; ++i) { + ranks[i] = makeRank(i, size, pattern); + } + return ranks; +} + +std::vector +makeFlatDepthRanks(size_t size, size_t depth, uint8_t rankBegin) +{ + std::vector ranks(size + ZL_PIVCO_HUFFMAN_SLOP, 0xA5); + uint8_t const symbolMask = (uint8_t)((1u << depth) - 1u); + for (size_t i = 0; i < size; ++i) { + ranks[i] = (uint8_t)(rankBegin + (uint8_t)(i & symbolMask)); + } + return ranks; +} + +struct PartitionReference { + std::vector bitmap; + std::vector lhs; + std::vector rhs; +}; + +PartitionReference referencePartition( + const std::vector& ranks, + size_t size, + uint8_t rightRank) +{ + PartitionReference out; + out.bitmap.assign(bitmapBytes(size), 0); + for (size_t i = 0; i < size; ++i) { + uint8_t const rank = ranks[i]; + if (rank >= rightRank) { + setBitmapBit(out.bitmap, i); + out.rhs.push_back(rank); + } else { + out.lhs.push_back(rank); + } + } + return out; +} + +std::vector referencePackFlatDepth( + const std::vector& ranks, + size_t size, + size_t depth, + uint8_t rankBegin) +{ + std::vector bitmap(bitmapBytes(size * depth), 0); + for (size_t i = 0; i < size; ++i) { + uint8_t const idx = (uint8_t)(ranks[i] - rankBegin); + for (size_t bit = 0; bit < depth; ++bit) { + if ((idx & (1u << bit)) != 0) { + setBitmapBit(bitmap, i * depth + bit); + } + } + } + return bitmap; +} + +std::vector referenceMergeFlatDepth( + const std::vector& bitmap, + size_t outSize, + size_t depth, + const std::vector& symbols) +{ + std::vector out(outSize); + uint8_t const mask = (uint8_t)((1u << depth) - 1); + for (size_t i = 0; i < outSize; ++i) { + size_t const bitOffset = i * depth; + size_t const byte = bitOffset / 8; + size_t const shift = bitOffset & 7; + uint16_t bits = bitmap[byte]; + if (byte + 1 < bitmap.size()) { + bits |= (uint16_t)(bitmap[byte + 1] << 8); + } + out[i] = symbols[(bits >> shift) & mask]; + } + return out; +} + +std::vector boundarySizes() +{ + return { + 0, 1, 2, 3, 4, 5, 7, 8, 9, 15, 16, 17, 31, 32, + 33, 63, 64, 65, 127, 128, 129, 255, 256, 257, 511, 512, 513, + }; +} + +std::vector makeMergeBits(size_t size, TopBitPattern pattern) +{ + std::vector bits(size); + for (size_t i = 0; i < size; ++i) { + switch (pattern) { + case TopBitPattern::AllZero: + bits[i] = false; + break; + case TopBitPattern::AllOne: + bits[i] = true; + break; + case TopBitPattern::Mixed: + bits[i] = (((i * 11) + size) % 5) < 2; + break; + } + } + return bits; +} + +std::vector bitmapFromBits( + const std::vector& bits, + size_t extraCapacity) +{ + size_t const exactBytes = bitmapBytes(bits.size()); + std::vector bitmap(exactBytes + extraCapacity, 0); + for (size_t i = 0; i < bits.size(); ++i) { + if (bits[i]) { + setBitmapBit(bitmap, i); + } + } + for (size_t i = exactBytes; i < bitmap.size(); ++i) { + bitmap[i] = 0xA5; + } + return bitmap; +} + +struct MergeInputs { + std::vector expected; + std::vector lhs; + std::vector rhs; + size_t ones = 0; +}; + +MergeInputs makeMergeInputs(const std::vector& bits) +{ + MergeInputs out; + out.expected.reserve(bits.size()); + for (size_t i = 0; i < bits.size(); ++i) { + if (bits[i]) { + uint8_t const value = + (uint8_t)(0x80u + ((out.rhs.size() * 13) & 0x7F)); + out.rhs.push_back(value); + out.expected.push_back(value); + ++out.ones; + } else { + uint8_t const value = + (uint8_t)(0x10u + ((out.lhs.size() * 17) & 0x6F)); + out.lhs.push_back(value); + out.expected.push_back(value); + } + } + out.lhs.resize(out.lhs.size() + ZL_PIVCO_HUFFMAN_SLOP, 0xEE); + out.rhs.resize(out.rhs.size() + ZL_PIVCO_HUFFMAN_SLOP, 0xDD); + return out; +} + +void expectBitmapPrefix( + const std::vector& actual, + const std::vector& expected) +{ + ASSERT_LE(expected.size(), actual.size()); + EXPECT_EQ(firstN(actual, expected.size()), expected); +} + +void expectBytesAfterPrefix( + const std::vector& actual, + size_t prefixSize, + uint8_t sentinel) +{ + ASSERT_LE(prefixSize, actual.size()); + std::vector const actualTail( + actual.begin() + prefixSize, actual.end()); + std::vector const expectedTail(actualTail.size(), sentinel); + EXPECT_EQ(actualTail, expectedTail); +} + +} // namespace + +TEST(PivCoHuffmanArchTest, PartitionKernelsMatchReference) +{ + for (auto const& arch : supportedEncodeArchs()) { + ASSERT_NE(arch.kernels->partitionFull, nullptr) << arch.name; + ASSERT_NE(arch.kernels->partitionLeft, nullptr) << arch.name; + ASSERT_NE(arch.kernels->partitionRight, nullptr) << arch.name; + ASSERT_NE(arch.kernels->partitionNone, nullptr) << arch.name; + + for (TopBitPattern pattern : { TopBitPattern::AllZero, + TopBitPattern::AllOne, + TopBitPattern::Mixed }) { + for (size_t size : boundarySizes()) { + SCOPED_TRACE( + std::string(arch.name) + + " size=" + std::to_string(size)); + uint8_t const rightRank = 0x80; + auto const ranks = makeRanks(size, pattern); + auto const ref = referencePartition(ranks, size, rightRank); + size_t const bytes = bitmapBytes(size); + + std::vector bitmap( + bytes + ZL_PIVCO_HUFFMAN_SLOP, 0xCC); + std::vector lhs(size + ZL_PIVCO_HUFFMAN_SLOP, 0xEE); + std::vector rhs(size + ZL_PIVCO_HUFFMAN_SLOP, 0xDD); + size_t const ones = arch.kernels->partitionFull( + bitmap.data(), + lhs.data(), + rhs.data(), + ranks.data(), + size, + rightRank); + EXPECT_EQ(ones, ref.rhs.size()); + expectBitmapPrefix(bitmap, ref.bitmap); + EXPECT_EQ(firstN(lhs, ref.lhs.size()), ref.lhs); + EXPECT_EQ(firstN(rhs, ref.rhs.size()), ref.rhs); + + std::fill(bitmap.begin(), bitmap.end(), 0xCC); + std::fill(lhs.begin(), lhs.end(), 0xEE); + EXPECT_EQ( + arch.kernels->partitionLeft( + bitmap.data(), + lhs.data(), + ranks.data(), + size, + rightRank), + ref.rhs.size()); + expectBitmapPrefix(bitmap, ref.bitmap); + EXPECT_EQ(firstN(lhs, ref.lhs.size()), ref.lhs); + + std::fill(bitmap.begin(), bitmap.end(), 0xCC); + std::fill(rhs.begin(), rhs.end(), 0xDD); + EXPECT_EQ( + arch.kernels->partitionRight( + bitmap.data(), + rhs.data(), + ranks.data(), + size, + rightRank), + ref.rhs.size()); + expectBitmapPrefix(bitmap, ref.bitmap); + EXPECT_EQ(firstN(rhs, ref.rhs.size()), ref.rhs); + + std::fill(bitmap.begin(), bitmap.end(), 0xCC); + arch.kernels->partitionNone( + bitmap.data(), ranks.data(), size, rightRank); + expectBitmapPrefix(bitmap, ref.bitmap); + } + } + } +} + +TEST(PivCoHuffmanArchTest, PartitionFullSupportsDocumentedInputAliasing) +{ + for (auto const& arch : supportedEncodeArchs()) { + for (TopBitPattern pattern : { TopBitPattern::AllZero, + TopBitPattern::AllOne, + TopBitPattern::Mixed }) { + for (size_t size : boundarySizes()) { + SCOPED_TRACE( + std::string(arch.name) + + " size=" + std::to_string(size)); + uint8_t const rightRank = 0x80; + auto const ranks = makeRanks(size, pattern); + auto const ref = referencePartition(ranks, size, rightRank); + size_t const bytes = bitmapBytes(size); + + std::vector bitmap( + bytes + ZL_PIVCO_HUFFMAN_SLOP, 0xCC); + std::vector lhsAlias = ranks; + std::vector rhs(size + ZL_PIVCO_HUFFMAN_SLOP, 0xDD); + EXPECT_EQ( + arch.kernels->partitionFull( + bitmap.data(), + lhsAlias.data(), + rhs.data(), + lhsAlias.data(), + size, + rightRank), + ref.rhs.size()); + expectBitmapPrefix(bitmap, ref.bitmap); + EXPECT_EQ(firstN(lhsAlias, ref.lhs.size()), ref.lhs); + EXPECT_EQ(firstN(rhs, ref.rhs.size()), ref.rhs); + + std::fill(bitmap.begin(), bitmap.end(), 0xCC); + std::vector lhs(size + ZL_PIVCO_HUFFMAN_SLOP, 0xEE); + std::vector rhsAlias = ranks; + EXPECT_EQ( + arch.kernels->partitionFull( + bitmap.data(), + lhs.data(), + rhsAlias.data(), + rhsAlias.data(), + size, + rightRank), + ref.rhs.size()); + expectBitmapPrefix(bitmap, ref.bitmap); + EXPECT_EQ(firstN(lhs, ref.lhs.size()), ref.lhs); + EXPECT_EQ(firstN(rhsAlias, ref.rhs.size()), ref.rhs); + } + } + } +} + +TEST(PivCoHuffmanArchTest, FlatDepthPackAndMergeMatchReference) +{ + auto const sizes = boundarySizes(); + for (size_t depth = 1; depth <= 8; ++depth) { + std::vector symbols((size_t)1 << depth); + for (size_t i = 0; i < symbols.size(); ++i) { + symbols[i] = (uint8_t)((i * 37 + depth * 11) & 0xFF); + } + + for (size_t size : sizes) { + uint8_t const rankBegin = depth == 8 ? 0 : 17; + auto const ranks = makeFlatDepthRanks(size, depth, rankBegin); + auto const packed = + referencePackFlatDepth(ranks, size, depth, rankBegin); + auto const merged = + referenceMergeFlatDepth(packed, size, depth, symbols); + + for (auto const& arch : supportedEncodeArchs()) { + ASSERT_NE(arch.kernels->packFlatDepth, nullptr) << arch.name; + SCOPED_TRACE( + std::string("pack ") + arch.name + + " depth=" + std::to_string(depth) + + " size=" + std::to_string(size)); + size_t const guardOffset = + packed.size() + ZL_PIVCO_HUFFMAN_SLOP; + std::vector bitmap( + guardOffset + ZL_PIVCO_HUFFMAN_SLOP, 0xCC); + arch.kernels->packFlatDepth( + bitmap.data(), depth, ranks.data(), size, rankBegin); + expectBitmapPrefix(bitmap, packed); + expectBytesAfterPrefix(bitmap, guardOffset, 0xCC); + } + + for (auto const& arch : supportedDecodeArchs()) { + ASSERT_NE(arch.kernels->mergeFlatDepth, nullptr) << arch.name; + SCOPED_TRACE( + std::string("merge ") + arch.name + + " depth=" + std::to_string(depth) + + " size=" + std::to_string(size)); + std::vector out(size + ZL_PIVCO_HUFFMAN_SLOP, 0xCC); + arch.kernels->mergeFlatDepth( + out.data(), + size, + out.size(), + packed.data(), + packed.size(), + depth, + symbols.data()); + EXPECT_EQ(firstN(out, size), merged); + } + } + } +} + +TEST(PivCoHuffmanArchTest, MergeKernelsMatchReference) +{ + for (auto const& arch : supportedDecodeArchs()) { + ASSERT_NE(arch.kernels->mergeVectorVector, nullptr) << arch.name; + ASSERT_NE(arch.kernels->mergeConstantVector, nullptr) << arch.name; + ASSERT_NE(arch.kernels->mergeVectorConstant, nullptr) << arch.name; + + for (TopBitPattern pattern : { TopBitPattern::AllZero, + TopBitPattern::AllOne, + TopBitPattern::Mixed }) { + for (size_t size : boundarySizes()) { + auto const bits = makeMergeBits(size, pattern); + auto const inputs = makeMergeInputs(bits); + + // The expected outputs depend only on `bits`/`inputs`, so build + // them once here rather than inside the capacity loops below. + uint8_t const lhsConstant = 0x3C; + uint8_t const rhsConstant = 0xC3; + std::vector expectedConstantVector(size); + std::vector expectedVectorConstant(size); + for (size_t i = 0, lhsPos = 0, rhsPos = 0; i < size; ++i) { + expectedConstantVector[i] = + bits[i] ? inputs.rhs[rhsPos++] : lhsConstant; + expectedVectorConstant[i] = + bits[i] ? rhsConstant : inputs.lhs[lhsPos++]; + } + + for (size_t outExtra : + { (size_t)0, (size_t)ZL_PIVCO_HUFFMAN_SLOP }) { + for (size_t bitmapExtra : + { (size_t)0, (size_t)ZL_PIVCO_HUFFMAN_SLOP }) { + SCOPED_TRACE( + std::string(arch.name) + + " size=" + std::to_string(size) + " outExtra=" + + std::to_string(outExtra) + " bitmapExtra=" + + std::to_string(bitmapExtra)); + auto const bitmap = bitmapFromBits(bits, bitmapExtra); + + std::vector out(size + outExtra, 0xCC); + EXPECT_EQ( + arch.kernels->mergeVectorVector( + out.data(), + out.size(), + bitmap.data(), + bitmap.size(), + inputs.lhs.data(), + inputs.lhs.size() + - ZL_PIVCO_HUFFMAN_SLOP, + inputs.rhs.data(), + inputs.rhs.size() + - ZL_PIVCO_HUFFMAN_SLOP), + inputs.ones); + EXPECT_EQ(firstN(out, size), inputs.expected); + + std::fill(out.begin(), out.end(), 0xCC); + EXPECT_EQ( + arch.kernels->mergeConstantVector( + out.data(), + out.size(), + bitmap.data(), + bitmap.size(), + lhsConstant, + inputs.lhs.size() + - ZL_PIVCO_HUFFMAN_SLOP, + inputs.rhs.data(), + inputs.rhs.size() + - ZL_PIVCO_HUFFMAN_SLOP), + inputs.ones); + EXPECT_EQ(firstN(out, size), expectedConstantVector); + + std::fill(out.begin(), out.end(), 0xCC); + EXPECT_EQ( + arch.kernels->mergeVectorConstant( + out.data(), + out.size(), + bitmap.data(), + bitmap.size(), + inputs.lhs.data(), + inputs.lhs.size() + - ZL_PIVCO_HUFFMAN_SLOP, + rhsConstant, + inputs.rhs.size() + - ZL_PIVCO_HUFFMAN_SLOP), + inputs.ones); + EXPECT_EQ(firstN(out, size), expectedVectorConstant); + } + } + } + } + } +}