compressor.h

Go to the documentation of this file.

 
#pragma once
 
#include "backend/types.h"
#include "advanced/dag.h"
#include "pipeline/device_buffer.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 <string>
#include <unordered_map>
#include <vector>
 
namespace fz {
 
enum class SpecializationPolicy { Off, Auto, Force };
 
using FusionPolicy = SpecializationPolicy;
 
struct SpecializationGroupInfo {
    std::string implementation;              
    std::vector<std::string> stages;         
    std::string execution_path;              
};
using FusionGroupInfo = SpecializationGroupInfo;   
 
struct SpecializationInfo {
    SpecializationPolicy policy = SpecializationPolicy::Off;
    size_t legal_group_count = 0;
    std::vector<SpecializationGroupInfo> installed_groups;
    std::vector<SpecializationGroupInfo> installed_inverse_groups;
    std::string fallback_reason;
};
using FusionInfo = SpecializationInfo;   
 
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 setSpecializationPolicy(SpecializationPolicy mode) { fusion_policy_ = mode; }
    SpecializationPolicy getSpecializationPolicy() const { return fusion_policy_; }
    size_t getSpecializedGroupCount() const { return dag_ ? dag_->getFusedGroupCount() : 0; }
    const SpecializationInfo& getSpecializationInfo() const { return fusion_info_; }
 
    void setFusionPolicy(SpecializationPolicy mode) { setSpecializationPolicy(mode); }
    SpecializationPolicy getFusionPolicy() const { return getSpecializationPolicy(); }
    size_t getFusedGroupCount() const { return getSpecializedGroupCount(); }
    const SpecializationInfo& getFusionInfo() const { return fusion_info_; }
 
    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 bindExternalInput(Stage* stage);
 
    void finalize();
 
    void warmup(fz::stream_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_; }
 
