adaptive_bitpack_stage.h

Go to the documentation of this file.

#pragma once
 
#include "stage/stage.h"
#include "fzm_format.h"
#include "backend/types.h"
#include "fused/fused_block/warp_op_params.h"   // kMaxWarpElemsPerLane
#include <cstdint>
#include <cstring>
#include <stdexcept>
#include <string>
#include <type_traits>
#include <unordered_map>
#include <vector>
 
namespace fz {
 
struct AdaptiveBitpackConfig {
    DataType data_type;        
    uint8_t  outlier_selection;
    uint8_t  _pad[2];          
    uint32_t block_size;       
    uint64_t num_elements;     
 
    AdaptiveBitpackConfig()
        : data_type(DataType::INT32), outlier_selection(0), _pad{},
          block_size(32), num_elements(0) {}
};
static_assert(sizeof(AdaptiveBitpackConfig) <= FZM_STAGE_CONFIG_SIZE,
              "AdaptiveBitpackConfig must fit in FZM_STAGE_CONFIG_SIZE");
 
template<typename T>
class AdaptiveBitpackStage : public Stage {
    static_assert(std::is_same_v<T, int16_t> || std::is_same_v<T, int32_t>,
                  "AdaptiveBitpackStage: T must be int16_t or int32_t.");
public:
    AdaptiveBitpackStage() = default;
    ~AdaptiveBitpackStage() override;
 
    // ── Stage control ──────────────────────────────────────────────────────
    void setInverse(bool inv) override { is_inverse_ = inv; }
    bool isInverse() const override    { return is_inverse_; }
    // Forward (compress) is graph-capturable: execute() enqueues only
    // device-side work and the data-dependent compressed-size readback is
    // deferred to postStreamSync() (run after the launch, outside any capture
    // window).  The inverse path keeps a per-execute layout and is left out of
    // graph capture, mirroring RZEStage.
    bool isGraphCompatible() const override { return !is_inverse_; }
 
    void setBlockSize(uint32_t n) {
        if (n == 0 || n > 1024)
            throw std::invalid_argument(
                "AdaptiveBitpackStage::setBlockSize: n must be in [1, 1024], got "
                + std::to_string(n));
        block_size_ = n;
    }
    uint32_t getBlockSize() const { return block_size_; }
 
    FusionSpec getFusionSpec() const override {
        if (is_inverse_ || block_size_ == 0 || block_size_ % 32u != 0 ||
            block_size_ / 32u > fused::warp::kMaxWarpElemsPerLane) return {};
        return FusionSpec{FusionAccess::SegmentCodec, block_size_};
    }
 
    FusedOpDecl getFusedOp() const override {
        if (!getFusionSpec().fusable()) return {};
        FusedOpDecl d;
        d.strategy       = FusionStrategy::WarpRegister;
        // Outlier mode -> the (swappable) adaptive coder; plain mode -> the true
        // 1-byte-meta plain format its own inverse reads. Either composes into
        // the warp chain with no other change.
        d.op_name        = outlier_selection_ ? fused_coder_ : "PlainRateCoder";
        d.include_header = "fused/fused_block/warp_fusion.cuh";
        // ThreadFixedRateCoderN<N>/ThreadPlainRateCoderN<N> mirror AdaptiveBitpackCoder
        // and PlainRateCoder byte-for-byte for any N (launcher fills in <N> from the
        // paired predictor's block size). Only known when fused_coder_ is left at its
        // default "AdaptiveBitpackCoder" -- a caller who swapped in some other
        // SegmentCodec sink via setFusedCoder() has no known TI policy to offer.
        if (outlier_selection_) {
            if (fused_coder_ == "AdaptiveBitpackCoder") d.ti_op_name = "ThreadFixedRateCoderN";
        } else {
            d.ti_op_name = "ThreadPlainRateCoderN";
        }
        return d;
    }
 
    FusionSpec getInverseFusionSpec() const override {
        if (!is_inverse_ || !std::is_same<T, int32_t>::value ||
            block_size_ == 0 || block_size_ % 32u != 0 ||
            block_size_ / 32u > fused::warp::kMaxWarpElemsPerLane) return {};
        return FusionSpec{FusionAccess::SegmentCodec, block_size_};
    }
    FusedOpDecl getInverseFusedOp() const override {
        if (!getInverseFusionSpec().fusable()) return {};
        FusedOpDecl d;
        d.strategy       = FusionStrategy::WarpRegister;
        // Inverse decode reads the true archive format: outlier -> adaptive 2-byte
        // meta, plain -> PlainRateCoder 1-byte meta. (fused_coder_ is a
        // forward-only encode choice; the archive it emits is AdaptiveBitpack-
        // decodable, i.e. the adaptive path.)
        d.op_name        = outlier_selection_ ? "AdaptiveBitpackCoder" : "PlainRateCoder";
        d.include_header = "fused/fused_block/warp_fusion.cuh";
        d.elems_per_lane = block_size_ / 32u;
        // Inverse always decodes the true AdaptiveBitpack-format archive regardless of
        // what fused_coder_ was on the forward encode side, so the TI decode policy is
        // always known here (no setFusedCoder ambiguity, unlike the forward decl above).
        d.ti_op_name     = outlier_selection_ ? "ThreadFixedRateCoderN" : "ThreadPlainRateCoderN";
        return d;
    }
    size_t getFusedInverseElementCount() const override { return num_elements_; }
 
