// tensor_gpu_lib.m — Metal tensor ops as a shared library (v2)
//
// Build:
//   clang -shared -fobjc-arc -framework Metal -framework Foundation \
//     -install_name ~/projects/rail/tools/metal/libtensor_gpu.dylib \
//     tensor_gpu_lib.m -o libtensor_gpu.dylib
//
// Exports (C ABI, all return 1 on success / -1 on failure):
//   tgl_init(int)                                       — idempotent init
//   tgl_matmul_f64       (A,B,C, M,K,N)                 — C = A @ B
//   tgl_matmul_relu_f64  (A,B,bias,C, M,K,N)            — C = relu(A@B + bias)  (fused)
//   tgl_add_f64          (A,B,C, N)                     — C = A + B
//   tgl_mul_f64          (A,B,C, N)                     — C = A * B  (hadamard)
//   tgl_scale_f64        (A,C, scalar, N)               — C = A * scalar        (scalar is first 8 bytes of a double*)
//   tgl_relu_f64         (X,Y, N)                       — Y = max(0,X)
//   tgl_relu_backward_f64(X,grad,out, N)                — out = (X>0) ? grad : 0
//   tgl_sigmoid_f64      (X,Y, N)
//   tgl_exp_f64          (X,Y, N)
//   tgl_tanh_f64         (X,Y, N)
//   tgl_softmax_rows_f64 (X,Y, rows, cols)              — row-wise softmax
//   tgl_transpose_f64    (A,B, M,N)                     — B = A^T, shape N×M
//   tgl_sgd_update_f64   (w, grad, lr, N)               — w -= lr*grad (in place)
//   tgl_adam_update_f64  (w, g, m, v, hyp, N)           — fused Adam; hyp=[lr,β1,β2,ε,bc1,bc2], m/v updated in place
//   tgl_cross_entropy_f64(probs, targets_f64, losses, batch, vocab)
//
// All *_f64 entry points accept pointers into Rail float_arr payload.
// Rail float_arr layout: [count@0, double@8, double@16, ...].
// We offset by +1 double so A actually points at data, not the count header.

#import <Foundation/Foundation.h>
#import <Metal/Metal.h>
#include <stdio.h>
#include <string.h>
#include <dlfcn.h>
#include <sys/time.h>

static id<MTLDevice>               g_device = nil;
static id<MTLCommandQueue>         g_queue = nil;
static id<MTLLibrary>              g_lib = nil;
static int g_initialized = 0;

// All pipeline states, keyed by name; lazy-built.
static NSMutableDictionary<NSString*, id<MTLComputePipelineState>> *g_pipes = nil;

// ────────────────────────────────────────────────────────────────────
// Tiered MTLBuffer pool — track 14
//
// Rationale (see docs/metal-pool-design.md): the prior single-pool design
// (POOL_MAX=32, one bucket) let many distinct small (M,K,N) shapes crowd
// out the few large shapes a transformer emits at seq=1499 (attention
// scores ≈ 9 MB).  When the pool was full, large allocations dropped on
// the floor and re-allocated per call — the documented RSS leak.
//
// We bucket by byte-size class, give each tier its own dynamic capacity,
// and let large tiers grow without evicting small ones.  Storage mode is
// part of the key so a Shared buffer never pretends to be Private.
//
// Tier boundaries are powers of two; T0 ≤ 4 KB, then ×16 / ×16 / ×16 / ∞.
// All three storage modes get parallel sub-pools (currently only Shared
// is used by these kernels — the other two stay zero-allocated until a
// caller asks).
// ────────────────────────────────────────────────────────────────────

#define NUM_TIERS         5
#define NUM_STORAGE_MODES 3        // Shared / Managed / Private
#define TIER_INITIAL_CAP  8
#define TIER_MAX_CAP      128
#define TIER_IDLE_SHRINK_SECS 30.0
#define TIER_MIN_KEEP     4        // never shrink a tier below this many slots

static const NSUInteger TIER_MAX_BYTES[NUM_TIERS] = {
    (NSUInteger)4096,                  // T0 ≤ 4 KB
    (NSUInteger)65536,                 // T1 ≤ 64 KB
    (NSUInteger)1048576,               // T2 ≤ 1 MB
    (NSUInteger)16777216,              // T3 ≤ 16 MB
    (NSUInteger)0xFFFFFFFFFFFFFFFFULL, // T4 unbounded
};

typedef struct {
    id<MTLBuffer> buf;        // retained while slot occupied
    NSUInteger    size;       // bytes (rounded)
    int           in_use;
    double        last_release_t;
} pool_slot_t;

typedef struct {
    pool_slot_t  slots[TIER_MAX_CAP];
    int          count;          // how many slots are populated
    int          capacity;       // current grown capacity (≤ TIER_MAX_CAP)
    int          in_use;         // currently checked-out slots
    int          peak_in_use;    // high-water mark since boot
    long long    hits;           // satisfied from pool
    long long    misses;         // newBufferWithLength fired
    long long    drops;          // misses that couldn't even be tracked
    long long    shrinks;        // slots freed by idle reaper
    long long    bytes_in_pool;  // sum of slot.size (in_use + idle)
} pool_tier_t;

static pool_tier_t g_pools[NUM_STORAGE_MODES][NUM_TIERS];
static int         g_pools_inited = 0;
static int         g_pool_disable = 0;  // RAIL_GPU_POOL_DISABLE=1 → always miss

static dispatch_once_t g_pool_once;

static void pool_init_once(void) {
    dispatch_once(&g_pool_once, ^{
        for (int m = 0; m < NUM_STORAGE_MODES; m++)
            for (int t = 0; t < NUM_TIERS; t++)
                g_pools[m][t].capacity = TIER_INITIAL_CAP;
        const char *d = getenv("RAIL_GPU_POOL_DISABLE");
        if (d && d[0] && d[0] != '0') {
            g_pool_disable = 1;
            fprintf(stderr, "tgl: RAIL_GPU_POOL_DISABLE=1 (pool best-fit bypassed)\n");
        }
        g_pools_inited = 1;
    });
}

// ObjC runtime warmup — Rail's process isn't a Cocoa app, so class lists need
// forcing. Touching a constant NSString triggers __objc_classlist processing.
__attribute__((constructor))
static void _tgl_ctor(void) {
    volatile NSString *s = @"tgl_init_warmup";
    (void)s;
    pool_init_once();
}

static inline int storage_mode_index(MTLResourceOptions opt) {
    NSUInteger sm = (opt & MTLResourceStorageModeMask) >> MTLResourceStorageModeShift;
    if (sm == MTLStorageModeShared)  return 0;
    if (sm == MTLStorageModeManaged) return 1;
    if (sm == MTLStorageModePrivate) return 2;
    return 0;
}

static const char *storage_mode_name(int idx) {
    switch (idx) { case 0: return "shared"; case 1: return "managed";
                   case 2: return "private"; default: return "?"; }
}

static inline int tier_for_bytes(NSUInteger bytes) {
    for (int t = 0; t < NUM_TIERS; t++) if (bytes <= TIER_MAX_BYTES[t]) return t;
    return NUM_TIERS - 1;
}

static inline NSUInteger round_up_for_tier(NSUInteger bytes, int tier) {
    // T4 (unbounded) — round to 64 KB to avoid fragmenting megabuffers.
    if (tier == NUM_TIERS - 1) return (bytes + 65535) & ~((NSUInteger)65535);
    // Within a bounded tier, round up to next power-of-two so siblings
    // are interchangeable (best-fit always finds an exact reuse).
    NSUInteger p = 4096;
    while (p < bytes) p <<= 1;
    return p;
}

static id<MTLBuffer> pool_acquire_mode(NSUInteger bytes, MTLResourceOptions opt) {
    pool_init_once();
    int sm  = storage_mode_index(opt);
    int t   = tier_for_bytes(bytes);
    pool_tier_t *pt = &g_pools[sm][t];

    // RAIL_GPU_POOL_DISABLE=1 → skip best-fit reuse, force a fresh allocation
    // each call (falsification test for sequential-collapse hypothesis).
    if (g_pool_disable) goto miss_path;

    // Best-fit within tier — pick smallest idle slot ≥ bytes.
    int best = -1;
    for (int i = 0; i < pt->count; i++) {
        if (pt->slots[i].in_use)         continue;
        if (pt->slots[i].size < bytes)   continue;
        if (best < 0 || pt->slots[i].size < pt->slots[best].size) best = i;
    }
    if (best >= 0) {
        pt->slots[best].in_use = 1;
        pt->hits++;
        pt->in_use++;
        if (pt->in_use > pt->peak_in_use) pt->peak_in_use = pt->in_use;
        return pt->slots[best].buf;
    }
miss_path:;

    // Miss — allocate.
    NSUInteger rounded = round_up_for_tier(bytes, t);
    id<MTLBuffer> b = [g_device newBufferWithLength:rounded options:opt];
    if (!b) return nil;
    pt->misses++;

    if (pt->count < pt->capacity) {
        // Track in pool.
        int idx = pt->count++;
        pt->slots[idx].buf            = b;
        pt->slots[idx].size           = rounded;
        pt->slots[idx].in_use         = 1;
        pt->slots[idx].last_release_t = 0.0;
        pt->bytes_in_pool            += rounded;
        pt->in_use++;
        if (pt->in_use > pt->peak_in_use) pt->peak_in_use = pt->in_use;
        return b;
    }

    // Capacity full — try to grow (×2 up to TIER_MAX_CAP).
    if (pt->capacity < TIER_MAX_CAP) {
        int new_cap = pt->capacity * 2;
        if (new_cap > TIER_MAX_CAP) new_cap = TIER_MAX_CAP;
        pt->capacity = new_cap;
        int idx = pt->count++;
        pt->slots[idx].buf            = b;
        pt->slots[idx].size           = rounded;
        pt->slots[idx].in_use         = 1;
        pt->slots[idx].last_release_t = 0.0;
        pt->bytes_in_pool            += rounded;
        pt->in_use++;
        if (pt->in_use > pt->peak_in_use) pt->peak_in_use = pt->in_use;
        return b;
    }

    // Tier maxed — return the buffer untracked (ARC frees on autoreleasepool exit).
    pt->drops++;
    return b;
}

// Convenience wrapper — kernels in this file all use Shared storage.
static id<MTLBuffer> pool_acquire(NSUInteger bytes) {
    return pool_acquire_mode(bytes, MTLResourceStorageModeShared);
}

// Idle shrink: free at most one stale slot per tier per release tick to
// avoid stutter.  Keep at least TIER_MIN_KEEP slots alive even when idle.
static void pool_idle_reap(pool_tier_t *pt, double now) {
    if (pt->count <= TIER_MIN_KEEP) return;
    int oldest = -1;
    double oldest_t = now;
    for (int i = 0; i < pt->count; i++) {
        if (pt->slots[i].in_use) continue;
        if (pt->slots[i].last_release_t == 0.0) continue;
        double age = now - pt->slots[i].last_release_t;
        if (age < TIER_IDLE_SHRINK_SECS) continue;
        if (pt->slots[i].last_release_t < oldest_t) {
            oldest_t = pt->slots[i].last_release_t;
            oldest = i;
        }
    }
    if (oldest < 0) return;
    pt->bytes_in_pool -= pt->slots[oldest].size;
    pt->slots[oldest].buf = nil;       // ARC drops the MTLBuffer
    // Swap-remove with last slot to keep the dense prefix invariant.
    int last = pt->count - 1;
    if (oldest != last) pt->slots[oldest] = pt->slots[last];
    pt->slots[last].buf = nil;
    pt->count--;
    pt->shrinks++;
}

static void pool_release(id<MTLBuffer> b) {
    if (!b) return;
    double now = [NSDate timeIntervalSinceReferenceDate];
    for (int sm = 0; sm < NUM_STORAGE_MODES; sm++) {
        for (int t = 0; t < NUM_TIERS; t++) {
            pool_tier_t *pt = &g_pools[sm][t];
            for (int i = 0; i < pt->count; i++) {
                if (pt->slots[i].buf == b) {
                    pt->slots[i].in_use = 0;
                    pt->slots[i].last_release_t = now;
                    pt->in_use--;
                    pool_idle_reap(pt, now);
                    return;
                }
            }
        }
    }
    // Untracked buffer (capacity-overflow miss) — ARC will free it.
}

// Public C ABI accessor — used by tensor_daemond's stats endpoint.
typedef struct {
    int       tier;
    const char *storage;
    int       capacity;
    int       count;
    int       in_use;
    int       peak_in_use;
    long long hits;
    long long misses;
    long long drops;
    long long shrinks;
    long long bytes_in_pool;
} tgl_pool_stat_t;

int tgl_pool_stats(tgl_pool_stat_t *out, int max) {
    pool_init_once();
    int n = 0;
    for (int sm = 0; sm < NUM_STORAGE_MODES; sm++) {
        for (int t = 0; t < NUM_TIERS; t++) {
            if (n >= max) return n;
            pool_tier_t *pt = &g_pools[sm][t];
            out[n].tier          = t;
            out[n].storage       = storage_mode_name(sm);
            out[n].capacity      = pt->capacity;
            out[n].count         = pt->count;
            out[n].in_use        = pt->in_use;
            out[n].peak_in_use   = pt->peak_in_use;
            out[n].hits          = pt->hits;
            out[n].misses        = pt->misses;
            out[n].drops         = pt->drops;
            out[n].shrinks       = pt->shrinks;
            out[n].bytes_in_pool = pt->bytes_in_pool;
            n++;
        }
    }
    return n;
}

