lorenzo.h

#pragma once
 
#include "stage/stage.h"
#include "fzm_format.h"
#include <cuda_runtime.h>
#include <array>
#include <cstdint>
#include <cmath>
#include <cstring>
 
namespace fz {
 
 
enum class ErrorBoundMode : uint8_t {
    ABS = 0, 
    REL = 1, 
    NOA = 2, 
};
 
struct LorenzoConfig {
    float    error_bound;   
    uint32_t quant_radius;  
    uint32_t num_elements;  
    uint32_t outlier_count; 
    DataType input_type;    
    DataType code_type;     
    uint8_t  ndim;          
    uint8_t  eb_mode;       
    uint32_t dim_x;         
    uint32_t dim_y;         
    uint32_t dim_z;         
    float    user_eb;       
    float    value_base;    
    uint8_t  zigzag_codes;  
    uint8_t  reserved[3];   
 
    // Total: 44 bytes (fits easily in 128B stage_config)
 
    LorenzoConfig()
        : error_bound(0.0f), quant_radius(0), num_elements(0), outlier_count(0),
          input_type(DataType::FLOAT32), code_type(DataType::UINT16),
          ndim(1), eb_mode(0), dim_x(0), dim_y(1), dim_z(1),
          user_eb(0.0f), value_base(0.0f), zigzag_codes(0), reserved{0, 0, 0} {}
};
static_assert(sizeof(LorenzoConfig) <= FZM_STAGE_CONFIG_SIZE, "LorenzoConfig must fit in FZM_STAGE_CONFIG_SIZE");
 
template<typename TInput = float, typename TCode = uint16_t>
class LorenzoStage : public Stage {
public:
    struct Config {
        float  error_bound       = 1e-3;    
        int    quant_radius      = 32768;   
        float  outlier_capacity  = 0.2f;    
        std::array<size_t, 3> dims = {0, 1, 1};
        ErrorBoundMode eb_mode = ErrorBoundMode::ABS;
        float precomputed_value_base = 0.0f;
        bool zigzag_codes = false;
        Config() = default;
        Config(TInput eb, TCode radius = 32768, float outlier_cap = 0.2f,
               std::array<size_t, 3> d = {0, 1, 1})
            : error_bound(eb), quant_radius(radius), outlier_capacity(outlier_cap),
              dims(d) {}
    };
    
    explicit LorenzoStage(const Config& config = Config());
 
    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;
    
    std::string getName() const override {
        switch (ndim()) {
            case 2:  return "Lorenzo2D";
            case 3:  return "Lorenzo3D";
            default: return "Lorenzo1D";
        }
    }
    size_t getNumInputs()  const override { return is_inverse_ ? 4 : 1; }
    size_t getNumOutputs() const override { return is_inverse_ ? 1 : 4; }
    
    std::vector<std::string> getOutputNames() const override {
        return {"codes", "outlier_errors", "outlier_indices", "outlier_count"};
    }
    
    std::vector<size_t> estimateOutputSizes(
        const std::vector<size_t>& input_sizes
    ) const override;
    
    std::unordered_map<std::string, size_t> getActualOutputSizesByName() const override {
        auto names = getOutputNames();
        std::unordered_map<std::string, size_t> result;
        for (size_t i = 0; i < names.size() && i < actual_output_sizes_.size(); i++) {
            result[names[i]] = actual_output_sizes_[i];
        }
        return result;
    }
    size_t getActualOutputSize(int index) const override {
        return (index >= 0 && index < static_cast<int>(actual_output_sizes_.size()))
            ? actual_output_sizes_[index] : 0;
    }
 
    // Preserve the forward-mode actual_output_sizes_ across decompression passes.
    // decompressMulti() calls saveState()/restoreState() around each inverse
    // execute() to prevent the inverse pass from permanently corrupting the
    // 4-element forward output-size vector (inverse sets it to a 1-element vector).
    void saveState()    override { saved_output_sizes_ = actual_output_sizes_; }
    void restoreState() override { actual_output_sizes_ = saved_output_sizes_; }
 
    // Configuration accessors
    void setErrorBound(TInput error_bound) { config_.error_bound = error_bound; }
    void setQuantRadius(TCode radius) { config_.quant_radius = radius; }
    void setOutlierCapacity(float capacity) { config_.outlier_capacity = capacity; }
    void setDims(const std::array<size_t, 3>& dims) override { config_.dims = dims; }
    void setErrorBoundMode(ErrorBoundMode mode) { config_.eb_mode = mode; }
    // Provide a pre-computed value_range (NOA) or max(|data|) (REL) to skip
    // the internal data scan during execute().  Pass 0 to re-enable auto-scan.
    void setValueBase(float value_base) { config_.precomputed_value_base = value_base; }
    void setZigzagCodes(bool enable) { config_.zigzag_codes = enable; }
    void setDims(size_t x, size_t y = 1, size_t z = 1) { config_.dims = {x, y, z}; }
 
