include/ling3/gdn_step.h

#pragma once

#include <cstddef>
#include <cstdint>
#include <memory>
#include <span>
#include <vector>

namespace ling3 {

struct GdnState {
    std::vector<std::uint16_t> fp16;
    std::vector<float> fp32;
    bool fp32_valid = false;
};

struct GdnRunTimings {
    double stage_ms = 0.0;
    double npu_ms = 0.0;
    double collect_ms = 0.0;
    double total_ms = 0.0;
};

class GdnStep {
public:
    GdnStep(
        std::span<const std::byte> heads6_model,
        std::span<const std::byte> heads5_model);
    ~GdnStep();

    GdnStep(const GdnStep &) = delete;
    GdnStep & operator=(const GdnStep &) = delete;
    GdnStep(GdnStep &&) noexcept;
    GdnStep & operator=(GdnStep &&) noexcept;

    void Reset();
    GdnState SaveState();
    void RestoreState(const GdnState & state);
    GdnRunTimings Run(
        std::span<const float> query,
        std::span<const float> key,
        std::span<const float> value,
        std::span<const float> decay,
        std::span<const float> beta,
        std::span<float> output);
    GdnRunTimings RunBatch16(
        std::span<const float> query,
        std::span<const float> key,
        std::span<const float> value,
        std::span<const float> decay,
        std::span<const float> beta,
        std::span<float> output);
    GdnRunTimings RunBatchCpu(
        std::span<const float> query,
        std::span<const float> key,
        std::span<const float> value,
        std::span<const float> decay,
        std::span<const float> beta,
        std::span<float> output);
    bool has_batch16() const noexcept;
    std::size_t state_bytes() const noexcept;

private:
    struct Impl;
    std::unique_ptr<Impl> impl_;
};

} // namespace ling3