huffman_stage.h

Go to the documentation of this file.

#pragma once
 
#include "stage/stage.h"
#include "fzm_format.h"
#include "coders/huffman/phf/hf.h"   // phf_header, phf_stream_t
 
#include "backend/types.h"
#include <cstdint>
#include <cstring>
#include <memory>
#include <stdexcept>
#include <string>
#include <type_traits>
#include <unordered_map>
#include <vector>
 
// Forward-declare phf::Buf<T> to avoid pulling hf_buf.h (CUDA-only) into this header.
// HuffmanStage<T>::~HuffmanStage() is defined in huffman_stage.cu where the type is complete.
namespace phf { template<typename E> struct Buf; }
 
namespace fz {
 
enum class HuffmanEncodeMode {
    Coarse, 
    Fine,   
};
 
enum class HuffmanBookSource {
    PerBlock,  
    Fixed,     
    Adaptive,  
};
 
enum class HuffmanBookModel {
    Gaussian,           
    Laplace,            
    GeneralizedNormal,  
    Uniform,            
};
 
struct HuffmanBookSpec {
    HuffmanBookModel model  = HuffmanBookModel::Gaussian;
    double           center = -1.0;
    double           scale  = 32.0;
    double           shape  = 2.0;
};
 
template <typename T>
class HuffmanStage : public Stage {
    static_assert(
        std::is_same_v<T, uint8_t>  ||
        std::is_same_v<T, uint16_t> ||
        std::is_same_v<T, uint32_t>,
        "HuffmanStage: T must be uint8_t, uint16_t, or uint32_t.");
 
public:
    // __host__-only: HIP's clang infers __host__ __device__ for defaulted
    // special members by default, but ~unique_ptr<phf::Buf<T>> (the type this
    // destructor implicitly invokes, defined in huffman_stage.cu) is
    // host-only, so that inference fails to compile under HIP. Neither
    // function touches device code at all — pin them host-only explicitly.
    __host__ HuffmanStage();
    __host__ ~HuffmanStage() override;
 
    // ── Configuration ─────────────────────────────────────────────────────────
 
    void setBklen(uint32_t bklen) {
        constexpr uint32_t kMul = (4u / sizeof(T)) ? (4u / sizeof(T)) : 1u;
        bklen_ = ((bklen + kMul - 1u) / kMul) * kMul;
    }
    uint32_t getBklen() const         { return bklen_; }
 
    void             setEncodeMode(HuffmanEncodeMode mode) { encode_mode_ = mode; }
    HuffmanEncodeMode getEncodeMode() const                { return encode_mode_; }
 
    // ── Pre-built codebooks ───────────────────────────────────────────────────
 
    void              setBookSource(HuffmanBookSource src) { book_source_ = src; }
    HuffmanBookSource getBookSource() const                { return book_source_; }
 
    void setFixedBookFromFreq(const uint32_t* h_freq, uint32_t n);
 
    void setFixedBookFromModel(const HuffmanBookSpec& spec);
 
    void setValidateSymbolRange(bool on) { validate_symbol_range_ = on; }
    bool getValidateSymbolRange() const  { return validate_symbol_range_; }
 
    const std::vector<uint32_t>& getFixedBookFreq() const { return fixed_freq_; }
 
    void    setAdaptiveFloorShift(uint8_t shift) { adaptive_floor_shift_ = shift; }
    uint8_t getAdaptiveFloorShift() const        { return adaptive_floor_shift_; }
 
    uint8_t getAdaptiveFloorShiftUsed() const    { return adaptive_shift_used_; }
 
    bool getAdaptiveFallbackUsed() const { return adaptive_fallback_; }
 
    std::vector<std::string> getRunNotes() const override {
        if (adaptive_fallback_) return {"huffman_adaptive_fallback"};
        return {};
    }
 
    void  setRefitThreshold(float ratio) { refit_threshold_ = ratio; }
    float getRefitThreshold() const      { return refit_threshold_; }
 
    void     setRefitInterval(uint32_t n) { refit_interval_ = n; }
    uint32_t getRefitInterval() const     { return refit_interval_; }
 
    uint32_t getRefitCount() const { return refit_count_; }
 
    double getFitBitsPerSymbol() const { return fit_bits_per_sym_; }
 
    bool getLastUsedFineEncode() const { return last_used_fine_; }
 
    uint8_t getLastMaxCodeLen() const { return last_max_codelen_; }
 
    bool                   hasBookSpec() const { return has_book_spec_; }
    const HuffmanBookSpec& getBookSpec() const { return book_spec_; }
 
    // ── Stage control ─────────────────────────────────────────────────────────
    void setInverse(bool inv) override { is_inverse_ = inv; }
    bool isInverse() const override    { return is_inverse_; }
 
