ginterp_stage.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 <string>
#include <stdexcept>
#include <type_traits>
#include <unordered_map>
#include <vector>
 
// Forward-declared in the public header so callers don't pull in the
// cuSZ-Hi type subset. The full definition lives in cusz_type_subset.h,
// included only inside the stage TU.
struct INTERPOLATION_PARAMS;
 
namespace fz {
 
enum class GInterp2DGeometry : uint8_t { Fine16 = 0, Native64 = 1 };
 
struct GInterpConfig {
    // ── identity / dims ─────────────────────────────────────────────────────
    double   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;             
    uint32_t anchor_dim_x;      
    uint32_t anchor_dim_y;      
    uint32_t anchor_dim_z;      
    float    user_eb;           
    float    value_base;        
 
    // ── resolved INTERPOLATION_PARAMS (phase 2 auto-tune output, 39 B) ──────
    // Both encoder and decoder must use the exact same intp_param values, so
    // the resolved params are written here on compress and consumed on decompress.
    // Layout chosen to mirror `INTERPOLATION_PARAMS` field-for-field; we don't
    // memcpy the struct directly because its padding is implementation-defined.
    double   intp_alpha;        
    double   intp_beta;         
    uint8_t  intp_use_md[6];    
    uint8_t  intp_use_natural[6];
    uint8_t  intp_reverse[6];
    uint8_t  auto_tuning_mode;  
    uint8_t  geometry_2d;       
    uint8_t  pad[7];            
 
    // ~96 bytes. Comfortable margin under FZM_STAGE_CONFIG_SIZE (128 B);
    // the static_assert below is the source of truth.
 
    GInterpConfig()
        : error_bound(0.0), quant_radius(0), num_elements(0), outlier_count(0),
          input_type(DataType::FLOAT32), code_type(DataType::UINT16),
          ndim(3), eb_mode(0),
          dim_x(0), dim_y(1), dim_z(1),
          anchor_dim_x(0), anchor_dim_y(0), anchor_dim_z(0),
          user_eb(0.0f), value_base(0.0f),
          intp_alpha(1.75), intp_beta(4.0),
          intp_use_md{1, 1, 0, 0, 0, 0},
          intp_use_natural{0, 0, 0, 0, 0, 0},
          intp_reverse{0, 0, 0, 0, 0, 0},
          auto_tuning_mode(0), geometry_2d(0), pad{} {}
};
static_assert(sizeof(GInterpConfig) <= FZM_STAGE_CONFIG_SIZE,
              "GInterpConfig must fit in FZM_STAGE_CONFIG_SIZE");
 
template <typename TInput = float, typename TCode = uint16_t>
class GInterpStage : public Stage {
public:
    struct Config {
        float error_bound        = 1e-3f;
        int   quant_radius       = 0;
        float outlier_capacity   = 0.10f;
        std::array<size_t, 3> dims = {0, 0, 0};
        ErrorBoundMode eb_mode   = ErrorBoundMode::ABS;
        float precomputed_value_base = 0.0f;
        uint8_t auto_tuning_mode = 0;
 
        double  manual_alpha = 0.0;
        double  manual_beta  = 0.0;
 
        GInterp2DGeometry geometry_2d = GInterp2DGeometry::Fine16;
 
        Config() = default;
    };
 
    explicit GInterpStage(const Config& cfg = Config()) : config_(cfg) {
        actual_output_sizes_.resize(4, 0);
    }
    ~GInterpStage() override;
 
    // ── Stage interface ──────────────────────────────────────────────────────
    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 needsProfilingScratch() ? kProfilingErrCount * sizeof(float) : 0;
    }
    size_t estimatePinnedFootprintBytes(size_t /*estimated_inlen*/) const override {
        return needsProfilingScratch() ? kProfilingErrCount * sizeof(float) : 0;
    }
 
