#!/usr/bin/env bash
# SPDX-License-Identifier: AGPL-3.0-or-later
# Copyright (C) 2026 David (Corey) Herrera
# build_train_gpu.sh — build the on-device GPU training driver (Phase 3, Chunk B).
#
# Runs on a K80 host (GPU 0), from any directory.  Compiles the CPU-oracle
# translation units with gcc-11 first (self-contained; no BLAS), then nvcc
# device-links the driver against them.  Mirrors the Phase-2/3 recipe.
#
#   bash cuda_kernels/build_train_gpu.sh     (outputs land in <repo>/build/)
#
# Requires data/enwik8 (exactly 100,000,000 bytes) for a run.  Then e.g.:
#   CUDA_VISIBLE_DEVICES=0 build/train_gpu --opt sr4   --steps 1000
#   CUDA_VISIBLE_DEVICES=0 build/train_gpu --opt adamw --steps 1000
set -euo pipefail
cd "$(dirname "$0")/.."

CUDA=${CUDA:-/usr/local/cuda-11.4}
NVCC=$CUDA/bin/nvcc
CC=${CC:-gcc-11}
ARCH=sm_37

mkdir -p build
OBJ=build

# CPU oracle objects (gcc-11; -O2, NO -ffast-math so the SR4 scale reduction is
# not reassociated — keeps the GPU SR4 step bit-exact vs the reference core).
$CC -O2 -I. -c rwkvx/rx_pretrain.c            -o $OBJ/rx_pretrain_tg.o
$CC -O2 -I. -c data/fd_byte_dataset.c         -o $OBJ/fd_byte_dataset_tg.o
$CC -O2 -I. -c ternary_optim/ternary_optim.c  -o $OBJ/ternary_optim_tg.o

# Driver (includes cuda_kernels/fullmodel.cu + cuda_kernels/optim.cu).
# opt#3: --default-stream per-thread routes every <<<...,0>>> kernel launch (and,
# via an explicit cublasSetStream, cuBLAS) onto the per-thread default stream,
# which is capturable — the prerequisite for capturing the whole train step as a
# CUDA graph.  The standalone test_*_cuda builds do NOT use this flag (legacy
# default stream), and the shared code is byte-identical either way.
$NVCC -O2 -arch=$ARCH -ccbin $CC -Wno-deprecated-gpu-targets \
  --default-stream per-thread \
  -isystem cuda_kernels/nvcc_shim -I. \
  cuda_kernels/train_gpu.cu \
  $OBJ/rx_pretrain_tg.o $OBJ/fd_byte_dataset_tg.o $OBJ/ternary_optim_tg.o \
  -L$CUDA/lib64 -lcublas -lm -o build/train_gpu

echo "built build/train_gpu"