420 lines
17 KiB
C
420 lines
17 KiB
C
/* SPDX-License-Identifier: AGPL-3.0-or-later */
|
|
/* Copyright (C) 2026 David (Corey) Herrera */
|
|
/**
|
|
* @file fm_metal.h
|
|
* @brief Metal GPU compute backend for Foundry — public C API (prefix fm_).
|
|
* Build flag USE_METAL=1. GPU tensor ops via Apple's Metal framework on a
|
|
* global singleton (MTLDevice + MTLCommandQueue + specialized pipeline cache),
|
|
* with unified memory (MTLResourceStorageModeShared) for zero-copy on Apple Silicon.
|
|
*
|
|
* This header owns the backend-wide API — lifecycle, function-constant
|
|
* specialization, buffers, no-copy regions, dispatch thresholds, batching —
|
|
* and then includes one header per kernel family from metal/kernels/. Include
|
|
* this file and you have the whole backend; a family's own header documents
|
|
* only that family. Kernel families are added without touching this file.
|
|
*
|
|
* @status Buffers, no-copy mmap regions with offset binding, MPS GEMM/GEMV
|
|
* (single + batched), custom kernels and function-constant pipeline
|
|
* specialization all work.
|
|
* @issues None
|
|
* @todo None
|
|
*/
|
|
|
|
#ifndef FOUNDRY_METAL_H
|
|
#define FOUNDRY_METAL_H
|
|
|
|
#include <stdbool.h>
|
|
#include <stddef.h>
|
|
#include <stdint.h>
|
|
|
|
#ifdef __cplusplus
|
|
extern "C" {
|
|
#endif
|
|
|
|
/* ================================================================
|
|
* Opaque types
|
|
* ================================================================ */
|
|
|
|
typedef struct fm_buffer fm_buffer;
|
|
|
|
/* ================================================================
|
|
* Lifecycle
|
|
* ================================================================ */
|
|
|
|
/* Initialize Metal device, command queue, and pipeline cache.
|
|
* Returns 0 on success, -1 on failure (no Metal device). */
|
|
int fm_init(void);
|
|
|
|
/* Cleanup all Metal resources. */
|
|
void fm_free(void);
|
|
|
|
/* Returns true if Metal backend is initialized and ready. */
|
|
bool fm_available(void);
|
|
|
|
/**
|
|
* @brief True when the custom-kernel shader library actually built.
|
|
*
|
|
* fm_available() only reports that a Metal device was initialized. Shader
|
|
* compilation is non-fatal, so when it fails every custom kernel silently
|
|
* becomes a no-op while fm_available() still returns true. All families share
|
|
* one MSL translation unit, so one bad fragment disables the whole backend.
|
|
*
|
|
* Tests and benchmarks MUST gate on this: a no-op backend makes GPU-vs-GPU
|
|
* comparisons agree on all-zeros and makes bandwidth numbers fictional.
|
|
*/
|
|
bool fm_shaders_available(void);
|
|
|
|
/* ================================================================
|
|
* Function-constant pipeline specialization
|
|
*
|
|
* A Metal function constant is a value bound when the pipeline is
|
|
* built, not when the kernel is dispatched. Inside the kernel it is a
|
|
* compile-time constant, so loop bounds specialized this way unroll and
|
|
* the local arrays those loops index stay in registers instead of
|
|
* spilling to thread-local memory. That is why every shape-dependent
|
|
* kernel in this backend takes its shape as a function constant rather
|
|
* than as a buffer argument.
|
|
*
|
|
* One pipeline is compiled and cached per (kernel name + constant
|
|
* values) combination. Repeating a call with the same values reuses
|
|
* the cached pipeline; changing any value builds and caches a new one.
|
|
* ================================================================ */
|
|
|
|
/* Maximum number of distinct constant indices in one specialization set. */
|
|
#define FM_FC_MAX_VALUES 16
|
|
|
|
/* Type tag for a function-constant value. Must match the declared type
|
|
* of the corresponding [[function_constant(index)]] in the shader — a
|
|
* mismatch makes pipeline creation fail. */
|
|
typedef enum fm_fc_type {
|
|
FM_FC_TYPE_NONE = 0,
|
|
FM_FC_TYPE_BOOL = 1, /* MSL bool */
|
|
FM_FC_TYPE_SHORT = 2, /* MSL short */
|
|
FM_FC_TYPE_INT = 3 /* MSL int */
|
|
} fm_fc_type;
|
|
|
|
/* A set of function-constant values, keyed by the shader's
|
|
* [[function_constant(index)]] index.
|
|
*
|
|
* Plain value type: declare it on the stack, copy it freely, nothing to
|
|
* free. Entries are kept sorted by index, so two sets holding the same
|
|
* values produce the same cache key regardless of the order they were
|
|
* set in. Treat the fields as private — use the fm_fc_* accessors. */
|
|
typedef struct fm_fc_values {
|
|
int count;
|
|
uint16_t index[FM_FC_MAX_VALUES];
|
|
uint8_t type[FM_FC_MAX_VALUES]; /* fm_fc_type */
|
|
int32_t value[FM_FC_MAX_VALUES];
|
|
} fm_fc_values;
|
|
|
|
/**
|
|
* Empties a constant set. Call this before the first fm_fc_set_* on a
|
|
* set that was not zero-initialized.
|
|
*
|
|
* @param values Set to clear (ignored if NULL)
|
|
*
|
|
* Thread-safe: yes (operates only on the caller's struct)
|
|
* Reentrant: yes
|
|
* Signal-safe: yes
|
|
*/
|
|
void fm_fc_reset(fm_fc_values *values);
|
|
|
|
/**
|
|
* Sets the MSL `int` constant at `index`, replacing any previous value
|
|
* at that index.
|
|
*
|
|
* @param values Constant set to modify
|
|
* @param index The shader's [[function_constant(index)]] index
|
|
* @param value Value to bind at pipeline-build time
|
|
* @return 0 on success, -1 if the set is full or arguments are invalid
|
|
*
|
|
* Thread-safe: yes (operates only on the caller's struct)
|
|
* Reentrant: yes
|
|
* Signal-safe: yes
|
|
*/
|
|
int fm_fc_set_int(fm_fc_values *values, unsigned index, int32_t value);
|
|
|
|
/* Sets the MSL `short` constant at `index`. See fm_fc_set_int. */
|
|
int fm_fc_set_short(fm_fc_values *values, unsigned index, int16_t value);
|
|
|
|
/* Sets the MSL `bool` constant at `index`. See fm_fc_set_int. */
|
|
int fm_fc_set_bool(fm_fc_values *values, unsigned index, bool value);
|
|
|
|
/**
|
|
* Builds (or fetches from cache) the pipeline for `kernel_name`
|
|
* specialized with `constants`, and returns it as an opaque handle.
|
|
*
|
|
* The handle is only meaningful for identity comparison: two calls with
|
|
* equal constant sets return the same pointer, calls with different
|
|
* constant sets return different pointers. The cache owns it — do not
|
|
* release it, and do not use it after fm_free().
|
|
*
|
|
* @param kernel_name Kernel name as declared in the shader source
|
|
* @param constants Constant set, or NULL for the unspecialized pipeline
|
|
* @return Opaque cached pipeline handle, or NULL on failure
|
|
*
|
|
* Thread-safe: no (shares the global pipeline cache)
|
|
* Reentrant: no
|
|
* Signal-safe: no
|
|
*/
|
|
const void *fm_pipeline_for_constants(const char *kernel_name,
|
|
const fm_fc_values *constants);
|
|
|
|
/* The diagnostic kernel that proves specialization reached the shader
|
|
* compiler lives in metal/kernels/fm_kernel_probe.h. */
|
|
|
|
/* ================================================================
|
|
* Buffers (unified memory, zero-copy on Apple Silicon)
|
|
* ================================================================ */
|
|
|
|
/* Wrap existing CPU pointer as a Metal shared buffer (zero-copy if page-aligned).
|
|
* If pointer is not page-aligned, allocates a new Metal buffer and copies data.
|
|
* The caller retains ownership of ptr — do not free it while buffer is in use. */
|
|
fm_buffer* fm_buffer_from_ptr(void *ptr, size_t bytes);
|
|
|
|
/* Allocate a new Metal buffer with shared memory. */
|
|
fm_buffer* fm_buffer_alloc(size_t bytes);
|
|
|
|
/* Free a Metal buffer wrapper. Does NOT free the underlying data
|
|
* if created via fm_buffer_from_ptr. */
|
|
void fm_buffer_free(fm_buffer *buf);
|
|
|
|
/* Get CPU-accessible pointer to buffer data. */
|
|
void* fm_buffer_data(fm_buffer *buf);
|
|
|
|
/* Get buffer size in bytes. */
|
|
size_t fm_buffer_size(fm_buffer *buf);
|
|
|
|
/* ================================================================
|
|
* No-copy regions and offset binding
|
|
*
|
|
* fm_buffer_from_ptr copies whenever the pointer is not page-aligned,
|
|
* which is almost always for a weight inside a GGUF: the file's tensor
|
|
* data starts at an arbitrary offset and tensors are only 32-byte
|
|
* aligned, so binding each weight as its own buffer means duplicating
|
|
* the whole model. That is unaffordable and defeats the point of
|
|
* mmap'ing the file in the first place.
|
|
*
|
|
* The fix is to stop making a buffer per weight. Wrap the whole mmap
|
|
* once — its base IS page-aligned, so that single buffer genuinely
|
|
* takes the no-copy path — and address each tensor as a byte offset
|
|
* into it via fm_bufref.
|
|
*
|
|
* Measured on Apple M4 / macOS 26:
|
|
* - A buffer offset only has to satisfy the alignment of the type the
|
|
* kernel reads (1 for uchar, 4 for float, 16 for float4/half4).
|
|
* There is NO page-alignment requirement: offsets that are multiples
|
|
* of 32 but not of the 16 KiB page read correctly for every type,
|
|
* under both API and GPU validation. GGUF's 32-byte tensor
|
|
* alignment therefore satisfies binding with room to spare.
|
|
* - maxBufferLength is 14302248960 bytes (13.32 GiB). One buffer
|
|
* cannot cover a larger mapping, so fm_region splits an oversized
|
|
* region into overlapping no-copy segments.
|
|
* ================================================================ */
|
|
|
|
/* Byte alignment a bind offset must satisfy: the strictest element
|
|
* alignment any kernel here reads (float4/half4). Not a page. */
|
|
#define FM_BUFFER_OFFSET_ALIGNMENT 16
|
|
|
|
/* A bindable location: a buffer plus a byte offset into it. Plain value
|
|
* type — copy it freely, nothing to free. It does not own the buffer. */
|
|
typedef struct fm_bufref {
|
|
fm_buffer *buffer;
|
|
size_t offset;
|
|
} fm_bufref;
|
|
|
|
/* Builds a reference to `offset` bytes into `buffer`. */
|
|
fm_bufref fm_bufref_at(fm_buffer *buffer, size_t offset);
|
|
|
|
/* FM_BUFFER_OFFSET_ALIGNMENT, for callers that need it at runtime. */
|
|
size_t fm_buffer_offset_alignment(void);
|
|
|
|
/* Largest single buffer this device supports, in bytes. A mapping
|
|
* larger than this needs fm_region. */
|
|
size_t fm_buffer_max_length(void);
|
|
|
|
/**
|
|
* Wraps an already page-aligned mapping in ONE buffer without copying.
|
|
*
|
|
* Unlike fm_buffer_from_ptr this never falls back to a copy: if the
|
|
* pointer is not page-aligned, the length exceeds fm_buffer_max_length(),
|
|
* or Metal refuses the mapping, it fails and returns NULL. That is
|
|
* deliberate — a silent copy of a weight region is the bug this exists
|
|
* to prevent.
|
|
*
|
|
* The caller keeps ownership of the mapping and must not unmap it while
|
|
* the buffer lives. Free with fm_buffer_free.
|
|
*
|
|
* @param mapped_base Page-aligned base of the mapping (e.g. from mmap)
|
|
* @param mapped_bytes Length in bytes; rounded up to a page internally,
|
|
* which is safe because mmap already rounds the
|
|
* mapping up to a page boundary
|
|
* @return Buffer aliasing the mapping, or NULL (see above)
|
|
*
|
|
* Thread-safe: no (touches the global device)
|
|
* Reentrant: no
|
|
* Signal-safe: no
|
|
*/
|
|
fm_buffer *fm_buffer_map_no_copy(void *mapped_base, size_t mapped_bytes);
|
|
|
|
/* A mapping covered by one or more no-copy buffer segments, so that any
|
|
* byte range inside it can be bound without copying. */
|
|
typedef struct fm_region fm_region;
|
|
|
|
/**
|
|
* Covers a page-aligned mapping with no-copy buffer segments.
|
|
*
|
|
* A mapping that fits in one buffer gets exactly one segment. A larger
|
|
* one is covered by overlapping segments: consecutive segments start
|
|
* `segment length - max_span_bytes` apart, so any range of at most
|
|
* `max_span_bytes` lies wholly inside a single segment and never
|
|
* straddles a boundary. Segments alias the same pages — they cost
|
|
* address space, not memory, and nothing is copied.
|
|
*
|
|
* Never copies. Returns NULL if the base is not page-aligned, if
|
|
* max_span_bytes leaves no room for forward progress, or if any segment
|
|
* fails to map.
|
|
*
|
|
* @param mapped_base Page-aligned base of the mapping (e.g. from mmap)
|
|
* @param mapped_bytes Length of the mapping in bytes
|
|
* @param max_span_bytes Largest single range that will ever be bound —
|
|
* the biggest tensor. 0 selects a 1 GiB default.
|
|
* @return Region handle, or NULL on failure
|
|
*
|
|
* Thread-safe: no
|
|
* Reentrant: no
|
|
* Signal-safe: no
|
|
*/
|
|
fm_region *fm_region_open(void *mapped_base, size_t mapped_bytes,
|
|
size_t max_span_bytes);
|
|
|
|
/* Releases the segment buffers. Does not unmap the caller's mapping. */
|
|
void fm_region_close(fm_region *region);
|
|
|
|
/**
|
|
* Pretends the device's maximum buffer length is `bytes`, so the
|
|
* multi-segment path can be exercised without a mapping larger than the
|
|
* real 13.32 GiB limit. Affects fm_region_open only; 0 restores the
|
|
* device limit. Diagnostic hook for tests — production code has no
|
|
* reason to call it.
|
|
*
|
|
* Thread-safe: no (global)
|
|
* Reentrant: no
|
|
* Signal-safe: no
|
|
*/
|
|
void fm_region_force_segment_span(size_t bytes);
|
|
|
|
/* Number of segments covering the region (1 when it fits in one buffer). */
|
|
int fm_region_segment_count(const fm_region *region);
|
|
|
|
/**
|
|
* Resolves a byte range in the region to a bindable buffer + offset.
|
|
*
|
|
* @param region Region to address into
|
|
* @param offset Byte offset of the range from the mapping base
|
|
* @param bytes Length of the range
|
|
* @param out Receives the buffer and the offset within it
|
|
* @return 0 on success; -1 if the range runs past the region, is
|
|
* larger than the region's max_span_bytes, or has an
|
|
* offset that violates FM_BUFFER_OFFSET_ALIGNMENT
|
|
*
|
|
* Thread-safe: yes (read-only once the region is open)
|
|
* Reentrant: yes
|
|
* Signal-safe: no (may log on failure)
|
|
*/
|
|
int fm_region_bind(const fm_region *region, size_t offset, size_t bytes,
|
|
fm_bufref *out);
|
|
|
|
/* ================================================================
|
|
* Dispatch thresholds
|
|
*
|
|
* Backend-wide policy, not a kernel family: the CPU callers in tensor/
|
|
* and rwkvx/ test an op's size against these before choosing Metal.
|
|
* ================================================================ */
|
|
|
|
/* Dispatch threshold: ops below this stay on CPU.
|
|
*
|
|
* On Apple Silicon with unified memory, single-token decode (M=1) should
|
|
* ALWAYS stay on CPU. Single-vector operations are memory-bandwidth bound
|
|
* and CPU/GPU share the same bus, so Metal can only add dispatch overhead
|
|
* (~100-200us per kernel). Metal wins for batched operations (prefill,
|
|
* training) where M >= ~8-16 amortizes the overhead across many rows.
|
|
*
|
|
* GEMM: M*K*N. For M=1, K=2560, N=2560 → 6.5M, well under threshold.
|
|
* For prefill M=32, 32*2560*2560 = 210M, above threshold → GPU.
|
|
* GEMV: M*K. Single-vector always stays on CPU.
|
|
* WKV7: Only the sequence variant (fm_wkv7_seq) benefits from Metal.
|
|
* The per-token fm_wkv7_one always stays on CPU.
|
|
* Q8 matvec: fast_matvec_q8_preq is inherently M=1, always CPU.
|
|
* Elementwise: softmax/rms_norm/silu on single vectors stay on CPU. */
|
|
#define FM_GEMM_THRESHOLD (128L * 1024 * 1024) /* ~128M ops: Metal wins for batched prefill (N>=8) */
|
|
#define FM_Q8_GEMM_THRESHOLD (128L * 1024 * 1024) /* ~128M ops: Metal INT8 GEMM for batched prefill */
|
|
#define FM_GEMV_THRESHOLD (4096L * 4096)
|
|
#define FM_WKV7_THRESHOLD (INT64_MAX) /* disabled: fm_wkv7_one is always per-token */
|
|
#define FM_Q8_THRESHOLD (INT64_MAX) /* disabled: fast_matvec_q8_preq is always M=1 */
|
|
#define FM_ELEM_THRESHOLD (1024 * 1024) /* 1M elements: only batch softmax/norm */
|
|
|
|
/* ================================================================
|
|
* Synchronization & Command batching
|
|
*
|
|
* GPU dispatches auto-accumulate in a persistent command buffer
|
|
* (MLX-style). Call fm_sync() when CPU needs GPU results.
|
|
* ================================================================ */
|
|
|
|
/* Commit pending GPU work and wait for completion.
|
|
* Call this at data dependency points where CPU needs GPU output. */
|
|
void fm_sync(void);
|
|
|
|
/* Enable/disable auto-batch mode.
|
|
* When enabled, GPU dispatches accumulate until fm_sync() is called.
|
|
* When disabled (default), each dispatch commits and waits immediately.
|
|
* Inference engines should enable this for performance. */
|
|
void fm_set_auto_batch(bool enable);
|
|
|
|
/* Begin accumulating GPU commands into a shared command buffer.
|
|
* All fm_* calls between begin/end encode without committing. */
|
|
int fm_batch_begin(void);
|
|
|
|
/* Commit the batched command buffer and wait for completion. */
|
|
void fm_batch_end(void);
|
|
|
|
/* Check if command batching is active. */
|
|
bool fm_batch_active(void);
|
|
|
|
#ifdef __cplusplus
|
|
}
|
|
#endif
|
|
|
|
/* ================================================================
|
|
* Kernel families
|
|
*
|
|
* One header per unit under metal/kernels/, included last so the types
|
|
* above are already declared. Each is self-sufficient (it includes this
|
|
* header back, which the guard makes a no-op), so a caller that wants
|
|
* only one family can include just that header.
|
|
*
|
|
* Adding a family: create the .h/.m pair, register its MSL fragment in
|
|
* metal/kernels/fm_kernel_registry.{h,m}, add it to METAL_SRC in the
|
|
* Makefile, and add one line here.
|
|
* ================================================================ */
|
|
|
|
#include "metal/kernels/fm_kernel_convert.h"
|
|
#include "metal/kernels/fm_kernel_elementwise.h"
|
|
#include "metal/kernels/fm_kernel_gemm_f32.h"
|
|
#include "metal/kernels/fm_kernel_matvec_q3_k.h"
|
|
#include "metal/kernels/fm_kernel_matvec_q4_k.h"
|
|
#include "metal/kernels/fm_kernel_matvec_q5_k.h"
|
|
#include "metal/kernels/fm_kernel_matvec_q6_k.h"
|
|
#include "metal/kernels/fm_kernel_matvec_q8.h"
|
|
#include "metal/kernels/fm_kernel_matvec_q8_0.h"
|
|
#include "metal/kernels/fm_kernel_norm.h"
|
|
#include "metal/kernels/fm_kernel_probe.h"
|
|
#include "metal/kernels/fm_kernel_rope.h"
|
|
#include "metal/kernels/fm_kernel_ssm.h"
|
|
#include "metal/kernels/fm_kernel_ssm_chunked.h"
|
|
#include "metal/kernels/fm_kernel_moe.h"
|
|
#include "metal/kernels/fm_kernel_attn.h"
|
|
#include "metal/kernels/fm_kernel_wkv7.h"
|
|
|
|
#endif /* FOUNDRY_METAL_H */
|