quantizer.h

Go to the documentation of this file.

#pragma once
 
#include "stage/stage.h"
#include "fzm_format.h"
#include "fused/lorenzo_quant/lorenzo_quant.h"  // for ErrorBoundMode
#include "backend/types.h"
#include <array>
#include <cmath>
#include <cstdint>
#include <memory>
#include <cstring>
#include <limits>
 
namespace fz {
 
struct QuantizerConfig {
    float    abs_error_bound;   
    float    user_error_bound;  
    float    value_base;        
    uint32_t quant_radius;      
    uint32_t num_elements;      
    uint32_t outlier_count;     
    DataType input_type;        
    DataType code_type;         
    uint8_t  eb_mode;           
    uint8_t  zigzag_codes;      
    float    outlier_threshold; 
    uint8_t  inplace_outliers;  
    uint8_t  linear_mode;       
    uint8_t  dither;            
    uint8_t  _pad[5];           
    uint64_t dither_seed;       
    float    dither_strength;   
    uint32_t _pad2;             
    double   abs_error_bound_f64;
    double   value_base_f64;    
 
    QuantizerConfig()
        : abs_error_bound(0.0f), user_error_bound(0.0f), value_base(0.0f),
          quant_radius(0), num_elements(0), outlier_count(0),
          input_type(DataType::FLOAT32), code_type(DataType::UINT16),
          eb_mode(0), zigzag_codes(0),
          outlier_threshold(std::numeric_limits<float>::infinity()),
          inplace_outliers(0), linear_mode(0), dither(0), _pad{}, dither_seed(0),
          dither_strength(1.0f), _pad2(0),
          abs_error_bound_f64(0.0), value_base_f64(0.0) {}
};
static_assert(sizeof(QuantizerConfig) <= FZM_STAGE_CONFIG_SIZE,
              "QuantizerConfig must fit in FZM_STAGE_CONFIG_SIZE");
 
template<typename TInput = float, typename TCode = uint16_t>
class QuantizerStage : public Stage {
public:
    struct Config {
        float  error_bound           = 1e-4f;   
        int    quant_radius          = 32768;   
        float  outlier_capacity      = 0.05f;   
        ErrorBoundMode eb_mode       = ErrorBoundMode::ABS;
        float precomputed_value_base = 0.0f;
        bool  zigzag_codes           = false;
        float outlier_threshold      = std::numeric_limits<float>::infinity();
        bool  inplace_outliers       = false;
        bool  linear_mode            = false;
        bool     dither              = false;
        uint64_t dither_seed         = 0;
        float    dither_strength     = 1.0f;
 
        Config() = default;
        Config(TInput eb, ErrorBoundMode mode = ErrorBoundMode::ABS,
               int radius = 32768, float outlier_cap = 0.05f)
            : error_bound(static_cast<float>(eb)), quant_radius(radius),
              outlier_capacity(outlier_cap), eb_mode(mode) {}
    };
 
    explicit QuantizerStage(const Config& config = Config());
    ~QuantizerStage() 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 (isLinearMode() || isInplaceMode()) ? 0 : sizeof(uint32_t);
    }
 
    std::string getName() const override { return "Quantizer"; }
 
    size_t getNumInputs() const override {
        if (!is_inverse_) return 1;
        return (isLinearMode() || isInplaceMode()) ? 1 : 3;
    }
    size_t getNumOutputs() const override {
        if (is_inverse_) return 1;
        return (isLinearMode() || isInplaceMode()) ? 1 : 3;
    }
 
