log_transform_stage.h

Go to the documentation of this file.

#pragma once
 
#include "stage/stage.h"
#include "fzm_format.h"
#include "backend/types.h"
#include <array>
#include <cmath>
#include <cstdint>
#include <memory>
#include <cstring>
#include <stdexcept>
#include <string>
#include <type_traits>
#include <unordered_map>
#include <vector>
 
namespace fz {
 
struct LogTransformConfig {
    float    error_bound;    
    float    threshold;      
    uint32_t num_elements;   
    uint32_t outlier_count;  
    DataType input_type;     
    uint8_t  reserved[3];    
 
    LogTransformConfig()
        : error_bound(1e-3f), threshold(0.0f), num_elements(0), outlier_count(0),
          input_type(DataType::FLOAT32), reserved{0, 0, 0} {}
};
static_assert(sizeof(LogTransformConfig) == 20,
              "LogTransformConfig must be 20 bytes");
 
template<typename TInput = float>
class LogTransformStage : public Stage {
    static_assert(std::is_floating_point<TInput>::value,
                  "LogTransformStage: TInput must be a floating-point type.");
 
public:
    static constexpr float kLogRoundTripSlack = 1e-6f;
 
    static constexpr size_t kMinOutlierSlots = 8;
 
    struct Config {
        float error_bound = 1e-3f;
        float threshold = 0.0f;
        float outlier_capacity = 0.05f;
        Config() = default;
        explicit Config(float eb, float thr = 0.0f, float cap = 0.05f)
            : error_bound(eb), threshold(thr), outlier_capacity(cap) {}
    };
 
    LogTransformStage() = default;
    explicit LogTransformStage(const Config& config) : config_(config) {}
    ~LogTransformStage() override;
 
    // ── Execution ────────────────────────────────────────────────────────────
    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(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);
    }
 
    bool isGraphCompatible() const override { return true; }
 
    // ── Metadata ─────────────────────────────────────────────────────────────
    std::string getName() const override { return "LogTransform"; }
 
    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 {
        if (is_inverse_) return {"reconstructed"};
        return {"output", "signs", "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_; }
 
    uint16_t getStageTypeId() const override {
        return static_cast<uint16_t>(StageType::LOG_TRANSFORM);
    }
 
    uint8_t getOutputDataType(size_t output_index) const override {
        if (is_inverse_) return static_cast<uint8_t>(inputDataType());
        switch (output_index) {
            case 0:  return static_cast<uint8_t>(inputDataType());   // log values
            case 1:  return static_cast<uint8_t>(DataType::UINT8);   // sign bitmap
            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 {
        if (!is_inverse_) return static_cast<uint8_t>(inputDataType());
        switch (input_index) {
            case 0:  return static_cast<uint8_t>(inputDataType());   // log values
            case 1:  return static_cast<uint8_t>(DataType::UINT8);   // sign bitmap
            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);
        }
    }
 
    // ── Serialization ────────────────────────────────────────────────────────
    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(LogTransformConfig); }
    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_actual_output_sizes_  = actual_output_sizes_;
    }
    void restoreState() override {
        config_               = saved_config_;
        num_elements_         = saved_num_elements_;
        actual_outlier_count_ = saved_actual_outlier_count_;
        actual_output_sizes_  = saved_actual_output_sizes_;
    }
 
    // ── Configuration ────────────────────────────────────────────────────────
    void setErrorBound(float delta)      { config_.error_bound = delta; }
    void setThreshold(float t)           { config_.threshold = t; }
    void setOutlierCapacity(float c)     { config_.outlier_capacity = c; }
 
    float getErrorBound()      const { return config_.error_bound; }
    float getThreshold()       const { return config_.threshold; }
    float getOutlierCapacity() const { return config_.outlier_capacity; }
    uint32_t getOutlierCount() const { return actual_outlier_count_; }
 
    float quantizerErrorBound() const {
        return quantizerErrorBoundFor(config_.error_bound);
    }
 
    static float quantizerErrorBoundFor(float delta) {
        return std::log2(1.0f + delta) - kLogRoundTripSlack;
    }
 
    static float minimumErrorBound() {
        return std::exp2(2.0f * kLogRoundTripSlack) - 1.0f;
    }
 
private:
    Config config_;
    Config saved_config_;
    bool   is_inverse_ = false;
 
    std::vector<size_t> actual_output_sizes_{0, 0, 0, 0};
    std::vector<size_t> saved_actual_output_sizes_{0, 0, 0, 0};
 
    size_t   num_elements_               = 0;
    size_t   saved_num_elements_         = 0;
    uint32_t actual_outlier_count_       = 0;
    uint32_t saved_actual_outlier_count_ = 0;
 
    uint32_t*   d_outlier_count_scratch_ = nullptr;
    MemoryPool* persistent_pool_         = nullptr;
    std::weak_ptr<const void> persistent_pool_alive_;
 
 
    void initOutlierCountScratch(MemoryPool* pool);
 
    size_t maxOutlierCount(size_t n) const {
        const size_t scaled =
            static_cast<size_t>(static_cast<double>(n) * config_.outlier_capacity);
        const size_t floored = (scaled < kMinOutlierSlots) ? kMinOutlierSlots : scaled;
        return (floored > n) ? n : floored;
    }
 
    static size_t signBytes(size_t n) { return (n + 7) / 8; }
 
    static DataType inputDataType() {
        return std::is_same<TInput, double>::value ? DataType::FLOAT64
                                                   : DataType::FLOAT32;
    }
};
 
extern template class LogTransformStage<float>;
 
} // namespace fz