/* SPDX-License-Identifier: AGPL-3.0-or-later */
/* Copyright (C) 2026 David (Corey) Herrera */
/*
 * fm_arena.cuh — device scratch bump-arena for the resident GPU trainer.
 *
 * Optimization #2 (resident scratch): the per-step forward/backward used to
 * cudaMalloc/cudaFree ~1500 transient device buffers per training step — each
 * cudaFree synchronizes the K80, which dominated step time at the 1M config.
 * This arena pre-allocates ONE device slab and hands out bump-pointer slices;
 * it is reset (rewound to empty) at the top of every fm_forward / fm_backward,
 * so no transient buffer is ever cudaMalloc'd/cudaFree'd in steady state.
 *
 * Parity: this changes ONLY memory management. Every kernel/GEMM sees the same
 * values in the same order, so the loss curve is bit-identical to the pre-arena
 * trainer (both SR4 and AdamW arms).
 *
 * Dual-mode (so the standalone parity/test builds are byte-identical):
 *   - active  (fmA_init called, train_gpu path): fmA_alloc bumps the slab,
 *     fmA_free is a no-op, fmA_reset rewinds. Buffers freed en-masse at reset.
 *   - inactive(default, test_*_cuda builds that include the merged .cu files
 *     without ever calling fmA_init): fmA_alloc == cudaMalloc, fmA_free ==
 *     cudaFree, fmA_reset is a no-op. Identical to the original code.
 *
 * Robustness: if the slab is exhausted while active, fmA_alloc falls back to a
 * tracked cudaMalloc (freed at the next fmA_reset), so correctness never depends
 * on the slab being large enough — only performance does.
 */
#ifndef FM_ARENA_CUH
#define FM_ARENA_CUH
#include <cuda_runtime.h>
#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>

#define FMA_OVF_MAX 512
typedef struct {
    char  *base;        /* device slab */
    size_t cap;         /* slab bytes  */
    size_t off;         /* bump offset */
    int    active;      /* 0 => fall back to cudaMalloc/cudaFree */
    void  *ovf[FMA_OVF_MAX];  /* slab-exhaustion fallbacks, freed at reset */
    int    novf;
    size_t hiwater;     /* peak off (diagnostic) */
} fm_arena_t;

static fm_arena_t g_fmA = {0,0,0,0,{0},0,0};

/* ---- CUDA Graph capture support (opt #3) --------------------------------
 * The resident train step is captured once as a CUDA graph and replayed each
 * step to escape the launch-bound floor.  Two globals steer capture-safety:
 *   g_fm_stream    : the work stream every per-step memcpy/memset is issued on.
 *                    0 in the standalone test builds (legacy default stream,
 *                    matches the kernels' <<<...,0>>>); set to cudaStreamPerThread
 *                    by the train_gpu driver (which compiles with
 *                    --default-stream per-thread, so its <<<...,0>>> kernels and
 *                    cuBLAS-on-per-thread all share one capturable stream).
 *   g_fm_capturing : nonzero WHILE a train step is being stream-captured.  Host
 *                    syncs and the per-step H2D input copies / D2H loss readback
 *                    are skipped during capture (a sync would abort the capture;
 *                    inputs/loss are marshalled by the driver around replay).
 * In every non-capture path (all test builds, the eval forwards, the non-graph
 * fallback) both are their defaults and behaviour is byte-identical to before. */
static cudaStream_t g_fm_stream    = 0;
static int          g_fm_capturing = 0;
/* skip the terminal device sync while capturing (replay is synced by the driver) */
static inline void fm_sync(void){ if(!g_fm_capturing) cudaDeviceSynchronize(); }

static inline void fmA_init(size_t cap_bytes){
    if(g_fmA.base){ cudaFree(g_fmA.base); g_fmA.base=NULL; }
    if(cudaMalloc((void**)&g_fmA.base, cap_bytes)!=cudaSuccess){
        fprintf(stderr,"fmA_init: cudaMalloc(%zu) failed\n",cap_bytes); exit(1);
    }
    g_fmA.cap=cap_bytes; g_fmA.off=0; g_fmA.active=1; g_fmA.novf=0; g_fmA.hiwater=0;
}
static inline void fmA_reset(void){
    if(!g_fmA.active) return;
    for(int i=0;i<g_fmA.novf;i++) cudaFree(g_fmA.ovf[i]);
    g_fmA.novf=0; g_fmA.off=0;
}
/* allocate n bytes of device scratch (256B-aligned). */
static inline void* fmA_alloc(size_t n){
    if(!g_fmA.active){
        void*p=NULL; if(cudaMalloc(&p,n)!=cudaSuccess){
            fprintf(stderr,"fmA_alloc(fallback) cudaMalloc(%zu) failed\n",n); exit(1);} return p;
    }
    size_t a=(n+255)&~(size_t)255;
    if(g_fmA.off+a<=g_fmA.cap){
        void*p=g_fmA.base+g_fmA.off; g_fmA.off+=a;
        if(g_fmA.off>g_fmA.hiwater) g_fmA.hiwater=g_fmA.off;
        return p;
    }
    /* slab exhausted: tracked cudaMalloc fallback (freed next reset) */
    void*p=NULL; if(cudaMalloc(&p,n)!=cudaSuccess){
        fprintf(stderr,"fmA_alloc(ovf) cudaMalloc(%zu) failed\n",n); exit(1);}
    if(g_fmA.novf<FMA_OVF_MAX) g_fmA.ovf[g_fmA.novf++]=p;
    return p;
}
/* free: no-op while active (arena owns it); real cudaFree when inactive. */
static inline void fmA_free(void*p){
    if(!g_fmA.active){ if(p) cudaFree(p); }
}
#endif /* FM_ARENA_CUH */