adaptive_lorenzo_stage.h

Go to the documentation of this file.

#pragma once
 
#include "stage/stage.h"
#include "fzm_format.h"
#include "backend/types.h"
#include <algorithm>
#include <array>
#include <cstdint>
#include <cstring>
#include "fused/common/data_type_of.h"
#include "fused/lorenzo_quant/lorenzo_quant.h"  // ErrorBoundMode, resolveApproxRelMode — FusedQuantAdaptiveLorenzoStage
 
#include <stdexcept>
#include <string>
#include <type_traits>
#include <unordered_map>
#include <vector>
 
namespace fz {
 
struct AdaptiveLorenzoConfig {
    DataType data_type;        
    uint8_t  coder_block_size; 
    uint8_t  blocks_per_tile;  
    uint8_t  enable_order2;    
    uint8_t  enable_centering; 
    uint8_t  reserved[3];      
    uint32_t num_elements;     
 
    AdaptiveLorenzoConfig()
        : data_type(DataType::INT32), coder_block_size(32), blocks_per_tile(8),
          enable_order2(1), enable_centering(1), reserved{0, 0, 0},
          num_elements(0) {}
};
static_assert(sizeof(AdaptiveLorenzoConfig) <= FZM_STAGE_CONFIG_SIZE,
              "AdaptiveLorenzoConfig must fit in FZM_STAGE_CONFIG_SIZE");
 
template<typename T>
class AdaptiveLorenzoStage : public Stage {
    static_assert(std::is_integral<T>::value && std::is_signed<T>::value,
                  "AdaptiveLorenzoStage requires a signed integer type");
public:
    struct Config {
        uint32_t coder_block_size = 32;
        uint32_t blocks_per_tile = 8;
        bool enable_order2    = true;   
        bool enable_centering = true;   
        Config() = default;
    };
 
    AdaptiveLorenzoStage() { validate(); }
    explicit AdaptiveLorenzoStage(const Config& config) : config_(config) { validate(); }
 
    void setInverse(bool inv) override { is_inverse_ = inv; }
    bool isInverse() const override    { return is_inverse_; }
 
    uint32_t getTileSize() const {
        return config_.coder_block_size * config_.blocks_per_tile;
    }
 
    bool bindDownstreamEncodingOracle(const EncodingOracleDecl& decl) override {
        if (!decl.valid() || !decl.additive ||
            (decl.kind != EncodingOracleKind::PlainFixedRateBitpack &&
             decl.kind != EncodingOracleKind::AdaptiveFixedRateBitpack) ||
            decl.input_data_type != static_cast<uint8_t>(getElementDataType()) ||
            decl.unit_elems != config_.coder_block_size) {
            return false;
        }
        bound_oracle_ = decl;
        has_bound_oracle_ = true;
        return true;
    }
 
    bool hasBoundEncodingOracle() const { return has_bound_oracle_; }
    EncodingOracleKind getBoundEncodingOracleKind() const {
        return has_bound_oracle_ ? bound_oracle_.kind
                                 : EncodingOracleKind::PlainFixedRateBitpack;
    }
 
    FusionSpec getFusionSpec() const override {
        if (is_inverse_ || !has_bound_oracle_) return {};
        return FusionSpec{FusionAccess::TileSelector, getTileSize(),
                          config_.coder_block_size};
    }
 
    std::vector<FusedAuxOutputDecl> getFusedAuxOutputs() const override {
        if (!getFusionSpec().fusable()) return {};
        return {
            FusedAuxOutputDecl{1, "modes", FusedAuxSizeKind::FixedBitsPerUnit,
                               static_cast<uint8_t>(DataType::UINT8), getTileSize(),
                               2u, 0u},
            FusedAuxOutputDecl{2, "means", FusedAuxSizeKind::CompactedElements,
                               static_cast<uint8_t>(getElementDataType()), getTileSize(),
                               0u, 1u},
        };
    }
 
    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;
 
    std::string getName() const override { return "AdaptiveLorenzo"; }
    size_t getNumInputs()  const override { return is_inverse_ ? 3 : 1; }
    size_t getNumOutputs() const override { return is_inverse_ ? 1 : 3; }
 
    std::vector<std::string> getOutputNames() const override {
        return {"output", "modes", "means"};
    }
 
    ~AdaptiveLorenzoStage() override { releaseScratch(); }
 
    void postStreamSync(fz::stream_t stream) override;
 
    bool isGraphCompatible() const override { return false; }
 
    size_t estimateScratchBytes(
        const std::vector<size_t>& input_sizes
    ) const override;
 
    std::vector<size_t> estimateOutputSizes(
        const std::vector<size_t>& input_sizes
    ) const override {
        if (input_sizes.empty()) return {0, 0, 0};
        if (is_inverse_) return {input_sizes[0]};
        const size_t n     = input_sizes[0] / sizeof(T);
        const size_t tiles = numTiles(n);
        // modes: 2 bits per tile. means: worst case one per tile — the real
        // length is trimmed in postStreamSync() once the scan total is known.
        return {input_sizes[0], (tiles + 3) / 4, tiles * sizeof(T)};
    }
 
    void saveState() override    { saved_output_sizes_ = actual_output_sizes_; }
    void restoreState() override {
        if (!saved_output_sizes_.empty()) actual_output_sizes_ = saved_output_sizes_;
    }
 
    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 setFusedSideOutput(int output_index, size_t bytes) override {
        if (actual_output_sizes_.size() < 3) actual_output_sizes_.resize(3, 0);
        if (output_index == 1 || output_index == 2)
            actual_output_sizes_[static_cast<size_t>(output_index)] = bytes;
    }
 