    void setPrimarySource(Stage* stage) { primary_source_stage_ = stage; }
 
    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,
        fz::stream_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,
        fz::stream_t stream = 0
    );
 
    void decompress(
        const void* d_input,
        size_t      input_size,
        void**      d_output,
        size_t*     output_size,
        fz::stream_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,
        fz::stream_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,
        fz::stream_t stream = 0
    );
 
    // ── Explicit-ownership execution API ──────────────────────────────────────
    //
    // Span-based wrappers over the pointer overloads above. Behavior is
    // identical — these exist so ownership is visible in the signature instead
    // of depending on `void**` vs `void*` and on setPoolManagedDecompOutput().
    // Prefer these in new code; the pointer overloads remain supported.
 
    BorrowedDeviceBuffer compress(ConstDeviceSpan input, fz::stream_t stream = 0);
 
    size_t compressInto(ConstDeviceSpan input, DeviceSpan output, fz::stream_t stream = 0);
 
    BorrowedDeviceBuffer decompressBorrowed(ConstDeviceSpan input, fz::stream_t stream = 0);
 
    OwnedDeviceBuffer decompressOwned(ConstDeviceSpan input, fz::stream_t stream = 0);
 
    size_t decompressInto(ConstDeviceSpan input, DeviceSpan output, fz::stream_t stream = 0);
 
    size_t decompressIntoAsync(ConstDeviceSpan input, DeviceSpan output, fz::stream_t stream = 0);
 
    void prepareInverse(size_t uncompressed_size);
 
    void reset(fz::stream_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) {
        coloring_enabled_ = enable;              // survives a setMemoryStrategy() DAG swap
        dag_->setColoringEnabled(enable);
    }
    bool isColoringEnabled() const       { return dag_->isColoringEnabled(); }
    size_t getColorRegionCount() const   { return dag_->getColorRegionCount(); }
 
    bool isColoringRequested() const     { return coloring_enabled_; }
 
    std::unordered_map<std::string, std::vector<std::string>> collectRunNotes() const {
        std::unordered_map<std::string, std::vector<std::string>> notes;
        for (const auto& s : stages_) {
            if (!s) continue;
            auto n = s->getRunNotes();
            if (!n.empty()) notes.emplace(s->getName(), std::move(n));
        }
        return notes;
    }
 
    // ── CUDA Graph Capture (compression-only) ─────────────────────────────────
 
    void enableGraphMode(bool enable);
    bool isGraphModeEnabled() const { return graph_mode_enabled_; }
 
    void captureGraph(fz::stream_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, fz::stream_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,
        fz::stream_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,
        fz::stream_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,
        fz::stream_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(fz::stream_t s) {
            if (ptr && pool) { pool->free(ptr, s); ptr = nullptr; capacity = 0; }
        }
        bool allocate(MemoryPool* p, size_t bytes, fz::stream_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,
        fz::stream_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 bindSemanticContracts();
    void planAndInstallFusion();
    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, fz::stream_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,
        fz::stream_t stream);
 
    std::vector<size_t> resolveDecompressSegmentSizes(
        const void* d_input, size_t input_size, fz::stream_t stream) const;
 
    PipelineOutputMap mapCompressedBufferPointers(
        const void* d_input, const std::vector<size_t>& seg_sizes) const;
 
    std::pair<Stage*, size_t> resolvePrimarySource() const;
 
    void* allocateDecompressOutput(
        size_t actual_size, void* caller_output, size_t caller_capacity,
        fz::stream_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,
        fz::stream_t stream);
 
    void buildStaticBufferMetadata();
 
    std::vector<size_t> readConcatSegmentSizes(
        const void* d_blob, size_t n, fz::stream_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,
        bool                                      enable_inverse_fusion = false,
        bool                                      caller_supplies_results = false
    );
 
    // ── 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,
        fz::stream_t stream
    ) const;
 
    void concatOutputs(void** d_output, size_t* output_size, fz::stream_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_;
    bool coloring_enabled_ = true;
    PipelinePerfResult last_perf_result_;
 
    std::vector<DAGNode*> input_nodes_;
    // Explicit choice of which input_nodes_ entry decompress() returns, for
    // pipelines with more than one source stage. Null = default to
    // input_nodes_[0]. Set via setPrimarySource(); resolved to an index lazily
    // (source stages aren't wired into input_nodes_ until finalize()).
    Stage* primary_source_stage_ = nullptr;
    // (node, buffer_id) pairs registered by bindExternalInput(), captured
    // immediately (before any later connect() calls add more buffer ids to
    // the same node). setupInputBuffers() appends these to input_nodes_/
    // input_buffer_ids_ after its own auto-discovery pass, since it clears
    // both first and auto-discovery alone would miss any stage that also
    // has real dependencies (getSourceStages() requires dependencies.empty()).
    std::vector<std::pair<DAGNode*, int>> explicit_external_bindings_;
    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_;
    bool       fused_source_unpadded_ = false;
 
    // 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::vector<size_t>                fwd_size_signature;
    };
    std::unique_ptr<InvDAGCache> inv_cache_;
    std::vector<size_t> forwardSizeSignature() const;
 
    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_;
    FusionPolicy fusion_policy_ = FusionPolicy::Off;
    FusionInfo fusion_info_;
 
    // 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_;
 
    fz::graph_t      captured_graph_;
    fz::graph_exec_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");
    }
 
    // if constexpr, not a runtime check: on an unsupported backend, StageT's
    // constructor may not exist in the build at all (its .cu translation
    // unit excluded — see Stage::isSupportedOnBackend()'s doc comment), so
    // every line below that references `new StageT()` must never be
    // instantiated at all, not merely never executed — hence the whole rest
    // of the function lives in the `if constexpr` branch rather than after
    // a standalone early-throw.
    if constexpr (!StageT::isSupportedOnBackend()) {
        throw std::runtime_error(
            "addStage(): this stage type is not supported on the current "
            "GPU backend (FZGMOD_BACKEND) this library was built for");
    } else {
        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