// crates/ml/src/cuda_pipeline/consume_lb_budget_kernel.cu
//
// SP7 (2026-05-03) — GPU dispatch for cql/c51 budget consumer.
//
// Background: the SP7 loss-balance controller writes the controller-derived
// budget per (head, branch) into ISV[BUDGET_{CQL,C51}_BASE+b] and the
// activation flag into ISV[LB_{CQL,C51}_ACTIVE_BASE+b]. The downstream
// consumer used to resolve `effective = (active >= 0.5) ? controller : bootstrap`
// host-side (Rust) inside `compute_adaptive_budgets`, then pass the resulting
// f32 scalars to `apply_c51_budget_scale` / `apply_cql_saxpy`. Because that
// host code was reached through the `aux_child` CUDA-Graph capture, the if/else
// resolved ONCE at capture time (step 0 of each fold, when the activation flag
// is the FoldReset sentinel `0.0`); the bootstrap branch baked the literal
// 0.02/0.05 into the captured kernel-arg buffers, freezing the budget for
// every replay regardless of what the controller wrote at runtime. Same
// disease as SP4 host-side EMA elimination 2026-05-01 (see
// `feedback_no_cpu_compute_strict.md`).
//
// Fix architecture: this kernel resolves the dispatch on-GPU each step. It
// runs inside the `aux_child` graph (kernel launch is graph-safe), reading
// ISV at replay time (current values), and writes the effective per-branch
// budget plus the trunk-mean and per-branch correction to a mapped-pinned
// scratch buffer. Downstream `dqn_scale_f32_dev_ptr_kernel` /
// `dqn_saxpy_f32_dev_ptr_kernel` (also in this translation unit) read the
// scalar alpha from device memory at execution time, so every replay sees
// fresh values — capture-safe end-to-end.
//
// Buffer layout (`lb_budget_effective_buf`, 10 f32):
//   [0]  cql_trunk_mean        — mean(cql_effective[0..4])
//   [1]  cql_correction[0]     — cql_effective[0] / cql_trunk_mean
//   [2]  cql_correction[1]     — cql_effective[1] / cql_trunk_mean
//   [3]  cql_correction[2]     — cql_effective[2] / cql_trunk_mean
//   [4]  cql_correction[3]     — cql_effective[3] / cql_trunk_mean
//   [5]  c51_trunk_mean        — mean(c51_effective[0..4])
//   [6]  c51_correction[0]     — c51_effective[0] / c51_trunk_mean
//   [7]  c51_correction[1]     — c51_effective[1] / c51_trunk_mean
//   [8]  c51_correction[2]     — c51_effective[2] / c51_trunk_mean
//   [9]  c51_correction[3]     — c51_effective[3] / c51_trunk_mean
//
// SAXPY callers (`apply_c51_budget_scale` / `apply_cql_saxpy`) issue a
// trunk-wide kernel reading `lb_buf[trunk_idx]`, then per-branch correction
// kernels reading `lb_buf[trunk_idx + 1 + b]`. Net effect: branch slice ends
// up scaled by `effective[b] = trunk × correction[b]`, identical to the prior
// host-side semantics, but every value derives from runtime ISV reads.
//
// Cold-start contract: bootstrap constants (CQL=0.02, C51=0.05) are passed as
// kernel args (not hardcoded literals) per `feedback_isv_for_adaptive_bounds`,
// keeping the values flowing through ISV/kernel-args. They MUST stay
// byte-identical to the host-side `CQL_BOOTSTRAP_BUDGET=0.02` /
// `C51_BOOTSTRAP_BUDGET=0.05` literals that lived in `compute_adaptive_budgets`
// (which themselves match `loss_balance_controller_kernel.cu`'s
// `COLD_START_FLOOR_CQL=0.02` / `COLD_START_FLOOR_C51=0.05`).
//
// Layout: single block, 8 threads (2 heads × 4 branches) for the dispatch
// kernel. No atomicAdd (per `feedback_no_atomicadd`). __threadfence_system()
// after writes so the host-visible mapped-pinned slice sees the updates after
// stream sync (HEALTH_DIAG path).

#include <cuda_runtime.h>

