// multi_head_policy_gate_forward.cu — gating head from raw regime features.
//
// Spec: docs/superpowers/specs/2026-06-02-multi-head-policy-with-r-multiple.md
//       ADDENDUM 2026-06-03 §R.3 Option B (regime-direct gating, bypass VSN)
// Plan: docs/superpowers/plans/2026-06-03-multi-head-policy-implementation.md (Phase 2A-A)
//
// gate_logits[b, k] = Σ_j W_gate[k, j] × regime_h[b, j] + b_gate[k]
// gate_probs[b, k]  = softmax_k(gate_logits[b])
//
// Reads regime_h DIRECTLY (parallel channel — does NOT go through VSN/Mamba2).
// This is the load-bearing architectural choice motivated by the empirical
// regime-attenuation finding (pearl_local_smoke_noise_floor_and_regime_concentration):
// VSN's softmax-over-40 gating structurally limits any single encoder-input
// feature to ~1/40 of attention budget, suppressing vol-axis routing. The gate
// reads regime[0..6] without that bottleneck.
//
// Block layout:
//   grid  = (B, 1, 1)
//   block = (K, 1, 1)    — typically K=3, ≤ MAX_K_HEADS=8
//   One block per batch. K threads cooperate via shared-mem reductions.
//
// Per feedback_no_atomicadd: no atomic ops.
// Per feedback_no_nvrtc: pre-compiled cubin via build.rs.

#include <stdint.h>

#define REGIME_DIM      6
#define MAX_K_HEADS     8

extern "C" __global__ void multi_head_policy_gate_forward(
    const float* __restrict__ regime_h,       // [B × REGIME_DIM=6]
    const float* __restrict__ W_gate,         // [K × REGIME_DIM]
    const float* __restrict__ b_gate,         // [K]
    float*       __restrict__ gate_logits,    // [B × K] (output, for backward)
    float*       __restrict__ gate_probs,     // [B × K] (output, post-softmax)
    const int B,
    const int K
) {
    const int b = blockIdx.x;
    if (b >= B) return;

    const int tid = threadIdx.x;
    __shared__ float s_logits[MAX_K_HEADS];     // pre-softmax logits
    __shared__ float s_exp[MAX_K_HEADS];        // exp(logit - max) for softmax
    __shared__ float s_max;
    __shared__ float s_sum;

    // ── Step 1: linear projection ────────────────────────────────────────
    // gate_logits[b, k] = Σ_j W_gate[k, j] × regime_h[b, j] + b_gate[k]
    // One thread per head k.
    if (tid < K) {
        float gl = b_gate[tid];
        #pragma unroll
        for (int j = 0; j < REGIME_DIM; ++j) {
            gl += W_gate[tid * REGIME_DIM + j] * regime_h[b * REGIME_DIM + j];
        }
        s_logits[tid] = gl;
        // Persist pre-softmax logits to gmem for the Phase 2A-B backward kernel.
        gate_logits[b * K + tid] = gl;
    }
    __syncthreads();

    // ── Step 2: softmax (max-subtract for stability) ─────────────────────
    if (tid == 0) {
        float mx = s_logits[0];
        for (int k = 1; k < K; ++k) {
            mx = fmaxf(mx, s_logits[k]);
        }
        s_max = mx;
    }
    __syncthreads();

    if (tid < K) {
        s_exp[tid] = expf(s_logits[tid] - s_max);
    }
    __syncthreads();

    if (tid == 0) {
        float sm = 0.0f;
        for (int k = 0; k < K; ++k) {
            sm += s_exp[k];
        }
        s_sum = sm;
    }
    __syncthreads();

    if (tid < K) {
        gate_probs[b * K + tid] = s_exp[tid] / s_sum;
    }
}