// ────────────────────────────────────────────────────────────────────
// Pipeline state lookup (lazy build)
// ────────────────────────────────────────────────────────────────────

static id<MTLComputePipelineState> pso(NSString *name) {
    id<MTLComputePipelineState> p = g_pipes[name];
    if (p) return p;
    NSError *err = nil;
    id<MTLFunction> fn = [g_lib newFunctionWithName:name];
    if (!fn) return nil;
    p = [g_device newComputePipelineStateWithFunction:fn error:&err];
    if (!p) return nil;
    g_pipes[name] = p;
    return p;
}

// ────────────────────────────────────────────────────────────────────
// Initialization: compile shader if needed, open library.
// ────────────────────────────────────────────────────────────────────

int tgl_init(int dummy) {
    (void)dummy;
    if (g_initialized) return 1;

    g_device = MTLCreateSystemDefaultDevice();
    if (!g_device) return -1;

    // Runtime-compile the shaders from tensor_gpu.metal source.  Locate
    // it relative to this dylib's own path on disk (via dladdr on tgl_init
    // itself) — no Xcode toolchain dependency, no env-var dependency, and
    // works regardless of what HOME the caller has (rail_native's child
    // process runs with an empty env).
    Dl_info dl;
    if (dladdr((const void *)&tgl_init, &dl) == 0 || !dl.dli_fname) {
        fprintf(stderr, "[tgl_init] dladdr failed\n"); return -1;
    }
    NSString *dylib_path = [NSString stringWithUTF8String:dl.dli_fname];
    NSString *src_path = [[dylib_path stringByDeletingLastPathComponent]
                          stringByAppendingPathComponent:@"tensor_gpu.metal"];
    fprintf(stderr, "[tgl_init] dylib=%s src_path=%s\n",
            [dylib_path UTF8String], [src_path UTF8String]);
    NSError *err = nil;
    NSString *src = [NSString stringWithContentsOfFile:src_path encoding:NSUTF8StringEncoding error:&err];
    if (!src) {
        fprintf(stderr, "[tgl_init] read failed: %s\n",
                err ? [[err localizedDescription] UTF8String] : "?");
        return -1;
    }
    MTLCompileOptions *opts = [MTLCompileOptions new];
    // MSL 3.1 needed for bfloat (Metal 3.1, macOS 14+). Default opts pin
    // to an older version — see lm_v3_chunked_bf16 launch debugging 2026-05-14.
    if (@available(macOS 14.0, *)) {
        opts.languageVersion = MTLLanguageVersion3_1;
    }
    g_lib = [g_device newLibraryWithSource:src options:opts error:&err];
    if (!g_lib) {
        fprintf(stderr, "[tgl_init] newLibraryWithSource failed: %s\n",
                err ? [[err localizedDescription] UTF8String] : "?");
        return -1;
    }
    fprintf(stderr, "[tgl_init] OK — %lu functions\n", (unsigned long)[g_lib functionNames].count);

    g_queue = [g_device newCommandQueue];
    g_pipes = [NSMutableDictionary dictionary];
    g_initialized = 1;
    return 1;
}

static inline int ensure_init(void) {
    if (!g_initialized) return tgl_init(0);
    return 1;
}

// ────────────────────────────────────────────────────────────────────
// Helpers for the common f64 → f32 → GPU → f32 → f64 pipeline.
// ────────────────────────────────────────────────────────────────────

static inline void f64_to_f32(const double *src, float *dst, int n) {
    for (int i = 0; i < n; i++) dst[i] = (float)src[i];
}
static inline void f32_to_f64(const float *src, double *dst, int n) {
    for (int i = 0; i < n; i++) dst[i] = (double)src[i];
}

// IEEE 754 double → half (flush subnormals to zero; saturate on overflow).
// Used by the fp16 matmul entry points. Bias buffers stay fp32 on the GPU
// side — only activations/weights go through half conversion.
static inline void f64_to_f16(const double *src, uint16_t *dst, int n) {
    for (int i = 0; i < n; i++) {
        float f = (float)src[i];
        union { float f; uint32_t u; } v; v.f = f;
        uint32_t sign = (v.u >> 31) & 0x1;
        int32_t  exp  = (int32_t)((v.u >> 23) & 0xFF) - 127 + 15;
        uint32_t mant = v.u & 0x7FFFFF;
        if (exp <= 0)       dst[i] = (uint16_t)(sign << 15);
        else if (exp >= 31) dst[i] = (uint16_t)((sign << 15) | (0x1F << 10));
        else                dst[i] = (uint16_t)((sign << 15) | ((uint32_t)exp << 10) | (mant >> 13));
    }
}

// IEEE 754 half → double. Subnormals flush to zero (kernels produce
// normalized outputs for trained workloads).
static inline void f16_to_f64(const uint16_t *src, double *dst, int n) {
    for (int i = 0; i < n; i++) {
        uint16_t h = src[i];
        uint32_t sign = (h >> 15) & 0x1;
        uint32_t exp  = (h >> 10) & 0x1F;
        uint32_t mant = h & 0x3FF;
        uint32_t f32;
        if (exp == 0)       f32 = sign << 31;
        else if (exp == 31) f32 = (sign << 31) | 0x7F800000 | (mant << 13);
        else                f32 = (sign << 31) | ((exp + 127 - 15) << 23) | (mant << 13);
        union { uint32_t u; float f; } v; v.u = f32;
        dst[i] = (double)v.f;
    }
}

// Dispatch a 1D "gid < n" kernel with N threads. Threadgroup 256 covers typical sizes.
static void dispatch_1d(id<MTLComputeCommandEncoder> enc, uint32_t n) {
    NSUInteger tg = 256;
    if (n < tg) tg = n;
    [enc dispatchThreads:MTLSizeMake(n, 1, 1)
      threadsPerThreadgroup:MTLSizeMake(tg, 1, 1)];
}

// ────────────────────────────────────────────────────────────────────
// MATMUL: C[M×N] = A[M×K] × B[K×N]
// ────────────────────────────────────────────────────────────────────

int tgl_matmul_f64(const double *A, const double *B, double *C, int M, int K, int N) {
    if (ensure_init() != 1) return -1;
    const double *Aptr = A + 1, *Bptr = B + 1;
    double *Cptr = C + 1;

    @autoreleasepool {
        uint32_t sA = M*K, sB = K*N, sC = M*N;
        id<MTLBuffer> bufA = pool_acquire(sA*4);
        id<MTLBuffer> bufB = pool_acquire(sB*4);
        id<MTLBuffer> bufC = pool_acquire(sC*4);

        f64_to_f32(Aptr, (float*)bufA.contents, sA);
        f64_to_f32(Bptr, (float*)bufB.contents, sB);

        id<MTLComputePipelineState> p = pso(@"matmul_blocked");
        if (!p) p = pso(@"matmul");
        if (!p) return -1;

        id<MTLCommandBuffer> cmd = [g_queue commandBuffer];
        id<MTLComputeCommandEncoder> enc = [cmd computeCommandEncoder];
        [enc setComputePipelineState:p];
        [enc setBuffer:bufA offset:0 atIndex:0];
        [enc setBuffer:bufB offset:0 atIndex:1];
        [enc setBuffer:bufC offset:0 atIndex:2];
        uint32_t mu=M, ku=K, nu=N;
        [enc setBytes:&mu length:4 atIndex:3];
        [enc setBytes:&ku length:4 atIndex:4];
        [enc setBytes:&nu length:4 atIndex:5];
        [enc dispatchThreadgroups:MTLSizeMake((N+63)/64, (M+63)/64, 1)
           threadsPerThreadgroup:MTLSizeMake(16, 16, 1)];
        [enc endEncoding];
        [cmd commit]; [cmd waitUntilCompleted];

        f32_to_f64((float*)bufC.contents, Cptr, sC);
        pool_release(bufA); pool_release(bufB); pool_release(bufC);
    }
    return 1;
}

// Fused matmul+bias+relu: C = relu(A@B + bias_broadcast_rowwise)
// bias is a length-N float array (one per output column). Pass NULL header too.
int tgl_matmul_relu_f64(const double *A, const double *B, const double *bias,
                        double *C, int M, int K, int N) {
    if (ensure_init() != 1) return -1;
    const double *Aptr = A + 1, *Bptr = B + 1, *Bsptr = bias + 1;
    double *Cptr = C + 1;

    @autoreleasepool {
        uint32_t sA = M*K, sB = K*N, sC = M*N;
        id<MTLBuffer> bufA = pool_acquire(sA*4);
        id<MTLBuffer> bufB = pool_acquire(sB*4);
        id<MTLBuffer> bufBias = pool_acquire(N*4);
        id<MTLBuffer> bufC = pool_acquire(sC*4);

        f64_to_f32(Aptr, (float*)bufA.contents, sA);
        f64_to_f32(Bptr, (float*)bufB.contents, sB);
        f64_to_f32(Bsptr, (float*)bufBias.contents, N);

        id<MTLComputePipelineState> p = pso(@"matmul_bias_relu");
        if (!p) return -1;

        id<MTLCommandBuffer> cmd = [g_queue commandBuffer];
        id<MTLComputeCommandEncoder> enc = [cmd computeCommandEncoder];
        [enc setComputePipelineState:p];
        [enc setBuffer:bufA offset:0 atIndex:0];
        [enc setBuffer:bufB offset:0 atIndex:1];
        [enc setBuffer:bufBias offset:0 atIndex:2];
        [enc setBuffer:bufC offset:0 atIndex:3];
        uint32_t mu=M, ku=K, nu=N;
        [enc setBytes:&mu length:4 atIndex:4];
        [enc setBytes:&ku length:4 atIndex:5];
        [enc setBytes:&nu length:4 atIndex:6];
        [enc dispatchThreadgroups:MTLSizeMake((N+15)/16, (M+15)/16, 1)
           threadsPerThreadgroup:MTLSizeMake(16, 16, 1)];
        [enc endEncoding];
        [cmd commit]; [cmd waitUntilCompleted];

        f32_to_f64((float*)bufC.contents, Cptr, sC);
        pool_release(bufA); pool_release(bufB);
        pool_release(bufBias); pool_release(bufC);
    }
    return 1;
}

// Fused matmul+bias+gelu — transformer FFN hidden layer.
int tgl_matmul_gelu_f64(const double *A, const double *B, const double *bias,
                        double *C, int M, int K, int N) {
    if (ensure_init() != 1) return -1;
    const double *Aptr = A+1, *Bptr = B+1, *Bsptr = bias+1;
    double *Cptr = C+1;
    @autoreleasepool {
        uint32_t sA = M*K, sB = K*N, sC = M*N;
        id<MTLBuffer> bA = pool_acquire(sA*4);
        id<MTLBuffer> bB = pool_acquire(sB*4);
        id<MTLBuffer> bBias = pool_acquire(N*4);
        id<MTLBuffer> bC = pool_acquire(sC*4);
        f64_to_f32(Aptr, (float*)bA.contents, sA);
        f64_to_f32(Bptr, (float*)bB.contents, sB);
        f64_to_f32(Bsptr, (float*)bBias.contents, N);

        id<MTLComputePipelineState> p = pso(@"matmul_bias_gelu");
        if (!p) return -1;
        id<MTLCommandBuffer> cmd = [g_queue commandBuffer];
        id<MTLComputeCommandEncoder> enc = [cmd computeCommandEncoder];
        [enc setComputePipelineState:p];
        [enc setBuffer:bA offset:0 atIndex:0];
        [enc setBuffer:bB offset:0 atIndex:1];
        [enc setBuffer:bBias offset:0 atIndex:2];
        [enc setBuffer:bC offset:0 atIndex:3];
        uint32_t mu=M, ku=K, nu=N;
        [enc setBytes:&mu length:4 atIndex:4];
        [enc setBytes:&ku length:4 atIndex:5];
        [enc setBytes:&nu length:4 atIndex:6];
        [enc dispatchThreadgroups:MTLSizeMake((N+15)/16,(M+15)/16,1)
           threadsPerThreadgroup:MTLSizeMake(16,16,1)];
        [enc endEncoding];
        [cmd commit]; [cmd waitUntilCompleted];

        f32_to_f64((float*)bC.contents, Cptr, sC);
        pool_release(bA); pool_release(bB); pool_release(bBias); pool_release(bC);
    }
    return 1;
}

