#![allow(unsafe_code)] // Required for CUDA kernel launches.

//! `GpuAuxHeads` — Plan 4 Task 6 Commit A.
//!
//! Multi-task auxiliary heads orchestrator (E.6). Owns the per-head
//! `CudaFunction` handles for both forward / backward / loss-reduce kernels
//! plus the regime-label builder, and exposes thin `pub(crate)` wrappers
//! that mirror the kernel ABIs in row-major device-pointer form.
//!
//! NO PRODUCTION CALLERS in this commit — Commit B's training-loop wire-up
//! invokes these methods. The orchestrator is built as additive scaffolding
//! per the same convention used for `gpu_grn.rs` (1B-i / 2c.2 pattern):
//! the kernels and orchestrator land first, the wire-up commit lights them
//! up.
//!
//! # Architecture
//!
//! Both heads share the `Linear(SH2 → 32) → ELU → Linear(32 → K)` shape:
//!
//!   - Next-bar regression head: K = 1, MSE loss.
//!   - Regime classification head: K = 5, cross-entropy loss.
//!
//! Param tensors live in the trainer's flat `params_buf` at indices
//! `[119..127)` (8 tensors total — w1/b1/w2/b2 per head). See
//! `compute_param_sizes()` in `gpu_dqn_trainer.rs` for the exact layout.
//!
//! # Pearls applied
//!
//! - `feedback_no_atomicadd.md` — backward kernels emit per-sample partials;
//!   the trailing `aux_param_grad_reduce` kernel collapses along the batch
//!   dim with shmem-tree reduction (no atomicAdd anywhere).
//! - `pearl_cold_path_no_exception_to_gpu_drives.md` — every reduction is
//!   GPU-side; the orchestrator does zero CPU/DtoH work.
//! - `feedback_wire_everything_up.md` — this orchestrator is intentionally
//!   built as a pre-wire scaffold for Commit B, mirroring the same
//!   commit-by-commit ordering used for VSN (1B-ii orchestrator → 1B-iii
//!   wire-up). The Commit B wire-up retires the orphan-by-design status.

use std::sync::Arc;

use cudarc::driver::{CudaFunction, CudaStream, LaunchConfig, PushKernelArg};

use crate::MLError;

// ── Cubin embedding (kernels live in aux_heads_kernel.cu / aux_heads_loss_ema_kernel.cu) ──

use super::gpu_dqn_trainer::{AUX_HEADS_CUBIN, AUX_HEADS_LOSS_EMA_CUBIN};

/// Hidden dim shared by both auxiliary heads (Linear_1 output width).
/// Kept in sync with the `AUX_HIDDEN_DIM` macro at the top of
/// `aux_heads_kernel.cu`.
pub(crate) const AUX_HIDDEN_DIM: usize = 32;

/// Output cardinality of the next-bar direction classification head
/// (SP13 B1.1a, 2026-05-05). Flipped 1 → 2: the head was a K=1 MSE
/// regression head; B1.1a flips it to a K=2 softmax classifier reading
/// i32 direction labels in {-1, 0, 1} where -1 = mask/skip. The
/// per-class softmax is structurally bounded in [0, 1] so downstream
/// ISV consumers (slot 375) see naturally-bounded inputs without a
/// runtime tanh squash per `pearl_bounded_modifier_outputs_require_structural_activation`.
pub(crate) const AUX_NEXT_BAR_K: usize = 2;

/// Output cardinality of the regime classification head (5-way softmax).
pub(crate) const AUX_REGIME_K: usize = 5;

/// Block dim used by the per-sample forward / backward kernels.
const AUX_BLOCK: u32 = 256;

