From c5078bbe1162ac12fbf6b53d2370d7bb40394dd8 Mon Sep 17 00:00:00 2001 From: Nick Terrell Date: Thu, 23 Jul 2026 14:54:28 -0700 Subject: [PATCH 1/2] Remove partitionLeft / mergeVectorConstant (#900) Summary: Canonical Huffman trees always emit shallower leaves left of deeper leaves. Therefore these methods are not needed, and are dead code. Differential Revision: D113463813 --- benchmark/micro/micro_pivco_huffman.cpp | 105 ------------------ .../pivco_huffman/arch/decode_pivco_arch.c | 20 ---- .../pivco_huffman/arch/decode_pivco_arch.h | 25 ----- .../pivco_huffman/arch/encode_pivco_arch.c | 11 -- .../pivco_huffman/arch/encode_pivco_arch.h | 20 ---- .../pivco_huffman/decode_pivco_kernel.c | 12 +- .../pivco_huffman/encode_pivco_kernel.c | 5 +- .../transforms/test_pivco_huffman.cpp | 37 +----- 8 files changed, 5 insertions(+), 230 deletions(-) diff --git a/benchmark/micro/micro_pivco_huffman.cpp b/benchmark/micro/micro_pivco_huffman.cpp index 36a27049..30e6ae22 100644 --- a/benchmark/micro/micro_pivco_huffman.cpp +++ b/benchmark/micro/micro_pivco_huffman.cpp @@ -216,30 +216,6 @@ void verifyPartitionFull(const EncodeArch& arch, const PartitionData& data) requireEqualPrefix(rhs, expectedRhs, expectedOnes); } -void verifyPartitionLeft(const EncodeArch& arch, const PartitionData& data) -{ - std::vector expectedBitmap(data.bitmap.size()); - std::vector expectedLhs(data.lhs.size()); - size_t const expectedOnes = ZL_PivCoHuffmanEncode_generic.partitionLeft( - expectedBitmap.data(), - expectedLhs.data(), - data.ranks.data(), - data.size, - data.rightRank); - - std::vector bitmap(data.bitmap.size()); - std::vector lhs(data.lhs.size()); - size_t const ones = arch.kernels->partitionLeft( - bitmap.data(), - lhs.data(), - data.ranks.data(), - data.size, - data.rightRank); - ZL_REQUIRE_EQ(ones, expectedOnes); - requireEqualPrefix(bitmap, expectedBitmap, bitmapBytes(data.size)); - requireEqualPrefix(lhs, expectedLhs, data.size - expectedOnes); -} - void verifyPartitionRight(const EncodeArch& arch, const PartitionData& data) { std::vector expectedBitmap(data.bitmap.size()); @@ -356,35 +332,6 @@ void verifyMergeConstantVector(const DecodeArch& arch, const MergeData& data) requireEqualPrefix(out, expectedOut, data.totalSize); } -void verifyMergeVectorConstant(const DecodeArch& arch, const MergeData& data) -{ - std::vector expectedOut(data.out.size()); - uint8_t const rhs = 0xC3; - size_t const expectedOnes = - ZL_PivCoHuffmanDecode_generic.mergeVectorConstant( - expectedOut.data(), - expectedOut.size(), - data.bitmap.data(), - data.bitmap.size(), - data.lhs.data(), - data.lhsSize, - rhs, - data.rhsSize); - - std::vector out(data.out.size()); - size_t const ones = arch.kernels->mergeVectorConstant( - out.data(), - out.size(), - data.bitmap.data(), - data.bitmap.size(), - data.lhs.data(), - data.lhsSize, - rhs, - data.rhsSize); - ZL_REQUIRE_EQ(ones, expectedOnes); - requireEqualPrefix(out, expectedOut, data.totalSize); -} - void verifyMergeFlatDepth(const DecodeArch& arch, const FlatData& data) { std::vector expectedOut(data.out.size()); @@ -437,28 +384,6 @@ void registerPartitionFullBenchmark(EncodeArch arch) }); } -void registerPartitionLeftBenchmark(EncodeArch arch) -{ - auto data = std::make_shared(64 * 1024); - verifyPartitionLeft(arch, *data); - RegisterBenchmark( - benchName(arch.name, "partitionLeft"), - [arch, data](benchmark::State& state) { - size_t ones = 0; - for (auto _ : state) { - ones += arch.kernels->partitionLeft( - data->bitmap.data(), - data->lhs.data(), - data->ranks.data(), - data->size, - data->rightRank); - } - benchmark::DoNotOptimize(ones); - state.SetBytesProcessed( - (int64_t)data->size * (int64_t)state.iterations()); - }); -} - void registerPartitionRightBenchmark(EncodeArch arch) { auto data = std::make_shared(64 * 1024); @@ -574,32 +499,6 @@ void registerMergeConstantVectorBenchmark(DecodeArch arch) }); } -void registerMergeVectorConstantBenchmark(DecodeArch arch) -{ - auto data = std::make_shared(64 * 1024); - verifyMergeVectorConstant(arch, *data); - RegisterBenchmark( - benchName(arch.name, "mergeVectorConstant"), - [arch, data](benchmark::State& state) { - size_t ones = 0; - uint8_t const rhs = 0xC3; - for (auto _ : state) { - ones += arch.kernels->mergeVectorConstant( - data->out.data(), - data->out.size(), - data->bitmap.data(), - data->bitmap.size(), - data->lhs.data(), - data->lhsSize, - rhs, - data->rhsSize); - } - benchmark::DoNotOptimize(ones); - state.SetBytesProcessed( - (int64_t)data->totalSize * (int64_t)state.iterations()); - }); -} - void registerMergeFlatDepthBenchmark(DecodeArch arch, size_t depth) { auto data = std::make_shared(64 * 1024, depth); @@ -628,12 +527,10 @@ void registerMergeFlatDepthBenchmark(DecodeArch arch, size_t depth) void registerEncodeBenchmarks(EncodeArch arch) { ZL_REQUIRE_NN(arch.kernels->partitionFull); - ZL_REQUIRE_NN(arch.kernels->partitionLeft); ZL_REQUIRE_NN(arch.kernels->partitionRight); ZL_REQUIRE_NN(arch.kernels->partitionNone); ZL_REQUIRE_NN(arch.kernels->packFlatDepth); registerPartitionFullBenchmark(arch); - registerPartitionLeftBenchmark(arch); registerPartitionRightBenchmark(arch); registerPartitionNoneBenchmark(arch); for (size_t depth = 1; depth <= 8; ++depth) { @@ -645,11 +542,9 @@ void registerDecodeBenchmarks(DecodeArch arch) { ZL_REQUIRE_NN(arch.kernels->mergeVectorVector); ZL_REQUIRE_NN(arch.kernels->mergeConstantVector); - ZL_REQUIRE_NN(arch.kernels->mergeVectorConstant); ZL_REQUIRE_NN(arch.kernels->mergeFlatDepth); registerMergeVectorVectorBenchmark(arch); registerMergeConstantVectorBenchmark(arch); - registerMergeVectorConstantBenchmark(arch); for (size_t depth = 1; depth <= 8; ++depth) { registerMergeFlatDepthBenchmark(arch, depth); } diff --git a/src/openzl/codecs/pivco_huffman/arch/decode_pivco_arch.c b/src/openzl/codecs/pivco_huffman/arch/decode_pivco_arch.c index f9e98455..d3ddc830 100644 --- a/src/openzl/codecs/pivco_huffman/arch/decode_pivco_arch.c +++ b/src/openzl/codecs/pivco_huffman/arch/decode_pivco_arch.c @@ -143,25 +143,6 @@ static size_t mergeConstantVector( 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( @@ -199,6 +180,5 @@ 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 index e90d6bb9..6e74ad72 100644 --- a/src/openzl/codecs/pivco_huffman/arch/decode_pivco_arch.h +++ b/src/openzl/codecs/pivco_huffman/arch/decode_pivco_arch.h @@ -69,31 +69,6 @@ typedef struct { 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. diff --git a/src/openzl/codecs/pivco_huffman/arch/encode_pivco_arch.c b/src/openzl/codecs/pivco_huffman/arch/encode_pivco_arch.c index 46457bef..a4f7dd46 100644 --- a/src/openzl/codecs/pivco_huffman/arch/encode_pivco_arch.c +++ b/src/openzl/codecs/pivco_huffman/arch/encode_pivco_arch.c @@ -76,16 +76,6 @@ static size_t partitionGeneric( 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, @@ -140,7 +130,6 @@ static bool supported(const ZL_cpuid_t* cpuid) 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 index 728ebbfd..511682d4 100644 --- a/src/openzl/codecs/pivco_huffman/arch/encode_pivco_arch.h +++ b/src/openzl/codecs/pivco_huffman/arch/encode_pivco_arch.h @@ -42,26 +42,6 @@ typedef struct { 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 diff --git a/src/openzl/codecs/pivco_huffman/decode_pivco_kernel.c b/src/openzl/codecs/pivco_huffman/decode_pivco_kernel.c index 7f057378..3854e4ce 100644 --- a/src/openzl/codecs/pivco_huffman/decode_pivco_kernel.c +++ b/src/openzl/codecs/pivco_huffman/decode_pivco_kernel.c @@ -93,17 +93,9 @@ static bool mergeDecodeResults( lhsSize, rhs->output, rhsSize); - } else if (rhs->isConstant) { - ones = kernels->mergeVectorConstant( - output, - outputCapacity, - bitmap, - bitmapBytes, - lhs->output, - lhsSize, - rhs->symbol, - rhsSize); } else { + // rhs cannot be constant in canonical Huffman codes. + assert(!rhs->isConstant); ones = kernels->mergeVectorVector( output, outputCapacity, diff --git a/src/openzl/codecs/pivco_huffman/encode_pivco_kernel.c b/src/openzl/codecs/pivco_huffman/encode_pivco_kernel.c index 59add326..2c561ee2 100644 --- a/src/openzl/codecs/pivco_huffman/encode_pivco_kernel.c +++ b/src/openzl/codecs/pivco_huffman/encode_pivco_kernel.c @@ -174,10 +174,9 @@ static bool encodeNode( } else if (lhsIsConstant) { numOnes = kernels->partitionRight( bitmap, rhsRanks, nodeRanks, numRanks, (uint8_t)splitRank); - } else if (rhsIsConstant) { - numOnes = kernels->partitionLeft( - bitmap, lhsRanks, nodeRanks, numRanks, (uint8_t)splitRank); } else { + // rhs cannot be constant in canonical Huffman codes. + assert(!rhsIsConstant); numOnes = kernels->partitionFull( bitmap, lhsRanks, diff --git a/tests/unittest/transforms/test_pivco_huffman.cpp b/tests/unittest/transforms/test_pivco_huffman.cpp index e0f0525b..7a037d17 100644 --- a/tests/unittest/transforms/test_pivco_huffman.cpp +++ b/tests/unittest/transforms/test_pivco_huffman.cpp @@ -431,7 +431,6 @@ 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; @@ -463,19 +462,6 @@ TEST(PivCoHuffmanArchTest, PartitionKernelsMatchReference) 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( @@ -609,7 +595,6 @@ 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, @@ -621,14 +606,10 @@ TEST(PivCoHuffmanArchTest, MergeKernelsMatchReference) // 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) { + for (size_t i = 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 : @@ -673,22 +654,6 @@ TEST(PivCoHuffmanArchTest, MergeKernelsMatchReference) - 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); } } } From 5c0d165b39c68851c12e518c8c7b4d13e557ac1a Mon Sep 17 00:00:00 2001 From: Nick Terrell Date: Thu, 23 Jul 2026 14:54:28 -0700 Subject: [PATCH 2/2] Add the SIMD table generator script (#858) Summary: Adds `scripts/gen_pivco_huffman_tables.py`, which renders the generated lookup-table headers consumed by the architecture-specific kernels (x86, neon, avx512). It is checked in ahead of the per-architecture kernels that include its output. Reviewed By: kevinjzhang Differential Revision: D109893569 --- scripts/gen_pivco_huffman_tables.py | 356 ++++++++++++++++++++++++++++ 1 file changed, 356 insertions(+) create mode 100644 scripts/gen_pivco_huffman_tables.py diff --git a/scripts/gen_pivco_huffman_tables.py b/scripts/gen_pivco_huffman_tables.py new file mode 100644 index 00000000..ee28be48 --- /dev/null +++ b/scripts/gen_pivco_huffman_tables.py @@ -0,0 +1,356 @@ +#!/usr/bin/env python3 +# Copyright (c) Meta Platforms, Inc. and affiliates. + +from __future__ import annotations + +from pathlib import Path + + +ROOT = Path(__file__).resolve().parents[1] +SCRIPT = "scripts/gen_pivco_huffman_tables.py" +GENERATED = "generated" + + +def header_guard(path: str) -> str: + return path.replace("src/", "").replace("/", "_").replace(".", "_").upper() + + +def hex8(value: int) -> str: + assert 0 <= value <= 0xFF + return f"0x{value:02x}" + + +def popcounts() -> list[int]: + return [mask.bit_count() for mask in range(256)] + + +def zeros_counts() -> list[int]: + return [8 - mask.bit_count() for mask in range(256)] + + +def compress_u16(mask: int, pad: int) -> list[int]: + out = [] + for lane in range(8): + if (mask & (1 << lane)) != 0: + out += [2 * lane, 2 * lane + 1] + out += [pad] * (16 - len(out)) + return out + + +def compress_u16_pair(pad: int) -> list[list[int]]: + return [ + compress_u16(mask, pad) + compress_u16((~mask) & 0xFF, pad) + for mask in range(256) + ] + + +def partition_lo16(mask: int) -> list[int]: + # Low-half gather control for the 16-wide partition: right ranks (set lanes) + # packed forward at the front, left ranks (unset lanes) packed reversed at the + # far back; the gap in between is filled by the high-half control. + out = [0x80] * 16 + pos = 0 + for lane in range(8): + if (mask >> lane) & 1: + out[pos] = lane + pos += 1 + pos = 15 + for lane in range(8): + if not ((mask >> lane) & 1): + out[pos] = lane + pos -= 1 + return out + + +def partition_hi_base16(mask: int) -> list[int]: + # High-half gather control before zero-padding: right ranks packed forward at + # the front, left ranks packed reversed ending at lane 7. Values are source + # byte indices (8 + lane), since the high half lives in bytes 8..15. + out = [0x80] * 16 + pos = 0 + for lane in range(8): + if (mask >> lane) & 1: + out[pos] = 8 + lane + pos += 1 + pos = 7 + for lane in range(8): + if not ((mask >> lane) & 1): + out[pos] = 8 + lane + pos -= 1 + return out + + +def partition_lo() -> list[list[int]]: + return [partition_lo16(mask) for mask in range(256)] + + +def partition_hi_padded() -> list[list[int]]: + # Pad each high-half control with 8 zero (0x80) bytes on either end so a + # misaligned load at offset (8 - popcount(loMask)) slides the high contribution + # into the correct output lanes. + pad = [0x80] * 8 + return [pad + partition_hi_base16(mask) + pad for mask in range(256)] + + +def merge_lo_ctrl16(mask: int) -> list[int]: + # Low-half merge control for the 16-wide merge. Bytes 0..7 hold the right + # shuffle for output lanes 0..7: set bits get the local right rank; unset + # bits get 0xFF ^ leftRank, so the high bit marks a left lane and XOR-ing the + # whole vector with 0xFF recovers the left shuffle. Bytes 8..15 broadcast + # popcount(mask) so it can be added into the high-half control. + out = [] + rhs = 0 + lhs = 0 + for bit in range(8): + if (mask >> bit) & 1: + out.append(rhs) + rhs += 1 + else: + out.append(0xFF ^ lhs) + lhs += 1 + return out + [mask.bit_count()] * 8 + + +def merge_hi_ctrl8(mask: int) -> list[int]: + # High-half merge control (8 bytes): set bits get the local right rank; unset + # bits get 0xFF - 8 - localLeftRank, so adding popcount(loMask) yields + # 0xFF ^ fullLeftRank for output lanes 8..15. + out = [] + rhs = 0 + lhs = 0 + for bit in range(8): + if (mask >> bit) & 1: + out.append(rhs) + rhs += 1 + else: + out.append(0xFF - 8 - lhs) + lhs += 1 + return out + + +def merge_lo_ctrl() -> list[list[int]]: + return [merge_lo_ctrl16(mask) for mask in range(256)] + + +def merge_hi_ctrl() -> list[list[int]]: + return [merge_hi_ctrl8(mask) for mask in range(256)] + + +def neon_expand8(mask: int, lhs_base: int, rhs_base: int) -> list[int]: + lhs = lhs_base + rhs = rhs_base + out = [] + for bit in range(8): + if (mask & (1 << bit)) != 0: + out.append(rhs) + rhs += 1 + else: + out.append(lhs) + lhs += 1 + return out + + +def neon_expand8_table(rhs_base: int) -> list[list[int]]: + return [neon_expand8(mask, lhs_base=0, rhs_base=rhs_base) for mask in range(256)] + + +def neon_expand8_pre() -> list[list[list[int]]]: + return [ + [ + neon_expand8(mask, lhs_base=8 - prior_ones, rhs_base=16 + prior_ones) + for mask in range(256) + ] + for prior_ones in range(9) + ] + + +def avx512_unpack_permute(depth: int) -> list[int]: + out = [] + for group in range(8): + for i in range(8): + out.append(group * depth + i if i < depth else 0) + return out + + +def avx512_shift_control(depth: int) -> list[int]: + return [i * depth for _group in range(8) for i in range(8)] + + +def avx512_unpack_permute_tables() -> list[list[int]]: + return [avx512_unpack_permute(depth) for depth in range(2, 8)] + + +def avx512_shift_control_tables() -> list[list[int]]: + return [avx512_shift_control(depth) for depth in range(2, 8)] + + +def format_u8_values(values: list[int], per_line: int = 16) -> str: + lines = [] + for i in range(0, len(values), per_line): + lines.append( + " " + ", ".join(str(value) for value in values[i : i + per_line]) + "," + ) + return "\n".join(lines) + + +def format_u8_table_1d(name: str, values: list[int], align: int) -> str: + return ( + f"static const ZL_ALIGNED({align}) uint8_t {name}[{len(values)}] = {{\n" + f"{format_u8_values(values)}\n" + "};" + ) + + +def format_u8_row(row: list[int]) -> str: + return "{ " + ", ".join(hex8(value) for value in row) + " }" + + +def format_u8_table_2d(name: str, rows: list[list[int]], align: int) -> str: + row_len = len(rows[0]) + body = ",\n".join(" " + format_u8_row(row) for row in rows) + return ( + f"static const ZL_ALIGNED({align}) uint8_t {name}[{len(rows)}][{row_len}] = {{\n" + f"{body}\n" + "};" + ) + + +def format_u8_table_3d(name: str, planes: list[list[list[int]]], align: int) -> str: + plane_len = len(planes[0]) + row_len = len(planes[0][0]) + formatted_planes = [] + for plane in planes: + rows = ",\n".join(" " + format_u8_row(row) for row in plane) + formatted_planes.append(f" {{\n{rows}\n }}") + body = ",\n".join(formatted_planes) + return ( + f"static const ZL_ALIGNED({align}) uint8_t " + f"{name}[{len(planes)}][{plane_len}][{row_len}] = {{\n" + f"{body},\n" + "};" + ) + + +def render_header(path: str, body: str) -> str: + guard = header_guard(path) + return f"""// Copyright (c) Meta Platforms, Inc. and affiliates. +#ifndef {guard} +#define {guard} + +/** + * @{GENERATED} by {SCRIPT} + * Please don't modify this file directly! + */ + +#include "openzl/shared/portability.h" + +ZL_BEGIN_C_DECLS + +// clang-format off +{body} +// clang-format on + +ZL_END_C_DECLS + +#endif // {guard} +""" + + +def render_avx2_tables() -> str: + body = "\n\n".join( + [ + format_u8_table_2d( + "ZL_kPivCoHuffmanMergeLoCtrl", merge_lo_ctrl(), align=16 + ), + format_u8_table_2d( + "ZL_kPivCoHuffmanMergeHiCtrl", merge_hi_ctrl(), align=16 + ), + format_u8_table_1d( + "ZL_kPivCoHuffmanNotPopcount8", zeros_counts(), align=16 + ), + format_u8_table_2d( + "ZL_kPivCoHuffmanCompressU16Pair", compress_u16_pair(0x80), align=32 + ), + format_u8_table_2d("ZL_kPivCoHuffmanPartitionLo", partition_lo(), align=16), + format_u8_table_2d( + "ZL_kPivCoHuffmanPartitionHiPadded", + partition_hi_padded(), + align=32, + ), + ] + ) + return render_header( + "src/openzl/codecs/pivco_huffman/arch/common_pivco_avx2_tables.h", + body, + ) + + +def render_neon_tables() -> str: + body = "\n\n".join( + [ + format_u8_table_1d("ZL_kPivCoHuffmanNeonPopcount", popcounts(), align=64), + format_u8_table_1d( + "ZL_kPivCoHuffmanNeonZeroCount", zeros_counts(), align=64 + ), + format_u8_table_2d( + "ZL_kPivCoHuffmanNeonCompressU16", + compress_u16_pair(0xFF), + align=32, + ), + format_u8_table_2d( + "ZL_kPivCoHuffmanNeonExpand8", + neon_expand8_table(rhs_base=8), + align=32, + ), + format_u8_table_2d( + "ZL_kPivCoHuffmanNeonExpand8Q", + neon_expand8_table(rhs_base=16), + align=32, + ), + format_u8_table_3d( + "ZL_kPivCoHuffmanNeonExpand8Pre", + neon_expand8_pre(), + align=64, + ), + ] + ) + return render_header( + "src/openzl/codecs/pivco_huffman/arch/common_pivco_neon_tables.h", + body, + ) + + +def render_avx512_tables() -> str: + body = "\n\n".join( + [ + format_u8_table_2d( + "ZL_kPivCoHuffmanAvx512UnpackPermute", + avx512_unpack_permute_tables(), + align=64, + ), + format_u8_table_2d( + "ZL_kPivCoHuffmanAvx512ShiftControl", + avx512_shift_control_tables(), + align=64, + ), + ] + ) + return render_header( + "src/openzl/codecs/pivco_huffman/arch/common_pivco_avx512_tables.h", + body, + ) + + +def main() -> None: + outputs = { + "src/openzl/codecs/pivco_huffman/arch/common_pivco_avx2_tables.h": render_avx2_tables(), + "src/openzl/codecs/pivco_huffman/arch/common_pivco_neon_tables.h": render_neon_tables(), + "src/openzl/codecs/pivco_huffman/arch/common_pivco_avx512_tables.h": render_avx512_tables(), + } + for path, contents in outputs.items(): + (ROOT / path).write_text(contents) + print(f"generated {path}") + + +if __name__ == "__main__": + main()