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 "fused/fused_block/warp_op_params.h"
#include <algorithm>
#include <array>
#include <cstdint>
#include <cstring>
#include <stdexcept>
#include <string>
#include <type_traits>
#include <unordered_map>
#include <vector>
 
namespace fz {
 
struct LorenzoConfig {
    DataType data_type;   
    uint8_t  ndim;        
    uint8_t  centering;   
    uint8_t  order;       
    uint32_t dim_x;       
    uint32_t dim_y;       
    uint32_t dim_z;       
    uint32_t block_size;  
 
    LorenzoConfig()
        : data_type(DataType::INT32), ndim(1), centering(0), order(1),
          dim_x(0), dim_y(1), dim_z(1), block_size(0) {}
};
static_assert(sizeof(LorenzoConfig) <= FZM_STAGE_CONFIG_SIZE,
              "LorenzoConfig must fit in FZM_STAGE_CONFIG_SIZE");
 
template<typename T>
class LorenzoStage : public Stage {
    static_assert(std::is_integral<T>::value && std::is_signed<T>::value,
                  "LorenzoStage requires a signed integer type");
public:
    LorenzoStage() = default;
 
    explicit LorenzoStage(uint32_t block_size, bool centering = false,
                          uint8_t order = 1)
        : block_size_(block_size), centering_(centering), order_(order) {
        if (block_size > 1024)
            throw std::invalid_argument(
                "LorenzoStage: block_size must be in [0, 1024], got "
                + std::to_string(block_size));
        if (centering && block_size == 0)
            throw std::invalid_argument(
                "LorenzoStage: centering requires block_size > 0");
        if (order != 1 && order != 2)
            throw std::invalid_argument(
                "LorenzoStage: order must be 1 or 2, got " + std::to_string(order));
        if (order == 2 && block_size == 0)
            throw std::invalid_argument(
                "LorenzoStage: order 2 requires block_size > 0");
    }
 
    void setInverse(bool inv) override { is_inverse_ = inv; }
    bool isInverse() const override    { return is_inverse_; }
 
    void setDims(const std::array<size_t, 3>& dims) override { dims_ = dims; }
    void setDims(size_t x, size_t y = 1, size_t z = 1) { dims_ = {x, y, z}; }
    std::array<size_t, 3> getDims() const { return dims_; }
 
    void setBlockSize(uint32_t n) {
        if (n > 1024)
            throw std::invalid_argument(
                "LorenzoStage::setBlockSize: n must be in [0, 1024], got "
                + std::to_string(n));
        block_size_ = n;
    }
    uint32_t getBlockSize() const { return block_size_; }
 
    FusionSpec getFusionSpec() const override {
        if (isInverse() || centeringActive() || block_size_ == 0) return {};
        return FusionSpec{FusionAccess::RegionLocal, block_size_};
    }
 
    FusedOpDecl getFusedOp() const override {
        // 1-D block-reset Lorenzo fuses for any block_size that is a multiple of
        // 32 up to the warp-register cap (block_size = 32*EPL). EPL == 1 is
        // cuSZp2; EPL == 4 (block 128) is SZp's composed chain. The device policy
        // runs the delta chain across the whole block for EPL > 1. Centering is
        // excluded via getFusionSpec() above (checked again here so a caller who
        // calls getFusedOp() directly, bypassing getFusionSpec(), can't get a
        // false-positive fused op either).
        if (!getFusionSpec().fusable() || block_size_ % 32u != 0 ||
            block_size_ / 32u > fused::warp::kMaxWarpElemsPerLane) return {};
        const uint32_t epl = block_size_ / 32u;
        FusedOpDecl d;
        d.strategy       = FusionStrategy::WarpRegister;
        d.op_name        = "Lorenzo1DPredictor";
        d.include_header = "fused/fused_block/warp_fusion.cuh";
        d.elems_per_lane = epl;
        d.n_ab           = 0;
        // ThreadLorenzo1DPredictor's predict(base, d[32]) is written for exactly one
        // 32-element block; EPL > 1 (block 64/96/128, e.g. SZp's composed chain) has
        // no TI policy today, so only declare it at EPL == 1.
        if (epl == 1u) d.ti_op_name = "ThreadLorenzo1DPredictor";
        fused::warp::Lorenzo1DParams p{0.0f, epl};
        d.params.resize(sizeof(p));
        std::memcpy(d.params.data(), &p, sizeof(p));
        return d;
    }
 