    // Not graph-compatible in any configuration, and not planned.  Fixed/Adaptive
    // remove the histogram D2H, but encode still returns total_nbit/total_ncell to
    // the host to assemble phf_header before the H2D merge.  Closing that would need
    // device-side header assembly and a device-side merge; graph capture has not
    // shown a measurable win elsewhere in the library, so it is not worth that.
    bool isGraphCompatible() const override { return false; }
 
    // ── Pool lifecycle ────────────────────────────────────────────────────────
 
    void onFinalize(size_t estimated_inlen, MemoryPool* pool) override;
 
    size_t estimateDeviceFootprintBytes(size_t inlen) const override;
    size_t estimatePinnedFootprintBytes(size_t inlen) const override;
 
    // ── Execution ─────────────────────────────────────────────────────────────
    void execute(
        fz::stream_t stream,
        MemoryPool*  pool,
        const std::vector<void*>&   inputs,
        const std::vector<void*>&   outputs,
        const std::vector<size_t>&  sizes
    ) override;
 
    // ── Metadata ──────────────────────────────────────────────────────────────
    std::string getName()      const override { return "Huffman"; }
    size_t      getNumInputs() const override { return 1; }
    size_t      getNumOutputs()const override { return 1; }
 
    std::vector<size_t> estimateOutputSizes(
        const std::vector<size_t>& input_sizes
    ) const override {
        if (input_sizes.empty()) return {0};
        if (!is_inverse_) {
            // Upper bound: generous 2× input for worst-case bitstream + header overhead.
            return {input_sizes[0] * 2 + 4096};
        }
        // Inverse: exact decoded size restored from the serialized header.
        return {original_len_ * sizeof(T)};
    }
 
    std::unordered_map<std::string, size_t>
    getActualOutputSizesByName() const override {
        return {{"output", actual_output_size_}};
    }
 
    size_t getActualOutputSize(int index) const override {
        return (index == 0) ? actual_output_size_ : 0;
    }
 
    // ── Type system ───────────────────────────────────────────────────────────
    uint16_t getStageTypeId() const override {
        return static_cast<uint16_t>(StageType::HUFFMAN);
    }
 
    // Byte-transparent output: opt out of pipeline type-compatibility checking.
    uint8_t getOutputDataType(size_t /*output_index*/) const override {
        return static_cast<uint8_t>(DataType::UNKNOWN);
    }
    uint8_t getInputDataType(size_t /*input_index*/) const override {
        return static_cast<uint8_t>(DataType::UNKNOWN);
    }
 
    // ── Serialization ─────────────────────────────────────────────────────────
    size_t serializeHeader(
        size_t /*output_index*/, uint8_t* buf, size_t max_size
    ) const override {
        if (max_size < 11) return 0;
        buf[0] = static_cast<uint8_t>(dataTypeOf<T>());
        uint16_t bk = static_cast<uint16_t>(bklen_);
        std::memcpy(buf + 1, &bk,           sizeof(uint16_t));
        std::memcpy(buf + 3, &original_len_, sizeof(uint64_t));
        return 11;
    }
 
    void deserializeHeader(const uint8_t* buf, size_t size) override {
        if (size >= 3) {
            uint16_t bk;
            std::memcpy(&bk, buf + 1, sizeof(uint16_t));
            // Take the archive's bklen verbatim — NOT through setBklen(). Decode has
            // to reproduce the layout the encoder actually used, so rounding here
            // would describe a different reverse codebook than the one in the stream.
            //
            // But an archive written before setBklen enforced 4-byte alignment can
            // carry a bklen that puts the bitstream at a non-4 offset, and the decode
            // kernel reads it as uint32*. That does not merely decode wrong: it
            // poisons the CUDA context, and every later allocation in the process
            // fails with "misaligned address" — a failure that points nowhere near
            // this stage. Refuse it here, where the cause is still legible.
            if ((sizeof(T) * static_cast<size_t>(bk)) % 4u != 0)
                throw std::runtime_error(
                    "HuffmanStage: archive declares bklen=" + std::to_string(bk) +
                    " with a " + std::to_string(sizeof(T)) + "-byte symbol, which puts "
                    "the bitstream at a non-4-byte offset. Such archives were written "
                    "by a build predating the alignment fix and cannot be decoded "
                    "correctly — the stream itself is malformed, not just this reader.");
            bklen_ = bk;
        }
        if (size >= 11)
            std::memcpy(&original_len_, buf + 3, sizeof(uint64_t));
    }
 
    size_t getMaxHeaderSize(size_t /*output_index*/) const override { return 11; }
 
    void saveState() override {
        saved_bklen_        = bklen_;
        saved_original_len_ = original_len_;
        saved_output_size_  = actual_output_size_;
    }
 
