| mps_add_f32 | 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. |
| mps_sub_f32 | 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. |
| mps_mul_f32 | |
| mps_div_f32 | |
| mps_pow_f32 | |
| mps_log_f32 | |
| mps_exp_f32 | |
| mps_sqrt_f32 | |
| mps_relu_f32 | |
| mps_scale_f32 | |
| mps_sum_f32 | 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. |
| mps_matmul_f32 | 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: |
| mps_matmul_abt_f32 | 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. |
| mps_matmul_atb_f32 | |
| mps_relu_backward_f32 | 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. |
| mps_available | 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. |
| mps_batch_begin | 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_end | |
| mps_batch_is_active | |
| mps_batch_submit | 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. |
| mps_batch_register_finalize | 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. |
| mps_matmul_f32_chained | 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). |
| mps_relu_f32_chained | |
| mps_matmul_f32_persistent | 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. |
| mps_sub_f32_chained | 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). |
| mps_mul_f32_chained | 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). |
| mps_square_f32_chained | 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. |
| mps_sum_f32_chained | 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. |
| mps_scale_f32_chained | 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. |
| mps_gated_gelu_f32_chained | Fused gated-activation (GeGLU/tanh-GELU) step: 'bufRaw' is a still-
pending GPU buffer holding one row-major [rows, 2intermediateSize]
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. |
| mps_gated_gelu_bf16_chained | 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. |
| mps_matmul_abt_f32_chained | 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. |
| mps_matmul_atb_f32_chained | |
| mps_relu_backward_f32_chained | |
| mps_relu_backward_f32_chained_ab | 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. |
| mps_matmul_atb_f32_chained_ab | 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. |
| mps_matmul_abt_f32_persistent | 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). |
| mps_sgd_update_f32_chained | 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. |
| mps_matmul_f16 | 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. |
| mps_matmul_bf16 | |
| mps_relu_f16 | |
| mps_relu_bf16 | |
| mps_matmul_f16_persistent | 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. |
| mps_matmul_bf16_persistent | |
| mps_relu_f16_chained | |
| mps_relu_bf16_chained | |
| mps_sub_f16_chained | 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. |
| mps_sub_bf16_chained | |
| mps_scale_f16_chained | 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). |
| mps_scale_bf16_chained | |
| mps_add_bf16_chained | 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. |
| mps_gelu_mul_bf16_chained | out[i] = gelu(bufA[i]) * bufB[i] -- mps_gated_gelu_bf16_chained with the
gate and up operands in two separate device buffers. |
| mps_rmsnorm_bf16_chained | 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. |
| mps_matmul_atb_f16_chained_ab | 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. |
| mps_matmul_atb_bf16_chained_ab | |
| mps_matmul_abt_f16_persistent | dA[M,K] = dC[M,N] . B^T[N,K], B stored [K,N] -- same shape as
mps_matmul_abt_f32_persistent. |
| mps_matmul_abt_bf16_persistent | |
| mps_relu_backward_f16_chained_ab | |
| mps_relu_backward_bf16_chained_ab | |
| mps_sgd_update_f16_chained | In-place SGD step against a persistent fp16/bf16 device buffer --
same shape as mps_sgd_update_f32_chained, weights never leave the GPU. |
| mps_sgd_update_bf16_chained | |
| mps_attn_batch_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. |
| mps_attn_batch_bf16_chained | 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. |
| mps_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. |
| mps_attention_head_dim_ok | |
| mps_flash_attention | |
| mps_flash_attention_ex | |
| mps_linear_attention | |
| mps_linear_attention_ex | |
| mps_flash_attention_f32 | The two fp32 ops on host arrays (upload, compute, read back), with the
rocm_*_f32 arguments. Batch-aware like mps_sub_f32. |
| mps_linear_attention_f32 | |
| mps_flash_attn_bf16 | The same flash and linear attention on one sequence laid out the way a
transformer's projections produce it -- q/out [seqLen, numHeadsheadDim],
k/v [kvLen, numKvHeadsheadDim], 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_chained | 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. |
| mps_linear_attn_bf16 | Linear attention on token-major bf16, non-causal or causal (causal needs
kvLen == seqLen here). Batch-aware like mps_flash_attn_bf16. |
| mps_linear_attn_bf16_chained | 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. |
| mps_linear_attn_state_bytes | Size of a linear-attention state: numHeads * (headDim^2 + headDim)
floats (per head, S^T as [headDim][headDim] then z). |
| mps_device_handle | 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. |
| mps_queue_handle | |
| mps_dispatch | 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. |