compressor.h

Go to the documentation of this file.

 
#pragma once
 
#include "pipeline/dag.h"
#include "pipeline/perf.h"
#include "pipeline/config.h"
#include "stage/stage.h"
#include "stage/stage_factory.h"
#include "mem/mempool.h"
#include "fzm_format.h"
 
#include <array>
#include <memory>
#include <stdexcept>
#include <unordered_map>
 
namespace fz {
 
class Pipeline {
public:
    explicit Pipeline(
        size_t input_data_size = 0,
        MemoryStrategy strategy = MemoryStrategy::MINIMAL,
        float pool_multiplier = 3.0f
    );
 
    explicit Pipeline(const std::string& config_path);
 
    ~Pipeline();
 
    // ── Configuration ─────────────────────────────────────────────────────────
 
    void setMemoryStrategy(MemoryStrategy strategy);
 
    void setNumStreams(int num_streams);
 
    void setDims(size_t x, size_t y = 1, size_t z = 1) { dims_ = {x, y, z}; }
    void setDims(std::array<size_t, 3> dims)             { dims_ = dims; }
    std::array<size_t, 3> getDims() const                { return dims_; }
 
    // ── Builder API ───────────────────────────────────────────────────────────
 
    template<typename StageT, typename... Args>
    StageT* addStage(Args&&... args);
 
    int connect(Stage* dependent, Stage* producer, const std::string& output_name = "output");
 
    int connect(Stage* dependent, const std::vector<Stage*>& producers);
 
    void finalize();
 
    void warmup(cudaStream_t stream = 0);
 
    void setWarmupOnFinalize(bool enable) { warmup_on_finalize_ = enable; }
    bool isWarmupOnFinalizeEnabled() const { return warmup_on_finalize_; }
 
    void setPoolManagedDecompOutput(bool enable) { pool_managed_decomp_ = enable; }
    bool isPoolManagedDecompOutput() const { return pool_managed_decomp_; }
 
    size_t getMaxCompressedSize(size_t input_bytes) const;
 
    size_t getLastUncompressedSize() const {
        return original_input_size_ > 0 ? original_input_size_ : input_size_;
    }
 
    // ── Execution ─────────────────────────────────────────────────────────────
 
    void compress(
        const void* d_input,
        size_t      input_size,
        void**      d_output,
        size_t*     output_size,
        cudaStream_t stream = 0
    );
 
    void compress(
        const void* d_input,
        size_t      input_size,
        void*       d_output_buf,
        size_t      output_buf_capacity,
        size_t*     actual_output_size,
        cudaStream_t stream = 0
    );
 
    void decompress(
        const void* d_input,
        size_t      input_size,
        void**      d_output,
        size_t*     output_size,
        cudaStream_t stream = 0
    );
 
    void decompress(
        const void* d_input,
        size_t      input_size,
        void*       d_output_buf,
        size_t      output_buf_capacity,
        size_t*     actual_output_size,
        cudaStream_t stream = 0
    );
 
    void decompressInto(
        const void* d_input,
        size_t      input_size,
        void*       d_output_buf,
        size_t      output_buf_capacity,
        size_t*     actual_output_size,
        cudaStream_t stream = 0
    );
 
    void prepareInverse(size_t uncompressed_size);
 
    void reset(cudaStream_t stream = 0);
 
    // ── Profiling ─────────────────────────────────────────────────────────────
 
    void enableProfiling(bool enable);
    bool isProfilingEnabled() const { return profiling_enabled_; }
 
    const PipelinePerfResult& getLastPerfResult() const { return last_perf_result_; }
 
    CompressionDAG* getDAG() { return dag_.get(); }
 
    size_t getPoolThreshold() const;
 
    bool isMemPoolFallbackMode() const;
 
    void enableBoundsCheck(bool enable) { dag_->enableBoundsCheck(enable); }
    bool isBoundsCheckEnabled() const   { return dag_->isBoundsCheckEnabled(); }
 
    void setColoringEnabled(bool enable) { dag_->setColoringEnabled(enable); }
    bool isColoringEnabled() const       { return dag_->isColoringEnabled(); }
    size_t getColorRegionCount() const   { return dag_->getColorRegionCount(); }
 
    // ── CUDA Graph Capture (compression-only) ─────────────────────────────────
 
    void enableGraphMode(bool enable);
    bool isGraphModeEnabled() const { return graph_mode_enabled_; }
 
    void captureGraph(cudaStream_t stream = 0);
    bool isGraphCaptured() const { return graph_captured_; }
 
