device_backend.hpp

Go to the documentation of this file.

 
#pragma once
 
#include <cstddef>
#include <cstdint>
 
namespace pulsatrix {
 
enum class DeviceType {
    Cpu,
    Cuda,
    Hip
};
 
enum class CopyDirection {
    HostToDevice,
    DeviceToHost,
    DeviceToDevice,
    HostToHost
};
 
enum class ElementwiseOp {
    Relu,
    Neg,
    Tanh,     
    Sigmoid,  
    Silu,     
    Exp       
};
 
enum class LrpGate {
    None,      
    Positive,  
    Negative   
};
 
enum class LogicOp {
    ConjunctionForward,
    ConjunctionBackward,
    ConjunctionLrp,
    DisjunctionForward,
    DisjunctionBackward,
    DisjunctionLrp
};
 
enum class RecurrentCellOp { RnnBackward, LstmForward, LstmBackward, LstmLrp, GruBackward, GruLrp };
 
struct RecurrentCellArgs {
    const float* in[10] = {};
    float* out[6] = {};
    float eps = 0.0f;  
};
 
enum class SsmPassOp {
    MambaForward,
    MambaBackward,
    MambaGradBC,
    MambaLrp,
    RwkvTokenShift,
    RwkvForward,
    RwkvBackward,
    RwkvShiftBackward,
    RwkvLrp,
    RwkvShiftLrp,
    StabilizedDiv,
    RetnetForward,
    RetnetStateGrad,
    RetnetGradQK,
    RetnetGradV,
    RetnetScores,
    RetnetReadout,
    RetnetLrpInput,
    ReverseTimeSum
};
 
struct SsmPassArgs {
    const float* in[12] = {};
    float* out[10] = {};
    int64_t n = 0;
    int64_t l = 0;
    int64_t d = 0;
    int64_t s = 0;
    float eps = 0.0f;    
    float gamma = 0.0f;  
};
 
struct ConvGeometry {
    size_t kh;
    size_t kw;
    size_t stride_h = 1;
    size_t stride_w = 1;
    size_t pad_h = 0;
    size_t pad_w = 0;
};
 
enum class RlRowOp { DqnLoss, DqnGrad, PgLoss, PgGrad, PpoLoss, PpoGrad, DqnTarget, PolyakBlend };
 
struct RlRowArgs {
    const float* in[5] = {};
    float* out[4] = {};
    int64_t rows = 0;
    int64_t cols = 0;
    float scale = 0.0f;  
    float lower = 0.0f;  
    float upper = 0.0f;  
    float gamma = 0.0f;  
    float tau = 0.0f;    
};
 
class DeviceBackend {
public:
    virtual ~DeviceBackend() = default;
 
    [[nodiscard]] virtual DeviceType device() const noexcept = 0;
 
    [[nodiscard]] virtual void* allocate(size_t bytes) = 0;
 
    virtual void free(void* ptr) noexcept = 0;
 
    virtual void copy(void* dst, const void* src, size_t bytes, CopyDirection dir) = 0;
 
    virtual void fill(void* ptr, float value, size_t n) = 0;
 
    virtual void gemm(const float* a, const float* b, float* out, size_t m, size_t k, size_t n) = 0;
 
    virtual void elementwise(ElementwiseOp op, const float* in, float* out, size_t n) = 0;
 
    virtual void add(const float* a, const float* b, float* out, size_t n) = 0;
 
    virtual void mul(const float* a, const float* b, float* out, size_t n) = 0;
 
    // ---- GPU-native-kernels Mission 1: primitives for device-resident training ----------
    // Every reduction below has a fixed summation order on every backend (no atomics), so
    // results are reproducible run to run.
 
    virtual void gemm_ex(const float* a, bool transpose_a, const float* b, bool transpose_b, float* out, size_t m,
                         size_t k, size_t n, float beta) = 0;
 
    virtual void column_sums(const float* in, float* out, size_t rows, size_t cols, float beta) = 0;
 
    virtual void add_row_vector(const float* in, const float* row, float* out, size_t rows, size_t cols) = 0;
 
    virtual void elementwise_backward(ElementwiseOp op, const float* x, const float* grad_out, float* grad_in,
                                      size_t n) = 0;
 
    virtual void axpby(float alpha, const float* x, float beta, const float* y, float* out, size_t n) = 0;
 
    [[nodiscard]] virtual float dot(const float* a, const float* b, size_t n) = 0;
 
    virtual void softmax_rows(const float* in, float* out, size_t rows, size_t cols) = 0;
 
    virtual void softmax_rows_backward(const float* y, const float* dy, float* dx, size_t rows, size_t cols) = 0;
 
    virtual void logsumexp_rows(const float* in, float* out, size_t rows, size_t cols) = 0;
 
    virtual void adam_step(float* param, const float* grad, float* m, float* v, size_t n, float lr, float beta1,
                           float beta2, float eps, float bias_correction1, float bias_correction2) = 0;
 
