tiled_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 TiledLorenzoConfig {
    DataType data_type;   
    uint8_t  ndim;        
    uint8_t  tile_x;      
    uint8_t  tile_y;      
    uint8_t  tile_z;      
    uint8_t  reserved[3]; 
    uint32_t dim_x;       
    uint32_t dim_y;       
    uint32_t dim_z;       
 
    TiledLorenzoConfig()
        : data_type(DataType::INT32), ndim(2),
          tile_x(8), tile_y(8), tile_z(1), reserved{0, 0, 0},
          dim_x(0), dim_y(1), dim_z(1) {}
};
static_assert(sizeof(TiledLorenzoConfig) <= FZM_STAGE_CONFIG_SIZE,
              "TiledLorenzoConfig must fit in FZM_STAGE_CONFIG_SIZE");
 
template<typename T>
class TiledLorenzoStage : public Stage {
    static_assert(std::is_same_v<T, int16_t> || std::is_same_v<T, int32_t>,
                  "TiledLorenzoStage: T must be int16_t or int32_t.");
public:
    TiledLorenzoStage() = default;
    ~TiledLorenzoStage() override = default;
 
    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 {
        if (!dims_pinned_) dims_ = dims;
    }
    void setDims(size_t x, size_t y = 1, size_t z = 1) {
        if (!dims_pinned_) dims_ = {x, y, z};
    }
    std::array<size_t, 3> getDims() const { return dims_; }
    bool hasDimsOverride() const { return dims_pinned_; }
 
    void setDimsOverride(size_t x, size_t y, size_t z) {
        dims_ = {x, y, z};
        dims_pinned_ = true;
    }
 
    void setTileShape(uint32_t tx, uint32_t ty = 1, uint32_t tz = 1) {
        auto chk = [](uint32_t v, const char* nm) {
            if (v > 255)
                throw std::invalid_argument(
                    std::string("TiledLorenzoStage::setTileShape: ") + nm
                    + " must be in [0, 255], got " + std::to_string(v));
        };
        chk(tx, "tx"); chk(ty, "ty"); chk(tz, "tz");
        const uint32_t prod = (tx ? tx : 1) * (ty ? ty : 1) * (tz ? tz : 1);
        if (prod > 1024)
            throw std::invalid_argument(
                "TiledLorenzoStage::setTileShape: tx*ty*tz must be in [1, 1024], got "
                + std::to_string(prod));
        tile_ = {tx, ty, tz};
        tile_set_ = true;
    }
    std::array<uint32_t, 3> getTileShape() const { return effectiveTile(); }
 