extern "C" __global__ void lb_budget_dispatch(
    // ISV signal bus (read-only).
    const float* __restrict__ isv_signals,
    int   active_cql_isv_base,    // LB_CQL_ACTIVE_BASE = 313
    int   active_c51_isv_base,    // LB_C51_ACTIVE_BASE = 317
    int   budget_cql_isv_base,    // BUDGET_CQL_BASE    = 198
    int   budget_c51_isv_base,    // BUDGET_C51_BASE    = 190
    // Cold-start bootstrap constants (passed as args, not hardcoded; matches
    // host-side CQL_BOOTSTRAP_BUDGET / C51_BOOTSTRAP_BUDGET in
    // compute_adaptive_budgets per the new dev-ptr architecture).
    float cql_bootstrap,
    float c51_bootstrap,
    // Output scratch buffer (10 f32). See top-of-file layout block.
    float* __restrict__ out
) {
    const float ACTIVE_THRESHOLD = 0.5f;
    const float SP7_EPS_DIV      = 1e-8f;
    const float TRUNK_EPS        = 1e-6f;  // matches host correction-skip threshold

    int tid = threadIdx.x;
    if (blockIdx.x != 0 || tid >= 8) return;   // 8 = 2 heads × 4 branches

    int head   = tid / 4;       // 0 = CQL, 1 = C51
    int branch = tid % 4;       // 0..3 (dir, mag, ord, urg)

    // ── Read activation flag + controller budget per (head, branch) ──
    int active_base = (head == 0) ? active_cql_isv_base : active_c51_isv_base;
    int budget_base = (head == 0) ? budget_cql_isv_base : budget_c51_isv_base;
    float bootstrap = (head == 0) ? cql_bootstrap : c51_bootstrap;

    float active = isv_signals[active_base + branch];
    float ctrl   = isv_signals[budget_base + branch];

    // ── Dispatch on activation flag ──────────────────────────────────
    float effective = (active >= ACTIVE_THRESHOLD) ? fmaxf(ctrl, SP7_EPS_DIV)
                                                   : bootstrap;

    // ── Stage in shared mem so thread 0 (per head) can compute trunk mean ─
    __shared__ float effective_smem[8];
    effective_smem[tid] = effective;
    __syncthreads();

    // Two threads (one per head) compute the trunk mean over their 4 branches
    // and write the trunk slot. The other six threads write only their
    // correction slot. This is a fixed 8-thread pattern — no atomics needed.
    int head_base_out = (head == 0) ? 0 : 5;     // [0..5)=cql, [5..10)=c51
    int trunk_idx_out = head_base_out + 0;        // slot 0 of each head block
    int corr_idx_out  = head_base_out + 1 + branch;

    if (branch == 0) {
        // Thread (head, 0) computes trunk mean.
        int sm_base = head * 4;
        float sum = effective_smem[sm_base + 0]
                  + effective_smem[sm_base + 1]
                  + effective_smem[sm_base + 2]
                  + effective_smem[sm_base + 3];
        float trunk_mean = sum * 0.25f;
        out[trunk_idx_out] = trunk_mean;
    }
    __syncthreads();

    // ── Compute and write per-branch correction ──────────────────────
    float trunk_mean = out[trunk_idx_out];
    // Mirror the host-side guard `if c51_trunk > 1e-6 { c51_branch[b] / c51_trunk } else { 1.0 }`.
    float correction = (trunk_mean > TRUNK_EPS) ? (effective / trunk_mean) : 1.0f;
    out[corr_idx_out] = correction;

    __threadfence_system();
}

/* ══════════════════════════════════════════════════════════════════════
 * F32 IN-PLACE SCALE KERNEL (alpha read from device memory)
 *
 * y[i] *= *alpha_ptr   for i = 0..n-1
 *
 * Mirror of `dqn_scale_f32_kernel` in `dqn_utility_kernels.cu`, but `alpha`
 * is a device pointer (single f32) read at execution time rather than a
 * scalar baked into the kernel-arg buffer at capture time. Required for
 * SP7 capture-safe per-branch budget consumption.
 *
 * Launch config: grid=(ceil(n/256), 1, 1), block=(256, 1, 1).
 * ══════════════════════════════════════════════════════════════════════ */
extern "C" __global__ void dqn_scale_f32_dev_ptr_kernel(
    float* y,
    const float* __restrict__ alpha_ptr,
    int n
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n) {
        float alpha = *alpha_ptr;
        // NaN-safe: matches dqn_scale_f32_kernel — IEEE 754 says 0*NaN=NaN;
        // when alpha=0 (cold-start / MSE warmup) the y buffer can contain
        // NaN, scaling by 0 must yield 0.
        y[i] = (alpha == 0.0f) ? 0.0f : y[i] * alpha;
    }
}

/* ══════════════════════════════════════════════════════════════════════
 * F32 SAXPY KERNEL (alpha read from device memory)
 *
 * y[i] += (*alpha_ptr) * x[i]   for i = 0..n-1
 *
 * Mirror of `dqn_saxpy_f32_kernel` in `dqn_utility_kernels.cu`, but `alpha`
 * is a device pointer (single f32) read at execution time. Used by
 * `apply_cql_saxpy` and per-branch variant in the SP7 capture-safe path.
 *
 * Launch config: grid=(ceil(n/256), 1, 1), block=(256, 1, 1).
 * ══════════════════════════════════════════════════════════════════════ */
extern "C" __global__ void dqn_saxpy_f32_dev_ptr_kernel(
    float* __restrict__ y,
    const float* __restrict__ x,
    const float* __restrict__ alpha_ptr,
    int n
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n) {
        float alpha = *alpha_ptr;
        y[i] = y[i] + alpha * x[i];
    }
}
