/* MoE mixture / load-balance / ISV-producer kernels for the Mixture-of-
 * Experts redesign. Spec:
 * docs/superpowers/specs/2026-04-27-moe-regime-redesign-design.md §6.2.
 *
 * GPU-resident, capture-friendly, no atomicAdd
 * (per `feedback_no_atomicadd.md`).
 */

extern "C" __global__ void moe_mixture_forward(
    const float* __restrict__ expert_outputs,  /* [K, B, C] */
    const float* __restrict__ gate,            /* [B, K] */
    float*       __restrict__ h_s2,            /* [B, C] */
    int B,
    int K,
    int C
) {
    /* One thread per (b, c) pair. */
    int total = B * C;
    int tid = blockIdx.x * blockDim.x + threadIdx.x;
    if (tid >= total) return;

    int b = tid / C;
    int c = tid % C;

    float acc = 0.0f;
    for (int k = 0; k < K; ++k) {
        float g  = gate[b * K + k];
        float ek = expert_outputs[k * (B * C) + b * C + c];
        acc += g * ek;
    }
    h_s2[tid] = acc;
    __threadfence_system();  /* make write PCIe-visible to mapped pinned host_ptr */
}

extern "C" __global__ void moe_mixture_backward(
    const float* __restrict__ dh_s2,           /* [B, C] */
    const float* __restrict__ gate,            /* [B, K] */
    float*       __restrict__ de_k,            /* [K, B, C] */
    int B,
    int K,
    int C
) {
    /* One thread per (k, b, c). de_k[k,b,c] = g[b,k] * dh_s2[b,c]. */
    int total = K * B * C;
    int tid = blockIdx.x * blockDim.x + threadIdx.x;
    if (tid >= total) return;

    int k = tid / (B * C);
    int rem = tid - k * (B * C);
    int b = rem / C;
    int c = rem - b * C;

    de_k[tid] = gate[b * K + k] * dh_s2[b * C + c];
    __threadfence_system();
}

extern "C" __global__ void moe_dgate_reduce(
    const float* __restrict__ dh_s2,           /* [B, C] */
    const float* __restrict__ expert_outputs,  /* [K, B, C] */
    float*       __restrict__ dg,              /* [B, K] */
    int B,
    int K,
    int C
) {
    /* Grid: (B, K, 1); block: 256 threads (or chosen by host). */
    int b = blockIdx.x;
    int k = blockIdx.y;
    if (b >= B || k >= K) return;

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

    float local = 0.0f;
    for (int c = tid; c < C; c += block) {
        local += dh_s2[b * C + c] * expert_outputs[k * (B * C) + b * C + c];
    }
    smem[tid] = local;
    __syncthreads();

    for (int s = block / 2; s > 0; s >>= 1) {
        if (tid < s) smem[tid] += smem[tid + s];
        __syncthreads();
    }

    if (tid == 0) {
        dg[b * K + k] = smem[0];
        __threadfence_system();
    }
}

extern "C" __global__ void moe_load_balance_loss(
    const float* __restrict__ gate,            /* [B, K] */
    float*       __restrict__ loss_per_k,      /* [K] */
    const float* __restrict__ isv_signals,     /* [ISV_TOTAL_DIM] */
    int B,
    int K,
    int isv_lambda_eff_idx                     /* MOE_LAMBDA_EFF_INDEX = 128 */
) {
    /* One block per k. Each block reduces over B (sum gate[b,k]) into smem
     * and writes lambda*K*(sum/B)² to loss_per_k[k].
     * 1-block-per-k design avoids atomicAdd.
     *
     * λ is read from ISV[isv_lambda_eff_idx] — the adaptive controller
     * (moe_lambda_eff_update kernel) updates this slot each step from the
     * gate-entropy EMA so the load-balance push scales inversely with
     * gate concentration. Cold-start is λ_floor (constructor write); fold
     * reset re-applies λ_floor; controller drives it up to
     * λ_floor + λ_max_extra under full entropy collapse.
     *
     * Per `feedback_isv_for_adaptive_bounds.md` / `pearl_cold_path_no_exception_to_gpu_drives.md`:
     * the consumer reads the ISV slot directly — no host roundtrip. */
    extern __shared__ float smem[];
    int k = blockIdx.x;
    if (k >= K) return;
    int tid = threadIdx.x;
    int block = blockDim.x;

    float lambda = isv_signals[isv_lambda_eff_idx];

    float local = 0.0f;
    for (int b = tid; b < B; b += block) {
        local += gate[b * K + k];
    }
    smem[tid] = local;
    __syncthreads();

    for (int s = block / 2; s > 0; s >>= 1) {
        if (tid < s) smem[tid] += smem[tid + s];
        __syncthreads();
    }

    if (tid == 0) {
        float col_mean = smem[0] / (float)B;
        loss_per_k[k] = lambda * (float)K * col_mean * col_mean;
        __threadfence_system();
    }
}

