lorenzo_quant.h

Go to the documentation of this file.

 
#pragma once
 
#include "stage/stage.h"
#include "fzm_format.h"
#include "backend/types.h"
#include "log.h"
#include "fused/common/data_type_of.h"
 
#include <array>
#include <cstdint>
#include <memory>
#include <cmath>
#include <cstring>
 
namespace fz {
 
 
enum class ErrorBoundMode : uint8_t {
    ABS  = 0, 
    REL  = 1, 
    NOA  = 2, 
    PREL = 3, 
};
 
inline ErrorBoundMode resolveApproxRelMode(ErrorBoundMode mode, const char* stage_name) {
    if (mode != ErrorBoundMode::REL) return mode;
    FZ_LOG(WARN,
        "%s: ErrorBoundMode::REL is deprecated for this stage and has been "
        "mapped to PREL (abs_eb = eb * max(|data|)). This does NOT guarantee a "
        "per-element relative bound. Use PREL explicitly to silence this, or "
        "QuantizerStage with REL for an exact point-wise bound.",
        stage_name);
    return ErrorBoundMode::PREL;
}
 
struct LorenzoQuantConfig {
    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  centering;     
    uint8_t  reserved[2];   
    double   error_bound_f64;
    double   value_base_f64;  
 
    // Total: 60 bytes (fits easily in 128B stage_config)
 
    LorenzoQuantConfig()
        : 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), centering(0),
          reserved{0, 0}, error_bound_f64(0.0), value_base_f64(0.0) {}
};
static_assert(sizeof(LorenzoQuantConfig) <= FZM_STAGE_CONFIG_SIZE, "LorenzoQuantConfig must fit in FZM_STAGE_CONFIG_SIZE");
 
template<typename TInput = float, typename TCode = uint16_t>
class LorenzoQuantStage : 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;
        bool centering = 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 LorenzoQuantStage(const Config& config = Config());
    ~LorenzoQuantStage() override;
 
    void execute(
        fz::stream_t stream,
        MemoryPool* pool,
        const std::vector<void*>& inputs,
        const std::vector<void*>& outputs,
        const std::vector<size_t>& sizes
    ) override;
 
    void postStreamSync(fz::stream_t stream) override;
 
    void onFinalize(size_t estimated_inlen, MemoryPool* pool) override;
 
    size_t estimateDeviceFootprintBytes(size_t /*estimated_inlen*/) const override {
        return sizeof(uint32_t);
    }
 
    std::string getName() const override { return "LorenzoQuant"; }
    size_t getNumInputs()  const override {
        return is_inverse_ ? (config_.centering ? 4 : 3) : 1;
    }
    size_t getNumOutputs() const override {
        return is_inverse_ ? 1 : (config_.centering ? 4 : 3);
    }
 
    std::vector<std::string> getOutputNames() const override {
        if (config_.centering)
            return {"codes", "outlier_errors", "outlier_indices", "means"};
        return {"codes", "outlier_errors", "outlier_indices"};
    }
    
    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 = resolveApproxRelMode(mode, "LorenzoQuantStage");
    }
    // Provide a pre-computed value_range (NOA) or max(|data|) (PREL) 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 setCentering(bool enable) { config_.centering = 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; }
    TInput getComputedAbsErrorBound() const { return computed_abs_eb_; }
    float getValueBase() const { return config_.precomputed_value_base; }
    bool  getZigzagCodes() const { return config_.zigzag_codes; }
    bool  getCentering() const { return config_.centering; }
 
    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 {
        return static_cast<uint16_t>(StageType::LORENZO_QUANT);
    }
    
    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
            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(LorenzoQuantConfig)) {
            throw std::runtime_error("Insufficient buffer for Lorenzo config");
        }
 
        LorenzoQuantConfig 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    = static_cast<float>(computed_value_base_);
        config.error_bound_f64 = static_cast<double>(computed_abs_eb_);
        config.value_base_f64  = static_cast<double>(computed_value_base_);
        config.zigzag_codes  = config_.zigzag_codes ? uint8_t{1} : uint8_t{0};
        config.centering     = config_.centering ? uint8_t{1} : uint8_t{0};
        config.reserved[0]   = 0; config.reserved[1] = 0;
 
        std::memcpy(header_buffer, &config, sizeof(LorenzoQuantConfig));
        return sizeof(LorenzoQuantConfig);
    }
    
    size_t getMaxHeaderSize(size_t output_index) const override {
        (void)output_index;
        return sizeof(LorenzoQuantConfig);
    }
    
    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");
        }
 
        LorenzoQuantConfig config;
        std::memcpy(&config, header_buffer, std::min(size, sizeof(LorenzoQuantConfig)));
 
        // error_bound in the header is always the absolute bound used at compression.
        // Prefer the full-precision copy; pre-2026-08-08 headers leave it 0.
        constexpr size_t kSizeBeforeF64 = 44;
        const bool has_f64 = (size > kSizeBeforeF64 && config.error_bound_f64 != 0.0);
        config_.error_bound  = config.error_bound;
        computed_abs_eb_     = has_f64 ? static_cast<TInput>(config.error_bound_f64)
                                       : 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) {
            // Files written before the PREL split stored eb_mode==REL for what
            // was always the approximate mode. Map it silently (no warning —
            // decode uses the stored absolute error_bound regardless of mode).
            auto stored = static_cast<ErrorBoundMode>(config.eb_mode);
            config_.eb_mode                = (stored == ErrorBoundMode::REL)
                                             ? ErrorBoundMode::PREL : stored;
            config_.precomputed_value_base = config.value_base;
            computed_value_base_           = has_f64
                                             ? static_cast<TInput>(config.value_base_f64)
                                             : static_cast<TInput>(config.value_base);
        } else {
            config_.eb_mode                = ErrorBoundMode::ABS;
            config_.precomputed_value_base = 0.0f;
            computed_value_base_           = static_cast<TInput>(0);
        }
        // zigzag_codes field added in v2 (≥44B). Compared against the literal v2
        // size, not sizeof(LorenzoQuantConfig) — the struct has grown since (the
        // f64 bound fields), and keying off sizeof would stop reading these
        // fields from every v2 archive.
        constexpr size_t kV2Size = 44;
        if (size >= kV2Size) {
            config_.zigzag_codes = (config.zigzag_codes != 0);
            // `centering` reuses a byte older writers zeroed as `reserved`, so
            // pre-centering archives decode as centering-off.
            config_.centering    = (config.centering != 0);
        } else {
            config_.zigzag_codes = false;
            config_.centering    = 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;
    TInput computed_value_base_ = static_cast<TInput>(0);
    uint32_t* d_outlier_count_scratch_ = nullptr;
    MemoryPool* persistent_pool_ = nullptr;
    std::weak_ptr<const void> persistent_pool_alive_;
 
 
    void initOutlierCountScratch(MemoryPool* pool);
    
    DataType getInputDataType() const { return fused::dataTypeOf<TInput>(); }
    DataType getCodeDataType()  const { return fused::dataTypeOf<TCode>(); }
    
    size_t getMaxOutlierCount(size_t num_elements) const {
        return static_cast<size_t>(std::ceil(num_elements * config_.outlier_capacity));
    }
};
 
