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 HuffmanExecutionMode {
    HostCoordinated, 
    DeviceResident,  
};
 
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;
        uint32_t rounded = ((bklen + kMul - 1u) / kMul) * kMul;
        // bklen_ is stored uint16_t in the serialized header (see the file
        // comment above) and threaded as uint16_t through the histogram
        // kernel API end to end. 65536 itself would silently wrap to 0 at
        // that boundary — validate here instead of producing a corrupt or
        // crashing stage.
        if (rounded > 0xFFFFu)
            throw std::invalid_argument(
                "HuffmanStage::setBklen: bklen (after rounding, " + std::to_string(rounded) +
                ") exceeds 65535 — bklen_ is a uint16_t end to end (serialized header and "
                "histogram kernel API), so the full symbol range of a 16+ bit type cannot be "
                "histogrammed in one pass. Reduce the upstream quant_radius/code range, or "
                "split the input across multiple Huffman stages.");
        bklen_ = rounded;
    }
    uint32_t getBklen() const         { return bklen_; }
 
    void setExecutionMode(HuffmanExecutionMode mode) { execution_mode_ = mode; }
    HuffmanExecutionMode getExecutionMode() const { return execution_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                   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_; }
 
    bool isGraphCompatible() const override {
        return !is_inverse_ && execution_mode_ == HuffmanExecutionMode::DeviceResident
               && book_source_ == HuffmanBookSource::Fixed && is_terminal_output_;
    }
 
    void setTerminalOutput(bool terminal) override { is_terminal_output_ = terminal; }
 
    void postStreamSync(fz::stream_t stream) override;
 
    // ── 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_) {
            const size_t n = input_sizes[0] / sizeof(T);
            if (n == 0) return {0};
            const size_t sublen = capi_phf_coarse_tune_sublen(n);
            const size_t pardeg = (n - 1) / sublen + 1;
            constexpr size_t kMaxCodeBits = 27;
            constexpr size_t kCellBits = 32;
            const size_t cells_per_partition =
                (sublen * kMaxCodeBits + kCellBits - 1) / kCellBits;
            const size_t bitstream_bytes = pardeg * cells_per_partition * sizeof(uint32_t);
            const size_t reverse_book_bytes =
                phf_reverse_book_bytes(static_cast<uint16_t>(bklen_), 4, sizeof(T));
            return {PHFHEADER_FORCED_ALIGN + reverse_book_bytes
                    + 2 * pardeg * sizeof(PHF_METADATA) + bitstream_bytes};
        }
        // 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();
    HuffmanExecutionMode execution_mode_ = HuffmanExecutionMode::HostCoordinated;
    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;
 
    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
 
    // 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;
 
    // cuSZ Huffman working buffers — allocated on first execute() or in onFinalize()
    std::unique_ptr<phf::Buf<T>> buf_;
    phf_header header_ {};
    uint8_t* pending_device_output_ = nullptr;
    bool pending_device_readback_ = false;
    size_t pending_device_inlen_ = 0;
    bool is_terminal_output_ = true;
 
    // 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);
 
    // Builds d_bk4 and d_revbk4 directly from buf_->d_freq. DeviceResident uses
    // this for PerBlock/Adaptive histograms and for a Fixed table copied H2D.
    void buildDeviceBook(fz::stream_t stream, MemoryPool* pool,
                         size_t expected_symbols, uint32_t source_mode);
 
    // Records the achieved bit rate and applies Adaptive refit policy. Called
    // immediately for HostCoordinated and from postStreamSync for DeviceResident.
    void updateAdaptiveRate(size_t inlen, size_t total_nbit);
 
    // 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