    uint32_t getTileElems() const {
        auto t = effectiveTile();
        return t[0] * t[1] * t[2];
    }
 
    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 "TiledLorenzo"; }
    size_t getNumInputs()  const override { return 1; }
    size_t getNumOutputs() const override { return 1; }
 
    // Block-local for fusion: forward, 2-D tile (tz == 1). The block a downstream
    // coder sees is one tile (tile_elems); the fused driver (cuSZp3) recomputes
    // the separable delta per element. tz > 1 is not fused yet.
    FusionSpec getFusionSpec() const override {
        auto t = effectiveTile();
        if (is_inverse_) return {};
        // Block-local for fusion; the block is one tile. 2-D (tz==1) and 3-D (tz>1)
        // both fuse — the warp op gate (getFusedOp) requires tile_elems==64 (EPL=2).
        return FusionSpec{FusionAccess::BlockLocal, t[0] * t[1] * t[2]};
    }
 
    FusedOpDecl getFusedOp() const override {
        const auto t = effectiveTile();
        if (is_inverse_ || t[0] * t[1] * t[2] != 64u) return {};   // only tile_elems==64 (EPL=2)
        const uint32_t dx = static_cast<uint32_t>(dims_[0]);
        const uint32_t dy = static_cast<uint32_t>(dims_[1]);
        const uint32_t dz = static_cast<uint32_t>(dims_[2]);
        const uint32_t tx = t[0], ty = t[1], tz = t[2];
        const uint32_t ntx = (dx + tx - 1u) / tx;
        const uint32_t nty = (dy + ty - 1u) / ty;
        const uint32_t ntz = (dz + tz - 1u) / tz;
        FusedOpDecl d;
        d.strategy       = FusionStrategy::WarpRegister;
        d.include_header = "fused/fused_block/warp_fusion.cuh";
        d.elems_per_lane = 2;
        if (tz == 1u) {   // 2-D
            d.op_name = "TiledLorenzo2DPredictor";
            d.n_ab    = static_cast<size_t>(ntx) * nty * tx * ty;
            fused::warp::TiledLorenzo2DParams p{0.0f, dx, dy, tx, ty, ntx};
            d.params.resize(sizeof(p)); std::memcpy(d.params.data(), &p, sizeof(p));
        } else {          // 3-D (PROTOTYPE)
            d.op_name = "TiledLorenzo3DPredictor";
            d.n_ab    = static_cast<size_t>(ntx) * nty * ntz * tx * ty * tz;
            fused::warp::TiledLorenzo3DParams p{0.0f, dx, dy, dz, tx, ty, tz, ntx, nty};
            d.params.resize(sizeof(p)); std::memcpy(d.params.data(), &p, sizeof(p));
        }
        return d;
    }
 
    std::vector<size_t> estimateOutputSizes(
        const std::vector<size_t>& input_sizes
    ) const override {
        // Forward: natural n -> padded tile-major (num_tiles * tile_elems).
        // Inverse: padded tile-major -> natural n.
        const size_t n = naturalElems(input_sizes);
        const size_t out_elems = is_inverse_ ? n : paddedElems(n);
        return {out_elems * sizeof(T)};
    }
 
    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::TILED_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(TiledLorenzoConfig))
            throw std::runtime_error("TiledLorenzoStage: header buffer too small");
        auto t = effectiveTile();
        TiledLorenzoConfig cfg;
        cfg.data_type = getElementDataType();
        cfg.ndim      = static_cast<uint8_t>(ndim());
        cfg.tile_x    = static_cast<uint8_t>(t[0]);
        cfg.tile_y    = static_cast<uint8_t>(t[1]);
        cfg.tile_z    = static_cast<uint8_t>(t[2]);
        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]);
        std::memcpy(buf, &cfg, sizeof(cfg));
        return sizeof(cfg);
    }
 
    void deserializeHeader(const uint8_t* buf, size_t size) override {
        if (size < sizeof(TiledLorenzoConfig))
            throw std::runtime_error("TiledLorenzoStage: header too small");
        TiledLorenzoConfig cfg;
        std::memcpy(&cfg, buf, sizeof(cfg));
        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;
        tile_ = {cfg.tile_x, cfg.tile_y, cfg.tile_z};
        tile_set_ = (cfg.tile_x != 0);
    }
 
    size_t getMaxHeaderSize(size_t /*output_index*/) const override {
        return sizeof(TiledLorenzoConfig);
    }
 
private:
    bool is_inverse_           = false;
    bool dims_pinned_          = false;  
    size_t actual_output_size_ = 0;
    std::array<size_t, 3>   dims_     = {0, 1, 1};
    std::array<uint32_t, 3> tile_     = {8, 8, 1};
    bool                    tile_set_ = false;
 
    std::array<uint32_t, 3> effectiveTile() const {
        if (tile_set_) {
            return {tile_[0] ? tile_[0] : 1u,
                    tile_[1] ? tile_[1] : 1u,
                    tile_[2] ? tile_[2] : 1u};
        }
        switch (ndim()) {
            case 3:  return {4, 4, 4};
            case 2:  return {8, 8, 1};
            default: return {64, 1, 1};
        }
    }
 
    size_t naturalElems(const std::vector<size_t>& input_sizes) const {
        if (dims_[0] > 0) return dims_[0] * dims_[1] * dims_[2];
        return input_sizes.empty() ? 0 : input_sizes[0] / sizeof(T);
    }
 
    size_t paddedElems(size_t natural_n) const {
        auto t = effectiveTile();
        const size_t dx = (dims_[0] > 0) ? dims_[0] : natural_n;
        const size_t dy = (dims_[0] > 0) ? dims_[1] : 1;
        const size_t dz = (dims_[0] > 0) ? dims_[2] : 1;
        if (dx == 0) return 0;
        const size_t ntx = (dx + t[0] - 1) / t[0];
        const size_t nty = (dy + t[1] - 1) / t[1];
        const size_t ntz = (dz + t[2] - 1) / t[2];
        return ntx * nty * ntz * (size_t)t[0] * t[1] * t[2];
    }
 
    static DataType getElementDataType() {
        if (std::is_same<T, int16_t>::value) return DataType::INT16;
        return DataType::INT32;
    }
};
 
extern template class TiledLorenzoStage<int16_t>;
extern template class TiledLorenzoStage<int32_t>;
 
} // namespace fz