    TInput getErrorBound() const { return config_.error_bound; }
    TCode  getQuantRadius() const { return config_.quant_radius; }
    float  getOutlierCapacity() const { return config_.outlier_capacity; }
    std::array<size_t, 3> getDims() const { return config_.dims; }
    ErrorBoundMode getErrorBoundMode() const { return config_.eb_mode; }
    float getValueBase() const { return config_.precomputed_value_base; }
    bool  getZigzagCodes() const { return config_.zigzag_codes; }
 
    int ndim() const {
        if (config_.dims[2] > 1) return 3;
        if (config_.dims[1] > 1) return 2;
        return 1;
    }
    
    void setInverse(bool inverse) { is_inverse_ = inverse; }
    bool isInverse() const { return is_inverse_; }
 
    // ── Serialization ─────────────────────────────────────────────────────────
    
    uint16_t getStageTypeId() const override {
        switch (ndim()) {
            case 2:  return static_cast<uint16_t>(StageType::LORENZO_2D);
            case 3:  return static_cast<uint16_t>(StageType::LORENZO_3D);
            default: return static_cast<uint16_t>(StageType::LORENZO_1D);
        }
    }
    
    uint8_t getOutputDataType(size_t output_index) const override {
        switch (output_index) {
            case 0: return static_cast<uint8_t>(getCodeDataType());      // codes
            case 1: return static_cast<uint8_t>(getInputDataType());     // outlier_errors
            case 2: return static_cast<uint8_t>(DataType::UINT32);       // outlier_indices
            case 3: return static_cast<uint8_t>(DataType::UINT32);       // outlier_count
            default: return static_cast<uint8_t>(DataType::UINT8);
        }
    }
 
    uint8_t getInputDataType(size_t /*input_index*/) const override {
        return static_cast<uint8_t>(getInputDataType());
    }
    
    size_t serializeHeader(size_t output_index, uint8_t* header_buffer, size_t max_size) const override {
        (void)output_index;  // Lorenzo uses same header for all outputs
 
        if (max_size < sizeof(LorenzoConfig)) {
            throw std::runtime_error("Insufficient buffer for Lorenzo config");
        }
 
        LorenzoConfig config;
        config.error_bound   = static_cast<float>(computed_abs_eb_);  // abs bound used by decompressor
        config.quant_radius  = static_cast<uint32_t>(config_.quant_radius);
        config.num_elements  = static_cast<uint32_t>(num_elements_);
        config.outlier_count = actual_outlier_count_;
        config.input_type    = getInputDataType();
        config.code_type     = getCodeDataType();
        config.ndim          = static_cast<uint8_t>(ndim());
        config.eb_mode       = static_cast<uint8_t>(config_.eb_mode);
        config.dim_x         = static_cast<uint32_t>(config_.dims[0]);
        config.dim_y         = static_cast<uint32_t>(config_.dims[1]);
        config.dim_z         = static_cast<uint32_t>(config_.dims[2]);
        config.user_eb       = static_cast<float>(config_.error_bound);  // original user-specified value
        config.value_base    = computed_value_base_;
        config.zigzag_codes  = config_.zigzag_codes ? uint8_t{1} : uint8_t{0};
        config.reserved[0]   = 0; config.reserved[1] = 0; config.reserved[2] = 0;
 
        std::memcpy(header_buffer, &config, sizeof(LorenzoConfig));
        return sizeof(LorenzoConfig);
    }
    
    size_t getMaxHeaderSize(size_t output_index) const override {
        (void)output_index;
        return sizeof(LorenzoConfig);
    }
    
    void deserializeHeader(const uint8_t* header_buffer, size_t size) override {
        // Minimum size is the original 32-byte layout (before user_eb/value_base were added).
        constexpr size_t kLegacySize = 32;
        if (size < kLegacySize) {
            throw std::runtime_error("Invalid Lorenzo config size");
        }
 
        LorenzoConfig config;
        std::memcpy(&config, header_buffer, std::min(size, sizeof(LorenzoConfig)));
 
        // error_bound in the header is always the absolute bound used at compression.
        config_.error_bound  = config.error_bound;
        computed_abs_eb_     = static_cast<TInput>(config.error_bound);
        config_.quant_radius = static_cast<TCode>(config.quant_radius);
        num_elements_        = config.num_elements;
        actual_outlier_count_= config.outlier_count;
        // New fields: present only in headers written by v1+ (≥40B, added user_eb/value_base/eb_mode).
        constexpr size_t kV1Size = 40;
        if (size >= kV1Size) {
            config_.eb_mode                = static_cast<ErrorBoundMode>(config.eb_mode);
            config_.precomputed_value_base = config.value_base;
            computed_value_base_           = config.value_base;
        } else {
            config_.eb_mode                = ErrorBoundMode::ABS;
            config_.precomputed_value_base = 0.0f;
            computed_value_base_           = 0.0f;
        }
        // zigzag_codes field added in v2 (≥44B).
        if (size >= sizeof(LorenzoConfig)) {
            config_.zigzag_codes = (config.zigzag_codes != 0);
        } else {
            config_.zigzag_codes = false;
        }
 
        // Restore spatial dimensions; handle old (pre-dims) files gracefully
        int eff_ndim = (config.ndim == 0) ? 1 : static_cast<int>(config.ndim);
        // dim_x: stored explicitly; fall back to derivation for old files
        if (config.dim_x > 0) {
            config_.dims[0] = config.dim_x;
        } else if (config.num_elements > 0) {
            size_t yz = std::max<size_t>(1, config.dim_y) * std::max<size_t>(1, config.dim_z);
            config_.dims[0] = config.num_elements / yz;
        } else {
            config_.dims[0] = 0;
        }
        if (eff_ndim >= 2) {
            config_.dims[1] = (config.dim_y > 0) ? config.dim_y : 1;
        } else {
            config_.dims[1] = 1;
        }
        if (eff_ndim >= 3) {
            config_.dims[2] = (config.dim_z > 0) ? config.dim_z : 1;
        } else {
            config_.dims[2] = 1;
        }
    }
    
private:
    Config config_;
    std::vector<size_t> actual_output_sizes_;
    std::vector<size_t> saved_output_sizes_;  // saved by saveState(), restored by restoreState()
    size_t num_elements_ = 0;              // Track for header
    uint32_t actual_outlier_count_ = 0;    // Track for header
    bool is_inverse_ = false;              // false = compress, true = decompress
    TInput computed_abs_eb_ = 0;
    float computed_value_base_ = 0.0f;
    const void* d_outlier_count_ptr_ = nullptr;
    