/// Forward orchestrator — owns the 5 forward / loss-reduce / label-builder
/// kernel handles for both auxiliary heads.
///
/// Dropping this struct releases the underlying `CudaFunction` handles via
/// cudarc's RAII; nothing else needs explicit teardown.
#[allow(missing_debug_implementations)]
pub(crate) struct AuxHeadsForwardOps {
    next_bar_forward_kernel: CudaFunction,
    regime_forward_kernel: CudaFunction,
    next_bar_loss_reduce_kernel: CudaFunction,
    regime_loss_reduce_kernel: CudaFunction,
    regime_label_kernel: CudaFunction,
    /// Aux loss EMA producer kernel (single thread, single block; mirrors
    /// `h_s2_rms_ema_update`'s shape). Loaded from `AUX_HEADS_LOSS_EMA_CUBIN`.
    ///
    /// SP13 B1.0 (2026-05-05): the sibling label-scale EMA kernel was
    /// retired together with the ISV[117] label-scale divisor in
    /// `aux_next_bar_loss_reduce` / `aux_next_bar_backward`. Labels are
    /// already z-normalised at the data layer; B1.1 will replace MSE
    /// with CE entirely.
    loss_ema_kernel: CudaFunction,
}

impl AuxHeadsForwardOps {
    /// Load all forward / reduce / EMA kernel handles from the precompiled
    /// cubins. Mirrors `GrnBlock::new`'s loader pattern.
    pub(crate) fn new(stream: &Arc<CudaStream>) -> Result<Self, MLError> {
        let context = stream.context();

        let aux_module = context
            .load_cubin(AUX_HEADS_CUBIN.to_vec())
            .map_err(|e| MLError::ModelError(format!("aux_heads cubin load: {e}")))?;
        let next_bar_forward_kernel = aux_module
            .load_function("aux_next_bar_forward")
            .map_err(|e| MLError::ModelError(format!("aux_next_bar_forward load: {e}")))?;
        let regime_forward_kernel = aux_module
            .load_function("aux_regime_forward")
            .map_err(|e| MLError::ModelError(format!("aux_regime_forward load: {e}")))?;
        let next_bar_loss_reduce_kernel = aux_module
            .load_function("aux_next_bar_loss_reduce")
            .map_err(|e| MLError::ModelError(format!("aux_next_bar_loss_reduce load: {e}")))?;
        let regime_loss_reduce_kernel = aux_module
            .load_function("aux_regime_loss_reduce")
            .map_err(|e| MLError::ModelError(format!("aux_regime_loss_reduce load: {e}")))?;
        let regime_label_kernel = aux_module
            .load_function("aux_regime_label_from_states")
            .map_err(|e| MLError::ModelError(format!("aux_regime_label_from_states load: {e}")))?;

        let ema_module = context
            .load_cubin(AUX_HEADS_LOSS_EMA_CUBIN.to_vec())
            .map_err(|e| MLError::ModelError(format!("aux_heads_loss_ema cubin load: {e}")))?;
        let loss_ema_kernel = ema_module
            .load_function("aux_heads_loss_ema_update")
            .map_err(|e| MLError::ModelError(format!("aux_heads_loss_ema_update load: {e}")))?;

        Ok(Self {
            next_bar_forward_kernel,
            regime_forward_kernel,
            next_bar_loss_reduce_kernel,
            regime_loss_reduce_kernel,
            regime_label_kernel,
            loss_ema_kernel,
        })
    }