    FusionSpec getInverseFusionSpec() const override {
        if (!isInverse() || centeringActive() || block_size_ == 0 ||
            block_size_ % 32u != 0 ||
            block_size_ / 32u > fused::warp::kMaxWarpElemsPerLane) return {};
        return FusionSpec{FusionAccess::RegionLocal, block_size_};
    }
    FusedOpDecl getInverseFusedOp() const override {
        if (!getInverseFusionSpec().fusable()) return {};
        FusedOpDecl d;
        d.strategy       = FusionStrategy::WarpRegister;
        d.op_name        = "Lorenzo1DPredictor";
        d.include_header = "fused/fused_block/warp_fusion.cuh";
        d.elems_per_lane = block_size_ / 32u;
        // ti_op_name intentionally left empty: measured 2026-09-08, TI decode's
        // serial de-delta chain (ThreadLorenzo1DPredictor::unpredict_and_write) is a
        // real regression vs the existing warp-cooperative decode (NYX/temperature
        // cuszp2 outlier: 374 GB/s vs ~651-897 before). Root-caused via `sudo ncu`
        // (same investigation as TiledLorenzoStage's delta modes): registers/
        // occupancy/shared-mem are near-identical to the fast identity decode, but
        // average latency per issued instruction is ~3x higher (52 vs 18 cycles) --
        // the chase is a genuine serial prefix sum (exact-order dependent), unlike
        // cost()/pack()'s reduction-shaped ops, which is why compress wins everywhere
        // but decode only wins where there is no chase. See
        // reports/w2_ti_decode_design.md UPDATE 2. Forward (compress) is unaffected.
        return d;
    }
 
    void setCentering(bool enable) { centering_ = enable; }
    bool getCentering() const { return centering_; }
 
    void setOrder(uint8_t k) {
        if (k != 1 && k != 2)
            throw std::invalid_argument(
                "LorenzoStage::setOrder: order must be 1 or 2, got " + std::to_string(k));
        order_ = k;
    }
    uint8_t getOrder() const { return order_; }
 
    int ndim() const {
        if (dims_[2] > 1) return 3;
        if (dims_[1] > 1) return 2;
        return 1;
    }
 
    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 "Lorenzo"; }
 
    size_t getNumInputs()  const override {
        return (is_inverse_ && centeringActive()) ? 2 : 1;
    }
    size_t getNumOutputs() const override {
        return (!is_inverse_ && centeringActive()) ? 2 : 1;
    }
 
    std::vector<std::string> getOutputNames() const override {
        if (centeringActive()) return {"output", "means"};
        return {"output"};
    }
 
    std::vector<size_t> estimateOutputSizes(
        const std::vector<size_t>& input_sizes
    ) const override {
        const size_t in = input_sizes.empty() ? 0 : input_sizes[0];
        if (!centeringActive()) return {in};
        return {in, numBlocks(in / sizeof(T)) * sizeof(T)};
    }
 
    std::unordered_map<std::string, size_t>
    getActualOutputSizesByName() const override {
        if (centeringActive())
            return {{"output", actual_output_size_}, {"means", actual_means_size_}};
        return {{"output", actual_output_size_}};
    }
 
    size_t getActualOutputSize(int index) const override {
        if (index == 0) return actual_output_size_;
        if (index == 1 && centeringActive()) return actual_means_size_;
        return 0;
    }
 
    uint16_t getStageTypeId() const override {
        return static_cast<uint16_t>(StageType::LORENZO);
    }
 
    uint8_t getOutputDataType(size_t /*output_index*/) const override {
        return static_cast<uint8_t>(getElementDataType());
    }
 
    uint8_t getInputDataType(size_t /*input_index*/) const override {
        return static_cast<uint8_t>(getElementDataType());
    }
 
    size_t serializeHeader(size_t /*output_index*/, uint8_t* buf, size_t max_size) const override {
        if (max_size < sizeof(LorenzoConfig))
            throw std::runtime_error("LorenzoStage: header buffer too small");
        LorenzoConfig cfg;
        cfg.data_type = getElementDataType();
        cfg.ndim      = static_cast<uint8_t>(ndim());
        cfg.dim_x     = static_cast<uint32_t>(dims_[0]);
        cfg.dim_y     = static_cast<uint32_t>(dims_[1]);
        cfg.dim_z     = static_cast<uint32_t>(dims_[2]);
        cfg.block_size = block_size_;
        cfg.centering  = centeringActive() ? 1u : 0u;
        cfg.order      = order_;
        std::memcpy(buf, &cfg, sizeof(LorenzoConfig));
        return sizeof(LorenzoConfig);
    }
 
