rze_stage.h

Go to the documentation of this file.

#pragma once
 
#include "stage/stage.h"
#include "fzm_format.h"
#include <cuda_runtime.h>
#include <cstdint>
#include <cstring>
#include <stdexcept>
#include <string>
#include <unordered_map>
#include <vector>
 
namespace fz {
 
class RZEStage : public Stage {
public:
    RZEStage()
        : is_inverse_(false)
        , chunk_size_(16384)
        , word_size_(1)
        , actual_output_size_(0)
        , cached_orig_bytes_(0)
        , d_scratch_(nullptr)
        , d_sizes_dev_(nullptr)
        , d_clean_dev_(nullptr)
        , d_dst_off_dev_(nullptr)
        , scratch_capacity_(0)
    {}
 
    ~RZEStage() 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);  }
 
    size_t getChunkSize()       const { return chunk_size_; }
    size_t getRequiredInputAlignment() const override { return chunk_size_; }
    int    getWordSize()        const { return static_cast<int>(word_size_); }
    uint32_t getCachedOrigBytes() const { return cached_orig_bytes_; }
 
    // ── Execution ──────────────────────────────────────────────────────────
    void execute(
        cudaStream_t stream,
        MemoryPool* pool,
        const std::vector<void*>& inputs,
        const std::vector<void*>& outputs,
        const std::vector<size_t>& sizes
    ) override;
    void postStreamSync(cudaStream_t stream) override;
 
    // ── Metadata ───────────────────────────────────────────────────────────
    std::string getName() const override { return "RZE"; }
    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 = original data + stream header.
        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;
        return {n_bytes + 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 (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 * (static_cast<size_t>(chunk_size_) + 3 * sizeof(uint32_t));
    }
 
    uint16_t getStageTypeId() const override {
        return static_cast<uint16_t>(StageType::RZE);
    }
 
    uint8_t getOutputDataType(size_t) const override {
        return static_cast<uint8_t>(DataType::UINT8);
    }
 
    // ── Serialization ──────────────────────────────────────────────────────
    size_t serializeHeader(
        size_t output_index, uint8_t* buf, size_t max_size
    ) const override {
        (void)output_index;
        if (max_size < 9) return 0;
        std::memcpy(buf,     &chunk_size_,        sizeof(uint32_t));
        buf[4] = word_size_;
        std::memcpy(buf + 5, &cached_orig_bytes_, sizeof(uint32_t));
        return 9;
    }
 
    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));
    }
 
    size_t getMaxHeaderSize(size_t) const override { return 9; }
 
    void saveState() override {
        saved_chunk_size_ = chunk_size_;
        saved_word_size_ = word_size_;
        saved_cached_orig_bytes_ = cached_orig_bytes_;
    }
 
    void restoreState() override {
        chunk_size_ = saved_chunk_size_;
        word_size_ = saved_word_size_;
        cached_orig_bytes_ = saved_cached_orig_bytes_;
    }
 
private:
    bool     is_inverse_;
    uint32_t chunk_size_;
    uint32_t saved_chunk_size_ = 0;
    uint8_t  word_size_;
    uint8_t  saved_word_size_ = 0;
    size_t   actual_output_size_;
    uint32_t cached_orig_bytes_ = 0;
    uint32_t saved_cached_orig_bytes_ = 0;
 
    // ── Persistent forward scratch buffers ───────────────────────────────────
    uint8_t*  d_scratch_;
    uint32_t* d_sizes_dev_;
    uint32_t* d_clean_dev_;
    uint32_t* d_dst_off_dev_;
    mutable bool         tail_readback_pending_ = false;
    mutable cudaStream_t tail_readback_stream_ = nullptr;
    mutable uint32_t     tail_last_index_ = 0;
    mutable uint8_t*     tail_output_ptr_ = nullptr;
    size_t    scratch_capacity_;
    MemoryPool* scratch_pool_owner_ = nullptr;
    bool        scratch_from_pool_ = false;
};
 
} // namespace fz