// Batched matmul: C[b] = A[b] @ B[b] for b in 0..B_DIM-1.
int tgl_matmul_batched_f64(const double *A, const double *B, double *C,
                           int B_DIM, int M, int K, int N) {
    if (ensure_init() != 1) return -1;
    const double *Aptr = A+1, *Bptr = B+1;
    double *Cptr = C+1;
    @autoreleasepool {
        uint32_t sA = B_DIM*M*K, sB = B_DIM*K*N, sC = B_DIM*M*N;
        id<MTLBuffer> bA = pool_acquire(sA*4);
        id<MTLBuffer> bB = pool_acquire(sB*4);
        id<MTLBuffer> bC = pool_acquire(sC*4);
        f64_to_f32(Aptr, (float*)bA.contents, sA);
        f64_to_f32(Bptr, (float*)bB.contents, sB);

        id<MTLComputePipelineState> p = pso(@"matmul_batched");
        if (!p) return -1;
        id<MTLCommandBuffer> cmd = [g_queue commandBuffer];
        id<MTLComputeCommandEncoder> enc = [cmd computeCommandEncoder];
        [enc setComputePipelineState:p];
        [enc setBuffer:bA offset:0 atIndex:0];
        [enc setBuffer:bB offset:0 atIndex:1];
        [enc setBuffer:bC offset:0 atIndex:2];
        uint32_t bu=B_DIM, mu=M, ku=K, nu=N;
        [enc setBytes:&bu length:4 atIndex:3];
        [enc setBytes:&mu length:4 atIndex:4];
        [enc setBytes:&ku length:4 atIndex:5];
        [enc setBytes:&nu length:4 atIndex:6];
        [enc dispatchThreads:MTLSizeMake(N, M, B_DIM)
          threadsPerThreadgroup:MTLSizeMake(N<8?N:8, M<8?M:8, 1)];
        [enc endEncoding];
        [cmd commit]; [cmd waitUntilCompleted];

        f32_to_f64((float*)bC.contents, Cptr, sC);
        pool_release(bA); pool_release(bB); pool_release(bC);
    }
    return 1;
}

// ────────────────────────────────────────────────────────────────────
// fp16 MATMUL variants (Phase 4a Option A, 2026-04-21).
// A/B/C stage through f64 → half on input, half → f64 on output.
// Bias (when present) stays fp32 on GPU — see docs/plans/LABRAT_FIRST_WIN.md
// for the "fp32-bias" hint that made the fused kernels land 1.7–1.8×.
// ────────────────────────────────────────────────────────────────────

int tgl_matmul_f16(const double *A, const double *B, double *C, int M, int K, int N) {
    if (ensure_init() != 1) return -1;
    const double *Aptr = A + 1, *Bptr = B + 1;
    double *Cptr = C + 1;

    @autoreleasepool {
        uint32_t sA = M*K, sB = K*N, sC = M*N;
        id<MTLBuffer> bufA = pool_acquire(sA*2);
        id<MTLBuffer> bufB = pool_acquire(sB*2);
        id<MTLBuffer> bufC = pool_acquire(sC*2);

        f64_to_f16(Aptr, (uint16_t*)bufA.contents, sA);
        f64_to_f16(Bptr, (uint16_t*)bufB.contents, sB);

        id<MTLComputePipelineState> p = pso(@"matmul_f16");
        if (!p) return -1;

        id<MTLCommandBuffer> cmd = [g_queue commandBuffer];
        id<MTLComputeCommandEncoder> enc = [cmd computeCommandEncoder];
        [enc setComputePipelineState:p];
        [enc setBuffer:bufA offset:0 atIndex:0];
        [enc setBuffer:bufB offset:0 atIndex:1];
        [enc setBuffer:bufC offset:0 atIndex:2];
        uint32_t mu=M, ku=K, nu=N;
        [enc setBytes:&mu length:4 atIndex:3];
        [enc setBytes:&ku length:4 atIndex:4];
        [enc setBytes:&nu length:4 atIndex:5];
        [enc dispatchThreadgroups:MTLSizeMake((N+15)/16, (M+15)/16, 1)
           threadsPerThreadgroup:MTLSizeMake(16, 16, 1)];
        [enc endEncoding];
        [cmd commit]; [cmd waitUntilCompleted];

        f16_to_f64((uint16_t*)bufC.contents, Cptr, sC);
        pool_release(bufA); pool_release(bufB); pool_release(bufC);
    }
    return 1;
}

int tgl_matmul_blocked_f16(const double *A, const double *B, double *C, int M, int K, int N) {
    if (ensure_init() != 1) return -1;
    const double *Aptr = A + 1, *Bptr = B + 1;
    double *Cptr = C + 1;

    @autoreleasepool {
        uint32_t sA = M*K, sB = K*N, sC = M*N;
        id<MTLBuffer> bufA = pool_acquire(sA*2);
        id<MTLBuffer> bufB = pool_acquire(sB*2);
        id<MTLBuffer> bufC = pool_acquire(sC*2);

        f64_to_f16(Aptr, (uint16_t*)bufA.contents, sA);
        f64_to_f16(Bptr, (uint16_t*)bufB.contents, sB);

        id<MTLComputePipelineState> p = pso(@"matmul_blocked_f16");
        if (!p) return -1;

        id<MTLCommandBuffer> cmd = [g_queue commandBuffer];
        id<MTLComputeCommandEncoder> enc = [cmd computeCommandEncoder];
        [enc setComputePipelineState:p];
        [enc setBuffer:bufA offset:0 atIndex:0];
        [enc setBuffer:bufB offset:0 atIndex:1];
        [enc setBuffer:bufC offset:0 atIndex:2];
        uint32_t mu=M, ku=K, nu=N;
        [enc setBytes:&mu length:4 atIndex:3];
        [enc setBytes:&ku length:4 atIndex:4];
        [enc setBytes:&nu length:4 atIndex:5];
        [enc dispatchThreadgroups:MTLSizeMake((N+63)/64, (M+63)/64, 1)
           threadsPerThreadgroup:MTLSizeMake(16, 16, 1)];
        [enc endEncoding];
        [cmd commit]; [cmd waitUntilCompleted];

        f16_to_f64((uint16_t*)bufC.contents, Cptr, sC);
        pool_release(bufA); pool_release(bufB); pool_release(bufC);
    }
    return 1;
}

int tgl_matmul_bias_relu_f16(const double *A, const double *B, const double *bias,
                             double *C, int M, int K, int N) {
    if (ensure_init() != 1) return -1;
    const double *Aptr = A + 1, *Bptr = B + 1, *Bsptr = bias + 1;
    double *Cptr = C + 1;

    @autoreleasepool {
        uint32_t sA = M*K, sB = K*N, sC = M*N;
        id<MTLBuffer> bufA = pool_acquire(sA*2);
        id<MTLBuffer> bufB = pool_acquire(sB*2);
        id<MTLBuffer> bufBias = pool_acquire(N*4);   // fp32 bias per task-spec hint
        id<MTLBuffer> bufC = pool_acquire(sC*2);

        f64_to_f16(Aptr, (uint16_t*)bufA.contents, sA);
        f64_to_f16(Bptr, (uint16_t*)bufB.contents, sB);
        f64_to_f32(Bsptr, (float*)bufBias.contents, N);

        id<MTLComputePipelineState> p = pso(@"matmul_bias_relu_f16");
        if (!p) return -1;

        id<MTLCommandBuffer> cmd = [g_queue commandBuffer];
        id<MTLComputeCommandEncoder> enc = [cmd computeCommandEncoder];
        [enc setComputePipelineState:p];
        [enc setBuffer:bufA offset:0 atIndex:0];
        [enc setBuffer:bufB offset:0 atIndex:1];
        [enc setBuffer:bufBias offset:0 atIndex:2];
        [enc setBuffer:bufC offset:0 atIndex:3];
        uint32_t mu=M, ku=K, nu=N;
        [enc setBytes:&mu length:4 atIndex:4];
        [enc setBytes:&ku length:4 atIndex:5];
        [enc setBytes:&nu length:4 atIndex:6];
        [enc dispatchThreadgroups:MTLSizeMake((N+15)/16, (M+15)/16, 1)
           threadsPerThreadgroup:MTLSizeMake(16, 16, 1)];
        [enc endEncoding];
        [cmd commit]; [cmd waitUntilCompleted];

        f16_to_f64((uint16_t*)bufC.contents, Cptr, sC);
        pool_release(bufA); pool_release(bufB);
        pool_release(bufBias); pool_release(bufC);
    }
    return 1;
}

int tgl_matmul_bias_gelu_f16(const double *A, const double *B, const double *bias,
                             double *C, int M, int K, int N) {
    if (ensure_init() != 1) return -1;
    const double *Aptr = A + 1, *Bptr = B + 1, *Bsptr = bias + 1;
    double *Cptr = C + 1;

    @autoreleasepool {
        uint32_t sA = M*K, sB = K*N, sC = M*N;
        id<MTLBuffer> bufA = pool_acquire(sA*2);
        id<MTLBuffer> bufB = pool_acquire(sB*2);
        id<MTLBuffer> bufBias = pool_acquire(N*4);
        id<MTLBuffer> bufC = pool_acquire(sC*2);

        f64_to_f16(Aptr, (uint16_t*)bufA.contents, sA);
        f64_to_f16(Bptr, (uint16_t*)bufB.contents, sB);
        f64_to_f32(Bsptr, (float*)bufBias.contents, N);

        id<MTLComputePipelineState> p = pso(@"matmul_bias_gelu_f16");
        if (!p) return -1;

        id<MTLCommandBuffer> cmd = [g_queue commandBuffer];
        id<MTLComputeCommandEncoder> enc = [cmd computeCommandEncoder];
        [enc setComputePipelineState:p];
        [enc setBuffer:bufA offset:0 atIndex:0];
        [enc setBuffer:bufB offset:0 atIndex:1];
        [enc setBuffer:bufBias offset:0 atIndex:2];
        [enc setBuffer:bufC offset:0 atIndex:3];
        uint32_t mu=M, ku=K, nu=N;
        [enc setBytes:&mu length:4 atIndex:4];
        [enc setBytes:&ku length:4 atIndex:5];
        [enc setBytes:&nu length:4 atIndex:6];
        [enc dispatchThreadgroups:MTLSizeMake((N+15)/16, (M+15)/16, 1)
           threadsPerThreadgroup:MTLSizeMake(16, 16, 1)];
        [enc endEncoding];
        [cmd commit]; [cmd waitUntilCompleted];

        f16_to_f64((uint16_t*)bufC.contents, Cptr, sC);
        pool_release(bufA); pool_release(bufB);
        pool_release(bufBias); pool_release(bufC);
    }
    return 1;
}

// Correct parallel sum: one threadgroup per 256 elements produces a
// partial; host finishes on the few remaining.
int tgl_sum_f64(const double *X, double *result_out, int N) {
    if (ensure_init() != 1) return -1;
    const double *Xp = X+1;
    uint32_t n_tg = (N + 255) / 256;
    @autoreleasepool {
        id<MTLBuffer> bX = pool_acquire(N*4);
        id<MTLBuffer> bP = pool_acquire(n_tg*4);
        f64_to_f32(Xp, (float*)bX.contents, N);
        // Zero the partials buffer (pool may hand back stale memory).
        memset(bP.contents, 0, n_tg*4);

        id<MTLComputePipelineState> p = pso(@"tensor_sum_partials");
        if (!p) return -1;
        id<MTLCommandBuffer> cmd = [g_queue commandBuffer];
        id<MTLComputeCommandEncoder> enc = [cmd computeCommandEncoder];
        [enc setComputePipelineState:p];
        [enc setBuffer:bX offset:0 atIndex:0];
        [enc setBuffer:bP offset:0 atIndex:1];
        uint32_t nu = N;
        [enc setBytes:&nu length:4 atIndex:2];
        [enc dispatchThreadgroups:MTLSizeMake(n_tg, 1, 1)
           threadsPerThreadgroup:MTLSizeMake(256, 1, 1)];
        [enc endEncoding];
        [cmd commit]; [cmd waitUntilCompleted];

        float *pP = (float*)bP.contents;
        double total = 0.0;
        for (uint32_t i = 0; i < n_tg; i++) total += pP[i];
        result_out[0] = total;  // single-element output, not a float_arr
        pool_release(bX); pool_release(bP);
    }
    return 1;
}

// Softmax backward: dx = y * (dy - Σ y_k dy_k)
int tgl_softmax_backward_f64(const double *y, const double *dy, double *dx,
                             int rows, int cols) {
    if (ensure_init() != 1) return -1;
    const double *yp = y+1, *dyp = dy+1;
    double *dxp = dx+1;
    uint32_t n = rows * cols;
    @autoreleasepool {
        id<MTLBuffer> by = pool_acquire(n*4);
        id<MTLBuffer> bdy = pool_acquire(n*4);
        id<MTLBuffer> bdx = pool_acquire(n*4);
        f64_to_f32(yp, (float*)by.contents, n);
        f64_to_f32(dyp, (float*)bdy.contents, n);
        id<MTLComputePipelineState> p = pso(@"softmax_backward");
        if (!p) return -1;
        id<MTLCommandBuffer> cmd = [g_queue commandBuffer];
        id<MTLComputeCommandEncoder> enc = [cmd computeCommandEncoder];
        [enc setComputePipelineState:p];
        [enc setBuffer:by offset:0 atIndex:0];
        [enc setBuffer:bdy offset:0 atIndex:1];
        [enc setBuffer:bdx offset:0 atIndex:2];
        uint32_t ru=rows, cu=cols;
        [enc setBytes:&ru length:4 atIndex:3];
        [enc setBytes:&cu length:4 atIndex:4];
        dispatch_1d(enc, rows);
        [enc endEncoding];
        [cmd commit]; [cmd waitUntilCompleted];
        f32_to_f64((float*)bdx.contents, dxp, n);
        pool_release(by); pool_release(bdy); pool_release(bdx);
    }
    return 1;
}