    uint16_t getStageTypeId() const override {
        return static_cast<uint16_t>(StageType::ADAPTIVE_LORENZO);
    }
 
    uint8_t getOutputDataType(size_t output_index) const override {
        // The mode map is a byte stream; the residuals and means are T.
        return static_cast<uint8_t>(output_index == 1 ? DataType::UINT8
                                                      : getElementDataType());
    }
    uint8_t getInputDataType(size_t input_index) const override {
        return static_cast<uint8_t>(input_index == 1 ? DataType::UINT8
                                                     : getElementDataType());
    }
 
    size_t serializeHeader(size_t /*output_index*/, uint8_t* buf, size_t max_size) const override {
        if (max_size < sizeof(AdaptiveLorenzoConfig))
            throw std::runtime_error("AdaptiveLorenzoStage: header buffer too small");
        AdaptiveLorenzoConfig cfg;
        cfg.data_type        = getElementDataType();
        cfg.coder_block_size = static_cast<uint8_t>(config_.coder_block_size);
        cfg.blocks_per_tile  = static_cast<uint8_t>(config_.blocks_per_tile);
        cfg.enable_order2    = config_.enable_order2 ? 1u : 0u;
        cfg.enable_centering = config_.enable_centering ? 1u : 0u;
        cfg.num_elements     = static_cast<uint32_t>(num_elements_);
        std::memcpy(buf, &cfg, sizeof(cfg));
        return sizeof(cfg);
    }
 
    void deserializeHeader(const uint8_t* buf, size_t size) override {
        if (size < sizeof(AdaptiveLorenzoConfig))
            throw std::runtime_error("AdaptiveLorenzoStage: header too small");
        AdaptiveLorenzoConfig cfg;
        std::memcpy(&cfg, buf, sizeof(cfg));
        config_.coder_block_size = cfg.coder_block_size;
        config_.blocks_per_tile  = cfg.blocks_per_tile;
        config_.enable_order2    = (cfg.enable_order2 != 0);
        config_.enable_centering = (cfg.enable_centering != 0);
        num_elements_            = cfg.num_elements;
        validate();
    }
 
    size_t getMaxHeaderSize(size_t /*output_index*/) const override {
        return sizeof(AdaptiveLorenzoConfig);
    }
 
private:
    Config config_;
    EncodingOracleDecl bound_oracle_;
    bool has_bound_oracle_ = false;
    bool   is_inverse_ = false;
    size_t num_elements_ = 0;
    std::vector<size_t> actual_output_sizes_{0, 0, 0};
    std::vector<size_t> saved_output_sizes_;
 
    // Persistent compaction scratch, grown when a larger input is seen. Held
    // across calls so postStreamSync() can read the scan total after the fact.
    uint8_t*    d_modes_dense_ = nullptr;   
    T*          d_means_dense_ = nullptr;   
    uint32_t*   d_flags_       = nullptr;   
    uint32_t*   d_offsets_     = nullptr;   
    size_t      scratch_tiles_ = 0;
    MemoryPool* scratch_pool_  = nullptr;
    size_t      pending_tiles_ = 0;         
 
    size_t ensureScratch(size_t num_tiles, MemoryPool* pool, fz::stream_t stream);
    void   releaseScratch();
 
    size_t numTiles(size_t n) const {
        const size_t t = getTileSize();
        return (n + t - 1) / t;
    }
 
    void validate() const {
        if (config_.coder_block_size != 32)
            throw std::invalid_argument(
                "AdaptiveLorenzoStage: coder_block_size must be 32 (the cost model "
                "and the per-block warp reduction both assume a 32-element block)");
        if (config_.blocks_per_tile < 1 || config_.blocks_per_tile > 32)
            throw std::invalid_argument(
                "AdaptiveLorenzoStage: blocks_per_tile must be in [1, 32] so the "
                "tile fits one CUDA block, got "
                + std::to_string(config_.blocks_per_tile));
    }
 
    static DataType getElementDataType() { return fused::dataTypeOf<T>(); }
};
 
extern template class AdaptiveLorenzoStage<int16_t>;
extern template class AdaptiveLorenzoStage<int32_t>;
 
template<typename T>
void launchAdaptiveLorenzoForward(
    const T* d_input, T* d_residuals, uint8_t* d_modes, T* d_means,
    uint32_t* d_flags, size_t n, uint32_t tile_size, bool enable_order2,
    bool enable_centering, EncodingOracleKind oracle_kind, fz::stream_t stream);
 
template<typename T>
void launchAdaptiveLorenzoInverse(
    const T* d_residuals, const uint8_t* d_modes, const T* d_means,
    const uint32_t* d_offsets, T* d_output,
    size_t n, uint32_t tile_size, fz::stream_t stream);
 
template<typename T>
void launchAdaptiveLorenzoCompact(
    const uint8_t* d_modes_dense, const T* d_means_dense, const uint32_t* d_offsets,
    uint8_t* d_modes_packed, T* d_means_compact, size_t num_tiles,
    fz::stream_t stream);
 
void launchAdaptiveLorenzoFlags(
    const uint8_t* d_modes_packed, uint32_t* d_flags, size_t num_tiles,
    fz::stream_t stream);
 
} // namespace fz