foxhunt/crates/ml-alpha/src/heads.rs

//! Multi-horizon sigmoid heads (128 → 5) and projection (128 → 8 + layer-norm).
//!
//! Per spec Section 2 + 2026-05-16 stacked amendment, the heads sit on
//! the CfC hidden output and emit 5 sigmoid probabilities of "up" at
//! horizons {30, 100, 300, 1000, 6000} snapshots forward.

use std::sync::Arc;

use anyhow::{Context, Result};
use cudarc::driver::{CudaSlice, CudaStream, DevicePtr, DevicePtrMut, LaunchConfig, PushKernelArg};
use ml_core::device::MlDevice;

use crate::pinned_mem::MappedF32Buffer;

const HEADS_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/multi_horizon_heads.cubin"));

pub const N_HORIZONS: usize = 5;
pub const HORIZONS: [usize; N_HORIZONS] = [30, 100, 300, 1000, 6000];
pub const HIDDEN_DIM: usize = 128;
pub const PROJ_DIM: usize = 8;

#[derive(Clone, Debug)]
pub struct HeadsWeights {
    pub w: Vec<f32>, // [5, 128] row-major
    pub b: Vec<f32>, // [5]
}

#[derive(Clone, Debug)]
pub struct ProjectionWeights {
    pub w: Vec<f32>,        // [8, 128]
    pub b: Vec<f32>,        // [8]
    pub ln_gain: Vec<f32>,  // [8]
    pub ln_bias: Vec<f32>,  // [8]
}

/// Placeholder marker — concrete trunk integration in Task 9.
pub struct MultiHorizonHeads;

pub struct Projection;

pub fn multi_horizon_heads_gpu(
    dev: &MlDevice,
    w: &HeadsWeights,
    h: &[f32],
) -> Result<[f32; N_HORIZONS]> {
    assert_eq!(h.len(), HIDDEN_DIM);
    assert_eq!(w.w.len(), N_HORIZONS * HIDDEN_DIM);
    assert_eq!(w.b.len(), N_HORIZONS);

    let stream: &Arc<CudaStream> = dev.cuda_stream().context("heads stream")?;
    let ctx = dev.cuda_context().context("heads ctx")?;
    let module = ctx.load_cubin(HEADS_CUBIN.to_vec()).context("load heads cubin")?;
    let func = module.load_function("multi_horizon_heads").context("load function")?;

    let w_d = upload(stream, &w.w)?;
    let b_d = upload(stream, &w.b)?;
    let h_d = upload(stream, h)?;
    let mut probs_d = stream.alloc_zeros::<f32>(N_HORIZONS).context("probs alloc")?;

    let cfg = LaunchConfig {
        grid_dim: (1, 1, 1),
        block_dim: (N_HORIZONS as u32, 1, 1),
        shared_mem_bytes: 0,
    };
    let mut launch = stream.launch_builder(&func);
    launch.arg(&w_d).arg(&b_d).arg(&h_d).arg(&mut probs_d);
    unsafe { launch.launch(cfg).context("heads launch")?; }

    let v = download(stream, &probs_d)?;
    let mut out = [0f32; N_HORIZONS];
    out.copy_from_slice(&v);
    Ok(out)
}

fn upload(stream: &Arc<CudaStream>, host: &[f32]) -> Result<CudaSlice<f32>> {
    let n = host.len();
    let staging = unsafe { MappedF32Buffer::new(n) }
        .map_err(|e| anyhow::anyhow!("heads upload staging: {e}"))?;
    staging.write_from_slice(host);
    let mut dst = stream.alloc_zeros::<f32>(n).context("heads upload alloc")?;
    if n > 0 {
        let nbytes = n * std::mem::size_of::<f32>();
        unsafe {
            let (dst_ptr, _g) = dst.device_ptr_mut(stream);
            cudarc::driver::result::memcpy_dtod_async(dst_ptr, staging.dev_ptr, nbytes, stream.cu_stream())
                .context("heads upload DtoD")?;
        }
    }
    Ok(dst)
}

fn download(stream: &Arc<CudaStream>, src: &CudaSlice<f32>) -> Result<Vec<f32>> {
    let n = src.len();
    let staging = unsafe { MappedF32Buffer::new(n) }
        .map_err(|e| anyhow::anyhow!("heads download staging: {e}"))?;
    let nbytes = n * std::mem::size_of::<f32>();
    unsafe {
        let (src_ptr, _g) = src.device_ptr(stream);
        cudarc::driver::result::memcpy_dtod_async(staging.dev_ptr, src_ptr, nbytes, stream.cu_stream())
            .context("heads download DtoD")?;
    }
    stream.synchronize().context("heads download sync")?;
    Ok(staging.read_all())
}