    /// Launch `aux_next_bar_forward`: `Linear(h_s2) → ELU → Linear → softmax`.
    ///
    /// SP13 B1.1a (2026-05-05): K-flipped 1 → 2 with softmax tile output.
    /// Forward writes BOTH the `[B, K]` logits buffer (saved for backward
    /// re-read; the previous regression-mode design only saved the scalar
    /// prediction) AND the `[B, K]` softmax tile (the new structurally-
    /// bounded prediction surface read by 3 downstream consumers: loss
    /// reduce, backward, dir-acc / isv-tanh producers).
    ///
    /// Caller-owned buffers (raw `u64` device pointers for graph-capture
    /// safety; mirrors `gpu_grn.rs::forward_raw`):
    ///   * `h_s2_ptr`        — `[B, SH2]` row-major.
    ///   * `w1_ptr`/`b1_ptr` — `[H, SH2]` / `[H]`. Slice from `params_buf[119]`/`[120]`.
    ///   * `w2_ptr`/`b2_ptr` — `[K, H]` / `[K]`.   Slice from `params_buf[121]`/`[122]`.
    ///   * `hidden_out_ptr`  — `[B, H]` SAVED post-ELU buffer (consumed by backward).
    ///   * `logits_out_ptr`  — `[B, K]` SAVED logits (consumed by backward
    ///                         only; the consumers below read the softmax tile).
    ///   * `softmax_out_ptr` — `[B, K]` SAVED softmax tile (consumed by
    ///                         loss reduce + backward + dir-acc + isv-tanh).
    ///
    /// `k` is `AUX_NEXT_BAR_K = 2` in B1.1a; passed as runtime arg for
    /// kernel-signature stability (mirrors `forward_regime`'s K-runtime
    /// pattern — a future K bump won't require an ABI change).
    ///
    /// Block: `AUX_BLOCK` threads. Grid: `B` blocks.
    /// Shared mem: `(H + K) * sizeof(f32)` bytes (h cache + softmax temp).
    #[allow(clippy::too_many_arguments)]
    pub(crate) fn forward_next_bar(
        &self,
        stream: &Arc<CudaStream>,
        h_s2_ptr: u64,
        w1_ptr: u64,
        b1_ptr: u64,
        w2_ptr: u64,
        b2_ptr: u64,
        b: usize,
        sh2: usize,
        k: usize,
        hidden_out_ptr: u64,
        logits_out_ptr: u64,
        softmax_out_ptr: u64,
    ) -> Result<(), MLError> {
        let b_i32 = b as i32;
        let sh2_i32 = sh2 as i32;
        let k_i32 = k as i32;
        let smem_bytes = (AUX_HIDDEN_DIM as u32 + k as u32) * std::mem::size_of::<f32>() as u32;
        unsafe {
            stream
                .launch_builder(&self.next_bar_forward_kernel)
                .arg(&h_s2_ptr)
                .arg(&w1_ptr)
                .arg(&b1_ptr)
                .arg(&w2_ptr)
                .arg(&b2_ptr)
                .arg(&b_i32)
                .arg(&sh2_i32)
                .arg(&k_i32)
                .arg(&hidden_out_ptr)
                .arg(&logits_out_ptr)
                .arg(&softmax_out_ptr)
                .launch(LaunchConfig {
                    grid_dim: (b as u32, 1, 1),
                    block_dim: (AUX_BLOCK, 1, 1),
                    shared_mem_bytes: smem_bytes,
                })
                .map_err(|e| MLError::ModelError(format!("aux_next_bar_forward: {e}")))?;
        }
        Ok(())
    }

    /// Launch `aux_regime_forward`: `Linear(h_s2) → ELU → Linear → logits[K]`.
    ///
    /// `K` is `AUX_REGIME_K = 5` in this commit's design; passed as runtime
    /// arg for kernel-signature stability.
    #[allow(clippy::too_many_arguments)]
    pub(crate) fn forward_regime(
        &self,
        stream: &Arc<CudaStream>,
        h_s2_ptr: u64,
        w1_ptr: u64,
        b1_ptr: u64,
        w2_ptr: u64,
        b2_ptr: u64,
        b: usize,
        sh2: usize,
        k: usize,
        hidden_out_ptr: u64,
        logits_out_ptr: u64,
    ) -> Result<(), MLError> {
        let b_i32 = b as i32;
        let sh2_i32 = sh2 as i32;
        let k_i32 = k as i32;
        let smem_bytes = (AUX_HIDDEN_DIM as u32) * std::mem::size_of::<f32>() as u32;
        unsafe {
            stream
                .launch_builder(&self.regime_forward_kernel)
                .arg(&h_s2_ptr)
                .arg(&w1_ptr)
                .arg(&b1_ptr)
                .arg(&w2_ptr)
                .arg(&b2_ptr)
                .arg(&b_i32)
                .arg(&sh2_i32)
                .arg(&k_i32)
                .arg(&hidden_out_ptr)
                .arg(&logits_out_ptr)
                .launch(LaunchConfig {
                    grid_dim: (b as u32, 1, 1),
                    block_dim: (AUX_BLOCK, 1, 1),
                    shared_mem_bytes: smem_bytes,
                })
                .map_err(|e| MLError::ModelError(format!("aux_regime_forward: {e}")))?;
        }
        Ok(())
    }