// Fused softmax→CE backward: (probs - one_hot(targets)) / batch.
int tgl_ce_softmax_backward_f64(const double *probs, const double *targets,
                                double *grad, int batch, int vocab) {
    if (ensure_init() != 1) return -1;
    const double *pp = probs+1, *tp = targets+1;
    double *gp = grad+1;
    uint32_t n = batch * vocab;
    @autoreleasepool {
        id<MTLBuffer> bP = pool_acquire(n*4);
        id<MTLBuffer> bT = pool_acquire(batch*4);
        id<MTLBuffer> bG = pool_acquire(n*4);
        f64_to_f32(pp, (float*)bP.contents, n);
        f64_to_f32(tp, (float*)bT.contents, batch);
        id<MTLComputePipelineState> p = pso(@"ce_softmax_backward");
        if (!p) return -1;
        id<MTLCommandBuffer> cmd = [g_queue commandBuffer];
        id<MTLComputeCommandEncoder> enc = [cmd computeCommandEncoder];
        [enc setComputePipelineState:p];
        [enc setBuffer:bP offset:0 atIndex:0];
        [enc setBuffer:bT offset:0 atIndex:1];
        [enc setBuffer:bG offset:0 atIndex:2];
        uint32_t bu=batch, vu=vocab;
        [enc setBytes:&bu length:4 atIndex:3];
        [enc setBytes:&vu length:4 atIndex:4];
        dispatch_1d(enc, n);
        [enc endEncoding];
        [cmd commit]; [cmd waitUntilCompleted];
        f32_to_f64((float*)bG.contents, gp, n);
        pool_release(bP); pool_release(bT); pool_release(bG);
    }
    return 1;
}

// LayerNorm backward. Caller precomputes mean + rstd per row on forward.
int tgl_layernorm_backward_f64(const double *x, const double *mean,
                               const double *rstd, const double *gamma,
                               const double *dy, double *dx,
                               int rows, int dim) {
    if (ensure_init() != 1) return -1;
    const double *xp = x+1, *mp = mean+1, *rp = rstd+1, *gp = gamma+1, *dyp = dy+1;
    double *dxp = dx+1;
    uint32_t n = rows * dim;
    @autoreleasepool {
        id<MTLBuffer> bX = pool_acquire(n*4);
        id<MTLBuffer> bM = pool_acquire(rows*4);
        id<MTLBuffer> bR = pool_acquire(rows*4);
        id<MTLBuffer> bG = pool_acquire(dim*4);
        id<MTLBuffer> bDy = pool_acquire(n*4);
        id<MTLBuffer> bDx = pool_acquire(n*4);
        f64_to_f32(xp, (float*)bX.contents, n);
        f64_to_f32(mp, (float*)bM.contents, rows);
        f64_to_f32(rp, (float*)bR.contents, rows);
        f64_to_f32(gp, (float*)bG.contents, dim);
        f64_to_f32(dyp, (float*)bDy.contents, n);
        id<MTLComputePipelineState> p = pso(@"layernorm_backward");
        if (!p) return -1;
        id<MTLCommandBuffer> cmd = [g_queue commandBuffer];
        id<MTLComputeCommandEncoder> enc = [cmd computeCommandEncoder];
        [enc setComputePipelineState:p];
        [enc setBuffer:bX offset:0 atIndex:0];
        [enc setBuffer:bM offset:0 atIndex:1];
        [enc setBuffer:bR offset:0 atIndex:2];
        [enc setBuffer:bG offset:0 atIndex:3];
        [enc setBuffer:bDy offset:0 atIndex:4];
        [enc setBuffer:bDx offset:0 atIndex:5];
        uint32_t ru=rows, du=dim;
        [enc setBytes:&ru length:4 atIndex:6];
        [enc setBytes:&du length:4 atIndex:7];
        dispatch_1d(enc, rows);
        [enc endEncoding];
        [cmd commit]; [cmd waitUntilCompleted];
        f32_to_f64((float*)bDx.contents, dxp, n);
        pool_release(bX); pool_release(bM); pool_release(bR);
        pool_release(bG); pool_release(bDy); pool_release(bDx);
    }
    return 1;
}

// ────────────────────────────────────────────────────────────────────
// Generic 1-in-1-out kernel dispatch (relu/sigmoid/exp/tanh_fwd)
// ────────────────────────────────────────────────────────────────────

static int unary_op(NSString *kname, const double *X, double *Y, int N) {
    if (ensure_init() != 1) return -1;
    const double *Xptr = X + 1;
    double *Yptr = Y + 1;
    @autoreleasepool {
        id<MTLBuffer> bX = pool_acquire(N*4);
        id<MTLBuffer> bY = pool_acquire(N*4);
        f64_to_f32(Xptr, (float*)bX.contents, N);

        id<MTLComputePipelineState> p = pso(kname);
        if (!p) return -1;
        id<MTLCommandBuffer> cmd = [g_queue commandBuffer];
        id<MTLComputeCommandEncoder> enc = [cmd computeCommandEncoder];
        [enc setComputePipelineState:p];
        [enc setBuffer:bX offset:0 atIndex:0];
        [enc setBuffer:bY offset:0 atIndex:1];
        uint32_t nu = N;
        [enc setBytes:&nu length:4 atIndex:2];
        dispatch_1d(enc, N);
        [enc endEncoding];
        [cmd commit]; [cmd waitUntilCompleted];

        f32_to_f64((float*)bY.contents, Yptr, N);
        pool_release(bX); pool_release(bY);
    }
    return 1;
}

int tgl_relu_f64   (const double *X, double *Y, int N) { return unary_op(@"tensor_relu",     X, Y, N); }
int tgl_sigmoid_f64(const double *X, double *Y, int N) { return unary_op(@"tensor_sigmoid",  X, Y, N); }
int tgl_exp_f64    (const double *X, double *Y, int N) { return unary_op(@"tensor_exp",      X, Y, N); }
int tgl_tanh_f64   (const double *X, double *Y, int N) { return unary_op(@"tensor_tanh_fwd", X, Y, N); }

// ────────────────────────────────────────────────────────────────────
// Generic 2-in-1-out kernel dispatch (add/mul)
// ────────────────────────────────────────────────────────────────────

static int binary_op(NSString *kname, const double *A, const double *B, double *C, int N) {
    if (ensure_init() != 1) return -1;
    const double *Ap = A+1, *Bp = B+1;
    double *Cp = C+1;
    @autoreleasepool {
        id<MTLBuffer> bA = pool_acquire(N*4);
        id<MTLBuffer> bB = pool_acquire(N*4);
        id<MTLBuffer> bC = pool_acquire(N*4);
        f64_to_f32(Ap, (float*)bA.contents, N);
        f64_to_f32(Bp, (float*)bB.contents, N);

        id<MTLComputePipelineState> p = pso(kname);
        if (!p) return -1;
        id<MTLCommandBuffer> cmd = [g_queue commandBuffer];
        id<MTLComputeCommandEncoder> enc = [cmd computeCommandEncoder];
        [enc setComputePipelineState:p];
        [enc setBuffer:bA offset:0 atIndex:0];
        [enc setBuffer:bB offset:0 atIndex:1];
        [enc setBuffer:bC offset:0 atIndex:2];
        uint32_t nu = N;
        [enc setBytes:&nu length:4 atIndex:3];
        dispatch_1d(enc, N);
        [enc endEncoding];
        [cmd commit]; [cmd waitUntilCompleted];

        f32_to_f64((float*)bC.contents, Cp, N);
        pool_release(bA); pool_release(bB); pool_release(bC);
    }
    return 1;
}

int tgl_add_f64(const double *A, const double *B, double *C, int N) { return binary_op(@"tensor_add", A,B,C,N); }
int tgl_mul_f64(const double *A, const double *B, double *C, int N) { return binary_op(@"tensor_mul", A,B,C,N); }

// ────────────────────────────────────────────────────────────────────
// Scale: C = A * scalar
// Rail foreign ABI untags ints; for floats we pass as an untagged double*
// pointer whose first 8 bytes hold the scalar.
// ────────────────────────────────────────────────────────────────────

int tgl_scale_f64(const double *A, double *C, const double *scalar_ptr, int N) {
    if (ensure_init() != 1) return -1;
    const double *Ap = A+1;
    double *Cp = C+1;
    // scalar_ptr points at a 1-element float_arr: [count@0, value@8]
    double sv = scalar_ptr[1];
    float sf = (float)sv;

    @autoreleasepool {
        id<MTLBuffer> bA = pool_acquire(N*4);
        id<MTLBuffer> bC = pool_acquire(N*4);
        f64_to_f32(Ap, (float*)bA.contents, N);

        id<MTLComputePipelineState> p = pso(@"tensor_scale");
        if (!p) return -1;
        id<MTLCommandBuffer> cmd = [g_queue commandBuffer];
        id<MTLComputeCommandEncoder> enc = [cmd computeCommandEncoder];
        [enc setComputePipelineState:p];
        [enc setBuffer:bA offset:0 atIndex:0];
        [enc setBuffer:bC offset:0 atIndex:1];
        [enc setBytes:&sf length:4 atIndex:2];
        uint32_t nu = N;
        [enc setBytes:&nu length:4 atIndex:3];
        dispatch_1d(enc, N);
        [enc endEncoding];
        [cmd commit]; [cmd waitUntilCompleted];

        f32_to_f64((float*)bC.contents, Cp, N);
        pool_release(bA); pool_release(bC);
    }
    return 1;
}

// ────────────────────────────────────────────────────────────────────
// ReLU backward: out = (X>0) ? grad : 0
// ────────────────────────────────────────────────────────────────────

int tgl_relu_backward_f64(const double *X, const double *grad, double *out, int N) {
    if (ensure_init() != 1) return -1;
    const double *Xp = X+1, *Gp = grad+1;
    double *Op = out+1;
    @autoreleasepool {
        id<MTLBuffer> bX = pool_acquire(N*4);
        id<MTLBuffer> bG = pool_acquire(N*4);
        id<MTLBuffer> bO = pool_acquire(N*4);
        f64_to_f32(Xp, (float*)bX.contents, N);
        f64_to_f32(Gp, (float*)bG.contents, N);

        id<MTLComputePipelineState> p = pso(@"tensor_relu_backward");
        if (!p) return -1;
        id<MTLCommandBuffer> cmd = [g_queue commandBuffer];
        id<MTLComputeCommandEncoder> enc = [cmd computeCommandEncoder];
        [enc setComputePipelineState:p];
        [enc setBuffer:bX offset:0 atIndex:0];
        [enc setBuffer:bG offset:0 atIndex:1];
        [enc setBuffer:bO offset:0 atIndex:2];
        uint32_t nu = N;
        [enc setBytes:&nu length:4 atIndex:3];
        dispatch_1d(enc, N);
        [enc endEncoding];
        [cmd commit]; [cmd waitUntilCompleted];

        f32_to_f64((float*)bO.contents, Op, N);
        pool_release(bX); pool_release(bG); pool_release(bO);
    }
    return 1;
}

// ────────────────────────────────────────────────────────────────────
// Row-wise softmax: Y[r,:] = softmax(X[r,:])
// ────────────────────────────────────────────────────────────────────

int tgl_softmax_rows_f64(const double *X, double *Y, int rows, int cols) {
    if (ensure_init() != 1) return -1;
    const double *Xp = X+1;
    double *Yp = Y+1;
    uint32_t n = rows * cols;

    @autoreleasepool {
        id<MTLBuffer> bX = pool_acquire(n*4);
        id<MTLBuffer> bMax = pool_acquire(rows*4);
        id<MTLBuffer> bY = pool_acquire(n*4);
        id<MTLBuffer> bSum = pool_acquire(rows*4);
        f64_to_f32(Xp, (float*)bX.contents, n);

        id<MTLComputePipelineState> pMax = pso(@"softmax_max");
        id<MTLComputePipelineState> pExp = pso(@"softmax_exp_sum");
        id<MTLComputePipelineState> pNorm = pso(@"softmax_normalize");
        if (!pMax || !pExp || !pNorm) return -1;

        id<MTLCommandBuffer> cmd = [g_queue commandBuffer];
        uint32_t ru = rows, cu = cols;

        id<MTLComputeCommandEncoder> e1 = [cmd computeCommandEncoder];
        [e1 setComputePipelineState:pMax];
        [e1 setBuffer:bX offset:0 atIndex:0];
        [e1 setBuffer:bMax offset:0 atIndex:1];
        [e1 setBytes:&ru length:4 atIndex:2];
        [e1 setBytes:&cu length:4 atIndex:3];
        dispatch_1d(e1, rows);
        [e1 endEncoding];

        id<MTLComputeCommandEncoder> e2 = [cmd computeCommandEncoder];
        [e2 setComputePipelineState:pExp];
        [e2 setBuffer:bX offset:0 atIndex:0];
        [e2 setBuffer:bY offset:0 atIndex:1];
        [e2 setBuffer:bSum offset:0 atIndex:2];
        [e2 setBuffer:bMax offset:0 atIndex:3];
        [e2 setBytes:&ru length:4 atIndex:4];
        [e2 setBytes:&cu length:4 atIndex:5];
        dispatch_1d(e2, rows);
        [e2 endEncoding];

        id<MTLComputeCommandEncoder> e3 = [cmd computeCommandEncoder];
        [e3 setComputePipelineState:pNorm];
        [e3 setBuffer:bY offset:0 atIndex:0];
        [e3 setBuffer:bSum offset:0 atIndex:1];
        [e3 setBytes:&ru length:4 atIndex:2];
        [e3 setBytes:&cu length:4 atIndex:3];
        [e3 dispatchThreads:MTLSizeMake(cols, rows, 1)
          threadsPerThreadgroup:MTLSizeMake(cols < 16 ? cols : 16, rows < 16 ? rows : 16, 1)];
        [e3 endEncoding];

        [cmd commit]; [cmd waitUntilCompleted];

        f32_to_f64((float*)bY.contents, Yp, n);
        pool_release(bX); pool_release(bMax); pool_release(bY); pool_release(bSum);
    }
    return 1;
}

