// Author: Simon-Pierre Boucher — contact@spboucher.ai #pragma once #include "core/tensor.h" #include namespace MTL { class CommandBuffer; class ComputeCommandEncoder; } // Metal dispatch wrappers over the batched execution model (RESEARCH.md §4): // ops encode into one long-lived serial compute encoder inside one command // buffer; nothing runs until the caller reaches a readback boundary and // calls sync(). Serial dispatch type means dispatch N+1 sees dispatch N's // writes — no barriers, no per-op waits. // // Contract: tensors passed to these functions must stay alive until sync() // returns (the training loop and tests naturally satisfy this; the // allocator-level retire list lands with M4's trainer). namespace forge::metal { class Stream { public: static Stream& get(); // Current encoder (lazily opens a command buffer + serial encoder). MTL::ComputeCommandEncoder* encoder(); // Encoder for a run of MUTUALLY INDEPENDENT dispatches. Metal's default // serial encoder orders every dispatch against the previous one, which // wastes the GPU when the work is genuinely parallel — measured 15x on a // batch of small independent dispatches (RESEARCH.md 4b). Switching // dispatch type ends the current encoder and opens a new one in the same // command buffer; Metal orders tracked resources across that boundary, so // the switch itself acts as the barrier. // // Caller's contract: everything encoded between two switches must be // free of read-after-write dependencies on each other. MTL::ComputeCommandEncoder* concurrent_encoder(); // Makes encoder() hand back a concurrent encoder; see ConcurrentRegion. void set_concurrent(bool on); // Readback boundary: end encoding, commit, wait. Returns immediately if // nothing is pending. GPU time of the completed buffer is accumulated // into gpu_seconds(). void sync(); double gpu_seconds() const { return gpu_seconds_; } private: Stream() = default; MTL::ComputeCommandEncoder* encoder_of(int dispatch_type); MTL::CommandBuffer* cmd_ = nullptr; MTL::ComputeCommandEncoder* enc_ = nullptr; int dispatch_type_ = -1; bool concurrent_ = false; double gpu_seconds_ = 0.0; }; inline void sync() { Stream::get().sync(); } // RAII: inside this scope, ops encode into a CONCURRENT compute encoder. // Everything encoded in the region must be mutually independent (no // read-after-write between them); dependencies across the region boundary are // fine because switching encoder kind ends the encoder, and Metal orders // tracked resources across encoders in a command buffer. struct ConcurrentRegion { ConcurrentRegion() { Stream::get().set_concurrent(true); } ~ConcurrentRegion() { Stream::get().set_concurrent(false); } ConcurrentRegion(const ConcurrentRegion&) = delete; ConcurrentRegion& operator=(const ConcurrentRegion&) = delete; }; // ---- f32 forward ops (parity-tested vs forge::cpu) ------------------------- enum class MatmulKernel { Auto, Naive, Tiled, Simdgroup }; void matmul(const Tensor& a, const Tensor& b, Tensor& c, bool transpose_a = false, bool transpose_b = false, bool accumulate = false, MatmulKernel kernel = MatmulKernel::Auto); void add(const Tensor& a, const Tensor& b, Tensor& out); void mul(const Tensor& a, const Tensor& b, Tensor& out); void scale(const Tensor& a, float s, Tensor& out); void add_bias(const Tensor& x, const Tensor& bias, Tensor& out); void silu(const Tensor& x, Tensor& out); void gelu(const Tensor& x, Tensor& out); void relu2(const Tensor& x, Tensor& out); // max(x,0)^2 void relu2_backward(const Tensor& x, const Tensor& dout, Tensor& dx); void softcap(const Tensor& x, float cap, Tensor& out); // cap*tanh(x/cap) void softcap_backward(const Tensor& y, const Tensor& dout, float cap, Tensor& dx); void softmax(const Tensor& x, Tensor& out); // rows = last dim void rmsnorm(const Tensor& x, const Tensor& w, float eps, Tensor& out); void layernorm(const Tensor& x, const Tensor& w, const Tensor& b, float eps, Tensor& out); // ---- f32 backward / training ops (ACCUMULATE into d* outputs) -------------- void accumulate(Tensor& dst, const Tensor& src); // dst += src void axpy(Tensor& dst, const Tensor& src, const Tensor& s); // dst += src * s[0] void silu_backward(const Tensor& x, const Tensor& dout, Tensor& dx); void gelu_backward(const Tensor& x, const Tensor& dout, Tensor& dx); void add_bias_backward(const Tensor& dout, Tensor& dbias); void rmsnorm_backward(const Tensor& x, const Tensor& w, float eps, const Tensor& dout, Tensor& dx, Tensor& dw); void layernorm_backward(const Tensor& x, const Tensor& w, float eps, const Tensor& dout, Tensor& dx, Tensor& dw, Tensor& db); // freqs: [head_dim/2] per-pair inverse frequencies, host-precomputed // (fixed theta, llama3 rope-scaling, per-layer theta — all just tables). void rope(const Tensor& x, int64_t n_heads, const Tensor& freqs, int64_t pos_offset, Tensor& out); void rope_backward(const Tensor& dout, int64_t n_heads, const Tensor& freqs, int64_t pos_offset, Tensor& dx); void embedding(const Tensor& weight, const Tensor& ids, Tensor& out); // ids i32 void embedding_backward(const Tensor& ids, const Tensor& dout, Tensor& dweight); // window > 0: sliding-window attention; attn_softcap > 0: cap·tanh pre-softmax // (unfused path only — the fused kernels don't support softcap). void attention(const Tensor& q, const Tensor& k, const Tensor& v, int64_t n_heads, int64_t n_kv_heads, bool causal, float scale, Tensor& out, Tensor* probs_out, int64_t window = 0, float attn_softcap = 0.0f); void attention_backward(const Tensor& q, const Tensor& k, const Tensor& v, const Tensor& probs, const Tensor& out, const Tensor& dout, int64_t n_heads, int64_t n_kv_heads, float scale, Tensor& dq, Tensor& dk, Tensor& dv, float attn_softcap = 0.0f); // Fused (flash) attention: stores no T x T probabilities, only the per-row // logsumexp `lse` [B, n_heads, T] that the backward re-expands. Supported for // a fixed set of head_dims — flash_supported() reports which. bool flash_supported(int64_t head_dim); // Scalar: one thread per query row. MMA: simdgroup_matrix 8x8 tiles. Same // outputs; Auto picks MMA where supported. enum class FlashKernel { Auto, Scalar, MMA }; // window > 0 requires the Scalar kernel (Auto routes there automatically); // out-of-window KV blocks are skipped, so cost scales with the window. void flash_attention(const Tensor& q, const Tensor& k, const Tensor& v, int64_t n_heads, int64_t n_kv_heads, bool causal, float scale, Tensor& out, Tensor& lse, FlashKernel kernel = FlashKernel::Auto, int64_t window = 0); void flash_attention_backward(const Tensor& q, const Tensor& k, const Tensor& v, const Tensor& out, const Tensor& lse, const Tensor& dout, int64_t n_heads, int64_t n_kv_heads, bool causal, float scale, Tensor& dq, Tensor& dk, Tensor& dv, FlashKernel kernel = FlashKernel::Auto, int64_t window = 0); // losses: [N] per-row buffer; loss_out: [1] mean over n_valid. dlogits // optional (accumulated). ids must be i32. void cross_entropy(const Tensor& logits, const Tensor& targets, int64_t n_valid, Tensor& losses, Tensor& loss_out, Tensor* dlogits); // ---- QAT / MoE -------------------------------------------------------------- // Per-row fake quantization of a 2-D weight; mode 0 = int8 (absmax/127), // mode 1 = ternary (absmean, BitNet-style). Backward is STE — no kernel. void fake_quant(const Tensor& w, int mode, Tensor& out); // dx += p ∘ (dout − dot(dout, p)) per row; thread-per-row, small last dims. void softmax_backward(const Tensor& p, const Tensor& dout, Tensor& dx); void sigmoid(const Tensor& x, Tensor& out); void sigmoid_backward(const Tensor& y, const Tensor& dout, Tensor& dx); // Keep top-k per row of p [N,E] — selected by p+bias, gated by p alone; // norm renormalizes kept gates to sum 1. void topk_renorm(const Tensor& p, const Tensor& bias, int64_t k, bool norm, Tensor& out); void topk_renorm_backward(const Tensor& p, const Tensor& bias, const Tensor& dout, int64_t k, bool norm, Tensor& dp); // counts[e] += #rows with nonzero gate for expert e void expert_counts(const Tensor& gates, Tensor& counts); // out[i,:] = x[i,:] * gates[i,e]; the _accumulate variant does dst += (dx path). void row_scale(const Tensor& x, const Tensor& gates, int64_t e, Tensor& out); void row_scale_accumulate(const Tensor& x, const Tensor& gates, int64_t e, Tensor& dst); // dgates[i,e] += dot(dout[i,:], x[i,:]) void row_scale_gate_backward(const Tensor& dout, const Tensor& x, int64_t e, Tensor& dgates); // Fused optimizer update for one tensor; all state f32. void adamw_step(Tensor& w, const Tensor& g, Tensor& m, Tensor& v, float lr, float beta1, float beta2, int64_t t, float eps, float wd, float grad_scale); // out[0] = sum(x^2) — single-threadgroup reduce (fine for M4). void sumsq(const Tensor& x, Tensor& out); // out[0] = sum(x) * mul void sum(const Tensor& x, Tensor& out, float mul); } // namespace forge::metal