    /// Launch `aux_next_bar_loss_reduce`: scalar mean cross-entropy +
    /// scalar valid-count.
    ///
    /// SP13 B1.1a (2026-05-05): MSE → softmax CE flip. Reads the `[B, K]`
    /// softmax tile (produced by `forward_next_bar`) and the `[B] i32`
    /// labels buffer (`-1 / 0..K`; -1 is the producer's mask sentinel).
    /// Writes mean CE = `numerator / max(B_valid, 1)` to `loss_out[1]`
    /// AND the `B_valid` count to `valid_count_out[1]` for the backward
    /// kernel to consume (so loss + backward share the same divisor —
    /// derivatives of the same scalar function).
    ///
    /// Two parallel partial-reduction strips (loss numerator + valid
    /// count); shmem matches the K-class `regime_loss_reduce` shape.
    pub(crate) fn next_bar_loss_reduce(
        &self,
        stream: &Arc<CudaStream>,
        softmax_ptr: u64,
        labels_ptr: u64,
        b: usize,
        k: usize,
        loss_out_ptr: u64,
        valid_count_out_ptr: u64,
    ) -> Result<(), MLError> {
        let b_i32 = b as i32;
        let k_i32 = k as i32;
        // Two parallel partial-reduction strips: loss + valid-count. See
        // `aux_next_bar_loss_reduce` shmem layout in the kernel.
        let smem_bytes = 2 * AUX_BLOCK * std::mem::size_of::<f32>() as u32;
        unsafe {
            stream
                .launch_builder(&self.next_bar_loss_reduce_kernel)
                .arg(&softmax_ptr)
                .arg(&labels_ptr)
                .arg(&b_i32)
                .arg(&k_i32)
                .arg(&loss_out_ptr)
                .arg(&valid_count_out_ptr)
                .launch(LaunchConfig {
                    grid_dim: (1, 1, 1),
                    block_dim: (AUX_BLOCK, 1, 1),
                    shared_mem_bytes: smem_bytes,
                })
                .map_err(|e| MLError::ModelError(format!("aux_next_bar_loss_reduce: {e}")))?;
        }
        Ok(())
    }

    /// Launch `aux_regime_loss_reduce`: scalar mean CE + scalar mean accuracy.
    pub(crate) fn regime_loss_reduce(
        &self,
        stream: &Arc<CudaStream>,
        logits_ptr: u64,
        label_ptr: u64,
        b: usize,
        k: usize,
        loss_out_ptr: u64,
        correct_out_ptr: u64,
    ) -> Result<(), MLError> {
        let b_i32 = b as i32;
        let k_i32 = k as i32;
        // Two parallel partial-reduction strips: loss + accuracy. See
        // `aux_regime_loss_reduce` shmem layout in the kernel.
        let smem_bytes = 2 * AUX_BLOCK * std::mem::size_of::<f32>() as u32;
        unsafe {
            stream
                .launch_builder(&self.regime_loss_reduce_kernel)
                .arg(&logits_ptr)
                .arg(&label_ptr)
                .arg(&b_i32)
                .arg(&k_i32)
                .arg(&loss_out_ptr)
                .arg(&correct_out_ptr)
                .launch(LaunchConfig {
                    grid_dim: (1, 1, 1),
                    block_dim: (AUX_BLOCK, 1, 1),
                    shared_mem_bytes: smem_bytes,
                })
                .map_err(|e| MLError::ModelError(format!("aux_regime_loss_reduce: {e}")))?;
        }
        Ok(())
    }

