stage.h

#pragma once
 
#include "fzm_format.h"
#include <array>
#include <cuda_runtime.h>
#include <cstdint>
#include <stdexcept>
#include <string>
#include <unordered_map>
#include <vector>
 
namespace fz {
 
// Forward declaration — avoids requiring mempool.h in every stage header
class MemoryPool;
 
class Stage {
public:
    virtual ~Stage() = default;
 
    virtual void execute(
        cudaStream_t stream,
        MemoryPool* pool,
        const std::vector<void*>& inputs,
        const std::vector<void*>& outputs,
        const std::vector<size_t>& sizes
    ) = 0;
 
    virtual std::string getName() const = 0;
 
    virtual size_t getNumInputs()  const = 0;
    virtual size_t getNumOutputs() const = 0;
 
    virtual size_t getRequiredInputAlignment() const { return 1; }
 
    virtual std::vector<std::string> getOutputNames() const {
        return {"output"};
    }
 
    int getOutputIndex(const std::string& name) const {
        auto names = getOutputNames();
        for (size_t i = 0; i < names.size(); i++) {
            if (names[i] == name) return static_cast<int>(i);
        }
        return -1;
    }
 
    virtual std::vector<size_t> estimateOutputSizes(
        const std::vector<size_t>& input_sizes
    ) const = 0;
 
    virtual std::unordered_map<std::string, size_t> getActualOutputSizesByName() const = 0;
 
    virtual size_t getActualOutputSize(int index) const {
        auto names = getOutputNames();
        if (index < 0 || index >= static_cast<int>(names.size())) return 0;
        auto m  = getActualOutputSizesByName();
        auto it = m.find(names[index]);
        return (it != m.end()) ? it->second : 0;
    }
 
    virtual void setInverse(bool inverse) { (void)inverse; }
    virtual bool isInverse() const { return false; }
 
    virtual uint16_t getStageTypeId() const = 0;
 
    virtual uint8_t getOutputDataType(size_t output_index) const = 0;
 
    virtual uint8_t getInputDataType(size_t /*input_index*/) const {
        return static_cast<uint8_t>(DataType::UNKNOWN);
    }
 
    virtual size_t serializeHeader(size_t output_index, uint8_t* header_buffer, size_t max_size) const {
        (void)output_index; (void)header_buffer; (void)max_size;
        return 0;
    }
 
    virtual void deserializeHeader(const uint8_t* header_buffer, size_t size) {
        (void)header_buffer; (void)size;
    }
 
    virtual void saveState()    {}
    virtual void restoreState() {}
 
    virtual void setDims(const std::array<size_t, 3>& dims) { (void)dims; }
 
    virtual void postStreamSync(cudaStream_t stream) { (void)stream; }
 
    virtual size_t getMaxHeaderSize(size_t output_index) const {
        (void)output_index;
        return 0;
    }
 
    virtual bool isGraphCompatible() const { return true; }
 
    virtual size_t estimateScratchBytes(const std::vector<size_t>& input_sizes) const {
        (void)input_sizes;
        return 0;
    }
};
 
} // namespace fz