// ────────────────────────────────────────────────────────────────────
// Transpose: B[N×M] = A[M×N]^T
// ────────────────────────────────────────────────────────────────────

int tgl_transpose_f64(const double *A, double *B, int M, int N) {
    if (ensure_init() != 1) return -1;
    const double *Ap = A+1;
    double *Bp = B+1;
    uint32_t n = M*N;
    @autoreleasepool {
        id<MTLBuffer> bA = pool_acquire(n*4);
        id<MTLBuffer> bB = pool_acquire(n*4);
        f64_to_f32(Ap, (float*)bA.contents, n);

        id<MTLComputePipelineState> p = pso(@"tensor_transpose");
        if (!p) return -1;
        id<MTLCommandBuffer> cmd = [g_queue commandBuffer];
        id<MTLComputeCommandEncoder> enc = [cmd computeCommandEncoder];
        [enc setComputePipelineState:p];
        [enc setBuffer:bA offset:0 atIndex:0];
        [enc setBuffer:bB offset:0 atIndex:1];
        uint32_t mu=M, nu=N;
        [enc setBytes:&mu length:4 atIndex:2];
        [enc setBytes:&nu length:4 atIndex:3];
        [enc dispatchThreads:MTLSizeMake(N, M, 1)
          threadsPerThreadgroup:MTLSizeMake(N<16?N:16, M<16?M:16, 1)];
        [enc endEncoding];
        [cmd commit]; [cmd waitUntilCompleted];

        f32_to_f64((float*)bB.contents, Bp, n);
        pool_release(bA); pool_release(bB);
    }
    return 1;
}

// ────────────────────────────────────────────────────────────────────
// SGD update: w -= lr * grad (in-place on weights)
// lr passed by address (scalar_ptr[1] = value).
// ────────────────────────────────────────────────────────────────────

int tgl_sgd_update_f64(double *weights, const double *grad, const double *lr_ptr, int N) {
    if (ensure_init() != 1) return -1;
    double *wp = weights + 1;
    const double *gp = grad + 1;
    float lrf = (float)lr_ptr[1];
    @autoreleasepool {
        id<MTLBuffer> bW = pool_acquire(N*4);
        id<MTLBuffer> bG = pool_acquire(N*4);
        f64_to_f32(wp, (float*)bW.contents, N);
        f64_to_f32(gp, (float*)bG.contents, N);

        id<MTLComputePipelineState> p = pso(@"sgd_update");
        if (!p) return -1;
        id<MTLCommandBuffer> cmd = [g_queue commandBuffer];
        id<MTLComputeCommandEncoder> enc = [cmd computeCommandEncoder];
        [enc setComputePipelineState:p];
        [enc setBuffer:bW offset:0 atIndex:0];
        [enc setBuffer:bG offset:0 atIndex:1];
        [enc setBytes:&lrf length:4 atIndex:2];
        uint32_t nu = N;
        [enc setBytes:&nu length:4 atIndex:3];
        dispatch_1d(enc, N);
        [enc endEncoding];
        [cmd commit]; [cmd waitUntilCompleted];

        f32_to_f64((float*)bW.contents, wp, N);
        pool_release(bW); pool_release(bG);
    }
    return 1;
}

// ────────────────────────────────────────────────────────────────────
// Adam fused in-place update.
//   hyp is a Rail float_arr of length ≥6: [lr, β1, β2, ε, bc1, bc2]
//   m and v are mutated in place (same side as weights).
//   One GPU dispatch per parameter tensor per step.
// ────────────────────────────────────────────────────────────────────

int tgl_adam_update_f64(double *w, const double *g,
                        double *m, double *v,
                        const double *hyp, int N) {
    if (ensure_init() != 1) return -1;
    double *wp = w + 1;
    const double *gp = g + 1;
    double *mp = m + 1;
    double *vp = v + 1;
    const double *hp = hyp + 1;

    float hyp_f[6];
    for (int i = 0; i < 6; i++) hyp_f[i] = (float)hp[i];

    @autoreleasepool {
        id<MTLBuffer> bW = pool_acquire(N*4);
        id<MTLBuffer> bG = pool_acquire(N*4);
        id<MTLBuffer> bM = pool_acquire(N*4);
        id<MTLBuffer> bV = pool_acquire(N*4);
        f64_to_f32(wp, (float*)bW.contents, N);
        f64_to_f32(gp, (float*)bG.contents, N);
        f64_to_f32(mp, (float*)bM.contents, N);
        f64_to_f32(vp, (float*)bV.contents, N);

        id<MTLComputePipelineState> p = pso(@"adam_update");
        if (!p) return -1;
        id<MTLCommandBuffer> cmd = [g_queue commandBuffer];
        id<MTLComputeCommandEncoder> enc = [cmd computeCommandEncoder];
        [enc setComputePipelineState:p];
        [enc setBuffer:bW offset:0 atIndex:0];
        [enc setBuffer:bG offset:0 atIndex:1];
        [enc setBuffer:bM offset:0 atIndex:2];
        [enc setBuffer:bV offset:0 atIndex:3];
        [enc setBytes:hyp_f length:sizeof(hyp_f) atIndex:4];
        uint32_t nu = N;
        [enc setBytes:&nu length:4 atIndex:5];
        dispatch_1d(enc, N);
        [enc endEncoding];
        [cmd commit]; [cmd waitUntilCompleted];

        f32_to_f64((float*)bW.contents, wp, N);
        f32_to_f64((float*)bM.contents, mp, N);
        f32_to_f64((float*)bV.contents, vp, N);
        pool_release(bW); pool_release(bG);
        pool_release(bM); pool_release(bV);
    }
    return 1;
}

// ────────────────────────────────────────────────────────────────────
// Cross-entropy: losses[i] = -log(probs[i, targets[i]])
// targets passed as double values (f64) — we cast to uint inside.
// ────────────────────────────────────────────────────────────────────

int tgl_cross_entropy_f64(const double *probs, const double *targets,
                          double *losses, int batch, int vocab) {
    if (ensure_init() != 1) return -1;
    const double *pp = probs + 1, *tp = targets + 1;
    double *lp = losses + 1;
    uint32_t np = batch * vocab;

    @autoreleasepool {
        id<MTLBuffer> bP = pool_acquire(np*4);
        id<MTLBuffer> bT = pool_acquire(batch*4);
        id<MTLBuffer> bL = pool_acquire(batch*4);
        f64_to_f32(pp, (float*)bP.contents, np);
        uint32_t *tg = (uint32_t*)bT.contents;
        for (int i = 0; i < batch; i++) tg[i] = (uint32_t)tp[i];

        id<MTLComputePipelineState> p = pso(@"cross_entropy");
        if (!p) return -1;
        id<MTLCommandBuffer> cmd = [g_queue commandBuffer];
        id<MTLComputeCommandEncoder> enc = [cmd computeCommandEncoder];
        [enc setComputePipelineState:p];
        [enc setBuffer:bP offset:0 atIndex:0];
        [enc setBuffer:bT offset:0 atIndex:1];
        [enc setBuffer:bL offset:0 atIndex:2];
        uint32_t bu = batch, vu = vocab;
        [enc setBytes:&bu length:4 atIndex:3];
        [enc setBytes:&vu length:4 atIndex:4];
        dispatch_1d(enc, batch);
        [enc endEncoding];
        [cmd commit]; [cmd waitUntilCompleted];

        f32_to_f64((float*)bL.contents, lp, batch);
        pool_release(bP); pool_release(bT); pool_release(bL);
    }
    return 1;
}

// ────────────────────────────────────────────────────────────────────
// Metal IR — JIT-compiled user kernels from Rail source.
//
// tgl_unary_from_source(src, name, X, Y, N):
//   Compile Metal source at runtime (cached by pointer identity of
//   src, assumed to be a compile-time-baked string literal) into a
//   library, look up `name` as a compute kernel, and dispatch it
//   over N float_arr elements.
//
// Kernel signature expected:
//   kernel void <name>(device const float *X [[buffer(0)]],
//                      device float *Y [[buffer(1)]],
//                      constant uint &n [[buffer(2)]],
//                      uint gid [[thread_position_in_grid]]);
//
// Caching: we key a NSMutableDictionary by the Metal source string.
// Recompile only on first sight; subsequent calls reuse the cached
// MTLLibrary and pipeline.
// ────────────────────────────────────────────────────────────────────

static NSMutableDictionary<NSString *, id<MTLLibrary>> *g_src_libs = nil;
static NSMutableDictionary<NSString *, id<MTLComputePipelineState>> *g_src_pipes = nil;

int tgl_unary_from_source(const char *src, const char *kname,
                           const double *X, double *Y, int N) {
    if (ensure_init() != 1) return -1;
    const double *Xp = X + 1;
    double *Yp = Y + 1;
    @autoreleasepool {
        if (!g_src_libs) g_src_libs = [NSMutableDictionary dictionary];
        if (!g_src_pipes) g_src_pipes = [NSMutableDictionary dictionary];

        NSString *nsSrc = [NSString stringWithUTF8String:src];
        NSString *nsName = [NSString stringWithUTF8String:kname];
        NSString *pipeKey = [nsSrc stringByAppendingString:nsName];

        id<MTLComputePipelineState> p = g_src_pipes[pipeKey];
        if (!p) {
            id<MTLLibrary> lib = g_src_libs[nsSrc];
            if (!lib) {
                NSError *err = nil;
                lib = [g_device newLibraryWithSource:nsSrc options:nil error:&err];
                if (!lib) {
                    fprintf(stderr, "tgl_unary_from_source: compile failed: %s\n",
                            err ? [[err localizedDescription] UTF8String] : "(null)");
                    return -1;
                }
                g_src_libs[nsSrc] = lib;
            }
            id<MTLFunction> fn = [lib newFunctionWithName:nsName];
            if (!fn) {
                fprintf(stderr, "tgl_unary_from_source: function `%s` not found\n", kname);
                return -1;
            }
            NSError *err = nil;
            p = [g_device newComputePipelineStateWithFunction:fn error:&err];
            if (!p) {
                fprintf(stderr, "tgl_unary_from_source: pipeline failed: %s\n",
                        err ? [[err localizedDescription] UTF8String] : "(null)");
                return -1;
            }
            g_src_pipes[pipeKey] = p;
        }

        id<MTLBuffer> bX = pool_acquire(N * 4);
        id<MTLBuffer> bY = pool_acquire(N * 4);
        f64_to_f32(Xp, (float *)bX.contents, N);

        id<MTLCommandBuffer> cmd = [g_queue commandBuffer];
        id<MTLComputeCommandEncoder> enc = [cmd computeCommandEncoder];
        [enc setComputePipelineState:p];
        [enc setBuffer:bX offset:0 atIndex:0];
        [enc setBuffer:bY offset:0 atIndex:1];
        uint32_t nu = N;
        [enc setBytes:&nu length:4 atIndex:2];
        dispatch_1d(enc, N);
        [enc endEncoding];
        [cmd commit]; [cmd waitUntilCompleted];

        f32_to_f64((float *)bY.contents, Yp, N);
        pool_release(bX); pool_release(bY);
    }
    return 1;
}

// ────────────────────────────────────────────────────────────────────
// Sub-ms wall clock for microbenches. Rail's stdlib/time.rail `time`
// is Unix-second granularity; this returns float ms since epoch.
// ────────────────────────────────────────────────────────────────────

double tgl_now_ms(int _ignored) {
    struct timeval tv;
    gettimeofday(&tv, NULL);
    return (double)tv.tv_sec * 1000.0 + (double)tv.tv_usec / 1000.0;
}

// ────────────────────────────────────────────────────────────────────
// HalfTensor path (Phase 4b follow-on, 2026-04-21).
//
// These entry points skip the f64↔f16 cast that tgl_matmul_f16 pays on
// every call. A HalfTensor's Rail-side storage is a float_arr whose
// data area is reinterpreted as packed uint16_t[n]. Pack/unpack crosses
// the Rail↔C boundary once at cast time (init, checkpoint) rather than
// per-matmul. Kernels still use the existing matmul_f16 shader — only
// the host-side entry changes.
//
// ABI: Rail's float_arr layout is [length(8), data...]. We pass the
// float_arr pointer as `double*` for ABI consistency and cast to
// `uint16_t*` internally. That cast plus the `+1` offset gives a
// pointer to the first packed half.
// ────────────────────────────────────────────────────────────────────

// f64 array → packed-half float_arr (host-side). Callers pass two Rail
// float_arrs; after the call, dst's data area holds n packed halfs.
int tgl_f64_to_half(const double *src, double *dst, int n) {
    const double *Sp = src + 1;
    uint16_t *Dp = (uint16_t *)(dst + 1);
    f64_to_f16(Sp, Dp, n);
    return 1;
}

// Inverse: packed-half float_arr → f64 array.
int tgl_half_to_f64(const double *src, double *dst, int n) {
    const uint16_t *Sp = (const uint16_t *)(src + 1);
    double *Dp = dst + 1;
    f16_to_f64(Sp, Dp, n);
    return 1;
}