    /// Launch `aux_regime_label_from_states`: discretise the regime score
    /// channel of the state vector into `[B] uint8` regime labels in `[0, K)`.
    ///
    /// `regime_score_idx` is a feature offset into the row-major `[B, F]`
    /// state buffer (typically `ml_core::state_layout::REGIME_STABILITY_INDEX`,
    /// resolved by the caller).
    #[allow(clippy::too_many_arguments)]
    pub(crate) fn regime_label_from_states(
        &self,
        stream: &Arc<CudaStream>,
        states_ptr: u64,
        b: usize,
        f: usize,
        regime_score_idx: usize,
        k: usize,
        label_out_ptr: u64,
    ) -> Result<(), MLError> {
        let b_i32 = b as i32;
        let f_i32 = f as i32;
        let r_i32 = regime_score_idx as i32;
        let k_i32 = k as i32;
        let blocks = ((b as u32) + AUX_BLOCK - 1) / AUX_BLOCK;
        unsafe {
            stream
                .launch_builder(&self.regime_label_kernel)
                .arg(&states_ptr)
                .arg(&b_i32)
                .arg(&f_i32)
                .arg(&r_i32)
                .arg(&k_i32)
                .arg(&label_out_ptr)
                .launch(LaunchConfig {
                    grid_dim: (blocks, 1, 1),
                    block_dim: (AUX_BLOCK, 1, 1),
                    shared_mem_bytes: 0,
                })
                .map_err(|e| MLError::ModelError(format!("aux_regime_label_from_states: {e}")))?;
        }
        Ok(())
    }

    /// Launch `aux_heads_loss_ema_update`: cold-path-cadence step
    /// observation for both scalar loss buffers into
    /// `producer_step_scratch_buf[scratch_idx_next_bar..+1)` /
    /// `[scratch_idx_regime]`. Single-thread, single-block — mirrors
    /// `h_s2_rms_ema_update`'s retrofit footprint.
    ///
    /// SP4 Layer A Task A13.1 retrofit (2026-05-01): kernel writes step
    /// observations only; the host launcher (`launch_aux_heads_loss_ema`
    /// on `GpuDqnTrainer`) syncs the stream and applies Pearls A+D
    /// host-side via `pearls_ad_update`.
    #[allow(clippy::too_many_arguments)]
    pub(crate) fn launch_loss_ema(
        &self,
        stream: &Arc<CudaStream>,
        next_bar_loss_scalar_ptr: u64,
        regime_loss_scalar_ptr: u64,
        scratch_dev_ptr: u64,
        scratch_idx_next_bar: i32,
        scratch_idx_regime: i32,
    ) -> Result<(), MLError> {
        unsafe {
            stream
                .launch_builder(&self.loss_ema_kernel)
                .arg(&next_bar_loss_scalar_ptr)
                .arg(&regime_loss_scalar_ptr)
                .arg(&scratch_dev_ptr)
                .arg(&scratch_idx_next_bar)
                .arg(&scratch_idx_regime)
                .launch(LaunchConfig {
                    grid_dim: (1, 1, 1),
                    block_dim: (1, 1, 1),
                    shared_mem_bytes: 0,
                })
                .map_err(|e| MLError::ModelError(format!("aux_heads_loss_ema_update: {e}")))?;
        }
        Ok(())
    }

}

/// Backward orchestrator — owns the 3 backward / final-reduce kernel
/// handles. Held separate from the forward ops because Commit B's wire-up
/// will dispatch forward + backward at distinct points in the training
/// step (forward runs alongside the main encoder forward; backward runs
/// after the main loss backward and SAXPY-accumulates `dh_s2` into the
/// trunk's main-path gradient).
#[allow(missing_debug_implementations)]
pub(crate) struct AuxHeadsBackwardOps {
    next_bar_backward_kernel: CudaFunction,
    regime_backward_kernel: CudaFunction,
    /// Final-reduce kernel for the per-sample param-grad partials emitted
    /// by both backward kernels. `aux_param_grad_reduce(partials [B, P], P,
    /// final [P])`.
    param_grad_reduce_kernel: CudaFunction,
}