    // ---- GPU-native-kernels Mission 1b ---------------------------------------------------
 
    [[nodiscard]] virtual float sum(const float* in, size_t n) = 0;
 
    virtual void dropout_forward(const float* in, float* out, float* mask, size_t n, float p, float scale,
                                 uint64_t seed, uint64_t offset) = 0;
 
    virtual void bce_with_logits(const float* logits, const float* target, float* out, size_t n) = 0;
 
    virtual void bce_with_logits_grad(const float* logits, const float* target, float* grad, size_t n,
                                      float scale) = 0;
 
    // ---- GPU-native-kernels Mission 2: transformer building blocks -------------------------
    // The fused row operations below run the same per-row source on every backend (src/
    // row_math.hpp): CPU loops over rows, GPUs run one thread per row. row_std / row_rms /
    // log_prob are per-row device buffers (rows entries). Index buffers hold whole numbers as
    // floats, exact below 2^24.
 
    virtual void layer_norm_forward(const float* in, const float* gamma, const float* beta, float* xhat, float* out,
                                    float* row_std, size_t rows, size_t cols, float eps) = 0;
 
    virtual void layer_norm_backward(const float* grad_out, const float* gamma, const float* xhat,
                                     const float* row_std, float* grad_in, size_t rows, size_t cols) = 0;
 
    virtual void rms_norm_forward(const float* in, const float* gamma, float* out, float* row_rms, size_t rows,
                                  size_t cols, float eps) = 0;
 
    virtual void rms_norm_backward(const float* grad_out, const float* gamma, const float* in, const float* row_rms,
                                   float* grad_in, float* gamma_terms, size_t rows, size_t cols) = 0;
 
    virtual void rope_rotate(const float* in, const float* cos_table, const float* sin_table, float* out,
                             size_t num_slices, size_t seq_len, size_t head_dim, bool inverse) = 0;
 
    virtual void permute_0213(const float* in, float* out, size_t d0, size_t d1, size_t d2, size_t d3) = 0;
 
    virtual void gather_rows(const float* table, const float* indices, float* out, size_t count, size_t dim) = 0;
 
    virtual void scatter_add_rows(const float* src, const float* indices, float* table, size_t count, size_t dim) = 0;
 
    virtual void tanh_gaussian_forward(const float* mean, const float* log_std, const float* eps, float* action,
                                       float* std_cache, float* log_prob, size_t rows, size_t cols,
                                       float stabilizer, double half_log_two_pi) = 0;
 
    virtual void tanh_gaussian_backward(const float* action, const float* std_cache, const float* eps,
                                        const float* grad_action, const float* grad_log_prob, float* grad_mean,
                                        float* grad_log_std, size_t n, float stabilizer) = 0;
 
    // ---- GPU-native-kernels Mission 3: LRP rules and logic modules --------------------------
    // Shared per-output source in src/lrp_math.hpp (CPU loops, GPU one thread per output).
    // Every reduction is owned by one thread and runs in the original loop order:
    // deterministic, no atomics.
 
    virtual void lrp_linear(const float* x, const float* w, const float* z, const float* r, float* r_in, size_t rows,
                            size_t in_features, size_t out_features, float eps) = 0;
 
    virtual void lrp_residual_split(const float* a, const float* b, const float* r, float* r_a, float* r_b, size_t n,
                                    float eps) = 0;
 
    virtual void lrp_bilinear_elementwise(const float* a, const float* b, const float* r, float* r_out, size_t n,
                                          float eps) = 0;
 
    virtual void lrp_bilinear_matmul(const float* a, const float* b, const float* o, const float* r_o, float* r_a,
                                     float* r_b, size_t slices, size_t m, size_t p, size_t q, float eps,
                                     bool b_transposed) = 0;
 
    virtual void lrp_softmax_rows(const float* x, const float* y, const float* r, float* r_in, size_t rows,
                                  size_t cols) = 0;
 
    virtual void lrp_rope(const float* x, const float* y, const float* r, const float* cos_table,
                          const float* sin_table, float* r_in, size_t slices, size_t seq_len, size_t head_dim,
                          float eps) = 0;
 
    virtual void logic_pointwise(LogicOp op, int norm, const float* a, const float* b, const float* g_or_r,
                                 const float* y, float* out_a, float* out_b, size_t n, float eps) = 0;
 
    virtual void aggregator_forward(const float* x, float* mean_pow, float* out, size_t n, size_t cols, float p) = 0;
    virtual void aggregator_backward(const float* x, const float* mean_pow, const float* grad_out, float* grad_in,
                                     size_t n, size_t cols, float p) = 0;
    virtual void aggregator_lrp(const float* x, const float* mean_pow, const float* r_out, float* r_in, size_t n,
                                size_t cols, float p, float eps) = 0;
 
    // ---- GPU-native-kernels Mission 4: convolution, pooling, spatial norms -----------------
    // Shared per-output source in src/cnn_math.hpp; deterministic, no atomics.
 