    DataType getInputDataType() const {
        if (std::is_same<TInput, float>::value) return DataType::FLOAT32;
        if (std::is_same<TInput, double>::value) return DataType::FLOAT64;
        return DataType::FLOAT32;
    }
    
    DataType getCodeDataType() const {
        if (std::is_same<TCode, uint8_t>::value) return DataType::UINT8;
        if (std::is_same<TCode, uint16_t>::value) return DataType::UINT16;
        if (std::is_same<TCode, uint32_t>::value) return DataType::UINT32;
        return DataType::UINT16;
    }
    
    size_t getMaxOutlierCount(size_t num_elements) const {
        return static_cast<size_t>(std::ceil(num_elements * config_.outlier_capacity));
    }
};
 
extern template class LorenzoStage<float, uint16_t>;
extern template class LorenzoStage<float, uint8_t>;
extern template class LorenzoStage<double, uint16_t>;
extern template class LorenzoStage<double, uint32_t>;
 
// Kernel launcher declarations — defined in lorenzo.cu.
 
template<typename TInput, typename TCode>
void launchLorenzoKernel(
    const TInput* d_input, size_t n,
    TInput ebx2_r, TCode quant_radius,
    TCode* d_codes, TInput* d_outlier_errors,
    uint32_t* d_outlier_indices, uint32_t* d_outlier_count,
    size_t max_outliers, int grid_size,
    bool zigzag_codes,
    cudaStream_t stream
);
 
template<typename TInput, typename TCode>
void launchLorenzoInverseKernel(
    const TCode* d_codes,
    const TInput* d_outlier_errors, const uint32_t* d_outlier_indices,
    const uint32_t* d_outlier_count,
    size_t n, size_t max_outliers,
    TInput ebx2, TCode quant_radius,
    TInput* d_output,
    bool zigzag_codes,
    cudaStream_t stream, MemoryPool* pool
);
 
template<typename TInput, typename TCode>
void launchLorenzoKernel2D(
    const TInput* d_input, size_t nx, size_t ny,
    TInput ebx2_r, TCode quant_radius,
    TCode* d_codes, TInput* d_outlier_errors,
    uint32_t* d_outlier_indices, uint32_t* d_outlier_count,
    size_t max_outliers,
    bool zigzag_codes,
    cudaStream_t stream
);
 
template<typename TInput, typename TCode>
void launchLorenzoInverseKernel2D(
    const TCode* d_codes,
    const TInput* d_outlier_errors, const uint32_t* d_outlier_indices,
    const uint32_t* d_outlier_count,
    size_t nx, size_t ny, size_t max_outliers,
    TInput ebx2, TCode quant_radius,
    TInput* d_output,
    bool zigzag_codes,
    cudaStream_t stream, MemoryPool* pool
);
 
template<typename TInput, typename TCode>
void launchLorenzoKernel3D(
    const TInput* d_input, size_t nx, size_t ny, size_t nz,
    TInput ebx2_r, TCode quant_radius,
    TCode* d_codes, TInput* d_outlier_errors,
    uint32_t* d_outlier_indices, uint32_t* d_outlier_count,
    size_t max_outliers,
    bool zigzag_codes,
    cudaStream_t stream
);
 
template<typename TInput, typename TCode>
void launchLorenzoInverseKernel3D(
    const TCode* d_codes,
    const TInput* d_outlier_errors, const uint32_t* d_outlier_indices,
    const uint32_t* d_outlier_count,
    size_t nx, size_t ny, size_t nz, size_t max_outliers,
    TInput ebx2, TCode quant_radius,
    TInput* d_output,
    bool zigzag_codes,
    cudaStream_t stream, MemoryPool* pool
);
 
} // namespace fz