    size_t getPeakMemoryUsage() const;
    size_t getCurrentMemoryUsage() const;
    void printPipeline() const;
 
    // ── File Serialization ────────────────────────────────────────────────────
 
    struct FZMFileHeader {
        FZMHeaderCore               core;
        std::vector<FZMStageInfo>   stages;
        std::vector<FZMBufferEntry> buffers;
    };
 
    void writeToFile(const std::string& filename, cudaStream_t stream = 0);
 
    static FZMFileHeader readHeader(const std::string& filename);
 
    FZMFileHeader buildHeader() const;
 
    // ── In-memory metadata header (decode without a prior compress) ────────────
 
    std::vector<uint8_t> serializeHeaderToMemory() const;
 
    void primeInverseFromHeader(const void* header_bytes, size_t header_size);
 
    static void decompressFromFile(
        const std::string&  filename,
        void**              d_output,
        size_t*             output_size,
        cudaStream_t        stream             = 0,
        PipelinePerfResult* perf_out           = nullptr,
        size_t              pool_override_bytes = 0
    );
 
    void decompressFromFileInstance(
        const std::string&  filename,
        void**              d_output,
        size_t*             output_size,
        cudaStream_t        stream   = 0,
        PipelinePerfResult* perf_out = nullptr
    );
 
    void decompressFromMemory(
        const void*  header_bytes,
        size_t       header_size,
        const void*  d_blob,
        size_t       blob_size,
        void**       d_output,
        size_t*      output_size,
        cudaStream_t stream = 0
    );
 
    // ── Config File ───────────────────────────────────────────────────────────
 
    void loadConfig(const std::string& path);
 
    void saveConfig(const std::string& path) const;
 
private:
    // ── RAII buffer wrappers (private implementation detail) ─────────────────
 
    // Pool-allocated persistent device buffer.
    struct PoolBuffer {
        void*       ptr      = nullptr;
        size_t      capacity = 0;
        MemoryPool* pool     = nullptr;
 
        ~PoolBuffer()                         { free(0); }
        PoolBuffer()                          = default;
        PoolBuffer(const PoolBuffer&)         = delete;
        PoolBuffer& operator=(const PoolBuffer&) = delete;
 
        void free(cudaStream_t s) {
            if (ptr && pool) { pool->free(ptr, s); ptr = nullptr; capacity = 0; }
        }
        bool allocate(MemoryPool* p, size_t bytes, cudaStream_t s,
                      const char* tag, bool persistent = false) {
            free(s);
            pool = p;
            ptr  = pool->allocate(bytes, s, tag, persistent);
            if (ptr) capacity = bytes;
            return ptr != nullptr;
        }
    };
 
    // cudaHostAlloc pinned host buffer — grows on demand, never shrinks.
    struct PinnedBuffer {
        void*  ptr      = nullptr;
        size_t capacity = 0;
 
        ~PinnedBuffer()                           { if (ptr) cudaFreeHost(ptr); }
        PinnedBuffer()                            = default;
        PinnedBuffer(const PinnedBuffer&)         = delete;
        PinnedBuffer& operator=(const PinnedBuffer&) = delete;
 
        // Returns false on CUDA allocation failure.
        bool ensureCapacity(size_t bytes) {
            if (capacity >= bytes) return true;
            if (ptr) { cudaFreeHost(ptr); ptr = nullptr; capacity = 0; }
            if (cudaHostAlloc(&ptr, bytes, cudaHostAllocDefault) != cudaSuccess) return false;
            capacity = bytes;
            return true;
        }
    };
 
    // cudaMalloc device buffer — grows on demand, never shrinks.
    struct DeviceBuffer {
        void*  ptr      = nullptr;
        size_t capacity = 0;
 
        ~DeviceBuffer()                           { if (ptr) cudaFree(ptr); }
        DeviceBuffer()                            = default;
        DeviceBuffer(const DeviceBuffer&)         = delete;
        DeviceBuffer& operator=(const DeviceBuffer&) = delete;
 
        // Returns false on CUDA allocation failure.
        bool ensureCapacity(size_t bytes) {
            if (capacity >= bytes) return true;
            if (ptr) { cudaFree(ptr); ptr = nullptr; capacity = 0; }
            if (cudaMalloc(&ptr, bytes) != cudaSuccess) return false;
            capacity = bytes;
            return true;
        }
    };
 
    // ── Internal helpers ──────────────────────────────────────────────────────
 
    Stage* addRawStage(Stage* stage);
 