// Zero-cast fp16 matmul. Inputs and output are packed-half float_arrs.
// Kernel is the existing matmul_f16 shader; only the host side changes:
// memcpy in, dispatch, memcpy out. No f64 staging buffer, no per-element
// cast. This is the Phase 4b "eliminate the language's need for the
// cast" line item.
int tgl_matmul_half_host(const double *A, const double *B, double *C,
                         int M, int K, int N) {
    if (ensure_init() != 1) return -1;
    const uint16_t *Aptr = (const uint16_t *)(A + 1);
    const uint16_t *Bptr = (const uint16_t *)(B + 1);
    uint16_t *Cptr       = (uint16_t       *)(C + 1);

    @autoreleasepool {
        uint32_t sA = M*K, sB = K*N, sC = M*N;
        id<MTLBuffer> bufA = pool_acquire(sA*2);
        id<MTLBuffer> bufB = pool_acquire(sB*2);
        id<MTLBuffer> bufC = pool_acquire(sC*2);

        memcpy(bufA.contents, Aptr, sA*2);
        memcpy(bufB.contents, Bptr, sB*2);

        id<MTLComputePipelineState> p = pso(@"matmul_f16");
        if (!p) return -1;

        id<MTLCommandBuffer> cmd = [g_queue commandBuffer];
        id<MTLComputeCommandEncoder> enc = [cmd computeCommandEncoder];
        [enc setComputePipelineState:p];
        [enc setBuffer:bufA offset:0 atIndex:0];
        [enc setBuffer:bufB offset:0 atIndex:1];
        [enc setBuffer:bufC offset:0 atIndex:2];
        uint32_t mu=M, ku=K, nu=N;
        [enc setBytes:&mu length:4 atIndex:3];
        [enc setBytes:&ku length:4 atIndex:4];
        [enc setBytes:&nu length:4 atIndex:5];
        [enc dispatchThreadgroups:MTLSizeMake((N+15)/16, (M+15)/16, 1)
           threadsPerThreadgroup:MTLSizeMake(16, 16, 1)];
        [enc endEncoding];
        [cmd commit]; [cmd waitUntilCompleted];

        memcpy(Cptr, bufC.contents, sC*2);
        pool_release(bufA); pool_release(bufB); pool_release(bufC);
    }
    return 1;
}

// Rail-native mixed precision matmul: f64 activations x fp16 weights -> f64.
//
// X arrives as a Rail Tensor (f64 float_arr); W arrives as a Rail HalfTensor
// (fp16 packed in a float_arr-shaped buffer). Y is a Rail Tensor (f64).
// Host casts f64->f32 for X on input and f32->f64 for Y on output; GPU side
// holds X/Y in fp32, W in fp16, dot product in fp64. The Rail-side caller
// stays entirely in f64 land — same surface as matmul_half but precision-
// preserving across sequential calls.
//
// See tools/metal/tensor_gpu.metal:matmul_f32x_halfw for the kernel and
// the precision rationale.
int tgl_matmul_f32x_halfw_host(const double *X_f64, const double *W_half,
                               double *Y_f64, int M, int K, int N) {
    if (ensure_init() != 1) return -1;
    const double   *Xptr = X_f64 + 1;
    const uint16_t *Wptr = (const uint16_t *)(W_half + 1);
    double         *Yptr = Y_f64 + 1;

    @autoreleasepool {
        uint32_t sX = M*K, sW = K*N, sY = M*N;
        // X and Y in fp32 (4 bytes/element); W in fp16 (2 bytes/element).
        id<MTLBuffer> bufX = pool_acquire(sX*4);
        id<MTLBuffer> bufW = pool_acquire(sW*2);
        id<MTLBuffer> bufY = pool_acquire(sY*4);

        // f64 -> f32 cast host-side at the GPU boundary.
        float *X_f32 = (float *)bufX.contents;
        for (uint32_t i = 0; i < sX; i++) X_f32[i] = (float)Xptr[i];

        // fp16 weights: zero-cast memcpy from Rail HalfTensor payload.
        memcpy(bufW.contents, Wptr, sW*2);

        id<MTLComputePipelineState> p = pso(@"matmul_f32x_halfw");
        if (!p) return -1;

        id<MTLCommandBuffer> cmd = [g_queue commandBuffer];
        id<MTLComputeCommandEncoder> enc = [cmd computeCommandEncoder];
        [enc setComputePipelineState:p];
        [enc setBuffer:bufX offset:0 atIndex:0];
        [enc setBuffer:bufW offset:0 atIndex:1];
        [enc setBuffer:bufY offset:0 atIndex:2];
        uint32_t mu=M, ku=K, nu=N;
        [enc setBytes:&mu length:4 atIndex:3];
        [enc setBytes:&ku length:4 atIndex:4];
        [enc setBytes:&nu length:4 atIndex:5];
        [enc dispatchThreadgroups:MTLSizeMake((N+15)/16, (M+15)/16, 1)
           threadsPerThreadgroup:MTLSizeMake(16, 16, 1)];
        [enc endEncoding];
        [cmd commit]; [cmd waitUntilCompleted];

        // f32 -> f64 cast host-side; Rail caller sees Tensor (f64) output.
        const float *Y_f32 = (const float *)bufY.contents;
        for (uint32_t i = 0; i < sY; i++) Yptr[i] = (double)Y_f32[i];

        pool_release(bufX); pool_release(bufW); pool_release(bufY);
    }
    return 1;
}

// Zero-cast fp16 element-wise add. Inputs + output are packed-half
// float_arrs; GPU keeps accumulator in fp32 and casts back to half at
// store. Uses the tensor_add_f16 kernel (gid-per-element, 1-D).
int tgl_add_half_host(const double *A, const double *B, double *C, int n) {
    if (ensure_init() != 1) return -1;
    const uint16_t *Aptr = (const uint16_t *)(A + 1);
    const uint16_t *Bptr = (const uint16_t *)(B + 1);
    uint16_t       *Cptr = (uint16_t       *)(C + 1);

    @autoreleasepool {
        NSUInteger bytes = (NSUInteger)n * 2;
        id<MTLBuffer> bufA = pool_acquire(bytes);
        id<MTLBuffer> bufB = pool_acquire(bytes);
        id<MTLBuffer> bufC = pool_acquire(bytes);

        memcpy(bufA.contents, Aptr, bytes);
        memcpy(bufB.contents, Bptr, bytes);

        id<MTLComputePipelineState> p = pso(@"tensor_add_f16");
        if (!p) return -1;

        id<MTLCommandBuffer> cmd = [g_queue commandBuffer];
        id<MTLComputeCommandEncoder> enc = [cmd computeCommandEncoder];
        [enc setComputePipelineState:p];
        [enc setBuffer:bufA offset:0 atIndex:0];
        [enc setBuffer:bufB offset:0 atIndex:1];
        [enc setBuffer:bufC offset:0 atIndex:2];
        uint32_t nu = n;
        [enc setBytes:&nu length:4 atIndex:3];

        NSUInteger tg = 256;
        [enc dispatchThreadgroups:MTLSizeMake((n + tg - 1) / tg, 1, 1)
           threadsPerThreadgroup:MTLSizeMake(tg, 1, 1)];
        [enc endEncoding];
        [cmd commit]; [cmd waitUntilCompleted];

        memcpy(Cptr, bufC.contents, bytes);
        pool_release(bufA); pool_release(bufB); pool_release(bufC);
    }
    return 1;
}

// Zero-cast fp16 element-wise scale. Scalar is passed via a 1-element
// Rail float_arr (same ABI as tgl_scale_f64): scalar_arr[1] holds the
// f64 value, narrowed to fp32 on the GPU side.
int tgl_scale_half_host(const double *A, double *C,
                        const double *scalar_arr, int n) {
    if (ensure_init() != 1) return -1;
    const uint16_t *Aptr = (const uint16_t *)(A + 1);
    uint16_t       *Cptr = (uint16_t       *)(C + 1);
    float sf = (float)scalar_arr[1];

    @autoreleasepool {
        NSUInteger bytes = (NSUInteger)n * 2;
        id<MTLBuffer> bufA = pool_acquire(bytes);
        id<MTLBuffer> bufC = pool_acquire(bytes);

        memcpy(bufA.contents, Aptr, bytes);

        id<MTLComputePipelineState> p = pso(@"tensor_scale_f16");
        if (!p) return -1;

        id<MTLCommandBuffer> cmd = [g_queue commandBuffer];
        id<MTLComputeCommandEncoder> enc = [cmd computeCommandEncoder];
        [enc setComputePipelineState:p];
        [enc setBuffer:bufA offset:0 atIndex:0];
        [enc setBuffer:bufC offset:0 atIndex:1];
        [enc setBytes:&sf length:4 atIndex:2];
        uint32_t nu = n;
        [enc setBytes:&nu length:4 atIndex:3];
        dispatch_1d(enc, n);
        [enc endEncoding];
        [cmd commit]; [cmd waitUntilCompleted];

        memcpy(Cptr, bufC.contents, bytes);
        pool_release(bufA); pool_release(bufC);
    }
    return 1;
}

// Zero-cast fp16 matrix transpose: B[N×M] = A[M×N]^T. Packed-half
// storage in + out, no f64 boundary cross.
int tgl_transpose_half_host(const double *A, double *B, int M, int N) {
    if (ensure_init() != 1) return -1;
    const uint16_t *Aptr = (const uint16_t *)(A + 1);
    uint16_t       *Bptr = (uint16_t       *)(B + 1);
    NSUInteger bytes = (NSUInteger)M * (NSUInteger)N * 2;

    @autoreleasepool {
        id<MTLBuffer> bufA = pool_acquire(bytes);
        id<MTLBuffer> bufB = pool_acquire(bytes);
        memcpy(bufA.contents, Aptr, bytes);

        id<MTLComputePipelineState> p = pso(@"tensor_transpose_f16");
        if (!p) return -1;

        id<MTLCommandBuffer> cmd = [g_queue commandBuffer];
        id<MTLComputeCommandEncoder> enc = [cmd computeCommandEncoder];
        [enc setComputePipelineState:p];
        [enc setBuffer:bufA offset:0 atIndex:0];
        [enc setBuffer:bufB offset:0 atIndex:1];
        uint32_t mu = M, nu = N;
        [enc setBytes:&mu length:4 atIndex:2];
        [enc setBytes:&nu length:4 atIndex:3];
        [enc dispatchThreads:MTLSizeMake(N, M, 1)
          threadsPerThreadgroup:MTLSizeMake(N < 16 ? N : 16, M < 16 ? M : 16, 1)];
        [enc endEncoding];
        [cmd commit]; [cmd waitUntilCompleted];

        memcpy(Bptr, bufB.contents, bytes);
        pool_release(bufA); pool_release(bufB);
    }
    return 1;
}

// Zero-cast fp16 row-wise softmax. Log-sum-exp reformulation keeps
// exp() inputs ≤ 0 on the GPU so no fp16 overflow even when raw
// logits would have saturated. Accumulators stay fp32.
int tgl_softmax_half_host(const double *A, double *C, int rows, int cols) {
    if (ensure_init() != 1) return -1;
    const uint16_t *Aptr = (const uint16_t *)(A + 1);
    uint16_t       *Cptr = (uint16_t       *)(C + 1);
    NSUInteger bytes = (NSUInteger)rows * (NSUInteger)cols * 2;

    @autoreleasepool {
        id<MTLBuffer> bufA = pool_acquire(bytes);
        id<MTLBuffer> bufC = pool_acquire(bytes);
        memcpy(bufA.contents, Aptr, bytes);

        id<MTLComputePipelineState> p = pso(@"tensor_softmax_f16");
        if (!p) return -1;

        id<MTLCommandBuffer> cmd = [g_queue commandBuffer];
        id<MTLComputeCommandEncoder> enc = [cmd computeCommandEncoder];
        [enc setComputePipelineState:p];
        [enc setBuffer:bufA offset:0 atIndex:0];
        [enc setBuffer:bufC offset:0 atIndex:1];
        uint32_t ru = rows, cu = cols;
        [enc setBytes:&ru length:4 atIndex:2];
        [enc setBytes:&cu length:4 atIndex:3];
        dispatch_1d(enc, rows);
        [enc endEncoding];
        [cmd commit]; [cmd waitUntilCompleted];

        memcpy(Cptr, bufC.contents, bytes);
        pool_release(bufA); pool_release(bufC);
    }
    return 1;
}

// ────────────────────────────────────────────────────────────────────
// Transformer fused-op wrappers (added 2026-05-14)
// CPU-equivalent in stdlib/transformer.rail:
//   rms_rows_save / rope_positions / silu_fwd_loop
// ────────────────────────────────────────────────────────────────────

