//! `CfcTrunk` — owns CfC weights, multi-horizon heads, projection, and a
//! ping-pong hidden buffer. `forward_snapshot` dispatches the four
//! perception kernels sequentially. The Mamba2 prefix (per the
//! 2026-05-16 spec amendment) lands in a follow-up task before Graph A
//! capture; this task validates the CfC + heads + projection forward
//! path standalone.
//!
//! Per the addendum API patterns: cache `Arc<CudaStream>` and the
//! `CudaFunction` handles in the struct so the hot path avoids module
//! reload.

use std::sync::Arc;

use anyhow::{Context, Result};
use cudarc::driver::sys::{CUgraphInstantiate_flags, CUstreamCaptureMode};
use cudarc::driver::{CudaFunction, CudaGraph, CudaModule, CudaSlice, CudaStream, DevicePtr, DevicePtrMut, LaunchConfig, PushKernelArg};
use ml_core::device::MlDevice;
use rand::{Rng, SeedableRng};
use rand_chacha::ChaCha8Rng;

use crate::cfc::snap_features::{Mbp10RawInput, ES_TICK_SIZE, FEATURE_DIM, REGIME_DIM};
use crate::heads::{HIDDEN_DIM, N_HORIZONS, PROJ_DIM};
use crate::pinned_mem::MappedF32Buffer;

const SNAP_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/snap_feature_assemble.cubin"));
const STEP_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/cfc_step.cubin"));
const HEADS_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/multi_horizon_heads.cubin"));
const PROJ_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/projection.cubin"));

#[derive(Clone, Debug)]
pub struct CfcConfig {
    pub n_in: usize,
    pub n_hid: usize,
}

impl Default for CfcConfig {
    fn default() -> Self {
        Self { n_in: FEATURE_DIM, n_hid: HIDDEN_DIM }
    }
}

pub struct CfcTrunk {
    cfg: CfcConfig,
    stream: Arc<CudaStream>,

    // Cubin modules (kept alive for the function refs below)
    _snap_module: Arc<CudaModule>,
    _step_module: Arc<CudaModule>,
    _heads_module: Arc<CudaModule>,
    _proj_module: Arc<CudaModule>,
    snap_fn: CudaFunction,
    step_fn: CudaFunction,
    heads_fn: CudaFunction,
    proj_fn: CudaFunction,

    // Device-resident weights
    w_in_d: CudaSlice<f32>,
    w_rec_d: CudaSlice<f32>,
    b_d: CudaSlice<f32>,
    tau_d: CudaSlice<f32>,
    heads_w_d: CudaSlice<f32>,
    heads_b_d: CudaSlice<f32>,
    proj_w_d: CudaSlice<f32>,
    proj_b_d: CudaSlice<f32>,
    proj_g_d: CudaSlice<f32>,
    proj_n_d: CudaSlice<f32>,

    // Ping-pong hidden state
    h_ping: CudaSlice<f32>,
    h_pong: CudaSlice<f32>,

    // Pre-allocated per-step scratch (raw input + features + outputs)
    bid_px_d: CudaSlice<f32>,
    bid_sz_d: CudaSlice<f32>,
    ask_px_d: CudaSlice<f32>,
    ask_sz_d: CudaSlice<f32>,
    prev_bid_sz_d: CudaSlice<f32>,
    prev_ask_sz_d: CudaSlice<f32>,
    regime_d: CudaSlice<f32>,
    snap_feat_d: CudaSlice<f32>,
    probs_d: CudaSlice<f32>,
    proj_out_d: CudaSlice<f32>,

    // Pre-allocated mapped-pinned staging slots for raw input. Reused
    // across upload_into calls so no host-malloc happens during or
    // around CUDA Graph capture (allocations on the capturing thread
    // cause CUDA_ERROR_STREAM_CAPTURE_ISOLATION).
    stg_bid_px: MappedF32Buffer,
    stg_bid_sz: MappedF32Buffer,
    stg_ask_px: MappedF32Buffer,
    stg_ask_sz: MappedF32Buffer,
    stg_regime: MappedF32Buffer,

