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 "fused/chunk_fusion/chunk_op_params.h" // shared fused-op Params (host packs)
#include "backend/types.h"
#include <array>
#include <cmath>
#include <cstdint>
#include <memory>
#include <cstring>
#include <limits>
#include <type_traits>
 
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  linear_high_precision; 
    uint8_t  power_of_two_bound;    
    uint8_t  _pad[3];           
    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),
          linear_high_precision(0), power_of_two_bound(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");
static_assert(sizeof(QuantizerConfig) == 72,
              "QuantizerConfig archive layout changed; add an explicit compatibility path");
 
template<typename TInput = float, typename TCode = uint16_t>
class QuantizerStage : public Stage {
public:
    struct Config {
        double error_bound           = 1e-4;    
        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  linear_high_precision  = false;
        bool  power_of_two_bound     = 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<double>(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_; }
 
    FusionSpec getFusionSpec() const override {
        if (is_inverse_) return {};
        if (isLinearMode() && !config_.linear_high_precision)
            return FusionSpec{FusionAccess::Elementwise, 0};
        if (isInplaceMode() && config_.zigzag_codes &&
            (config_.eb_mode == ErrorBoundMode::ABS || config_.eb_mode == ErrorBoundMode::NOA))
            return FusionSpec{FusionAccess::Elementwise, 0};
        // Split-outlier (3-port) ABS/NOA zigzag quant is Elementwise too: codes flow to the
        // next stage (port 0), outlier_vals/idxs (ports 1,2) escape as pipeline leaves
        // that the fused runner fills and the pipeline auto-concatenates. Dither is
        // excluded (the fused op has no per-element dither path).
        if (isSplitOutlierFusable()) return FusionSpec{FusionAccess::Elementwise, 0};
        return {};
    }
 
    void primeAbsEbForFusion() {
        if (config_.eb_mode == ErrorBoundMode::ABS) {
            computed_abs_eb_ = resolveUniformBound();
            if (config_.power_of_two_bound)
                computed_abs_eb_ = floorPowerOfTwo(computed_abs_eb_);
        }
    }
 
    void primeComputedAbsEb(const void* d_in, size_t scan_n,
                            MemoryPool* pool, fz::stream_t stream);
    TInput getComputedAbsEb() const { return computed_abs_eb_; }
    uint32_t getActualOutlierCount() const { return actual_outlier_count_; }
    bool supportsChunkInverseFusion() const {
        return std::is_same<TInput, float>::value &&
               std::is_same<TCode, uint32_t>::value &&
               !isLinearMode() && config_.zigzag_codes && !config_.dither &&
               (config_.eb_mode == ErrorBoundMode::ABS ||
                config_.eb_mode == ErrorBoundMode::NOA);
    }
    bool supportsWarpInverseFusion() const {
        return std::is_same<TInput, float>::value &&
               std::is_same<TCode, uint32_t>::value &&
               isLinearMode() && !config_.linear_high_precision &&
               (config_.eb_mode == ErrorBoundMode::ABS ||
                config_.eb_mode == ErrorBoundMode::NOA);
    }
    void setFusedInverseResult(size_t output_bytes) override {
        actual_output_sizes_ = {output_bytes};
    }
 
    FusionSpec getInverseFusionSpec() const override {
        if (!is_inverse_ || !supportsWarpInverseFusion()) return {};
        return FusionSpec{FusionAccess::Elementwise, 0};
    }
    FusedOpDecl getInverseFusedOp() const override {
        if (!getInverseFusionSpec().fusable()) return {};
        // Marker op: the linear `code * 2*abs_eb` dequant is a built-in tail of the
        // warp inverse harness, not a composed device policy. The generic runner
        // reads the step from getFusedInverseDequantStep(), not from this name.
        FusedOpDecl d;
        d.strategy = FusionStrategy::WarpRegister;
        d.op_name  = "LinearDequant";
        return d;
    }
    double getFusedInverseDequantStep() const override {
        return 2.0 * static_cast<double>(computed_abs_eb_);
    }
    double getFusedForwardQuantStep() const override {
        return 2.0 * static_cast<double>(computed_abs_eb_);
    }
 
    FusedOpDecl getFusedOp() const override {
        if (!std::is_same<TInput, float>::value || is_inverse_) return {};  // device ops read float
        // Warp-register (cuSZp): linear float quant is the Elementwise loader, for ABS or NOA
        // (both resolve to one uniform-step absolute bound — the runner primes the NOA
        // range scan and passes the resolved abs_eb). REL is excluded: it is log-domain
        // per-value quant, not a single abs_eb, so it can't ride the fused kernel.
        if (isLinearMode() && !config_.linear_high_precision &&
            (config_.eb_mode == ErrorBoundMode::ABS || config_.eb_mode == ErrorBoundMode::NOA))
            return FusedOpDecl{FusionStrategy::WarpRegister, "LinearQuant", "", {}};
        // Chunk-cooperative (PFPL): inplace+zigzag ABS/NOA float quant.
        if (isInplaceMode() && config_.zigzag_codes &&
            (config_.eb_mode == ErrorBoundMode::ABS || config_.eb_mode == ErrorBoundMode::NOA)) {
            fused::chunk::QuantInplaceZigzagParams p;
            p.ebx2_r    = 1.0f / (2.0f * static_cast<float>(computed_abs_eb_));
            p.radius    = static_cast<uint32_t>(config_.quant_radius);
            p.threshold = config_.outlier_threshold;
            FusedOpDecl d;
            d.strategy       = FusionStrategy::ChunkCooperative;
            d.op_name        = "QuantInplaceZigzag";
            d.include_header = "fused/chunk_fusion/chunk_fusion.cuh";
            d.params.resize(sizeof(p));
            std::memcpy(d.params.data(), &p, sizeof(p));
            return d;
        }
        // Chunk-cooperative (3-port): split-outlier ABS/NOA zigzag float quant. Same
        // uniform-step params as inplace; the op appends outliers to side buffers
        // instead of inlining raw bits (see QuantSplitOutlier in chunk_fusion.cuh).
        if (isSplitOutlierFusable()) {
            fused::chunk::QuantSplitOutlierParams p;
            p.ebx2_r    = 1.0f / (2.0f * static_cast<float>(computed_abs_eb_));
            p.radius    = static_cast<uint32_t>(config_.quant_radius);
            p.threshold = config_.outlier_threshold;
            FusedOpDecl d;
            d.strategy       = FusionStrategy::ChunkCooperative;
            d.op_name        = "QuantSplitOutlier";
            d.include_header = "fused/chunk_fusion/chunk_fusion.cuh";
            d.params.resize(sizeof(p));
            std::memcpy(d.params.data(), &p, sizeof(p));
            return d;
        }
        return {};
    }
 
    std::vector<FusedAuxOutputDecl> getFusedAuxOutputs() const override {
        if (is_inverse_ || !isSplitOutlierFusable()) return {};
        return {
            FusedAuxOutputDecl{1, "outlier_vals", FusedAuxSizeKind::CompactedElements,
                               static_cast<uint8_t>(getInputDataType()), 1u, 0u, 1u},
            FusedAuxOutputDecl{2, "outlier_idxs", FusedAuxSizeKind::CompactedElements,
                               static_cast<uint8_t>(DataType::UINT32), 1u, 0u, 1u},
        };
    }
 
    void primeFusedForwardState(const FusedPrimeContext& c) override {
        fused_outlier_count_set_ = false;   // reset per fused compress; set by setFusedSideOutput
        num_elements_ = c.input_bytes / sizeof(TInput);
        primeComputedAbsEb(c.d_input, c.input_bytes / sizeof(TInput), c.pool, c.stream);
    }
 
    void setFusedSideOutput(int output_index, size_t bytes) override {
        if (actual_output_sizes_.size() < 3) actual_output_sizes_.resize(3, 0);
        if (output_index == 1) {
            actual_outlier_count_   = static_cast<uint32_t>(bytes / sizeof(TInput));
            actual_output_sizes_[1] = bytes;
            fused_outlier_count_set_ = true;
        } else if (output_index == 2) {
            actual_outlier_count_   = static_cast<uint32_t>(bytes / sizeof(uint32_t));
            actual_output_sizes_[2] = bytes;
            fused_outlier_count_set_ = true;
        }
    }
 
    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(double eb)             { config_.error_bound = 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 setLinearHighPrecision(bool enable) { config_.linear_high_precision = enable; }
    void setPowerOfTwoBound(bool enable)     { config_.power_of_two_bound = 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           getLinearHighPrecision() const { return config_.linear_high_precision; }
    bool           getPowerOfTwoBound()   const { return config_.power_of_two_bound; }
    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     fused_outlier_count_set_ = false;
    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;
    uint32_t* d_linear_overflow_scratch_ = nullptr;
    MemoryPool* persistent_pool_ = nullptr;
    std::weak_ptr<const void> persistent_pool_alive_;
 
 
    void initOutlierCountScratch(MemoryPool* pool);
    void initLinearOverflowScratch(MemoryPool* pool);
 
    // ── execute() kernel-launch helpers ────────────────────────────────────
    // Each wraps exactly one forward/inverse kernel-instantiation path so
    // execute() reduces to mode resolution + dispatch. Pure call-site
    // extraction: no kernel logic, math, or launch parameters changed.
 
    void launchRelForward(fz::stream_t stream, int grid, int block,
                           const TInput* in, size_t num_elements,
                           float log2eb, float log2eb_r, float opp_eb, float oopp_eb,
                           size_t max_outliers,
                           TCode* codes, TInput* outlier_vals, uint32_t* outlier_idxs);
    void launchLinearForward(fz::stream_t stream, int grid, int block,
                              const TInput* in, size_t num_elements,
                              TInput ebx2_r, double ebx2_r_f64, TCode* codes);
    void launchInplaceForward(fz::stream_t stream, int grid, int block,
                               const TInput* in, size_t num_elements,
                               TInput ebx2_r, TCode* codes);
    void launchAbsNoaForward(fz::stream_t stream, int grid, int block,
                              const TInput* in, size_t num_elements,
                              TInput ebx2_r, TInput ebx2, TInput abs_eb,
                              TInput dither_amp, uint64_t dither_seed,
                              size_t max_outliers,
                              TCode* codes, TInput* outlier_vals, uint32_t* outlier_idxs);
 
    void launchRelInverse(fz::stream_t stream, int grid, int block,
                           const TCode* codes, size_t num_elements,
                           float log2eb, TInput* out);
    void launchLinearInverse(fz::stream_t stream, int grid, int block,
                              const TCode* codes, size_t num_elements,
                              TInput ebx2, double ebx2_f64, TInput* out);
    void launchInplaceInverse(fz::stream_t stream, int grid, int block,
                               const TCode* codes, size_t num_elements,
                               TInput ebx2, TInput* out);
    void launchAbsNoaInverse(fz::stream_t stream, int grid, int block,
                              const TCode* codes, size_t num_elements,
                              TInput ebx2, TInput dither_amp, uint64_t dither_seed,
                              TInput* out);
 
    bool isInplaceMode() const {
        return config_.inplace_outliers
            && config_.eb_mode != ErrorBoundMode::REL;
    }
 
    bool isLinearMode() const { return config_.linear_mode; }
 
    static TInput floorPowerOfTwo(TInput value) {
        if (!(value > TInput(0)) || !std::isfinite(value))
            throw std::runtime_error(
                "QuantizerStage: power_of_two_bound requires a finite positive absolute bound");
        int exponent = 0;
        std::frexp(value, &exponent);
        return std::ldexp(TInput(1), exponent - 1);
    }
 
    TInput resolveUniformBound(TInput scale = TInput(1)) const {
        const double resolved = config_.error_bound * static_cast<double>(scale);
        TInput bound = static_cast<TInput>(resolved);
        // The strict path must never loosen a decimal/TOML request merely because
        // its TInput representation rounded upward. f64 compares equal here;
        // this primarily protects f32.
        if (config_.linear_high_precision && static_cast<double>(bound) > resolved)
            bound = std::nextafter(bound, TInput(0));
        return bound;
    }
 
    void applyUniformBoundPolicy(TInput data_abs_max, bool have_data_abs_max);
 
    bool isSplitOutlierFusable() const {
        return std::is_same<TInput, float>::value && !is_inverse_
            && !isInplaceMode() && !isLinearMode()
            && config_.zigzag_codes && !config_.dither
            && (config_.eb_mode == ErrorBoundMode::ABS ||
                config_.eb_mode == ErrorBoundMode::NOA);
    }
 
    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