    struct OutputBuffer {
        void*       d_ptr;
        size_t      actual_size;
        size_t      allocated_size;
        std::string name;
        int         buffer_id;
    };
    std::vector<OutputBuffer> getOutputBuffers() const;
 
    static void* loadCompressedData(
        const std::string&   filename,
        const FZMFileHeader& header,
        cudaStream_t         stream = 0,
        MemoryPool*          pool   = nullptr
    );
 
    void validate();
    std::pair<std::vector<Stage*>, std::vector<Stage*>> identifyTopology();
    void setupInputBuffers(const std::vector<Stage*>& sources);
    int  autoDetectUnconnectedOutputs();
    void detectMultiOutputScenario(int pipeline_outputs);
    void configureStreamsIfNeeded();
 
    // finalize() sub-steps
    void typeCheckConnections();
    void computeInputAlignment();
    void notifyStagesFinalizeHooks();
    void refinePoolSize();
    void setupGraphModeInput();
    void preallocatePadBuffer();
    void preallocateConcatBuffers();
 
    // compress() helper: handles graph-mode copy or alignment padding.
    // Returns the effective source pointer and padded source size.
    std::pair<const void*, size_t> prepareInputSource(
        const void* d_input, size_t input_size, cudaStream_t stream);
 
    void propagateBufferSizes(bool force_from_current_inputs = false);
 
    std::vector<Stage*> getSourceStages() const;
    std::vector<Stage*> getSinkStages() const;
 
    // ── Inverse DAG helpers ───────────────────────────────────────────────────
 
    struct FwdStageDesc {
        Stage*           stage;
        std::vector<int> output_buf_ids;
        std::vector<int> input_buf_ids;
    };
 
    using PipelineOutputMap = std::unordered_map<int, std::pair<void*, size_t>>;
 
    // decompress() helper: builds or reuses the inverse DAG cache.
    void buildOrReuseInvCache(
        const PipelineOutputMap& po_map,
        Stage*       src_stage,
        size_t       src_sz,
        cudaStream_t stream);
 
    void decompressCore(
        const void* d_input,
        size_t      input_size,
        void*       caller_output,
        size_t      caller_capacity,
        bool        synchronize,
        void**      d_output,
        size_t*     output_size,
        cudaStream_t stream);
 
    void buildStaticBufferMetadata();
 
    std::vector<size_t> readConcatSegmentSizes(
        const void* d_blob, size_t n, cudaStream_t stream) const;
 
    // decompressFromFile() helpers.
    static FZMFileHeader parseHeaderFromMemory(const void* data, size_t size);
    static size_t computeFilePoolSize(const FZMFileHeader& fh, size_t pool_override_bytes);
    static std::pair<std::vector<std::unique_ptr<Stage>>, std::vector<FwdStageDesc>>
        reconstructForwardTopology(const FZMFileHeader& fh);
    static std::unordered_map<Stage*, size_t> buildSourceSizesFromHeader(
        const FZMFileHeader& fh, const std::vector<FwdStageDesc>& fwd_topology);
 
    static std::pair<std::unique_ptr<CompressionDAG>,
                     std::unordered_map<Stage*, int>>
    buildInverseDAG(
        const std::vector<FwdStageDesc>&          fwd_stages,
        const PipelineOutputMap&                  pipeline_outputs,
        MemoryPool*                               pool,
        MemoryStrategy                            strategy,
        const std::unordered_map<Stage*, size_t>& source_sizes,
        bool                                      enable_profiling
    );
 
    // ── Concat helpers ────────────────────────────────────────────────────────
 
    struct OutputBufferInfo {
        int         buffer_id;
        void*       d_ptr;
        size_t      actual_size;
        std::string stage_name;
        std::string output_name;
    };
 
    std::vector<OutputBufferInfo> collectOutputBuffers() const;
 
    size_t calculateConcatSize(const std::vector<OutputBufferInfo>& outputs) const;
 
    size_t writeConcatBuffer(
        const std::vector<OutputBufferInfo>& outputs,
        uint8_t*     d_concat_bytes,
        cudaStream_t stream
    ) const;
 
    void concatOutputs(void** d_output, size_t* output_size, cudaStream_t stream);
 
    // ── Member variables ──────────────────────────────────────────────────────
 
    std::unique_ptr<MemoryPool>      mem_pool_;
    std::unique_ptr<CompressionDAG>  dag_;
    MemoryStrategy                   strategy_;
 
    std::vector<std::unique_ptr<Stage>> stages_;
    std::unordered_map<Stage*, DAGNode*> stage_to_node_;
 