extern template class LorenzoQuantStage<float, uint16_t>;
extern template class LorenzoQuantStage<float, uint8_t>;
extern template class LorenzoQuantStage<double, uint16_t>;
extern template class LorenzoQuantStage<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,
    fz::stream_t stream,
    // Per-tile means output (one per 1024-element tile), or nullptr to disable
    // adaptive centering.
    TInput* d_means = nullptr
);
 
template<typename TInput, typename TCode>
void launchLorenzoInverseKernel(
    const TCode* d_codes,
    const TInput* d_outlier_errors, const uint32_t* d_outlier_indices,
    uint32_t outlier_n,
    size_t n,
    TInput ebx2, TCode quant_radius,
    TInput* d_output,
    bool zigzag_codes,
    fz::stream_t stream, MemoryPool* pool,
    // Per-tile means from the forward pass, or nullptr if centering is off.
    const TInput* d_means = nullptr
);
 
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,
    fz::stream_t stream
);
 
template<typename TInput, typename TCode>
void launchLorenzoInverseKernel2D(
    const TCode* d_codes,
    const TInput* d_outlier_errors, const uint32_t* d_outlier_indices,
    uint32_t outlier_n,
    size_t nx, size_t ny,
    TInput ebx2, TCode quant_radius,
    TInput* d_output,
    bool zigzag_codes,
    fz::stream_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,
    fz::stream_t stream
);
 
template<typename TInput, typename TCode>
void launchLorenzoInverseKernel3D(
    const TCode* d_codes,
    const TInput* d_outlier_errors, const uint32_t* d_outlier_indices,
    uint32_t outlier_n,
    size_t nx, size_t ny, size_t nz,
    TInput ebx2, TCode quant_radius,
    TInput* d_output,
    bool zigzag_codes,
    fz::stream_t stream, MemoryPool* pool
);
 
} // namespace fz