gpulz_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 <memory>
#include <stdexcept>
#include <string>
#include <unordered_map>
#include <vector>
 
namespace fz {
 
class GPULZStage : public Stage {
public:
    GPULZStage()
        : is_inverse_(false)
        , chunk_size_(2048)
        , word_size_(4)
        , actual_output_size_(0)
        , cached_orig_bytes_(0)
    {}
 
    ~GPULZStage() 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); }
    void setWordSize(size_t bytes)  { word_size_  = static_cast<uint8_t>(bytes);  }
 
    void setMatchLevel(int level) { match_level_ = static_cast<uint8_t>(level); }
    int  getMatchLevel() const    { return static_cast<int>(match_level_); }
 
    void setSplitMode(bool on) { split_mode_ = on; }
    bool getSplitMode() const  { return split_mode_; }
 
    size_t getChunkSize()              const { return chunk_size_; }
 
    size_t getRequiredInputAlignment() const override { return 1; }
    int    getWordSize()               const { return static_cast<int>(word_size_); }
    uint32_t getCachedOrigBytes()      const { return cached_orig_bytes_; }
 
    // ── 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 "GPULZ"; }
    size_t getNumInputs()  const override {
        return (is_inverse_ && split_mode_) ? 4 : 1;
    }
    size_t getNumOutputs() const override {
        return (!is_inverse_ && split_mode_) ? 4 : 1;
    }
 
    std::vector<std::string> getOutputNames() const override {
        if (!is_inverse_ && split_mode_)
            return {"literals", "lengths", "offsets", "meta"};
        return {"output"};
    }
 
    std::vector<size_t> estimateOutputSizes(
        const std::vector<size_t>& input_sizes
    ) const override {
        if (is_inverse_) {
            // Capacity must cover the padded extent the decode kernel writes,
            // even though execute() reports the unpadded size afterwards.
            if (cached_orig_bytes_ > 0)
                return {static_cast<size_t>(cached_orig_bytes_)};
            return {input_sizes.empty() ? 0 : input_sizes[0]};
        }
        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 + 8 * n_chunks;
        // Every bound below is against the PADDED extent, not the input size.
        // execute() zero-pads a partial tail chunk up to chunk_size_ and encodes
        // it as a full chunk, so the tail can contribute chunk_size_ bytes of
        // output from fewer than chunk_size_ bytes of input. Bounding by
        // `n_bytes` therefore under-reserves by exactly the tail padding, and
        // the encode writes past the buffer the DAG allocated (E22 in the
        // benchmarking repo: overruns of 4-40 B observed, silent when this
        // stage is mid-pipeline).
        const size_t padded = n_chunks * chunk_size_;
 
        if (split_mode_) {
            const size_t block_elems = chunk_size_ / word_size_;
            const size_t flag_stride = (block_elems + 7) / 8;
            // literals: every element a literal (or every chunk raw) -> padded.
            // lengths/offsets: one byte per match, at most one match per element.
            // meta: header + every chunk's full-width bitmap.
            return {align4(padded),
                    align4(n_chunks * block_elems),
                    align4(n_chunks * block_elems),
                    align4(hdr + n_chunks * flag_stride)};
        }
        // Forward: worst case = padded data (every chunk falls back to raw
        // storage) + stream header (two uint32_t per chunk).
        // postStreamSync() rounds the final size up to a 4-byte boundary and
        // zero-fills the pad; reserve that pad here too (see RREStage).
        return {align4(padded + hdr)};
    }
 
    std::unordered_map<std::string, size_t>
    getActualOutputSizesByName() const override;
    size_t getActualOutputSize(int index) const override;
 
    size_t estimateScratchBytes(
        const std::vector<size_t>& input_sizes
    ) const override {
        if (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_;
        const size_t block_elems    = chunk_size_ / word_size_;
        const size_t flag_bytes_max = (block_elems + 7) / 8;
        if (is_inverse_) {
            // Split inverse restripes the four ports back into the packed
            // single-stream form before running the normal decode path.
            return split_mode_ ? (in_bytes + n_chunks * flag_bytes_max
                                  + 4 * n_chunks * sizeof(uint32_t))
                               : 0;
        }
        size_t bytes = n_chunks * (static_cast<size_t>(chunk_size_)
                                   + flag_bytes_max + 4 * sizeof(uint32_t));
        if (split_mode_) bytes += n_chunks * 5 * sizeof(uint32_t) + 16;
        return bytes;
    }
 
    uint16_t getStageTypeId() const override {
        return static_cast<uint16_t>(StageType::GPULZ);
    }
 
    uint8_t getOutputDataType(size_t) const override {
        return static_cast<uint8_t>(DataType::UINT8);
    }
    uint8_t getInputDataType(size_t) 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 {
        (void)output_index;
        if (max_size < 14) return 0;
        std::memcpy(buf,      &chunk_size_,          sizeof(uint32_t));
        buf[4] = word_size_;
        std::memcpy(buf + 5,  &cached_orig_bytes_,   sizeof(uint32_t));
        buf[9] = split_mode_ ? 1u : 0u;
        std::memcpy(buf + 10, &orig_unpadded_bytes_, sizeof(uint32_t));
        return 14;
    }
 
    void deserializeHeader(const uint8_t* buf, size_t size) override {
        if (size >= 4)  std::memcpy(&chunk_size_,        buf,     sizeof(uint32_t));
        if (size >= 5)  word_size_ = buf[4];
        if (size >= 9)  std::memcpy(&cached_orig_bytes_, buf + 5, sizeof(uint32_t));
        if (size >= 10) split_mode_ = (buf[9] != 0);
        if (size >= 14) std::memcpy(&orig_unpadded_bytes_, buf + 10, sizeof(uint32_t));
    }
 
    size_t getMaxHeaderSize(size_t) const override { return 14; }
 
    void saveState() override {
        saved_chunk_size_        = chunk_size_;
        saved_word_size_         = word_size_;
        saved_cached_orig_bytes_ = cached_orig_bytes_;
        saved_split_mode_        = split_mode_;
        saved_orig_unpadded_bytes_ = orig_unpadded_bytes_;
    }
 
    void restoreState() override {
        chunk_size_        = saved_chunk_size_;
        word_size_         = saved_word_size_;
        cached_orig_bytes_ = saved_cached_orig_bytes_;
        split_mode_        = saved_split_mode_;
        orig_unpadded_bytes_ = saved_orig_unpadded_bytes_;
    }
 
private:
    static constexpr size_t align4(size_t n) { return (n + 3) & ~size_t(3); }
 
    void finishSplitReadback(fz::stream_t stream) const;
 
    void freeForwardScratch(fz::stream_t stream, bool sync_before_cuda_free);
 
    void executeForward(fz::stream_t stream, MemoryPool* pool,
                         const std::vector<void*>& inputs,
                         const std::vector<void*>& outputs,
                         size_t in_bytes);
    void executeInverse(fz::stream_t stream, MemoryPool* pool,
                         const std::vector<void*>& inputs,
                         const std::vector<void*>& outputs,
                         size_t in_bytes);
 
    bool     is_inverse_;
    uint32_t chunk_size_;
    uint32_t saved_chunk_size_ = 0;
    uint8_t  word_size_;
    uint8_t  saved_word_size_ = 0;
    uint8_t  match_level_ = 1;
    bool     split_mode_ = false;
    bool     saved_split_mode_ = false;
    size_t   actual_output_size_;
    // Split mode: per-port actual sizes, in getOutputNames() order.
    size_t   actual_split_sizes_[4] = {0, 0, 0, 0};
    uint32_t cached_orig_bytes_ = 0;          // chunk-padded extent the codec works on
    uint32_t saved_cached_orig_bytes_ = 0;
    // True input size before tail-chunk padding. The inverse decodes the full
    // padded extent but must *report* this, or a downstream stage that derives
    // an element count from its input size (LorenzoQuantStage's inverse does
    // exactly that) inflates its own output past the allocated buffer.
    uint32_t orig_unpadded_bytes_ = 0;
    uint32_t saved_orig_unpadded_bytes_ = 0;
 
    // ── Persistent forward scratch buffers ─────────────────────────────────
    uint8_t*  d_data_scratch_ = nullptr;
    uint8_t*  d_flag_scratch_ = nullptr;
    uint32_t* d_flag_size_    = nullptr;
    uint32_t* d_data_size_    = nullptr;
    uint32_t* d_clean_dev_    = nullptr;
    uint32_t* d_dst_off_dev_  = nullptr;
    // Split mode: per-chunk destination offsets into the literals / token
    // streams, plus a 4-entry device totals array read back in postStreamSync.
    uint32_t* d_lit_off_dev_  = nullptr;
    uint32_t* d_tok_off_dev_  = nullptr;
    uint32_t* d_meta_off_dev_ = nullptr;
    uint32_t* d_lit_cnt_dev_  = nullptr;
    uint32_t* d_tok_cnt_dev_  = nullptr;
    uint32_t* d_totals_dev_   = nullptr;
    mutable uint8_t*     split_out_ptr_[4] = {nullptr, nullptr, nullptr, nullptr};
    mutable bool         split_readback_pending_ = false;
    mutable bool         tail_readback_pending_ = false;
    mutable fz::stream_t tail_readback_stream_  = nullptr;
    mutable uint32_t     tail_last_index_       = 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