spb/forge Public MIT

1// Author: Simon-Pierre Boucher — contact@spboucher.ai2#pragma once34#include "core/tensor.h"56#include <cstdint>78namespace MTL {9class CommandBuffer;10class ComputeCommandEncoder;11}1213// Metal dispatch wrappers over the batched execution model (RESEARCH.md §4):14// ops encode into one long-lived serial compute encoder inside one command15// buffer; nothing runs until the caller reaches a readback boundary and16// calls sync(). Serial dispatch type means dispatch N+1 sees dispatch N's17// writes — no barriers, no per-op waits.18//19// Contract: tensors passed to these functions must stay alive until sync()20// returns (the training loop and tests naturally satisfy this; the21// allocator-level retire list lands with M4's trainer).22namespace forge::metal {2324class Stream {25public:26    static Stream& get();2728    // Current encoder (lazily opens a command buffer + serial encoder).29    MTL::ComputeCommandEncoder* encoder();3031    // Encoder for a run of MUTUALLY INDEPENDENT dispatches. Metal's default32    // serial encoder orders every dispatch against the previous one, which33    // wastes the GPU when the work is genuinely parallel — measured 15x on a34    // batch of small independent dispatches (RESEARCH.md 4b). Switching35    // dispatch type ends the current encoder and opens a new one in the same36    // command buffer; Metal orders tracked resources across that boundary, so37    // the switch itself acts as the barrier.38    //39    // Caller's contract: everything encoded between two switches must be40    // free of read-after-write dependencies on each other.41    MTL::ComputeCommandEncoder* concurrent_encoder();4243    // Makes encoder() hand back a concurrent encoder; see ConcurrentRegion.44    void set_concurrent(bool on);4546    // Readback boundary: end encoding, commit, wait. Returns immediately if47    // nothing is pending. GPU time of the completed buffer is accumulated48    // into gpu_seconds().49    void sync();5051    double gpu_seconds() const { return gpu_seconds_; }5253private:54    Stream() = default;55    MTL::ComputeCommandEncoder* encoder_of(int dispatch_type);5657    MTL::CommandBuffer* cmd_ = nullptr;58    MTL::ComputeCommandEncoder* enc_ = nullptr;59    int dispatch_type_ = -1;60    bool concurrent_ = false;61    double gpu_seconds_ = 0.0;62};6364inline void sync() { Stream::get().sync(); }6566// RAII: inside this scope, ops encode into a CONCURRENT compute encoder.67// Everything encoded in the region must be mutually independent (no68// read-after-write between them); dependencies across the region boundary are69// fine because switching encoder kind ends the encoder, and Metal orders70// tracked resources across encoders in a command buffer.71struct ConcurrentRegion {72    ConcurrentRegion() { Stream::get().set_concurrent(true); }73    ~ConcurrentRegion() { Stream::get().set_concurrent(false); }74    ConcurrentRegion(const ConcurrentRegion&) = delete;75    ConcurrentRegion& operator=(const ConcurrentRegion&) = delete;76};7778// ---- f32 forward ops (parity-tested vs forge::cpu) -------------------------79enum class MatmulKernel { Auto, Naive, Tiled, Simdgroup };8081void matmul(const Tensor& a, const Tensor& b, Tensor& c,82            bool transpose_a = false, bool transpose_b = false,83            bool accumulate = false, MatmulKernel kernel = MatmulKernel::Auto);8485void add(const Tensor& a, const Tensor& b, Tensor& out);86void mul(const Tensor& a, const Tensor& b, Tensor& out);87void scale(const Tensor& a, float s, Tensor& out);88void add_bias(const Tensor& x, const Tensor& bias, Tensor& out);89void silu(const Tensor& x, Tensor& out);90void gelu(const Tensor& x, Tensor& out);91void relu2(const Tensor& x, Tensor& out);                         // max(x,0)^292void relu2_backward(const Tensor& x, const Tensor& dout, Tensor& dx);93void softcap(const Tensor& x, float cap, Tensor& out);            // cap*tanh(x/cap)94void softcap_backward(const Tensor& y, const Tensor& dout, float cap, Tensor& dx);9596void softmax(const Tensor& x, Tensor& out);                      // rows = last dim97void rmsnorm(const Tensor& x, const Tensor& w, float eps, Tensor& out);98void layernorm(const Tensor& x, const Tensor& w, const Tensor& b, float eps, Tensor& out);99100// ---- f32 backward / training ops (ACCUMULATE into d* outputs) --------------101void accumulate(Tensor& dst, const Tensor& src);          // dst += src102void axpy(Tensor& dst, const Tensor& src, const Tensor& s); // dst += src * s[0]103void silu_backward(const Tensor& x, const Tensor& dout, Tensor& dx);104void gelu_backward(const Tensor& x, const Tensor& dout, Tensor& dx);105void add_bias_backward(const Tensor& dout, Tensor& dbias);106void rmsnorm_backward(const Tensor& x, const Tensor& w, float eps,107                      const Tensor& dout, Tensor& dx, Tensor& dw);108void layernorm_backward(const Tensor& x, const Tensor& w, float eps,109                        const Tensor& dout, Tensor& dx, Tensor& dw, Tensor& db);110111// freqs: [head_dim/2] per-pair inverse frequencies, host-precomputed112// (fixed theta, llama3 rope-scaling, per-layer theta — all just tables).113void rope(const Tensor& x, int64_t n_heads, const Tensor& freqs, int64_t pos_offset,114          Tensor& out);115void rope_backward(const Tensor& dout, int64_t n_heads, const Tensor& freqs,116                   int64_t pos_offset, Tensor& dx);117118void embedding(const Tensor& weight, const Tensor& ids, Tensor& out); // ids i32119void embedding_backward(const Tensor& ids, const Tensor& dout, Tensor& dweight);120121// window > 0: sliding-window attention; attn_softcap > 0: cap·tanh pre-softmax122// (unfused path only — the fused kernels don't support softcap).123void attention(const Tensor& q, const Tensor& k, const Tensor& v,124               int64_t n_heads, int64_t n_kv_heads, bool causal, float scale,125               Tensor& out, Tensor* probs_out, int64_t window = 0,126               float attn_softcap = 0.0f);127void attention_backward(const Tensor& q, const Tensor& k, const Tensor& v,128                        const Tensor& probs, const Tensor& out, const Tensor& dout,129                        int64_t n_heads, int64_t n_kv_heads, float scale,130                        Tensor& dq, Tensor& dk, Tensor& dv, float attn_softcap = 0.0f);131132// Fused (flash) attention: stores no T x T probabilities, only the per-row133// logsumexp `lse` [B, n_heads, T] that the backward re-expands. Supported for134// a fixed set of head_dims — flash_supported() reports which.135bool flash_supported(int64_t head_dim);136// Scalar: one thread per query row. MMA: simdgroup_matrix 8x8 tiles. Same137// outputs; Auto picks MMA where supported.138enum class FlashKernel { Auto, Scalar, MMA };139// window > 0 requires the Scalar kernel (Auto routes there automatically);140// out-of-window KV blocks are skipped, so cost scales with the window.141void flash_attention(const Tensor& q, const Tensor& k, const Tensor& v,142                     int64_t n_heads, int64_t n_kv_heads, bool causal, float scale,143                     Tensor& out, Tensor& lse, FlashKernel kernel = FlashKernel::Auto,144                     int64_t window = 0);145void flash_attention_backward(const Tensor& q, const Tensor& k, const Tensor& v,146                              const Tensor& out, const Tensor& lse, const Tensor& dout,147                              int64_t n_heads, int64_t n_kv_heads, bool causal,148                              float scale, Tensor& dq, Tensor& dk, Tensor& dv,149                              FlashKernel kernel = FlashKernel::Auto,150                              int64_t window = 0);151152// losses: [N] per-row buffer; loss_out: [1] mean over n_valid. dlogits153// optional (accumulated). ids must be i32.154void cross_entropy(const Tensor& logits, const Tensor& targets, int64_t n_valid,155                   Tensor& losses, Tensor& loss_out, Tensor* dlogits);156157// ---- QAT / MoE --------------------------------------------------------------158// Per-row fake quantization of a 2-D weight; mode 0 = int8 (absmax/127),159// mode 1 = ternary (absmean, BitNet-style). Backward is STE — no kernel.160void fake_quant(const Tensor& w, int mode, Tensor& out);161// dx += p ∘ (dout − dot(dout, p)) per row; thread-per-row, small last dims.162void softmax_backward(const Tensor& p, const Tensor& dout, Tensor& dx);163void sigmoid(const Tensor& x, Tensor& out);164void sigmoid_backward(const Tensor& y, const Tensor& dout, Tensor& dx);165// Keep top-k per row of p [N,E] — selected by p+bias, gated by p alone;166// norm renormalizes kept gates to sum 1.167void topk_renorm(const Tensor& p, const Tensor& bias, int64_t k, bool norm,168                 Tensor& out);169void topk_renorm_backward(const Tensor& p, const Tensor& bias, const Tensor& dout,170                          int64_t k, bool norm, Tensor& dp);171// counts[e] += #rows with nonzero gate for expert e172void expert_counts(const Tensor& gates, Tensor& counts);173// out[i,:] = x[i,:] * gates[i,e]; the _accumulate variant does dst += (dx path).174void row_scale(const Tensor& x, const Tensor& gates, int64_t e, Tensor& out);175void row_scale_accumulate(const Tensor& x, const Tensor& gates, int64_t e, Tensor& dst);176// dgates[i,e] += dot(dout[i,:], x[i,:])177void row_scale_gate_backward(const Tensor& dout, const Tensor& x, int64_t e,178                             Tensor& dgates);179180// Fused optimizer update for one tensor; all state f32.181void adamw_step(Tensor& w, const Tensor& g, Tensor& m, Tensor& v,182                float lr, float beta1, float beta2, int64_t t, float eps, float wd,183                float grad_scale);184// out[0] = sum(x^2) — single-threadgroup reduce (fine for M4).185void sumsq(const Tensor& x, Tensor& out);186// out[0] = sum(x) * mul187void sum(const Tensor& x, Tensor& out, float mul);188189} // namespace forge::metal190