golomb_rice_stage.h

Go to the documentation of this file.

#pragma once
 
#include "stage/stage.h"
#include "fzm_format.h"
#include "backend/types.h"
#include "fused/common/data_type_of.h"
#include <cstdint>
#include <cstring>
#include <memory>
#include <stdexcept>
#include <string>
#include <type_traits>
#include <unordered_map>
#include <vector>
 
namespace fz {
 
template<typename T>
class GolombRiceStage : public Stage {
    static_assert(std::is_integral<T>::value && std::is_signed<T>::value,
                  "GolombRiceStage requires a signed integer type");
public:
    GolombRiceStage() = default;
    ~GolombRiceStage() override;
 
    // ── Stage control ──────────────────────────────────────────────────────
    void setInverse(bool inv) override { is_inverse_ = inv; }
    bool isInverse() const override    { return is_inverse_; }
 
    bool isGraphCompatible() const override { return !is_inverse_; }
 
    void setChunkSize(size_t bytes) { chunk_size_ = static_cast<uint32_t>(bytes); }
    size_t getChunkSize() const { return chunk_size_; }
    size_t getRequiredInputAlignment() const override { return chunk_size_; }
    uint32_t getCachedOrigBytes() const { return cached_orig_bytes_; }
 
    FusionSpec getFusionSpec() const override {
        if (is_inverse_ || chunk_size_ != 16384u || !std::is_same<T, int32_t>::value)
            return {};
        return FusionSpec{FusionAccess::SegmentCodec, chunk_size_};
    }
 
    FusedOpDecl getFusedOp() const override {
        if (!getFusionSpec().fusable()) return {};
        return FusedOpDecl{FusionStrategy::ChunkCooperative, "GolombRiceCoder",
                           "fused/chunk_fusion/chunk_fusion.cuh", {}};
    }
    void setFusedArchiveResult(size_t archive_bytes, size_t orig_bytes) override {
        setFusedResult(archive_bytes, orig_bytes);
    }
    void setFusedResult(size_t archive_bytes, size_t orig_bytes) {
        actual_output_size_    = archive_bytes;
        cached_orig_bytes_     = static_cast<uint32_t>(orig_bytes);
        tail_readback_pending_ = false;
    }
 
    // ── 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;
    void postStreamSync(fz::stream_t stream) override;
 
    // ── Metadata ───────────────────────────────────────────────────────────
    std::string getName() const override { return "GolombRice"; }
    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 (is_inverse_) {
            if (cached_orig_bytes_ > 0)
                return {static_cast<size_t>(cached_orig_bytes_)};
            return {input_sizes.empty() ? 0 : input_sizes[0]};
        }
        // Forward worst case: every chunk falls back to raw storage (original
        // bytes + 1 flagged size word each) plus the stream header, plus the
        // tail safety pad postStreamSync may need to zero (see execute()).
        const size_t n_bytes  = input_sizes.empty() ? 0 : input_sizes[0];
        const size_t n_chunks = (n_bytes + chunk_size_ - 1) / chunk_size_;
        const size_t hdr      = 4 + 4 + 4 * n_chunks;
        const size_t worst    = n_bytes + hdr + kTailPad;
        return {(worst + 3) & ~size_t(3)};
    }
 
    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;
    }
 
    size_t estimateScratchBytes(
        const std::vector<size_t>& input_sizes
    ) const override {
        if (is_inverse_ || input_sizes.empty()) return 0;
        const size_t in_bytes = input_sizes[0];
        const size_t n_chunks = (in_bytes + chunk_size_ - 1) / chunk_size_;
        return n_chunks * (chunkScratchStride() + 3 * sizeof(uint32_t));
    }
 
    uint16_t getStageTypeId() const override {
        return static_cast<uint16_t>(StageType::GOLOMB_RICE);
    }
 
    uint8_t getOutputDataType(size_t) const override {
        return static_cast<uint8_t>(DataType::UINT8);
    }
 
    // ── Serialization ──────────────────────────────────────────────────────
    static DataType getElementDataType() { return fused::dataTypeOf<T>(); }
 
    size_t serializeHeader(size_t, uint8_t* buf, size_t max_size) const override {
        if (max_size < 5) return 0;
        std::memcpy(buf, &chunk_size_, sizeof(uint32_t));
        buf[4] = static_cast<uint8_t>(getElementDataType());
        return 5;
    }
    void deserializeHeader(const uint8_t* buf, size_t size) override {
        if (size >= 4) std::memcpy(&chunk_size_, buf, sizeof(uint32_t));
    }
    size_t getMaxHeaderSize(size_t) const override { return 5; }
 
    void saveState() override { saved_chunk_size_ = chunk_size_; }
    void restoreState() override { chunk_size_ = saved_chunk_size_; }
 
    uint32_t elemsPerChunk() const { return chunk_size_ / static_cast<uint32_t>(sizeof(T)); }
 
    static constexpr uint32_t kEscapeQ = 24;
    static constexpr uint32_t bitWidth() { return 8u * sizeof(T); }
    static constexpr uint32_t kMaxCandidate() { return bitWidth() > 24 ? 24u : bitWidth() - 1; }
 
    static constexpr uint32_t kIntervalsPerChunk = 16;
 
    size_t chunkScratchStride() const {
        const size_t worst_bits  = static_cast<size_t>(elemsPerChunk()) * (kEscapeQ + bitWidth());
        const size_t worst_bytes = (4 + 4 * kIntervalsPerChunk) /*k + offset table*/
                                 + kIntervalsPerChunk /*per-interval rounding slack*/
                                 + (worst_bits + 7) / 8;
        return ((worst_bytes + 3) & ~size_t(3)) + kTailPad;
    }
    static constexpr size_t kTailPad = 16;
 
private:
    bool     is_inverse_ = false;
    uint32_t chunk_size_ = 16384;
    uint32_t saved_chunk_size_ = 0;
    size_t   actual_output_size_ = 0;
    uint32_t cached_orig_bytes_ = 0;
    uint32_t saved_cached_orig_bytes_ = 0;
 
    // ── Persistent forward scratch buffers ───────────────────────────────────
    uint8_t*  d_scratch_ = nullptr;
    uint32_t* d_sizes_dev_ = nullptr;
    uint32_t* d_dst_off_dev_ = nullptr;
    mutable bool         tail_readback_pending_ = false;
    mutable uint32_t     tail_last_index_ = 0;
    mutable uint32_t     tail_header_size_ = 0;
    mutable uint8_t*     tail_output_ptr_ = nullptr;
    size_t    scratch_capacity_ = 0;
    MemoryPool* scratch_pool_owner_ = nullptr;
    bool        scratch_from_pool_ = false;
    std::weak_ptr<const void> scratch_alive_;
};
 
} // namespace fz