    void deserializeHeader(const uint8_t* buf, size_t size) override {
        // Accept legacy 16-byte headers (no block_size field).
        constexpr size_t kMinSize = 16;
        if (size < kMinSize)
            throw std::runtime_error("LorenzoStage: header too small");
        LorenzoConfig cfg;  // default-constructed: block_size = 0
        std::memcpy(&cfg, buf, std::min(size, sizeof(LorenzoConfig)));
        int eff_ndim = (cfg.ndim == 0) ? 1 : static_cast<int>(cfg.ndim);
        dims_[0] = cfg.dim_x;
        dims_[1] = (eff_ndim >= 2) ? cfg.dim_y : 1;
        dims_[2] = (eff_ndim >= 3) ? cfg.dim_z : 1;
        block_size_ = (size >= sizeof(LorenzoConfig)) ? cfg.block_size : 0;
        // `centering` occupies a byte that legacy writers zeroed as `reserved`,
        // so old archives decode as centering-off without a version bump.
        centering_  = (cfg.centering != 0);
        // `order` reuses a byte legacy writers zeroed as `reserved`; 0 reads as
        // first order, so pre-LZ2 archives decode unchanged.
        order_      = (cfg.order == 2) ? 2u : 1u;
    }
 
    size_t getMaxHeaderSize(size_t /*output_index*/) const override {
        return sizeof(LorenzoConfig);
    }
 
private:
    void executeBlockMode(
        fz::stream_t stream,
        const T* in, T* out, size_t n, size_t byte_size,
        const std::vector<void*>& inputs,
        const std::vector<void*>& outputs);
 
    void executeNDMode(
        fz::stream_t stream,
        const T* in, T* out, size_t n, size_t byte_size);
 
    bool is_inverse_         = false;
    size_t actual_output_size_ = 0;
    size_t actual_means_size_  = 0;
    std::array<size_t, 3> dims_ = {0, 1, 1};
    uint32_t block_size_ = 0;  
    bool centering_ = false;   
    uint8_t order_ = 1;        
 
    bool centeringActive() const { return centering_ && block_size_ > 0; }
 
    size_t numBlocks(size_t n) const {
        return (block_size_ == 0) ? 0 : (n + block_size_ - 1) / block_size_;
    }
 
    static DataType getElementDataType() {
        if (std::is_same<T, int8_t>::value)  return DataType::INT8;
        if (std::is_same<T, int16_t>::value) return DataType::INT16;
        if (std::is_same<T, int32_t>::value) return DataType::INT32;
        if (std::is_same<T, int64_t>::value) return DataType::INT64;
        return DataType::INT32;
    }
};
 
extern template class LorenzoStage<int8_t>;
extern template class LorenzoStage<int16_t>;
extern template class LorenzoStage<int32_t>;
extern template class LorenzoStage<int64_t>;
 
// Kernel launcher declarations — defined in lorenzo_stage.cu.
 
template<typename T>
void launchLorenzoDeltaKernel1D(
    const T* d_input, T* d_output, size_t n, fz::stream_t stream,
    unsigned block_threads = 256);
 
template<typename T>
void launchLorenzoPrefixSumKernel1D(
    const T* d_input, T* d_output, size_t n, fz::stream_t stream,
    unsigned block_threads = 256);
 
template<typename T>
void launchLorenzoDeltaCentered1D(
    const T* d_input, T* d_output, T* d_means, size_t n, fz::stream_t stream,
    unsigned block_threads);
 
template<typename T>
void launchLorenzoSegmentedScan(
    const T* d_input, const T* d_means, T* d_output, size_t n, fz::stream_t stream,
    unsigned block_threads, int passes);
 
template<typename T>
void launchLorenzo2Delta1D(
    const T* d_input, T* d_output, T* d_means, size_t n, fz::stream_t stream,
    unsigned block_threads);
 
template<typename T>
void launchLorenzoDeltaKernel2D(
    const T* d_input, T* d_output, size_t nx, size_t ny, fz::stream_t stream);
 
template<typename T>
void launchLorenzoPrefixSumKernel2D(
    const T* d_input, T* d_output, size_t nx, size_t ny, fz::stream_t stream);
 
template<typename T>
void launchLorenzoDeltaKernel3D(
    const T* d_input, T* d_output, size_t nx, size_t ny, size_t nz, fz::stream_t stream);
 
template<typename T>
void launchLorenzoPrefixSumKernel3D(
    const T* d_input, T* d_output, size_t nx, size_t ny, size_t nz, fz::stream_t stream);
 
} // namespace fz