    // CUDA Graph A — captured perception forward. Scalars (dt_s, ts_ns,
    // prev_mid, ...) are frozen at capture time per cudarc 0.19 semantics
    // (`launch_builder.arg(&scalar)` records the value, not the pointer).
    // The trunk re-captures via `capture_graph_a` whenever those scalars
    // would change meaningfully. A follow-up task moves scalars into a
    // device-resident buffer so the graph stays valid across changes.
    graph_a: Option<CudaGraph>,
}

impl CfcTrunk {
    pub fn new_random(dev: &MlDevice, cfg: &CfcConfig, seed: u64) -> Result<Self> {
        let stream: Arc<CudaStream> = dev.cuda_stream().context("trunk stream")?.clone();
        let ctx = dev.cuda_context().context("trunk ctx")?;

        let snap_module = ctx.load_cubin(SNAP_CUBIN.to_vec()).context("load snap cubin")?;
        let step_module = ctx.load_cubin(STEP_CUBIN.to_vec()).context("load step cubin")?;
        let heads_module = ctx.load_cubin(HEADS_CUBIN.to_vec()).context("load heads cubin")?;
        let proj_module = ctx.load_cubin(PROJ_CUBIN.to_vec()).context("load proj cubin")?;
        let snap_fn = snap_module.load_function("snap_feature_assemble").context("snap_fn")?;
        let step_fn = step_module.load_function("cfc_step").context("step_fn")?;
        let heads_fn = heads_module.load_function("multi_horizon_heads").context("heads_fn")?;
        let proj_fn = proj_module.load_function("projection_kernel").context("proj_fn")?;

        // Random init.
        let mut r = ChaCha8Rng::seed_from_u64(seed);
        let scale_in = (1.0_f32 / cfg.n_in as f32).sqrt();
        let scale_rec = (1.0_f32 / cfg.n_hid as f32).sqrt();
        let w_in: Vec<f32> = (0..cfg.n_hid * cfg.n_in).map(|_| r.gen_range(-scale_in..scale_in)).collect();
        let w_rec: Vec<f32> = (0..cfg.n_hid * cfg.n_hid).map(|_| r.gen_range(-scale_rec..scale_rec)).collect();
        let b: Vec<f32> = vec![0.0; cfg.n_hid];
        // tau log-uniform over [10ms, 1000s] per spec Section 2 (ln(0.01)=-4.6, ln(1000)=6.9).
        let tau: Vec<f32> = (0..cfg.n_hid)
            .map(|_| {
                let u: f32 = r.gen_range(0.0..1.0);
                (-4.6 + u * (4.6 + 6.9)).exp()
            })
            .collect();
        let head_scale = scale_rec;
        let heads_w: Vec<f32> = (0..N_HORIZONS * cfg.n_hid).map(|_| r.gen_range(-head_scale..head_scale)).collect();
        let heads_b: Vec<f32> = vec![0.0; N_HORIZONS];
        let proj_w: Vec<f32> = (0..PROJ_DIM * cfg.n_hid).map(|_| r.gen_range(-head_scale..head_scale)).collect();
        let proj_b: Vec<f32> = vec![0.0; PROJ_DIM];
        let proj_g: Vec<f32> = vec![1.0; PROJ_DIM];
        let proj_n: Vec<f32> = vec![0.0; PROJ_DIM];

        let stg_bid_px = unsafe { MappedF32Buffer::new(10) }
            .map_err(|e| anyhow::anyhow!("stg bid_px alloc: {e}"))?;
        let stg_bid_sz = unsafe { MappedF32Buffer::new(10) }
            .map_err(|e| anyhow::anyhow!("stg bid_sz alloc: {e}"))?;
        let stg_ask_px = unsafe { MappedF32Buffer::new(10) }
            .map_err(|e| anyhow::anyhow!("stg ask_px alloc: {e}"))?;
        let stg_ask_sz = unsafe { MappedF32Buffer::new(10) }
            .map_err(|e| anyhow::anyhow!("stg ask_sz alloc: {e}"))?;
        let stg_regime = unsafe { MappedF32Buffer::new(REGIME_DIM) }
            .map_err(|e| anyhow::anyhow!("stg regime alloc: {e}"))?;

        Ok(Self {
            cfg: cfg.clone(),
            w_in_d: upload(&stream, &w_in)?,
            w_rec_d: upload(&stream, &w_rec)?,
            b_d: upload(&stream, &b)?,
            tau_d: upload(&stream, &tau)?,
            heads_w_d: upload(&stream, &heads_w)?,
            heads_b_d: upload(&stream, &heads_b)?,
            proj_w_d: upload(&stream, &proj_w)?,
            proj_b_d: upload(&stream, &proj_b)?,
            proj_g_d: upload(&stream, &proj_g)?,
            proj_n_d: upload(&stream, &proj_n)?,
            h_ping: stream.alloc_zeros::<f32>(cfg.n_hid).context("h_ping alloc")?,
            h_pong: stream.alloc_zeros::<f32>(cfg.n_hid).context("h_pong alloc")?,
            bid_px_d: stream.alloc_zeros::<f32>(10).context("bid_px alloc")?,
            bid_sz_d: stream.alloc_zeros::<f32>(10).context("bid_sz alloc")?,
            ask_px_d: stream.alloc_zeros::<f32>(10).context("ask_px alloc")?,
            ask_sz_d: stream.alloc_zeros::<f32>(10).context("ask_sz alloc")?,
            prev_bid_sz_d: stream.alloc_zeros::<f32>(10).context("prev_bid_sz alloc")?,
            prev_ask_sz_d: stream.alloc_zeros::<f32>(10).context("prev_ask_sz alloc")?,
            regime_d: stream.alloc_zeros::<f32>(REGIME_DIM).context("regime alloc")?,
            snap_feat_d: stream.alloc_zeros::<f32>(FEATURE_DIM).context("snap_feat alloc")?,
            probs_d: stream.alloc_zeros::<f32>(N_HORIZONS).context("probs alloc")?,
            proj_out_d: stream.alloc_zeros::<f32>(PROJ_DIM).context("proj_out alloc")?,
            stream,
            _snap_module: snap_module,
            _step_module: step_module,
            _heads_module: heads_module,
            _proj_module: proj_module,
            snap_fn,
            step_fn,
            heads_fn,
            proj_fn,
            stg_bid_px,
            stg_bid_sz,
            stg_ask_px,
            stg_ask_sz,
            stg_regime,
            graph_a: None,
        })
    }