    virtual void im2col(const float* in, float* col, size_t n, size_t c, size_t h, size_t w,
                        const ConvGeometry& geometry) = 0;
 
    virtual void col2im_add(const float* col, float* out, size_t n, size_t c, size_t h, size_t w,
                            const ConvGeometry& geometry) = 0;
 
    virtual void add_channel_vector(const float* in, const float* vec, float* out, size_t n, size_t c, size_t inner) =
                                    0;
 
    virtual void lrp_conv(const float* col, const float* kernel, const float* pre_bias, const float* r, float* r_col,
                          size_t n, size_t out_channels, size_t p, size_t q, float eps) = 0;
 
    // ---- LRP-rules campaign Mission 2: Zennit-compatible Gamma / AlphaBeta / ZBox ----------
    virtual void lrp_stabilized_divide(const float* r, const float* denom, const float* gate, float* out, size_t n,
                                       float eps, LrpGate gate_mode) = 0;
 
    virtual void max_pool_forward(const float* in, float* out, float* argmax, size_t planes, size_t h, size_t w, size_t
                                  kh, size_t kw) = 0;
 
    virtual void max_unpool(const float* src, const float* argmax, float* dst, size_t planes, size_t h, size_t w, size_t
                            kh, size_t kw) = 0;
 
    virtual void avg_pool_forward(const float* in, float* out, size_t planes, size_t h, size_t w, size_t kh, size_t kw)
                                  = 0;
 
    virtual void avg_pool_backward(const float* grad_out, float* grad_in, size_t planes, size_t h, size_t w, size_t kh,
                                   size_t kw) = 0;
 
    virtual void lrp_avg_pool(const float* x, const float* r, float* r_in, size_t planes, size_t h, size_t w, size_t kh,
                              size_t kw, float eps) = 0;
 
    virtual void batch_norm_forward(const float* in, const float* gamma, const float* beta, float* xhat, float* out,
                                    float* channel_std, size_t n, size_t c, size_t spatial, float eps) = 0;
 
    virtual void batch_norm_backward(const float* grad_out, const float* gamma, const float* xhat, const float*
                                     channel_std, float* grad_in, float* gamma_grad, float* beta_grad, size_t n, size_t
                                     c, size_t spatial) = 0;
 
    virtual void batch_norm_update_running(const float* in, float* running_mean, float* running_var, size_t n,
                                           size_t c, size_t spatial, float momentum) = 0;
 
    virtual void batch_norm_eval_forward(const float* in, const float* gamma, const float* beta,
                                         const float* running_mean, const float* running_var, float* xhat, float* out,
                                         float* channel_std, size_t n, size_t c, size_t spatial, float eps) = 0;
 
    virtual void batch_norm_eval_backward(const float* grad_out, const float* gamma, const float* xhat,
                                          const float* channel_std, float* grad_in, float* gamma_grad,
                                          float* beta_grad, size_t n, size_t c, size_t spatial) = 0;
 
    virtual void group_norm_forward(const float* in, const float* gamma, const float* beta, float* xhat, float* out,
                                    float* group_std, size_t n, size_t c, size_t spatial, size_t num_groups, float eps)
                                    = 0;
 
    virtual void group_norm_backward(const float* grad_out, const float* gamma, const float* xhat, const float*
                                     group_std, float* grad_in, float* gamma_grad, float* beta_grad, size_t n, size_t c,
                                     size_t spatial, size_t num_groups) = 0;
 
    // ---- GPU-native-kernels Mission 5: recurrent networks ------------------------------------
 
    virtual void copy_2d(float* dst, size_t dst_stride, const float* src, size_t src_stride, size_t rows,
                         size_t cols) = 0;
 
    virtual void accumulate_rows(const float* in, float* out, size_t rows, size_t cols) = 0;
 
    virtual void recurrent_cell(RecurrentCellOp op, const RecurrentCellArgs& args, size_t n) = 0;
 
    virtual void gru_lrp_hprev(const float* h_prev, const float* w_hn, const float* hn, const float* r_term_b,
                               const float* direct, float* r_hprev, size_t rows, size_t hidden, float eps) = 0;
 
    // ---- GPU-native-kernels Mission 6: state-space models -------------------------------------
    // Shared per-lane source in src/ssm_math.hpp; deterministic, no atomics.
 
    virtual void ssm_pass(SsmPassOp op, const SsmPassArgs& args) = 0;
 
    // ---- GPU-native-kernels Mission 7: reinforcement learning -----------------------------------
    // Shared per-row source in src/rl_math.hpp; deterministic, no atomics.
 
    virtual void rl_rows(RlRowOp op, const RlRowArgs& args) = 0;
 
    // ---- FND-3: selection ------------------------------------------------------------------
 
    virtual void top_k_rows(const float* in, float* values, float* indices, size_t rows, size_t cols, size_t k,
                            bool largest) = 0;
};
 
}  // namespace pulsatrix