    EncodingOracleDecl getEncodingOracle() const override {
        if (is_inverse_ || block_size_ == 0) return {};
        EncodingOracleDecl d;
        d.kind = outlier_selection_
            ? EncodingOracleKind::AdaptiveFixedRateBitpack
            : EncodingOracleKind::PlainFixedRateBitpack;
        d.op_name = outlier_selection_ ? "AdaptiveBitpackCoder" : "PlainBitpackCoder";
        d.include_header = "fused/fused_block/warp_fusion.cuh";
        d.input_data_type = static_cast<uint8_t>(getElementDataType());
        d.unit_elems = block_size_;
        d.exact = true;
        d.additive = true;
        return d;
    }
 
    void setFusedCoder(std::string name) { fused_coder_ = std::move(name); }
    const std::string& getFusedCoder() const { return fused_coder_; }
 
    void setFusedResult(size_t num_elements, size_t archive_bytes) {
        num_elements_       = num_elements;
        actual_output_size_ = archive_bytes;
    }
    void setFusedArchiveResult(size_t archive_bytes, size_t orig_bytes) override {
        setFusedResult(orig_bytes / sizeof(T), archive_bytes);
    }
 
    void setOutlierSelection(bool enable) { outlier_selection_ = enable; }
    bool getOutlierSelection() const { return outlier_selection_; }
 
    // ── Execution ──────────────────────────────────────────────────────────
    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;
 
    // ── Metadata ───────────────────────────────────────────────────────────
    std::string getName() const override { return "AdaptiveBitpack"; }
    size_t getNumInputs()  const override { return 1; }
    size_t getNumOutputs() const override { return 1; }
 
    std::vector<size_t> estimateOutputSizes(
        const std::vector<size_t>& input_sizes
    ) const override;
 
    size_t estimateScratchBytes(
        const std::vector<size_t>& input_sizes
    ) const override;
 
    std::unordered_map<std::string, size_t>
    getActualOutputSizesByName() const override {
        return {{"output", actual_output_size_}};
    }
    size_t getActualOutputSize(int index) const override {
        return (index == 0) ? actual_output_size_ : 0;
    }
 
    uint16_t getStageTypeId() const override {
        return static_cast<uint16_t>(StageType::ADAPTIVE_BITPACK);
    }
 
    // Forward: signed codes -> uint8 archive. Inverse: archive -> signed codes.
    uint8_t getOutputDataType(size_t) const override {
        return static_cast<uint8_t>(is_inverse_ ? getElementDataType()
                                                 : DataType::UINT8);
    }
    uint8_t getInputDataType(size_t) const override {
        return static_cast<uint8_t>(is_inverse_ ? DataType::UINT8
                                                 : getElementDataType());
    }
 
    // ── Serialization ──────────────────────────────────────────────────────
    size_t serializeHeader(size_t, uint8_t* buf, size_t max_size) const override {
        if (max_size < sizeof(AdaptiveBitpackConfig)) return 0;
        AdaptiveBitpackConfig cfg;
        cfg.data_type         = getElementDataType();
        cfg.outlier_selection = outlier_selection_ ? uint8_t{1} : uint8_t{0};
        cfg.block_size        = block_size_;
        cfg.num_elements      = static_cast<uint64_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(AdaptiveBitpackConfig))
            throw std::runtime_error("AdaptiveBitpackStage: header too small");
        AdaptiveBitpackConfig cfg;
        std::memcpy(&cfg, buf, sizeof(cfg));
        block_size_        = cfg.block_size ? cfg.block_size : 32u;
        num_elements_      = static_cast<size_t>(cfg.num_elements);
        outlier_selection_ = (cfg.outlier_selection != 0);
    }
    size_t getMaxHeaderSize(size_t) const override {
        return sizeof(AdaptiveBitpackConfig);
    }
 
    void saveState() override {
        saved_block_size_       = block_size_;
        saved_num_elements_     = num_elements_;
        saved_actual_size_      = actual_output_size_;
        saved_outlier_select_   = outlier_selection_;
    }
    void restoreState() override {
        block_size_         = saved_block_size_;
        num_elements_       = saved_num_elements_;
        actual_output_size_ = saved_actual_size_;
        outlier_selection_  = saved_outlier_select_;
    }
 
    size_t getNumElements() const { return num_elements_; }
 
private:
    bool     is_inverse_         = false;
    uint32_t block_size_         = 32;
    bool     outlier_selection_  = false;
    size_t   num_elements_       = 0;
    size_t   actual_output_size_ = 0;
    std::string fused_coder_ = "AdaptiveBitpackCoder";
 
    // Forward-path persistent scratch (kept alive across execute() so the
    // compressed-size readback can be deferred to postStreamSync(), and so no
    // allocation happens inside a captured graph replay). Grown lazily when a
    // larger input is seen; freed in the destructor. Pool-managed (persistent).
    uint32_t*   d_cost_          = nullptr;  
    uint32_t*   d_offset_        = nullptr;  
    size_t      scratch_blocks_  = 0;        
    MemoryPool* scratch_pool_    = nullptr;  
    size_t      fwd_num_blocks_  = 0;        
    size_t      fwd_meta_region_ = 0;        
 
    uint32_t saved_block_size_   = 32;
    bool     saved_outlier_select_ = false;
    size_t   saved_num_elements_ = 0;
    size_t   saved_actual_size_  = 0;
 
    static DataType getElementDataType() {
        if (std::is_same<T, int16_t>::value) return DataType::INT16;
        return DataType::INT32;
    }
};
 
extern template class AdaptiveBitpackStage<int16_t>;
extern template class AdaptiveBitpackStage<int32_t>;
 
} // namespace fz