#pragma once
#include <std/erased.h>
// SafeC Standard Library — GPU tensor ops via Metal (Apple Silicon/macOS).
//
// The "unified memory" piece of the ML story (MLX's actual selling point):
// a Metal buffer created with the default storage mode is CPU+GPU shared
// memory on Apple Silicon — no explicit host<->device copy step the way
// CUDA/ROCm need (see gpu_cuda.h/gpu_rocm.h's discrete-memory model) —
// 'contents' on an MTLBuffer just IS a pointer the CPU can read straight
// out of after the GPU finishes writing it.
//
// Same "no real Objective-C support in safec" situation as
// std/gui/gui_cocoa.sc — every call here goes through the plain-C
// Objective-C runtime (objc_msgSend et al.) and the Metal framework's own
// C entry point (MTLCreateSystemDefaultDevice). See that file's header
// comment for the full technique (typedef fn-cast objc_msgSend per
// message shape). Only linkable/runnable on macOS with a Metal-capable
// GPU (every Apple Silicon Mac, and most Intel Macs from the last
// decade) — link with '-framework Metal -framework Foundation'.
//
// Gated behind the 'mps' feature (see Package.toml's [features] — same
// convention as every other backend-specific piece of std/ this session
// added): only compile/link this on a build that opted in, since a
// non-Apple target obviously can't provide these frameworks.
namespace std {
// Elementwise 'out[i] = a[i] + b[i]' for i in [0, n), computed on the GPU
// via a small inline Metal Shading Language compute kernel compiled at
// runtime. Returns 1 on success, 0 if no Metal device is available or any
// step of the pipeline setup fails (compile error in the embedded MSL
// source, no GPU present, etc.) — 'out' is left untouched on failure.
//
// Verified working end-to-end on real Apple Silicon hardware (correct
// output for a real GPU dispatch). Previously segfaulted at the
// dispatchThreadgroups:threadsPerThreadgroup: call specifically — root
// cause was a genuine SafeC compiler codegen bug (fixed in CodeGen.cpp's
// genCall), not anything specific to this file: any indirect call passing
// two or more consecutive struct-by-value arguments over 16 bytes (and not
// a Homogeneous Floating-point Aggregate) through a function-pointer cast
// hit it, because this compiler represented such arguments as raw LLVM
// aggregate values instead of the AAPCS64-correct "caller copies to a
// stack temporary, passes a plain pointer" form — self-consistent for
// SafeC-to-SafeC calls (which is why that never showed up before), but not
// what a real, externally-compiled function like objc_msgSend's target IMP
// expects to receive in its argument registers.
int mps_add_f32([]const float a, []const float b, []float out);
// internal: the (reference, length) core of mps_add_f32()
int mps_add_f32_n_(const &float a, const &float b, &float out, unsigned long n);
// Same shape/verification status as mps_add_f32 (elementwise, GPU-executed,
// unified-memory readback) — sub/mul/div/pow are two-input, log/exp/sqrt
// are one-input. mul/div/pow/log/exp/sqrt share mps_add_f32's setup/
// dispatch/readback path (see __mps_run_binary_kernel/__mps_run_unary_kernel
// in gpu_mps.sc), so the fix that made mps_add_f32's GPU dispatch work
// applies to all of them identically — verified directly for
// mps_sub_f32/mps_mul_f32 as representative binary/no-special-case cases;
// div/pow/log/exp/sqrt follow the exact same code path with only the
// one-line kernel body differing. mps_sub_f32 itself is NOT routed through
// __mps_run_binary_kernel (unlike mul/div/pow) — it's batch-aware, same
// reasoning as mps_relu_f32/mps_matmul_f32 (see mps_relu_f32's comment):
// a per-op caller needs to encode it onto a shared batch's
// command buffer for the fused forward+loss+backward training-step path,
// not force its own commit+wait mid-batch.
int mps_sub_f32([]const float a, []const float b, []float out);
// internal: the (reference, length) core of mps_sub_f32()
int mps_sub_f32_n_(const &float a, const &float b, &float out, unsigned long n);
int mps_mul_f32([]const float a, []const float b, []float out);
// internal: the (reference, length) core of mps_mul_f32()
int mps_mul_f32_n_(const &float a, const &float b, &float out, unsigned long n);
int mps_div_f32([]const float a, []const float b, []float out);
// internal: the (reference, length) core of mps_div_f32()
int mps_div_f32_n_(const &float a, const &float b, &float out, unsigned long n);
int mps_pow_f32([]const float a, []const float b, []float out);
// internal: the (reference, length) core of mps_pow_f32()
int mps_pow_f32_n_(const &float a, const &float b, &float out, unsigned long n);
int mps_log_f32([]const float a, []float out);
// internal: the (reference, length) core of mps_log_f32()
int mps_log_f32_n_(const &float a, &float out, unsigned long n);
int mps_exp_f32([]const float a, []float out);
// internal: the (reference, length) core of mps_exp_f32()
int mps_exp_f32_n_(const &float a, &float out, unsigned long n);
int mps_sqrt_f32([]const float a, []float out);
// internal: the (reference, length) core of mps_sqrt_f32()
int mps_sqrt_f32_n_(const &float a, &float out, unsigned long n);
int mps_relu_f32([]const float a, []float out);
// internal: the (reference, length) core of mps_relu_f32()
int mps_relu_f32_n_(const &float a, &float out, unsigned long n);
int mps_scale_f32([]const float a, float k, []float out);
// internal: the (reference, length) core of mps_scale_f32()
int mps_scale_f32_n_(const &float a, float k, &float out, unsigned long n);
// out[0] = sum(a[0..n)) -- two-stage parallel reduction (grid-stride
// accumulate + threadgroup tree reduction, see gpu_mps_kernels.metal's
// sum_kernel comment), with the small per-threadgroup partial-sums array
// combined on the CPU after one readback. Batch-aware like mps_sub_f32/
// mps_relu_f32 above: under an open batch, the CPU-side partial-sum
// combine can't run until mps_batch_end()'s single wait, so it's deferred
// into a registered finalize callback instead of happening inline here —
// see gpu_mps.sc's implementation comment.
int mps_sum_f32([]const float a, []float out);
// internal: the (reference, length) core of mps_sum_f32()
int mps_sum_f32_n_(const &float a, &float out, unsigned long n);
// out[M,N] = a[M,K] . b[K,N], computed on the GPU. float32-only (Metal has
// no 'double' type at all). Internally picks between three kernels
// depending on shape, most-specialized first:
// - M, K, N all exact multiples of 32: matmul_kernel_smma_multi --
// threadgroup-memory tiling (data reuse, like matmul_kernel) AND
// simdgroup_matrix hardware matrix units (like matmul_kernel_smma)
// together, the same combination PyTorch's MPSMatrixMultiplication and
// MLX's fastest GEMM paths use.
// - otherwise, M, K, N all exact multiples of 8: matmul_kernel_smma --
// simdgroup_matrix alone, no threadgroup-memory tiling.
// - otherwise: matmul_kernel's ordinary threadgroup-memory tiling, no
// alignment requirement at all.
// The two simdgroup_matrix kernels require their alignment because
// simdgroup_load/store past a matrix's real bounds is undefined behavior,
// unlike matmul_kernel's explicit boundary checks. See gpu_mps.sc's
// comment on the implementation, and gpu_mps_kernels.metal's comments on
// each kernel, for detail on all three paths.
int mps_matmul_f32([]const float a, []const float b, []float out, unsigned long M, unsigned long K, unsigned long N);
// internal: the (reference, length) core of mps_matmul_f32()
int mps_matmul_f32_n_(const &float a, const &float b, &float out,
unsigned long M, unsigned long K, unsigned long N);
// GPU-backed matmul backward: dA = dC . B^T, dB = A^T . dC (same math
// tensor_blas.sc's __matmul_backward_blas passes to cblas_dgemm via a
// transpose flag; here each is its own kernel instead). Used by
// a GPU matmul's backward pass, so a forward pass on the GPU gets
// a matmul_gpu backward pass too, instead of falling back to the CPU
// backward the way it used to — see that function's comment for why that
// fallback was the actual dominant cost in a training loop, not dispatch
// overhead.
// mps_matmul_abt_f32: out[M,K] = a[M,N] . b^T, b stored [K,N] (dA's shape)
// mps_matmul_atb_f32: out[K,N] = a^T . b, a stored [M,K] (dB's shape)
// Both are batch-aware (see "Batched dispatch" below) the same way
// mps_matmul_f32/mps_relu_f32 are -- a training step's backward pass is
// its own sequential chain (matmul2-backward -> relu-backward ->
// matmul1-backward, each waiting on the previous), so it pays the same
// per-op command-buffer round-trip cost forward used to before batching.
// Same three-tier smma-multi/smma/tiled shape-based kernel choice as
// mps_matmul_f32 (see its comment) internally, too.
int mps_matmul_abt_f32(const &float a, const &float b, &float out,
unsigned long M, unsigned long N, unsigned long K);
int mps_matmul_atb_f32(const &float a, const &float b, &float out,
unsigned long M, unsigned long K, unsigned long N);
// GPU-backed relu backward: out[i] = a[i] > 0 ? selfGrad[i] : 0. Same
// shape/verification path as mps_add_f32 (see __mps_run_binary_kernel).
// Batch-aware, same as the two matmul-backward entry points above.
int mps_relu_backward_f32([]const float a, []const float selfGrad, []float out);
// internal: the (reference, length) core of mps_relu_backward_f32()
int mps_relu_backward_f32_n_(const &float a, const &float selfGrad, &float out, unsigned long n);
// True if a Metal device is available on this machine at all (checks
// MTLCreateSystemDefaultDevice() != nil) — lets callers fall back to the
// CPU Tensor ops (tensor.h) when there's no GPU, without needing to
// attempt an mps_* call first just to find out.
int mps_available();
// ── Device buffers ────────────────────────────────────────────────────────────
// Element types of a buffer (MpsBuf.dtype) and of the attention ops' dtype
// argument: 0 = fp32, or raw 16-bit values.
#define MPS_FP16 1 // same values as gpu_cuda.h's CUDA_FP16 / CUDA_BF16
#define MPS_BF16 2
// A Metal buffer the GPU ops below read and write directly -- the MPS twin
// of gpu_rocm.h's RocmBuf / RocmHalfBuf. 'handle' is the MTLBuffer (CPU/GPU
// shared memory, never dereferenced by callers) and null for an invalid
// (failed or freed) buffer; every op returns 0 for one. 'bytes' is the size
// it was created with -- what free() returns to the buffer pool and the
// most read() copies. 'dtype' is its element type: 0 = fp32, MPS_FP16 or
// MPS_BF16 (informational: ops don't check it, so a buffer can hold any
// 4- or 2-byte data, e.g. a uint32 length table uploaded through upload()).
//
// Persistent buffers (alloc / upload* / linear_attn_state) are owned by the
// caller: upload once, then pass the same buffer to as many dispatches and
// batches as you like, instead of re-uploading a CPU array on every call.
// Construct only while no batch that might still be reading a freed buffer
// is in flight (constructors reuse pool slots); free() exactly once, after
// every batch reading it has ended. All constructors return an invalid
// buffer without a Metal device.
// An Objective-C / Metal object (id<MTLBuffer>, id<MTLDevice>, ...): opaque,
// never constructed, only passed back to the Metal APIs it comes from.
struct MtlObject { char opaque_; };
struct MpsBuf {
?&MtlObject handle;
unsigned long bytes;
int dtype;
int valid() const; // handle != null
void free(); // back to the buffer pool; handle := null
// Copies the first dst.len (<= this.bytes) bytes into 'dst'. Valid once every
// batch writing the buffer has ended (shared memory: no transfer). 1 = done.
int read([]unsigned char dst) const;
// Copies src.len bytes (<= this.bytes) from 'src' into the buffer: valid while no
// batch that reads or writes it is in flight. 1 = done.
int write([]const unsigned char src);
// The buffer's memory as the CPU sees it (shared storage: the same bytes
// the GPU reads and writes), null for an invalid buffer or without Metal.
// Read or write it only while no batch using the buffer is in flight.
?&float contents() const;
static struct MpsBuf alloc(unsigned long n, int dtype); // n uninitialized elements
static struct MpsBuf upload([]const float host);
// internal: the (reference, length) core of upload()
static struct MpsBuf upload_n_(const &float host, unsigned long n); // n fp32 values
// n fp32 values rounded on the CPU (float16.h's f32_to_fp16_bulk /
// f32_to_bf16_bulk) and stored as fp16 / bf16.
static struct MpsBuf upload_f16([]const float host);
// internal: the (reference, length) core of upload_f16()
static struct MpsBuf upload_f16_n_(const &float host, unsigned long n);
static struct MpsBuf upload_bf16([]const float host);
// internal: the (reference, length) core of upload_bf16()
static struct MpsBuf upload_bf16_n_(const &float host, unsigned long n);
// n raw fp16 or bf16 values (dtype MPS_FP16 / MPS_BF16), copied as is.
static struct MpsBuf upload_bits([]const unsigned short bits, int dtype);
// internal: the (reference, length) core of upload_bits()
static struct MpsBuf upload_bits_n_(const &unsigned short bits, unsigned long n, int dtype);
// A zeroed recurrent linear-attention state (mps_linear_attn_bf16_chained):
// mps_linear_attn_state_bytes(numHeads, headDim) bytes of fp32.
static struct MpsBuf linear_attn_state(unsigned long numHeads, unsigned long headDim);
// The output of the op most recently encoded into the open batch (its
// bytes and dtype), invalid if no batch is active or nothing has been
// encoded yet. Call right after a batched op returns to chain its
// still-pending result into a following op without waiting for
// mps_batch_end(). The batch owns it: don't free() it; it returns to the
// pool when the batch ends.
static struct MpsBuf batch_last_output();
};
// ── Batched dispatch ──────────────────────────────────────────────────────────
// Every mps_*_f32 call above independently creates its own command
// buffer, dispatches, and blocks on waitUntilCompleted before returning
// -- correct, but for a chain of several ops (e.g. matmul -> relu ->
// matmul in a forward pass), paying one full command-buffer-submission +
// host/device sync round trip PER OP is the dominant cost once shader
// compilation is already cached (see mps_matmul_f32's implementation
// comment), especially for small workloads where dispatch latency swamps
// actual compute time.
//
// mps_batch_begin()/mps_batch_end() let a caller wrap several mps_*_f32
// calls so they share ONE command buffer and ONE sync point:
//
// mps_batch_begin();
// mps_matmul_f32(x, w1, h, ...); // encoded onto the shared buffer,
// // NOT dispatched yet -- 'h' is not
// // valid at this point
// mps_relu_f32(h, r, ...); // reads 'h' — safe: this dispatch
// // is encoded AFTER matmul's on the
// // same buffer, so Metal orders it
// // after even though the CPU
// // hasn't waited yet
// mps_matmul_f32(r, w2, y, ...);
// mps_batch_end(); // commits once, waits once, THEN
// // every op's 'out' pointer above
// // becomes valid
//
// Every batched op's 'out' buffer is invalid to read until
// mps_batch_end() returns — the same contract as reading a
// cudaMemcpyAsync destination before synchronizing.
//
// IMPORTANT: that example is only safe because none of mps_matmul_f32's
// or mps_relu_f32's *inputs* above are themselves another batched op's
// still-pending output — 'x'/'w1'/'w2' are all already-populated,
// ordinary CPU float arrays. mps_matmul_f32/mps_relu_f32 always upload
// their input from a CPU float* into a fresh GPU buffer, so feeding one
// batched op's 'out' pointer into a later op's 'a'/'b' *within the same
// unflushed batch* silently reads whatever was in that CPU memory before
// the batch started (found and fixed via direct testing, not something
// caught by inspection: a naive matmul -> relu chain like this computed
// 0 instead of the real value). Use mps_matmul_f32_chained/
// mps_relu_f32_chained (below) instead for that case — they take the
// device buffer a previous batched op already wrote to
// (MpsBuf::batch_last_output()) directly, skipping the CPU round trip
// entirely rather than reading stale CPU data.
void mps_batch_begin();
void mps_batch_end();
int mps_batch_is_active();
// Commits what the open batch has encoded so far without waiting and keeps
// the batch open: the GPU starts on that work while the caller encodes the
// next. The command buffers run in order on the one queue, and
// mps_batch_end() still waits for all of them (a trainer submits
// every few steps of an epoch and waits once at its end). No-op without a
// batch.
void mps_batch_submit();
// Type-erased form behind mps_batch_register_finalize (below; see std/erased.h).
typedef fn void(?&Erased ctx) MpsFinalizeFn;
void mps_batch_register_finalize_raw_(MpsFinalizeFn finalizeFn, ?&Erased ctx);
// Register a callback to run once, in registration order, after
// mps_batch_end()'s single wait and raw GPU->CPU readback of every
// batched op's output buffer. Exists for callers that
// need to do their own CPU-side post-processing on a batched op's result
// (e.g. float32->float64 conversion) — that processing has to wait for
// the same sync point the raw readback does, so it can't just run
// immediately after the mps_*_f32 call returns the way it does outside
// a batch.
generic<T>
void mps_batch_register_finalize(fn void(?&T) finalizeFn, ?&T ctx) {
unsafe { mps_batch_register_finalize_raw_((MpsFinalizeFn)finalizeFn, (?&Erased)ctx); }
}
// Same computation as mps_matmul_f32/mps_relu_f32, but the primary input
// ('bufA') is a device buffer — typically MpsBuf::batch_last_output()'s
// return from a previous batched op in the same still-open batch —
// instead of a CPU float array, so no CPU<->GPU round trip happens for
// that input at all. Only usable while a batch is active
// (mps_batch_is_active()); the secondary input 'b' (matmul's second
// operand) is still an ordinary CPU float array here — see
// mps_matmul_f32_persistent below for the case where BOTH operands are
// already device buffers (e.g. a weight matrix uploaded once with
// MpsBuf::upload instead of fresh every call).
int mps_matmul_f32_chained(struct MpsBuf bufA, const &float b, &float out,
unsigned long M, unsigned long K, unsigned long N);
int mps_relu_f32_chained(struct MpsBuf bufA, &float out, unsigned long n);
// out[M,N] = bufA[M,K] . bufB[K,N] -- like mps_matmul_f32_chained, but
// 'bufB' is ALSO a pre-existing device buffer instead of a plain CPU
// array to upload fresh. For a matmul run repeatedly against the same
// weight matrix (e.g. an inference loop doing the same forward pass
// hundreds of times), mps_matmul_f32_chained would re-upload (memcpy)
// that weight matrix from CPU on every single call even though it never
// changes — this skips that entirely for both operands.
int mps_matmul_f32_persistent(struct MpsBuf bufA, struct MpsBuf bufB, &float out,
unsigned long M, unsigned long K, unsigned long N);
// out[i] = bufA[i] - b[i] -- same chaining shape as mps_matmul_f32_chained/
// mps_relu_f32_chained above, used by tensor_sub_gpu for the loss
// computation's diff = pred - target (pred, e.g. a matmul's output, is
// still GPU-pending; target is an ordinary already-populated Tensor).
int mps_sub_f32_chained(struct MpsBuf bufA, const &float b, &float out, unsigned long n);
// out[i] = bufA[i] * b[i] -- generic chained elementwise multiply (one
// pending GPU input, one plain CPU array), used by Tensor::square_gpu's
// backward: 2*diff (still GPU-pending from the forward pass) times the
// incoming gradient (an ordinary CPU array -- see __sum_backward, which
// writes it eagerly with no GPU involvement at all).
int mps_mul_f32_chained(struct MpsBuf bufA, const &float b, &float out, unsigned long n);
// out[i] = buf[i] * buf[i] -- self-multiply of a single still-pending GPU
// buffer (both mul_kernel operand slots bound to the same buffer; reading
// the same device memory twice is safe, there's no write). Used by
// Tensor::square_gpu's forward pass for the loss computation's
// sq = diff * diff, where 'diff' is itself still GPU-pending -- the
// generic mps_mul_f32_chained above can't express this shape since its
// second operand has to already be a plain CPU array.
int mps_square_f32_chained(struct MpsBuf buf, &float out, unsigned long n);
// out[0] = sum(bufA[0..n)) -- chained version of mps_sum_f32 above, input
// already a still-pending GPU buffer. Same deferred-partial-sum-combine
// finalize as the non-chained batched path.
int mps_sum_f32_chained(struct MpsBuf bufA, &float out, unsigned long n);
// out[i] = buf[i] * k -- chained version of mps_scale_f32, input already a
// still-pending GPU buffer. Used by tensor_sub_gpu's backward for
// dTarget = -1 * (incoming gradient), when that incoming gradient is
// itself still GPU-pending.
int mps_scale_f32_chained(struct MpsBuf buf, float k, &float out, unsigned long n);
// Fused gated-activation (GeGLU/tanh-GELU) step: 'bufRaw' is a still-
// pending GPU buffer holding one row-major [rows, 2*intermediateSize]
// matmul output (gate half then up half per row -- e.g. the result of a
// matmul against a row-concatenated [gateW; upW] weight); computes
// out[row,col] = gelu(bufRaw[row,col]) * bufRaw[row, intermediateSize+col]
// for col in [0, intermediateSize), same tanh-approx GELU formula as
// activations.sc.gelu_fwd(), in one dispatch -- see
// gpu_mps_kernels.metal's gated_gelu_kernel. 'out' has rows*
// intermediateSize elements, same throwaway-CPU-scratch-until-batch-end
// contract as every other *_chained function here.
int mps_gated_gelu_f32_chained(struct MpsBuf bufRaw, &float out, unsigned long rows, unsigned long intermediateSize);
// bf16 counterpart of mps_gated_gelu_f32_chained above -- 'bufRaw'/'out'
// are raw bf16 bits, same fused GeGLU step. Lets a whole bf16
// matmul -> gated-GELU -> matmul FFN run as one GPU round trip.
int mps_gated_gelu_bf16_chained(struct MpsBuf bufRaw, &unsigned short out, unsigned long rows, unsigned long intermediateSize);
// ── Backward chained variants ─────────────────────────────────────────────────
// Same idea as mps_matmul_f32_chained/mps_relu_f32_chained above, for the
// backward chain (matmul2-backward's dA=dH output feeds relu-backward's
// selfGrad input; relu-backward's output feeds matmul1-backward's dC
// input) -- each takes the *upstream op's pending GPU buffer* for
// whichever argument is the flowing gradient, so no CPU round trip
// happens for it. The other argument (always an ordinary already-
// populated Tensor -- a forward activation or a weight matrix) stays a
// plain CPU float array, same restriction as the forward chained pair.
// mps_matmul_abt_f32_chained: 'bufDC' replaces 'a' (the incoming
// gradient) -- for a deeper graph where dC is itself still-pending.
// mps_matmul_atb_f32_chained: 'bufDC' replaces 'b' (matmul_atb_f32's
// second argument is the incoming gradient) -- this is the one this
// library's own 2-layer MLP training loop actually needs, for
// dW1 = X^T . dH_pre where dH_pre is relu-backward's pending output.
// mps_relu_backward_f32_chained: 'bufSelfGrad' replaces 'selfGrad' --
// this library's own MLP needs this one too, for dH_pre = relu'(H) *
// dH where dH is matmul2-backward's pending dA output.
int mps_matmul_abt_f32_chained(struct MpsBuf bufDC, const &float b, &float out,
unsigned long M, unsigned long N, unsigned long K);
int mps_matmul_atb_f32_chained(const &float a, struct MpsBuf bufDC, &float out,
unsigned long M, unsigned long K, unsigned long N);
int mps_relu_backward_f32_chained(const &float a, struct MpsBuf bufSelfGrad, &float out, unsigned long n);
// Same as mps_relu_backward_f32_chained, but 'a' is ALSO a still-pending
// GPU buffer -- needed when relu's own input is itself an activation
// computed earlier in the SAME still-open batch (the fused
// forward+loss+backward training step). See gpu_mps.sc's comment.
int mps_relu_backward_f32_chained_ab(struct MpsBuf bufA, struct MpsBuf bufSelfGrad, &float out, unsigned long n);
// Same as mps_matmul_atb_f32_chained, but 'a' is ALSO a still-pending GPU
// buffer (not a plain CPU array) -- needed when the matmul's first operand
// is itself an activation computed earlier in the SAME still-open batch
// (e.g. dW2 = H^T . dY in the fused forward+loss+backward training step,
// where H hasn't been read back to CPU yet either). See gpu_mps.sc's
// comment on the implementation.
int mps_matmul_atb_f32_chained_ab(struct MpsBuf bufA, struct MpsBuf bufDC, &float out,
unsigned long M, unsigned long K, unsigned long N);
// Same as mps_matmul_abt_f32_chained, but 'b' is ALSO a pre-existing
// device buffer (typically MpsBuf::upload, e.g. a weight matrix
// that stays GPU-resident across a whole training run) instead of a
// plain CPU array. Needed for dH = dY . W2^T when W2 is never read back
// to a CPU array between steps (see mps_sgd_update_f32_chained below).
int mps_matmul_abt_f32_persistent(struct MpsBuf bufDC, struct MpsBuf bufB, &float out,
unsigned long M, unsigned long N, unsigned long K);
// In-place SGD step against a persistent device buffer: w[i] -= lr*
// grad[i]. Both 'bufW' and 'bufGrad' are already device buffers (bufW
// typically from MpsBuf::upload, bufGrad typically another op's
// still-pending output from the same batch's backward pass) -- no 'out'
// parameter, since the whole point is that the updated weight never
// leaves the GPU. See gpu_mps.sc's comment and gpu_mps_kernels.metal's
// sgd_update_kernel.
int mps_sgd_update_f32_chained(struct MpsBuf bufW, struct MpsBuf bufGrad, float lr, unsigned long n);
// ── fp16 / bf16 tier ─────────────────────────────────────────────────────
// Real half/bfloat compute on the GPU, not just storage: gpu_mps_kernels.metal's
// matmul_kernel_f16/matmul_kernel_bf16/relu_kernel_f16/relu_kernel_bf16
// hold/multiply in the reduced-precision type and accumulate in float
// (standard mixed-precision-GEMM practice) — verified to compile through
// this machine's real Metal toolchain (both 'half' and 'bfloat' are
// natively supported MSL types on this SDK). The functions below take
// and return ordinary 'float*' host arrays, same shape as mps_matmul_f32/
// mps_relu_f32, so a caller doesn't need to touch fp16/bf16 bit patterns
// directly: conversion (see float16.h's f32_to_fp16_bulk/f32_to_bf16_bulk
// and their inverses) happens on the CPU side of the upload/readback,
// invisibly, inside these functions.
//
// This is the naive (TILE_SIZE=16, threadgroup-memory-tiled) kernel tier
// only — the smma/multi-simdgroup tiers (see mps_matmul_f32's comment)
// use simdgroup_float8x8, a float-only type on this Metal SDK; porting
// those to half/bfloat accumulation is future work, not attempted here.
// Every call here does its own upload/dispatch/readback (no persistent-
// buffer or batching integration) — the simplest, most directly testable
// shape, matching mps_matmul_f32's non-batched fallback path.
//
// Expected precision loss is real and by design: fp16's ~3-4 significant
// decimal digits (10 mantissa bits) and bf16's ~2-3 (7 mantissa bits) are
// both coarser than float32's ~7 — see float16.h's precision notes.
// Returns 1 on success, 0 on failure (no Metal device, pipeline
// compilation failure, or allocation failure) — 'out' is left untouched
// on failure.
int mps_matmul_f16(const &float a, const &float b, &float out,
unsigned long M, unsigned long K, unsigned long N);
int mps_matmul_bf16(const &float a, const &float b, &float out,
unsigned long M, unsigned long K, unsigned long N);
int mps_relu_f16([]const float a, []float out);
// internal: the (reference, length) core of mps_relu_f16()
int mps_relu_f16_n_(const &float a, &float out, unsigned long n);
int mps_relu_bf16([]const float a, []float out);
// internal: the (reference, length) core of mps_relu_bf16()
int mps_relu_bf16_n_(const &float a, &float out, unsigned long n);
// ── fp16 / bf16 persistent-buffer / batched tier ────────────────────────
// Mirrors MpsBuf::upload/mps_matmul_f32_persistent/mps_relu_f32_
// chained/mps_sub_f32_chained/mps_scale_f32_chained/mps_matmul_atb_f32_
// chained_ab/mps_matmul_abt_f32_persistent/mps_relu_backward_f32_chained_
// ab/mps_sgd_update_f32_chained's whole family, in fp16/bf16 — enough ops
// to run a full GPU-resident forward+backward+SGD step (see
// gpu_mps_kernels.metal's matching kernel-tier comment) with weights that
// never leave the GPU as float32 at all, the same shape as the float32
// GPU-resident training rewrite this session's MPS work already
// delivered, just in half the bytes per weight/activation.
//
// Key difference from mps_matmul_f16/mps_relu_f16 above: those take/
// return 'float*' and do a full convert-upload-dispatch-convert round
// trip on every call, meant for one-shot use. Everything below instead
// passes 'unsigned short*' scratch buffers (raw fp16/bf16 bits) for any
// CPU-side throwaway output, exactly mirroring how mps_matmul_f32_
// persistent's chained scratch arrays (scratchZ1/scratchH/etc. in
// train_gpu.sc) are never meaningfully read except the one value that
// matters -- so intermediate activations/gradients here never get
// converted back to float32 at all, only read once as raw bits into a
// scratch array mps_batch_end() overwrites into, matching the float32
// version's "throwaway CPU scratch, GPU buffer is what actually chains"
// design. A caller uploads host-side f32 data with MpsBuf::upload_f16/
// upload_bf16 (rounded with float16.h's f32_to_fp16_bulk/f32_to_bf16_bulk
// on the way up), or data already in 16-bit form with MpsBuf::upload_bits,
// and converts back with fp16_to_f32_bulk/bf16_to_f32_bulk only for values
// it actually needs on the CPU (e.g. a loss scalar for monitoring).
//
// Same naive (TILE_SIZE=16 tiled) kernel tier as mps_matmul_f16/
// mps_matmul_bf16 above for every matmul-shaped op here -- no smma/
// multi-simdgroup fast path (see that comment).
// out = bufA[M,K] . bufB[K,N] -- both operands already device buffers
// (e.g. from MpsBuf::upload_f16). Must be called within an open
// batch (mps_batch_begin()); 'out' is throwaway scratch (raw fp16/bf16
// bits, not valid until mps_batch_end()) -- read the real result via
// MpsBuf::batch_last_output() and chain it into the next op, same
// pattern as mps_matmul_f32_persistent.
int mps_matmul_f16_persistent(struct MpsBuf bufA, struct MpsBuf bufB, &unsigned short out,
unsigned long M, unsigned long K, unsigned long N);
int mps_matmul_bf16_persistent(struct MpsBuf bufA, struct MpsBuf bufB, &unsigned short out,
unsigned long M, unsigned long K, unsigned long N);
int mps_relu_f16_chained(struct MpsBuf bufA, &unsigned short out, unsigned long n);
int mps_relu_bf16_chained(struct MpsBuf bufA, &unsigned short out, unsigned long n);
// out[i] = bufA[i] - bHalf[i]/bBf[i] -- 'bHalf'/'bBf' is a plain CPU
// array of raw fp16/bf16 bits (e.g. a training target, pre-converted
// once with f32_to_fp16_bulk/f32_to_bf16_bulk before the loop), uploaded
// fresh each call, same shape as mps_sub_f32_chained.
int mps_sub_f16_chained(struct MpsBuf bufA, const &unsigned short bHalf, &unsigned short out, unsigned long n);
int mps_sub_bf16_chained(struct MpsBuf bufA, const &unsigned short bBf, &unsigned short out, unsigned long n);
// out[i] = buf[i] * k -- 'k' is a plain float scalar (see gpu_mps_
// kernels.metal's scale_kernel_f16/bf16 comment on why the scalar itself
// stays float-precision).
int mps_scale_f16_chained(struct MpsBuf buf, float k, &unsigned short out, unsigned long n);
int mps_scale_bf16_chained(struct MpsBuf buf, float k, &unsigned short out, unsigned long n);
// out[i] = bufA[i] + bufB[i] with BOTH operands device buffers (pending
// or persistent) -- unlike mps_sub_bf16_chained above, whose second
// operand is a CPU array. For residual connections between chained ops.
int mps_add_bf16_chained(struct MpsBuf bufA, struct MpsBuf bufB, &unsigned short out, unsigned long n);
// out[i] = gelu(bufA[i]) * bufB[i] -- mps_gated_gelu_bf16_chained with the
// gate and up operands in two separate device buffers.
int mps_gelu_mul_bf16_chained(struct MpsBuf bufA, struct MpsBuf bufB, &unsigned short out, unsigned long n);
// Row-wise RMSNorm: out[row,c] = x[row,c] * weight[c] /
// sqrt(mean(x[row,:]^2) + eps). 'bufX' is a pending or persistent device
// buffer of [rows, cols] bf16; 'weightBuf' is the [cols] weight, typically
// uploaded once via MpsBuf::upload and cached by the caller.
int mps_rmsnorm_bf16_chained(struct MpsBuf bufX, struct MpsBuf weightBuf, &unsigned short out,
unsigned long rows, unsigned long cols, float eps);
// dB[K,N] = A^T[K,M] . dC[M,N], A stored [M,K] -- both operands already
// device buffers, same shape as mps_matmul_atb_f32_chained_ab.
int mps_matmul_atb_f16_chained_ab(struct MpsBuf bufA, struct MpsBuf bufDC, &unsigned short out,
unsigned long M, unsigned long K, unsigned long N);
int mps_matmul_atb_bf16_chained_ab(struct MpsBuf bufA, struct MpsBuf bufDC, &unsigned short out,
unsigned long M, unsigned long K, unsigned long N);
// dA[M,K] = dC[M,N] . B^T[N,K], B stored [K,N] -- same shape as
// mps_matmul_abt_f32_persistent.
int mps_matmul_abt_f16_persistent(struct MpsBuf bufDC, struct MpsBuf bufB, &unsigned short out,
unsigned long M, unsigned long N, unsigned long K);
int mps_matmul_abt_bf16_persistent(struct MpsBuf bufDC, struct MpsBuf bufB, &unsigned short out,
unsigned long M, unsigned long N, unsigned long K);
int mps_relu_backward_f16_chained_ab(struct MpsBuf bufA, struct MpsBuf bufSelfGrad, &unsigned short out, unsigned long n);
int mps_relu_backward_bf16_chained_ab(struct MpsBuf bufA, struct MpsBuf bufSelfGrad, &unsigned short out, unsigned long n);
// In-place SGD step against a persistent fp16/bf16 device buffer --
// same shape as mps_sgd_update_f32_chained, weights never leave the GPU.
int mps_sgd_update_f16_chained(struct MpsBuf bufW, struct MpsBuf bufGrad, float lr, unsigned long n);
int mps_sgd_update_bf16_chained(struct MpsBuf bufW, struct MpsBuf bufGrad, float lr, unsigned long n);
// ── Batched multi-head attention (bf16) ─────────────────────────────────────
// Scaled-dot-product attention for one new query token per sequence (a
// decode step) against each sequence's own K/V history, across the whole
// batch and every head in one dispatch: grouped-query heads, optional
// sliding window, softmax, weighted sum over V. No 1/sqrt(headDim) factor
// is applied -- fold any scaling into q. See gpu_mps_kernels.metal's
// attn_batch_bf16_kernel for the kernel design.
//
// q: [batchSize, numHeads*headDim] raw bf16 bits.
// kAll/vAll: [batchSize, maxKvLen, numKvHeads*headDim] raw bf16 bits --
// each sequence's full K/V history including the
// just-appended token, gathered contiguously by the caller.
// cacheLens: [batchSize] -- each sequence's valid KV length
// (<= maxKvLen); rows past it are never read.
// out: [batchSize, numHeads*headDim] raw bf16 bits (pre-o_proj).
//
// One-shot dispatch (own commit+wait+readback). Returns 1 on success, 0 if
// no Metal device is available, pipeline compilation fails, headDim isn't
// 64, 128, 256 or 512, or numHeads isn't a multiple of numKvHeads -- callers
// fall back to a CPU path, same contract as every other mps_*_bf16 entry
// point in this file. There is no KV-length limit.
int mps_attn_batch_bf16(const &unsigned short q, const &unsigned short kAll, const &unsigned short vAll,
const &unsigned int cacheLens, &unsigned short out,
unsigned long batchSize, unsigned long numHeads, unsigned long numKvHeads,
unsigned long headDim, unsigned long maxKvLen,
int isSliding, unsigned long slidingWindow);
// Chained variant of mps_attn_batch_bf16: 'bufQ'/'bufK'/'bufV'/
// 'bufCacheLens' are device buffers (e.g. from MpsBuf::upload) and
// this MUST run inside an open batch (mps_batch_begin()). The output stays
// a pending GPU buffer (MpsBuf::batch_last_output()) so it can feed a
// following op -- e.g. o_proj via mps_matmul_abt_bf16_persistent -- in the
// same commit+wait. 'out' is throwaway CPU scratch, same contract as every
// other chained op. Same failure conditions as the one-shot variant, plus
// returning 0 outside a batch.
int mps_attn_batch_bf16_chained(struct MpsBuf bufQ, struct MpsBuf bufK, struct MpsBuf bufV, struct MpsBuf bufCacheLens,
&unsigned short out,
unsigned long batchSize, unsigned long numHeads, unsigned long numKvHeads,
unsigned long headDim, unsigned long maxKvLen,
int isSliding, unsigned long slidingWindow);
// ── Device-buffer GEMM ──────────────────────────────────────────────────────
// c[m,n] = op(a)[m,k] . op(b)[k,n] on fp32 buffers, row-major, written into
// the caller's buffer 'c' -- gpu_rocm.h's rocm_gemm. 'ta' / 'tb' = 1 means
// that operand is stored transposed (a as k x m, b as n x k); at most one of
// them may be (the three cases are mps_matmul_f32's, mps_matmul_abt_f32's and
// mps_matmul_atb_f32's kernels, with the same shape-based choice of tier).
// Like the device-buffer attention ops below: encoded onto the open batch
// (c valid after mps_batch_end()) or, outside one, run and waited for.
// 1 = done/encoded, 0 = bad shape, a buffer too small, or no Metal device.
int mps_gemm(struct MpsBuf a, int ta, struct MpsBuf b, int tb, struct MpsBuf c,
unsigned long m, unsigned long n, unsigned long k);
// ── Flash and linear attention ──────────────────────────────────────────────
// Batched multi-head versions of attention.h's flash_attention_forward and
// linear_attention_forward on the GPU, in fp32, fp16 or bf16 (all math in
// fp32). Two calling forms share the same kernels:
// - mps_flash_attention / mps_linear_attention (and the host-array _f32
// versions) take the same arguments as gpu_rocm.h's and
// gpu_cuda_attention.h's rocm_/cuda_ ops, so code moves between AMD,
// NVIDIA and Apple GPUs unchanged: fp32 device buffers (from
// MpsBuf::upload), q and out [heads][seqLen][headDim], k and v
// [heads][kvLen][headDim] (fold a batch into heads).
// - The _ex versions take an MpsAttnShape (a batch, grouped-query K/V
// with kvHeads dividing heads -- query head h reads kv head
// h / (heads / kvHeads) -- and a sliding window) and a dtype: 0 = fp32,
// MPS_FP16, MPS_BF16 (raw 16-bit values, as from
// MpsBuf::upload_f16/upload_bf16). Q and out are then
// [batch][heads][seqLen][headDim], K and V [batch][kvHeads][kvLen][headDim].
// All buffers are row-major and contiguous.
//
// Flash attention: softmax(q k^T * scale) v without forming the score matrix
// -- FlashAttention-2 on the GPU's simdgroup matrix units for head dims up
// to 256, a simdgroup-reduction kernel for 512. 'scale' <= 0 means
// 1/sqrt(headDim). Linear attention: phi(x) = elu(x)+1,
// out_i = phi(q_i) S / phi(q_i) z with S = sum_j phi(k_j) v_j^T and
// z = sum_j phi(k_j) -- no scale, since the division cancels it.
//
// Masking: query i sits at position pos = i + (kvLen - seqLen) -- the last
// seqLen keys are the queries' own tokens, any before them a KV cache.
// causal hides keys j > pos (needs kvLen >= seqLen); a window w > 0 (flash
// only) also hides keys j <= pos - w, i.e. the query sees its last w
// positions. A query with no visible key gets zeros.
//
// Head dims: 32, 64, 96, 128, 256 or 512 (mps_attention_head_dim_ok). The
// device-buffer forms encode onto the open batch (results valid after
// mps_batch_end()) or, outside one, run and wait. 1 = done/encoded, 0 = bad
// shape or no Metal device -- callers fall back to a CPU path.
struct MpsAttnShape {
unsigned long batch;
unsigned long heads; // query heads
unsigned long kvHeads; // key/value heads; divides heads
unsigned long seqLen; // query length
unsigned long kvLen; // key/value length
unsigned long headDim; // 32, 64, 96, 128, 256 or 512
unsigned long window; // flash attention sliding window; 0 = none
};
int mps_attention_head_dim_ok(unsigned long headDim);
int mps_flash_attention(struct MpsBuf q, struct MpsBuf k, struct MpsBuf v, struct MpsBuf out, unsigned long heads, unsigned long seqLen,
unsigned long kvLen, unsigned long headDim, float scale, int causal);
int mps_flash_attention_ex(struct MpsBuf q, struct MpsBuf k, struct MpsBuf v, struct MpsBuf out, struct MpsAttnShape s, int dtype,
float scale, int causal);
int mps_linear_attention(struct MpsBuf q, struct MpsBuf k, struct MpsBuf v, struct MpsBuf out, unsigned long heads, unsigned long seqLen,
unsigned long kvLen, unsigned long headDim, int causal);
int mps_linear_attention_ex(struct MpsBuf q, struct MpsBuf k, struct MpsBuf v, struct MpsBuf out, struct MpsAttnShape s, int dtype, int causal);
// The two fp32 ops on host arrays (upload, compute, read back), with the
// rocm_*_f32 arguments. Batch-aware like mps_sub_f32.
int mps_flash_attention_f32(const &float q, const &float k, const &float v, &float out,
unsigned long heads, unsigned long seqLen, unsigned long kvLen,
unsigned long headDim, float scale, int causal);
int mps_linear_attention_f32(const &float q, const &float k, const &float v, &float out,
unsigned long heads, unsigned long seqLen, unsigned long kvLen,
unsigned long headDim, int causal);
// ── Token-major bf16 attention ──────────────────────────────────────────────
// The same flash and linear attention on one sequence laid out the way a
// transformer's projections produce it -- q/out [seqLen, numHeads*headDim],
// k/v [kvLen, numKvHeads*headDim], raw bf16 bits -- plus a recurrent linear
// attention whose state stays on the GPU. Use mps_attn_batch_bf16 above for
// one-query-per-sequence decode steps across a batch.
//
// mps_flash_attn_bf16: 'scale' as above (pass 1.0 for models that pre-scale
// q); isSliding/slidingWindow as the window w. Batch-aware like
// mps_sub_f32: inside an open batch the result lands in 'out' at
// mps_batch_end(); otherwise on return. Returns 0 on the conditions above.
int mps_flash_attn_bf16(const &unsigned short q, const &unsigned short k, const &unsigned short v,
&unsigned short out,
unsigned long seqLen, unsigned long kvLen,
unsigned long numHeads, unsigned long numKvHeads, unsigned long headDim,
float scale, int causal, int isSliding, unsigned long slidingWindow);
// Chained variant: 'bufQ'/'bufK'/'bufV' are device buffers (pending or
// persistent), must run inside an open batch, and leaves its output as
// MpsBuf::batch_last_output() -- 'out' is throwaway scratch, same
// contract as every other chained op.
int mps_flash_attn_bf16_chained(struct MpsBuf bufQ, struct MpsBuf bufK, struct MpsBuf bufV, &unsigned short out,
unsigned long seqLen, unsigned long kvLen,
unsigned long numHeads, unsigned long numKvHeads, unsigned long headDim,
float scale, int causal, int isSliding, unsigned long slidingWindow);
// Linear attention on token-major bf16, non-causal or causal (causal needs
// kvLen == seqLen here). Batch-aware like mps_flash_attn_bf16.
int mps_linear_attn_bf16(const &unsigned short q, const &unsigned short k, const &unsigned short v,
&unsigned short out,
unsigned long seqLen, unsigned long kvLen,
unsigned long numHeads, unsigned long numKvHeads, unsigned long headDim,
int causal);
// Recurrent (causal) linear attention with the state kept on the GPU:
// 'bufState' holds S and z for every head, starts from the sequence so
// far and is advanced in place by these 'seqLen' tokens (k/v as
// [seqLen, ...]). Feed a prompt in one call, then each decode step with
// seqLen 1 -- O(headDim^2) per token, independent of how long the
// sequence already is. Must run inside an open batch; output as for
// mps_flash_attn_bf16_chained. Calls on the same state must stay in
// order, which a single batch (or batches ended in sequence) guarantees.
int mps_linear_attn_bf16_chained(struct MpsBuf bufQ, struct MpsBuf bufK, struct MpsBuf bufV, struct MpsBuf bufState, &unsigned short out,
unsigned long seqLen,
unsigned long numHeads, unsigned long numKvHeads, unsigned long headDim);
// Size of a linear-attention state: numHeads * (headDim^2 + headDim)
// floats (per head, S^T as [headDim][headDim] then z).
unsigned long mps_linear_attn_state_bytes(unsigned long numHeads, unsigned long headDim);
// A zeroed state for a new sequence: MpsBuf::linear_attn_state (free() it
// when the sequence is done).
// ── Interop: for libraries layered on this backend (device_mps.sc) ───────────────
// The MTLDevice and the MTLCommandQueue every op above runs on; null without
// a Metal device. Command buffers another library commits to the queue run
// in commit order with this backend's batches.
?&MtlObject mps_device_handle();
?&MtlObject mps_queue_handle();
// Runs kernel 'kernel' of the embedded library (gpu_mps_kernels.metal) on
// 'groups' threadgroups of 'groupSize' threads (one dimension): buffers
// bufs bound at indices 0 .. bufs.len-1 and, when params isn't empty, its
// bytes (copied at encode time) at index bufs.len. Encoded onto the open
// batch, or run and waited for outside one, like mps_gemm. The pipeline is
// built on first use and cached by name (at most 64 names; the cache keeps
// the 'kernel' pointer, so pass a string literal). 1 = done/encoded,
// 0 = unknown kernel, an invalid buffer, a group larger than the kernel
// allows, or no Metal device.
int mps_dispatch(const &char kernel, []const MpsBuf bufs, []const unsigned char params,
unsigned long groups, unsigned long groupSize);
} // namespace std