// X: [n_rows, dim] f64; Y: same shape; G: [dim] f64 gamma; RSTD: [n_rows] f64 out.
// eps fixed at 1e-5 (matches stdlib/transformer.rail:rmsnorm_save default).
// Rail's FFI doesn't route a mid-signature `double` to d0 — hardcoding sidesteps that.
int tgl_rmsnorm_save_f64(const double *X, double *Y, const double *G,
                         double *RSTD, int dim, int n_rows) {
    if (ensure_init() != 1) return -1;
    const double *Xp = X+1; const double *Gp = G+1;
    double *Yp = Y+1; double *Rp = RSTD+1;
    uint32_t total = (uint32_t)(n_rows * dim);
    @autoreleasepool {
        id<MTLBuffer> bX = pool_acquire(total*4);
        id<MTLBuffer> bY = pool_acquire(total*4);
        id<MTLBuffer> bG = pool_acquire(dim*4);
        id<MTLBuffer> bR = pool_acquire(n_rows*4);
        f64_to_f32(Xp, (float*)bX.contents, total);
        f64_to_f32(Gp, (float*)bG.contents, dim);

        id<MTLComputePipelineState> p = pso(@"rmsnorm_save");
        if (!p) return -1;
        id<MTLCommandBuffer> cmd = [g_queue commandBuffer];
        id<MTLComputeCommandEncoder> enc = [cmd computeCommandEncoder];
        [enc setComputePipelineState:p];
        [enc setBuffer:bX offset:0 atIndex:0];
        [enc setBuffer:bY offset:0 atIndex:1];
        [enc setBuffer:bG offset:0 atIndex:2];
        [enc setBuffer:bR offset:0 atIndex:3];
        uint32_t du = (uint32_t)dim; float epsf = 1.0e-5f;
        [enc setBytes:&du length:4 atIndex:4];
        [enc setBytes:&epsf length:4 atIndex:5];
        // one threadgroup per row, 64 threads (matches kernel's threadgroup partial[64])
        [enc dispatchThreadgroups:MTLSizeMake(n_rows, 1, 1)
          threadsPerThreadgroup:MTLSizeMake(64, 1, 1)];
        [enc endEncoding];
        [cmd commit]; [cmd waitUntilCompleted];
        f32_to_f64((float*)bY.contents, Yp, total);
        f32_to_f64((float*)bR.contents, Rp, n_rows);
        pool_release(bX); pool_release(bY); pool_release(bG); pool_release(bR);
    }
    return 1;
}

// X: [seq, d] f64, mutated in place.  sign = +1 forward, -1 inverse.
int tgl_rope_apply_f64(double *X, int seq, int d, double sign) {
    if (ensure_init() != 1) return -1;
    double *Xp = X+1;
    uint32_t total = (uint32_t)(seq * d);
    @autoreleasepool {
        id<MTLBuffer> bX = pool_acquire(total*4);
        f64_to_f32(Xp, (float*)bX.contents, total);

        id<MTLComputePipelineState> p = pso(@"rope_apply");
        if (!p) return -1;
        id<MTLCommandBuffer> cmd = [g_queue commandBuffer];
        id<MTLComputeCommandEncoder> enc = [cmd computeCommandEncoder];
        [enc setComputePipelineState:p];
        [enc setBuffer:bX offset:0 atIndex:0];
        uint32_t su = (uint32_t)seq, du = (uint32_t)d;
        float sf = (float)sign;
        [enc setBytes:&su length:4 atIndex:1];
        [enc setBytes:&du length:4 atIndex:2];
        [enc setBytes:&sf length:4 atIndex:3];
        uint half_d = d / 2;
        [enc dispatchThreads:MTLSizeMake(half_d, seq, 1)
          threadsPerThreadgroup:MTLSizeMake(half_d < 16 ? half_d : 16,
                                            seq    < 16 ? seq    : 16, 1)];
        [enc endEncoding];
        [cmd commit]; [cmd waitUntilCompleted];
        f32_to_f64((float*)bX.contents, Xp, total);
        pool_release(bX);
    }
    return 1;
}

// X: [N] f64.  Y: [N] f64 = silu(X).  SIG: [N] f64 = σ(X) for backward.
int tgl_silu_fwd_f64(const double *X, double *Y, double *SIG, int N) {
    if (ensure_init() != 1) return -1;
    const double *Xp = X+1;
    double *Yp = Y+1; double *Sp = SIG+1;
    @autoreleasepool {
        id<MTLBuffer> bX = pool_acquire(N*4);
        id<MTLBuffer> bY = pool_acquire(N*4);
        id<MTLBuffer> bS = pool_acquire(N*4);
        f64_to_f32(Xp, (float*)bX.contents, N);

        id<MTLComputePipelineState> p = pso(@"silu_fwd");
        if (!p) return -1;
        id<MTLCommandBuffer> cmd = [g_queue commandBuffer];
        id<MTLComputeCommandEncoder> enc = [cmd computeCommandEncoder];
        [enc setComputePipelineState:p];
        [enc setBuffer:bX offset:0 atIndex:0];
        [enc setBuffer:bY offset:0 atIndex:1];
        [enc setBuffer:bS offset:0 atIndex:2];
        uint32_t nu = (uint32_t)N;
        [enc setBytes:&nu length:4 atIndex:3];
        dispatch_1d(enc, N);
        [enc endEncoding];
        [cmd commit]; [cmd waitUntilCompleted];
        f32_to_f64((float*)bY.contents, Yp, N);
        f32_to_f64((float*)bS.contents, Sp, N);
        pool_release(bX); pool_release(bY); pool_release(bS);
    }
    return 1;
}

// ────────────────────────────────────────────────────────────────────
// Phase 0 — JIT compile foundation (rail-jit-fused-kernels-plan.md)
// Read MSL source from /tmp/rail_jit_kernel.metal, compile via
// newLibraryWithSource at runtime, cache the MTLComputePipelineState
// in an array.  Rail receives back an integer kernel ID.  Proves the
// "Rail emits MSL → Rail dispatches" loop end-to-end.
// ────────────────────────────────────────────────────────────────────
static NSMutableArray *g_jit_pipes = nil;

int tgl_jit_compile_from_tmp_file(int dummy) {
    (void)dummy;
    if (ensure_init() != 1) return -1;
    if (!g_jit_pipes) g_jit_pipes = [NSMutableArray array];
    NSError *err = nil;
    NSString *src = [NSString stringWithContentsOfFile:@"/tmp/rail_jit_kernel.metal"
                                              encoding:NSUTF8StringEncoding
                                                 error:&err];
    if (!src) {
        fprintf(stderr, "tgl_jit_compile: read /tmp/rail_jit_kernel.metal failed: %s\n",
                err ? [[err localizedDescription] UTF8String] : "?");
        return -1;
    }
    MTLCompileOptions *opts = [MTLCompileOptions new];
    id<MTLLibrary> lib = [g_device newLibraryWithSource:src options:opts error:&err];
    if (!lib) {
        fprintf(stderr, "tgl_jit_compile: newLibraryWithSource failed: %s\n",
                err ? [[err localizedDescription] UTF8String] : "?");
        return -1;
    }
    NSArray *names = [lib functionNames];
    if (names.count == 0) {
        fprintf(stderr, "tgl_jit_compile: no functions in compiled library\n");
        return -1;
    }
    id<MTLFunction> fn = [lib newFunctionWithName:names[0]];
    if (!fn) return -1;
    id<MTLComputePipelineState> p = [g_device newComputePipelineStateWithFunction:fn error:&err];
    if (!p) return -1;
    [g_jit_pipes addObject:p];
    return (int)(g_jit_pipes.count - 1);
}

// Minimal 1-in-1-out dispatcher for the JIT'd kernel.  Same buffer/grid
// shape as tgl_silu_fwd minus the sigmoid-save output, so the smoke test
// can feed it the same MSL signature.
int tgl_jit_dispatch_1in1out(int kid, const double *X, double *Y, int N) {
    if (kid < 0 || !g_jit_pipes || kid >= (int)g_jit_pipes.count) return -1;
    if (ensure_init() != 1) return -1;
    const double *Xp = X + 1;
    double *Yp = Y + 1;
    @autoreleasepool {
        id<MTLBuffer> bX = pool_acquire(N * 4);
        id<MTLBuffer> bY = pool_acquire(N * 4);
        f64_to_f32(Xp, (float*)bX.contents, N);
        id<MTLComputePipelineState> p = g_jit_pipes[kid];
        id<MTLCommandBuffer> cmd = [g_queue commandBuffer];
        id<MTLComputeCommandEncoder> enc = [cmd computeCommandEncoder];
        [enc setComputePipelineState:p];
        [enc setBuffer:bX offset:0 atIndex:0];
        [enc setBuffer:bY offset:0 atIndex:1];
        uint32_t nu = (uint32_t)N;
        [enc setBytes:&nu length:4 atIndex:2];
        dispatch_1d(enc, N);
        [enc endEncoding];
        [cmd commit]; [cmd waitUntilCompleted];
        f32_to_f64((float*)bY.contents, Yp, N);
        pool_release(bX); pool_release(bY);
    }
    return 1;
}

// ────────────────────────────────────────────────────────────────────
// bf16 host-side staging helpers (added 2026-05-14).
// bf16 layout = upper 16 bits of f32 representation; conversion is
// just a bit-shift with round-to-nearest-even.
// ────────────────────────────────────────────────────────────────────
static inline void f64_to_bf16(const double *src, uint16_t *dst, int n) {
    for (int i = 0; i < n; i++) {
        float f = (float)src[i];
        uint32_t bits;
        memcpy(&bits, &f, 4);
        uint32_t lsb = (bits >> 16) & 1u;
        uint32_t rounding_bias = 0x7FFFu + lsb;
        bits += rounding_bias;
        dst[i] = (uint16_t)(bits >> 16);
    }
}

static inline void bf16_to_f64(const uint16_t *src, double *dst, int n) {
    for (int i = 0; i < n; i++) {
        uint32_t bits = ((uint32_t)src[i]) << 16;
        float f;
        memcpy(&f, &bits, 4);
        dst[i] = (double)f;
    }
}

// Drop-in for tgl_matmul_f16 — same Rail-side signature, bf16 internally.
int tgl_matmul_bf16(const double *A, const double *B, double *C, int M, int K, int N) {
    if (ensure_init() != 1) return -1;
    const double *Aptr = A + 1, *Bptr = B + 1;
    double *Cptr = C + 1;

    @autoreleasepool {
        uint32_t sA = M*K, sB = K*N, sC = M*N;
        id<MTLBuffer> bufA = pool_acquire(sA*2);
        id<MTLBuffer> bufB = pool_acquire(sB*2);
        id<MTLBuffer> bufC = pool_acquire(sC*2);

        f64_to_bf16(Aptr, (uint16_t*)bufA.contents, sA);
        f64_to_bf16(Bptr, (uint16_t*)bufB.contents, sB);

        id<MTLComputePipelineState> p = pso(@"matmul_bf16");
        if (!p) return -1;

        id<MTLCommandBuffer> cmd = [g_queue commandBuffer];
        id<MTLComputeCommandEncoder> enc = [cmd computeCommandEncoder];
        [enc setComputePipelineState:p];
        [enc setBuffer:bufA offset:0 atIndex:0];
        [enc setBuffer:bufB offset:0 atIndex:1];
        [enc setBuffer:bufC offset:0 atIndex:2];
        uint32_t mu=M, ku=K, nu=N;
        [enc setBytes:&mu length:4 atIndex:3];
        [enc setBytes:&ku length:4 atIndex:4];
        [enc setBytes:&nu length:4 atIndex:5];
        [enc dispatchThreadgroups:MTLSizeMake((N+15)/16, (M+15)/16, 1)
           threadsPerThreadgroup:MTLSizeMake(16, 16, 1)];
        [enc endEncoding];
        [cmd commit]; [cmd waitUntilCompleted];

        bf16_to_f64((uint16_t*)bufC.contents, Cptr, sC);
        pool_release(bufA); pool_release(bufB); pool_release(bufC);
    }
    return 1;
}

