luboslenco
323 lines
12 KiB
C
323 lines
12 KiB
C
/*
|
||
* Iris Math Kernels - Header
|
||
*
|
||
* Low-level math operations for the Iris inference engine.
|
||
* All operations work on float32 tensors in row-major order.
|
||
*/
|
||
|
||
#ifndef IRIS_KERNELS_H
|
||
#define IRIS_KERNELS_H
|
||
|
||
#include <math.h>
|
||
#include <stddef.h>
|
||
#include <stdint.h>
|
||
|
||
/* Forward declarations */
|
||
struct iris_image;
|
||
|
||
/* Fast exponential approximation using range reduction + degree-5 polynomial.
|
||
* Relative error < 2e-6 across the full float range.
|
||
* Inlined so the compiler can auto-vectorize loops that use it. */
|
||
static inline float fast_expf(float x) {
|
||
if (x < -87.3f)
|
||
return 0.0f;
|
||
if (x > 88.7f)
|
||
return 1e38f;
|
||
float n = floorf(x * 1.4426950408889634f + 0.5f);
|
||
float r = x - n * 0.6931471805599453f;
|
||
float p = 1.0f + r * (1.0f + r * (0.5f + r * (0.16666667f + r * (0.04166667f + r * 0.00833333f))));
|
||
union {
|
||
float f;
|
||
int32_t i;
|
||
} v;
|
||
v.f = p;
|
||
v.i += (int32_t)n << 23;
|
||
return v.f;
|
||
}
|
||
|
||
/* ========================================================================
|
||
* Basic Operations
|
||
* ======================================================================== */
|
||
|
||
/* Element-wise operations */
|
||
void iris_add(float *out, const float *a, const float *b, int n);
|
||
|
||
/* In-place variants */
|
||
void iris_add_inplace(float *a, const float *b, int n);
|
||
void iris_mul_inplace(float *a, const float *b, int n);
|
||
|
||
/* Accumulate: a += scale * b */
|
||
void iris_axpy(float *a, float scale, const float *b, int n);
|
||
|
||
/* ========================================================================
|
||
* Matrix Operations
|
||
* ======================================================================== */
|
||
|
||
/*
|
||
* General matrix multiplication: C = A @ B
|
||
* A: [M, K], B: [K, N], C: [M, N]
|
||
*/
|
||
void iris_matmul(float *C, const float *A, const float *B, int M, int K, int N);
|
||
|
||
/*
|
||
* Matrix multiplication with transposed B: C = A @ B^T
|
||
* A: [M, K], B: [N, K], C: [M, N]
|
||
*/
|
||
void iris_matmul_t(float *C, const float *A, const float *B, int M, int K, int N);
|
||
|
||
/*
|
||
* Linear layer: y = x @ W^T + b (if b != NULL)
|
||
* x: [seq_len, in_dim], W: [out_dim, in_dim], b: [out_dim], y: [seq_len, out_dim]
|
||
*/
|
||
void iris_linear(float *y, const float *x, const float *W, const float *b, int seq_len, int in_dim, int out_dim);
|
||
|
||
/*
|
||
* Linear layer without bias: y = x @ W^T
|
||
*/
|
||
void iris_linear_nobias(float *y, const float *x, const float *W, int seq_len, int in_dim, int out_dim);
|
||
|
||
/*
|
||
* Linear layer without bias using bf16 weights
|
||
* x: [seq_len, in_dim] (f32), W: [out_dim, in_dim] (bf16), y: [seq_len, out_dim] (f32)
|
||
*/
|
||
void iris_linear_nobias_bf16(float *y, const float *x, const uint16_t *W_bf16, int seq_len, int in_dim, int out_dim);
|
||
|
||
/*
|
||
* Linear layer without bias using GGML Q8_0 block-quantized weights.
|
||
* x: [seq_len, in_dim] (f32), W: [out_dim, in_dim] (Q8_0 blocks), y: [seq_len, out_dim] (f32).
|
||
* Accelerated on the Vulkan backend (in-shader dequant); CPU dequant fallback
|
||
* otherwise. W points at the raw 34-byte-per-32-element block stream.
|
||
*/
|
||
void iris_linear_nobias_q8(float *y, const float *x, const void *W_q8, int seq_len, int in_dim, int out_dim);
|
||
|
||
/* ========================================================================
|
||
* GPU Batch Operations
|
||
* These functions allow batching multiple GPU operations to reduce sync overhead.
|
||
* On non-GPU builds, these are no-ops.
|
||
* ======================================================================== */
|
||
|
||
/*
|
||
* Begin a batch of GPU operations.
|
||
* Operations after this call are queued but not executed until iris_gpu_end_batch().
|
||
* NOTE: Only use for INDEPENDENT operations (outputs don't feed into subsequent inputs).
|
||
*/
|
||
void iris_gpu_begin_batch(void);
|
||
|
||
/*
|
||
* End a batch of GPU operations.
|
||
* Executes all queued operations and waits for completion.
|
||
*/
|
||
void iris_gpu_end_batch(void);
|
||
|
||
/* ========================================================================
|
||
* Convolution Operations
|
||
* ======================================================================== */
|
||
|
||
/*
|
||
* 2D Convolution: out = conv2d(in, weight, bias)
|
||
* in: [batch, in_ch, H, W]
|
||
* weight: [out_ch, in_ch, kH, kW]
|
||
* bias: [out_ch] (can be NULL)
|
||
* out: [batch, out_ch, outH, outW]
|
||
*/
|
||
void iris_conv2d(float *out, const float *in, const float *weight, const float *bias, int batch, int in_ch, int out_ch, int H, int W, int kH, int kW,
|
||
int stride, int padding);
|
||
|
||
/* ========================================================================
|
||
* Normalization
|
||
* ======================================================================== */
|
||
|
||
/*
|
||
* RMS Normalization (no mean centering, no bias)
|
||
* x: [seq_len, hidden], weight: [hidden]
|
||
*/
|
||
void iris_rms_norm(float *out, const float *x, const float *weight, int seq_len, int hidden, float eps);
|
||
|
||
/*
|
||
* Group Normalization
|
||
* x: [batch, channels, H, W], gamma/beta: [channels]
|
||
*/
|
||
void iris_group_norm(float *out, const float *x, const float *gamma, const float *beta, int batch, int channels, int H, int W, int num_groups, float eps);
|
||
|
||
/*
|
||
* Batch Normalization (inference mode with running stats)
|
||
* x: [batch, channels, H, W]
|
||
*/
|
||
void iris_batch_norm(float *out, const float *x, const float *running_mean, const float *running_var, const float *gamma, const float *beta, int batch,
|
||
int channels, int H, int W, float eps);
|
||
|
||
/* ========================================================================
|
||
* Activation Functions
|
||
* ======================================================================== */
|
||
|
||
/* SiLU / Swish activation: x * sigmoid(x) */
|
||
void iris_silu(float *x, int n);
|
||
|
||
/* Fused SiLU(gate) * up - single pass for SwiGLU */
|
||
void iris_silu_mul(float *gate, const float *up, int n);
|
||
|
||
/* Softmax over last dimension */
|
||
void iris_softmax(float *x, int rows, int cols);
|
||
void iris_softmax_cpu(float *x, int rows, int cols);
|
||
|
||
/* ========================================================================
|
||
* Attention Operations
|
||
* ======================================================================== */
|
||
|
||
/*
|
||
* Scaled dot-product attention
|
||
* Q: [batch, heads, seq_q, head_dim]
|
||
* K: [batch, heads, seq_k, head_dim]
|
||
* V: [batch, heads, seq_k, head_dim]
|
||
* out: [batch, heads, seq_q, head_dim]
|
||
* scale: typically 1/sqrt(head_dim)
|
||
*/
|
||
void iris_attention(float *out, const float *Q, const float *K, const float *V, int batch, int heads, int seq_q, int seq_k, int head_dim, float scale);
|
||
|
||
/*
|
||
* Flash attention - memory-efficient tiled attention.
|
||
* Uses online softmax to avoid materializing O(n²) attention matrix.
|
||
* Memory: O(seq_q + tile_size²) instead of O(seq_q × seq_k).
|
||
*
|
||
* Works on [seq, heads*head_dim] layout (same as transformer tensors).
|
||
* Q: [seq_q, heads * head_dim]
|
||
* K: [seq_k, heads * head_dim]
|
||
* V: [seq_k, heads * head_dim]
|
||
* out: [seq_q, heads * head_dim]
|
||
*/
|
||
void iris_flash_attention(float *out, const float *Q, const float *K, const float *V, int seq_q, int seq_k, int heads, int head_dim, float scale);
|
||
|
||
/*
|
||
* Apply rotary position embeddings (RoPE)
|
||
* x: [batch, seq, heads, head_dim]
|
||
* freqs: [seq, head_dim/2, 2] (cos, sin pairs)
|
||
*/
|
||
void iris_apply_rope(float *x, const float *freqs, int batch, int seq, int heads, int head_dim);
|
||
|
||
/*
|
||
* Compute RoPE frequencies
|
||
* pos: position indices [seq]
|
||
* freqs: output [seq, dim/2, 2]
|
||
*/
|
||
void iris_compute_rope_freqs(float *freqs, const int *pos, int seq, int dim, float theta);
|
||
|
||
/* ========================================================================
|
||
* Pooling and Reshape
|
||
* ======================================================================== */
|
||
|
||
/* Upsample with nearest neighbor */
|
||
void iris_upsample_nearest(float *out, const float *in, int batch, int channels, int H, int W, int scale_h, int scale_w);
|
||
|
||
/* Patchify: [B, C, H, W] -> [B, C*p*p, H/p, W/p] */
|
||
void iris_patchify(float *out, const float *in, int batch, int channels, int H, int W, int patch_size);
|
||
|
||
/* Unpatchify: [B, C*p*p, H, W] -> [B, C, H*p, W*p] */
|
||
void iris_unpatchify(float *out, const float *in, int batch, int channels, int H, int W, int patch_size);
|
||
|
||
/* ========================================================================
|
||
* Random Number Generation
|
||
* ======================================================================== */
|
||
|
||
/* Initialize RNG with seed */
|
||
void iris_rng_seed(uint64_t seed);
|
||
|
||
/* Generate uniform random [0, 1) */
|
||
float iris_random_uniform(void);
|
||
|
||
/* Generate standard normal using Box-Muller */
|
||
float iris_random_normal(void);
|
||
|
||
/* Fill tensor with random normal values */
|
||
void iris_randn(float *out, int n);
|
||
|
||
/* Fill tensor with uniform random [0, 1) */
|
||
void iris_rand(float *out, int n);
|
||
|
||
/* ========================================================================
|
||
* Utility Functions
|
||
* ======================================================================== */
|
||
|
||
/* Copy tensor */
|
||
void iris_copy(float *dst, const float *src, int n);
|
||
|
||
/* ========================================================================
|
||
* Progress Callbacks
|
||
* ======================================================================== */
|
||
|
||
/* Substep types during transformer forward pass */
|
||
typedef enum {
|
||
IRIS_SUBSTEP_DOUBLE_BLOCK, /* Double-stream block completed */
|
||
IRIS_SUBSTEP_SINGLE_BLOCK, /* Single-stream block completed */
|
||
IRIS_SUBSTEP_FINAL_LAYER, /* Final layer completed */
|
||
} iris_substep_type_t;
|
||
|
||
/*
|
||
* Substep callback - called during transformer forward pass.
|
||
* type: which operation completed
|
||
* index: 0-based index of this substep within its type
|
||
* total: total count for this substep type
|
||
*/
|
||
typedef void (*iris_substep_callback_t)(iris_substep_type_t type, int index, int total);
|
||
|
||
/*
|
||
* Step callback - called at sampling step boundaries.
|
||
* step: current step (1-based), or 0 to indicate sampling is starting
|
||
* total: total number of steps
|
||
*/
|
||
typedef void (*iris_step_callback_t)(int step, int total);
|
||
|
||
/* Global callback pointers - set by caller before inference */
|
||
extern iris_substep_callback_t iris_substep_callback;
|
||
extern iris_step_callback_t iris_step_callback;
|
||
|
||
/*
|
||
* Phase callback - called at major phase boundaries.
|
||
* phase: descriptive name ("encoding text", "decoding image", etc.)
|
||
* done: 0 when starting, 1 when finished
|
||
*/
|
||
typedef void (*iris_phase_callback_t)(const char *phase, int done);
|
||
extern iris_phase_callback_t iris_phase_callback;
|
||
|
||
/*
|
||
* Step image callback - called after each denoising step with decoded image.
|
||
* step: current step (1-based)
|
||
* total: total number of steps
|
||
* img: decoded image at this step (caller must NOT free)
|
||
*
|
||
* To use: set both iris_step_image_callback and iris_step_image_vae before
|
||
* calling the sampling function. The callback is only invoked when both are set.
|
||
*/
|
||
typedef void (*iris_step_image_callback_t)(int step, int total, const struct iris_image *img);
|
||
extern iris_step_image_callback_t iris_step_image_callback;
|
||
extern void *iris_step_image_vae; /* Set to iris_vae_t* for step image decoding */
|
||
|
||
/*
|
||
* Text encoder progress callback - called once per Qwen3 layer.
|
||
* layer: current layer (0-based)
|
||
* total: total number of layers (36)
|
||
*/
|
||
typedef void (*iris_text_progress_callback_t)(int layer, int total);
|
||
extern iris_text_progress_callback_t iris_text_progress_callback;
|
||
|
||
/*
|
||
* VAE progress callback - called once per resblock/attention block.
|
||
* block: current block (0-based)
|
||
* total: total number of blocks (11 for encoder, 15 for decoder)
|
||
*/
|
||
typedef void (*iris_vae_progress_callback_t)(int block, int total);
|
||
extern iris_vae_progress_callback_t iris_vae_progress_callback;
|
||
|
||
/* Global verbose flag - when 0, library code suppresses diagnostic output */
|
||
extern int iris_verbose;
|
||
|
||
/* Global circular flag - when 1, convolutions use circular (toroidal wrap-around)
|
||
* padding instead of zero padding, producing seamless / tileable images. */
|
||
extern int iris_circular;
|
||
|
||
/* Global VAE tiling flag - when 1, the VAE decoder splits the latent into
|
||
* smaller overlapping tiles, decodes each separately and blends them, trading
|
||
* a little speed for much lower peak (GPU) memory on large images. */
|
||
extern int iris_vae_tiling;
|
||
|
||
#endif /* IRIS_KERNELS_H */
|