bitpack_stage.h

Go to the documentation of this file.

#pragma once
 
#include "stage/stage.h"
#include "fzm_format.h"
#include "backend/types.h"
#include <cstdint>
#include <cstring>
#include <stdexcept>
#include <string>
#include <type_traits>
#include <unordered_map>
#include <vector>
 
namespace fz {
 
template<typename T>
class BitpackStage : 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>,
        "BitpackStage: T must be uint8_t, uint16_t, or uint32_t.");
 
public:
    BitpackStage() = default;
 
    // ── Stage control ──────────────────────────────────────────────────────────
    void setInverse(bool inv) override { is_inverse_ = inv; }
    bool isInverse() const override    { return is_inverse_; }
 
    // ── Configuration ──────────────────────────────────────────────────────────
 
    void setNBits(uint8_t nbits) {
        if (nbits == 0 || nbits > 8 * sizeof(T) || (nbits & (nbits - 1)) != 0)
            throw std::invalid_argument(
                "BitpackStage::setNBits: nbits must be a power of two "
                "in [1, " + std::to_string(8 * sizeof(T)) + "], got "
                + std::to_string(nbits));
        nbits_ = nbits;
    }
    uint8_t getNBits() const { return nbits_; }
 
    void setBase(T base) { base_ = base; }
    T getBase() const { return base_; }
 
    void setShift(uint8_t shift) {
        if (shift >= 8 * sizeof(T))
            throw std::invalid_argument(
                "BitpackStage::setShift: shift must be in [0, "
                + std::to_string(8 * sizeof(T) - 1) + "], got "
                + std::to_string(shift));
        shift_ = shift;
    }
    uint8_t getShift() const { return shift_; }
 
    void setAutoDetect(bool enable) { auto_detect_ = enable; }
    bool isAutoDetect() const { return auto_detect_; }
 
    void setAutoBase(bool enable) { auto_base_ = enable; }
    bool isAutoBase() const { return auto_base_; }
 
    void setAutoShift(bool enable) { auto_shift_ = enable; }
    bool isAutoShift() const { return auto_shift_; }
 
    void setAdaptive(bool enable) {
        auto_base_ = auto_shift_ = auto_detect_ = enable;
    }
 
    // ── 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 "Bitpack"; }
    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_) {
            if (auto_detect_) {
                // nbits is unknown until execute() scans the data; return worst
                // case (full-width, no compression) so PREALLOCATE has enough room.
                return {input_sizes[0]};
            }
            // Forward: packed output is ceil(n * nbits / 8) bytes.
            const size_t n = input_sizes[0] / sizeof(T);
            return {(n * nbits_ + 7) / 8};
        } else {
            // Inverse: worst case — every packed bit expands to a full element.
            // input_sizes[0] is the packed byte count; max elements = bytes * (8/nbits).
            const size_t max_elems = (input_sizes[0] * 8 + nbits_ - 1) / nbits_;
            return {max_elems * 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::BITPACK);
    }
 
    // Packed byte stream has no meaningful element type; opt out of type 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 < 15) return 0;
        buf[0] = static_cast<uint8_t>(dataTypeOf<T>());
        buf[1] = nbits_;
        std::memcpy(buf + 2, &num_elements_, sizeof(uint64_t));
        buf[10] = shift_;
        const uint32_t base32 = static_cast<uint32_t>(base_);
        std::memcpy(buf + 11, &base32, sizeof(uint32_t));
        return 15;
    }
 
    void deserializeHeader(const uint8_t* buf, size_t size) override {
        // buf[0] (DataType) is used by the factory to pick the right instantiation.
        // We only need nbits, num_elements, shift, and base here.
        if (size >= 2)  nbits_ = buf[1];
        if (size >= 10) std::memcpy(&num_elements_, buf + 2, sizeof(uint64_t));
        // Pre-shift archives stop at 10 bytes; leave shift/base at their defaults.
        if (size >= 15) {
            shift_ = buf[10];
            uint32_t base32 = 0;
            std::memcpy(&base32, buf + 11, sizeof(uint32_t));
            base_ = static_cast<T>(base32);
        }
    }
 
    size_t getMaxHeaderSize(size_t /*output_index*/) const override { return 15; }
 
    // saveState/restoreState: deserializeHeader (called during decompression
    // setup) overwrites num_elements with the value from the file header.
    // Save the forward-pass values so they can be restored afterward.
    void saveState() override {
        saved_nbits_        = nbits_;
        saved_num_elements_ = num_elements_;
        saved_output_size_  = actual_output_size_;
        saved_shift_        = shift_;
        saved_base_         = base_;
    }
 
    void restoreState() override {
        nbits_               = saved_nbits_;
        num_elements_        = saved_num_elements_;
        actual_output_size_  = saved_output_size_;
        shift_               = saved_shift_;
        base_                = saved_base_;
    }
 
    // Auto-detect requires a D2H sync to read the scanned min/max/OR, so it
    // cannot be recorded inside a CUDA Graph.
    bool isGraphCompatible() const override {
        return !(auto_detect_ || auto_base_ || auto_shift_);
    }
 
private:
    bool     is_inverse_        = false;
    bool     auto_detect_       = false;
    bool     auto_base_         = false;
    bool     auto_shift_        = false;
    uint8_t  nbits_             = 8 * sizeof(T);   // default: keep all bits (identity)
    uint8_t  shift_             = 0;               // low bits dropped before packing
    T        base_              = T(0);            // frame-of-reference offset
    uint64_t num_elements_      = 0;               // set by forward execute; used by inverse
    size_t   actual_output_size_ = 0;
 
    // saveState snapshots
    uint8_t  saved_nbits_        = 8 * sizeof(T);
    uint8_t  saved_shift_        = 0;
    T        saved_base_         = T(0);
    uint64_t saved_num_elements_ = 0;
    size_t   saved_output_size_  = 0;
 
    template<typename U>
    static constexpr DataType dataTypeOf() {
        if (std::is_same_v<U,  uint8_t>) return DataType::UINT8;
        if (std::is_same_v<U, uint16_t>) return DataType::UINT16;
        if (std::is_same_v<U, uint32_t>) return DataType::UINT32;
        return DataType::UINT8; // unreachable
    }
};
 
extern template class BitpackStage<uint8_t>;
extern template class BitpackStage<uint16_t>;
extern template class BitpackStage<uint32_t>;
 
} // namespace fz