impl AuxHeadsBackwardOps {
    /// Load the 3 backward kernel handles from the precompiled aux-heads
    /// cubin.
    pub(crate) fn new(stream: &Arc<CudaStream>) -> Result<Self, MLError> {
        let context = stream.context();
        let aux_module = context
            .load_cubin(AUX_HEADS_CUBIN.to_vec())
            .map_err(|e| MLError::ModelError(format!("aux_heads cubin load (bwd): {e}")))?;

        let next_bar_backward_kernel = aux_module
            .load_function("aux_next_bar_backward")
            .map_err(|e| MLError::ModelError(format!("aux_next_bar_backward load: {e}")))?;
        let regime_backward_kernel = aux_module
            .load_function("aux_regime_backward")
            .map_err(|e| MLError::ModelError(format!("aux_regime_backward load: {e}")))?;
        let param_grad_reduce_kernel = aux_module
            .load_function("aux_param_grad_reduce")
            .map_err(|e| MLError::ModelError(format!("aux_param_grad_reduce load: {e}")))?;

        Ok(Self {
            next_bar_backward_kernel,
            regime_backward_kernel,
            param_grad_reduce_kernel,
        })
    }

    /// Launch `aux_next_bar_backward` — emits per-sample param-grad
    /// partials + per-sample `dh_s2 [B, SH2]`. Caller follows up with
    /// `param_grad_reduce` for each of {`dW1`, `db1`, `dW2`, `db2`} to
    /// collapse along the batch dim.
    ///
    /// SP13 B1.1a (2026-05-05): MSE → softmax CE flip. Reads the saved
    /// `[B, K]` softmax tile (produced by `forward_next_bar`) plus the
    /// `[B] i32` labels buffer and the `[1]` `valid_count` scalar
    /// (written by `next_bar_loss_reduce`). Mirrors `aux_regime_backward`'s
    /// K-fanout structure but with mean-over-valid-rows derivative
    /// `d_logits[b, kc] = (softmax[b, kc] - one_hot(labels[b], kc)) / B_valid`
    /// for valid rows, zeroed across the K-vector for masked rows
    /// (label == -1) so partial-grad writes are no-ops.
    #[allow(clippy::too_many_arguments)]
    pub(crate) fn backward_next_bar(
        &self,
        stream: &Arc<CudaStream>,
        h_s2_ptr: u64,
        w1_ptr: u64,
        w2_ptr: u64,
        hidden_post_ptr: u64,
        softmax_ptr: u64,
        labels_ptr: u64,
        valid_count_ptr: u64,
        b: usize,
        sh2: usize,
        k: usize,
        dw1_partial_ptr: u64,
        db1_partial_ptr: u64,
        dw2_partial_ptr: u64,
        db2_partial_ptr: u64,
        dh_s2_out_ptr: u64,
    ) -> Result<(), MLError> {
        let b_i32 = b as i32;
        let sh2_i32 = sh2 as i32;
        let k_i32 = k as i32;
        // Shared mem: 2*H (d_h_pre + h_post caches) + K (d_logits cache).
        let smem_bytes =
            ((2 * AUX_HIDDEN_DIM as u32) + (k as u32)) * std::mem::size_of::<f32>() as u32;
        unsafe {
            stream
                .launch_builder(&self.next_bar_backward_kernel)
                .arg(&h_s2_ptr)
                .arg(&w1_ptr)
                .arg(&w2_ptr)
                .arg(&hidden_post_ptr)
                .arg(&softmax_ptr)
                .arg(&labels_ptr)
                .arg(&valid_count_ptr)
                .arg(&b_i32)
                .arg(&sh2_i32)
                .arg(&k_i32)
                .arg(&dw1_partial_ptr)
                .arg(&db1_partial_ptr)
                .arg(&dw2_partial_ptr)
                .arg(&db2_partial_ptr)
                .arg(&dh_s2_out_ptr)
                .launch(LaunchConfig {
                    grid_dim: (b as u32, 1, 1),
                    block_dim: (AUX_BLOCK, 1, 1),
                    shared_mem_bytes: smem_bytes,
                })
                .map_err(|e| MLError::ModelError(format!("aux_next_bar_backward: {e}")))?;
        }
        Ok(())
    }