extern "C" __global__ void moe_load_balance_reduce(
    const float* __restrict__ loss_per_k,
    float*       __restrict__ loss_total,
    int K
) {
    /* Single-thread reduce; K is small (=8). */
    if (threadIdx.x != 0 || blockIdx.x != 0) return;
    float acc = 0.0f;
    for (int k = 0; k < K; ++k) acc += loss_per_k[k];
    loss_total[0] = acc;
    __threadfence_system();
}

/* Row-wise softmax: out[b,k] = exp(in[b,k]) / sum_j exp(in[b,j]).
 * Grid: one block per row (b). Single-thread per block — K is small (=8).
 * No atomicAdd, no shared memory needed for K<=8 (all in registers). */
extern "C" __global__ void moe_row_softmax(
    const float* __restrict__ in_pre,   /* [B, K] pre-softmax logits */
    float*       __restrict__ out_soft, /* [B, K] softmax output */
    int B,
    int K
) {
    int b = blockIdx.x * blockDim.x + threadIdx.x;
    if (b >= B) return;

    /* Numerically stable: subtract row max before exponentiation. */
    float row_max = in_pre[b * K];
    for (int k = 1; k < K; ++k) {
        float v = in_pre[b * K + k];
        if (v > row_max) row_max = v;
    }
    float sum_exp = 0.0f;
    for (int k = 0; k < K; ++k) {
        float e = expf(in_pre[b * K + k] - row_max);
        out_soft[b * K + k] = e;
        sum_exp += e;
    }
    float inv_sum = 1.0f / sum_exp;
    for (int k = 0; k < K; ++k) {
        out_soft[b * K + k] *= inv_sum;
    }
    __threadfence_system();
}

/* Softmax backward per row:
 *   dg_pre[b,k] = g_soft[b,k] * (dg[b,k] - sum_j g_soft[b,j] * dg[b,j])
 * Grid: one block per row (b). Single-thread per block — K is small (=8). */
extern "C" __global__ void moe_softmax_backward(
    const float* __restrict__ g_soft,   /* [B, K] saved softmax output from forward */
    const float* __restrict__ dg,       /* [B, K] upstream gradient w.r.t. gate */
    float*       __restrict__ dg_pre,   /* [B, K] gradient w.r.t. pre-softmax logits */
    int B,
    int K
) {
    int b = blockIdx.x * blockDim.x + threadIdx.x;
    if (b >= B) return;

    /* dot = sum_j g_soft[b,j] * dg[b,j] */
    float dot = 0.0f;
    for (int k = 0; k < K; ++k) {
        dot += g_soft[b * K + k] * dg[b * K + k];
    }
    for (int k = 0; k < K; ++k) {
        dg_pre[b * K + k] = g_soft[b * K + k] * (dg[b * K + k] - dot);
    }
    __threadfence_system();
}

extern "C" __global__ void moe_expert_util_ema_update(
    const float* __restrict__ gate,            /* [B, K] */
    float*       __restrict__ scratch_buf,     /* producer_step_scratch_buf */
    int B,
    int K,
    int scratch_idx_util_base,                 /* slot 44 — first of K util scratch slots */
    int scratch_idx_entropy                    /* slot 52 — gate-entropy scratch slot */
) {
    /* SP4 Layer A Task A13.2 retrofit (2026-05-01); GPU-Pearls refactor
     * (2026-05-01): write step observations only; the host launcher
     * follows up with the GPU `apply_pearls_ad_kernel` on the same stream
     * (single launch, n_slots=K+1) which writes ISV[118..127). Writes one
     * scratch slot per expert util (`scratch_idx_util_base..+K`) plus one
     * for gate entropy (`scratch_idx_entropy`). **α dropped per SP4 —
     * Pearls A+D adapt α from per-slot signal-vs-noise variance.**
     * Single-block, single-thread cold-path cadence (matches retrofit
     * family shape). */
    if (blockIdx.x != 0 || threadIdx.x != 0) return;

    /* Single forward pass: compute col_mean[k] for each expert, write to
     * scratch util slot, accumulate Shannon entropy. Avoids the 2× pass
     * over `gate` the prior implementation did (one for util, one for
     * entropy reread). */
    float entropy = 0.0f;
    for (int k = 0; k < K; ++k) {
        float sum = 0.0f;
        for (int b = 0; b < B; ++b) sum += gate[b * K + k];
        const float col_mean = sum / (float)B;
        scratch_buf[scratch_idx_util_base + k] = col_mean;
        if (col_mean > 1e-9f) entropy += -col_mean * logf(col_mean);
    }
    scratch_buf[scratch_idx_entropy] = entropy;
    __threadfence_system();
}