    struct ConnectionInfo {
        Stage*      dependent;
        Stage*      producer;
        std::string output_name;
        int         output_index;
    };
    std::vector<ConnectionInfo> connections_;
 
    int  num_streams_;
    bool is_finalized_;
    bool warmup_on_finalize_;
    bool pool_managed_decomp_;
 
    // is_compressed_: true after the first successful compress() (gates writeToFile).
    // was_compressed_: true between compress() and the next reset() (gates captureGraph).
    bool is_compressed_;
    bool was_compressed_;
 
    bool profiling_enabled_;
    PipelinePerfResult last_perf_result_;
 
    std::vector<DAGNode*> input_nodes_;
    std::vector<DAGNode*> output_nodes_;
    std::vector<int>      input_buffer_ids_;
    std::vector<int>      output_buffer_ids_;
 
    PoolBuffer   d_concat_buffer_;
    bool         needs_concat_;
 
    // Pool-persistent decompress output buffers (one per source stage).
    // Only used when pool_managed_decomp_ == true.
    std::vector<void*> d_decomp_outputs_;
 
    // Pinned host buffer for concat header (one H2D copy instead of N).
    PinnedBuffer h_concat_header_;
    // Persistent pinned host + device descriptor buffers for the gather kernel.
    PinnedBuffer h_copy_descs_;
    DeviceBuffer d_copy_descs_;
 
    size_t input_size_;
 
    // Per-source input sizes from the most recent compress(), ordered to match
    // input_nodes_. Used by decompress() to size each inverse result buffer.
    std::vector<size_t> source_input_sizes_;
 
    // Input alignment in bytes — LCM of all stage getRequiredInputAlignment() values.
    // compress() zero-pads to this boundary transparently.
    size_t     input_alignment_bytes_;
    PoolBuffer d_pad_buf_;
 
    // Original (pre-padding) input size. decompress() uses this to trim the
    // reported output back to what the caller provided. 0 when no padding.
    size_t original_input_size_;
 
    size_t input_size_hint_;
    float  pool_multiplier_;
 
    // Dataset dimensions (x=fast, y, z). Pushed to each stage on addStage() and
    // again at finalize(). Default {0,1,1} = 1-D, infer x from input size.
    std::array<size_t, 3> dims_;
 
    struct InvDAGCache {
        std::unique_ptr<CompressionDAG>    inv_dag;
        std::unordered_map<Stage*, int>    inv_result_map;
        std::unordered_map<int, int>       fwd_to_inv_ext_buf;
        std::unordered_map<Stage*, size_t> source_sizes;
    };
    std::unique_ptr<InvDAGCache> inv_cache_;
 
    struct BufferMetadata {
        int         buffer_id;
        size_t      actual_size;
        size_t      allocated_size;
        std::string name;
        DAGNode*    producer;
        int         output_index;
    };
    std::vector<BufferMetadata> buffer_metadata_;
 
    bool graph_mode_enabled_;
    bool graph_captured_;
 
    // Fixed device input buffer whose address is baked into the captured graph.
    // compress() copies user input here before cudaGraphLaunch().
    PoolBuffer d_graph_input_;
    size_t     d_graph_input_size_;
 
    cudaGraph_t     captured_graph_;
    cudaGraphExec_t graph_exec_;
};
 
// ── Template implementation ───────────────────────────────────────────────────
 
template<typename StageT, typename... Args>
StageT* Pipeline::addStage(Args&&... args) {
    if (is_finalized_) {
        throw std::runtime_error("Cannot add stages after finalization");
    }
 
    auto stage_ptr = std::make_unique<StageT>(std::forward<Args>(args)...);
    StageT* stage  = stage_ptr.get();
 
    stage->setDims(dims_);
 
    DAGNode* node        = dag_->addStage(stage, stage->getName());
    size_t   num_outputs = stage->getNumOutputs();
    auto     output_names = stage->getOutputNames();
 
    // Pre-allocate all output slots as unconnected (size=1 placeholder).
    // connect() will promote any that get wired to downstream stages.
    for (size_t i = 0; i < num_outputs; i++) {
        std::string out_name = i < output_names.size() ? output_names[i] : std::to_string(i);
        dag_->addUnconnectedOutput(node, 1, i, stage->getName() + "." + out_name + "_unconnected");
    }
 
    stage_to_node_[stage] = node;
    stages_.push_back(std::move(stage_ptr));
    return stage;
}
 
} // namespace fz