compressor.h

#pragma once
 
#include "pipeline/dag.h"
#include "pipeline/perf.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
    );
 
    ~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_; }
 
    // ── Execution ─────────────────────────────────────────────────────────────
 
    struct InputSpec {
        Stage*      source;
        const void* d_data;
        size_t      size;
    };
 
    void compress(
        const void* d_input,
        size_t      input_size,
        void**      d_output,
        size_t*     output_size,
        cudaStream_t stream = 0
    );
 
    void compress(
        const std::vector<InputSpec>& inputs,
        void**  d_output,
        size_t* output_size,
        cudaStream_t stream = 0
    );
 
    void setInputSizeHint(Stage* source, size_t size) {
        per_source_hints_[source] = size;
    }
 
    void decompress(
        const void* d_input,
        size_t      input_size,
        void**      d_output,
        size_t*     output_size,
        cudaStream_t stream = 0
    );
 
    std::vector<std::pair<void*, size_t>> decompressMulti(
        const void* d_input    = nullptr,
        size_t      input_size = 0,
        cudaStream_t stream    = 0
    );
 
    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;
 
    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;
 
    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
    );
 
private:
    // ── 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();
 
    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>>;
 
    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_;
 
    void*  d_concat_buffer_;
    size_t concat_buffer_capacity_;
    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. Allocated once at finalize() to
    // avoid N+1 tiny H2D copies — CPU packs the header then issues one H2D copy.
    void*  h_concat_header_;
    size_t h_concat_header_capacity_;
 
    // Persistent pinned host + device descriptor buffers for the gather kernel.
    // Sized at finalize() for max_outputs; reused every compress call.
    void*  h_copy_descs_;
    void*  d_copy_descs_;
    size_t copy_descs_capacity_;
 
    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,
    // computed at finalize(). compress() zero-pads to this boundary transparently.
    size_t input_alignment_bytes_;
    void*  d_pad_buf_;
    size_t d_pad_buf_size_;
 
    // 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_;
 
    // Per-source size hints (set via setInputSizeHint()). Override input_size_hint_
    // for the matching source during propagateBufferSizes().
    std::unordered_map<Stage*, size_t> per_source_hints_;
 
    // Dataset dimensions (x=fast, y, z). Used by convenience.h addLorenzo() to
    // select 1-D/2-D/3-D automatically. Default {0,1,1} = 1-D, infer x from input.
    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().
    void*  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();
 
    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