100 lines
3.7 KiB
C
100 lines
3.7 KiB
C
/* SPDX-License-Identifier: AGPL-3.0-or-later */
|
|
/* Copyright (C) 2026 David (Corey) Herrera */
|
|
/**
|
|
* @file fm_kernel_elementwise.h
|
|
* @brief Elementwise family: activations and binary ops over flat float32
|
|
* arrays. Bandwidth-bound, one thread per element (scalar form) or per
|
|
* four elements (`_4` float4 form).
|
|
*
|
|
* @status Complete: silu, sigmoid, add, mul, each in a scalar and a float4
|
|
* form. The float4 forms accept any n — the last thread handles a 1-3
|
|
* element tail scalar-wise.
|
|
* @issues None
|
|
* @todo None
|
|
*
|
|
* FUNCTION-CONSTANT INDICES OWNED BY THIS FAMILY: 120-129 (none used —
|
|
* these kernels have no shape to specialize on).
|
|
* Function constants share ONE index space across the whole assembled
|
|
* shader library, so each family reserves a block and two families
|
|
* colliding on an index fails the compile of every kernel in the backend,
|
|
* not just theirs. Allocated so far: 0-2 probe · 30 Q3_K · 50 Q5_K ·
|
|
* 60 Q6_K · 70/80 Q8_0+Q4_K · 100-109 rope · 110-119 norm ·
|
|
* 120-129 elementwise.
|
|
*/
|
|
|
|
#ifndef FOUNDRY_METAL_KERNEL_ELEMENTWISE_H
|
|
#define FOUNDRY_METAL_KERNEL_ELEMENTWISE_H
|
|
|
|
#include "metal/fm_metal.h"
|
|
|
|
#ifdef __cplusplus
|
|
extern "C" {
|
|
#endif
|
|
|
|
/* ================================================================
|
|
* Scalar forms — one element per thread.
|
|
* ================================================================ */
|
|
|
|
/**
|
|
* SiLU activation: out[i] = x[i] / (1 + exp(-x[i])).
|
|
*
|
|
* @param x Input, n floats
|
|
* @param out Output, n floats (may alias x)
|
|
* @param n Element count
|
|
*
|
|
* Thread-safe: no (encodes into the global command buffer)
|
|
* Reentrant: no
|
|
* Signal-safe: no
|
|
*/
|
|
void fm_silu(fm_buffer *x, fm_buffer *out, int n);
|
|
|
|
/**
|
|
* Logistic sigmoid: out[i] = 1 / (1 + exp(-x[i])).
|
|
*
|
|
* @param x Input, n floats
|
|
* @param out Output, n floats (may alias x)
|
|
* @param n Element count
|
|
*
|
|
* Thread-safe: no
|
|
* Reentrant: no
|
|
* Signal-safe: no
|
|
*/
|
|
void fm_sigmoid(fm_buffer *x, fm_buffer *out, int n);
|
|
|
|
/* Elementwise binary ops: out[i] = a[i] (+|*) b[i], n floats each. */
|
|
void fm_add(fm_buffer *a, fm_buffer *b, fm_buffer *out, int n);
|
|
void fm_mul(fm_buffer *a, fm_buffer *b, fm_buffer *out, int n);
|
|
|
|
/* ================================================================
|
|
* float4 forms — four elements per thread.
|
|
*
|
|
* Identical results to the scalar forms, a quarter of the load/store
|
|
* instructions. n need not be a multiple of 4: the grid is rounded up and
|
|
* the final thread falls back to a scalar loop over the 1-3 leftover
|
|
* elements. Every buffer is bound at offset 0 and a Metal allocation is
|
|
* page aligned, so the 16-byte alignment a float4 access needs holds by
|
|
* construction — there is no alignment precondition on the caller.
|
|
*
|
|
* MEASURED (M4, 10 GPU cores, macOS 26): these are NOT faster than the
|
|
* scalar forms — 3.8 vs 4.8 GB/s at n=2048, 74.6 vs 78.9 at n=65536,
|
|
* 117.2 vs 117.4 at n=1M. Fewer, wider loads buy nothing at either end:
|
|
* past ~1M elements the scalar form already sits on the ~120 GB/s memory
|
|
* roof, and below that the kernel is launch- and occupancy-bound, where
|
|
* a quarter as many threads means a quarter as many threadgroups and so
|
|
* fewer of the ten cores fed. Prefer the scalar forms; these exist
|
|
* because the same reduction may pay on a wider GPU, and they are kept
|
|
* correct and tested for that day. Do not reach for them as a speed
|
|
* lever.
|
|
* ================================================================ */
|
|
|
|
void fm_silu_4(fm_buffer *x, fm_buffer *out, int n);
|
|
void fm_sigmoid_4(fm_buffer *x, fm_buffer *out, int n);
|
|
void fm_add_4(fm_buffer *a, fm_buffer *b, fm_buffer *out, int n);
|
|
void fm_mul_4(fm_buffer *a, fm_buffer *b, fm_buffer *out, int n);
|
|
|
|
#ifdef __cplusplus
|
|
}
|
|
#endif
|
|
|
|
#endif /* FOUNDRY_METAL_KERNEL_ELEMENTWISE_H */
|