    /// Launch `aux_regime_backward` — emits per-sample param-grad partials
    /// + per-sample `dh_s2 [B, SH2]`. Caller follows up with
    /// `param_grad_reduce` for each of {`dW1`, `db1`, `dW2`, `db2`}.
    #[allow(clippy::too_many_arguments)]
    pub(crate) fn backward_regime(
        &self,
        stream: &Arc<CudaStream>,
        h_s2_ptr: u64,
        w1_ptr: u64,
        w2_ptr: u64,
        hidden_post_ptr: u64,
        logits_ptr: u64,
        label_ptr: u64,
        b: usize,
        sh2: usize,
        k: usize,
        dw1_partial_ptr: u64,
        db1_partial_ptr: u64,
        dw2_partial_ptr: u64,
        db2_partial_ptr: u64,
        dh_s2_out_ptr: u64,
    ) -> Result<(), MLError> {
        let b_i32 = b as i32;
        let sh2_i32 = sh2 as i32;
        let k_i32 = k as i32;
        let smem_bytes =
            ((2 * AUX_HIDDEN_DIM as u32) + (k as u32)) * std::mem::size_of::<f32>() as u32;
        unsafe {
            stream
                .launch_builder(&self.regime_backward_kernel)
                .arg(&h_s2_ptr)
                .arg(&w1_ptr)
                .arg(&w2_ptr)
                .arg(&hidden_post_ptr)
                .arg(&logits_ptr)
                .arg(&label_ptr)
                .arg(&b_i32)
                .arg(&sh2_i32)
                .arg(&k_i32)
                .arg(&dw1_partial_ptr)
                .arg(&db1_partial_ptr)
                .arg(&dw2_partial_ptr)
                .arg(&db2_partial_ptr)
                .arg(&dh_s2_out_ptr)
                .launch(LaunchConfig {
                    grid_dim: (b as u32, 1, 1),
                    block_dim: (AUX_BLOCK, 1, 1),
                    shared_mem_bytes: smem_bytes,
                })
                .map_err(|e| MLError::ModelError(format!("aux_regime_backward: {e}")))?;
        }
        Ok(())
    }

    /// Launch `aux_param_grad_reduce`: collapse `partials [B, P]` along
    /// the batch dim into `final_out [P]`. One block per output element
    /// (`P` blocks); each block does a shmem-tree reduction over `B`
    /// samples. No atomicAdd.
    pub(crate) fn param_grad_reduce(
        &self,
        stream: &Arc<CudaStream>,
        partials_ptr: u64,
        b: usize,
        p: usize,
        final_out_ptr: u64,
    ) -> Result<(), MLError> {
        let b_i32 = b as i32;
        let p_i32 = p as i32;
        let smem_bytes = AUX_BLOCK * std::mem::size_of::<f32>() as u32;
        unsafe {
            stream
                .launch_builder(&self.param_grad_reduce_kernel)
                .arg(&partials_ptr)
                .arg(&b_i32)
                .arg(&p_i32)
                .arg(&final_out_ptr)
                .launch(LaunchConfig {
                    grid_dim: (p as u32, 1, 1),
                    block_dim: (AUX_BLOCK, 1, 1),
                    shared_mem_bytes: smem_bytes,
                })
                .map_err(|e| MLError::ModelError(format!("aux_param_grad_reduce: {e}")))?;
        }
        Ok(())
    }
}
