/* vsn_mask_ema — GPU-driven per-group mask step observation into
 * `producer_step_scratch_buf[scratch_first..scratch_first+num_groups)` for
 * Pearls A+D consumption. SP4 Layer A Task A13.3 retrofit (2026-05-01).
 *
 * Plan 4 Task 1B-i. Producer kernel that mirrors `h_s2_rms_ema_kernel`'s
 * single-block shmem-reduce pattern. Consumes the per-sample mask buffer
 * `mask [B, num_groups]` written by `vsn_softmax_and_gate_forward` and
 * writes `mean_b(mask[b, g])` to one scratch slot per feature group.
 *
 * Layout: row-major `[B, num_groups]`. `num_groups` is bounded above by
 * `SL_NUM_FEATURE_GROUPS = 6` (from `state_layout.cuh`) and is passed in
 * for graph-capture symmetry with the forward kernel; the inner loop
 * runs `num_groups` separate reductions sharing the same 256-float
 * scratchpad — total smem footprint is identical to
 * `h_s2_rms_ema_kernel.cu`'s 256 floats.
 *
 * GPU-drives-CPU-reads (spec §4.C.6), pearl_cold_path_no_exception_to
 * _gpu_drives.md, and feedback_no_cpu_compute_strict.md compliant. No
 * atomicAdd (per feedback_no_atomicadd.md); the in-block tree reduction is
 * standard `s >>= 1` with __syncthreads. The writeback path is fully GPU-
 * side via the chained `apply_pearls_ad_kernel` (2026-05-01 GPU-Pearls
 * refactor — same-stream sequential, no host sync).
 *
 * The consumer side (HEALTH_DIAG / VSN focus monitor) reads
 * ISV[VSN_MASK_GROUP_0_EMA_INDEX..+num_groups); only the EMA mechanism
 * behind the slots changes here. **α dropped per SP4 — Pearls A+D adapt
 * α from per-slot signal-vs-noise variance.**
 *
 * Cost: one launch per training step (cold-path cadence — same site as
 * `reward_component_ema` / `h_s2_rms_ema`). 256 threads × ceil(B/256)
 * strided loads per group × num_groups groups, plus 6 scratch writes by
 * thread 0. For typical B=512, num_groups=6, this is < 1 µs total.
 */

extern "C" __global__ void vsn_mask_ema_update(
    const float* __restrict__ mask,         /* [B, num_groups] saved by VSN forward */
    int          B,
    int          num_groups,                /* == SL_NUM_FEATURE_GROUPS = 6 */
    float*       __restrict__ scratch_buf,  /* producer_step_scratch_buf */
    int          scratch_first_index        /* slot 53 — first of num_groups scratch slots */
) {
    /* Single-block contract — match h_s2_rms_ema's grid guard. */
    if (blockIdx.x != 0) return;

    extern __shared__ float smem[];
    const int tid   = (int)threadIdx.x;
    const int block = (int)blockDim.x;

    /* Run num_groups separate shmem-tree reductions. Reusing the same
     * scratchpad keeps shmem at 256 floats regardless of num_groups; the
     * cost is `num_groups` __syncthreads passes, which at num_groups=6
     * is well under any kernel-launch overhead we'd save by fusing. */
    for (int g = 0; g < num_groups; ++g) {

        /* Per-thread sum over a strided slice of the column for group g. */
        float local_sum = 0.0f;
        for (int i = tid; i < B; i += block) {
            local_sum += mask[(size_t)i * num_groups + g];
        }
        smem[tid] = local_sum;
        __syncthreads();

        /* In-block tree reduction (no atomicAdd). */
        for (int s = block / 2; s > 0; s >>= 1) {
            if (tid < s) smem[tid] += smem[tid + s];
            __syncthreads();
        }

        /* Thread 0 writes the step observation. The host caller's
         * invariant requires `B > 0`, so the divide is safe; the NaN
         * from divide-by-zero would propagate loudly through Pearls A+D. */
        if (tid == 0) {
            const float mean = smem[0] / (float)B;
            scratch_buf[scratch_first_index + g] = mean;
        }
        __syncthreads();   /* fence before reusing smem for the next group */
    }
    /* Single fence after all groups have written — PCIe-visible to
     * mapped-pinned host_ptr. */
    if (tid == 0) {
        __threadfence_system();
    }
}