    fn upload_pre_allocated(
        stream: &Arc<CudaStream>,
        host: &[f32],
        staging: &MappedF32Buffer,
        dst: &mut CudaSlice<f32>,
    ) -> Result<()> {
        let n = host.len();
        assert_eq!(n, dst.len());
        assert_eq!(n, staging.len);
        staging.write_from_slice(host);
        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("upload_pre_allocated DtoD")?;
            }
        }
        Ok(())
    }

    /// Capture the four perception kernels into CUDA Graph A. The
    /// scalars (dt_s, ts_ns, prev_mid, trade flow, trade count) at the
    /// time of capture are baked in; production usage must re-capture
    /// when those would change. For Phase A training the scalars come
    /// from the trainer's per-batch context and capture happens fresh
    /// each call.
    pub fn capture_graph_a(&mut self, input: &Mbp10RawInput) -> Result<()> {
        // Stage raw input into device buffers BEFORE capture using
        // pre-allocated mapped-pinned staging (host-malloc during
        // capture trips CUDA_ERROR_STREAM_CAPTURE_ISOLATION).
        Self::upload_pre_allocated(&self.stream, &input.bid_px, &self.stg_bid_px, &mut self.bid_px_d)?;
        Self::upload_pre_allocated(&self.stream, &input.bid_sz, &self.stg_bid_sz, &mut self.bid_sz_d)?;
        Self::upload_pre_allocated(&self.stream, &input.ask_px, &self.stg_ask_px, &mut self.ask_px_d)?;
        Self::upload_pre_allocated(&self.stream, &input.ask_sz, &self.stg_ask_sz, &mut self.ask_sz_d)?;
        Self::upload_pre_allocated(&self.stream, &input.regime, &self.stg_regime, &mut self.regime_d)?;
        self.stream.synchronize().context("graph A: pre-capture sync")?;

        // cudarc's default event tracking creates cross-stream deps that
        // the capture engine rejects with STREAM_CAPTURE_ISOLATION.
        // Mirror the working pattern in crates/ml/.../fused_training.rs.
        let ctx = self.stream.context().clone();
        let _ = ctx.check_err();
        unsafe { ctx.disable_event_tracking(); }

        let begin_result = self
            .stream
            .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED);
        if let Err(e) = begin_result {
            unsafe { ctx.enable_event_tracking(); }
            let _ = ctx.check_err();
            return Err(anyhow::anyhow!("graph A begin_capture: {e}"));
        }

        let dispatch_result = self.dispatch_perception(input);

        let graph_result = self
            .stream
            .end_capture(CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH);

        unsafe { ctx.enable_event_tracking(); }
        let _ = ctx.check_err();
        dispatch_result?;

        let graph = graph_result
            .context("graph A end_capture")?
            .ok_or_else(|| anyhow::anyhow!("end_capture returned None — no work was captured"))?;
        self.graph_a = Some(graph);
        Ok(())
    }

    /// Replay the captured graph. Returns probs[5] + projection[8].
    /// Caller is responsible for updating the trunk's pre-allocated
    /// input buffers (via `update_input_buffers`) before calling this.
    pub fn perception_forward_captured(&mut self) -> Result<([f32; N_HORIZONS], [f32; PROJ_DIM])> {
        let graph = self
            .graph_a
            .as_ref()
            .context("Graph A not captured — call capture_graph_a first")?;
        graph.launch().context("graph A launch")?;
        self.stream.synchronize().context("graph A sync")?;

        let probs_vec = download(&self.stream, &self.probs_d)?;
        let proj_vec = download(&self.stream, &self.proj_out_d)?;
        let mut probs = [0f32; N_HORIZONS];
        probs.copy_from_slice(&probs_vec);
        let mut proj = [0f32; PROJ_DIM];
        proj.copy_from_slice(&proj_vec);
        Ok((probs, proj))
    }

    fn dispatch_perception(&mut self, input: &Mbp10RawInput) -> Result<()> {
        let prev_mid = input.prev_mid;
        let trade_signed_vol = input.trade_signed_vol;
        let trade_count = input.trade_count as i32;
        let ts_ns = input.ts_ns as i64;
        let prev_ts_ns = input.prev_ts_ns as i64;
        let tick_size = ES_TICK_SIZE;
        let cfg1 = LaunchConfig { grid_dim: (1, 1, 1), block_dim: (1, 1, 1), shared_mem_bytes: 0 };
        {
            let mut launch = self.stream.launch_builder(&self.snap_fn);
            launch
                .arg(&self.bid_px_d).arg(&self.bid_sz_d).arg(&self.ask_px_d).arg(&self.ask_sz_d)
                .arg(&self.prev_bid_sz_d).arg(&self.prev_ask_sz_d)
                .arg(&self.regime_d)
                .arg(&prev_mid).arg(&trade_signed_vol).arg(&trade_count)
                .arg(&ts_ns).arg(&prev_ts_ns).arg(&tick_size)
                .arg(&mut self.snap_feat_d);
            unsafe { launch.launch(cfg1).context("snap launch")?; }
        }

        let dt_s = input.ts_ns.saturating_sub(input.prev_ts_ns) as f32 * 1e-9;
        let n_in_i = self.cfg.n_in as i32;
        let n_hid_i = self.cfg.n_hid as i32;
        let block_dim = 128.min(self.cfg.n_hid as u32);
        let grid_dim = ((self.cfg.n_hid as u32) + block_dim - 1) / block_dim;
        let cfg2 = LaunchConfig { grid_dim: (grid_dim, 1, 1), block_dim: (block_dim, 1, 1), shared_mem_bytes: 0 };
        {
            let mut launch = self.stream.launch_builder(&self.step_fn);
            launch
                .arg(&self.w_in_d).arg(&self.w_rec_d).arg(&self.b_d).arg(&self.tau_d)
                .arg(&self.snap_feat_d).arg(&self.h_ping)
                .arg(&dt_s).arg(&n_in_i).arg(&n_hid_i)
                .arg(&mut self.h_pong);
            unsafe { launch.launch(cfg2).context("step launch")?; }
        }
        // NOTE: ping-pong swap intentionally OMITTED inside the captured
        // path. Inside Graph A we always write to h_pong from h_ping; the
        // swap happens via a DtoD copy outside the captured region after
        // each replay. Doing the swap as `std::mem::swap` would change
        // pointer identity, which breaks captured kernel args.
        let cfg3 = LaunchConfig { grid_dim: (1, 1, 1), block_dim: (N_HORIZONS as u32, 1, 1), shared_mem_bytes: 0 };
        {
            let mut launch = self.stream.launch_builder(&self.heads_fn);
            launch.arg(&self.heads_w_d).arg(&self.heads_b_d).arg(&self.h_pong).arg(&mut self.probs_d);
            unsafe { launch.launch(cfg3).context("heads launch")?; }
        }
        let cfg4 = LaunchConfig { grid_dim: (1, 1, 1), block_dim: (PROJ_DIM as u32, 1, 1), shared_mem_bytes: 0 };
        {
            let mut launch = self.stream.launch_builder(&self.proj_fn);
            launch
                .arg(&self.proj_w_d).arg(&self.proj_b_d)
                .arg(&self.proj_g_d).arg(&self.proj_n_d)
                .arg(&self.h_pong).arg(&mut self.proj_out_d);
            unsafe { launch.launch(cfg4).context("proj launch")?; }
        }
        Ok(())
    }

    /// Mirror snapshot input into the pre-allocated device buffers
    /// without dispatching the kernels (used between Graph A replays).
    pub fn update_input_buffers(&mut self, input: &Mbp10RawInput) -> Result<()> {
        Self::upload_pre_allocated(&self.stream, &input.bid_px, &self.stg_bid_px, &mut self.bid_px_d)?;
        Self::upload_pre_allocated(&self.stream, &input.bid_sz, &self.stg_bid_sz, &mut self.bid_sz_d)?;
        Self::upload_pre_allocated(&self.stream, &input.ask_px, &self.stg_ask_px, &mut self.ask_px_d)?;
        Self::upload_pre_allocated(&self.stream, &input.ask_sz, &self.stg_ask_sz, &mut self.ask_sz_d)?;
        Ok(())
    }

    /// Sequential dispatch: snap_features -> cfc_step -> heads -> projection.
    /// Returns (probs[5], projection[8]). Updates h_ping/h_pong in place.
    pub fn forward_snapshot(
        &mut self,
        input: &Mbp10RawInput,
    ) -> Result<([f32; N_HORIZONS], [f32; PROJ_DIM])> {
        // Stage raw input into pre-allocated device buffers via mapped-pinned.
        upload_into(&self.stream, &input.bid_px, &mut self.bid_px_d)?;
        upload_into(&self.stream, &input.bid_sz, &mut self.bid_sz_d)?;
        upload_into(&self.stream, &input.ask_px, &mut self.ask_px_d)?;
        upload_into(&self.stream, &input.ask_sz, &mut self.ask_sz_d)?;
        upload_into(&self.stream, &input.regime, &mut self.regime_d)?;

        // Step 1: snap_feature_assemble.
        let prev_mid = input.prev_mid;
        let trade_signed_vol = input.trade_signed_vol;
        let trade_count = input.trade_count as i32;
        let ts_ns = input.ts_ns as i64;
        let prev_ts_ns = input.prev_ts_ns as i64;
        let tick_size = ES_TICK_SIZE;
        let cfg1 = LaunchConfig { grid_dim: (1, 1, 1), block_dim: (1, 1, 1), shared_mem_bytes: 0 };
        {
            let mut launch = self.stream.launch_builder(&self.snap_fn);
            launch
                .arg(&self.bid_px_d).arg(&self.bid_sz_d).arg(&self.ask_px_d).arg(&self.ask_sz_d)
                .arg(&self.prev_bid_sz_d).arg(&self.prev_ask_sz_d)
                .arg(&self.regime_d)
                .arg(&prev_mid).arg(&trade_signed_vol).arg(&trade_count)
                .arg(&ts_ns).arg(&prev_ts_ns).arg(&tick_size)
                .arg(&mut self.snap_feat_d);
            unsafe { launch.launch(cfg1).context("snap launch")?; }
        }

        // Step 2: cfc_step (h_ping -> h_pong).
        let dt_s = input.ts_ns.saturating_sub(input.prev_ts_ns) as f32 * 1e-9;
        let n_in_i = self.cfg.n_in as i32;
        let n_hid_i = self.cfg.n_hid as i32;
        let block_dim = 128.min(self.cfg.n_hid as u32);
        let grid_dim = ((self.cfg.n_hid as u32) + block_dim - 1) / block_dim;
        let cfg2 = LaunchConfig { grid_dim: (grid_dim, 1, 1), block_dim: (block_dim, 1, 1), shared_mem_bytes: 0 };
        {
            let mut launch = self.stream.launch_builder(&self.step_fn);
            launch
                .arg(&self.w_in_d).arg(&self.w_rec_d).arg(&self.b_d).arg(&self.tau_d)
                .arg(&self.snap_feat_d).arg(&self.h_ping)
                .arg(&dt_s).arg(&n_in_i).arg(&n_hid_i)
                .arg(&mut self.h_pong);
            unsafe { launch.launch(cfg2).context("step launch")?; }
        }
        std::mem::swap(&mut self.h_ping, &mut self.h_pong);

        // Step 3: heads.
        let cfg3 = LaunchConfig { grid_dim: (1, 1, 1), block_dim: (N_HORIZONS as u32, 1, 1), shared_mem_bytes: 0 };
        {
            let mut launch = self.stream.launch_builder(&self.heads_fn);
            launch.arg(&self.heads_w_d).arg(&self.heads_b_d).arg(&self.h_ping).arg(&mut self.probs_d);
            unsafe { launch.launch(cfg3).context("heads launch")?; }
        }

        // Step 4: projection.
        let cfg4 = LaunchConfig { grid_dim: (1, 1, 1), block_dim: (PROJ_DIM as u32, 1, 1), shared_mem_bytes: 0 };
        {
            let mut launch = self.stream.launch_builder(&self.proj_fn);
            launch
                .arg(&self.proj_w_d).arg(&self.proj_b_d)
                .arg(&self.proj_g_d).arg(&self.proj_n_d)
                .arg(&self.h_ping).arg(&mut self.proj_out_d);
            unsafe { launch.launch(cfg4).context("proj launch")?; }
        }

        // Copy current bid_sz/ask_sz into prev_*_sz for next-step OFI baseline (DtoD).
        copy_dtod(&self.stream, &self.bid_sz_d, &mut self.prev_bid_sz_d)?;
        copy_dtod(&self.stream, &self.ask_sz_d, &mut self.prev_ask_sz_d)?;

        self.stream.synchronize().context("trunk forward sync")?;

        // Readback (slow-path for now; Task 11 captures into Graph A and
        // keeps everything on-device).
        let probs_vec = download(&self.stream, &self.probs_d)?;
        let proj_vec = download(&self.stream, &self.proj_out_d)?;
        let mut probs = [0f32; N_HORIZONS];
        probs.copy_from_slice(&probs_vec);
        let mut proj = [0f32; PROJ_DIM];
        proj.copy_from_slice(&proj_vec);
        Ok((probs, proj))
    }

    pub fn snapshot_hidden(&self) -> Result<Vec<f32>> {
        download(&self.stream, &self.h_ping)
    }

    pub fn config(&self) -> &CfcConfig { &self.cfg }
}