// ────────────────────────────────────────────────────────────────────
// Phase 1.1 — JIT dispatcher for the fused rmsnorm+QKV pattern.
// 5 input buffers (X, G, WQ, WK, WV f64), 3 output buffers (Q, K, V f64),
// 2 constants (SEQ, D).  Stages f64↔f32 host-side, dispatches one
// threadgroup per row with 64 threads.
// ────────────────────────────────────────────────────────────────────
// 8-arg form: caller passes a single packed output buffer
// (3*SEQ*D + SEQ doubles: Q | K | V | R) and SEQ/D packed as (SEQ<<16)|D.
// R is the rstd row vector needed by the backward pass.
int tgl_jit_dispatch_rmsnorm_qkv(int kid,
                                  const double *X, const double *G,
                                  const double *WQ, const double *WK, const double *WV,
                                  double *QKV_packed,
                                  int SEQ_D) {
    if (kid < 0 || !g_jit_pipes || kid >= (int)g_jit_pipes.count) return -1;
    if (ensure_init() != 1) return -1;
    int SEQ = (SEQ_D >> 16) & 0xFFFF;
    int D   =  SEQ_D        & 0xFFFF;
    const double *Xp=X+1, *Gp=G+1, *WQp=WQ+1, *WKp=WK+1, *WVp=WV+1;
    double *QKVp = QKV_packed + 1;
    uint32_t sX=SEQ*D, sW=D*D, sQ=SEQ*D, sQKV=3*SEQ*D, sQKVR=sQKV+SEQ, sQKVRL=sQKVR+sQ;
    @autoreleasepool {
        id<MTLBuffer> bX = pool_acquire(sX*4);
        id<MTLBuffer> bG = pool_acquire(D*4);
        id<MTLBuffer> bWQ = pool_acquire(sW*4);
        id<MTLBuffer> bWK = pool_acquire(sW*4);
        id<MTLBuffer> bWV = pool_acquire(sW*4);
        // Single Metal output buffer; kernel sees Q/K/V/R/LN1 as 5 device
        // pointers via setBuffer:offset: at five byte-offsets into one buffer.
        // Layout: Q | K | V | R | LN1  (total 4*SEQ*D + SEQ floats)
        id<MTLBuffer> bQKV = pool_acquire(sQKVRL*4);
        f64_to_f32(Xp,  (float*)bX.contents,  sX);
        f64_to_f32(Gp,  (float*)bG.contents,  D);
        f64_to_f32(WQp, (float*)bWQ.contents, sW);
        f64_to_f32(WKp, (float*)bWK.contents, sW);
        f64_to_f32(WVp, (float*)bWV.contents, sW);

        id<MTLComputePipelineState> p = g_jit_pipes[kid];
        id<MTLCommandBuffer> cmd = [g_queue commandBuffer];
        id<MTLComputeCommandEncoder> enc = [cmd computeCommandEncoder];
        [enc setComputePipelineState:p];
        [enc setBuffer:bX  offset:0          atIndex:0];
        [enc setBuffer:bG  offset:0          atIndex:1];
        [enc setBuffer:bWQ offset:0          atIndex:2];
        [enc setBuffer:bWK offset:0          atIndex:3];
        [enc setBuffer:bWV offset:0          atIndex:4];
        [enc setBuffer:bQKV offset:0          atIndex:5];  // Q range
        [enc setBuffer:bQKV offset:sQ*4       atIndex:6];  // K range
        [enc setBuffer:bQKV offset:2*sQ*4     atIndex:7];  // V range
        uint32_t su=SEQ, du=D;
        [enc setBytes:&su length:4 atIndex:8];
        [enc setBytes:&du length:4 atIndex:9];
        [enc setBuffer:bQKV offset:3*sQ*4     atIndex:10]; // R range (rstd, SEQ floats)
        [enc setBuffer:bQKV offset:(3*sQ+SEQ)*4 atIndex:11]; // LN1 range (normalized X, SEQ*D floats)
        [enc dispatchThreadgroups:MTLSizeMake(SEQ, 1, 1)
          threadsPerThreadgroup:MTLSizeMake(64, 1, 1)];
        [enc endEncoding];
        [cmd commit]; [cmd waitUntilCompleted];
        f32_to_f64((float*)bQKV.contents, QKVp, sQKVRL);
        pool_release(bX); pool_release(bG);
        pool_release(bWQ); pool_release(bWK); pool_release(bWV);
        pool_release(bQKV);
    }
    return 1;
}

// ────────────────────────────────────────────────────────────────────
// Phase 1.4 — Generic 2-in-1-out JIT dispatcher.
// Used by silu+hadamard (SwiGLU's elementwise piece) and any future
// elementwise fusion that takes two same-shape inputs and produces one
// same-shape output.  4 args fit easily in Rail's FFI register window.
// ────────────────────────────────────────────────────────────────────
int tgl_jit_dispatch_2in1out(int kid, const double *A, const double *B,
                              double *C, int N) {
    if (kid < 0 || !g_jit_pipes || kid >= (int)g_jit_pipes.count) return -1;
    if (ensure_init() != 1) return -1;
    const double *Ap = A+1, *Bp = B+1;
    double *Cp = C+1;
    @autoreleasepool {
        id<MTLBuffer> bA = pool_acquire(N*4);
        id<MTLBuffer> bB = pool_acquire(N*4);
        id<MTLBuffer> bC = pool_acquire(N*4);
        f64_to_f32(Ap, (float*)bA.contents, N);
        f64_to_f32(Bp, (float*)bB.contents, N);
        id<MTLComputePipelineState> p = g_jit_pipes[kid];
        id<MTLCommandBuffer> cmd = [g_queue commandBuffer];
        id<MTLComputeCommandEncoder> enc = [cmd computeCommandEncoder];
        [enc setComputePipelineState:p];
        [enc setBuffer:bA offset:0 atIndex:0];
        [enc setBuffer:bB offset:0 atIndex:1];
        [enc setBuffer:bC offset:0 atIndex:2];
        uint32_t nu = (uint32_t)N;
        [enc setBytes:&nu length:4 atIndex:3];
        dispatch_1d(enc, N);
        [enc endEncoding];
        [cmd commit]; [cmd waitUntilCompleted];
        f32_to_f64((float*)bC.contents, Cp, N);
        pool_release(bA); pool_release(bB); pool_release(bC);
    }
    return 1;
}

// ────────────────────────────────────────────────────────────────────
// Specialized silu+hadamard dispatcher.  Like 2in1out, but the output
// buffer is 2*N doubles: out[0..N) holds h_act = silu(G)*U, and
// sig[N..2N) holds sigmoid(G) so the backward pass can reuse it.
// Caller-side FFI is still 5 args (kid, G, U, OutSig, N).
// ────────────────────────────────────────────────────────────────────
int tgl_jit_dispatch_silu_hadamard(int kid, const double *A, const double *B,
                                    double *C, int N) {
    if (kid < 0 || !g_jit_pipes || kid >= (int)g_jit_pipes.count) return -1;
    if (ensure_init() != 1) return -1;
    const double *Ap = A+1, *Bp = B+1;
    double *Cp = C+1;
    @autoreleasepool {
        id<MTLBuffer> bA = pool_acquire(N*4);
        id<MTLBuffer> bB = pool_acquire(N*4);
        // Output buffer holds 2*N floats: out followed by sigmoid.
        id<MTLBuffer> bC = pool_acquire(2*N*4);
        f64_to_f32(Ap, (float*)bA.contents, N);
        f64_to_f32(Bp, (float*)bB.contents, N);
        id<MTLComputePipelineState> p = g_jit_pipes[kid];
        id<MTLCommandBuffer> cmd = [g_queue commandBuffer];
        id<MTLComputeCommandEncoder> enc = [cmd computeCommandEncoder];
        [enc setComputePipelineState:p];
        [enc setBuffer:bA offset:0   atIndex:0];
        [enc setBuffer:bB offset:0   atIndex:1];
        [enc setBuffer:bC offset:0   atIndex:2];  // O range
        uint32_t nu = (uint32_t)N;
        [enc setBytes:&nu length:4 atIndex:3];
        [enc setBuffer:bC offset:N*4 atIndex:4];  // S range (sigmoid)
        dispatch_1d(enc, N);
        [enc endEncoding];
        [cmd commit]; [cmd waitUntilCompleted];
        f32_to_f64((float*)bC.contents, Cp, 2*N);
        pool_release(bA); pool_release(bB); pool_release(bC);
    }
    return 1;
}

// ────────────────────────────────────────────────────────────────────
// EXACT INTEGER MATMUL — bit-reproducible fixed-point accumulation
// Aq,Bq are Rail float_arr holding EXACT-INTEGER doubles (the fixed-point
// quantized inputs, |q| < 2^31, so exactly representable in f64 AND in
// int32). They are packed into int32 GPU buffers WITHOUT going through the
// f64->f32 path (which would corrupt magnitudes above 2^24). The kernel
// returns the per-output 2-limb accumulator as int64 [hi, lo]; those are
// written back into the Rail float_arr Hi,Lo as exact-integer doubles
// (hi, lo in [0,2^31), both within f64's 2^53 exact-integer range for the
// fixed-point regime). Bit-identical to the CPU reference in bx2.
// ────────────────────────────────────────────────────────────────────
int tgl_exact_matmul(const double *Aq, const double *Bq,
                     double *Hi, double *Lo, int M, int K, int N) {
    if (ensure_init() != 1) return -1;
    const double *Aptr = Aq + 1, *Bptr = Bq + 1;
    double *Hptr = Hi + 1, *Lptr = Lo + 1;

    @autoreleasepool {
        uint32_t sA = (uint32_t)M*K, sB = (uint32_t)K*N, sC = (uint32_t)M*N;
        id<MTLBuffer> bA = pool_acquire((NSUInteger)sA * 4);   // int32 in
        id<MTLBuffer> bB = pool_acquire((NSUInteger)sB * 4);   // int32 in
        id<MTLBuffer> bH = pool_acquire((NSUInteger)sC * 8);   // int64 out
        id<MTLBuffer> bL = pool_acquire((NSUInteger)sC * 8);   // int64 out

        int32_t *Ai = (int32_t*)bA.contents;
        int32_t *Bi = (int32_t*)bB.contents;
        for (uint32_t t = 0; t < sA; t++) Ai[t] = (int32_t)Aptr[t];
        for (uint32_t t = 0; t < sB; t++) Bi[t] = (int32_t)Bptr[t];

        id<MTLComputePipelineState> p = pso(@"exact_matmul");
        if (!p) return -1;

        id<MTLCommandBuffer> cmd = [g_queue commandBuffer];
        id<MTLComputeCommandEncoder> enc = [cmd computeCommandEncoder];
        [enc setComputePipelineState:p];
        [enc setBuffer:bA offset:0 atIndex:0];
        [enc setBuffer:bB offset:0 atIndex:1];
        [enc setBuffer:bH offset:0 atIndex:2];
        [enc setBuffer:bL offset:0 atIndex:3];
        uint32_t mu=M, ku=K, nu=N;
        [enc setBytes:&mu length:4 atIndex:4];
        [enc setBytes:&ku length:4 atIndex:5];
        [enc setBytes:&nu length:4 atIndex:6];

        // one thread per output element; non-uniform grid (N across, M down)
        NSUInteger tw = 16, th = 16;
        if ((NSUInteger)N < tw) tw = (NSUInteger)N;
        if ((NSUInteger)M < th) th = (NSUInteger)M;
        [enc dispatchThreads:MTLSizeMake((NSUInteger)N, (NSUInteger)M, 1)
          threadsPerThreadgroup:MTLSizeMake(tw, th, 1)];
        [enc endEncoding];
        [cmd commit]; [cmd waitUntilCompleted];

        long *Hl = (long*)bH.contents;
        long *Ll = (long*)bL.contents;
        for (uint32_t t = 0; t < sC; t++) {
            Hptr[t] = (double)Hl[t];
            Lptr[t] = (double)Ll[t];
        }
        pool_release(bA); pool_release(bB); pool_release(bH); pool_release(bL);
    }
    return 1;
}

// Threadgroup-tiled variant of tgl_exact_matmul. Same int32-in / [hi,lo]-out
// ABI; dispatches the exact_matmul_tiled kernel with full 16x16 threadgroups
// (grid padded up to tile multiples, kernel bounds-guards loads + store).
// Result is bit-for-bit identical to tgl_exact_matmul — tiling is a memory
// optimization, the per-output accumulation order is unchanged.
int tgl_exact_matmul_tiled(const double *Aq, const double *Bq,
                           double *Hi, double *Lo, int M, int K, int N) {
    if (ensure_init() != 1) return -1;
    const double *Aptr = Aq + 1, *Bptr = Bq + 1;
    double *Hptr = Hi + 1, *Lptr = Lo + 1;

    @autoreleasepool {
        uint32_t sA = (uint32_t)M*K, sB = (uint32_t)K*N, sC = (uint32_t)M*N;
        id<MTLBuffer> bA = pool_acquire((NSUInteger)sA * 4);   // int32 in
        id<MTLBuffer> bB = pool_acquire((NSUInteger)sB * 4);   // int32 in
        id<MTLBuffer> bH = pool_acquire((NSUInteger)sC * 8);   // int64 out
        id<MTLBuffer> bL = pool_acquire((NSUInteger)sC * 8);   // int64 out

        int32_t *Ai = (int32_t*)bA.contents;
        int32_t *Bi = (int32_t*)bB.contents;
        for (uint32_t t = 0; t < sA; t++) Ai[t] = (int32_t)Aptr[t];
        for (uint32_t t = 0; t < sB; t++) Bi[t] = (int32_t)Bptr[t];

        id<MTLComputePipelineState> p = pso(@"exact_matmul_tiled");
        if (!p) return -1;

        id<MTLCommandBuffer> cmd = [g_queue commandBuffer];
        id<MTLComputeCommandEncoder> enc = [cmd computeCommandEncoder];
        [enc setComputePipelineState:p];
        [enc setBuffer:bA offset:0 atIndex:0];
        [enc setBuffer:bB offset:0 atIndex:1];
        [enc setBuffer:bH offset:0 atIndex:2];
        [enc setBuffer:bL offset:0 atIndex:3];
        uint32_t mu=M, ku=K, nu=N;
        [enc setBytes:&mu length:4 atIndex:4];
        [enc setBytes:&ku length:4 atIndex:5];
        [enc setBytes:&nu length:4 atIndex:6];

        // Tiled: full 16x16 threadgroups; grid padded up to tile multiples so
        // every lane in each group exists (cooperative load + barrier need it).
        NSUInteger gx = ((NSUInteger)N + 15) / 16;
        NSUInteger gy = ((NSUInteger)M + 15) / 16;
        [enc dispatchThreadgroups:MTLSizeMake(gx, gy, 1)
           threadsPerThreadgroup:MTLSizeMake(16, 16, 1)];
        [enc endEncoding];
        [cmd commit]; [cmd waitUntilCompleted];

        long *Hl = (long*)bH.contents;
        long *Ll = (long*)bL.contents;
        for (uint32_t t = 0; t < sC; t++) {
            Hptr[t] = (double)Hl[t];
            Lptr[t] = (double)Ll[t];
        }
        pool_release(bA); pool_release(bB); pool_release(bH); pool_release(bL);
    }
    return 1;
}