    std::vector<std::string> getOutputNames() const override {
        if (is_inverse_) return {"reconstructed"};
        if (isLinearMode() || isInplaceMode()) return {"codes"};
        return {"codes", "outlier_vals", "outlier_idxs"};
    }
 
    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;
    }
 
    void setInverse(bool inverse) override { is_inverse_ = inverse; }
    bool isInverse() const override        { return is_inverse_; }
 
    void setDims(const std::array<size_t, 3>& dims) override { dims_ = dims; }
 
    uint16_t getStageTypeId() const override {
        return static_cast<uint16_t>(StageType::QUANTIZER);
    }
 
    uint8_t getOutputDataType(size_t output_index) const override {
        if (is_inverse_) return static_cast<uint8_t>(getInputDataType());
        // Linear mode: codes hold two's-complement signed q — declare the signed type
        // so the DAG connects cleanly to LorenzoStage<intN>.
        if (isLinearMode()) return static_cast<uint8_t>(signedOf(getCodeDataType()));
        if (isInplaceMode()) return static_cast<uint8_t>(getCodeDataType()); // only codes
        switch (output_index) {
            case 0: return static_cast<uint8_t>(getCodeDataType());
            case 1: return static_cast<uint8_t>(getInputDataType());
            case 2: return static_cast<uint8_t>(DataType::UINT32);
            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* buf, size_t max_size) const override;
    size_t getMaxHeaderSize(size_t) const override { return sizeof(QuantizerConfig); }
    void deserializeHeader(const uint8_t* buf, size_t size) override;
 
    void saveState() override {
        saved_config_ = config_;
        saved_num_elements_ = num_elements_;
        saved_actual_outlier_count_ = actual_outlier_count_;
        saved_computed_abs_eb_ = computed_abs_eb_;
        saved_computed_value_base_ = computed_value_base_;
        saved_actual_output_sizes_ = actual_output_sizes_;
    }
 
    void restoreState() override {
        config_ = saved_config_;
        num_elements_ = saved_num_elements_;
        actual_outlier_count_ = saved_actual_outlier_count_;
        computed_abs_eb_ = saved_computed_abs_eb_;
        computed_value_base_ = saved_computed_value_base_;
        actual_output_sizes_ = saved_actual_output_sizes_;
    }
 
    void setErrorBound(TInput eb)            { config_.error_bound = static_cast<float>(eb); }
    void setQuantRadius(int r)               { config_.quant_radius = r; }
    void setOutlierCapacity(float c)         { config_.outlier_capacity = c; }
    void setErrorBoundMode(ErrorBoundMode m) { config_.eb_mode = m; }
    void setValueBase(float vb)              { config_.precomputed_value_base = vb; }
    void setZigzagCodes(bool enable)         { config_.zigzag_codes = enable; }
    void setOutlierThreshold(float t)        { config_.outlier_threshold = t; }
    void setInplaceOutliers(bool enable)     { config_.inplace_outliers = enable; }
    void setLinearMode(bool enable)          { config_.linear_mode = enable; }
    void setDither(bool enable)              { config_.dither = enable; }
    void setDitherSeed(uint64_t seed)        { config_.dither_seed = seed; }
    void setDitherStrength(float strength)   { config_.dither_strength = strength; }
 
    TInput         getErrorBound()        const { return static_cast<TInput>(config_.error_bound); }
    int            getQuantRadius()       const { return config_.quant_radius; }
    ErrorBoundMode getErrorBoundMode()    const { return config_.eb_mode; }
    float          getValueBase()         const { return config_.precomputed_value_base; }
    float          getOutlierCapacity()   const { return config_.outlier_capacity; }
    bool           getZigzagCodes()       const { return config_.zigzag_codes; }
    float          getOutlierThreshold()  const { return config_.outlier_threshold; }
    bool           getInplaceOutliers()   const { return config_.inplace_outliers; }
    bool           getLinearMode()        const { return config_.linear_mode; }
    bool           getDither()            const { return config_.dither; }
    uint64_t       getDitherSeed()        const { return config_.dither_seed; }
    float          getDitherStrength()    const { return config_.dither_strength; }
 
private:
    Config config_;
    Config saved_config_;
    std::array<size_t, 3> dims_ = {0, 1, 1};  
    std::vector<size_t> actual_output_sizes_;
    std::vector<size_t> saved_actual_output_sizes_;
    size_t   num_elements_        = 0;
    size_t   saved_num_elements_  = 0;
    uint32_t actual_outlier_count_= 0;
    uint32_t saved_actual_outlier_count_ = 0;
    bool     is_inverse_          = false;
    TInput   computed_abs_eb_     = static_cast<TInput>(1e-4);
    TInput   saved_computed_abs_eb_ = static_cast<TInput>(1e-4);
    TInput   computed_value_base_ = static_cast<TInput>(0);
    TInput   saved_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);
 
    bool isInplaceMode() const {
        return config_.inplace_outliers
            && config_.eb_mode != ErrorBoundMode::REL;
    }
 
    bool isLinearMode() const { return config_.linear_mode; }
 
    static DataType signedOf(DataType d) {
        switch (d) {
            case DataType::UINT8:  return DataType::INT8;
            case DataType::UINT16: return DataType::INT16;
            case DataType::UINT32: return DataType::INT32;
            default:               return d;
        }
    }
 
    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 n) const {
        return static_cast<size_t>(std::ceil(n * config_.outlier_capacity));
    }
};
 
extern template class QuantizerStage<float,  uint16_t>;
extern template class QuantizerStage<float,  uint32_t>;
extern template class QuantizerStage<double, uint16_t>;
extern template class QuantizerStage<double, uint32_t>;
 
} // namespace fz