    void restoreState() override {
        bklen_              = saved_bklen_;
        original_len_       = saved_original_len_;
        actual_output_size_ = saved_output_size_;
    }
 
private:
    bool              is_inverse_   = false;
    uint32_t          bklen_        = defaultBklen();
    HuffmanEncodeMode encode_mode_  = HuffmanEncodeMode::Coarse;
    HuffmanBookSource book_source_  = HuffmanBookSource::PerBlock;
 
    // Frequency table defining the fixed codebook (host-only; see setFixedBookFromFreq).
    std::vector<uint32_t> fixed_freq_;
    // Model that generated fixed_freq_, when it came from setFixedBookFromModel().
    HuffmanBookSpec       book_spec_     {};
    bool                  has_book_spec_ = false;
    uint8_t               adaptive_floor_shift_ = 24;
    uint8_t               adaptive_shift_used_  = 24;
    bool                  adaptive_fallback_    = false;
 
    // Refit bookkeeping (Adaptive only).
    float    refit_threshold_  = 1.2f;
    bool     validate_symbol_range_ = true;
    uint32_t refit_interval_   = 0;     // 0 = only the bit-rate trigger
    uint32_t calls_since_fit_  = 0;
    double   fit_bits_per_sym_ = 0.0;   // rate the resident book achieved when fitted
    bool     just_fitted_      = false; // next encode establishes fit_bits_per_sym_
    uint32_t refit_count_      = 0;
    // True when buf_->d_bk4 / d_revbk4 currently hold the fixed book.  Cleared by
    // initBuf(), which allocates fresh (uninitialized) codebook buffers.
    bool                  fixed_book_resident_ = false;
 
    // Mirrored out of phf::Buf after each forward encode so callers can see which
    // encode path actually ran (Fine silently degrades to Coarse for long codes).
    bool    last_used_fine_   = false;
    uint8_t last_max_codelen_ = 0;
    // Max code length the Fine-fallback warning last fired for; suppresses one warning
    // per compress call when a Fixed/Adaptive book stays resident.
    uint8_t warned_max_codelen_ = 0;
 
    uint64_t original_len_       = 0;   // element count set by forward execute
    size_t   actual_output_size_ = 0;
    size_t            cap_inlen_        = 0;                          // allocated capacity (elements); grow-only
    uint32_t          last_bklen_       = 0;                          // bklen_ when buf_ was last allocated
    HuffmanEncodeMode last_encode_mode_ = HuffmanEncodeMode::Coarse;  // encode_mode_ when buf_ was last allocated
 
    // Histogram launch params — computed once in initBuf(), reused every execute()
    int hist_grid_dim_    = 0;
    int hist_block_dim_   = 0;
    int hist_shmem_use_   = 0;
    int hist_r_per_block_ = 0;
 
    // PHF working buffers — allocated from pool_ on first execute() or in onFinalize()
    std::unique_ptr<phf::Buf<T>> buf_;
    phf_header header_ {};
 
    // Pool used for buf_ allocations.  Set by onFinalize() or captured from the
    // pool parameter on the first execute() call.  Raw non-owning pointer; the
    // pool outlives the stage when used inside a Pipeline.
    MemoryPool* pool_ = nullptr;
 
    // saveState / restoreState snapshots
    uint32_t saved_bklen_        = defaultBklen();
    uint64_t saved_original_len_ = 0;
    size_t   saved_output_size_  = 0;
 
    static constexpr uint32_t defaultBklen() {
        if constexpr (std::is_same_v<T, uint8_t>) return 256;
        return 1024;
    }
 
    template<typename U>
    static constexpr DataType dataTypeOf() {
        if constexpr (std::is_same_v<U,  uint8_t>)  return DataType::UINT8;
        if constexpr (std::is_same_v<U, uint16_t>)  return DataType::UINT16;
        return DataType::UINT32;
    }
 
    // Allocates buf_ from pool and computes histogram launch params for the given
    // element count.  Must be in huffman_stage.cu: calls cudaFuncSetAttribute with
    // a __global__ pointer.  If buf_ already exists, destroys it first (returning
    // its allocations to the pool) before creating the new one.
    void initBuf(size_t inlen, MemoryPool* pool);
 
    // Builds the canonical codebook from fixed_freq_ into buf_ and H2Ds it.
    // Requires buf_ to exist.  Must be in huffman_stage.cu (needs phf::Buf<T>).
    void buildFixedBook(fz::stream_t stream);
 
    // Builds a reusable codebook from a sampled histogram, flooring frequencies and
    // retrying with a flatter floor until the book fits the 27-bit code field.
    void buildAdaptiveBook(const uint32_t* h_hist, fz::stream_t stream);
 
    // Index of the first symbol with freq > 0 whose built code is unusable, or -1.
    int findUnusableCode(const uint32_t* freq) const;
};
 
extern template class HuffmanStage<uint8_t>;
extern template class HuffmanStage<uint16_t>;
extern template class HuffmanStage<uint32_t>;
 
} // namespace fz