1// RAID Metal kernels — FP32 with double-float ("compensated") accumulation.2// Apple GPUs have no FP64. Every accumulator is a pair (hi, lo) maintained with error-free transforms (TwoSum /3// Fast2Sum); partial results per (series, chunk) are written as float2 and combined in FP64 on the CPU in a fixed4// order, which makes results deterministic for a given chunk size and threadgroup width.5#include <metal_stdlib>6using namespace metal;78// ---- double-float arithmetic -----------------------------------------------------------------------------------9struct df { float hi; float lo; };10inline df df_zero() { return df{0.0f, 0.0f}; }11inline df two_sum(float a, float b) { float s = a + b; float bb = s - a; float e = (a - (s - bb)) + (b - bb); return df{s, e}; }12inline df df_add_f(df a, float b) { df s = two_sum(a.hi, b); s.lo += a.lo; float h = s.hi + s.lo; float l = s.lo - (h - s.hi); return df{h, l}; }13inline df df_add(df a, df b) { df s = two_sum(a.hi, b.hi); s.lo += a.lo + b.lo; float h = s.hi + s.lo; float l = s.lo - (h - s.hi); return df{h, l}; }14inline float2 df_pack(df a) { return float2(a.hi, a.lo); }1516// Threadgroup tree reduction of double-float values (deterministic order).17inline df tg_reduce(df v, threadgroup df* sh, uint lid, uint n) {18 sh[lid] = v; threadgroup_barrier(mem_flags::mem_threadgroup);19 for (uint s = n / 2; s > 0; s >>= 1) { if (lid < s) sh[lid] = df_add(sh[lid], sh[lid + s]); threadgroup_barrier(mem_flags::mem_threadgroup); }20 return sh[0];21}2223// ---- 1. fused realized-measure accumulators ---------------------------------------------------------------------24// One threadgroup per (segment, chunk). Inputs: x (float), offsets (int64 as two uints packed → we use uint offsets,25// the host guarantees < 2^32 elements), chunk descriptors (segment id, begin, end). Outputs per chunk: 5 float2:26// Σx², Σ|x_i||x_{i-1}| (pairs with i-1 inside the chunk; the host adds the boundary products), Σx²1{x<0}, Σx²1{x>0}, Σx⁴.27struct Chunk { uint seg; uint begin; uint end; uint pad; };28kernel void realized_fused(device const float* x [[buffer(0)]], device const Chunk* chunks [[buffer(1)]], device float2* out [[buffer(2)]],29 uint gid [[threadgroup_position_in_grid]], uint lid [[thread_position_in_threadgroup]], uint tsz [[threads_per_threadgroup]]) {30 threadgroup df sh[256];31 Chunk c = chunks[gid];32 df s2 = df_zero(), sadj = df_zero(), sneg = df_zero(), spos = df_zero(), s4 = df_zero();33 for (uint i = c.begin + lid; i < c.end; i += tsz) {34 float v = x[i]; float q = v * v;35 s2 = df_add_f(s2, q); s4 = df_add_f(s4, q * q);36 if (v < 0.0f) sneg = df_add_f(sneg, q); else if (v > 0.0f) spos = df_add_f(spos, q);37 if (i > c.begin) sadj = df_add_f(sadj, fabs(v) * fabs(x[i - 1]));38 }39 df r2 = tg_reduce(s2, sh, lid, tsz); threadgroup_barrier(mem_flags::mem_threadgroup);40 df radj = tg_reduce(sadj, sh, lid, tsz); threadgroup_barrier(mem_flags::mem_threadgroup);41 df rneg = tg_reduce(sneg, sh, lid, tsz); threadgroup_barrier(mem_flags::mem_threadgroup);42 df rpos = tg_reduce(spos, sh, lid, tsz); threadgroup_barrier(mem_flags::mem_threadgroup);43 df r4 = tg_reduce(s4, sh, lid, tsz);44 if (lid == 0) { out[gid * 5 + 0] = df_pack(r2); out[gid * 5 + 1] = df_pack(radj); out[gid * 5 + 2] = df_pack(rneg); out[gid * 5 + 3] = df_pack(rpos); out[gid * 5 + 4] = df_pack(r4); }45}4647// Single-statistic variants (for the fusion ablation): op 0 = Σx², 1 = Σ|x||x_{-1}|, 2 = Σx² 1{x<0}, 3 = Σx² 1{x>0}, 4 = Σx⁴48kernel void segment_reduce(device const float* x [[buffer(0)]], device const Chunk* chunks [[buffer(1)]], device float2* out [[buffer(2)]], constant uint& op [[buffer(3)]],49 uint gid [[threadgroup_position_in_grid]], uint lid [[thread_position_in_threadgroup]], uint tsz [[threads_per_threadgroup]]) {50 threadgroup df sh[256];51 Chunk c = chunks[gid]; df acc = df_zero();52 for (uint i = c.begin + lid; i < c.end; i += tsz) {53 float v = x[i]; float q = v * v; float t = 0.0f;54 switch (op) { case 0: t = q; break; case 1: t = (i > c.begin) ? fabs(v) * fabs(x[i - 1]) : 0.0f; break; case 2: t = v < 0.0f ? q : 0.0f; break; case 3: t = v > 0.0f ? q : 0.0f; break; default: t = q * q; }55 acc = df_add_f(acc, t);56 }57 df r = tg_reduce(acc, sh, lid, tsz);58 if (lid == 0) out[gid] = df_pack(r);59}6061// ---- 2. rolling sum of squares (direct window, O(w) per output) ---------------------------------------------------62// out[i] = Σ_{j=i-w+1..i} x_j² within the segment; NaN where the window is incomplete. One thread per element.63kernel void rolling_sumsq(device const float* x [[buffer(0)]], device const uint* seg_begin [[buffer(1)]], device const uint* seg_of [[buffer(2)]], device float* out [[buffer(3)]],64 constant uint& w [[buffer(4)]], constant uint& n [[buffer(5)]], uint i [[thread_position_in_grid]]) {65 if (i >= n) return;66 uint b = seg_begin[seg_of[i]];67 if (i + 1 < b + w) { out[i] = NAN; return; }68 df acc = df_zero();69 for (uint j = i + 1 - w; j <= i; ++j) { float v = x[j]; acc = df_add_f(acc, v * v); }70 out[i] = acc.hi + acc.lo;71}7273// ---- 3. regression sufficient statistics (X'X upper triangle, X'y, y'y) per chunk ------------------------------74// X row-major (rows × k), k ≤ 8. One threadgroup per chunk; each thread accumulates its rows in double-float registers,75// then a deterministic tree reduction per statistic.76#define RAID_MAXK 877kernel void regression_stats(device const float* X [[buffer(0)]], device const float* y [[buffer(1)]], device const Chunk* chunks [[buffer(2)]], device float2* out [[buffer(3)]],78 constant uint& k [[buffer(4)]], uint gid [[threadgroup_position_in_grid]], uint lid [[thread_position_in_threadgroup]], uint tsz [[threads_per_threadgroup]]) {79 threadgroup df sh[256];80 Chunk c = chunks[gid];81 const uint ntri = k * (k + 1) / 2; const uint nstat = ntri + k + 1;82 df acc[RAID_MAXK * (RAID_MAXK + 1) / 2 + RAID_MAXK + 1];83 for (uint s = 0; s < nstat; ++s) acc[s] = df_zero();84 for (uint r = c.begin + lid; r < c.end; r += tsz) {85 float xr[RAID_MAXK]; for (uint a = 0; a < k; ++a) xr[a] = X[r * k + a];86 float yr = y[r]; uint s = 0;87 for (uint a = 0; a < k; ++a) for (uint b = a; b < k; ++b) { acc[s] = df_add_f(acc[s], xr[a] * xr[b]); ++s; }88 for (uint a = 0; a < k; ++a) { acc[s] = df_add_f(acc[s], xr[a] * yr); ++s; }89 acc[s] = df_add_f(acc[s], yr * yr);90 }91 for (uint s = 0; s < nstat; ++s) { df r = tg_reduce(acc[s], sh, lid, tsz); if (lid == 0) out[gid * nstat + s] = df_pack(r); threadgroup_barrier(mem_flags::mem_threadgroup); }92}9394// ---- 4. batched GARCH-family negative log-likelihood ---------------------------------------------------------------95// One thread per (series, parameter vector). params layout per thread (stride 8): [mu, omega, alpha, gamma, beta, nu,96// kconst, eabs] — gamma=0 for GARCH, nu<=2 → Gaussian; kconst = density constant and eabs = E|z| are computed in FP6497// on the host (Metal has no lgamma). var0 = FP64-computed initial variance passed from the host (mixed precision).98// model: 0 garch, 1 gjr, 2 egarch. Returns NLL (double-float summed) as float; +INF for invalid parameters.99struct GarchJob { uint seg; uint pidx; uint pad0; uint pad1; };100kernel void garch_nll(device const float* r [[buffer(0)]], device const uint* seg_begin [[buffer(1)]], device const uint* seg_end [[buffer(2)]], device const float* params [[buffer(3)]],101 device const float* var0 [[buffer(4)]], device const GarchJob* jobs [[buffer(5)]], device float* out [[buffer(6)]], constant uint& model [[buffer(7)]], constant uint& njobs [[buffer(8)]],102 uint t [[thread_position_in_grid]]) {103 if (t >= njobs) return;104 GarchJob j = jobs[t]; const uint b = seg_begin[j.seg], e = seg_end[j.seg]; const uint n = e - b;105 const float mu = params[j.pidx * 8 + 0], omega = params[j.pidx * 8 + 1], alpha = params[j.pidx * 8 + 2], gamma = params[j.pidx * 8 + 3], beta = params[j.pidx * 8 + 4], nu = params[j.pidx * 8 + 5];106 const float kconst = params[j.pidx * 8 + 6], Eabs = params[j.pidx * 8 + 7];107 const bool student = nu > 2.0f;108 bool valid = true;109 if (model == 0u) valid = omega > 0.0f && alpha >= 0.0f && beta >= 0.0f && alpha + beta < 1.0f;110 else if (model == 1u) valid = omega > 0.0f && alpha >= 0.0f && beta >= 0.0f && alpha + gamma >= 0.0f && alpha + beta + 0.5f * gamma < 1.0f;111 else valid = fabs(beta) < 1.0f;112 if (!valid || n < 2) { out[t] = INFINITY; return; }113 float s2 = var0[j.seg];114 df nll = df_zero();115 float e_prev = r[b] - mu;116 for (uint i = 0; i < n; ++i) {117 if (i > 0) {118 if (model == 2u) { float z = e_prev / sqrt(s2); float ls = omega + alpha * (fabs(z) - Eabs) + gamma * z + beta * log(s2); s2 = exp(ls); }119 else { float e2 = e_prev * e_prev; s2 = omega + (alpha + (e_prev < 0.0f ? gamma : 0.0f)) * e2 + beta * s2; }120 if (!(s2 > 0.0f) || !isfinite(s2)) { out[t] = INFINITY; return; }121 }122 float ev = r[b + i] - mu; float ll;123 if (student) ll = kconst - 0.5f * log(s2) - 0.5f * (nu + 1.0f) * log(1.0f + ev * ev / ((nu - 2.0f) * s2)); // no log1p in MSL; argument is O(1)124 else ll = kconst - 0.5f * log(s2) - 0.5f * ev * ev / s2;125 nll = df_add_f(nll, -ll);126 e_prev = ev;127 }128 out[t] = nll.hi + nll.lo;129}130