// Helpers ---------------------------------------------------------------

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!("trunk upload staging: {e}"))?;
    staging.write_from_slice(host);
    let mut dst = stream.alloc_zeros::<f32>(n).context("trunk 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("trunk upload DtoD")?;
        }
    }
    Ok(dst)
}

fn upload_into(stream: &Arc<CudaStream>, host: &[f32], dst: &mut CudaSlice<f32>) -> Result<()> {
    let n = host.len();
    assert_eq!(n, dst.len());
    let staging = unsafe { MappedF32Buffer::new(n) }
        .map_err(|e| anyhow::anyhow!("upload_into staging: {e}"))?;
    staging.write_from_slice(host);
    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("upload_into DtoD")?;
    }
    Ok(())
}

fn copy_dtod(stream: &Arc<CudaStream>, src: &CudaSlice<f32>, dst: &mut CudaSlice<f32>) -> Result<()> {
    let n = src.len();
    assert_eq!(n, dst.len());
    let nbytes = n * std::mem::size_of::<f32>();
    unsafe {
        let (src_ptr, _g1) = src.device_ptr(stream);
        let (dst_ptr, _g2) = dst.device_ptr_mut(stream);
        cudarc::driver::result::memcpy_dtod_async(dst_ptr, src_ptr, nbytes, stream.cu_stream())
            .context("copy_dtod")?;
    }
    Ok(())
}

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!("trunk 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("trunk download DtoD")?;
    }
    stream.synchronize().context("trunk download sync")?;
    Ok(staging.read_all())
}