    std::string getName() const override { return "GInterp"; }
    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", "anchor", "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> r;
        for (size_t i = 0; i < names.size() && i < actual_output_sizes_.size(); i++)
            r[names[i]] = actual_output_sizes_[i];
        return r;
    }
    size_t getActualOutputSize(int index) const override {
        return (index >= 0 && index < static_cast<int>(actual_output_sizes_.size()))
            ? actual_output_sizes_[index] : 0;
    }
 
    void saveState()    override { saved_output_sizes_ = actual_output_sizes_; }
    void restoreState() override { actual_output_sizes_ = saved_output_sizes_; }
 
    // ── Setters ──────────────────────────────────────────────────────────────
    void setErrorBound(float eb)              { config_.error_bound = eb; }
    void setQuantRadius(int radius)           { config_.quant_radius = radius; }
    void setOutlierCapacity(float cap)        { config_.outlier_capacity = cap; }
    void setErrorBoundMode(ErrorBoundMode m)  {
        config_.eb_mode = resolveApproxRelMode(m, "GInterpStage");
    }
    void setValueBase(float v)                { config_.precomputed_value_base = v; }
    void setAutoTuning(uint8_t mode)          { config_.auto_tuning_mode = mode; }
    void set2DGeometry(GInterp2DGeometry geometry) {
        if (geometry != GInterp2DGeometry::Fine16 &&
            geometry != GInterp2DGeometry::Native64)
            throw std::invalid_argument("GInterpStage: unsupported 2-D geometry");
        config_.geometry_2d = geometry;
    }
    void setManualAlphaBeta(double alpha, double beta) {
        config_.manual_alpha = alpha;
        config_.manual_beta  = beta;
    }
    void setDims(const std::array<size_t, 3>& dims) override;
    void setDims(size_t x, size_t y, size_t z) {
        setDims(std::array<size_t, 3>{x, y, z});
    }
 
    float          getErrorBound()       const { return config_.error_bound; }
    int            getQuantRadius()      const { return config_.quant_radius; }
    float          getOutlierCapacity()  const { return config_.outlier_capacity; }
    ErrorBoundMode getErrorBoundMode()   const { return config_.eb_mode; }
    float          getValueBase()        const { return config_.precomputed_value_base; }
    uint8_t        getAutoTuningMode()   const { return config_.auto_tuning_mode; }
    GInterp2DGeometry get2DGeometry()   const { return config_.geometry_2d; }
    std::array<size_t, 3> getDims()      const { return config_.dims; }
 
    void setInverse(bool inv) override { is_inverse_ = inv; }
    bool isInverse()           const override { return is_inverse_; }
 
    int ndim() const {
        if (config_.dims[2] > 1) return 3;
        if (config_.dims[1] > 1) return 2;
        return 1;
    }
 
    // ── Type / Serialization ─────────────────────────────────────────────────
    uint16_t getStageTypeId() const override {
        return static_cast<uint16_t>(StageType::G_INTERP);
    }
 
    uint8_t getOutputDataType(size_t output_index) const override {
        switch (output_index) {
            case 0: return static_cast<uint8_t>(codeDataType());    // codes
            case 1: return static_cast<uint8_t>(inputDataType());   // anchor
            case 2: return static_cast<uint8_t>(inputDataType());   // outlier_vals
            case 3: return static_cast<uint8_t>(DataType::UINT32);  // outlier_idxs
            default: return static_cast<uint8_t>(DataType::UINT8);
        }
    }
    uint8_t getInputDataType(size_t /*input_index*/) const override {
        return static_cast<uint8_t>(inputDataType());
    }
 
    size_t serializeHeader(size_t output_index, uint8_t* buf, size_t max_size) const override;
    void deserializeHeader(const uint8_t* buf, size_t size) override;
    size_t getMaxHeaderSize(size_t /*output_index*/) const override {
        return sizeof(GInterpConfig);
    }
 
    bool isGraphCompatible() const override {
        const bool tune_ok = (config_.auto_tuning_mode == 0 ||
                              config_.auto_tuning_mode == 5);
        const bool radius_ok = config_.quant_radius > 0;
        const bool eb_ok = (config_.eb_mode == ErrorBoundMode::ABS) ||
                           (config_.precomputed_value_base > 0.0f);
        const bool mode5_alpha_ok = (config_.auto_tuning_mode != 5) ||
                                    (config_.manual_alpha > 0.0);
        return tune_ok && radius_ok && eb_ok && mode5_alpha_ok;
    }
 
private:
    Config config_;
    std::vector<size_t> actual_output_sizes_;
    std::vector<size_t> saved_output_sizes_;
 
    bool   is_inverse_ = false;
    size_t num_elements_ = 0;
    uint32_t actual_outlier_count_ = 0;
    uint32_t* d_outlier_count_scratch_ = nullptr;
 
    TInput computed_abs_eb_ = 0;
    float computed_value_base_ = 0.0f;
 
    std::array<size_t, 3> anchor_dims_ = {0, 0, 0};
 
    // ── Phase 2: auto-tuning state ──────────────────────────────────────────
    static constexpr size_t kProfilingErrCount = 36;
    float*   d_profiling_errors_ = nullptr;
    float*   h_profiling_errors_ = nullptr;
    MemoryPool* persistent_pool_ = nullptr;
    std::weak_ptr<const void> persistent_pool_alive_;
 
 
    double  resolved_alpha_ = 1.75;
    double  resolved_beta_  = 4.0;
    uint8_t resolved_use_md_[6]      = {1, 1, 0, 0, 0, 0};
    uint8_t resolved_use_natural_[6] = {0, 0, 0, 0, 0, 0};
    uint8_t resolved_reverse_[6]     = {0, 0, 0, 0, 0, 0};
 
    bool needsProfilingScratch() const {
        const uint8_t m = config_.auto_tuning_mode;
        return m == 1 || m == 2 || m == 3 || m == 4;
    }
 
    void initProfilingScratch(MemoryPool* pool);
    void initOutlierCountScratch(MemoryPool* pool);
    INTERPOLATION_PARAMS buildIntpParam() const;
 
    static DataType inputDataType() { return fused::dataTypeOf<TInput>(); }
    static DataType codeDataType()  { return fused::dataTypeOf<TCode>(); }
    size_t getMaxOutlierCount(size_t n) const {
        return static_cast<size_t>(std::ceil(n * config_.outlier_capacity));
    }
};
 
extern template class GInterpStage<float, uint8_t>;
extern template class GInterpStage<float, uint16_t>;
extern template class GInterpStage<float, uint32_t>;
extern template class GInterpStage<double, uint8_t>;
extern template class GInterpStage<double, uint16_t>;
extern template class GInterpStage<double, uint32_t>;
 
} // namespace fz