// SafeC Standard Library — Metal GPU backend implementation (see
// gpu_mps.h). See std/gui/gui_cocoa.sc's header comment for the general
// "no native Objective-C support" technique this reuses. Every
// objc_msgSend reinterpret-cast is assigned to a named variable before
// being called (rather than called inline as '((SomeMsgSend)objc_msgSend)
// (...)') — the inline form compiles but a real link/run against Metal
// showed it silently keeps objc_msgSend's *original* 2-argument type for
// the call site instead of the cast's, producing a wrong-arg-count LLVM
// verifier failure; the assign-then-call form (matching gui_cocoa.sc's
// existing convention throughout) reliably picks up the cast type.
//
// Three things this file used to do at runtime, on every process or every
// dispatch, that are now either done once or not at all:
// 1. Kernel compilation. Every mps_*_f32 call used to hand a Metal
// Shading Language *source string* to newLibraryWithSource:options:
// error:, a real MSL-to-AIR compile. gpu_mps_kernels.metal is now
// compiled ahead of time by gen_mps_metallib.sh into a .metallib,
// embedded as a byte array (gpu_mps_metallib.h), and loaded once via
// newLibraryWithData:error: — no MSL compiler runs in-process.
// 2. Selector lookup. Every objc_msgSend call used to re-resolve its
// selector via sel_registerName(name) on every single dispatch,
// resolved once now (__mps_init_selectors) into static SEL variables.
// 3. Buffer allocation — the one that actually mattered. Every op used
// to call newBufferWithBytes:length:options:/newBufferWithLength:
// options: fresh, on *every single dispatch*, and never released
// them: a real Metal buffer allocation (IOSurface/IOAccelerator
// resource creation, not a cheap malloc) on every op call, growing
// without bound over a training loop. Measured directly: this, not
// shader compilation, was the actual dominant per-call cost once
// pipelines were already cached — fixing (1)/(2) alone left the
// bigger-model training loop's GPU time unchanged, because neither
// addressed the thing actually being paid for on every op.
// PyTorch's MPSAllocator and MLX's Metal buffer cache both solve
// this the same way: a caching allocator over MTLBuffer objects,
// keyed by size, so a same-shape op reuses a buffer instead of
// asking the driver for a new one — see the size-class buffer pool
// below (__mps_buf_get/__mps_buf_put), which applies that exact
// pattern here, the same size-class-cache design std::alloc/dealloc
// already uses on the CPU side (see std/mem.sc).
// safeguard: link macos -framework Metal -framework Foundation -lobjc
#pragma once
#include <std/stddef.h>
#include <std/ml/gpu_mps.h>
#include <std/ml/mps_ffi.h>
#include <std/ml/gpu_mps_metallib.h>
#include <std/ml/float16.h>
#include <std/mem.h>
namespace std {
extern unsigned char* malloc(size_t size);
extern void free(unsigned char* ptr);
#if defined(__APPLE__)
extern MtlObject* objc_getClass(const char* name);
extern ObjcSel_* sel_registerName(const char* name);
extern MtlObject* objc_msgSend(MtlObject* recv, ObjcSel_* op); // cast to each message's signature
extern unsigned char* memcpy(unsigned char* dst, const unsigned char* src, size_t n);
extern unsigned char* memset(unsigned char* dst, int c, size_t n);
// The real Metal framework entry point — a plain C function (not a
// message send): returns the default MTLDevice, or NULL if this machine
// has no Metal-capable GPU.
extern MtlObject* MTLCreateSystemDefaultDevice();
// libdispatch, part of libSystem on every macOS target (no extra link
// flag needed). Wraps the embedded metallib bytes for
// newLibraryWithData:error: below. queue=NULL and destructor=NULL together
// mean DISPATCH_DATA_DESTRUCTOR_DEFAULT: dispatch_data_create copies the
// buffer into memory it manages, so no block/destructor callback is ever
// invoked — safe to call from a language with no Objective-C block
// literal support.
extern MtlObject* dispatch_data_create(const unsigned char* buffer, unsigned long size, DispatchQueue_* queue,
DispatchBlock_* destructor);
// libobjc's autorelease pools (what @autoreleasepool compiles to): each batch
// runs inside one, see mps_batch_begin.
extern ObjcPool_* objc_autoreleasePoolPush();
extern void objc_autoreleasePoolPop(ObjcPool_* pool);
#else
// No Metal outside Apple platforms: null stand-ins for the Objective-C /
// Metal / libdispatch entry points, so this file still compiles and links
// into std everywhere (mps_available() is 0 and every op reports failure).
extern unsigned char* memcpy(unsigned char* dst, const unsigned char* src, size_t n);
extern unsigned char* memset(unsigned char* dst, int c, size_t n);
static MtlObject* objc_getClass(const char* name) { return null; }
static ObjcSel_* sel_registerName(const char* name) { return null; }
static MtlObject* objc_msgSend(MtlObject* recv, ObjcSel_* op) { return null; }
static MtlObject* MTLCreateSystemDefaultDevice() { return null; }
static MtlObject* dispatch_data_create(const unsigned char* buffer, unsigned long size, DispatchQueue_* queue,
DispatchBlock_* destructor) { return null; }
static ObjcPool_* objc_autoreleasePoolPush() { return null; }
static void objc_autoreleasePoolPop(ObjcPool_* pool) { }
#endif
struct MTLSize {
unsigned long width;
unsigned long height;
unsigned long depth;
};
typedef fn MtlObject*(MtlObject*, ObjcSel_*) MMsg0;
typedef fn MtlObject*(MtlObject*, ObjcSel_*, MtlObject*) MMsg1;
typedef fn MtlObject*(MtlObject*, ObjcSel_*, const char*) MMsgCStr;
// ...:error: messages: the last argument is where an NSError* goes
typedef fn MtlObject*(MtlObject*, ObjcSel_*, MtlObject*, struct MtlObject**) MMsg2;
typedef fn MtlObject*(MtlObject*, ObjcSel_*, unsigned long, unsigned long) MMsgNewBufferLen;
typedef fn void(MtlObject*, ObjcSel_*, MtlObject*) MMsgVoid1;
typedef fn void(MtlObject*, ObjcSel_*, MtlObject*, unsigned long, unsigned long) MMsgSetBuffer;
typedef fn void(MtlObject*, ObjcSel_*, const unsigned char*, unsigned long, unsigned long) MMsgSetBytes;
typedef fn unsigned char*(MtlObject*, ObjcSel_*) MMsgBytes0;
typedef fn void(MtlObject*, ObjcSel_*, struct MTLSize, struct MTLSize) MMsgDispatch;
typedef fn void(MtlObject*, ObjcSel_*) MMsgVoid0;
// ── Cached selectors ──────────────────────────────────────────────────────
// Filled once by __mps_init_selectors (called from __mps_ensure_device_queue
// the first time any op runs), then read directly everywhere below instead
// of calling sel_registerName(name) per dispatch.
static ObjcSel_* __sel_alloc = null;
static ObjcSel_* __sel_initWithUTF8String = null;
static ObjcSel_* __sel_newCommandQueue = null;
static ObjcSel_* __sel_newLibraryWithData = null;
static ObjcSel_* __sel_newFunctionWithName = null;
static ObjcSel_* __sel_newComputePipelineState = null;
static ObjcSel_* __sel_newBufferWithLength = null;
static ObjcSel_* __sel_commandBuffer = null;
static ObjcSel_* __sel_computeCommandEncoder = null;
static ObjcSel_* __sel_setComputePipelineState = null;
static ObjcSel_* __sel_setBuffer = null;
static ObjcSel_* __sel_setBytes = null;
static ObjcSel_* __sel_dispatchThreadgroups = null;
// Exact-grid dispatch for the one-thread-per-element kernels: their bodies
// have no bounds check, so rounding the grid up to whole threadgroups would
// run threads past 'n' (out-of-bounds writes when n isn't a multiple of the
// threadgroup width). Non-uniform threadgroups are supported on every Apple
// GPU this backend runs on.
static ObjcSel_* __sel_dispatchThreads = null;
static ObjcSel_* __sel_endEncoding = null;
static ObjcSel_* __sel_commit = null;
static ObjcSel_* __sel_waitUntilCompleted = null;
static ObjcSel_* __sel_contents = null;
static ObjcSel_* __sel_release = null;
// A buffer's -contents: its bytes in host memory (shared storage).
static unsigned char* __mps_contents(MtlObject* buf) {
unsafe {
MMsgBytes0 f = (MMsgBytes0)objc_msgSend;
return f(buf, __sel_contents);
}
}
static void __mps_init_selectors() {
unsafe {
__sel_alloc = sel_registerName("alloc");
__sel_initWithUTF8String = sel_registerName("initWithUTF8String:");
__sel_newCommandQueue = sel_registerName("newCommandQueue");
__sel_newLibraryWithData = sel_registerName("newLibraryWithData:error:");
__sel_newFunctionWithName = sel_registerName("newFunctionWithName:");
__sel_newComputePipelineState = sel_registerName("newComputePipelineStateWithFunction:error:");
__sel_newBufferWithLength = sel_registerName("newBufferWithLength:options:");
__sel_commandBuffer = sel_registerName("commandBuffer");
__sel_computeCommandEncoder = sel_registerName("computeCommandEncoder");
__sel_setComputePipelineState = sel_registerName("setComputePipelineState:");
__sel_setBuffer = sel_registerName("setBuffer:offset:atIndex:");
__sel_setBytes = sel_registerName("setBytes:length:atIndex:");
__sel_dispatchThreadgroups = sel_registerName("dispatchThreadgroups:threadsPerThreadgroup:");
__sel_dispatchThreads = sel_registerName("dispatchThreads:threadsPerThreadgroup:");
__sel_endEncoding = sel_registerName("endEncoding");
__sel_commit = sel_registerName("commit");
__sel_waitUntilCompleted = sel_registerName("waitUntilCompleted");
__sel_contents = sel_registerName("contents");
__sel_release = sel_registerName("release");
}
}
static MtlObject* __mps_nsstring(const char* cstr) {
unsafe {
MtlObject* cls = objc_getClass("NSString");
MMsg0 allocFn = (MMsg0)objc_msgSend;
MtlObject* obj = allocFn(cls, __sel_alloc);
MMsgCStr initFn = (MMsgCStr)objc_msgSend;
return initFn(obj, __sel_initWithUTF8String, cstr);
}
}
int mps_available() {
MtlObject* dev;
unsafe { dev = MTLCreateSystemDefaultDevice(); }
return dev != null;
}
// Cached device/queue/library, shared by every op in this file. Lazily
// created on whichever op runs first, then reused forever.
static MtlObject* __mps_device = null;
static MtlObject* __mps_queue = null;
static MtlObject* __mps_library = null;
static int __mps_ensure_device_queue() {
if (__mps_device != null) return 1;
unsafe {
MtlObject* dev = MTLCreateSystemDefaultDevice();
if (dev == null) return 0;
__mps_init_selectors();
MMsg0 msg0init = (MMsg0)objc_msgSend;
MtlObject* q = msg0init(dev, __sel_newCommandQueue);
if (q == null) return 0;
__mps_device = dev;
__mps_queue = q;
return 1;
}
}
// Loads the whole kernel set (gpu_mps_kernels.metal, precompiled to AIR
// bytecode by gen_mps_metallib.sh and embedded in gpu_mps_metallib.h) as
// one MTLLibrary, once. Every op below then just looks up its own function
// by name from this shared library — no MSL source, no runtime compile.
static int __mps_ensure_library() {
if (__mps_library != null) return 1;
if (!__mps_ensure_device_queue()) return 0;
unsafe {
MtlObject* data = dispatch_data_create((const unsigned char*)__mps_metallib_bytes, __mps_metallib_size,
null, null);
if (data == null) return 0;
MtlObject* errPtr = null;
MMsg2 newLibFn = (MMsg2)objc_msgSend;
MtlObject* library = newLibFn(__mps_device, __sel_newLibraryWithData, data, &errPtr);
if (library == null) return 0;
__mps_library = library;
return 1;
}
}
// Looks up 'kernelName' in the shared library and builds a compute
// pipeline for it. Shared by every op's pipeline-cache-fill branch below.
static MtlObject* __mps_make_pipeline(const char* kernelName) {
unsafe {
if (!__mps_ensure_library()) return null;
MtlObject* fnNameStr = __mps_nsstring(kernelName);
MMsg1 msg1 = (MMsg1)objc_msgSend;
MtlObject* kernelFn = msg1(__mps_library, __sel_newFunctionWithName, fnNameStr);
if (kernelFn == null) return null;
MtlObject* errPtr = null;
MMsg2 newPipelineFn = (MMsg2)objc_msgSend;
MtlObject* pipeline = newPipelineFn(__mps_device, __sel_newComputePipelineState, kernelFn, &errPtr);
return pipeline;
}
}
// ── Size-class MTLBuffer pool ────────────────────────────────────────────────
// Same idea as std::alloc/dealloc's size-class cache (std/mem.sc), applied
// to Metal buffers instead of heap blocks: a freed buffer goes back into a
// thread-local(-ish; this file was never thread-safe — see __mps_device's
// own lack of synchronization) free list bucketed by power-of-two size
// class, and the next request for that class is satisfied from there
// instead of asking the Metal driver for a brand new buffer. Classes run
// 4KB to 8MB (12 classes) — small enough to cover every buffer this file's
// actual op set ever allocates (a batch of matmul/relu/elementwise calls
// over a training-loop-sized MLP, not arbitrary user data), large enough
// that a single class comfortably holds this bigger benchmark's biggest
// tensor (W1 at 512x1024 doubles-converted-to-float32 = 2MB). Anything
// larger just allocates directly, uncached, same as std::alloc's oversize
// path. Per-class capacity is memory-budget-scaled (roughly 128MB worth of
// buffers per class, floored at 2 slots, capped at 128 slots) rather than a
// flat count, so a class of huge buffers doesn't retain as many as a class
// of small ones — same reasoning as alloc_class_cap_ in std/mem.sc. Raised
// from an original 8-slot/16MB cap: with many independent ops sharing one
// batch (see train_gpu.sc's inference loop — dozens of matmul/relu passes
// batched together purely to amortize one command-buffer submit + wait
// across all of them), every one of those ops needs its own buffer of the
// same handful of sizes (the fixed X/W1/W2/activation shapes), and NONE of
// them can return to the pool until the whole batch's wait completes — a
// small pool that's fine for one training step's ~9 ops was measured
// forcing hundreds of real newBufferWithLength allocations (not just pool
// misses) once a single batch's own concurrent demand for one size class
// exceeded 8, which cost far more than the extra command-buffer submits it
// was meant to avoid. A slot that overflows its class's cap gets an
// explicit 'release' instead of being silently dropped (leaked): Metal
// buffers are real, non-trivial GPU resources, not a few bytes of heap, so
// an unbounded leak here is worth avoiding even in
// the overflow case std::alloc's own quarantine ring tolerates for plain
// memory.
#define MPS_BUF_MIN_CLASS_ ((unsigned long)4096)
#define MPS_BUF_NUM_CLASSES_ ((unsigned long)12)
#define MPS_BUF_MAX_SLOTS_ ((unsigned long)128)
static MtlObject* __mps_bufSlot[12][128];
static unsigned long __mps_bufCount[12];
static long __mps_buf_class_for(unsigned long bytes) {
unsigned long clsBytes = MPS_BUF_MIN_CLASS_;
long idx = 0L;
while (clsBytes < bytes) {
if (idx >= (long)(MPS_BUF_NUM_CLASSES_ - (unsigned long)1)) { return -1L; }
clsBytes = clsBytes << 1UL;
idx = idx + 1L;
}
return idx;
}
static unsigned long __mps_buf_class_bytes(unsigned long idx) {
return MPS_BUF_MIN_CLASS_ << idx;
}
// ~128MB budget per class, floored at 2 slots, capped at the array's
// physical width (MPS_BUF_MAX_SLOTS_) -- same shift-based, division-free
// approach as std/mem.sc's alloc_class_cap_.
static unsigned long __mps_buf_class_cap(unsigned long cls) {
unsigned long classBytes = __mps_buf_class_bytes(cls);
unsigned long cap = ((unsigned long)128 * (unsigned long)1024 * (unsigned long)1024) / classBytes;
if (cap > MPS_BUF_MAX_SLOTS_) { cap = MPS_BUF_MAX_SLOTS_; }
if (cap < (unsigned long)2) { cap = (unsigned long)2; }
return cap;
}
// Returns an MTLBuffer of at least 'bytes', either popped from its size
// class's free list or freshly allocated (at the class's rounded-up size,
// so it's reusable by any future request that rounds to the same class).
static MtlObject* __mps_buf_get(unsigned long bytes) {
unsafe {
long clsSigned = __mps_buf_class_for(bytes);
MMsgNewBufferLen newBufLenFn = (MMsgNewBufferLen)objc_msgSend;
if (clsSigned >= 0L) {
unsigned long cls = (unsigned long)clsSigned;
if (__mps_bufCount[cls] > 0UL) {
__mps_bufCount[cls] = __mps_bufCount[cls] - 1UL;
return __mps_bufSlot[cls][__mps_bufCount[cls]];
}
return newBufLenFn(__mps_device, __sel_newBufferWithLength, __mps_buf_class_bytes(cls), 0UL);
}
return newBufLenFn(__mps_device, __sel_newBufferWithLength, bytes, 0UL);
}
}
// Same as __mps_buf_get, but also copies 'src' into the buffer's
// CPU-visible 'contents' (storage mode 0 = MTLResourceStorageModeShared,
// set on every allocation above) -- replaces the old
// newBufferWithBytes:length:options: call, which allocated and copied in
// one step but could never be satisfied from the pool.
static MtlObject* __mps_buf_get_with_bytes(const unsigned char* src, unsigned long bytes) {
unsafe {
MtlObject* buf = __mps_buf_get(bytes);
if (buf == null) return null;
MMsg0 msg0 = (MMsg0)objc_msgSend;
unsigned char* contents = __mps_contents(buf);
if (contents != null) { memcpy(contents, src, bytes); }
return buf;
}
}
// Returns 'buf' (of size 'bytes') to its size class's free list, or
// releases it outright if that class is already at capacity. Only valid
// to call once the GPU is provably done with the buffer -- i.e. after the
// command buffer that used it has completed (waitUntilCompleted returned,
// whether that wait happened here or, for batched dispatches, in
// mps_batch_end()).
static void __mps_buf_put(MtlObject* buf, unsigned long bytes) {
if (buf == null) return;
unsafe {
long clsSigned = __mps_buf_class_for(bytes);
if (clsSigned >= 0L) {
unsigned long cls = (unsigned long)clsSigned;
if (__mps_bufCount[cls] < __mps_buf_class_cap(cls)) {
__mps_bufSlot[cls][__mps_bufCount[cls]] = buf;
__mps_bufCount[cls] = __mps_bufCount[cls] + 1UL;
return;
}
}
MMsgVoid0 releaseFn = (MMsgVoid0)objc_msgSend;
releaseFn(buf, __sel_release);
}
}
// Shared setup/dispatch/readback for every elementwise binary kernel
// (add/sub/mul/div/pow) — the only thing that varies between them is which
// compiled kernel function they bind. Previously, the
// dispatchThreadgroups:threadsPerThreadgroup: call inside here segfaulted
// on real Apple Silicon hardware: it passes *two* struct arguments (each
// 'struct MTLSize', 3x unsigned long = 24 bytes) by value through an
// objc_msgSend cast. Root-caused and fixed in CodeGen.cpp's genCall:
// indirect calls (through a function-pointer cast, exactly this pattern)
// now lower non-HFA struct-by-value arguments over 16 bytes to the real
// AAPCS64 form (caller copies to a stack temporary, passes a plain
// pointer) instead of a raw aggregate value.
//
// 'pipelineCache' is a pointer to the CALLER's own static MtlObject* (one per
// op — see e.g. __mps_add_pipeline below): each op needs its own compiled
// MTLComputePipelineState, but the compile-once-reuse-forever pattern is
// identical for all of them.
static int __mps_run_binary_kernel(struct MtlObject** pipelineCache, const char* kernelName,
const float* a, const float* b, float* out, unsigned long n) {
unsafe {
if (!__mps_ensure_device_queue()) return 0;
MtlObject* device = __mps_device;
MtlObject* queue = __mps_queue;
MMsg0 msg0 = (MMsg0)objc_msgSend;
if (*pipelineCache == null) {
MtlObject* pipeline = __mps_make_pipeline(kernelName);
if (pipeline == null) return 0;
*pipelineCache = pipeline;
}
MtlObject* pipeline = *pipelineCache;
unsigned long bytes = n * sizeof(float);
MtlObject* bufA = __mps_buf_get_with_bytes((const unsigned char*)a, bytes);
MtlObject* bufB = __mps_buf_get_with_bytes((const unsigned char*)b, bytes);
MtlObject* bufOut = __mps_buf_get(bytes);
if (bufA == null || bufB == null || bufOut == null) return 0;
MtlObject* cmdBuf = msg0(queue, __sel_commandBuffer);
MtlObject* encoder = msg0(cmdBuf, __sel_computeCommandEncoder);
MMsgVoid1 msgVoid1 = (MMsgVoid1)objc_msgSend;
msgVoid1(encoder, __sel_setComputePipelineState, pipeline);
MMsgSetBuffer setBufFn = (MMsgSetBuffer)objc_msgSend;
setBufFn(encoder, __sel_setBuffer, bufA, 0UL, 0UL);
setBufFn(encoder, __sel_setBuffer, bufB, 0UL, 1UL);
setBufFn(encoder, __sel_setBuffer, bufOut, 0UL, 2UL);
struct MTLSize gridSize;
gridSize.width = n; gridSize.height = 1UL; gridSize.depth = 1UL;
unsigned long tgWidth = n < 64UL ? n : 64UL;
if (tgWidth == 0UL) tgWidth = 1UL;
struct MTLSize tgSize;
tgSize.width = tgWidth; tgSize.height = 1UL; tgSize.depth = 1UL;
MMsgDispatch dispatchFn = (MMsgDispatch)objc_msgSend;
dispatchFn(encoder, __sel_dispatchThreads, gridSize, tgSize);
MMsgVoid0 msgVoid0 = (MMsgVoid0)objc_msgSend;
msgVoid0(encoder, __sel_endEncoding);
msgVoid0(cmdBuf, __sel_commit);
msgVoid0(cmdBuf, __sel_waitUntilCompleted);
unsigned char* outPtr = __mps_contents(bufOut);
if (outPtr == null) return 0;
memcpy((unsigned char*)out, outPtr, bytes);
__mps_buf_put(bufA, bytes);
__mps_buf_put(bufB, bytes);
__mps_buf_put(bufOut, bytes);
return 1;
}
}
// Same as __mps_run_binary_kernel but for a single-input elementwise kernel
// (log/exp/sqrt) — one input buffer instead of two, otherwise identical
// setup/dispatch/readback.
static int __mps_run_unary_kernel(struct MtlObject** pipelineCache, const char* kernelName,
const float* a, float* out, unsigned long n) {
unsafe {
if (!__mps_ensure_device_queue()) return 0;
MtlObject* device = __mps_device;
MtlObject* queue = __mps_queue;
MMsg0 msg0 = (MMsg0)objc_msgSend;
if (*pipelineCache == null) {
MtlObject* pipeline = __mps_make_pipeline(kernelName);
if (pipeline == null) return 0;
*pipelineCache = pipeline;
}
MtlObject* pipeline = *pipelineCache;
unsigned long bytes = n * sizeof(float);
MtlObject* bufA = __mps_buf_get_with_bytes((const unsigned char*)a, bytes);
MtlObject* bufOut = __mps_buf_get(bytes);
if (bufA == null || bufOut == null) return 0;
MtlObject* cmdBuf = msg0(queue, __sel_commandBuffer);
MtlObject* encoder = msg0(cmdBuf, __sel_computeCommandEncoder);
MMsgVoid1 msgVoid1 = (MMsgVoid1)objc_msgSend;
msgVoid1(encoder, __sel_setComputePipelineState, pipeline);
MMsgSetBuffer setBufFn = (MMsgSetBuffer)objc_msgSend;
setBufFn(encoder, __sel_setBuffer, bufA, 0UL, 0UL);
setBufFn(encoder, __sel_setBuffer, bufOut, 0UL, 1UL);
struct MTLSize gridSize;
gridSize.width = n; gridSize.height = 1UL; gridSize.depth = 1UL;
unsigned long tgWidth = n < 64UL ? n : 64UL;
if (tgWidth == 0UL) tgWidth = 1UL;
struct MTLSize tgSize;
tgSize.width = tgWidth; tgSize.height = 1UL; tgSize.depth = 1UL;
MMsgDispatch dispatchFn = (MMsgDispatch)objc_msgSend;
dispatchFn(encoder, __sel_dispatchThreads, gridSize, tgSize);
MMsgVoid0 msgVoid0 = (MMsgVoid0)objc_msgSend;
msgVoid0(encoder, __sel_endEncoding);
msgVoid0(cmdBuf, __sel_commit);
msgVoid0(cmdBuf, __sel_waitUntilCompleted);
unsigned char* outPtr = __mps_contents(bufOut);
if (outPtr == null) return 0;
memcpy((unsigned char*)out, outPtr, bytes);
__mps_buf_put(bufA, bytes);
__mps_buf_put(bufOut, bytes);
return 1;
}
}
// One compiled-pipeline cache per elementwise op, each fed to
// __mps_run_binary_kernel/__mps_run_unary_kernel by address so it's filled
// (once) and reused (every call after) independently of every other op's
// cache.
static MtlObject* __mps_add_pipeline = null;
static MtlObject* __mps_sub_pipeline = null;
static MtlObject* __mps_mul_pipeline = null;
static MtlObject* __mps_div_pipeline = null;
static MtlObject* __mps_pow_pipeline = null;
static MtlObject* __mps_log_pipeline = null;
static MtlObject* __mps_exp_pipeline = null;
static MtlObject* __mps_sqrt_pipeline = null;
static MtlObject* __mps_relu_backward_pipeline = null;
int mps_add_f32_n_(const &float a, const &float b, &float out, unsigned long n) {
return __mps_run_binary_kernel(&__mps_add_pipeline, "add_kernel", a, b, out, n);
}
int mps_add_f32([]const float a_s, []const float b_s, []float out_s) {
unsigned long n = out_s.len;
if (a_s.len < n || b_s.len < n) { return 0; }
float a_d_[1];
const &float a = &a_d_[0];
if (a_s.len > 0UL) { unsafe { a = (const &float)a_s.ptr; } }
float b_d_[1];
const &float b = &b_d_[0];
if (b_s.len > 0UL) { unsafe { b = (const &float)b_s.ptr; } }
float out_d_[1];
&float out = &out_d_[0];
if (out_s.len > 0UL) { unsafe { out = (&float)out_s.ptr; } }
return mps_add_f32_n_(a, b, out, n);
}
// mps_relu_backward_f32's own definition (batch-aware, not routed through
// __mps_run_binary_kernel) lives further down, next to the matmul-backward
// pair it's used alongside in a real backward pass -- see that section's
// header comment.
// mps_sub_f32's own definition (batch-aware, not routed through
// __mps_run_binary_kernel) lives further down, next to mps_matmul_f32 and
// the batching primitives it needs to reference -- same reasoning as
// mps_relu_f32/mps_relu_backward_f32 above (see mps_relu_f32's comment).
int mps_mul_f32_n_(const &float a, const &float b, &float out, unsigned long n) {
return __mps_run_binary_kernel(&__mps_mul_pipeline, "mul_kernel", a, b, out, n);
}
int mps_mul_f32([]const float a_s, []const float b_s, []float out_s) {
unsigned long n = out_s.len;
if (a_s.len < n || b_s.len < n) { return 0; }
float a_d_[1];
const &float a = &a_d_[0];
if (a_s.len > 0UL) { unsafe { a = (const &float)a_s.ptr; } }
float b_d_[1];
const &float b = &b_d_[0];
if (b_s.len > 0UL) { unsafe { b = (const &float)b_s.ptr; } }
float out_d_[1];
&float out = &out_d_[0];
if (out_s.len > 0UL) { unsafe { out = (&float)out_s.ptr; } }
return mps_mul_f32_n_(a, b, out, n);
}
int mps_div_f32_n_(const &float a, const &float b, &float out, unsigned long n) {
return __mps_run_binary_kernel(&__mps_div_pipeline, "div_kernel", a, b, out, n);
}
int mps_div_f32([]const float a_s, []const float b_s, []float out_s) {
unsigned long n = out_s.len;
if (a_s.len < n || b_s.len < n) { return 0; }
float a_d_[1];
const &float a = &a_d_[0];
if (a_s.len > 0UL) { unsafe { a = (const &float)a_s.ptr; } }
float b_d_[1];
const &float b = &b_d_[0];
if (b_s.len > 0UL) { unsafe { b = (const &float)b_s.ptr; } }
float out_d_[1];
&float out = &out_d_[0];
if (out_s.len > 0UL) { unsafe { out = (&float)out_s.ptr; } }
return mps_div_f32_n_(a, b, out, n);
}
int mps_pow_f32_n_(const &float a, const &float b, &float out, unsigned long n) {
return __mps_run_binary_kernel(&__mps_pow_pipeline, "pow_kernel", a, b, out, n);
}
int mps_pow_f32([]const float a_s, []const float b_s, []float out_s) {
unsigned long n = out_s.len;
if (a_s.len < n || b_s.len < n) { return 0; }
float a_d_[1];
const &float a = &a_d_[0];
if (a_s.len > 0UL) { unsafe { a = (const &float)a_s.ptr; } }
float b_d_[1];
const &float b = &b_d_[0];
if (b_s.len > 0UL) { unsafe { b = (const &float)b_s.ptr; } }
float out_d_[1];
&float out = &out_d_[0];
if (out_s.len > 0UL) { unsafe { out = (&float)out_s.ptr; } }
return mps_pow_f32_n_(a, b, out, n);
}
int mps_log_f32_n_(const &float a, &float out, unsigned long n) {
return __mps_run_unary_kernel(&__mps_log_pipeline, "log_kernel", a, out, n);
}
int mps_log_f32([]const float a_s, []float out_s) {
unsigned long n = out_s.len;
if (a_s.len < n) { return 0; }
float a_d_[1];
const &float a = &a_d_[0];
if (a_s.len > 0UL) { unsafe { a = (const &float)a_s.ptr; } }
float out_d_[1];
&float out = &out_d_[0];
if (out_s.len > 0UL) { unsafe { out = (&float)out_s.ptr; } }
return mps_log_f32_n_(a, out, n);
}
int mps_exp_f32_n_(const &float a, &float out, unsigned long n) {
return __mps_run_unary_kernel(&__mps_exp_pipeline, "exp_kernel", a, out, n);
}
int mps_exp_f32([]const float a_s, []float out_s) {
unsigned long n = out_s.len;
if (a_s.len < n) { return 0; }
float a_d_[1];
const &float a = &a_d_[0];
if (a_s.len > 0UL) { unsafe { a = (const &float)a_s.ptr; } }
float out_d_[1];
&float out = &out_d_[0];
if (out_s.len > 0UL) { unsafe { out = (&float)out_s.ptr; } }
return mps_exp_f32_n_(a, out, n);
}
int mps_sqrt_f32_n_(const &float a, &float out, unsigned long n) {
return __mps_run_unary_kernel(&__mps_sqrt_pipeline, "sqrt_kernel", a, out, n);
}
int mps_sqrt_f32([]const float a_s, []float out_s) {
unsigned long n = out_s.len;
if (a_s.len < n) { return 0; }
float a_d_[1];
const &float a = &a_d_[0];
if (a_s.len > 0UL) { unsafe { a = (const &float)a_s.ptr; } }
float out_d_[1];
&float out = &out_d_[0];
if (out_s.len > 0UL) { unsafe { out = (&float)out_s.ptr; } }
return mps_sqrt_f32_n_(a, out, n);
}
// mps_relu_f32's own definition lives further down, alongside
// mps_matmul_f32 and the batching primitives it (like matmul) needs to
// check — see the comment there for why it's not routed through
// __mps_run_unary_kernel like log/exp/sqrt above it.
// out[i] = a[i] * k for a compile-time-unknown scalar k -- same shape as
// __mps_run_unary_kernel but with one extra setBytes:length:atIndex: call
// to pass k into buffer index 1 (out moves to index 2) before dispatch --
// kept as its own function (rather than folded into __mps_run_unary_kernel)
// because of that extra setBytes call.
static MtlObject* __mps_scale_pipeline = null;
static MtlObject* __mps_gated_gelu_pipeline = null;
int mps_scale_f32_n_(const &float a, float k, &float out, unsigned long n) {
unsafe {
if (!__mps_ensure_device_queue()) return 0;
MtlObject* device = __mps_device;
MtlObject* queue = __mps_queue;
MMsg0 msg0 = (MMsg0)objc_msgSend;
if (__mps_scale_pipeline == null) {
MtlObject* pipeline = __mps_make_pipeline("scale_kernel");
if (pipeline == null) return 0;
__mps_scale_pipeline = pipeline;
}
MtlObject* pipeline = __mps_scale_pipeline;
unsigned long bytes = n * sizeof(float);
MtlObject* bufA = __mps_buf_get_with_bytes((const unsigned char*)a, bytes);
MtlObject* bufOut = __mps_buf_get(bytes);
if (bufA == null || bufOut == null) return 0;
MtlObject* cmdBuf = msg0(queue, __sel_commandBuffer);
MtlObject* encoder = msg0(cmdBuf, __sel_computeCommandEncoder);
MMsgVoid1 msgVoid1 = (MMsgVoid1)objc_msgSend;
msgVoid1(encoder, __sel_setComputePipelineState, pipeline);
MMsgSetBuffer setBufFn = (MMsgSetBuffer)objc_msgSend;
setBufFn(encoder, __sel_setBuffer, bufA, 0UL, 0UL);
// setBytes:length:atIndex: shares setBuffer:offset:atIndex:'s
// (recv, sel, ptr, NSUInteger, NSUInteger) shape -- see
// mps_matmul_f32's comment on reusing MMsgSetBuffer for it.
MMsgSetBytes setBytesFn = (MMsgSetBytes)objc_msgSend;
float kVal = k;
setBytesFn(encoder, __sel_setBytes, (const unsigned char*)&kVal, 4UL, 1UL);
setBufFn(encoder, __sel_setBuffer, bufOut, 0UL, 2UL);
struct MTLSize gridSize;
gridSize.width = n; gridSize.height = 1UL; gridSize.depth = 1UL;
unsigned long tgWidth = n < 64UL ? n : 64UL;
if (tgWidth == 0UL) tgWidth = 1UL;
struct MTLSize tgSize;
tgSize.width = tgWidth; tgSize.height = 1UL; tgSize.depth = 1UL;
MMsgDispatch dispatchFn = (MMsgDispatch)objc_msgSend;
dispatchFn(encoder, __sel_dispatchThreads, gridSize, tgSize);
MMsgVoid0 msgVoid0 = (MMsgVoid0)objc_msgSend;
msgVoid0(encoder, __sel_endEncoding);
msgVoid0(cmdBuf, __sel_commit);
msgVoid0(cmdBuf, __sel_waitUntilCompleted);
unsigned char* outPtr = __mps_contents(bufOut);
if (outPtr == null) return 0;
memcpy((unsigned char*)out, outPtr, bytes);
__mps_buf_put(bufA, bytes);
__mps_buf_put(bufOut, bytes);
return 1;
}
}
int mps_scale_f32([]const float a_s, float k, []float out_s) {
unsigned long n = out_s.len;
if (a_s.len < n) { return 0; }
float a_d_[1];
const &float a = &a_d_[0];
if (a_s.len > 0UL) { unsafe { a = (const &float)a_s.ptr; } }
float out_d_[1];
&float out = &out_d_[0];
if (out_s.len > 0UL) { unsafe { out = (&float)out_s.ptr; } }
return mps_scale_f32_n_(a, k, out, n);
}
// out[0] = sum(a[0..n)) -- two-stage parallel reduction (see
// gpu_mps_kernels.metal's sum_kernel comment for the full design): a
// bounded number of threadgroups each grid-stride-accumulate then
// tree-reduce their own share of the input in threadgroup memory
// (O(log SUM_TG_SIZE) span instead of the old single-thread version's
// O(n)), emitting one partial sum per threadgroup; this function combines
// those (at most SUM_MAX_GROUPS of them, however large n is) on the CPU
// after one readback, which is cheap enough not to need a second GPU pass.
#define SUM_TG_SIZE ((unsigned long)256)
#define SUM_MAX_GROUPS ((unsigned long)256)
static MtlObject* __mps_sum_pipeline = null;
// mps_sum_f32's own definition (batch-aware, deferred-finalize partial-sum
// combine) lives further down, next to mps_matmul_f32/mps_sub_f32 and the
// batching primitives it needs to reference -- same reasoning as
// mps_relu_f32's comment above.
// ── Batched dispatch (see gpu_mps.h's mps_batch_begin comment for the
// full contract) ───────────────────────────────────────────────────────────
// __mps_batch_active/_cmdBuf/_encoder track one in-flight shared command
// buffer. __mps_pending_* is a fixed-size array of "GPU buffer -> CPU
// pointer" raw-readback entries collected while a batch is open — originally
// sized for a single training step's op count (a handful), raised to cover
// many independent ops sharing one batch too (e.g. dozens of inference
// passes batched together purely to amortize one command-buffer submit +
// wait across all of them -- see train_gpu.sc's inference loop). A real
// growable Vec would work too, but this file doesn't otherwise depend on
// std/collections, and a bounds-checked fixed array covers every actual
// use here without adding that dependency. __mps_finalize_* is the
// parallel list of caller-supplied post-readback callbacks (e.g.
// the old per-op path's float32->float64 conversion). __mps_poolbuf_* is a
// third, separate list: every buffer __mps_buf_get/_with_bytes hands out
// *while a batch is open* stays in flight until the whole batch's command
// buffer completes (the GPU may still be reading or writing it right up
// until waitUntilCompleted returns in mps_batch_end() below), so it can't
// go back to the pool at the point it's used the way the non-batched path
// does -- it's tracked here instead and returned in bulk once the wait is
// done. A buffer obtained via a _chained call's 'bufA' parameter (i.e.
// someone else's still-pending output, not a fresh __mps_buf_get) is never
// double-tracked, because tracking only happens at the __mps_buf_get call
// site itself, and chained functions never call it for their input.
#define MPS_BATCH_MAX_PENDING 1024
#define MPS_BATCH_MAX_FINALIZE 64
#define MPS_BATCH_MAX_POOLBUFS 1024
static int __mps_batch_active = 0;
static ObjcPool_* __mps_batch_pool = null;
static MtlObject* __mps_batch_cmdBuf = null;
static MtlObject* __mps_batch_encoder = null;
static MtlObject* __mps_pending_gpuBuf[MPS_BATCH_MAX_PENDING];
static unsigned char* __mps_pending_cpuOut[MPS_BATCH_MAX_PENDING];
static unsigned long __mps_pending_bytes[MPS_BATCH_MAX_PENDING];
static int __mps_pending_dtype[MPS_BATCH_MAX_PENDING]; // MpsBuf::batch_last_output's dtype
static unsigned long __mps_pending_count = 0UL;
static MpsFinalizeFn __mps_finalize_fn[MPS_BATCH_MAX_FINALIZE];
static ?&Erased __mps_finalize_ctx[MPS_BATCH_MAX_FINALIZE];
static unsigned long __mps_finalize_count = 0UL;
static MtlObject* __mps_poolbuf_buf[MPS_BATCH_MAX_POOLBUFS];
static unsigned long __mps_poolbuf_bytes[MPS_BATCH_MAX_POOLBUFS];
static unsigned long __mps_poolbuf_count = 0UL;
static void __mps_batch_track_pool_buf(MtlObject* buf, unsigned long bytes) {
if (buf == null) return;
if (__mps_poolbuf_count < (unsigned long)MPS_BATCH_MAX_POOLBUFS) {
__mps_poolbuf_buf[__mps_poolbuf_count] = buf;
__mps_poolbuf_bytes[__mps_poolbuf_count] = bytes;
__mps_poolbuf_count = __mps_poolbuf_count + 1UL;
}
}
// Batch-scoped versions of __mps_buf_get/_with_bytes: same allocation, but
// also register the buffer for deferred pool-return in mps_batch_end().
static MtlObject* __mps_buf_get_batched(unsigned long bytes) {
MtlObject* buf = __mps_buf_get(bytes);
__mps_batch_track_pool_buf(buf, bytes);
return buf;
}
static MtlObject* __mps_buf_get_with_bytes_batched(const unsigned char* src, unsigned long bytes) {
MtlObject* buf = __mps_buf_get_with_bytes(src, bytes);
__mps_batch_track_pool_buf(buf, bytes);
return buf;
}
int mps_batch_is_active() { return __mps_batch_active; }
void mps_batch_begin() {
unsafe {
if (!__mps_ensure_device_queue()) return;
// The command buffers and encoders below come back autoreleased, and
// a plain C program has no pool to drain them: without this one they
// were never freed (a training loop grew ~30 KB per epoch).
// mps_batch_end drains it once the GPU is done with them.
__mps_batch_pool = objc_autoreleasePoolPush();
MMsg0 msg0 = (MMsg0)objc_msgSend;
__mps_batch_cmdBuf = msg0(__mps_queue, __sel_commandBuffer);
__mps_batch_encoder = msg0(__mps_batch_cmdBuf, __sel_computeCommandEncoder);
__mps_pending_count = 0UL;
__mps_finalize_count = 0UL;
__mps_poolbuf_count = 0UL;
__mps_batch_active = 1;
}
}
void mps_batch_register_finalize_raw_(MpsFinalizeFn finalizeFn, ?&Erased ctx) {
if (__mps_finalize_count >= (unsigned long)MPS_BATCH_MAX_FINALIZE) return;
__mps_finalize_fn[__mps_finalize_count] = finalizeFn;
unsafe { __mps_finalize_ctx[__mps_finalize_count] = ctx; }
__mps_finalize_count = __mps_finalize_count + 1UL;
}
void mps_batch_end() {
if (!__mps_batch_active) return;
unsafe {
MMsgVoid0 msgVoid0 = (MMsgVoid0)objc_msgSend;
msgVoid0(__mps_batch_encoder, __sel_endEncoding);
msgVoid0(__mps_batch_cmdBuf, __sel_commit);
msgVoid0(__mps_batch_cmdBuf, __sel_waitUntilCompleted);
MMsg0 msg0 = (MMsg0)objc_msgSend;
unsigned long i = 0UL;
while (i < __mps_pending_count) {
unsigned char* outPtr = __mps_contents(__mps_pending_gpuBuf[i]);
if (outPtr != null) {
memcpy(__mps_pending_cpuOut[i], outPtr, __mps_pending_bytes[i]);
}
i = i + 1UL;
}
__mps_pending_count = 0UL;
// Every buffer used in this batch is now provably done (the GPU
// finished reading/writing it before waitUntilCompleted returned)
// -- safe to return the whole set to the pool at once.
i = 0UL;
while (i < __mps_poolbuf_count) {
__mps_buf_put(__mps_poolbuf_buf[i], __mps_poolbuf_bytes[i]);
i = i + 1UL;
}
__mps_poolbuf_count = 0UL;
// Deactivate before running finalizers: a finalize callback must
// never itself try to enqueue more batched work against a buffer
// that's already been committed.
__mps_batch_active = 0;
__mps_batch_cmdBuf = null;
__mps_batch_encoder = null;
i = 0UL;
while (i < __mps_finalize_count) {
__mps_finalize_fn[i](__mps_finalize_ctx[i]);
i = i + 1UL;
}
__mps_finalize_count = 0UL;
ObjcPool_* pool = __mps_batch_pool;
__mps_batch_pool = null;
if (pool != null) objc_autoreleasePoolPop(pool);
}
}
void mps_batch_submit() {
if (!__mps_batch_active) return;
unsafe {
MMsgVoid0 msgVoid0 = (MMsgVoid0)objc_msgSend;
msgVoid0(__mps_batch_encoder, __sel_endEncoding);
msgVoid0(__mps_batch_cmdBuf, __sel_commit);
// the queue runs its command buffers in order, so mps_batch_end()'s
// wait on the last one covers this one too
MMsg0 msg0 = (MMsg0)objc_msgSend;
__mps_batch_cmdBuf = msg0(__mps_queue, __sel_commandBuffer);
__mps_batch_encoder = msg0(__mps_batch_cmdBuf, __sel_computeCommandEncoder);
}
}
// mps_sub_f32: own dedicated dispatch (not routed through
// __mps_run_binary_kernel like mul/div/pow) -- same reasoning as
// mps_relu_f32 below (see its comment): batch-aware, so tensor_sub_gpu can
// encode it onto a shared batch's command buffer for the fused
// forward+loss+backward training-step path instead of forcing its own
// commit+wait mid-batch.
int mps_sub_f32_n_(const &float a, const &float b, &float out, unsigned long n) {
unsafe {
if (!__mps_ensure_device_queue()) return 0;
MtlObject* device = __mps_device;
MtlObject* queue = __mps_queue;
MMsg0 msg0 = (MMsg0)objc_msgSend;
if (__mps_sub_pipeline == null) {
MtlObject* pipeline = __mps_make_pipeline("sub_kernel");
if (pipeline == null) return 0;
__mps_sub_pipeline = pipeline;
}
MtlObject* pipeline = __mps_sub_pipeline;
unsigned long bytes = n * sizeof(float);
int batched = __mps_batch_active;
MtlObject* bufA = batched ? __mps_buf_get_with_bytes_batched((const unsigned char*)a, bytes)
: __mps_buf_get_with_bytes((const unsigned char*)a, bytes);
MtlObject* bufB = batched ? __mps_buf_get_with_bytes_batched((const unsigned char*)b, bytes)
: __mps_buf_get_with_bytes((const unsigned char*)b, bytes);
MtlObject* bufOut = batched ? __mps_buf_get_batched(bytes) : __mps_buf_get(bytes);
if (bufA == null || bufB == null || bufOut == null) return 0;
MtlObject* cmdBuf = null;
MtlObject* encoder;
if (batched) {
encoder = __mps_batch_encoder;
} else {
cmdBuf = msg0(queue, __sel_commandBuffer);
encoder = msg0(cmdBuf, __sel_computeCommandEncoder);
}
MMsgVoid1 msgVoid1 = (MMsgVoid1)objc_msgSend;
msgVoid1(encoder, __sel_setComputePipelineState, pipeline);
MMsgSetBuffer setBufFn = (MMsgSetBuffer)objc_msgSend;
setBufFn(encoder, __sel_setBuffer, bufA, 0UL, 0UL);
setBufFn(encoder, __sel_setBuffer, bufB, 0UL, 1UL);
setBufFn(encoder, __sel_setBuffer, bufOut, 0UL, 2UL);
struct MTLSize gridSize;
gridSize.width = n; gridSize.height = 1UL; gridSize.depth = 1UL;
unsigned long tgWidth = n < 64UL ? n : 64UL;
if (tgWidth == 0UL) tgWidth = 1UL;
struct MTLSize tgSize;
tgSize.width = tgWidth; tgSize.height = 1UL; tgSize.depth = 1UL;
MMsgDispatch dispatchFn = (MMsgDispatch)objc_msgSend;
dispatchFn(encoder, __sel_dispatchThreads, gridSize, tgSize);
if (batched) {
if (__mps_pending_count < (unsigned long)MPS_BATCH_MAX_PENDING) {
__mps_pending_gpuBuf[__mps_pending_count] = bufOut;
__mps_pending_cpuOut[__mps_pending_count] = (unsigned char*)out;
__mps_pending_bytes[__mps_pending_count] = bytes;
__mps_pending_dtype[__mps_pending_count] = 0;
__mps_pending_count = __mps_pending_count + 1UL;
}
return 1;
}
MMsgVoid0 msgVoid0 = (MMsgVoid0)objc_msgSend;
msgVoid0(encoder, __sel_endEncoding);
msgVoid0(cmdBuf, __sel_commit);
msgVoid0(cmdBuf, __sel_waitUntilCompleted);
unsigned char* outPtr = __mps_contents(bufOut);
if (outPtr == null) return 0;
memcpy((unsigned char*)out, outPtr, bytes);
__mps_buf_put(bufA, bytes);
__mps_buf_put(bufB, bytes);
__mps_buf_put(bufOut, bytes);
return 1;
}
}
int mps_sub_f32([]const float a_s, []const float b_s, []float out_s) {
unsigned long n = out_s.len;
if (a_s.len < n || b_s.len < n) { return 0; }
float a_d_[1];
const &float a = &a_d_[0];
if (a_s.len > 0UL) { unsafe { a = (const &float)a_s.ptr; } }
float b_d_[1];
const &float b = &b_d_[0];
if (b_s.len > 0UL) { unsafe { b = (const &float)b_s.ptr; } }
float out_d_[1];
&float out = &out_d_[0];
if (out_s.len > 0UL) { unsafe { out = (&float)out_s.ptr; } }
return mps_sub_f32_n_(a, b, out, n);
}
// mps_sum_f32: batch-aware. Under an open batch, the CPU-side "combine the
// partial-sums array" step (the while-loop below) can't run inline the way
// the non-batched path does -- bufOut isn't readable until mps_batch_end()'s
// single wait, same reasoning as a gradient-accumulation context's
// __mps_grad_accum_finalize. Same fix: register a finalize callback instead,
// pointed at a small scratch copy of the partials (populated by
// mps_batch_end()'s own generic pending-readback, same as any other batched
// op's output) rather than reading bufOut's Metal contents directly.
struct __MpsSumFinalizeCtx {
float* partials;
float* out;
unsigned long numGroups;
};
static void __mps_sum_finalize(?&__MpsSumFinalizeCtx ctxPtr) {
unsafe {
struct __MpsSumFinalizeCtx* ctx = (struct __MpsSumFinalizeCtx*)ctxPtr;
float acc = 0.0f;
unsigned long g = 0UL;
while (g < ctx->numGroups) { acc = acc + ctx->partials[g]; g = g + 1UL; }
*(ctx->out) = acc;
dealloc(ctx->partials);
dealloc(ctx);
}
}
int mps_sum_f32_n_(const &float a, &float out, unsigned long n) {
unsafe {
if (!__mps_ensure_device_queue()) return 0;
MtlObject* device = __mps_device;
MtlObject* queue = __mps_queue;
MMsg0 msg0 = (MMsg0)objc_msgSend;
if (__mps_sum_pipeline == null) {
MtlObject* pipeline = __mps_make_pipeline("sum_kernel");
if (pipeline == null) return 0;
__mps_sum_pipeline = pipeline;
}
MtlObject* pipeline = __mps_sum_pipeline;
unsigned long bytes = n * sizeof(float);
unsigned long numGroups = (n + SUM_TG_SIZE - 1UL) / SUM_TG_SIZE;
if (numGroups > SUM_MAX_GROUPS) numGroups = SUM_MAX_GROUPS;
if (numGroups == 0UL) numGroups = 1UL;
unsigned long partialBytes = numGroups * sizeof(float);
int batched = __mps_batch_active;
MtlObject* bufA = batched ? __mps_buf_get_with_bytes_batched((const unsigned char*)a, bytes)
: __mps_buf_get_with_bytes((const unsigned char*)a, bytes);
MtlObject* bufOut = batched ? __mps_buf_get_batched(partialBytes) : __mps_buf_get(partialBytes);
if (bufA == null || bufOut == null) return 0;
MtlObject* cmdBuf = null;
MtlObject* encoder;
if (batched) {
encoder = __mps_batch_encoder;
} else {
cmdBuf = msg0(queue, __sel_commandBuffer);
encoder = msg0(cmdBuf, __sel_computeCommandEncoder);
}
MMsgVoid1 msgVoid1 = (MMsgVoid1)objc_msgSend;
msgVoid1(encoder, __sel_setComputePipelineState, pipeline);
MMsgSetBuffer setBufFn = (MMsgSetBuffer)objc_msgSend;
setBufFn(encoder, __sel_setBuffer, bufA, 0UL, 0UL);
setBufFn(encoder, __sel_setBuffer, bufOut, 0UL, 1UL);
unsigned int nu = (unsigned int)n;
MMsgSetBytes setBytesFn = (MMsgSetBytes)objc_msgSend;
setBytesFn(encoder, __sel_setBytes, (const unsigned char*)&nu, 4UL, 2UL);
struct MTLSize gridSize;
gridSize.width = numGroups; gridSize.height = 1UL; gridSize.depth = 1UL;
struct MTLSize tgSize;
tgSize.width = SUM_TG_SIZE; tgSize.height = 1UL; tgSize.depth = 1UL;
MMsgDispatch dispatchFn = (MMsgDispatch)objc_msgSend;
dispatchFn(encoder, __sel_dispatchThreadgroups, gridSize, tgSize);
if (batched) {
float* partials = (float*)alloc_raw_(sizeof(float) * numGroups);
if (__mps_pending_count < (unsigned long)MPS_BATCH_MAX_PENDING) {
__mps_pending_gpuBuf[__mps_pending_count] = bufOut;
__mps_pending_cpuOut[__mps_pending_count] = (unsigned char*)partials;
__mps_pending_bytes[__mps_pending_count] = partialBytes;
__mps_pending_dtype[__mps_pending_count] = 0;
__mps_pending_count = __mps_pending_count + 1UL;
}
struct __MpsSumFinalizeCtx* ctx = (struct __MpsSumFinalizeCtx*)alloc_raw_(sizeof(struct __MpsSumFinalizeCtx));
ctx->partials = partials; ctx->out = out; ctx->numGroups = numGroups;
mps_batch_register_finalize(__mps_sum_finalize, (?&__MpsSumFinalizeCtx)ctx);
return 1;
}
MMsgVoid0 msgVoid0 = (MMsgVoid0)objc_msgSend;
msgVoid0(encoder, __sel_endEncoding);
msgVoid0(cmdBuf, __sel_commit);
msgVoid0(cmdBuf, __sel_waitUntilCompleted);
unsigned char* outPtr = __mps_contents(bufOut);
if (outPtr == null) return 0;
float* partials = (float*)outPtr;
float acc = 0.0f;
unsigned long g = 0UL;
while (g < numGroups) { acc = acc + partials[g]; g = g + 1UL; }
*out = acc;
__mps_buf_put(bufA, bytes);
__mps_buf_put(bufOut, partialBytes);
return 1;
}
}
int mps_sum_f32([]const float a_s, []float out_s) {
unsigned long n = a_s.len;
if (out_s.len < 1UL) { return 0; }
float a_d_[1];
const &float a = &a_d_[0];
if (a_s.len > 0UL) { unsafe { a = (const &float)a_s.ptr; } }
float out_d_[1];
&float out = &out_d_[0];
if (out_s.len > 0UL) { unsafe { out = (&float)out_s.ptr; } }
return mps_sum_f32_n_(a, out, n);
}
int mps_sum_f32_chained(struct MpsBuf bufA, &float out, unsigned long n) {
if (bufA.handle == null) return 0;
if (!__mps_batch_active) return 0;
unsafe {
if (__mps_sum_pipeline == null) {
MtlObject* pipeline = __mps_make_pipeline("sum_kernel");
if (pipeline == null) return 0;
__mps_sum_pipeline = pipeline;
}
unsigned long numGroups = (n + SUM_TG_SIZE - 1UL) / SUM_TG_SIZE;
if (numGroups > SUM_MAX_GROUPS) numGroups = SUM_MAX_GROUPS;
if (numGroups == 0UL) numGroups = 1UL;
unsigned long partialBytes = numGroups * sizeof(float);
MtlObject* bufOut = __mps_buf_get_batched(partialBytes);
if (bufOut == null) return 0;
MMsgVoid1 msgVoid1 = (MMsgVoid1)objc_msgSend;
msgVoid1(__mps_batch_encoder, __sel_setComputePipelineState, __mps_sum_pipeline);
MMsgSetBuffer setBufFn = (MMsgSetBuffer)objc_msgSend;
setBufFn(__mps_batch_encoder, __sel_setBuffer, bufA.handle, 0UL, 0UL);
setBufFn(__mps_batch_encoder, __sel_setBuffer, bufOut, 0UL, 1UL);
unsigned int nu = (unsigned int)n;
MMsgSetBytes setBytesFn = (MMsgSetBytes)objc_msgSend;
setBytesFn(__mps_batch_encoder, __sel_setBytes, (const unsigned char*)&nu, 4UL, 2UL);
struct MTLSize gridSize;
gridSize.width = numGroups; gridSize.height = 1UL; gridSize.depth = 1UL;
struct MTLSize tgSize;
tgSize.width = SUM_TG_SIZE; tgSize.height = 1UL; tgSize.depth = 1UL;
MMsgDispatch dispatchFn = (MMsgDispatch)objc_msgSend;
dispatchFn(__mps_batch_encoder, __sel_dispatchThreadgroups, gridSize, tgSize);
float* partials = (float*)alloc_raw_(sizeof(float) * numGroups);
if (__mps_pending_count < (unsigned long)MPS_BATCH_MAX_PENDING) {
__mps_pending_gpuBuf[__mps_pending_count] = bufOut;
__mps_pending_cpuOut[__mps_pending_count] = (unsigned char*)partials;
__mps_pending_bytes[__mps_pending_count] = partialBytes;
__mps_pending_dtype[__mps_pending_count] = 0;
__mps_pending_count = __mps_pending_count + 1UL;
}
struct __MpsSumFinalizeCtx* ctx = (struct __MpsSumFinalizeCtx*)alloc_raw_(sizeof(struct __MpsSumFinalizeCtx));
ctx->partials = partials; ctx->out = out; ctx->numGroups = numGroups;
mps_batch_register_finalize(__mps_sum_finalize, (?&__MpsSumFinalizeCtx)ctx);
return 1;
}
}
// out[M,N] = a[M,K] . b[K,N] -- naive one-thread-per-output-element kernel
// (no threadgroup-memory tiling/blocking), dispatched over a 2D grid. Not
// competitive with a tuned GEMM (MPSMatrixMultiplication, or a tiled kernel
// that reuses each loaded value across multiple output elements via
// threadgroup memory) -- this is the GPU-dispatch-pipeline-works version,
// same spirit as mps_add_f32 originally being "the simplest possible real
// GPU-executed op" before anything fancier. Metal has no 'double' type at
// all (confirmed: 'double' is a hard compile error in MSL), so this is
// float32-only; std::ml::Tensor stores 'double' data, so the GPU-backed
// Tensor path (tensor_matmul_gpu in tensor.sc) converts in and out of
// float32 around this call, trading precision for the ability to run on
// the GPU at all.
static MtlObject* __mps_matmul_pipeline = null;
static MtlObject* __mps_matmul_smma_pipeline = null;
static MtlObject* __mps_matmul_multi_pipeline = null;
static MtlObject* __mps_relu_pipeline = null;
int mps_matmul_f32_n_(const &float a, const &float b, &float out,
unsigned long M, unsigned long K, unsigned long N) {
unsafe {
if (!__mps_ensure_device_queue()) return 0;
MtlObject* device = __mps_device;
MtlObject* queue = __mps_queue;
MMsg0 msg0 = (MMsg0)objc_msgSend;
// Prefer the simdgroup_matrix (hardware matrix unit) kernel
// whenever the shape allows it -- see gpu_mps_kernels.metal's
// comment on matmul_kernel_smma for why it can't handle
// non-8-aligned shapes at all. Every other detail below (buffer
// setup, batching, readback) is identical between the two paths;
// only the pipeline object and the threadgroup/grid sizing differ.
int useMulti = (M % 32UL == 0UL) && (K % 32UL == 0UL) && (N % 32UL == 0UL);
int useSmma = (M % 8UL == 0UL) && (K % 8UL == 0UL) && (N % 8UL == 0UL);
MtlObject* pipeline;
if (useMulti) {
if (__mps_matmul_multi_pipeline == null) {
MtlObject* p = __mps_make_pipeline("matmul_kernel_smma_multi");
if (p == null) { useMulti = 0; } else { __mps_matmul_multi_pipeline = p; }
}
if (useMulti) pipeline = __mps_matmul_multi_pipeline;
}
if (!useMulti && useSmma) {
if (__mps_matmul_smma_pipeline == null) {
MtlObject* p = __mps_make_pipeline("matmul_kernel_smma");
if (p == null) { useSmma = 0; } else { __mps_matmul_smma_pipeline = p; }
}
if (useSmma) pipeline = __mps_matmul_smma_pipeline;
}
if (!useMulti && !useSmma) {
if (__mps_matmul_pipeline == null) {
MtlObject* p = __mps_make_pipeline("matmul_kernel");
if (p == null) return 0;
__mps_matmul_pipeline = p;
}
pipeline = __mps_matmul_pipeline;
}
unsigned long bytesA = M * K * sizeof(float);
unsigned long bytesB = K * N * sizeof(float);
unsigned long bytesOut = M * N * sizeof(float);
// Batched: reuse the shared command buffer/encoder that
// mps_batch_begin() already opened instead of creating our own
// (and don't end/commit/wait/readback here — mps_batch_end()
// does all of that once, for every op encoded since begin()).
// Buffers obtained while batched are tracked for deferred
// pool-return there too, instead of being put back immediately.
int batched = __mps_batch_active;
MtlObject* bufA = batched ? __mps_buf_get_with_bytes_batched((const unsigned char*)a, bytesA)
: __mps_buf_get_with_bytes((const unsigned char*)a, bytesA);
MtlObject* bufB = batched ? __mps_buf_get_with_bytes_batched((const unsigned char*)b, bytesB)
: __mps_buf_get_with_bytes((const unsigned char*)b, bytesB);
MtlObject* bufOut = batched ? __mps_buf_get_batched(bytesOut) : __mps_buf_get(bytesOut);
if (bufA == null || bufB == null || bufOut == null) return 0;
unsigned int Mu = (unsigned int)M;
unsigned int Ku = (unsigned int)K;
unsigned int Nu = (unsigned int)N;
MtlObject* cmdBuf = null;
MtlObject* encoder;
if (batched) {
encoder = __mps_batch_encoder;
} else {
cmdBuf = msg0(queue, __sel_commandBuffer);
encoder = msg0(cmdBuf, __sel_computeCommandEncoder);
}
MMsgVoid1 msgVoid1 = (MMsgVoid1)objc_msgSend;
msgVoid1(encoder, __sel_setComputePipelineState, pipeline);
MMsgSetBuffer setBufFn = (MMsgSetBuffer)objc_msgSend;
setBufFn(encoder, __sel_setBuffer, bufA, 0UL, 0UL);
setBufFn(encoder, __sel_setBuffer, bufB, 0UL, 1UL);
setBufFn(encoder, __sel_setBuffer, bufOut, 0UL, 2UL);
// setBytes:length:atIndex: has the exact same (recv, sel, ptr,
// NSUInteger, NSUInteger) shape as setBuffer:offset:atIndex:, so
// MMsgSetBuffer's typedef works for it too -- only the selector
// and what the middle argument means (raw bytes vs a buffer
// object) differ.
MMsgSetBytes setBytesFn = (MMsgSetBytes)objc_msgSend;
setBytesFn(encoder, __sel_setBytes, (const unsigned char*)&Mu, 4UL, 3UL);
setBytesFn(encoder, __sel_setBytes, (const unsigned char*)&Ku, 4UL, 4UL);
setBytesFn(encoder, __sel_setBytes, (const unsigned char*)&Nu, 4UL, 5UL);
// Dispatch: one threadgroup per 32x32 output tile (128 threads, 4
// simdgroups) for the multi path; one threadgroup per 8x8 output
// tile (32 threads, one simdgroup) for the single-simdgroup smma
// path; TILE_SIZE(16) x TILE_SIZE(16) threadgroups for the tiled
// fallback, matching matmul_kernel's threadgroup-memory tiling
// exactly -- see mps_matmul_f32's header comment (above the
// function) for why the tiled dispatch can't shrink for small
// matrices.
struct MTLSize tgSize;
struct MTLSize numThreadgroups;
if (useMulti) {
tgSize.width = 128UL; tgSize.height = 1UL; tgSize.depth = 1UL;
numThreadgroups.width = N / 32UL; numThreadgroups.height = M / 32UL; numThreadgroups.depth = 1UL;
} else if (useSmma) {
tgSize.width = 32UL; tgSize.height = 1UL; tgSize.depth = 1UL;
numThreadgroups.width = N / 8UL; numThreadgroups.height = M / 8UL; numThreadgroups.depth = 1UL;
} else {
unsigned long tgW = 16UL;
unsigned long tgH = 16UL;
tgSize.width = tgW; tgSize.height = tgH; tgSize.depth = 1UL;
unsigned long groupsX = (N + tgW - 1UL) / tgW;
unsigned long groupsY = (M + tgH - 1UL) / tgH;
numThreadgroups.width = groupsX; numThreadgroups.height = groupsY; numThreadgroups.depth = 1UL;
}
MMsgDispatch dispatchFn = (MMsgDispatch)objc_msgSend;
dispatchFn(encoder, __sel_dispatchThreadgroups, numThreadgroups, tgSize);
if (batched) {
if (__mps_pending_count < (unsigned long)MPS_BATCH_MAX_PENDING) {
__mps_pending_gpuBuf[__mps_pending_count] = bufOut;
__mps_pending_cpuOut[__mps_pending_count] = (unsigned char*)out;
__mps_pending_bytes[__mps_pending_count] = bytesOut;
__mps_pending_dtype[__mps_pending_count] = 0;
__mps_pending_count = __mps_pending_count + 1UL;
}
return 1; // deferred success -- 'out' not valid until mps_batch_end()
}
MMsgVoid0 msgVoid0 = (MMsgVoid0)objc_msgSend;
msgVoid0(encoder, __sel_endEncoding);
msgVoid0(cmdBuf, __sel_commit);
msgVoid0(cmdBuf, __sel_waitUntilCompleted);
unsigned char* outPtr = __mps_contents(bufOut);
if (outPtr == null) return 0;
memcpy((unsigned char*)out, outPtr, bytesOut);
__mps_buf_put(bufA, bytesA);
__mps_buf_put(bufB, bytesB);
__mps_buf_put(bufOut, bytesOut);
return 1;
}
}
int mps_matmul_f32([]const float a_s, []const float b_s, []float out_s, unsigned long M, unsigned long K, unsigned long N) {
if (a_s.len < M * K || b_s.len < K * N || out_s.len < M * N) { return 0; }
float a_d_[1];
const &float a = &a_d_[0];
if (a_s.len > 0UL) { unsafe { a = (const &float)a_s.ptr; } }
float b_d_[1];
const &float b = &b_d_[0];
if (b_s.len > 0UL) { unsafe { b = (const &float)b_s.ptr; } }
float out_d_[1];
&float out = &out_d_[0];
if (out_s.len > 0UL) { unsafe { out = (&float)out_s.ptr; } }
return mps_matmul_f32_n_(a, b, out, M, K, N);
}
// ── matmul backward (see gpu_mps.h's own comment on these two) ───────────────
// Same dispatch shape as mps_matmul_f32 (2D grid, cached pipeline, pooled
// buffers) but against matmul_abt_kernel/matmul_atb_kernel instead of
// matmul_kernel -- not routed through mps_matmul_f32 itself since the
// output shape and which operand reads transposed differ between the two,
// and neither needs batching (Tensor::backward() always runs after any
// forward-pass batch has already ended).
static MtlObject* __mps_matmul_abt_pipeline = null;
static MtlObject* __mps_matmul_abt_smma_pipeline = null;
static MtlObject* __mps_matmul_abt_multi_pipeline = null;
int mps_matmul_abt_f32(const &float a, const &float b, &float out,
unsigned long M, unsigned long N, unsigned long K) {
unsafe {
if (!__mps_ensure_device_queue()) return 0;
MtlObject* device = __mps_device;
MtlObject* queue = __mps_queue;
MMsg0 msg0 = (MMsg0)objc_msgSend;
// See mps_matmul_f32's comment on the smma/tiled choice.
int useMulti = (M % 32UL == 0UL) && (N % 32UL == 0UL) && (K % 32UL == 0UL);
int useSmma = (M % 8UL == 0UL) && (N % 8UL == 0UL) && (K % 8UL == 0UL);
MtlObject* pipeline;
if (useMulti) {
if (__mps_matmul_abt_multi_pipeline == null) {
MtlObject* p = __mps_make_pipeline("matmul_abt_kernel_smma_multi");
if (p == null) { useMulti = 0; } else { __mps_matmul_abt_multi_pipeline = p; }
}
if (useMulti) pipeline = __mps_matmul_abt_multi_pipeline;
}
if (!useMulti && useSmma) {
if (__mps_matmul_abt_smma_pipeline == null) {
MtlObject* p = __mps_make_pipeline("matmul_abt_kernel_smma");
if (p == null) { useSmma = 0; } else { __mps_matmul_abt_smma_pipeline = p; }
}
if (useSmma) pipeline = __mps_matmul_abt_smma_pipeline;
}
if (!useMulti && !useSmma) {
if (__mps_matmul_abt_pipeline == null) {
MtlObject* p = __mps_make_pipeline("matmul_abt_kernel");
if (p == null) return 0;
__mps_matmul_abt_pipeline = p;
}
pipeline = __mps_matmul_abt_pipeline;
}
unsigned long bytesA = M * N * sizeof(float);
unsigned long bytesB = K * N * sizeof(float);
unsigned long bytesOut = M * K * sizeof(float);
// Batched: same "reuse the shared encoder, defer commit/wait/
// readback to mps_batch_end()" pattern as mps_matmul_f32 -- see
// its comment.
int batched = __mps_batch_active;
MtlObject* bufA = batched ? __mps_buf_get_with_bytes_batched((const unsigned char*)a, bytesA)
: __mps_buf_get_with_bytes((const unsigned char*)a, bytesA);
MtlObject* bufB = batched ? __mps_buf_get_with_bytes_batched((const unsigned char*)b, bytesB)
: __mps_buf_get_with_bytes((const unsigned char*)b, bytesB);
MtlObject* bufOut = batched ? __mps_buf_get_batched(bytesOut) : __mps_buf_get(bytesOut);
if (bufA == null || bufB == null || bufOut == null) return 0;
unsigned int Mu = (unsigned int)M;
unsigned int Nu = (unsigned int)N;
unsigned int Ku = (unsigned int)K;
MtlObject* cmdBuf = null;
MtlObject* encoder;
if (batched) {
encoder = __mps_batch_encoder;
} else {
cmdBuf = msg0(queue, __sel_commandBuffer);
encoder = msg0(cmdBuf, __sel_computeCommandEncoder);
}
MMsgVoid1 msgVoid1 = (MMsgVoid1)objc_msgSend;
msgVoid1(encoder, __sel_setComputePipelineState, pipeline);
MMsgSetBuffer setBufFn = (MMsgSetBuffer)objc_msgSend;
setBufFn(encoder, __sel_setBuffer, bufA, 0UL, 0UL);
setBufFn(encoder, __sel_setBuffer, bufB, 0UL, 1UL);
setBufFn(encoder, __sel_setBuffer, bufOut, 0UL, 2UL);
MMsgSetBytes setBytesFn = (MMsgSetBytes)objc_msgSend;
setBytesFn(encoder, __sel_setBytes, (const unsigned char*)&Mu, 4UL, 3UL);
setBytesFn(encoder, __sel_setBytes, (const unsigned char*)&Nu, 4UL, 4UL);
setBytesFn(encoder, __sel_setBytes, (const unsigned char*)&Ku, 4UL, 5UL);
// See mps_matmul_f32's comment on the smma/tiled dispatch shapes.
struct MTLSize tgSize;
struct MTLSize numThreadgroups;
if (useMulti) {
tgSize.width = 128UL; tgSize.height = 1UL; tgSize.depth = 1UL;
numThreadgroups.width = K / 32UL; numThreadgroups.height = M / 32UL; numThreadgroups.depth = 1UL;
} else if (useSmma) {
tgSize.width = 32UL; tgSize.height = 1UL; tgSize.depth = 1UL;
numThreadgroups.width = K / 8UL; numThreadgroups.height = M / 8UL; numThreadgroups.depth = 1UL;
} else {
unsigned long tgW = 16UL;
unsigned long tgH = 16UL;
tgSize.width = tgW; tgSize.height = tgH; tgSize.depth = 1UL;
unsigned long groupsX = (K + tgW - 1UL) / tgW;
unsigned long groupsY = (M + tgH - 1UL) / tgH;
numThreadgroups.width = groupsX; numThreadgroups.height = groupsY; numThreadgroups.depth = 1UL;
}
MMsgDispatch dispatchFn = (MMsgDispatch)objc_msgSend;
dispatchFn(encoder, __sel_dispatchThreadgroups, numThreadgroups, tgSize);
if (batched) {
if (__mps_pending_count < (unsigned long)MPS_BATCH_MAX_PENDING) {
__mps_pending_gpuBuf[__mps_pending_count] = bufOut;
__mps_pending_cpuOut[__mps_pending_count] = (unsigned char*)out;
__mps_pending_bytes[__mps_pending_count] = bytesOut;
__mps_pending_dtype[__mps_pending_count] = 0;
__mps_pending_count = __mps_pending_count + 1UL;
}
return 1;
}
MMsgVoid0 msgVoid0 = (MMsgVoid0)objc_msgSend;
msgVoid0(encoder, __sel_endEncoding);
msgVoid0(cmdBuf, __sel_commit);
msgVoid0(cmdBuf, __sel_waitUntilCompleted);
unsigned char* outPtr = __mps_contents(bufOut);
if (outPtr == null) return 0;
memcpy((unsigned char*)out, outPtr, bytesOut);
__mps_buf_put(bufA, bytesA);
__mps_buf_put(bufB, bytesB);
__mps_buf_put(bufOut, bytesOut);
return 1;
}
}
static MtlObject* __mps_matmul_atb_pipeline = null;
static MtlObject* __mps_matmul_atb_smma_pipeline = null;
static MtlObject* __mps_matmul_atb_multi_pipeline = null;
int mps_matmul_atb_f32(const &float a, const &float b, &float out,
unsigned long M, unsigned long K, unsigned long N) {
unsafe {
if (!__mps_ensure_device_queue()) return 0;
MtlObject* device = __mps_device;
MtlObject* queue = __mps_queue;
MMsg0 msg0 = (MMsg0)objc_msgSend;
// See mps_matmul_f32's comment on the smma/tiled choice.
int useMulti = (M % 32UL == 0UL) && (K % 32UL == 0UL) && (N % 32UL == 0UL);
int useSmma = (M % 8UL == 0UL) && (K % 8UL == 0UL) && (N % 8UL == 0UL);
MtlObject* pipeline;
if (useMulti) {
if (__mps_matmul_atb_multi_pipeline == null) {
MtlObject* p = __mps_make_pipeline("matmul_atb_kernel_smma_multi");
if (p == null) { useMulti = 0; } else { __mps_matmul_atb_multi_pipeline = p; }
}
if (useMulti) pipeline = __mps_matmul_atb_multi_pipeline;
}
if (!useMulti && useSmma) {
if (__mps_matmul_atb_smma_pipeline == null) {
MtlObject* p = __mps_make_pipeline("matmul_atb_kernel_smma");
if (p == null) { useSmma = 0; } else { __mps_matmul_atb_smma_pipeline = p; }
}
if (useSmma) pipeline = __mps_matmul_atb_smma_pipeline;
}
if (!useMulti && !useSmma) {
if (__mps_matmul_atb_pipeline == null) {
MtlObject* p = __mps_make_pipeline("matmul_atb_kernel");
if (p == null) return 0;
__mps_matmul_atb_pipeline = p;
}
pipeline = __mps_matmul_atb_pipeline;
}
unsigned long bytesA = M * K * sizeof(float);
unsigned long bytesB = M * N * sizeof(float);
unsigned long bytesOut = K * N * sizeof(float);
int batched = __mps_batch_active;
MtlObject* bufA = batched ? __mps_buf_get_with_bytes_batched((const unsigned char*)a, bytesA)
: __mps_buf_get_with_bytes((const unsigned char*)a, bytesA);
MtlObject* bufB = batched ? __mps_buf_get_with_bytes_batched((const unsigned char*)b, bytesB)
: __mps_buf_get_with_bytes((const unsigned char*)b, bytesB);
MtlObject* bufOut = batched ? __mps_buf_get_batched(bytesOut) : __mps_buf_get(bytesOut);
if (bufA == null || bufB == null || bufOut == null) return 0;
unsigned int Mu = (unsigned int)M;
unsigned int Ku = (unsigned int)K;
unsigned int Nu = (unsigned int)N;
MtlObject* cmdBuf = null;
MtlObject* encoder;
if (batched) {
encoder = __mps_batch_encoder;
} else {
cmdBuf = msg0(queue, __sel_commandBuffer);
encoder = msg0(cmdBuf, __sel_computeCommandEncoder);
}
MMsgVoid1 msgVoid1 = (MMsgVoid1)objc_msgSend;
msgVoid1(encoder, __sel_setComputePipelineState, pipeline);
MMsgSetBuffer setBufFn = (MMsgSetBuffer)objc_msgSend;
setBufFn(encoder, __sel_setBuffer, bufA, 0UL, 0UL);
setBufFn(encoder, __sel_setBuffer, bufB, 0UL, 1UL);
setBufFn(encoder, __sel_setBuffer, bufOut, 0UL, 2UL);
MMsgSetBytes setBytesFn = (MMsgSetBytes)objc_msgSend;
setBytesFn(encoder, __sel_setBytes, (const unsigned char*)&Mu, 4UL, 3UL);
setBytesFn(encoder, __sel_setBytes, (const unsigned char*)&Ku, 4UL, 4UL);
setBytesFn(encoder, __sel_setBytes, (const unsigned char*)&Nu, 4UL, 5UL);
// See mps_matmul_f32's comment on the smma/tiled dispatch shapes.
struct MTLSize tgSize;
struct MTLSize numThreadgroups;
if (useMulti) {
tgSize.width = 128UL; tgSize.height = 1UL; tgSize.depth = 1UL;
numThreadgroups.width = N / 32UL; numThreadgroups.height = K / 32UL; numThreadgroups.depth = 1UL;
} else if (useSmma) {
tgSize.width = 32UL; tgSize.height = 1UL; tgSize.depth = 1UL;
numThreadgroups.width = N / 8UL; numThreadgroups.height = K / 8UL; numThreadgroups.depth = 1UL;
} else {
unsigned long tgW = 16UL;
unsigned long tgH = 16UL;
tgSize.width = tgW; tgSize.height = tgH; tgSize.depth = 1UL;
unsigned long groupsX = (N + tgW - 1UL) / tgW;
unsigned long groupsY = (K + tgH - 1UL) / tgH;
numThreadgroups.width = groupsX; numThreadgroups.height = groupsY; numThreadgroups.depth = 1UL;
}
MMsgDispatch dispatchFn = (MMsgDispatch)objc_msgSend;
dispatchFn(encoder, __sel_dispatchThreadgroups, numThreadgroups, tgSize);
if (batched) {
if (__mps_pending_count < (unsigned long)MPS_BATCH_MAX_PENDING) {
__mps_pending_gpuBuf[__mps_pending_count] = bufOut;
__mps_pending_cpuOut[__mps_pending_count] = (unsigned char*)out;
__mps_pending_bytes[__mps_pending_count] = bytesOut;
__mps_pending_dtype[__mps_pending_count] = 0;
__mps_pending_count = __mps_pending_count + 1UL;
}
return 1;
}
MMsgVoid0 msgVoid0 = (MMsgVoid0)objc_msgSend;
msgVoid0(encoder, __sel_endEncoding);
msgVoid0(cmdBuf, __sel_commit);
msgVoid0(cmdBuf, __sel_waitUntilCompleted);
unsigned char* outPtr = __mps_contents(bufOut);
if (outPtr == null) return 0;
memcpy((unsigned char*)out, outPtr, bytesOut);
__mps_buf_put(bufA, bytesA);
__mps_buf_put(bufB, bytesB);
__mps_buf_put(bufOut, bytesOut);
return 1;
}
}
// mps_relu_backward_f32: own dedicated dispatch (not routed through
// __mps_run_binary_kernel, same reasoning as mps_relu_f32 above it in the
// file -- batching needs to skip the always-commit-and-wait path that
// helper takes). Same 1D-grid dispatch shape as mps_relu_f32/mps_add_f32.
int mps_relu_backward_f32_n_(const &float a, const &float selfGrad, &float out, unsigned long n) {
unsafe {
if (!__mps_ensure_device_queue()) return 0;
MtlObject* device = __mps_device;
MtlObject* queue = __mps_queue;
MMsg0 msg0 = (MMsg0)objc_msgSend;
if (__mps_relu_backward_pipeline == null) {
MtlObject* pipeline = __mps_make_pipeline("relu_backward_kernel");
if (pipeline == null) return 0;
__mps_relu_backward_pipeline = pipeline;
}
MtlObject* pipeline = __mps_relu_backward_pipeline;
unsigned long bytes = n * sizeof(float);
int batched = __mps_batch_active;
MtlObject* bufA = batched ? __mps_buf_get_with_bytes_batched((const unsigned char*)a, bytes)
: __mps_buf_get_with_bytes((const unsigned char*)a, bytes);
MtlObject* bufSelfGrad = batched ? __mps_buf_get_with_bytes_batched((const unsigned char*)selfGrad, bytes)
: __mps_buf_get_with_bytes((const unsigned char*)selfGrad, bytes);
MtlObject* bufOut = batched ? __mps_buf_get_batched(bytes) : __mps_buf_get(bytes);
if (bufA == null || bufSelfGrad == null || bufOut == null) return 0;
MtlObject* cmdBuf = null;
MtlObject* encoder;
if (batched) {
encoder = __mps_batch_encoder;
} else {
cmdBuf = msg0(queue, __sel_commandBuffer);
encoder = msg0(cmdBuf, __sel_computeCommandEncoder);
}
MMsgVoid1 msgVoid1 = (MMsgVoid1)objc_msgSend;
msgVoid1(encoder, __sel_setComputePipelineState, pipeline);
MMsgSetBuffer setBufFn = (MMsgSetBuffer)objc_msgSend;
setBufFn(encoder, __sel_setBuffer, bufA, 0UL, 0UL);
setBufFn(encoder, __sel_setBuffer, bufSelfGrad, 0UL, 1UL);
setBufFn(encoder, __sel_setBuffer, bufOut, 0UL, 2UL);
struct MTLSize gridSize;
gridSize.width = n; gridSize.height = 1UL; gridSize.depth = 1UL;
unsigned long tgWidth = n < 64UL ? n : 64UL;
if (tgWidth == 0UL) tgWidth = 1UL;
struct MTLSize tgSize;
tgSize.width = tgWidth; tgSize.height = 1UL; tgSize.depth = 1UL;
MMsgDispatch dispatchFn = (MMsgDispatch)objc_msgSend;
dispatchFn(encoder, __sel_dispatchThreads, gridSize, tgSize);
if (batched) {
if (__mps_pending_count < (unsigned long)MPS_BATCH_MAX_PENDING) {
__mps_pending_gpuBuf[__mps_pending_count] = bufOut;
__mps_pending_cpuOut[__mps_pending_count] = (unsigned char*)out;
__mps_pending_bytes[__mps_pending_count] = bytes;
__mps_pending_dtype[__mps_pending_count] = 0;
__mps_pending_count = __mps_pending_count + 1UL;
}
return 1;
}
MMsgVoid0 msgVoid0 = (MMsgVoid0)objc_msgSend;
msgVoid0(encoder, __sel_endEncoding);
msgVoid0(cmdBuf, __sel_commit);
msgVoid0(cmdBuf, __sel_waitUntilCompleted);
unsigned char* outPtr = __mps_contents(bufOut);
if (outPtr == null) return 0;
memcpy((unsigned char*)out, outPtr, bytes);
__mps_buf_put(bufA, bytes);
__mps_buf_put(bufSelfGrad, bytes);
__mps_buf_put(bufOut, bytes);
return 1;
}
}
int mps_relu_backward_f32([]const float a_s, []const float selfGrad_s, []float out_s) {
unsigned long n = out_s.len;
if (a_s.len < n || selfGrad_s.len < n) { return 0; }
float a_d_[1];
const &float a = &a_d_[0];
if (a_s.len > 0UL) { unsafe { a = (const &float)a_s.ptr; } }
float selfGrad_d_[1];
const &float selfGrad = &selfGrad_d_[0];
if (selfGrad_s.len > 0UL) { unsafe { selfGrad = (const &float)selfGrad_s.ptr; } }
float out_d_[1];
&float out = &out_d_[0];
if (out_s.len > 0UL) { unsafe { out = (&float)out_s.ptr; } }
return mps_relu_backward_f32_n_(a, selfGrad, out, n);
}
// out[i] = max(a[i], 0) -- cached device/queue/pipeline and batching
// support, same pattern as mps_matmul_f32 above (not routed through
// __mps_run_unary_kernel like log/exp/sqrt: that helper always creates
// its own command buffer and blocks on waitUntilCompleted, which is
// exactly what a batched relu — encoded onto a matmul's shared command
// buffer, e.g. for a matmul -> relu -> matmul forward pass with no CPU
// step in between — needs to skip).
int mps_relu_f32_n_(const &float a, &float out, unsigned long n) {
unsafe {
if (!__mps_ensure_device_queue()) return 0;
MtlObject* device = __mps_device;
MtlObject* queue = __mps_queue;
MMsg0 msg0 = (MMsg0)objc_msgSend;
if (__mps_relu_pipeline == null) {
MtlObject* pipeline = __mps_make_pipeline("relu_kernel");
if (pipeline == null) return 0;
__mps_relu_pipeline = pipeline;
}
MtlObject* pipeline = __mps_relu_pipeline;
unsigned long bytes = n * sizeof(float);
int batched = __mps_batch_active;
MtlObject* bufA = batched ? __mps_buf_get_with_bytes_batched((const unsigned char*)a, bytes)
: __mps_buf_get_with_bytes((const unsigned char*)a, bytes);
MtlObject* bufOut = batched ? __mps_buf_get_batched(bytes) : __mps_buf_get(bytes);
if (bufA == null || bufOut == null) return 0;
MtlObject* cmdBuf = null;
MtlObject* encoder;
if (batched) {
encoder = __mps_batch_encoder;
} else {
cmdBuf = msg0(queue, __sel_commandBuffer);
encoder = msg0(cmdBuf, __sel_computeCommandEncoder);
}
MMsgVoid1 msgVoid1 = (MMsgVoid1)objc_msgSend;
msgVoid1(encoder, __sel_setComputePipelineState, pipeline);
MMsgSetBuffer setBufFn = (MMsgSetBuffer)objc_msgSend;
setBufFn(encoder, __sel_setBuffer, bufA, 0UL, 0UL);
setBufFn(encoder, __sel_setBuffer, bufOut, 0UL, 1UL);
struct MTLSize gridSize;
gridSize.width = n; gridSize.height = 1UL; gridSize.depth = 1UL;
unsigned long tgWidth = n < 64UL ? n : 64UL;
if (tgWidth == 0UL) tgWidth = 1UL;
struct MTLSize tgSize;
tgSize.width = tgWidth; tgSize.height = 1UL; tgSize.depth = 1UL;
MMsgDispatch dispatchFn = (MMsgDispatch)objc_msgSend;
dispatchFn(encoder, __sel_dispatchThreads, gridSize, tgSize);
if (batched) {
if (__mps_pending_count < (unsigned long)MPS_BATCH_MAX_PENDING) {
__mps_pending_gpuBuf[__mps_pending_count] = bufOut;
__mps_pending_cpuOut[__mps_pending_count] = (unsigned char*)out;
__mps_pending_bytes[__mps_pending_count] = bytes;
__mps_pending_dtype[__mps_pending_count] = 0;
__mps_pending_count = __mps_pending_count + 1UL;
}
return 1;
}
MMsgVoid0 msgVoid0 = (MMsgVoid0)objc_msgSend;
msgVoid0(encoder, __sel_endEncoding);
msgVoid0(cmdBuf, __sel_commit);
msgVoid0(cmdBuf, __sel_waitUntilCompleted);
unsigned char* outPtr = __mps_contents(bufOut);
if (outPtr == null) return 0;
memcpy((unsigned char*)out, outPtr, bytes);
__mps_buf_put(bufA, bytes);
__mps_buf_put(bufOut, bytes);
return 1;
}
}
int mps_relu_f32([]const float a_s, []float out_s) {
unsigned long n = out_s.len;
if (a_s.len < n) { return 0; }
float a_d_[1];
const &float a = &a_d_[0];
if (a_s.len > 0UL) { unsafe { a = (const &float)a_s.ptr; } }
float out_d_[1];
&float out = &out_d_[0];
if (out_s.len > 0UL) { unsafe { out = (&float)out_s.ptr; } }
return mps_relu_f32_n_(a, out, n);
}
// ── GPU-buffer-chained variants (see gpu_mps.h's own comment on these) ───────
// Only meaningful while a batch is active: 'bufA' is used directly as
// buffer(0) instead of uploading from a CPU float array, so the dispatch
// this encodes reads whatever a PREVIOUS dispatch on the same command
// buffer already wrote there -- correct without any CPU round trip
// specifically because Metal executes dispatches encoded on the same
// encoder in encoding order. 'bufA' itself is never pool-tracked here: it
// was already tracked by whichever op produced it as its own bufOut.
int mps_matmul_f32_chained(struct MpsBuf bufA, const &float b, &float out,
unsigned long M, unsigned long K, unsigned long N) {
if (bufA.handle == null) return 0;
if (!__mps_batch_active) return 0;
unsafe {
// See mps_matmul_f32's comment on the smma/tiled choice.
int useMulti = (M % 32UL == 0UL) && (K % 32UL == 0UL) && (N % 32UL == 0UL);
int useSmma = (M % 8UL == 0UL) && (K % 8UL == 0UL) && (N % 8UL == 0UL);
MtlObject* pipeline;
if (useMulti) {
if (__mps_matmul_multi_pipeline == null) {
MtlObject* p = __mps_make_pipeline("matmul_kernel_smma_multi");
if (p == null) { useMulti = 0; } else { __mps_matmul_multi_pipeline = p; }
}
if (useMulti) pipeline = __mps_matmul_multi_pipeline;
}
if (!useMulti && useSmma) {
if (__mps_matmul_smma_pipeline == null) {
MtlObject* p = __mps_make_pipeline("matmul_kernel_smma");
if (p == null) { useSmma = 0; } else { __mps_matmul_smma_pipeline = p; }
}
if (useSmma) pipeline = __mps_matmul_smma_pipeline;
}
if (!useMulti && !useSmma) {
if (__mps_matmul_pipeline == null) {
MtlObject* p = __mps_make_pipeline("matmul_kernel");
if (p == null) return 0;
__mps_matmul_pipeline = p;
}
pipeline = __mps_matmul_pipeline;
}
unsigned long bytesB = K * N * sizeof(float);
unsigned long bytesOut = M * N * sizeof(float);
MtlObject* bufB = __mps_buf_get_with_bytes_batched((const unsigned char*)b, bytesB);
MtlObject* bufOut = __mps_buf_get_batched(bytesOut);
if (bufB == null || bufOut == null) return 0;
unsigned int Mu = (unsigned int)M;
unsigned int Ku = (unsigned int)K;
unsigned int Nu = (unsigned int)N;
MMsgVoid1 msgVoid1 = (MMsgVoid1)objc_msgSend;
msgVoid1(__mps_batch_encoder, __sel_setComputePipelineState, pipeline);
MMsgSetBuffer setBufFn = (MMsgSetBuffer)objc_msgSend;
setBufFn(__mps_batch_encoder, __sel_setBuffer, bufA.handle, 0UL, 0UL);
setBufFn(__mps_batch_encoder, __sel_setBuffer, bufB, 0UL, 1UL);
setBufFn(__mps_batch_encoder, __sel_setBuffer, bufOut, 0UL, 2UL);
MMsgSetBytes setBytesFn = (MMsgSetBytes)objc_msgSend;
setBytesFn(__mps_batch_encoder, __sel_setBytes, (const unsigned char*)&Mu, 4UL, 3UL);
setBytesFn(__mps_batch_encoder, __sel_setBytes, (const unsigned char*)&Ku, 4UL, 4UL);
setBytesFn(__mps_batch_encoder, __sel_setBytes, (const unsigned char*)&Nu, 4UL, 5UL);
struct MTLSize tgSize;
struct MTLSize numThreadgroups;
if (useMulti) {
tgSize.width = 128UL; tgSize.height = 1UL; tgSize.depth = 1UL;
numThreadgroups.width = N / 32UL; numThreadgroups.height = M / 32UL; numThreadgroups.depth = 1UL;
} else if (useSmma) {
tgSize.width = 32UL; tgSize.height = 1UL; tgSize.depth = 1UL;
numThreadgroups.width = N / 8UL; numThreadgroups.height = M / 8UL; numThreadgroups.depth = 1UL;
} else {
unsigned long tgW = 16UL;
unsigned long tgH = 16UL;
tgSize.width = tgW; tgSize.height = tgH; tgSize.depth = 1UL;
unsigned long groupsX = (N + tgW - 1UL) / tgW;
unsigned long groupsY = (M + tgH - 1UL) / tgH;
numThreadgroups.width = groupsX; numThreadgroups.height = groupsY; numThreadgroups.depth = 1UL;
}
MMsgDispatch dispatchFn = (MMsgDispatch)objc_msgSend;
dispatchFn(__mps_batch_encoder, __sel_dispatchThreadgroups, numThreadgroups, tgSize);
if (__mps_pending_count < (unsigned long)MPS_BATCH_MAX_PENDING) {
__mps_pending_gpuBuf[__mps_pending_count] = bufOut;
__mps_pending_cpuOut[__mps_pending_count] = (unsigned char*)out;
__mps_pending_bytes[__mps_pending_count] = bytesOut;
__mps_pending_dtype[__mps_pending_count] = 0;
__mps_pending_count = __mps_pending_count + 1UL;
}
return 1;
}
}
// out[M,N] = bufA[M,K] . bufB[K,N] -- same computation as mps_matmul_f32_
// chained, but 'bufB' is ALSO a pre-existing device buffer (typically from
// MpsBuf::upload below) instead of a plain CPU array to upload
// fresh. Needed for e.g. an inference loop that runs the same matmul
// against the same weight matrix 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 between calls. 'bufA' can be either
// another persistent buffer or a previous op's still-pending chained
// output (MpsBuf::batch_last_output()) -- this function doesn't care
// which, both are just "already a device buffer".
int mps_matmul_f32_persistent(struct MpsBuf bufA, struct MpsBuf bufB, &float out,
unsigned long M, unsigned long K, unsigned long N) {
if (bufA.handle == null || bufB.handle == null) return 0;
if (!__mps_batch_active) return 0;
unsafe {
int useMulti = (M % 32UL == 0UL) && (K % 32UL == 0UL) && (N % 32UL == 0UL);
int useSmma = (M % 8UL == 0UL) && (K % 8UL == 0UL) && (N % 8UL == 0UL);
MtlObject* pipeline;
if (useMulti) {
if (__mps_matmul_multi_pipeline == null) {
MtlObject* p = __mps_make_pipeline("matmul_kernel_smma_multi");
if (p == null) { useMulti = 0; } else { __mps_matmul_multi_pipeline = p; }
}
if (useMulti) pipeline = __mps_matmul_multi_pipeline;
}
if (!useMulti && useSmma) {
if (__mps_matmul_smma_pipeline == null) {
MtlObject* p = __mps_make_pipeline("matmul_kernel_smma");
if (p == null) { useSmma = 0; } else { __mps_matmul_smma_pipeline = p; }
}
if (useSmma) pipeline = __mps_matmul_smma_pipeline;
}
if (!useMulti && !useSmma) {
if (__mps_matmul_pipeline == null) {
MtlObject* p = __mps_make_pipeline("matmul_kernel");
if (p == null) return 0;
__mps_matmul_pipeline = p;
}
pipeline = __mps_matmul_pipeline;
}
unsigned long bytesOut = M * N * sizeof(float);
MtlObject* bufOut = __mps_buf_get_batched(bytesOut);
if (bufOut == null) return 0;
unsigned int Mu = (unsigned int)M;
unsigned int Ku = (unsigned int)K;
unsigned int Nu = (unsigned int)N;
MMsgVoid1 msgVoid1 = (MMsgVoid1)objc_msgSend;
msgVoid1(__mps_batch_encoder, __sel_setComputePipelineState, pipeline);
MMsgSetBuffer setBufFn = (MMsgSetBuffer)objc_msgSend;
setBufFn(__mps_batch_encoder, __sel_setBuffer, bufA.handle, 0UL, 0UL);
setBufFn(__mps_batch_encoder, __sel_setBuffer, bufB.handle, 0UL, 1UL);
setBufFn(__mps_batch_encoder, __sel_setBuffer, bufOut, 0UL, 2UL);
MMsgSetBytes setBytesFn = (MMsgSetBytes)objc_msgSend;
setBytesFn(__mps_batch_encoder, __sel_setBytes, (const unsigned char*)&Mu, 4UL, 3UL);
setBytesFn(__mps_batch_encoder, __sel_setBytes, (const unsigned char*)&Ku, 4UL, 4UL);
setBytesFn(__mps_batch_encoder, __sel_setBytes, (const unsigned char*)&Nu, 4UL, 5UL);
struct MTLSize tgSize;
struct MTLSize numThreadgroups;
if (useMulti) {
tgSize.width = 128UL; tgSize.height = 1UL; tgSize.depth = 1UL;
numThreadgroups.width = N / 32UL; numThreadgroups.height = M / 32UL; numThreadgroups.depth = 1UL;
} else if (useSmma) {
tgSize.width = 32UL; tgSize.height = 1UL; tgSize.depth = 1UL;
numThreadgroups.width = N / 8UL; numThreadgroups.height = M / 8UL; numThreadgroups.depth = 1UL;
} else {
unsigned long tgW = 16UL;
unsigned long tgH = 16UL;
tgSize.width = tgW; tgSize.height = tgH; tgSize.depth = 1UL;
unsigned long groupsX = (N + tgW - 1UL) / tgW;
unsigned long groupsY = (M + tgH - 1UL) / tgH;
numThreadgroups.width = groupsX; numThreadgroups.height = groupsY; numThreadgroups.depth = 1UL;
}
MMsgDispatch dispatchFn = (MMsgDispatch)objc_msgSend;
dispatchFn(__mps_batch_encoder, __sel_dispatchThreadgroups, numThreadgroups, tgSize);
if (__mps_pending_count < (unsigned long)MPS_BATCH_MAX_PENDING) {
__mps_pending_gpuBuf[__mps_pending_count] = bufOut;
__mps_pending_cpuOut[__mps_pending_count] = (unsigned char*)out;
__mps_pending_bytes[__mps_pending_count] = bytesOut;
__mps_pending_dtype[__mps_pending_count] = 0;
__mps_pending_count = __mps_pending_count + 1UL;
}
return 1;
}
}
// ── MpsBuf: persistent device buffers (see gpu_mps.h) ────────────────────────
// A persistent buffer is a pool buffer (__mps_buf_get) the caller owns: it
// is NOT tracked by the batching machinery the way __mps_buf_get_batched's
// return value is, so it survives across dispatches and batches until
// MpsBuf::free returns it to the size-class pool (see __mps_buf_put's
// comment) -- which must only happen once every batch that might still
// read it has completed. A constructor may reuse a pool slot, so it is only
// safe while no GPU work that might still read that slot's previous
// contents is in flight (i.e. after any earlier batch's mps_batch_end()).
//
// Every constructor calls __mps_ensure_device_queue() first: one is
// typically the FIRST GPU call in a program (uploading a training loop's
// weights before any forward pass has run), and without the check
// __mps_buf_get would message a nil device (Objective-C's "message to nil
// returns nil", not a crash) and hand back a nil buffer that reads back as
// all zeros forever after -- found by a real correctness test, not by
// inspection.
static unsigned long __mps_dtype_bytes(int dtype) {
if (dtype == 0) return 4UL;
if (dtype == MPS_FP16 || dtype == MPS_BF16) return 2UL;
return 0UL;
}
static struct MpsBuf __mps_buf_make(MtlObject* handle, unsigned long bytes, int dtype) {
struct MpsBuf b;
b.handle = handle;
b.bytes = bytes;
b.dtype = dtype;
unsafe { if (handle == null) { b.bytes = 0UL; } }
return b;
}
static struct MpsBuf __mps_buf_upload(const unsigned char* src, unsigned long bytes, int dtype) {
if (!__mps_ensure_device_queue()) return __mps_buf_make(null, 0UL, dtype);
unsafe { return __mps_buf_make(__mps_buf_get_with_bytes(src, bytes), bytes, dtype); }
}
int MpsBuf::valid() const {
unsafe { return self.handle != null; }
}
void MpsBuf::free() {
unsafe {
if (self.handle != null) { __mps_buf_put(self.handle, self.bytes); }
self.handle = null;
}
self.bytes = 0UL;
}
int MpsBuf::read([]unsigned char dsts) const {
unsigned long bytes = dsts.len;
unsafe {
unsigned char* dst = (unsigned char*)dsts.ptr;
if (self.handle == null || bytes > self.bytes) return 0;
MMsg0 msg0 = (MMsg0)objc_msgSend;
unsigned char* src = __mps_contents(self.handle);
if (src == null) return 0;
memcpy((unsigned char*)dst, src, bytes);
return 1;
}
}
?&float MpsBuf::contents() const {
unsafe {
if (self.handle == null) return null;
MMsg0 msg0 = (MMsg0)objc_msgSend;
unsigned char* p = __mps_contents(self.handle);
if (p == null) return null;
return (?&float)(float*)p;
}
}
int MpsBuf::write([]const unsigned char srcs) {
unsigned long bytes = srcs.len;
unsafe {
const unsigned char* src = (const unsigned char*)srcs.ptr;
if (self.handle == null || bytes > self.bytes) return 0;
MMsg0 msg0 = (MMsg0)objc_msgSend;
unsigned char* dst = __mps_contents(self.handle);
if (dst == null) return 0;
memcpy(dst, (const unsigned char*)src, bytes);
return 1;
}
}
struct MpsBuf MpsBuf::alloc(unsigned long n, int dtype) {
unsigned long esize = __mps_dtype_bytes(dtype);
if (esize == 0UL || !__mps_ensure_device_queue()) return __mps_buf_make(null, 0UL, dtype);
unsafe { return __mps_buf_make(__mps_buf_get(n * esize), n * esize, dtype); }
}
struct MpsBuf MpsBuf::upload_n_(const &float host, unsigned long n) {
unsafe { return __mps_buf_upload((const unsigned char*)host, n * 4UL, 0); }
}
struct MpsBuf MpsBuf::upload([]const float host_s) {
unsigned long n = host_s.len;
float host_d_[1];
const &float host = &host_d_[0];
if (host_s.len > 0UL) { unsafe { host = (const &float)host_s.ptr; } }
return MpsBuf::upload_n_(host, n);
}
struct MpsBuf MpsBuf::upload_bits_n_(const &unsigned short bits, unsigned long n, int dtype) {
if (dtype != MPS_FP16 && dtype != MPS_BF16) return __mps_buf_make(null, 0UL, dtype);
unsafe { return __mps_buf_upload((const unsigned char*)bits, n * 2UL, dtype); }
}
struct MpsBuf MpsBuf::upload_bits([]const unsigned short bits_s, int dtype) {
unsigned long n = bits_s.len;
unsigned short bits_d_[1];
const &unsigned short bits = &bits_d_[0];
if (bits_s.len > 0UL) { unsafe { bits = (const &unsigned short)bits_s.ptr; } }
return MpsBuf::upload_bits_n_(bits, n, dtype);
}
struct MpsBuf MpsBuf::batch_last_output() {
if (!__mps_batch_active || __mps_pending_count == 0UL) return __mps_buf_make(null, 0UL, 0);
unsigned long i = __mps_pending_count - 1UL;
unsafe { return __mps_buf_make(__mps_pending_gpuBuf[i], __mps_pending_bytes[i], __mps_pending_dtype[i]); }
}
int mps_relu_f32_chained(struct MpsBuf bufA, &float out, unsigned long n) {
if (bufA.handle == null) return 0;
if (!__mps_batch_active) return 0;
unsafe {
if (__mps_relu_pipeline == null) {
MtlObject* pipeline = __mps_make_pipeline("relu_kernel");
if (pipeline == null) return 0;
__mps_relu_pipeline = pipeline;
}
unsigned long bytes = n * sizeof(float);
MtlObject* bufOut = __mps_buf_get_batched(bytes);
if (bufOut == null) return 0;
MMsgVoid1 msgVoid1 = (MMsgVoid1)objc_msgSend;
msgVoid1(__mps_batch_encoder, __sel_setComputePipelineState, __mps_relu_pipeline);
MMsgSetBuffer setBufFn = (MMsgSetBuffer)objc_msgSend;
setBufFn(__mps_batch_encoder, __sel_setBuffer, bufA.handle, 0UL, 0UL);
setBufFn(__mps_batch_encoder, __sel_setBuffer, bufOut, 0UL, 1UL);
struct MTLSize gridSize;
gridSize.width = n; gridSize.height = 1UL; gridSize.depth = 1UL;
unsigned long tgWidth = n < 64UL ? n : 64UL;
if (tgWidth == 0UL) tgWidth = 1UL;
struct MTLSize tgSize;
tgSize.width = tgWidth; tgSize.height = 1UL; tgSize.depth = 1UL;
MMsgDispatch dispatchFn = (MMsgDispatch)objc_msgSend;
dispatchFn(__mps_batch_encoder, __sel_dispatchThreads, gridSize, tgSize);
if (__mps_pending_count < (unsigned long)MPS_BATCH_MAX_PENDING) {
__mps_pending_gpuBuf[__mps_pending_count] = bufOut;
__mps_pending_cpuOut[__mps_pending_count] = (unsigned char*)out;
__mps_pending_bytes[__mps_pending_count] = bytes;
__mps_pending_dtype[__mps_pending_count] = 0;
__mps_pending_count = __mps_pending_count + 1UL;
}
return 1;
}
}
// out[i] = bufA[i] - b[i] -- see gpu_mps.h's comment.
int mps_sub_f32_chained(struct MpsBuf bufA, const &float b, &float out, unsigned long n) {
if (bufA.handle == null) return 0;
if (!__mps_batch_active) return 0;
unsafe {
if (__mps_sub_pipeline == null) {
MtlObject* pipeline = __mps_make_pipeline("sub_kernel");
if (pipeline == null) return 0;
__mps_sub_pipeline = pipeline;
}
unsigned long bytes = n * sizeof(float);
MtlObject* bufB = __mps_buf_get_with_bytes_batched((const unsigned char*)b, bytes);
MtlObject* bufOut = __mps_buf_get_batched(bytes);
if (bufB == null || bufOut == null) return 0;
MMsgVoid1 msgVoid1 = (MMsgVoid1)objc_msgSend;
msgVoid1(__mps_batch_encoder, __sel_setComputePipelineState, __mps_sub_pipeline);
MMsgSetBuffer setBufFn = (MMsgSetBuffer)objc_msgSend;
setBufFn(__mps_batch_encoder, __sel_setBuffer, bufA.handle, 0UL, 0UL);
setBufFn(__mps_batch_encoder, __sel_setBuffer, bufB, 0UL, 1UL);
setBufFn(__mps_batch_encoder, __sel_setBuffer, bufOut, 0UL, 2UL);
struct MTLSize gridSize;
gridSize.width = n; gridSize.height = 1UL; gridSize.depth = 1UL;
unsigned long tgWidth = n < 64UL ? n : 64UL;
if (tgWidth == 0UL) tgWidth = 1UL;
struct MTLSize tgSize;
tgSize.width = tgWidth; tgSize.height = 1UL; tgSize.depth = 1UL;
MMsgDispatch dispatchFn = (MMsgDispatch)objc_msgSend;
dispatchFn(__mps_batch_encoder, __sel_dispatchThreads, gridSize, tgSize);
if (__mps_pending_count < (unsigned long)MPS_BATCH_MAX_PENDING) {
__mps_pending_gpuBuf[__mps_pending_count] = bufOut;
__mps_pending_cpuOut[__mps_pending_count] = (unsigned char*)out;
__mps_pending_bytes[__mps_pending_count] = bytes;
__mps_pending_dtype[__mps_pending_count] = 0;
__mps_pending_count = __mps_pending_count + 1UL;
}
return 1;
}
}
// out[i] = bufA[i] * b[i] -- see gpu_mps.h's comment.
int mps_mul_f32_chained(struct MpsBuf bufA, const &float b, &float out, unsigned long n) {
if (bufA.handle == null) return 0;
if (!__mps_batch_active) return 0;
unsafe {
if (__mps_mul_pipeline == null) {
MtlObject* pipeline = __mps_make_pipeline("mul_kernel");
if (pipeline == null) return 0;
__mps_mul_pipeline = pipeline;
}
unsigned long bytes = n * sizeof(float);
MtlObject* bufB = __mps_buf_get_with_bytes_batched((const unsigned char*)b, bytes);
MtlObject* bufOut = __mps_buf_get_batched(bytes);
if (bufB == null || bufOut == null) return 0;
MMsgVoid1 msgVoid1 = (MMsgVoid1)objc_msgSend;
msgVoid1(__mps_batch_encoder, __sel_setComputePipelineState, __mps_mul_pipeline);
MMsgSetBuffer setBufFn = (MMsgSetBuffer)objc_msgSend;
setBufFn(__mps_batch_encoder, __sel_setBuffer, bufA.handle, 0UL, 0UL);
setBufFn(__mps_batch_encoder, __sel_setBuffer, bufB, 0UL, 1UL);
setBufFn(__mps_batch_encoder, __sel_setBuffer, bufOut, 0UL, 2UL);
struct MTLSize gridSize;
gridSize.width = n; gridSize.height = 1UL; gridSize.depth = 1UL;
unsigned long tgWidth = n < 64UL ? n : 64UL;
if (tgWidth == 0UL) tgWidth = 1UL;
struct MTLSize tgSize;
tgSize.width = tgWidth; tgSize.height = 1UL; tgSize.depth = 1UL;
MMsgDispatch dispatchFn = (MMsgDispatch)objc_msgSend;
dispatchFn(__mps_batch_encoder, __sel_dispatchThreads, gridSize, tgSize);
if (__mps_pending_count < (unsigned long)MPS_BATCH_MAX_PENDING) {
__mps_pending_gpuBuf[__mps_pending_count] = bufOut;
__mps_pending_cpuOut[__mps_pending_count] = (unsigned char*)out;
__mps_pending_bytes[__mps_pending_count] = bytes;
__mps_pending_dtype[__mps_pending_count] = 0;
__mps_pending_count = __mps_pending_count + 1UL;
}
return 1;
}
}
// out[i] = buf[i] * buf[i] -- see gpu_mps.h's comment. Same mul_kernel as
// mps_mul_f32_chained above, just bound to buffer indices 0 AND 1 with the
// same buffer -- reading the same device memory location from two bound
// slots is safe (both reads, no write), so this needs no dedicated kernel.
int mps_square_f32_chained(struct MpsBuf buf, &float out, unsigned long n) {
if (buf.handle == null) return 0;
if (!__mps_batch_active) return 0;
unsafe {
if (__mps_mul_pipeline == null) {
MtlObject* pipeline = __mps_make_pipeline("mul_kernel");
if (pipeline == null) return 0;
__mps_mul_pipeline = pipeline;
}
unsigned long bytes = n * sizeof(float);
MtlObject* bufOut = __mps_buf_get_batched(bytes);
if (bufOut == null) return 0;
MMsgVoid1 msgVoid1 = (MMsgVoid1)objc_msgSend;
msgVoid1(__mps_batch_encoder, __sel_setComputePipelineState, __mps_mul_pipeline);
MMsgSetBuffer setBufFn = (MMsgSetBuffer)objc_msgSend;
setBufFn(__mps_batch_encoder, __sel_setBuffer, buf.handle, 0UL, 0UL);
setBufFn(__mps_batch_encoder, __sel_setBuffer, buf.handle, 0UL, 1UL);
setBufFn(__mps_batch_encoder, __sel_setBuffer, bufOut, 0UL, 2UL);
struct MTLSize gridSize;
gridSize.width = n; gridSize.height = 1UL; gridSize.depth = 1UL;
unsigned long tgWidth = n < 64UL ? n : 64UL;
if (tgWidth == 0UL) tgWidth = 1UL;
struct MTLSize tgSize;
tgSize.width = tgWidth; tgSize.height = 1UL; tgSize.depth = 1UL;
MMsgDispatch dispatchFn = (MMsgDispatch)objc_msgSend;
dispatchFn(__mps_batch_encoder, __sel_dispatchThreads, gridSize, tgSize);
if (__mps_pending_count < (unsigned long)MPS_BATCH_MAX_PENDING) {
__mps_pending_gpuBuf[__mps_pending_count] = bufOut;
__mps_pending_cpuOut[__mps_pending_count] = (unsigned char*)out;
__mps_pending_bytes[__mps_pending_count] = bytes;
__mps_pending_dtype[__mps_pending_count] = 0;
__mps_pending_count = __mps_pending_count + 1UL;
}
return 1;
}
}
// out[i] = buf[i] * k -- see gpu_mps.h's comment.
int mps_scale_f32_chained(struct MpsBuf buf, float k, &float out, unsigned long n) {
if (buf.handle == null) return 0;
if (!__mps_batch_active) return 0;
unsafe {
if (__mps_scale_pipeline == null) {
MtlObject* pipeline = __mps_make_pipeline("scale_kernel");
if (pipeline == null) return 0;
__mps_scale_pipeline = pipeline;
}
unsigned long bytes = n * sizeof(float);
MtlObject* bufOut = __mps_buf_get_batched(bytes);
if (bufOut == null) return 0;
MMsgVoid1 msgVoid1 = (MMsgVoid1)objc_msgSend;
msgVoid1(__mps_batch_encoder, __sel_setComputePipelineState, __mps_scale_pipeline);
MMsgSetBuffer setBufFn = (MMsgSetBuffer)objc_msgSend;
setBufFn(__mps_batch_encoder, __sel_setBuffer, buf.handle, 0UL, 0UL);
MMsgSetBytes setBytesFn = (MMsgSetBytes)objc_msgSend;
float kVal = k;
setBytesFn(__mps_batch_encoder, __sel_setBytes, (const unsigned char*)&kVal, 4UL, 1UL);
setBufFn(__mps_batch_encoder, __sel_setBuffer, bufOut, 0UL, 2UL);
struct MTLSize gridSize;
gridSize.width = n; gridSize.height = 1UL; gridSize.depth = 1UL;
unsigned long tgWidth = n < 64UL ? n : 64UL;
if (tgWidth == 0UL) tgWidth = 1UL;
struct MTLSize tgSize;
tgSize.width = tgWidth; tgSize.height = 1UL; tgSize.depth = 1UL;
MMsgDispatch dispatchFn = (MMsgDispatch)objc_msgSend;
dispatchFn(__mps_batch_encoder, __sel_dispatchThreads, gridSize, tgSize);
if (__mps_pending_count < (unsigned long)MPS_BATCH_MAX_PENDING) {
__mps_pending_gpuBuf[__mps_pending_count] = bufOut;
__mps_pending_cpuOut[__mps_pending_count] = (unsigned char*)out;
__mps_pending_bytes[__mps_pending_count] = bytes;
__mps_pending_dtype[__mps_pending_count] = 0;
__mps_pending_count = __mps_pending_count + 1UL;
}
return 1;
}
}
// See gpu_mps.h's comment. Same dispatch shape as mps_scale_f32_chained
// above (one buffer input, one scalar constant, one buffer output), just
// with the scalar as a 4-byte 'uint' (matching the kernel's 'constant
// uint&') instead of a float, and a fused gate/up-split + GELU + multiply
// kernel body instead of a plain multiply-by-scalar. Shared by the f32 and
// bf16 variants; 'elemBytes' is the storage width (4 or 2).
static int __mps_gated_gelu_chained(const char* kernelName, struct MtlObject** pipelineSlot,
MtlObject* bufRaw, unsigned char* out, unsigned long rows,
unsigned long intermediateSize, unsigned long elemBytes) {
if (!__mps_batch_active) return 0;
unsafe {
if (*pipelineSlot == null) {
MtlObject* p = __mps_make_pipeline(kernelName);
if (p == null) return 0;
*pipelineSlot = p;
}
unsigned long n = rows * intermediateSize;
unsigned long bytes = n * elemBytes;
MtlObject* bufOut = __mps_buf_get_batched(bytes);
if (bufOut == null) return 0;
MMsgVoid1 msgVoid1 = (MMsgVoid1)objc_msgSend;
msgVoid1(__mps_batch_encoder, __sel_setComputePipelineState, *pipelineSlot);
MMsgSetBuffer setBufFn = (MMsgSetBuffer)objc_msgSend;
setBufFn(__mps_batch_encoder, __sel_setBuffer, bufRaw, 0UL, 0UL);
MMsgSetBytes setBytesFn = (MMsgSetBytes)objc_msgSend;
unsigned int iSizeVal = (unsigned int)intermediateSize;
setBytesFn(__mps_batch_encoder, __sel_setBytes, (const unsigned char*)&iSizeVal, 4UL, 1UL);
setBufFn(__mps_batch_encoder, __sel_setBuffer, bufOut, 0UL, 2UL);
struct MTLSize gridSize;
gridSize.width = n; gridSize.height = 1UL; gridSize.depth = 1UL;
unsigned long tgWidth = n < 64UL ? n : 64UL;
if (tgWidth == 0UL) tgWidth = 1UL;
struct MTLSize tgSize;
tgSize.width = tgWidth; tgSize.height = 1UL; tgSize.depth = 1UL;
MMsgDispatch dispatchFn = (MMsgDispatch)objc_msgSend;
dispatchFn(__mps_batch_encoder, __sel_dispatchThreads, gridSize, tgSize);
if (__mps_pending_count < (unsigned long)MPS_BATCH_MAX_PENDING) {
__mps_pending_gpuBuf[__mps_pending_count] = bufOut;
__mps_pending_cpuOut[__mps_pending_count] = out;
__mps_pending_bytes[__mps_pending_count] = bytes;
__mps_pending_dtype[__mps_pending_count] = elemBytes == 2UL ? MPS_BF16 : 0;
__mps_pending_count = __mps_pending_count + 1UL;
}
return 1;
}
}
int mps_gated_gelu_f32_chained(struct MpsBuf bufRaw, &float out, unsigned long rows, unsigned long intermediateSize) {
if (bufRaw.handle == null) return 0;
unsafe {
return __mps_gated_gelu_chained("gated_gelu_kernel", &__mps_gated_gelu_pipeline,
bufRaw.handle, (unsigned char*)out, rows, intermediateSize, sizeof(float));
}
}
static MtlObject* __mps_gated_gelu_bf16_pipeline = null;
int mps_gated_gelu_bf16_chained(struct MpsBuf bufRaw, &unsigned short out, unsigned long rows, unsigned long intermediateSize) {
if (bufRaw.handle == null) return 0;
unsafe {
return __mps_gated_gelu_chained("gated_gelu_kernel_bf16", &__mps_gated_gelu_bf16_pipeline,
bufRaw.handle, (unsigned char*)out, rows, intermediateSize, 2UL);
}
}
// ── Backward chained variants (see gpu_mps.h's own comment on these) ─────────
// Same "read a previous batched op's still-pending GPU buffer directly,
// no CPU round trip" idea as mps_matmul_f32_chained/mps_relu_f32_chained
// above, applied to the backward chain instead of the forward one.
int mps_matmul_abt_f32_chained(struct MpsBuf bufDC, const &float b, &float out,
unsigned long M, unsigned long N, unsigned long K) {
if (bufDC.handle == null) return 0;
if (!__mps_batch_active) return 0;
unsafe {
int useMulti = (M % 32UL == 0UL) && (N % 32UL == 0UL) && (K % 32UL == 0UL);
int useSmma = (M % 8UL == 0UL) && (N % 8UL == 0UL) && (K % 8UL == 0UL);
MtlObject* pipeline;
if (useMulti) {
if (__mps_matmul_abt_multi_pipeline == null) {
MtlObject* p = __mps_make_pipeline("matmul_abt_kernel_smma_multi");
if (p == null) { useMulti = 0; } else { __mps_matmul_abt_multi_pipeline = p; }
}
if (useMulti) pipeline = __mps_matmul_abt_multi_pipeline;
}
if (!useMulti && useSmma) {
if (__mps_matmul_abt_smma_pipeline == null) {
MtlObject* p = __mps_make_pipeline("matmul_abt_kernel_smma");
if (p == null) { useSmma = 0; } else { __mps_matmul_abt_smma_pipeline = p; }
}
if (useSmma) pipeline = __mps_matmul_abt_smma_pipeline;
}
if (!useMulti && !useSmma) {
if (__mps_matmul_abt_pipeline == null) {
MtlObject* p = __mps_make_pipeline("matmul_abt_kernel");
if (p == null) return 0;
__mps_matmul_abt_pipeline = p;
}
pipeline = __mps_matmul_abt_pipeline;
}
unsigned long bytesB = K * N * sizeof(float);
unsigned long bytesOut = M * K * sizeof(float);
MtlObject* bufB = __mps_buf_get_with_bytes_batched((const unsigned char*)b, bytesB);
MtlObject* bufOut = __mps_buf_get_batched(bytesOut);
if (bufB == null || bufOut == null) return 0;
unsigned int Mu = (unsigned int)M;
unsigned int Nu = (unsigned int)N;
unsigned int Ku = (unsigned int)K;
MMsgVoid1 msgVoid1 = (MMsgVoid1)objc_msgSend;
msgVoid1(__mps_batch_encoder, __sel_setComputePipelineState, pipeline);
MMsgSetBuffer setBufFn = (MMsgSetBuffer)objc_msgSend;
setBufFn(__mps_batch_encoder, __sel_setBuffer, bufDC.handle, 0UL, 0UL);
setBufFn(__mps_batch_encoder, __sel_setBuffer, bufB, 0UL, 1UL);
setBufFn(__mps_batch_encoder, __sel_setBuffer, bufOut, 0UL, 2UL);
MMsgSetBytes setBytesFn = (MMsgSetBytes)objc_msgSend;
setBytesFn(__mps_batch_encoder, __sel_setBytes, (const unsigned char*)&Mu, 4UL, 3UL);
setBytesFn(__mps_batch_encoder, __sel_setBytes, (const unsigned char*)&Nu, 4UL, 4UL);
setBytesFn(__mps_batch_encoder, __sel_setBytes, (const unsigned char*)&Ku, 4UL, 5UL);
struct MTLSize tgSize;
struct MTLSize numThreadgroups;
if (useMulti) {
tgSize.width = 128UL; tgSize.height = 1UL; tgSize.depth = 1UL;
numThreadgroups.width = K / 32UL; numThreadgroups.height = M / 32UL; numThreadgroups.depth = 1UL;
} else if (useSmma) {
tgSize.width = 32UL; tgSize.height = 1UL; tgSize.depth = 1UL;
numThreadgroups.width = K / 8UL; numThreadgroups.height = M / 8UL; numThreadgroups.depth = 1UL;
} else {
unsigned long tgW = 16UL;
unsigned long tgH = 16UL;
tgSize.width = tgW; tgSize.height = tgH; tgSize.depth = 1UL;
unsigned long groupsX = (K + tgW - 1UL) / tgW;
unsigned long groupsY = (M + tgH - 1UL) / tgH;
numThreadgroups.width = groupsX; numThreadgroups.height = groupsY; numThreadgroups.depth = 1UL;
}
MMsgDispatch dispatchFn = (MMsgDispatch)objc_msgSend;
dispatchFn(__mps_batch_encoder, __sel_dispatchThreadgroups, numThreadgroups, tgSize);
if (__mps_pending_count < (unsigned long)MPS_BATCH_MAX_PENDING) {
__mps_pending_gpuBuf[__mps_pending_count] = bufOut;
__mps_pending_cpuOut[__mps_pending_count] = (unsigned char*)out;
__mps_pending_bytes[__mps_pending_count] = bytesOut;
__mps_pending_dtype[__mps_pending_count] = 0;
__mps_pending_count = __mps_pending_count + 1UL;
}
return 1;
}
}
// Same as mps_matmul_abt_f32_chained above, but 'b' is ALSO a pre-existing
// device buffer (typically MpsBuf::upload, e.g. a weight matrix
// that never leaves the GPU across a whole training run) instead of a
// plain CPU array to upload fresh -- same reasoning as mps_matmul_f32_
// persistent. Needed for dH = dY . W2^T when W2 stays GPU-resident and
// is updated in place by sgd_update_kernel every step instead of ever
// being read back to a CPU array.
int mps_matmul_abt_f32_persistent(struct MpsBuf bufDC, struct MpsBuf bufB, &float out,
unsigned long M, unsigned long N, unsigned long K) {
if (bufDC.handle == null || bufB.handle == null) return 0;
if (!__mps_batch_active) return 0;
unsafe {
int useMulti = (M % 32UL == 0UL) && (N % 32UL == 0UL) && (K % 32UL == 0UL);
int useSmma = (M % 8UL == 0UL) && (N % 8UL == 0UL) && (K % 8UL == 0UL);
MtlObject* pipeline;
if (useMulti) {
if (__mps_matmul_abt_multi_pipeline == null) {
MtlObject* p = __mps_make_pipeline("matmul_abt_kernel_smma_multi");
if (p == null) { useMulti = 0; } else { __mps_matmul_abt_multi_pipeline = p; }
}
if (useMulti) pipeline = __mps_matmul_abt_multi_pipeline;
}
if (!useMulti && useSmma) {
if (__mps_matmul_abt_smma_pipeline == null) {
MtlObject* p = __mps_make_pipeline("matmul_abt_kernel_smma");
if (p == null) { useSmma = 0; } else { __mps_matmul_abt_smma_pipeline = p; }
}
if (useSmma) pipeline = __mps_matmul_abt_smma_pipeline;
}
if (!useMulti && !useSmma) {
if (__mps_matmul_abt_pipeline == null) {
MtlObject* p = __mps_make_pipeline("matmul_abt_kernel");
if (p == null) return 0;
__mps_matmul_abt_pipeline = p;
}
pipeline = __mps_matmul_abt_pipeline;
}
unsigned long bytesOut = M * K * sizeof(float);
MtlObject* bufOut = __mps_buf_get_batched(bytesOut);
if (bufOut == null) return 0;
unsigned int Mu = (unsigned int)M;
unsigned int Nu = (unsigned int)N;
unsigned int Ku = (unsigned int)K;
MMsgVoid1 msgVoid1 = (MMsgVoid1)objc_msgSend;
msgVoid1(__mps_batch_encoder, __sel_setComputePipelineState, pipeline);
MMsgSetBuffer setBufFn = (MMsgSetBuffer)objc_msgSend;
setBufFn(__mps_batch_encoder, __sel_setBuffer, bufDC.handle, 0UL, 0UL);
setBufFn(__mps_batch_encoder, __sel_setBuffer, bufB.handle, 0UL, 1UL);
setBufFn(__mps_batch_encoder, __sel_setBuffer, bufOut, 0UL, 2UL);
MMsgSetBytes setBytesFn = (MMsgSetBytes)objc_msgSend;
setBytesFn(__mps_batch_encoder, __sel_setBytes, (const unsigned char*)&Mu, 4UL, 3UL);
setBytesFn(__mps_batch_encoder, __sel_setBytes, (const unsigned char*)&Nu, 4UL, 4UL);
setBytesFn(__mps_batch_encoder, __sel_setBytes, (const unsigned char*)&Ku, 4UL, 5UL);
struct MTLSize tgSize;
struct MTLSize numThreadgroups;
if (useMulti) {
tgSize.width = 128UL; tgSize.height = 1UL; tgSize.depth = 1UL;
numThreadgroups.width = K / 32UL; numThreadgroups.height = M / 32UL; numThreadgroups.depth = 1UL;
} else if (useSmma) {
tgSize.width = 32UL; tgSize.height = 1UL; tgSize.depth = 1UL;
numThreadgroups.width = K / 8UL; numThreadgroups.height = M / 8UL; numThreadgroups.depth = 1UL;
} else {
unsigned long tgW = 16UL;
unsigned long tgH = 16UL;
tgSize.width = tgW; tgSize.height = tgH; tgSize.depth = 1UL;
unsigned long groupsX = (K + tgW - 1UL) / tgW;
unsigned long groupsY = (M + tgH - 1UL) / tgH;
numThreadgroups.width = groupsX; numThreadgroups.height = groupsY; numThreadgroups.depth = 1UL;
}
MMsgDispatch dispatchFn = (MMsgDispatch)objc_msgSend;
dispatchFn(__mps_batch_encoder, __sel_dispatchThreadgroups, numThreadgroups, tgSize);
if (__mps_pending_count < (unsigned long)MPS_BATCH_MAX_PENDING) {
__mps_pending_gpuBuf[__mps_pending_count] = bufOut;
__mps_pending_cpuOut[__mps_pending_count] = (unsigned char*)out;
__mps_pending_bytes[__mps_pending_count] = bytesOut;
__mps_pending_dtype[__mps_pending_count] = 0;
__mps_pending_count = __mps_pending_count + 1UL;
}
return 1;
}
}
int mps_matmul_atb_f32_chained(const &float a, struct MpsBuf bufDC, &float out,
unsigned long M, unsigned long K, unsigned long N) {
if (bufDC.handle == null) return 0;
if (!__mps_batch_active) return 0;
unsafe {
int useMulti = (M % 32UL == 0UL) && (K % 32UL == 0UL) && (N % 32UL == 0UL);
int useSmma = (M % 8UL == 0UL) && (K % 8UL == 0UL) && (N % 8UL == 0UL);
MtlObject* pipeline;
if (useMulti) {
if (__mps_matmul_atb_multi_pipeline == null) {
MtlObject* p = __mps_make_pipeline("matmul_atb_kernel_smma_multi");
if (p == null) { useMulti = 0; } else { __mps_matmul_atb_multi_pipeline = p; }
}
if (useMulti) pipeline = __mps_matmul_atb_multi_pipeline;
}
if (!useMulti && useSmma) {
if (__mps_matmul_atb_smma_pipeline == null) {
MtlObject* p = __mps_make_pipeline("matmul_atb_kernel_smma");
if (p == null) { useSmma = 0; } else { __mps_matmul_atb_smma_pipeline = p; }
}
if (useSmma) pipeline = __mps_matmul_atb_smma_pipeline;
}
if (!useMulti && !useSmma) {
if (__mps_matmul_atb_pipeline == null) {
MtlObject* p = __mps_make_pipeline("matmul_atb_kernel");
if (p == null) return 0;
__mps_matmul_atb_pipeline = p;
}
pipeline = __mps_matmul_atb_pipeline;
}
unsigned long bytesA = M * K * sizeof(float);
unsigned long bytesOut = K * N * sizeof(float);
MtlObject* bufA = __mps_buf_get_with_bytes_batched((const unsigned char*)a, bytesA);
MtlObject* bufOut = __mps_buf_get_batched(bytesOut);
if (bufA == null || bufOut == null) return 0;
unsigned int Mu = (unsigned int)M;
unsigned int Ku = (unsigned int)K;
unsigned int Nu = (unsigned int)N;
MMsgVoid1 msgVoid1 = (MMsgVoid1)objc_msgSend;
msgVoid1(__mps_batch_encoder, __sel_setComputePipelineState, pipeline);
MMsgSetBuffer setBufFn = (MMsgSetBuffer)objc_msgSend;
setBufFn(__mps_batch_encoder, __sel_setBuffer, bufA, 0UL, 0UL);
setBufFn(__mps_batch_encoder, __sel_setBuffer, bufDC.handle, 0UL, 1UL);
setBufFn(__mps_batch_encoder, __sel_setBuffer, bufOut, 0UL, 2UL);
MMsgSetBytes setBytesFn = (MMsgSetBytes)objc_msgSend;
setBytesFn(__mps_batch_encoder, __sel_setBytes, (const unsigned char*)&Mu, 4UL, 3UL);
setBytesFn(__mps_batch_encoder, __sel_setBytes, (const unsigned char*)&Ku, 4UL, 4UL);
setBytesFn(__mps_batch_encoder, __sel_setBytes, (const unsigned char*)&Nu, 4UL, 5UL);
struct MTLSize tgSize;
struct MTLSize numThreadgroups;
if (useMulti) {
tgSize.width = 128UL; tgSize.height = 1UL; tgSize.depth = 1UL;
numThreadgroups.width = N / 32UL; numThreadgroups.height = K / 32UL; numThreadgroups.depth = 1UL;
} else if (useSmma) {
tgSize.width = 32UL; tgSize.height = 1UL; tgSize.depth = 1UL;
numThreadgroups.width = N / 8UL; numThreadgroups.height = K / 8UL; numThreadgroups.depth = 1UL;
} else {
unsigned long tgW = 16UL;
unsigned long tgH = 16UL;
tgSize.width = tgW; tgSize.height = tgH; tgSize.depth = 1UL;
unsigned long groupsX = (N + tgW - 1UL) / tgW;
unsigned long groupsY = (K + tgH - 1UL) / tgH;
numThreadgroups.width = groupsX; numThreadgroups.height = groupsY; numThreadgroups.depth = 1UL;
}
MMsgDispatch dispatchFn = (MMsgDispatch)objc_msgSend;
dispatchFn(__mps_batch_encoder, __sel_dispatchThreadgroups, numThreadgroups, tgSize);
if (__mps_pending_count < (unsigned long)MPS_BATCH_MAX_PENDING) {
__mps_pending_gpuBuf[__mps_pending_count] = bufOut;
__mps_pending_cpuOut[__mps_pending_count] = (unsigned char*)out;
__mps_pending_bytes[__mps_pending_count] = bytesOut;
__mps_pending_dtype[__mps_pending_count] = 0;
__mps_pending_count = __mps_pending_count + 1UL;
}
return 1;
}
}
// Same as mps_matmul_atb_f32_chained above, but 'a' is ALSO a still-pending
// GPU buffer instead of a plain CPU array. Needed for the fused
// forward+loss+backward training step: dW2 = H^T . dY needs BOTH H (an
// activation computed earlier in the SAME still-open batch by
// Tensor::relu_gpu, not yet read back to CPU) and dY (matmul2-backward's
// own incoming gradient, also still-pending) read straight from GPU
// memory -- mps_matmul_atb_f32_chained's plain-CPU-array 'a' parameter
// can't express that. (Doesn't arise for matmul1's own backward: its 'a'
// is X, a real input tensor, always CPU-resident already.)
int mps_matmul_atb_f32_chained_ab(struct MpsBuf bufA, struct MpsBuf bufDC, &float out,
unsigned long M, unsigned long K, unsigned long N) {
if (bufA.handle == null || bufDC.handle == null) return 0;
if (!__mps_batch_active) return 0;
unsafe {
int useMulti = (M % 32UL == 0UL) && (K % 32UL == 0UL) && (N % 32UL == 0UL);
int useSmma = (M % 8UL == 0UL) && (K % 8UL == 0UL) && (N % 8UL == 0UL);
MtlObject* pipeline;
if (useMulti) {
if (__mps_matmul_atb_multi_pipeline == null) {
MtlObject* p = __mps_make_pipeline("matmul_atb_kernel_smma_multi");
if (p == null) { useMulti = 0; } else { __mps_matmul_atb_multi_pipeline = p; }
}
if (useMulti) pipeline = __mps_matmul_atb_multi_pipeline;
}
if (!useMulti && useSmma) {
if (__mps_matmul_atb_smma_pipeline == null) {
MtlObject* p = __mps_make_pipeline("matmul_atb_kernel_smma");
if (p == null) { useSmma = 0; } else { __mps_matmul_atb_smma_pipeline = p; }
}
if (useSmma) pipeline = __mps_matmul_atb_smma_pipeline;
}
if (!useMulti && !useSmma) {
if (__mps_matmul_atb_pipeline == null) {
MtlObject* p = __mps_make_pipeline("matmul_atb_kernel");
if (p == null) return 0;
__mps_matmul_atb_pipeline = p;
}
pipeline = __mps_matmul_atb_pipeline;
}
unsigned long bytesOut = K * N * sizeof(float);
MtlObject* bufOut = __mps_buf_get_batched(bytesOut);
if (bufOut == null) return 0;
unsigned int Mu = (unsigned int)M;
unsigned int Ku = (unsigned int)K;
unsigned int Nu = (unsigned int)N;
MMsgVoid1 msgVoid1 = (MMsgVoid1)objc_msgSend;
msgVoid1(__mps_batch_encoder, __sel_setComputePipelineState, pipeline);
MMsgSetBuffer setBufFn = (MMsgSetBuffer)objc_msgSend;
setBufFn(__mps_batch_encoder, __sel_setBuffer, bufA.handle, 0UL, 0UL);
setBufFn(__mps_batch_encoder, __sel_setBuffer, bufDC.handle, 0UL, 1UL);
setBufFn(__mps_batch_encoder, __sel_setBuffer, bufOut, 0UL, 2UL);
MMsgSetBytes setBytesFn = (MMsgSetBytes)objc_msgSend;
setBytesFn(__mps_batch_encoder, __sel_setBytes, (const unsigned char*)&Mu, 4UL, 3UL);
setBytesFn(__mps_batch_encoder, __sel_setBytes, (const unsigned char*)&Ku, 4UL, 4UL);
setBytesFn(__mps_batch_encoder, __sel_setBytes, (const unsigned char*)&Nu, 4UL, 5UL);
struct MTLSize tgSize;
struct MTLSize numThreadgroups;
if (useMulti) {
tgSize.width = 128UL; tgSize.height = 1UL; tgSize.depth = 1UL;
numThreadgroups.width = N / 32UL; numThreadgroups.height = K / 32UL; numThreadgroups.depth = 1UL;
} else if (useSmma) {
tgSize.width = 32UL; tgSize.height = 1UL; tgSize.depth = 1UL;
numThreadgroups.width = N / 8UL; numThreadgroups.height = K / 8UL; numThreadgroups.depth = 1UL;
} else {
unsigned long tgW = 16UL;
unsigned long tgH = 16UL;
tgSize.width = tgW; tgSize.height = tgH; tgSize.depth = 1UL;
unsigned long groupsX = (N + tgW - 1UL) / tgW;
unsigned long groupsY = (K + tgH - 1UL) / tgH;
numThreadgroups.width = groupsX; numThreadgroups.height = groupsY; numThreadgroups.depth = 1UL;
}
MMsgDispatch dispatchFn = (MMsgDispatch)objc_msgSend;
dispatchFn(__mps_batch_encoder, __sel_dispatchThreadgroups, numThreadgroups, tgSize);
if (__mps_pending_count < (unsigned long)MPS_BATCH_MAX_PENDING) {
__mps_pending_gpuBuf[__mps_pending_count] = bufOut;
__mps_pending_cpuOut[__mps_pending_count] = (unsigned char*)out;
__mps_pending_bytes[__mps_pending_count] = bytesOut;
__mps_pending_dtype[__mps_pending_count] = 0;
__mps_pending_count = __mps_pending_count + 1UL;
}
return 1;
}
}
int mps_relu_backward_f32_chained(const &float a, struct MpsBuf bufSelfGrad, &float out, unsigned long n) {
if (bufSelfGrad.handle == null) return 0;
if (!__mps_batch_active) return 0;
unsafe {
if (__mps_relu_backward_pipeline == null) {
MtlObject* pipeline = __mps_make_pipeline("relu_backward_kernel");
if (pipeline == null) return 0;
__mps_relu_backward_pipeline = pipeline;
}
unsigned long bytes = n * sizeof(float);
MtlObject* bufA = __mps_buf_get_with_bytes_batched((const unsigned char*)a, bytes);
MtlObject* bufOut = __mps_buf_get_batched(bytes);
if (bufA == null || bufOut == null) return 0;
MMsgVoid1 msgVoid1 = (MMsgVoid1)objc_msgSend;
msgVoid1(__mps_batch_encoder, __sel_setComputePipelineState, __mps_relu_backward_pipeline);
MMsgSetBuffer setBufFn = (MMsgSetBuffer)objc_msgSend;
setBufFn(__mps_batch_encoder, __sel_setBuffer, bufA, 0UL, 0UL);
setBufFn(__mps_batch_encoder, __sel_setBuffer, bufSelfGrad.handle, 0UL, 1UL);
setBufFn(__mps_batch_encoder, __sel_setBuffer, bufOut, 0UL, 2UL);
struct MTLSize gridSize;
gridSize.width = n; gridSize.height = 1UL; gridSize.depth = 1UL;
unsigned long tgWidth = n < 64UL ? n : 64UL;
if (tgWidth == 0UL) tgWidth = 1UL;
struct MTLSize tgSize;
tgSize.width = tgWidth; tgSize.height = 1UL; tgSize.depth = 1UL;
MMsgDispatch dispatchFn = (MMsgDispatch)objc_msgSend;
dispatchFn(__mps_batch_encoder, __sel_dispatchThreads, gridSize, tgSize);
if (__mps_pending_count < (unsigned long)MPS_BATCH_MAX_PENDING) {
__mps_pending_gpuBuf[__mps_pending_count] = bufOut;
__mps_pending_cpuOut[__mps_pending_count] = (unsigned char*)out;
__mps_pending_bytes[__mps_pending_count] = bytes;
__mps_pending_dtype[__mps_pending_count] = 0;
__mps_pending_count = __mps_pending_count + 1UL;
}
return 1;
}
}
// Same as mps_relu_backward_f32_chained above, but 'a' is ALSO a
// still-pending GPU buffer instead of a plain CPU array -- same reasoning
// as mps_matmul_atb_f32_chained_ab (see its comment): in the fused
// forward+loss+backward training step, relu's own input (Z1, matmul1's
// output) hasn't been read back to CPU yet either when relu-backward runs.
int mps_relu_backward_f32_chained_ab(struct MpsBuf bufA, struct MpsBuf bufSelfGrad, &float out, unsigned long n) {
if (bufA.handle == null || bufSelfGrad.handle == null) return 0;
if (!__mps_batch_active) return 0;
unsafe {
if (__mps_relu_backward_pipeline == null) {
MtlObject* pipeline = __mps_make_pipeline("relu_backward_kernel");
if (pipeline == null) return 0;
__mps_relu_backward_pipeline = pipeline;
}
unsigned long bytes = n * sizeof(float);
MtlObject* bufOut = __mps_buf_get_batched(bytes);
if (bufOut == null) return 0;
MMsgVoid1 msgVoid1 = (MMsgVoid1)objc_msgSend;
msgVoid1(__mps_batch_encoder, __sel_setComputePipelineState, __mps_relu_backward_pipeline);
MMsgSetBuffer setBufFn = (MMsgSetBuffer)objc_msgSend;
setBufFn(__mps_batch_encoder, __sel_setBuffer, bufA.handle, 0UL, 0UL);
setBufFn(__mps_batch_encoder, __sel_setBuffer, bufSelfGrad.handle, 0UL, 1UL);
setBufFn(__mps_batch_encoder, __sel_setBuffer, bufOut, 0UL, 2UL);
struct MTLSize gridSize;
gridSize.width = n; gridSize.height = 1UL; gridSize.depth = 1UL;
unsigned long tgWidth = n < 64UL ? n : 64UL;
if (tgWidth == 0UL) tgWidth = 1UL;
struct MTLSize tgSize;
tgSize.width = tgWidth; tgSize.height = 1UL; tgSize.depth = 1UL;
MMsgDispatch dispatchFn = (MMsgDispatch)objc_msgSend;
dispatchFn(__mps_batch_encoder, __sel_dispatchThreads, gridSize, tgSize);
if (__mps_pending_count < (unsigned long)MPS_BATCH_MAX_PENDING) {
__mps_pending_gpuBuf[__mps_pending_count] = bufOut;
__mps_pending_cpuOut[__mps_pending_count] = (unsigned char*)out;
__mps_pending_bytes[__mps_pending_count] = bytes;
__mps_pending_dtype[__mps_pending_count] = 0;
__mps_pending_count = __mps_pending_count + 1UL;
}
return 1;
}
}
// In-place SGD step against a persistent device buffer (see gpu_mps_
// kernels.metal's sgd_update_kernel): w[i] -= lr*grad[i], both already
// device buffers -- 'bufW' typically from MpsBuf::upload (a weight
// matrix uploaded once at the start of training and never read back to a
// CPU array until the caller is actually done with it), 'bufGrad'
// typically another op's still-pending chained output from the SAME
// batch's backward pass. No 'out'/pending-readback parameter: unlike
// every other batched op here, this one's result is never meant to reach
// CPU memory as part of the batch ending -- that's the entire point of
// keeping a weight GPU-resident across a training run instead of
// re-uploading it (mps_matmul_f32/_persistent) every step.
static MtlObject* __mps_sgd_update_pipeline = null;
int mps_sgd_update_f32_chained(struct MpsBuf bufW, struct MpsBuf bufGrad, float lr, unsigned long n) {
if (bufW.handle == null || bufGrad.handle == null) return 0;
if (!__mps_batch_active) return 0;
unsafe {
if (__mps_sgd_update_pipeline == null) {
MtlObject* pipeline = __mps_make_pipeline("sgd_update_kernel");
if (pipeline == null) return 0;
__mps_sgd_update_pipeline = pipeline;
}
MMsgVoid1 msgVoid1 = (MMsgVoid1)objc_msgSend;
msgVoid1(__mps_batch_encoder, __sel_setComputePipelineState, __mps_sgd_update_pipeline);
MMsgSetBuffer setBufFn = (MMsgSetBuffer)objc_msgSend;
setBufFn(__mps_batch_encoder, __sel_setBuffer, bufW.handle, 0UL, 0UL);
setBufFn(__mps_batch_encoder, __sel_setBuffer, bufGrad.handle, 0UL, 1UL);
MMsgSetBytes setBytesFn = (MMsgSetBytes)objc_msgSend;
float lrVal = lr;
setBytesFn(__mps_batch_encoder, __sel_setBytes, (const unsigned char*)&lrVal, 4UL, 2UL);
struct MTLSize gridSize;
gridSize.width = n; gridSize.height = 1UL; gridSize.depth = 1UL;
unsigned long tgWidth = n < 64UL ? n : 64UL;
if (tgWidth == 0UL) tgWidth = 1UL;
struct MTLSize tgSize;
tgSize.width = tgWidth; tgSize.height = 1UL; tgSize.depth = 1UL;
MMsgDispatch dispatchFn = (MMsgDispatch)objc_msgSend;
dispatchFn(__mps_batch_encoder, __sel_dispatchThreads, gridSize, tgSize);
return 1;
}
}
// ── fp16 / bf16 tier (see gpu_mps.h) ─────────────────────────────────────────
static MtlObject* __mps_matmul_f16_pipeline = null;
static MtlObject* __mps_matmul_bf16_pipeline = null;
static MtlObject* __mps_relu_f16_pipeline = null;
static MtlObject* __mps_relu_bf16_pipeline = null;
int mps_matmul_f16(const &float a, const &float b, &float out,
unsigned long M, unsigned long K, unsigned long N) {
unsafe {
if (!__mps_ensure_device_queue()) return 0;
MtlObject* queue = __mps_queue;
MMsg0 msg0 = (MMsg0)objc_msgSend;
if (__mps_matmul_f16_pipeline == null) {
MtlObject* p = __mps_make_pipeline("matmul_kernel_f16");
if (p == null) return 0;
__mps_matmul_f16_pipeline = p;
}
unsigned long nA = M * K;
unsigned long nB = K * N;
unsigned long nOut = M * N;
unsigned short* aHalf = (unsigned short*)malloc(nA * 2UL);
unsigned short* bHalf = (unsigned short*)malloc(nB * 2UL);
unsigned short* outHalf = (unsigned short*)malloc(nOut * 2UL);
if (aHalf == null || bHalf == null || outHalf == null) {
if (aHalf != null) free((unsigned char*)aHalf);
if (bHalf != null) free((unsigned char*)bHalf);
if (outHalf != null) free((unsigned char*)outHalf);
return 0;
}
f32_to_fp16_bulk_n_(a, aHalf, nA);
f32_to_fp16_bulk_n_(b, bHalf, nB);
unsigned long bytesA = nA * 2UL;
unsigned long bytesB = nB * 2UL;
unsigned long bytesOut = nOut * 2UL;
MtlObject* bufA = __mps_buf_get_with_bytes((const unsigned char*)aHalf, bytesA);
MtlObject* bufB = __mps_buf_get_with_bytes((const unsigned char*)bHalf, bytesB);
MtlObject* bufOut = __mps_buf_get(bytesOut);
free((unsigned char*)aHalf);
free((unsigned char*)bHalf);
if (bufA == null || bufB == null || bufOut == null) {
free((unsigned char*)outHalf);
return 0;
}
unsigned int Mu = (unsigned int)M;
unsigned int Ku = (unsigned int)K;
unsigned int Nu = (unsigned int)N;
MtlObject* cmdBuf = msg0(queue, __sel_commandBuffer);
MtlObject* encoder = msg0(cmdBuf, __sel_computeCommandEncoder);
MMsgVoid1 msgVoid1 = (MMsgVoid1)objc_msgSend;
msgVoid1(encoder, __sel_setComputePipelineState, __mps_matmul_f16_pipeline);
MMsgSetBuffer setBufFn = (MMsgSetBuffer)objc_msgSend;
setBufFn(encoder, __sel_setBuffer, bufA, 0UL, 0UL);
setBufFn(encoder, __sel_setBuffer, bufB, 0UL, 1UL);
setBufFn(encoder, __sel_setBuffer, bufOut, 0UL, 2UL);
MMsgSetBytes setBytesFn = (MMsgSetBytes)objc_msgSend;
setBytesFn(encoder, __sel_setBytes, (const unsigned char*)&Mu, 4UL, 3UL);
setBytesFn(encoder, __sel_setBytes, (const unsigned char*)&Ku, 4UL, 4UL);
setBytesFn(encoder, __sel_setBytes, (const unsigned char*)&Nu, 4UL, 5UL);
unsigned long tgW = 16UL;
unsigned long tgH = 16UL;
struct MTLSize tgSize;
tgSize.width = tgW; tgSize.height = tgH; tgSize.depth = 1UL;
struct MTLSize numThreadgroups;
numThreadgroups.width = (N + tgW - 1UL) / tgW;
numThreadgroups.height = (M + tgH - 1UL) / tgH;
numThreadgroups.depth = 1UL;
MMsgDispatch dispatchFn = (MMsgDispatch)objc_msgSend;
dispatchFn(encoder, __sel_dispatchThreadgroups, numThreadgroups, tgSize);
MMsgVoid0 msgVoid0 = (MMsgVoid0)objc_msgSend;
msgVoid0(encoder, __sel_endEncoding);
msgVoid0(cmdBuf, __sel_commit);
msgVoid0(cmdBuf, __sel_waitUntilCompleted);
unsigned char* outPtr = __mps_contents(bufOut);
if (outPtr == null) {
free((unsigned char*)outHalf);
__mps_buf_put(bufA, bytesA);
__mps_buf_put(bufB, bytesB);
__mps_buf_put(bufOut, bytesOut);
return 0;
}
memcpy((unsigned char*)outHalf, outPtr, bytesOut);
fp16_to_f32_bulk_n_(outHalf, out, nOut);
free((unsigned char*)outHalf);
__mps_buf_put(bufA, bytesA);
__mps_buf_put(bufB, bytesB);
__mps_buf_put(bufOut, bytesOut);
return 1;
}
}
int mps_matmul_bf16(const &float a, const &float b, &float out,
unsigned long M, unsigned long K, unsigned long N) {
unsafe {
if (!__mps_ensure_device_queue()) return 0;
MtlObject* queue = __mps_queue;
MMsg0 msg0 = (MMsg0)objc_msgSend;
if (__mps_matmul_bf16_pipeline == null) {
MtlObject* p = __mps_make_pipeline("matmul_kernel_bf16");
if (p == null) return 0;
__mps_matmul_bf16_pipeline = p;
}
unsigned long nA = M * K;
unsigned long nB = K * N;
unsigned long nOut = M * N;
unsigned short* aBf = (unsigned short*)malloc(nA * 2UL);
unsigned short* bBf = (unsigned short*)malloc(nB * 2UL);
unsigned short* outBf = (unsigned short*)malloc(nOut * 2UL);
if (aBf == null || bBf == null || outBf == null) {
if (aBf != null) free((unsigned char*)aBf);
if (bBf != null) free((unsigned char*)bBf);
if (outBf != null) free((unsigned char*)outBf);
return 0;
}
f32_to_bf16_bulk_n_(a, aBf, nA);
f32_to_bf16_bulk_n_(b, bBf, nB);
unsigned long bytesA = nA * 2UL;
unsigned long bytesB = nB * 2UL;
unsigned long bytesOut = nOut * 2UL;
MtlObject* bufA = __mps_buf_get_with_bytes((const unsigned char*)aBf, bytesA);
MtlObject* bufB = __mps_buf_get_with_bytes((const unsigned char*)bBf, bytesB);
MtlObject* bufOut = __mps_buf_get(bytesOut);
free((unsigned char*)aBf);
free((unsigned char*)bBf);
if (bufA == null || bufB == null || bufOut == null) {
free((unsigned char*)outBf);
return 0;
}
unsigned int Mu = (unsigned int)M;
unsigned int Ku = (unsigned int)K;
unsigned int Nu = (unsigned int)N;
MtlObject* cmdBuf = msg0(queue, __sel_commandBuffer);
MtlObject* encoder = msg0(cmdBuf, __sel_computeCommandEncoder);
MMsgVoid1 msgVoid1 = (MMsgVoid1)objc_msgSend;
msgVoid1(encoder, __sel_setComputePipelineState, __mps_matmul_bf16_pipeline);
MMsgSetBuffer setBufFn = (MMsgSetBuffer)objc_msgSend;
setBufFn(encoder, __sel_setBuffer, bufA, 0UL, 0UL);
setBufFn(encoder, __sel_setBuffer, bufB, 0UL, 1UL);
setBufFn(encoder, __sel_setBuffer, bufOut, 0UL, 2UL);
MMsgSetBytes setBytesFn = (MMsgSetBytes)objc_msgSend;
setBytesFn(encoder, __sel_setBytes, (const unsigned char*)&Mu, 4UL, 3UL);
setBytesFn(encoder, __sel_setBytes, (const unsigned char*)&Ku, 4UL, 4UL);
setBytesFn(encoder, __sel_setBytes, (const unsigned char*)&Nu, 4UL, 5UL);
unsigned long tgW = 16UL;
unsigned long tgH = 16UL;
struct MTLSize tgSize;
tgSize.width = tgW; tgSize.height = tgH; tgSize.depth = 1UL;
struct MTLSize numThreadgroups;
numThreadgroups.width = (N + tgW - 1UL) / tgW;
numThreadgroups.height = (M + tgH - 1UL) / tgH;
numThreadgroups.depth = 1UL;
MMsgDispatch dispatchFn = (MMsgDispatch)objc_msgSend;
dispatchFn(encoder, __sel_dispatchThreadgroups, numThreadgroups, tgSize);
MMsgVoid0 msgVoid0 = (MMsgVoid0)objc_msgSend;
msgVoid0(encoder, __sel_endEncoding);
msgVoid0(cmdBuf, __sel_commit);
msgVoid0(cmdBuf, __sel_waitUntilCompleted);
unsigned char* outPtr = __mps_contents(bufOut);
if (outPtr == null) {
free((unsigned char*)outBf);
__mps_buf_put(bufA, bytesA);
__mps_buf_put(bufB, bytesB);
__mps_buf_put(bufOut, bytesOut);
return 0;
}
memcpy((unsigned char*)outBf, outPtr, bytesOut);
bf16_to_f32_bulk_n_(outBf, out, nOut);
free((unsigned char*)outBf);
__mps_buf_put(bufA, bytesA);
__mps_buf_put(bufB, bytesB);
__mps_buf_put(bufOut, bytesOut);
return 1;
}
}
// Shared elementwise-relu dispatch body for mps_relu_f16/mps_relu_bf16 --
// only the pipeline name and the conversion functions differ, so this
// takes both as parameters rather than duplicating the whole
// upload/dispatch/readback sequence a third and fourth time.
typedef fn void(const float*, unsigned short*, unsigned long) __MpsToReducedFn;
typedef fn void(const unsigned short*, float*, unsigned long) __MpsToFloatFn;
static int __mps_relu_reduced_precision(const char* kernelName, struct MtlObject** pipelineSlot,
const float* a, float* out, unsigned long n,
__MpsToReducedFn toReduced, __MpsToFloatFn toFloat) {
unsafe {
if (!__mps_ensure_device_queue()) return 0;
MtlObject* queue = __mps_queue;
MMsg0 msg0 = (MMsg0)objc_msgSend;
if (*pipelineSlot == null) {
MtlObject* p = __mps_make_pipeline(kernelName);
if (p == null) return 0;
*pipelineSlot = p;
}
unsigned short* aRed = (unsigned short*)malloc(n * 2UL);
unsigned short* outRed = (unsigned short*)malloc(n * 2UL);
if (aRed == null || outRed == null) {
if (aRed != null) free((unsigned char*)aRed);
if (outRed != null) free((unsigned char*)outRed);
return 0;
}
toReduced(a, aRed, n);
unsigned long bytes = n * 2UL;
MtlObject* bufA = __mps_buf_get_with_bytes((const unsigned char*)aRed, bytes);
MtlObject* bufOut = __mps_buf_get(bytes);
free((unsigned char*)aRed);
if (bufA == null || bufOut == null) {
free((unsigned char*)outRed);
return 0;
}
MtlObject* cmdBuf = msg0(queue, __sel_commandBuffer);
MtlObject* encoder = msg0(cmdBuf, __sel_computeCommandEncoder);
MMsgVoid1 msgVoid1 = (MMsgVoid1)objc_msgSend;
msgVoid1(encoder, __sel_setComputePipelineState, *pipelineSlot);
MMsgSetBuffer setBufFn = (MMsgSetBuffer)objc_msgSend;
setBufFn(encoder, __sel_setBuffer, bufA, 0UL, 0UL);
setBufFn(encoder, __sel_setBuffer, bufOut, 0UL, 1UL);
MMsgSetBytes setBytesFn = (MMsgSetBytes)objc_msgSend;
unsigned int nParam = (unsigned int)n;
setBytesFn(encoder, __sel_setBytes, (const unsigned char*)&nParam, 4UL, 2UL);
struct MTLSize gridSize;
gridSize.width = n; gridSize.height = 1UL; gridSize.depth = 1UL;
unsigned long tgWidth = n < 256UL ? n : 256UL;
if (tgWidth == 0UL) tgWidth = 1UL;
struct MTLSize tgSize;
tgSize.width = tgWidth; tgSize.height = 1UL; tgSize.depth = 1UL;
struct MTLSize numThreadgroups;
numThreadgroups.width = (n + tgWidth - 1UL) / tgWidth;
numThreadgroups.height = 1UL; numThreadgroups.depth = 1UL;
MMsgDispatch dispatchFn = (MMsgDispatch)objc_msgSend;
dispatchFn(encoder, __sel_dispatchThreadgroups, numThreadgroups, tgSize);
MMsgVoid0 msgVoid0 = (MMsgVoid0)objc_msgSend;
msgVoid0(encoder, __sel_endEncoding);
msgVoid0(cmdBuf, __sel_commit);
msgVoid0(cmdBuf, __sel_waitUntilCompleted);
unsigned char* outPtr = __mps_contents(bufOut);
if (outPtr == null) {
free((unsigned char*)outRed);
__mps_buf_put(bufA, bytes);
__mps_buf_put(bufOut, bytes);
return 0;
}
memcpy((unsigned char*)outRed, outPtr, bytes);
toFloat(outRed, out, n);
free((unsigned char*)outRed);
__mps_buf_put(bufA, bytes);
__mps_buf_put(bufOut, bytes);
return 1;
}
}
int mps_relu_f16_n_(const &float a, &float out, unsigned long n) {
return __mps_relu_reduced_precision("relu_kernel_f16", &__mps_relu_f16_pipeline, a, out, n,
f32_to_fp16_bulk_n_, fp16_to_f32_bulk_n_);
}
int mps_relu_f16([]const float a_s, []float out_s) {
unsigned long n = out_s.len;
if (a_s.len < n) { return 0; }
float a_d_[1];
const &float a = &a_d_[0];
if (a_s.len > 0UL) { unsafe { a = (const &float)a_s.ptr; } }
float out_d_[1];
&float out = &out_d_[0];
if (out_s.len > 0UL) { unsafe { out = (&float)out_s.ptr; } }
return mps_relu_f16_n_(a, out, n);
}
int mps_relu_bf16_n_(const &float a, &float out, unsigned long n) {
return __mps_relu_reduced_precision("relu_kernel_bf16", &__mps_relu_bf16_pipeline, a, out, n,
f32_to_bf16_bulk_n_, bf16_to_f32_bulk_n_);
}
int mps_relu_bf16([]const float a_s, []float out_s) {
unsigned long n = out_s.len;
if (a_s.len < n) { return 0; }
float a_d_[1];
const &float a = &a_d_[0];
if (a_s.len > 0UL) { unsafe { a = (const &float)a_s.ptr; } }
float out_d_[1];
&float out = &out_d_[0];
if (out_s.len > 0UL) { unsafe { out = (&float)out_s.ptr; } }
return mps_relu_bf16_n_(a, out, n);
}
// ── fp16 / bf16 persistent-buffer / batched tier (see gpu_mps.h) ────────────
struct MpsBuf MpsBuf::upload_f16_n_(const &float host, unsigned long n) {
unsafe {
unsigned short* tmp = (unsigned short*)malloc(n * 2UL);
if (tmp == null) return __mps_buf_make(null, 0UL, MPS_FP16);
f32_to_fp16_bulk_n_(host, tmp, n);
struct MpsBuf b = __mps_buf_upload((const unsigned char*)tmp, n * 2UL, MPS_FP16);
free((unsigned char*)tmp);
return b;
}
}
struct MpsBuf MpsBuf::upload_f16([]const float host_s) {
unsigned long n = host_s.len;
float host_d_[1];
const &float host = &host_d_[0];
if (host_s.len > 0UL) { unsafe { host = (const &float)host_s.ptr; } }
return MpsBuf::upload_f16_n_(host, n);
}
struct MpsBuf MpsBuf::upload_bf16_n_(const &float host, unsigned long n) {
unsafe {
unsigned short* tmp = (unsigned short*)malloc(n * 2UL);
if (tmp == null) return __mps_buf_make(null, 0UL, MPS_BF16);
f32_to_bf16_bulk_n_(host, tmp, n);
struct MpsBuf b = __mps_buf_upload((const unsigned char*)tmp, n * 2UL, MPS_BF16);
free((unsigned char*)tmp);
return b;
}
}
struct MpsBuf MpsBuf::upload_bf16([]const float host_s) {
unsigned long n = host_s.len;
float host_d_[1];
const &float host = &host_d_[0];
if (host_s.len > 0UL) { unsafe { host = (const &float)host_s.ptr; } }
return MpsBuf::upload_bf16_n_(host, n);
}
// Shared dispatch body for every matmul-shaped fp16/bf16 chained op
// (matmul_kernel_f16/bf16, matmul_atb_kernel_f16/bf16, matmul_abt_kernel_
// f16/bf16 all take the same (bufA, bufB, bufOut, uint, uint, uint)
// binding shape from the host side -- only the kernel name and what the
// three uint constants/output shape mean differ between the three ops).
static int __mps_matmul3_chained(const char* kernelName, struct MtlObject** pipelineSlot, int dtype,
MtlObject* bufA, MtlObject* bufB, unsigned short* out,
unsigned int p0, unsigned int p1, unsigned int p2,
unsigned long outRows, unsigned long outCols) {
if (!__mps_batch_active) return 0;
unsafe {
if (*pipelineSlot == null) {
MtlObject* p = __mps_make_pipeline(kernelName);
if (p == null) return 0;
*pipelineSlot = p;
}
unsigned long bytesOut = outRows * outCols * 2UL;
MtlObject* bufOut = __mps_buf_get_batched(bytesOut);
if (bufOut == null) return 0;
MMsgVoid1 msgVoid1 = (MMsgVoid1)objc_msgSend;
msgVoid1(__mps_batch_encoder, __sel_setComputePipelineState, *pipelineSlot);
MMsgSetBuffer setBufFn = (MMsgSetBuffer)objc_msgSend;
setBufFn(__mps_batch_encoder, __sel_setBuffer, bufA, 0UL, 0UL);
setBufFn(__mps_batch_encoder, __sel_setBuffer, bufB, 0UL, 1UL);
setBufFn(__mps_batch_encoder, __sel_setBuffer, bufOut, 0UL, 2UL);
MMsgSetBytes setBytesFn = (MMsgSetBytes)objc_msgSend;
setBytesFn(__mps_batch_encoder, __sel_setBytes, (const unsigned char*)&p0, 4UL, 3UL);
setBytesFn(__mps_batch_encoder, __sel_setBytes, (const unsigned char*)&p1, 4UL, 4UL);
setBytesFn(__mps_batch_encoder, __sel_setBytes, (const unsigned char*)&p2, 4UL, 5UL);
unsigned long tgW = 16UL;
unsigned long tgH = 16UL;
struct MTLSize tgSize;
tgSize.width = tgW; tgSize.height = tgH; tgSize.depth = 1UL;
struct MTLSize numThreadgroups;
numThreadgroups.width = (outCols + tgW - 1UL) / tgW;
numThreadgroups.height = (outRows + tgH - 1UL) / tgH;
numThreadgroups.depth = 1UL;
MMsgDispatch dispatchFn = (MMsgDispatch)objc_msgSend;
dispatchFn(__mps_batch_encoder, __sel_dispatchThreadgroups, numThreadgroups, tgSize);
if (__mps_pending_count < (unsigned long)MPS_BATCH_MAX_PENDING) {
__mps_pending_gpuBuf[__mps_pending_count] = bufOut;
__mps_pending_cpuOut[__mps_pending_count] = (unsigned char*)out;
__mps_pending_bytes[__mps_pending_count] = bytesOut;
__mps_pending_dtype[__mps_pending_count] = dtype;
__mps_pending_count = __mps_pending_count + 1UL;
}
return 1;
}
}
static MtlObject* __mps_matmul_f16_persistent_pipeline = null;
int mps_matmul_f16_persistent(struct MpsBuf bufA, struct MpsBuf bufB, &unsigned short out,
unsigned long M, unsigned long K, unsigned long N) {
if (bufA.handle == null || bufB.handle == null) return 0;
return __mps_matmul3_chained("matmul_kernel_f16", &__mps_matmul_f16_persistent_pipeline, MPS_FP16,
bufA.handle, bufB.handle, out, (unsigned int)M, (unsigned int)K, (unsigned int)N, M, N);
}
static MtlObject* __mps_matmul_bf16_persistent_pipeline = null;
int mps_matmul_bf16_persistent(struct MpsBuf bufA, struct MpsBuf bufB, &unsigned short out,
unsigned long M, unsigned long K, unsigned long N) {
if (bufA.handle == null || bufB.handle == null) return 0;
return __mps_matmul3_chained("matmul_kernel_bf16", &__mps_matmul_bf16_persistent_pipeline, MPS_BF16,
bufA.handle, bufB.handle, out, (unsigned int)M, (unsigned int)K, (unsigned int)N, M, N);
}
static MtlObject* __mps_matmul_atb_f16_pipeline = null;
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) {
if (bufA.handle == null || bufDC.handle == null) return 0;
return __mps_matmul3_chained("matmul_atb_kernel_f16", &__mps_matmul_atb_f16_pipeline, MPS_FP16,
bufA.handle, bufDC.handle, out, (unsigned int)M, (unsigned int)K, (unsigned int)N, K, N);
}
static MtlObject* __mps_matmul_atb_bf16_pipeline = null;
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) {
if (bufA.handle == null || bufDC.handle == null) return 0;
return __mps_matmul3_chained("matmul_atb_kernel_bf16", &__mps_matmul_atb_bf16_pipeline, MPS_BF16,
bufA.handle, bufDC.handle, out, (unsigned int)M, (unsigned int)K, (unsigned int)N, K, N);
}
static MtlObject* __mps_matmul_abt_f16_pipeline = null;
int mps_matmul_abt_f16_persistent(struct MpsBuf bufDC, struct MpsBuf bufB, &unsigned short out,
unsigned long M, unsigned long N, unsigned long K) {
if (bufDC.handle == null || bufB.handle == null) return 0;
return __mps_matmul3_chained("matmul_abt_kernel_f16", &__mps_matmul_abt_f16_pipeline, MPS_FP16,
bufDC.handle, bufB.handle, out, (unsigned int)M, (unsigned int)N, (unsigned int)K, M, K);
}
// M==1 GEMV path for mps_matmul_abt_bf16_persistent -- see matmul_abt_
// kernel_bf16_gemv's comment in gpu_mps_kernels.metal. Dispatch is 1D over
// K, one output per simdgroup, GEMV_SIMDGROUPS (must match the kernel's)
// per threadgroup (buffer(3)=N, buffer(4)=K -- no M buffer, since this
// path only exists for M==1).
#define GEMV_SIMDGROUPS ((unsigned long)4)
static MtlObject* __mps_matmul_abt_bf16_gemv_pipeline = null;
static int __mps_matmul_abt_bf16_gemv_chained(MtlObject* bufA, MtlObject* bufB, unsigned short* out,
unsigned int N, unsigned int K) {
if (!__mps_batch_active) return 0;
unsafe {
if (__mps_matmul_abt_bf16_gemv_pipeline == null) {
MtlObject* p = __mps_make_pipeline("matmul_abt_kernel_bf16_gemv");
if (p == null) return 0;
__mps_matmul_abt_bf16_gemv_pipeline = p;
}
unsigned long bytesOut = (unsigned long)K * 2UL;
MtlObject* bufOut = __mps_buf_get_batched(bytesOut);
if (bufOut == null) return 0;
MMsgVoid1 msgVoid1 = (MMsgVoid1)objc_msgSend;
msgVoid1(__mps_batch_encoder, __sel_setComputePipelineState, __mps_matmul_abt_bf16_gemv_pipeline);
MMsgSetBuffer setBufFn = (MMsgSetBuffer)objc_msgSend;
setBufFn(__mps_batch_encoder, __sel_setBuffer, bufA, 0UL, 0UL);
setBufFn(__mps_batch_encoder, __sel_setBuffer, bufB, 0UL, 1UL);
setBufFn(__mps_batch_encoder, __sel_setBuffer, bufOut, 0UL, 2UL);
MMsgSetBytes setBytesFn = (MMsgSetBytes)objc_msgSend;
setBytesFn(__mps_batch_encoder, __sel_setBytes, (const unsigned char*)&N, 4UL, 3UL);
setBytesFn(__mps_batch_encoder, __sel_setBytes, (const unsigned char*)&K, 4UL, 4UL);
struct MTLSize tgSize;
tgSize.width = GEMV_SIMDGROUPS * 32UL; tgSize.height = 1UL; tgSize.depth = 1UL;
struct MTLSize numThreadgroups;
numThreadgroups.width = ((unsigned long)K + GEMV_SIMDGROUPS - 1UL) / GEMV_SIMDGROUPS;
numThreadgroups.height = 1UL;
numThreadgroups.depth = 1UL;
MMsgDispatch dispatchFn = (MMsgDispatch)objc_msgSend;
dispatchFn(__mps_batch_encoder, __sel_dispatchThreadgroups, numThreadgroups, tgSize);
if (__mps_pending_count < (unsigned long)MPS_BATCH_MAX_PENDING) {
__mps_pending_gpuBuf[__mps_pending_count] = bufOut;
__mps_pending_cpuOut[__mps_pending_count] = (unsigned char*)out;
__mps_pending_bytes[__mps_pending_count] = bytesOut;
__mps_pending_dtype[__mps_pending_count] = MPS_BF16;
__mps_pending_count = __mps_pending_count + 1UL;
}
return 1;
}
}
static MtlObject* __mps_matmul_abt_bf16_pipeline = null;
int mps_matmul_abt_bf16_persistent(struct MpsBuf bufDC, struct MpsBuf bufB, &unsigned short out,
unsigned long M, unsigned long N, unsigned long K) {
if (bufDC.handle == null || bufB.handle == null) return 0;
// M==1 (a decode step's projections) routes to the dedicated GEMV
// kernel; the tiled kernel would compute and discard 15 of every 16
// rows. See gpu_mps_kernels.metal for why a simdgroup_matrix bf16 tier
// isn't used for M>1 either.
if (M == 1UL) {
int rc = __mps_matmul_abt_bf16_gemv_chained(bufDC.handle, bufB.handle, out, (unsigned int)N, (unsigned int)K);
if (rc) { return rc; }
}
return __mps_matmul3_chained("matmul_abt_kernel_bf16", &__mps_matmul_abt_bf16_pipeline, MPS_BF16,
bufDC.handle, bufB.handle, out, (unsigned int)M, (unsigned int)N, (unsigned int)K, M, K);
}
// Shared dispatch body for a single-device-input elementwise op (relu).
static int __mps_unary_chained(const char* kernelName, struct MtlObject** pipelineSlot, int dtype,
MtlObject* bufA, unsigned short* out, unsigned long n) {
if (!__mps_batch_active) return 0;
unsafe {
if (*pipelineSlot == null) {
MtlObject* p = __mps_make_pipeline(kernelName);
if (p == null) return 0;
*pipelineSlot = p;
}
unsigned long bytes = n * 2UL;
MtlObject* bufOut = __mps_buf_get_batched(bytes);
if (bufOut == null) return 0;
MMsgVoid1 msgVoid1 = (MMsgVoid1)objc_msgSend;
msgVoid1(__mps_batch_encoder, __sel_setComputePipelineState, *pipelineSlot);
MMsgSetBuffer setBufFn = (MMsgSetBuffer)objc_msgSend;
setBufFn(__mps_batch_encoder, __sel_setBuffer, bufA, 0UL, 0UL);
setBufFn(__mps_batch_encoder, __sel_setBuffer, bufOut, 0UL, 1UL);
struct MTLSize gridSize;
gridSize.width = n; gridSize.height = 1UL; gridSize.depth = 1UL;
unsigned long tgWidth = n < 64UL ? n : 64UL;
if (tgWidth == 0UL) tgWidth = 1UL;
struct MTLSize tgSize;
tgSize.width = tgWidth; tgSize.height = 1UL; tgSize.depth = 1UL;
MMsgDispatch dispatchFn = (MMsgDispatch)objc_msgSend;
dispatchFn(__mps_batch_encoder, __sel_dispatchThreads, gridSize, tgSize);
if (__mps_pending_count < (unsigned long)MPS_BATCH_MAX_PENDING) {
__mps_pending_gpuBuf[__mps_pending_count] = bufOut;
__mps_pending_cpuOut[__mps_pending_count] = (unsigned char*)out;
__mps_pending_bytes[__mps_pending_count] = bytes;
__mps_pending_dtype[__mps_pending_count] = dtype;
__mps_pending_count = __mps_pending_count + 1UL;
}
return 1;
}
}
static MtlObject* __mps_relu_f16_chained_pipeline = null;
int mps_relu_f16_chained(struct MpsBuf bufA, &unsigned short out, unsigned long n) {
if (bufA.handle == null) return 0;
return __mps_unary_chained("relu_kernel_f16", &__mps_relu_f16_chained_pipeline, MPS_FP16, bufA.handle, out, n);
}
static MtlObject* __mps_relu_bf16_chained_pipeline = null;
int mps_relu_bf16_chained(struct MpsBuf bufA, &unsigned short out, unsigned long n) {
if (bufA.handle == null) return 0;
return __mps_unary_chained("relu_kernel_bf16", &__mps_relu_bf16_chained_pipeline, MPS_BF16, bufA.handle, out, n);
}
// Shared dispatch body for a two-device-input elementwise op (relu
// backward: dZ1 = relu'(Z1) * dH, both operands already GPU-pending).
static int __mps_binary_dev_chained(const char* kernelName, struct MtlObject** pipelineSlot, int dtype,
MtlObject* bufA, MtlObject* bufB, unsigned short* out, unsigned long n) {
if (!__mps_batch_active) return 0;
unsafe {
if (*pipelineSlot == null) {
MtlObject* p = __mps_make_pipeline(kernelName);
if (p == null) return 0;
*pipelineSlot = p;
}
unsigned long bytes = n * 2UL;
MtlObject* bufOut = __mps_buf_get_batched(bytes);
if (bufOut == null) return 0;
MMsgVoid1 msgVoid1 = (MMsgVoid1)objc_msgSend;
msgVoid1(__mps_batch_encoder, __sel_setComputePipelineState, *pipelineSlot);
MMsgSetBuffer setBufFn = (MMsgSetBuffer)objc_msgSend;
setBufFn(__mps_batch_encoder, __sel_setBuffer, bufA, 0UL, 0UL);
setBufFn(__mps_batch_encoder, __sel_setBuffer, bufB, 0UL, 1UL);
setBufFn(__mps_batch_encoder, __sel_setBuffer, bufOut, 0UL, 2UL);
struct MTLSize gridSize;
gridSize.width = n; gridSize.height = 1UL; gridSize.depth = 1UL;
unsigned long tgWidth = n < 64UL ? n : 64UL;
if (tgWidth == 0UL) tgWidth = 1UL;
struct MTLSize tgSize;
tgSize.width = tgWidth; tgSize.height = 1UL; tgSize.depth = 1UL;
MMsgDispatch dispatchFn = (MMsgDispatch)objc_msgSend;
dispatchFn(__mps_batch_encoder, __sel_dispatchThreads, gridSize, tgSize);
if (__mps_pending_count < (unsigned long)MPS_BATCH_MAX_PENDING) {
__mps_pending_gpuBuf[__mps_pending_count] = bufOut;
__mps_pending_cpuOut[__mps_pending_count] = (unsigned char*)out;
__mps_pending_bytes[__mps_pending_count] = bytes;
__mps_pending_dtype[__mps_pending_count] = dtype;
__mps_pending_count = __mps_pending_count + 1UL;
}
return 1;
}
}
static MtlObject* __mps_relu_backward_f16_pipeline = null;
int mps_relu_backward_f16_chained_ab(struct MpsBuf bufA, struct MpsBuf bufSelfGrad, &unsigned short out, unsigned long n) {
if (bufA.handle == null || bufSelfGrad.handle == null) return 0;
return __mps_binary_dev_chained("relu_backward_kernel_f16", &__mps_relu_backward_f16_pipeline, MPS_FP16, bufA.handle, bufSelfGrad.handle, out, n);
}
static MtlObject* __mps_relu_backward_bf16_pipeline = null;
int mps_relu_backward_bf16_chained_ab(struct MpsBuf bufA, struct MpsBuf bufSelfGrad, &unsigned short out, unsigned long n) {
if (bufA.handle == null || bufSelfGrad.handle == null) return 0;
return __mps_binary_dev_chained("relu_backward_kernel_bf16", &__mps_relu_backward_bf16_pipeline, MPS_BF16, bufA.handle, bufSelfGrad.handle, out, n);
}
// out[i] = bufA[i] - bHost[i] -- 'bHost' is a plain CPU array of raw
// fp16/bf16 bits, uploaded fresh each call (see gpu_mps.h's comment).
static int __mps_sub_chained(const char* kernelName, struct MtlObject** pipelineSlot, int dtype,
MtlObject* bufA, const unsigned short* bHost, unsigned short* out, unsigned long n) {
if (!__mps_batch_active) return 0;
unsafe {
if (*pipelineSlot == null) {
MtlObject* p = __mps_make_pipeline(kernelName);
if (p == null) return 0;
*pipelineSlot = p;
}
unsigned long bytes = n * 2UL;
MtlObject* bufB = __mps_buf_get_with_bytes_batched((const unsigned char*)bHost, bytes);
MtlObject* bufOut = __mps_buf_get_batched(bytes);
if (bufB == null || bufOut == null) return 0;
MMsgVoid1 msgVoid1 = (MMsgVoid1)objc_msgSend;
msgVoid1(__mps_batch_encoder, __sel_setComputePipelineState, *pipelineSlot);
MMsgSetBuffer setBufFn = (MMsgSetBuffer)objc_msgSend;
setBufFn(__mps_batch_encoder, __sel_setBuffer, bufA, 0UL, 0UL);
setBufFn(__mps_batch_encoder, __sel_setBuffer, bufB, 0UL, 1UL);
setBufFn(__mps_batch_encoder, __sel_setBuffer, bufOut, 0UL, 2UL);
struct MTLSize gridSize;
gridSize.width = n; gridSize.height = 1UL; gridSize.depth = 1UL;
unsigned long tgWidth = n < 64UL ? n : 64UL;
if (tgWidth == 0UL) tgWidth = 1UL;
struct MTLSize tgSize;
tgSize.width = tgWidth; tgSize.height = 1UL; tgSize.depth = 1UL;
MMsgDispatch dispatchFn = (MMsgDispatch)objc_msgSend;
dispatchFn(__mps_batch_encoder, __sel_dispatchThreads, gridSize, tgSize);
if (__mps_pending_count < (unsigned long)MPS_BATCH_MAX_PENDING) {
__mps_pending_gpuBuf[__mps_pending_count] = bufOut;
__mps_pending_cpuOut[__mps_pending_count] = (unsigned char*)out;
__mps_pending_bytes[__mps_pending_count] = bytes;
__mps_pending_dtype[__mps_pending_count] = dtype;
__mps_pending_count = __mps_pending_count + 1UL;
}
return 1;
}
}
static MtlObject* __mps_sub_f16_pipeline = null;
int mps_sub_f16_chained(struct MpsBuf bufA, const &unsigned short bHalf, &unsigned short out, unsigned long n) {
if (bufA.handle == null) return 0;
return __mps_sub_chained("sub_kernel_f16", &__mps_sub_f16_pipeline, MPS_FP16, bufA.handle, bHalf, out, n);
}
static MtlObject* __mps_sub_bf16_pipeline = null;
int mps_sub_bf16_chained(struct MpsBuf bufA, const &unsigned short bBf, &unsigned short out, unsigned long n) {
if (bufA.handle == null) return 0;
return __mps_sub_chained("sub_kernel_bf16", &__mps_sub_bf16_pipeline, MPS_BF16, bufA.handle, bBf, out, n);
}
// out[i] = buf[i] * k -- 'k' stays a plain float scalar (see gpu_mps_
// kernels.metal's scale_kernel_f16/bf16 comment).
static int __mps_scale_chained(const char* kernelName, struct MtlObject** pipelineSlot, int dtype,
MtlObject* buf, float k, unsigned short* out, unsigned long n) {
if (!__mps_batch_active) return 0;
unsafe {
if (*pipelineSlot == null) {
MtlObject* p = __mps_make_pipeline(kernelName);
if (p == null) return 0;
*pipelineSlot = p;
}
unsigned long bytes = n * 2UL;
MtlObject* bufOut = __mps_buf_get_batched(bytes);
if (bufOut == null) return 0;
MMsgVoid1 msgVoid1 = (MMsgVoid1)objc_msgSend;
msgVoid1(__mps_batch_encoder, __sel_setComputePipelineState, *pipelineSlot);
MMsgSetBuffer setBufFn = (MMsgSetBuffer)objc_msgSend;
setBufFn(__mps_batch_encoder, __sel_setBuffer, buf, 0UL, 0UL);
MMsgSetBytes setBytesFn = (MMsgSetBytes)objc_msgSend;
float kVal = k;
setBytesFn(__mps_batch_encoder, __sel_setBytes, (const unsigned char*)&kVal, 4UL, 1UL);
setBufFn(__mps_batch_encoder, __sel_setBuffer, bufOut, 0UL, 2UL);
struct MTLSize gridSize;
gridSize.width = n; gridSize.height = 1UL; gridSize.depth = 1UL;
unsigned long tgWidth = n < 64UL ? n : 64UL;
if (tgWidth == 0UL) tgWidth = 1UL;
struct MTLSize tgSize;
tgSize.width = tgWidth; tgSize.height = 1UL; tgSize.depth = 1UL;
MMsgDispatch dispatchFn = (MMsgDispatch)objc_msgSend;
dispatchFn(__mps_batch_encoder, __sel_dispatchThreads, gridSize, tgSize);
if (__mps_pending_count < (unsigned long)MPS_BATCH_MAX_PENDING) {
__mps_pending_gpuBuf[__mps_pending_count] = bufOut;
__mps_pending_cpuOut[__mps_pending_count] = (unsigned char*)out;
__mps_pending_bytes[__mps_pending_count] = bytes;
__mps_pending_dtype[__mps_pending_count] = dtype;
__mps_pending_count = __mps_pending_count + 1UL;
}
return 1;
}
}
static MtlObject* __mps_scale_f16_pipeline = null;
int mps_scale_f16_chained(struct MpsBuf buf, float k, &unsigned short out, unsigned long n) {
if (buf.handle == null) return 0;
return __mps_scale_chained("scale_kernel_f16", &__mps_scale_f16_pipeline, MPS_FP16, buf.handle, k, out, n);
}
static MtlObject* __mps_scale_bf16_pipeline = null;
int mps_scale_bf16_chained(struct MpsBuf buf, float k, &unsigned short out, unsigned long n) {
if (buf.handle == null) return 0;
return __mps_scale_chained("scale_kernel_bf16", &__mps_scale_bf16_pipeline, MPS_BF16, buf.handle, k, out, n);
}
// Both operands device buffers -- same binding shape as relu-backward's
// __mps_binary_dev_chained above.
static MtlObject* __mps_add_bf16_pipeline = null;
int mps_add_bf16_chained(struct MpsBuf bufA, struct MpsBuf bufB, &unsigned short out, unsigned long n) {
if (bufA.handle == null || bufB.handle == null) return 0;
return __mps_binary_dev_chained("add_kernel_bf16", &__mps_add_bf16_pipeline, MPS_BF16, bufA.handle, bufB.handle, out, n);
}
static MtlObject* __mps_gelu_mul_bf16_pipeline = null;
int mps_gelu_mul_bf16_chained(struct MpsBuf bufA, struct MpsBuf bufB, &unsigned short out, unsigned long n) {
if (bufA.handle == null || bufB.handle == null) return 0;
return __mps_binary_dev_chained("gelu_mul_kernel_bf16", &__mps_gelu_mul_bf16_pipeline, MPS_BF16, bufA.handle, bufB.handle, out, n);
}
// Row-wise RMSNorm -- see rmsnorm_kernel_bf16's comment in gpu_mps_
// kernels.metal. One RMSNORM_TG-thread threadgroup per row; RMSNORM_TG
// must match the kernel's.
#define RMSNORM_TG ((unsigned long)256)
static MtlObject* __mps_rmsnorm_bf16_pipeline = null;
int mps_rmsnorm_bf16_chained(struct MpsBuf bufX, struct MpsBuf weightBuf, &unsigned short out,
unsigned long rows, unsigned long cols, float eps) {
if (bufX.handle == null || weightBuf.handle == null) return 0;
if (!__mps_batch_active) return 0;
unsafe {
if (__mps_rmsnorm_bf16_pipeline == null) {
MtlObject* p = __mps_make_pipeline("rmsnorm_kernel_bf16");
if (p == null) return 0;
__mps_rmsnorm_bf16_pipeline = p;
}
unsigned long bytes = rows * cols * 2UL;
MtlObject* bufOut = __mps_buf_get_batched(bytes);
if (bufOut == null) return 0;
MMsgVoid1 msgVoid1 = (MMsgVoid1)objc_msgSend;
msgVoid1(__mps_batch_encoder, __sel_setComputePipelineState, __mps_rmsnorm_bf16_pipeline);
MMsgSetBuffer setBufFn = (MMsgSetBuffer)objc_msgSend;
setBufFn(__mps_batch_encoder, __sel_setBuffer, bufX.handle, 0UL, 0UL);
setBufFn(__mps_batch_encoder, __sel_setBuffer, weightBuf.handle, 0UL, 1UL);
setBufFn(__mps_batch_encoder, __sel_setBuffer, bufOut, 0UL, 2UL);
MMsgSetBytes setBytesFn = (MMsgSetBytes)objc_msgSend;
unsigned int colsVal = (unsigned int)cols;
setBytesFn(__mps_batch_encoder, __sel_setBytes, (const unsigned char*)&colsVal, 4UL, 3UL);
float epsVal = eps;
setBytesFn(__mps_batch_encoder, __sel_setBytes, (const unsigned char*)&epsVal, 4UL, 4UL);
struct MTLSize tgSize;
tgSize.width = RMSNORM_TG; tgSize.height = 1UL; tgSize.depth = 1UL;
struct MTLSize numThreadgroups;
numThreadgroups.width = rows; numThreadgroups.height = 1UL; numThreadgroups.depth = 1UL;
MMsgDispatch dispatchFn = (MMsgDispatch)objc_msgSend;
dispatchFn(__mps_batch_encoder, __sel_dispatchThreadgroups, numThreadgroups, tgSize);
if (__mps_pending_count < (unsigned long)MPS_BATCH_MAX_PENDING) {
__mps_pending_gpuBuf[__mps_pending_count] = bufOut;
__mps_pending_cpuOut[__mps_pending_count] = (unsigned char*)out;
__mps_pending_bytes[__mps_pending_count] = bytes;
__mps_pending_dtype[__mps_pending_count] = MPS_BF16;
__mps_pending_count = __mps_pending_count + 1UL;
}
return 1;
}
}
// In-place SGD step against a persistent fp16/bf16 device buffer -- no
// 'out'/pending-readback registration, result stays on the GPU (see
// mps_sgd_update_f32_chained's comment).
static int __mps_sgd_chained(const char* kernelName, struct MtlObject** pipelineSlot,
MtlObject* bufW, MtlObject* bufGrad, float lr, unsigned long n) {
if (!__mps_batch_active) return 0;
unsafe {
if (*pipelineSlot == null) {
MtlObject* p = __mps_make_pipeline(kernelName);
if (p == null) return 0;
*pipelineSlot = p;
}
MMsgVoid1 msgVoid1 = (MMsgVoid1)objc_msgSend;
msgVoid1(__mps_batch_encoder, __sel_setComputePipelineState, *pipelineSlot);
MMsgSetBuffer setBufFn = (MMsgSetBuffer)objc_msgSend;
setBufFn(__mps_batch_encoder, __sel_setBuffer, bufW, 0UL, 0UL);
setBufFn(__mps_batch_encoder, __sel_setBuffer, bufGrad, 0UL, 1UL);
MMsgSetBytes setBytesFn = (MMsgSetBytes)objc_msgSend;
float lrVal = lr;
setBytesFn(__mps_batch_encoder, __sel_setBytes, (const unsigned char*)&lrVal, 4UL, 2UL);
struct MTLSize gridSize;
gridSize.width = n; gridSize.height = 1UL; gridSize.depth = 1UL;
unsigned long tgWidth = n < 64UL ? n : 64UL;
if (tgWidth == 0UL) tgWidth = 1UL;
struct MTLSize tgSize;
tgSize.width = tgWidth; tgSize.height = 1UL; tgSize.depth = 1UL;
MMsgDispatch dispatchFn = (MMsgDispatch)objc_msgSend;
dispatchFn(__mps_batch_encoder, __sel_dispatchThreads, gridSize, tgSize);
return 1;
}
}
static MtlObject* __mps_sgd_update_f16_pipeline = null;
int mps_sgd_update_f16_chained(struct MpsBuf bufW, struct MpsBuf bufGrad, float lr, unsigned long n) {
if (bufW.handle == null || bufGrad.handle == null) return 0;
return __mps_sgd_chained("sgd_update_kernel_f16", &__mps_sgd_update_f16_pipeline, bufW.handle, bufGrad.handle, lr, n);
}
static MtlObject* __mps_sgd_update_bf16_pipeline = null;
int mps_sgd_update_bf16_chained(struct MpsBuf bufW, struct MpsBuf bufGrad, float lr, unsigned long n) {
if (bufW.handle == null || bufGrad.handle == null) return 0;
return __mps_sgd_chained("sgd_update_kernel_bf16", &__mps_sgd_update_bf16_pipeline, bufW.handle, bufGrad.handle, lr, n);
}
// ── Batched multi-head attention (bf16) ─────────────────────────────────────
// See attn_batch_bf16_kernel's comment in gpu_mps_kernels.metal. The kernel
// is instantiated per head dim (attn_batch_bf16_kernel_d<headDim>, one
// cached pipeline each) and runs one threadgroup of ATTN_SIMDGROUPS(D)
// simdgroups per (batch, head) -- __mps_attn_batch_simdgroups mirrors the
// kernel's formula. Other
// head dims, or a head count that isn't a whole multiple of the KV head
// count, are refused up front so callers fall back to their CPU path.
static unsigned long __mps_attn_batch_simdgroups(unsigned long headDim) {
unsigned long sg = 4096UL / headDim;
return sg < 32UL ? sg : 32UL;
}
// The attention kernels below (decode, flash, linear) are all templated
// on head dim and instantiated for 64/128/256/512, as
// '<kernel>_d<headDim>'; each op keeps one cached pipeline per head dim.
// Slot for 'headDim', or -1 if no instantiation exists.
static long __mps_head_dim_slot(unsigned long headDim) {
if (headDim == 64UL) return 0L;
if (headDim == 128UL) return 1L;
if (headDim == 256UL) return 2L;
if (headDim == 512UL) return 3L;
return -1L;
}
static int __mps_attn_shape_ok(unsigned long numHeads, unsigned long numKvHeads, unsigned long headDim) {
return __mps_head_dim_slot(headDim) >= 0L && numKvHeads != 0UL && numHeads % numKvHeads == 0UL;
}
// Cached pipeline for a head-dim-templated kernel, given its four
// instantiations' names. Callers have already checked __mps_attn_shape_ok.
static MtlObject* __mps_head_dim_pipeline(struct MtlObject** slots, unsigned long headDim, const char* name64,
const char* name128, const char* name256, const char* name512) {
unsafe {
long slot = __mps_head_dim_slot(headDim);
if (slots[slot] == null) {
const char* name = name64;
if (slot == 1L) name = name128;
if (slot == 2L) name = name256;
if (slot == 3L) name = name512;
slots[slot] = __mps_make_pipeline(name);
}
return slots[slot];
}
}
// Binds 'count' 4-byte constants to consecutive buffer indices starting at
// 'firstIndex'.
static void __mps_set_uints(MtlObject* encoder, unsigned int* params, unsigned long count, unsigned long firstIndex) {
unsafe {
MMsgSetBytes setBytesFn = (MMsgSetBytes)objc_msgSend;
unsigned long i = 0UL;
while (i < count) {
setBytesFn(encoder, __sel_setBytes, (const unsigned char*)¶ms[i], 4UL, firstIndex + i);
i = i + 1UL;
}
}
}
// Registers 'bufOut' as a chained op's pending readback into 'out' -- the
// same bookkeeping every *_chained op above does inline.
static void __mps_pending_add(MtlObject* bufOut, unsigned char* out, unsigned long bytes, int dtype) {
if (__mps_pending_count < (unsigned long)MPS_BATCH_MAX_PENDING) {
__mps_pending_gpuBuf[__mps_pending_count] = bufOut;
__mps_pending_cpuOut[__mps_pending_count] = out;
__mps_pending_bytes[__mps_pending_count] = bytes;
__mps_pending_dtype[__mps_pending_count] = dtype;
__mps_pending_count = __mps_pending_count + 1UL;
}
}
static MtlObject* __mps_attn_batch_pipelines[4];
// Binds the pipeline, buffers 0-4 and the five uint constants, then
// dispatches -- shared by the one-shot and chained entry points, which
// differ only in which encoder they target and how 'bufOut' is read back.
// Callers have already checked __mps_attn_shape_ok.
static int __mps_attn_batch_encode(MtlObject* encoder, MtlObject* bufQ, MtlObject* bufK, MtlObject* bufV, MtlObject* bufCacheLens,
MtlObject* bufOut, unsigned long batchSize, unsigned long numHeads,
unsigned long numKvHeads, unsigned long headDim, unsigned long maxKvLen,
int isSliding, unsigned long slidingWindow) {
unsafe {
MtlObject* pipeline = __mps_head_dim_pipeline(__mps_attn_batch_pipelines, headDim,
"attn_batch_bf16_kernel_d64", "attn_batch_bf16_kernel_d128",
"attn_batch_bf16_kernel_d256", "attn_batch_bf16_kernel_d512");
if (pipeline == null) return 0;
MMsgVoid1 msgVoid1 = (MMsgVoid1)objc_msgSend;
msgVoid1(encoder, __sel_setComputePipelineState, pipeline);
MMsgSetBuffer setBufFn = (MMsgSetBuffer)objc_msgSend;
setBufFn(encoder, __sel_setBuffer, bufQ, 0UL, 0UL);
setBufFn(encoder, __sel_setBuffer, bufK, 0UL, 1UL);
setBufFn(encoder, __sel_setBuffer, bufV, 0UL, 2UL);
setBufFn(encoder, __sel_setBuffer, bufCacheLens, 0UL, 3UL);
setBufFn(encoder, __sel_setBuffer, bufOut, 0UL, 4UL);
unsigned int params[5];
params[0] = (unsigned int)numHeads;
params[1] = (unsigned int)numKvHeads;
params[2] = (unsigned int)maxKvLen;
params[3] = isSliding ? 1U : 0U;
params[4] = (unsigned int)slidingWindow;
__mps_set_uints(encoder, params, 5UL, 5UL);
struct MTLSize tgSize;
tgSize.width = __mps_attn_batch_simdgroups(headDim) * 32UL; tgSize.height = 1UL; tgSize.depth = 1UL;
struct MTLSize numThreadgroups;
numThreadgroups.width = batchSize; numThreadgroups.height = numHeads; numThreadgroups.depth = 1UL;
MMsgDispatch dispatchFn = (MMsgDispatch)objc_msgSend;
dispatchFn(encoder, __sel_dispatchThreadgroups, numThreadgroups, tgSize);
return 1;
}
}
// Standalone one-shot dispatch (own commit+wait+readback, same shape as
// __mps_run_binary_kernel above) -- every input changes each decode step,
// so nothing here benefits from staying GPU-resident across calls.
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) {
if (!__mps_attn_shape_ok(numHeads, numKvHeads, headDim)) return 0;
unsafe {
if (!__mps_ensure_device_queue()) return 0;
MMsg0 msg0 = (MMsg0)objc_msgSend;
unsigned long qBytes = batchSize * numHeads * headDim * 2UL;
unsigned long kvBytes = batchSize * maxKvLen * numKvHeads * headDim * 2UL;
unsigned long cacheLensBytes = batchSize * 4UL;
MtlObject* bufQ = __mps_buf_get_with_bytes((const unsigned char*)q, qBytes);
MtlObject* bufK = __mps_buf_get_with_bytes((const unsigned char*)kAll, kvBytes);
MtlObject* bufV = __mps_buf_get_with_bytes((const unsigned char*)vAll, kvBytes);
MtlObject* bufCacheLens = __mps_buf_get_with_bytes((const unsigned char*)cacheLens, cacheLensBytes);
MtlObject* bufOut = __mps_buf_get(qBytes);
int ok = bufQ != null && bufK != null && bufV != null &&
bufCacheLens != null && bufOut != null;
if (ok) {
MtlObject* cmdBuf = msg0(__mps_queue, __sel_commandBuffer);
MtlObject* encoder = msg0(cmdBuf, __sel_computeCommandEncoder);
ok = __mps_attn_batch_encode(encoder, bufQ, bufK, bufV, bufCacheLens, bufOut,
batchSize, numHeads, numKvHeads, headDim, maxKvLen,
isSliding, slidingWindow);
MMsgVoid0 msgVoid0 = (MMsgVoid0)objc_msgSend;
msgVoid0(encoder, __sel_endEncoding);
if (ok) {
msgVoid0(cmdBuf, __sel_commit);
msgVoid0(cmdBuf, __sel_waitUntilCompleted);
unsigned char* outPtr = __mps_contents(bufOut);
ok = outPtr != null;
if (ok) { memcpy((unsigned char*)out, outPtr, qBytes); }
}
}
if (bufQ != null) { __mps_buf_put(bufQ, qBytes); }
if (bufK != null) { __mps_buf_put(bufK, kvBytes); }
if (bufV != null) { __mps_buf_put(bufV, kvBytes); }
if (bufCacheLens != null) { __mps_buf_put(bufCacheLens, cacheLensBytes); }
if (bufOut != null) { __mps_buf_put(bufOut, qBytes); }
return ok;
}
}
// Chained variant -- see gpu_mps.h's comment. Encodes onto the shared batch
// encoder and registers 'out' as the pending readback target like every
// other *_chained op, leaving the result as MpsBuf::batch_last_output().
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) {
if (bufQ.handle == null || bufK.handle == null || bufV.handle == null || bufCacheLens.handle == null) return 0;
if (!__mps_batch_active) return 0;
if (!__mps_attn_shape_ok(numHeads, numKvHeads, headDim)) return 0;
unsafe {
unsigned long outBytes = batchSize * numHeads * headDim * 2UL;
MtlObject* bufOut = __mps_buf_get_batched(outBytes);
if (bufOut == null) return 0;
if (!__mps_attn_batch_encode(__mps_batch_encoder, bufQ.handle, bufK.handle, bufV.handle, bufCacheLens.handle, bufOut,
batchSize, numHeads, numKvHeads, headDim, maxKvLen,
isSliding, slidingWindow)) return 0;
__mps_pending_add(bufOut, (unsigned char*)out, outBytes, MPS_BF16);
return 1;
}
}
// ── Flash and linear attention ──────────────────────────────────────────────
// See the "Flash and linear attention" kernels in gpu_mps_kernels.metal.
// Every call fills one MpsAttnParams (the kernels' AttnParams, field for
// field) with the problem size and element strides, so the head-major
// [batch][heads][seq][dim] API (mps_*_attention_ex) and the token-major
// [seq][heads*dim] bf16 API (mps_*_attn_bf16) share the same kernels.
struct MpsAttnParams {
unsigned int seqLen;
unsigned int kvLen;
unsigned int heads;
unsigned int kvHeads;
unsigned int causal;
unsigned int window;
unsigned int hasState;
unsigned int seg;
float scale;
unsigned int qRow;
unsigned int qHead;
unsigned int qBatch;
unsigned int kvRow;
unsigned int kvHead;
unsigned int kvBatch;
};
// These must match gpu_mps_kernels.metal: the matrix kernel (head dim
// <= 256) runs __mps_flash_simdgroups(dtype, headDim) simdgroups of 8 query rows,
// the head-dim-512 kernel FLASH_BQ simdgroups of one row, linear attention
// LIN_SG simdgroups of LIN_COLS(headDim) output columns.
#define FLASH_BQ ((unsigned long)16)
#define LIN_SG ((unsigned long)4)
// Simdgroups per threadgroup of flash_attn_mma_<dtype>_d<headDim> (its SG
// template argument in gpu_mps_kernels.metal).
static unsigned long __mps_flash_simdgroups(int dtype, unsigned long headDim) {
if (headDim == 96UL || headDim == 128UL) return 8UL;
if (headDim == 64UL && dtype == MPS_FP16) return 8UL;
return 4UL;
}
int mps_attention_head_dim_ok(unsigned long headDim) {
return headDim == 32UL || headDim == 64UL || headDim == 96UL || headDim == 128UL ||
headDim == 256UL || headDim == 512UL;
}
static unsigned long __mps_attn_dim_index(unsigned long headDim) {
if (headDim == 32UL) return 0UL;
if (headDim == 64UL) return 1UL;
if (headDim == 96UL) return 2UL;
if (headDim == 128UL) return 3UL;
if (headDim == 256UL) return 4UL;
return 5UL;
}
// One cached pipeline per (dtype, head dim): slot dtype * 6 + dim index.
static MtlObject* __mps_flash_pipelines[18];
static MtlObject* __mps_linear_pipelines[18];
static MtlObject* __mps_linear_z_pipelines[3];
// "<prefix><f32|f16|bf16>_d<headDim>" into 'buf' (>= 48 bytes); headDim 0
// leaves off the "_d" suffix.
static void __mps_attn_kernel_name(&char buf, const char* prefix, int dtype, unsigned long headDim) {
unsafe {
char* out = (char*)buf;
unsigned long n = 0UL;
while (prefix[n] != (char)0) { out[n] = prefix[n]; n = n + 1UL; }
const char* dt = dtype == MPS_FP16 ? "f16" : (dtype == MPS_BF16 ? "bf16" : "f32");
unsigned long i = 0UL;
while (dt[i] != (char)0) { out[n] = dt[i]; n = n + 1UL; i = i + 1UL; }
if (headDim != 0UL) {
out[n] = '_'; out[n + 1UL] = 'd'; n = n + 2UL;
char digits[8];
unsigned long nd = 0UL;
unsigned long x = headDim;
while (x != 0UL) { digits[nd] = (char)(48UL + x % 10UL); nd = nd + 1UL; x = x / 10UL; }
while (nd != 0UL) { nd = nd - 1UL; out[n] = digits[nd]; n = n + 1UL; }
}
out[n] = (char)0;
}
}
static MtlObject* __mps_attn_pipeline(struct MtlObject** slots, unsigned long slot, const char* prefix, int dtype,
unsigned long headDim) {
unsafe {
if (slots[slot] == null) {
char name[48];
__mps_attn_kernel_name(&name[0], prefix, dtype, headDim);
slots[slot] = __mps_make_pipeline(&name[0]);
}
return slots[slot];
}
}
// Shared shape checks: supported head dim, kvHeads dividing heads, a
// non-empty problem, and kvLen >= seqLen whenever positions matter.
static int __mps_attn_check(struct MpsAttnShape s, int dtype, int causal) {
if (dtype != 0 && dtype != MPS_FP16 && dtype != MPS_BF16) return 0;
if (!mps_attention_head_dim_ok(s.headDim)) return 0;
if (s.batch == 0UL || s.seqLen == 0UL || s.kvLen == 0UL) return 0;
if (s.kvHeads == 0UL || s.heads % s.kvHeads != 0UL) return 0;
if ((causal || s.window != 0UL) && s.kvLen < s.seqLen) return 0;
return 1;
}
// Head-major strides: q/out [batch][heads][seqLen][D], k/v [batch][kvHeads][kvLen][D].
static struct MpsAttnParams __mps_attn_params_head_major(struct MpsAttnShape s, int causal) {
struct MpsAttnParams p;
p.seqLen = (unsigned int)s.seqLen; p.kvLen = (unsigned int)s.kvLen;
p.heads = (unsigned int)s.heads; p.kvHeads = (unsigned int)s.kvHeads;
p.causal = causal ? 1U : 0U; p.window = (unsigned int)s.window; p.hasState = 0U; p.seg = 0U;
p.scale = (float)0.0;
p.qRow = (unsigned int)s.headDim;
p.qHead = (unsigned int)(s.seqLen * s.headDim);
p.qBatch = (unsigned int)(s.heads * s.seqLen * s.headDim);
p.kvRow = (unsigned int)s.headDim;
p.kvHead = (unsigned int)(s.kvLen * s.headDim);
p.kvBatch = (unsigned int)(s.kvHeads * s.kvLen * s.headDim);
return p;
}
// Token-major strides (batch 1): q/out [seqLen][heads*D], k/v [kvLen][kvHeads*D].
static struct MpsAttnParams __mps_attn_params_token_major(struct MpsAttnShape s, int causal) {
struct MpsAttnParams p = __mps_attn_params_head_major(s, causal);
p.qRow = (unsigned int)(s.heads * s.headDim);
p.qHead = (unsigned int)s.headDim;
p.kvRow = (unsigned int)(s.kvHeads * s.headDim);
p.kvHead = (unsigned int)s.headDim;
return p;
}
// Encodes flash attention onto 'encoder' ('p.scale' <= 0: the kernel uses 1/sqrt(headDim)).
static int __mps_flash_encode(MtlObject* encoder, MtlObject* bufQ, MtlObject* bufK, MtlObject* bufV, MtlObject* bufOut,
struct MpsAttnShape s, int dtype, struct MpsAttnParams p) {
unsafe {
int wide = s.headDim > 256UL;
unsigned long slot = (unsigned long)dtype * 6UL + __mps_attn_dim_index(s.headDim);
MtlObject* pipeline = __mps_attn_pipeline(__mps_flash_pipelines, slot,
wide ? "flash_attn_simd_" : "flash_attn_mma_", dtype, s.headDim);
if (pipeline == null) return 0;
MMsgVoid1 msgVoid1 = (MMsgVoid1)objc_msgSend;
msgVoid1(encoder, __sel_setComputePipelineState, pipeline);
MMsgSetBuffer setBufFn = (MMsgSetBuffer)objc_msgSend;
setBufFn(encoder, __sel_setBuffer, bufQ, 0UL, 0UL);
setBufFn(encoder, __sel_setBuffer, bufK, 0UL, 1UL);
setBufFn(encoder, __sel_setBuffer, bufV, 0UL, 2UL);
setBufFn(encoder, __sel_setBuffer, bufOut, 0UL, 3UL);
MMsgSetBytes setBytesFn = (MMsgSetBytes)objc_msgSend;
setBytesFn(encoder, __sel_setBytes, (const unsigned char*)&p, sizeof(struct MpsAttnParams), 4UL);
unsigned long sgs = wide ? FLASH_BQ : __mps_flash_simdgroups(dtype, s.headDim);
unsigned long rowsPerGroup = wide ? FLASH_BQ : sgs * 8UL;
struct MTLSize tgSize;
tgSize.width = sgs * 32UL; tgSize.height = 1UL; tgSize.depth = 1UL;
struct MTLSize numThreadgroups;
numThreadgroups.width = (s.seqLen + rowsPerGroup - 1UL) / rowsPerGroup;
numThreadgroups.height = s.heads;
numThreadgroups.depth = s.batch;
MMsgDispatch dispatchFn = (MMsgDispatch)objc_msgSend;
dispatchFn(encoder, __sel_dispatchThreadgroups, numThreadgroups, tgSize);
return 1;
}
}
// Chunk-parallel linear attention (see linear_chunk_state in
// gpu_mps_kernels.metal): head dims 64 and 128, whole sequences only.
#define LINC_ROWS ((unsigned long)64) // must match the kernel source
static MtlObject* __mps_linc_state_pipelines[6];
static MtlObject* __mps_linc_out_pipelines[12]; // [causal][dtype][dim]
static MtlObject* __mps_linc_prefix_pipelines[2];
// The per-segment states live in one scratch buffer this library keeps and
// grows. Calls in one batch may share it: the serial encoder orders their
// dispatches. A buffer replaced mid-batch is released when the batch ends.
static MtlObject* __mps_linc_scratch = null;
static unsigned long __mps_linc_scratch_bytes = 0UL;
static MtlObject* __mps_linc_scratch_get(unsigned long bytes) {
if (bytes <= __mps_linc_scratch_bytes) return __mps_linc_scratch;
if (__mps_linc_scratch != null) {
if (__mps_batch_active) { __mps_batch_track_pool_buf(__mps_linc_scratch, __mps_linc_scratch_bytes); }
else { __mps_buf_put(__mps_linc_scratch, __mps_linc_scratch_bytes); }
}
__mps_linc_scratch = __mps_buf_get(bytes);
__mps_linc_scratch_bytes = __mps_linc_scratch == null ? 0UL : bytes;
return __mps_linc_scratch;
}
static int __mps_linear_chunked_encode(MtlObject* encoder, MtlObject* bufQ, MtlObject* bufK, MtlObject* bufV, MtlObject* bufOut,
struct MpsAttnShape s, int dtype, struct MpsAttnParams p) {
unsafe {
unsigned long di = s.headDim == 64UL ? 0UL : 1UL;
unsigned long slot = (unsigned long)dtype * 2UL + di;
MtlObject* stateP = __mps_attn_pipeline(__mps_linc_state_pipelines, slot, "linear_chunk_state_", dtype, s.headDim);
MtlObject* outP = p.causal ? __mps_attn_pipeline(__mps_linc_out_pipelines, 6UL + slot, "linear_chunk_out_causal_", dtype, s.headDim)
: __mps_attn_pipeline(__mps_linc_out_pipelines, slot, "linear_chunk_out_", dtype, s.headDim);
if (__mps_linc_prefix_pipelines[di] == null) {
__mps_linc_prefix_pipelines[di] = __mps_make_pipeline(di == 0UL ? "linear_chunk_prefix_d64"
: "linear_chunk_prefix_d128");
}
MtlObject* prefixP = __mps_linc_prefix_pipelines[di];
if (stateP == null || outP == null || prefixP == null) return 0;
// Causal needs a state per LINC_ROWS-row chunk; non-causal only a
// total, summed from at most 16 segments.
unsigned long seg = LINC_ROWS;
if (!p.causal) {
unsigned long per = (s.kvLen + 15UL) / 16UL;
per = (per + LINC_ROWS - 1UL) / LINC_ROWS * LINC_ROWS;
if (per > seg) seg = per;
}
p.seg = (unsigned int)seg;
unsigned long nSeg = (s.kvLen + seg - 1UL) / seg;
unsigned long elems = s.headDim * s.headDim + s.headDim;
MtlObject* states = __mps_linc_scratch_get(s.batch * s.heads * nSeg * elems * 4UL);
if (states == null) return 0;
MMsgVoid1 msgVoid1 = (MMsgVoid1)objc_msgSend;
MMsgSetBuffer setBufFn = (MMsgSetBuffer)objc_msgSend;
MMsgSetBytes setBytesFn = (MMsgSetBytes)objc_msgSend;
MMsgDispatch dispatchFn = (MMsgDispatch)objc_msgSend;
unsigned long slabs = s.headDim / 64UL;
struct MTLSize tgSize;
struct MTLSize grid;
tgSize.height = 1UL; tgSize.depth = 1UL;
msgVoid1(encoder, __sel_setComputePipelineState, stateP);
setBufFn(encoder, __sel_setBuffer, bufK, 0UL, 0UL);
setBufFn(encoder, __sel_setBuffer, bufV, 0UL, 1UL);
setBufFn(encoder, __sel_setBuffer, states, 0UL, 2UL);
setBytesFn(encoder, __sel_setBytes, (const unsigned char*)&p, sizeof(struct MpsAttnParams), 3UL);
tgSize.width = s.headDim * 2UL;
grid.width = nSeg; grid.height = s.heads * slabs; grid.depth = s.batch;
dispatchFn(encoder, __sel_dispatchThreadgroups, grid, tgSize);
msgVoid1(encoder, __sel_setComputePipelineState, prefixP);
setBufFn(encoder, __sel_setBuffer, states, 0UL, 0UL);
setBytesFn(encoder, __sel_setBytes, (const unsigned char*)&p, sizeof(struct MpsAttnParams), 1UL);
tgSize.width = 256UL;
grid.width = (elems + 255UL) / 256UL; grid.height = s.batch * s.heads; grid.depth = 1UL;
dispatchFn(encoder, __sel_dispatchThreadgroups, grid, tgSize);
msgVoid1(encoder, __sel_setComputePipelineState, outP);
setBufFn(encoder, __sel_setBuffer, bufQ, 0UL, 0UL);
setBufFn(encoder, __sel_setBuffer, bufK, 0UL, 1UL);
setBufFn(encoder, __sel_setBuffer, bufV, 0UL, 2UL);
setBufFn(encoder, __sel_setBuffer, bufOut, 0UL, 3UL);
setBufFn(encoder, __sel_setBuffer, states, 0UL, 4UL);
setBytesFn(encoder, __sel_setBytes, (const unsigned char*)&p, sizeof(struct MpsAttnParams), 5UL);
tgSize.width = 256UL;
grid.width = (s.seqLen + LINC_ROWS - 1UL) / LINC_ROWS; grid.height = s.heads * slabs; grid.depth = s.batch;
dispatchFn(encoder, __sel_dispatchThreadgroups, grid, tgSize);
return 1;
}
}
// Encodes linear attention (and, with a state, the z update after it).
static int __mps_linear_encode(MtlObject* encoder, MtlObject* bufQ, MtlObject* bufK, MtlObject* bufV, MtlObject* bufOut, MtlObject* bufState,
struct MpsAttnShape s, int dtype, struct MpsAttnParams p) {
if (bufState == null && (s.headDim == 64UL || s.headDim == 128UL) && (!p.causal || s.kvLen == s.seqLen)) {
return __mps_linear_chunked_encode(encoder, bufQ, bufK, bufV, bufOut, s, dtype, p);
}
unsafe {
unsigned long slot = (unsigned long)dtype * 6UL + __mps_attn_dim_index(s.headDim);
MtlObject* pipeline = __mps_attn_pipeline(__mps_linear_pipelines, slot, "linear_attn_", dtype, s.headDim);
if (pipeline == null) return 0;
int hasState = bufState != null;
MtlObject* zPipeline = null;
if (hasState) {
zPipeline = __mps_attn_pipeline(__mps_linear_z_pipelines, (unsigned long)dtype,
"linear_attn_z_update_", dtype, 0UL);
if (zPipeline == null) return 0;
}
p.hasState = hasState ? 1U : 0U;
MMsgVoid1 msgVoid1 = (MMsgVoid1)objc_msgSend;
msgVoid1(encoder, __sel_setComputePipelineState, pipeline);
MMsgSetBuffer setBufFn = (MMsgSetBuffer)objc_msgSend;
setBufFn(encoder, __sel_setBuffer, bufQ, 0UL, 0UL);
setBufFn(encoder, __sel_setBuffer, bufK, 0UL, 1UL);
setBufFn(encoder, __sel_setBuffer, bufV, 0UL, 2UL);
setBufFn(encoder, __sel_setBuffer, bufOut, 0UL, 3UL);
// The kernel never touches buffer(4) without a state; any buffer
// satisfies the binding.
setBufFn(encoder, __sel_setBuffer, hasState ? bufState : bufOut, 0UL, 4UL);
MMsgSetBytes setBytesFn = (MMsgSetBytes)objc_msgSend;
setBytesFn(encoder, __sel_setBytes, (const unsigned char*)&p, sizeof(struct MpsAttnParams), 5UL);
unsigned long cols = 1UL; // LIN_COLS(D) in the kernel source
if (s.headDim <= 256UL) cols = 2UL;
if (s.headDim <= 128UL) cols = 4UL;
if (s.headDim <= 64UL) cols = 8UL;
if (s.headDim <= 32UL) cols = 16UL;
unsigned long groups = s.headDim / cols;
struct MTLSize tgSize;
tgSize.width = LIN_SG * 32UL; tgSize.height = 1UL; tgSize.depth = 1UL;
struct MTLSize numThreadgroups;
numThreadgroups.width = (groups + LIN_SG - 1UL) / LIN_SG;
numThreadgroups.height = s.heads;
numThreadgroups.depth = s.batch;
MMsgDispatch dispatchFn = (MMsgDispatch)objc_msgSend;
dispatchFn(encoder, __sel_dispatchThreadgroups, numThreadgroups, tgSize);
if (hasState) {
msgVoid1(encoder, __sel_setComputePipelineState, zPipeline);
setBufFn(encoder, __sel_setBuffer, bufK, 0UL, 0UL);
setBufFn(encoder, __sel_setBuffer, bufState, 0UL, 1UL);
setBytesFn(encoder, __sel_setBytes, (const unsigned char*)&p, sizeof(struct MpsAttnParams), 2UL);
unsigned int dim = (unsigned int)s.headDim;
setBytesFn(encoder, __sel_setBytes, (const unsigned char*)&dim, 4UL, 3UL);
unsigned long n = s.heads * s.headDim;
tgSize.width = 64UL;
numThreadgroups.width = (n + 63UL) / 64UL;
numThreadgroups.height = 1UL;
dispatchFn(encoder, __sel_dispatchThreadgroups, numThreadgroups, tgSize);
}
return 1;
}
}
// Runs 'lin' (0 flash, 1 linear) into the open batch, or in a one-op batch
// that is waited for when none is open.
static int __mps_attn_run(int lin, MtlObject* bufQ, MtlObject* bufK, MtlObject* bufV, MtlObject* bufOut, struct MpsAttnShape s,
int dtype, struct MpsAttnParams p) {
int ownBatch = !__mps_batch_active;
if (ownBatch) {
mps_batch_begin();
if (!__mps_batch_active) return 0;
}
int ok = lin ? __mps_linear_encode(__mps_batch_encoder, bufQ, bufK, bufV, bufOut, null, s, dtype, p)
: __mps_flash_encode(__mps_batch_encoder, bufQ, bufK, bufV, bufOut, s, dtype, p);
if (ownBatch) { mps_batch_end(); }
return ok;
}
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) {
if (q.handle == null || k.handle == null || v.handle == null || out.handle == null) return 0;
if (!__mps_attn_check(s, dtype, causal)) return 0;
struct MpsAttnParams p = __mps_attn_params_head_major(s, causal);
p.scale = scale;
unsafe { return __mps_attn_run(0, q.handle, k.handle, v.handle, out.handle, s, dtype, p); }
}
int mps_linear_attention_ex(struct MpsBuf q, struct MpsBuf k, struct MpsBuf v, struct MpsBuf out, struct MpsAttnShape s, int dtype, int causal) {
if (q.handle == null || k.handle == null || v.handle == null || out.handle == null) return 0;
if (s.window != 0UL || !__mps_attn_check(s, dtype, causal)) return 0;
struct MpsAttnParams p = __mps_attn_params_head_major(s, causal);
unsafe { return __mps_attn_run(1, q.handle, k.handle, v.handle, out.handle, s, dtype, p); }
}
static struct MpsAttnShape __mps_attn_shape_flat(unsigned long heads, unsigned long seqLen, unsigned long kvLen,
unsigned long headDim) {
struct MpsAttnShape s;
s.batch = 1UL; s.heads = heads; s.kvHeads = heads; s.seqLen = seqLen; s.kvLen = kvLen;
s.headDim = headDim; s.window = 0UL;
return s;
}
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) {
return mps_flash_attention_ex(q, k, v, out, __mps_attn_shape_flat(heads, seqLen, kvLen, headDim), 0, scale, 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) {
return mps_linear_attention_ex(q, k, v, out, __mps_attn_shape_flat(heads, seqLen, kvLen, headDim), 0, causal);
}
// Host-array forms: upload, run, read back. Inside an open batch, 'out' is
// filled at mps_batch_end() like every other batched op.
static int __mps_attn_host_f32(int lin, 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) {
struct MpsAttnShape s = __mps_attn_shape_flat(heads, seqLen, kvLen, headDim);
if (!__mps_attn_check(s, 0, causal)) return 0;
struct MpsAttnParams p = __mps_attn_params_head_major(s, causal);
p.scale = scale;
unsafe {
int ownBatch = !__mps_batch_active;
if (ownBatch) {
mps_batch_begin();
if (!__mps_batch_active) return 0;
}
unsigned long qBytes = heads * seqLen * headDim * 4UL;
unsigned long kvBytes = heads * kvLen * headDim * 4UL;
MtlObject* bufQ = __mps_buf_get_with_bytes_batched((const unsigned char*)q, qBytes);
MtlObject* bufK = __mps_buf_get_with_bytes_batched((const unsigned char*)k, kvBytes);
MtlObject* bufV = __mps_buf_get_with_bytes_batched((const unsigned char*)v, kvBytes);
MtlObject* bufOut = __mps_buf_get_batched(qBytes);
int ok = bufQ != null && bufK != null && bufV != null && bufOut != null;
if (ok) {
ok = lin ? __mps_linear_encode(__mps_batch_encoder, bufQ, bufK, bufV, bufOut, null, s, 0, p)
: __mps_flash_encode(__mps_batch_encoder, bufQ, bufK, bufV, bufOut, s, 0, p);
}
if (ok) { __mps_pending_add(bufOut, (unsigned char*)out, qBytes, 0); }
if (ownBatch) { mps_batch_end(); }
return ok;
}
}
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) {
return __mps_attn_host_f32(0, q, k, v, out, heads, seqLen, kvLen, headDim, scale, 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) {
return __mps_attn_host_f32(1, q, k, v, out, heads, seqLen, kvLen, headDim, (float)0.0, causal);
}
// ── Token-major bf16 API (flash) ────────────────────────────────────────────
static struct MpsAttnShape __mps_attn_shape_tm(unsigned long seqLen, unsigned long kvLen, unsigned long numHeads,
unsigned long numKvHeads, unsigned long headDim,
unsigned long window) {
struct MpsAttnShape s;
s.batch = 1UL; s.heads = numHeads; s.kvHeads = numKvHeads; s.seqLen = seqLen; s.kvLen = kvLen;
s.headDim = headDim; s.window = window;
return s;
}
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) {
if (bufQ.handle == null || bufK.handle == null || bufV.handle == null) return 0;
if (!__mps_batch_active) return 0;
struct MpsAttnShape s = __mps_attn_shape_tm(seqLen, kvLen, numHeads, numKvHeads, headDim,
isSliding ? slidingWindow : 0UL);
if (!__mps_attn_check(s, MPS_BF16, causal)) return 0;
struct MpsAttnParams p = __mps_attn_params_token_major(s, causal);
p.scale = scale;
unsafe {
unsigned long outBytes = seqLen * numHeads * headDim * 2UL;
MtlObject* bufOut = __mps_buf_get_batched(outBytes);
if (bufOut == null) return 0;
if (!__mps_flash_encode(__mps_batch_encoder, bufQ.handle, bufK.handle, bufV.handle, bufOut, s, MPS_BF16, p)) return 0;
__mps_pending_add(bufOut, (unsigned char*)out, outBytes, MPS_BF16);
return 1;
}
}
// Batch-aware like mps_sub_f32: inside an open batch it encodes there and
// 'out' is filled at that batch's mps_batch_end(); otherwise it runs its
// own one-op batch and 'out' is filled on return.
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) {
unsafe {
int ownBatch = !__mps_batch_active;
if (ownBatch) {
mps_batch_begin();
if (!__mps_batch_active) return 0;
}
unsigned long qBytes = seqLen * numHeads * headDim * 2UL;
unsigned long kvBytes = kvLen * numKvHeads * headDim * 2UL;
MtlObject* bufQ = __mps_buf_get_with_bytes_batched((const unsigned char*)q, qBytes);
MtlObject* bufK = __mps_buf_get_with_bytes_batched((const unsigned char*)k, kvBytes);
MtlObject* bufV = __mps_buf_get_with_bytes_batched((const unsigned char*)v, kvBytes);
int ok = bufQ != null && bufK != null && bufV != null &&
mps_flash_attn_bf16_chained(__mps_buf_make(bufQ, qBytes, MPS_BF16),
__mps_buf_make(bufK, kvBytes, MPS_BF16),
__mps_buf_make(bufV, kvBytes, MPS_BF16), out, seqLen, kvLen,
numHeads, numKvHeads, headDim, scale, causal, isSliding,
slidingWindow);
if (ownBatch) { mps_batch_end(); }
return ok;
}
}
// ── Token-major bf16 API (linear) ───────────────────────────────────────────
unsigned long mps_linear_attn_state_bytes(unsigned long numHeads, unsigned long headDim) {
return numHeads * (headDim * headDim + headDim) * 4UL;
}
struct MpsBuf MpsBuf::linear_attn_state(unsigned long numHeads, unsigned long headDim) {
unsafe {
if (!__mps_ensure_device_queue()) return __mps_buf_make(null, 0UL, 0);
unsigned long bytes = mps_linear_attn_state_bytes(numHeads, headDim);
MtlObject* buf = __mps_buf_get(bytes);
if (buf == null) return __mps_buf_make(null, 0UL, 0);
MMsg0 msg0 = (MMsg0)objc_msgSend;
unsigned char* ptr = __mps_contents(buf);
if (ptr == null) { __mps_buf_put(buf, bytes); return __mps_buf_make(null, 0UL, 0); }
memset(ptr, 0, bytes);
return __mps_buf_make(buf, bytes, 0);
}
}
// Encodes token-major linear attention onto the open batch; 'bufState'
// NULL means a zero start state that is discarded.
static int __mps_linear_tm_encode(MtlObject* bufQ, MtlObject* bufK, MtlObject* bufV, MtlObject* bufState, unsigned char* out,
unsigned long seqLen, unsigned long kvLen,
unsigned long numHeads, unsigned long numKvHeads, unsigned long headDim,
int causal) {
struct MpsAttnShape s = __mps_attn_shape_tm(seqLen, kvLen, numHeads, numKvHeads, headDim, 0UL);
if (!__mps_attn_check(s, MPS_BF16, causal)) return 0;
struct MpsAttnParams p = __mps_attn_params_token_major(s, causal);
unsafe {
unsigned long outBytes = seqLen * numHeads * headDim * 2UL;
MtlObject* bufOut = __mps_buf_get_batched(outBytes);
if (bufOut == null) return 0;
if (!__mps_linear_encode(__mps_batch_encoder, bufQ, bufK, bufV, bufOut, bufState, s, MPS_BF16, p)) return 0;
__mps_pending_add(bufOut, out, outBytes, MPS_BF16);
return 1;
}
}
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) {
if (bufQ.handle == null || bufK.handle == null || bufV.handle == null || bufState.handle == null) return 0;
if (!__mps_batch_active) return 0;
unsafe {
return __mps_linear_tm_encode(bufQ.handle, bufK.handle, bufV.handle, bufState.handle, (unsigned char*)out,
seqLen, seqLen, numHeads, numKvHeads, headDim, 1);
}
}
// Batch-aware like mps_flash_attn_bf16 above.
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) {
if (causal && kvLen != seqLen) return 0;
unsafe {
int ownBatch = !__mps_batch_active;
if (ownBatch) {
mps_batch_begin();
if (!__mps_batch_active) return 0;
}
unsigned long qBytes = seqLen * numHeads * headDim * 2UL;
unsigned long kvBytes = kvLen * numKvHeads * headDim * 2UL;
MtlObject* bufQ = __mps_buf_get_with_bytes_batched((const unsigned char*)q, qBytes);
MtlObject* bufK = __mps_buf_get_with_bytes_batched((const unsigned char*)k, kvBytes);
MtlObject* bufV = __mps_buf_get_with_bytes_batched((const unsigned char*)v, kvBytes);
int ok = bufQ != null && bufK != null && bufV != null &&
__mps_linear_tm_encode(bufQ, bufK, bufV, null, (unsigned char*)out,
seqLen, kvLen, numHeads, numKvHeads, headDim, causal);
if (ownBatch) { mps_batch_end(); }
return ok;
}
}
// ── Device-buffer GEMM (see gpu_mps.h's mps_gemm) ───────────────────────────
// The three matmul kernel families (ab, abt, atb) with the tier choice every
// mps_matmul_* entry point above makes, writing a caller-owned buffer instead
// of a batch-owned one with a CPU readback.
// Opens a one-op batch when none is active; *own says to end it.
static int __mps_encoder_begin(&int own) {
*own = !__mps_batch_active;
if (*own) mps_batch_begin();
return __mps_batch_active;
}
static void __mps_encoder_end(int own) {
if (own) mps_batch_end();
}
// variant 0: out[d0,d2] = a[d0,d1] . b[d1,d2] (matmul_kernel: M K N)
// variant 1: out[d0,d2] = a[d0,d1] . b^T, b [d2,d1] (matmul_abt_kernel: M N K)
// variant 2: out[d1,d2] = a^T . b, a [d0,d1], b [d0,d2] (matmul_atb_kernel: M K N)
static int __mps_gemm_encode(int variant, MtlObject* a, MtlObject* b, MtlObject* out, unsigned long d0, unsigned long d1,
unsigned long d2) {
unsafe {
int useMulti = d0 % 32UL == 0UL && d1 % 32UL == 0UL && d2 % 32UL == 0UL;
int useSmma = d0 % 8UL == 0UL && d1 % 8UL == 0UL && d2 % 8UL == 0UL;
struct MtlObject** multi = variant == 0 ? &__mps_matmul_multi_pipeline
: (variant == 1 ? &__mps_matmul_abt_multi_pipeline : &__mps_matmul_atb_multi_pipeline);
struct MtlObject** smma = variant == 0 ? &__mps_matmul_smma_pipeline
: (variant == 1 ? &__mps_matmul_abt_smma_pipeline : &__mps_matmul_atb_smma_pipeline);
struct MtlObject** tiled = variant == 0 ? &__mps_matmul_pipeline
: (variant == 1 ? &__mps_matmul_abt_pipeline : &__mps_matmul_atb_pipeline);
const char* base = variant == 0 ? "matmul_kernel" : (variant == 1 ? "matmul_abt_kernel" : "matmul_atb_kernel");
const char* multiName = variant == 0 ? "matmul_kernel_smma_multi"
: (variant == 1 ? "matmul_abt_kernel_smma_multi" : "matmul_atb_kernel_smma_multi");
const char* smmaName = variant == 0 ? "matmul_kernel_smma"
: (variant == 1 ? "matmul_abt_kernel_smma" : "matmul_atb_kernel_smma");
if (useMulti && *multi == null) *multi = __mps_make_pipeline(multiName);
if (useMulti && *multi == null) useMulti = 0;
if (!useMulti && useSmma && *smma == null) *smma = __mps_make_pipeline(smmaName);
if (!useMulti && useSmma && *smma == null) useSmma = 0;
if (!useMulti && !useSmma && *tiled == null) *tiled = __mps_make_pipeline(base);
MtlObject* pipeline = useMulti ? *multi : (useSmma ? *smma : *tiled);
if (pipeline == null) return 0;
// the output's rows and columns
unsigned long rows = variant == 2 ? d1 : d0;
unsigned long cols = d2;
unsigned int c0 = (unsigned int)d0;
unsigned int c1 = (unsigned int)d1;
unsigned int c2 = (unsigned int)d2;
MtlObject* enc = __mps_batch_encoder;
MMsgVoid1 msgVoid1 = (MMsgVoid1)objc_msgSend;
msgVoid1(enc, __sel_setComputePipelineState, pipeline);
MMsgSetBuffer setBufFn = (MMsgSetBuffer)objc_msgSend;
setBufFn(enc, __sel_setBuffer, a, 0UL, 0UL);
setBufFn(enc, __sel_setBuffer, b, 0UL, 1UL);
setBufFn(enc, __sel_setBuffer, out, 0UL, 2UL);
MMsgSetBytes setBytesFn = (MMsgSetBytes)objc_msgSend;
setBytesFn(enc, __sel_setBytes, (const unsigned char*)&c0, 4UL, 3UL);
setBytesFn(enc, __sel_setBytes, (const unsigned char*)&c1, 4UL, 4UL);
setBytesFn(enc, __sel_setBytes, (const unsigned char*)&c2, 4UL, 5UL);
struct MTLSize tgSize;
struct MTLSize groups;
tgSize.depth = 1UL; groups.depth = 1UL;
if (useMulti) {
tgSize.width = 128UL; tgSize.height = 1UL;
groups.width = cols / 32UL; groups.height = rows / 32UL;
} else if (useSmma) {
tgSize.width = 32UL; tgSize.height = 1UL;
groups.width = cols / 8UL; groups.height = rows / 8UL;
} else {
tgSize.width = 16UL; tgSize.height = 16UL;
groups.width = (cols + 15UL) / 16UL; groups.height = (rows + 15UL) / 16UL;
}
MMsgDispatch dispatchFn = (MMsgDispatch)objc_msgSend;
dispatchFn(enc, __sel_dispatchThreadgroups, groups, tgSize);
return 1;
}
}
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) {
if (a.handle == null || b.handle == null || c.handle == null) return 0;
if ((ta && tb) || m == 0UL || n == 0UL || k == 0UL) return 0;
if (a.bytes < m * k * 4UL || b.bytes < k * n * 4UL || c.bytes < m * n * 4UL) return 0;
int own = 0;
if (!__mps_encoder_begin(&own)) return 0;
int ok = 0;
unsafe {
if (ta) ok = __mps_gemm_encode(2, a.handle, b.handle, c.handle, k, m, n);
else if (tb) ok = __mps_gemm_encode(1, a.handle, b.handle, c.handle, m, k, n);
else ok = __mps_gemm_encode(0, a.handle, b.handle, c.handle, m, k, n);
}
__mps_encoder_end(own);
return ok;
}
// ── Interop (see gpu_mps.h's mps_device_handle and mps_dispatch) ────────────
?&MtlObject mps_device_handle() {
if (!__mps_ensure_device_queue()) return null;
unsafe { return (?&MtlObject)__mps_device; }
}
?&MtlObject mps_queue_handle() {
if (!__mps_ensure_device_queue()) return null;
unsafe { return (?&MtlObject)__mps_queue; }
}
#define MPS_DISPATCH_MAX_ 64UL
static const char* __mps_dispatch_names[64];
static MtlObject* __mps_dispatch_pipes[64];
static unsigned long __mps_dispatch_limit[64]; // the pipeline's maxTotalThreadsPerThreadgroup
static unsigned long __mps_dispatch_count = 0UL;
static ObjcSel_* __sel_maxThreads = null;
typedef fn unsigned long(MtlObject*, ObjcSel_*) MMsgULong;
static int __mps_str_eq(const char* a, const char* b) {
unsafe {
unsigned long i = 0UL;
while (a[i] != (char)0 && a[i] == b[i]) i = i + 1UL;
return a[i] == b[i];
}
}
// The cached slot for 'kernel', built on first use; -1 when it can't be.
static long __mps_dispatch_slot(const char* kernel) {
unsafe {
unsigned long i = 0UL;
while (i < __mps_dispatch_count) {
const char* name = __mps_dispatch_names[i];
if (name == kernel || __mps_str_eq(name, kernel)) return (long)i;
i = i + 1UL;
}
if (__mps_dispatch_count >= MPS_DISPATCH_MAX_) return -1L;
MtlObject* pipe = __mps_make_pipeline(kernel);
if (pipe == null) return -1L;
if (__sel_maxThreads == null) __sel_maxThreads = sel_registerName("maxTotalThreadsPerThreadgroup");
MMsgULong maxFn = (MMsgULong)objc_msgSend;
__mps_dispatch_names[i] = kernel;
__mps_dispatch_pipes[i] = pipe;
__mps_dispatch_limit[i] = maxFn(pipe, __sel_maxThreads);
__mps_dispatch_count = i + 1UL;
return (long)i;
}
}
int mps_dispatch(const &char kernel, []const MpsBuf bufs, []const unsigned char params,
unsigned long groups, unsigned long groupSize) {
unsigned long count = bufs.len;
unsigned long paramBytes = params.len;
if (groups == 0UL || groupSize == 0UL || count > 30UL || count == 0UL) return 0;
if (!__mps_ensure_library()) return 0;
unsafe {
const struct MpsBuf* b = (const struct MpsBuf*)bufs.ptr;
unsigned long i = 0UL;
while (i < count) { if (b[i].handle == null) return 0; i = i + 1UL; }
long slot = __mps_dispatch_slot((const char*)kernel);
if (slot < 0L || groupSize > __mps_dispatch_limit[slot]) return 0;
int own = 0;
if (!__mps_encoder_begin(&own)) return 0;
MtlObject* enc = __mps_batch_encoder;
MMsgVoid1 msgVoid1 = (MMsgVoid1)objc_msgSend;
msgVoid1(enc, __sel_setComputePipelineState, __mps_dispatch_pipes[slot]);
MMsgSetBuffer setBufFn = (MMsgSetBuffer)objc_msgSend;
i = 0UL;
while (i < count) { setBufFn(enc, __sel_setBuffer, b[i].handle, 0UL, i); i = i + 1UL; }
if (paramBytes > 0UL) {
MMsgSetBytes setBytesFn = (MMsgSetBytes)objc_msgSend;
setBytesFn(enc, __sel_setBytes, (unsigned char*)params.ptr, paramBytes, count);
}
struct MTLSize tgSize;
tgSize.width = groupSize; tgSize.height = 1UL; tgSize.depth = 1UL;
struct MTLSize grid;
grid.width = groups; grid.height = 1UL; grid.depth = 1UL;
MMsgDispatch dispatchFn = (MMsgDispatch)objc_msgSend;
dispatchFn(enc, __sel_dispatchThreadgroups, grid, tgSize);
__mps_encoder_end(own);
return 1;
}
}
} // namespace std