foxhunt/crates/ml-alpha/cuda/bce_loss_multi_horizon.cu

// bce_loss_multi_horizon.cu — Kendall σ-weighted multi-horizon BCE.
//
// Single source of truth for the multi-horizon BCE. Implements
// Kendall homoscedastic-uncertainty weighting per [Kendall+Gal+Cipolla
// 2018]: each horizon has a learnable `log_sigma_h` scalar that
// balances its loss contribution. Replaces the prior plain-pooled BCE.
//
// Math:
//
//   raw_bce_h  = Σ_{i in horizon h, m_i = 1} L_i        (unweighted sum)
//   count_h    = #{i in horizon h : m_i = 1}
//   mean_bce_h = raw_bce_h / count_h
//   w_h        = base_weight_h / (2 * exp(2 * log_sigma_h))
//   total_loss = Σ_h [ w_h * mean_bce_h + log_sigma_h ]
//
// Gradients:
//   d L / d p_i           = m_i * (w_h / count_h) * (p − y) / (p (1 − p))
//   d L / d log_sigma_h   = 1 − 2 * w_h * mean_bce_h
//
// Per `feedback_nvidia_grade_perf_for_kernels`:
//   - Warp-shuffle reduction (`__shfl_xor_sync`), NOT block tree-reduce.
//   - One `__syncthreads` for the cross-warp aggregate, total = O(1) barriers.
//   - Inactive lanes contribute 0 via ternary; never divergent shuffles.
//   - Hard-coded N_HORIZONS_BCE = 5 → per-warp 5-element accumulator stays in registers.
//
// Block layout: grid = (1), block = (256) = 8 warps. Single block over
// the full [n_pos, n_horizons] grid is fine — we never have > 256 ×
// (per-thread fan-in) positions in the trainer hot loop (max K × B =
// 64 × 8 = 512, well within stride-loop reach).

#define BCS_N_HORIZONS 5
#define BCS_BLOCK      256
#define BCS_N_WARPS    (BCS_BLOCK / 32)   // == 8

__device__ __forceinline__ float warp_reduce_sum(float v) {
    #pragma unroll
    for (int s = 16; s > 0; s >>= 1) {
        v += __shfl_xor_sync(0xffffffff, v, s);
    }
    return v;
}

__device__ __forceinline__ int warp_reduce_sum_int(int v) {
    #pragma unroll
    for (int s = 16; s > 0; s >>= 1) {
        v += __shfl_xor_sync(0xffffffff, v, s);
    }
    return v;
}

extern "C" __global__ void bce_multi_horizon_forward_backward(
    const float* __restrict__ probs,         // [n_pos * n_horizons]
    const float* __restrict__ labels,        // [n_pos * n_horizons]
    const float* __restrict__ base_weights,  // [n_horizons]  (per-horizon λ from ISV; constants OK)
    const float* __restrict__ log_sigma_h,   // [n_horizons]  (learnable Kendall σ logarithm)
    int n_pos,
    int n_horizons,
    float* __restrict__ loss_out,            // [1]                  — total Kendall loss
    float* __restrict__ mean_bce_per_h_out,  // [n_horizons]         — unweighted per-horizon mean (ISV signal)
    float* __restrict__ grad_probs,          // [n_pos * n_horizons] — ∂L/∂p
    int*   __restrict__ valid_count_out,     // [1]                  — Σ_i m_i (unweighted)
    float* __restrict__ d_log_sigma_h        // [n_horizons]         — ∂L/∂log_sigma_h
) {
    const int tid  = threadIdx.x;
    const int lane = tid & 31;
    const int warp = tid >> 5;
    const int total = n_pos * n_horizons;

    // Per-thread per-horizon partials in registers.
    float local_raw[BCS_N_HORIZONS];
    float local_cnt[BCS_N_HORIZONS];
    #pragma unroll
    for (int h = 0; h < BCS_N_HORIZONS; ++h) {
        local_raw[h] = 0.0f;
        local_cnt[h] = 0.0f;
    }
    int local_valid = 0;

    // Pass 1 — per-thread strided accumulation. Same loop also zeros
    // grad_probs at masked positions so the second pass can skip those.
    for (int i = tid; i < total; i += BCS_BLOCK) {
        const float y = labels[i];
        if (isnan(y)) {
            grad_probs[i] = 0.0f;
            continue;
        }
        const int h = i % n_horizons;
        const float p = fminf(fmaxf(probs[i], 1e-6f), 1.0f - 1e-6f);
        const float L = -(y * logf(p) + (1.0f - y) * logf(1.0f - p));
        local_raw[h] += L;
        local_cnt[h] += 1.0f;
        local_valid  += 1;
    }

    // Warp-shuffle reduce each of the 5 horizon sums + the per-horizon
    // count + the global valid count. Each lane in warp gets the full
    // warp-sum at the end (no need to broadcast manually).
    float warp_raw[BCS_N_HORIZONS];
    float warp_cnt[BCS_N_HORIZONS];
    #pragma unroll
    for (int h = 0; h < BCS_N_HORIZONS; ++h) {
        warp_raw[h] = warp_reduce_sum(local_raw[h]);
        warp_cnt[h] = warp_reduce_sum(local_cnt[h]);
    }
    int warp_valid = warp_reduce_sum_int(local_valid);

    // Cross-warp reduce — only lane 0 of each warp writes its partial.
    __shared__ float s_raw[BCS_N_WARPS * BCS_N_HORIZONS];
    __shared__ float s_cnt[BCS_N_WARPS * BCS_N_HORIZONS];
    __shared__ int   s_valid[BCS_N_WARPS];
    if (lane == 0) {
        #pragma unroll
        for (int h = 0; h < BCS_N_HORIZONS; ++h) {
            s_raw[warp * BCS_N_HORIZONS + h] = warp_raw[h];
            s_cnt[warp * BCS_N_HORIZONS + h] = warp_cnt[h];
        }
        s_valid[warp] = warp_valid;
    }
    __syncthreads();

    // Warp 0 finishes: reduce N_WARPS partials. All 32 lanes participate;
    // inactive lanes (tid ≥ N_WARPS) contribute 0 via ternary (NOT a
    // divergent shuffle).
    __shared__ float mean_bce_h_shared[BCS_N_HORIZONS];
    __shared__ float w_eff_shared[BCS_N_HORIZONS];
    __shared__ int valid_count_shared;

    if (warp == 0) {
        #pragma unroll
        for (int h = 0; h < BCS_N_HORIZONS; ++h) {
            float v_raw = (lane < BCS_N_WARPS) ? s_raw[lane * BCS_N_HORIZONS + h] : 0.0f;
            float v_cnt = (lane < BCS_N_WARPS) ? s_cnt[lane * BCS_N_HORIZONS + h] : 0.0f;
            v_raw = warp_reduce_sum(v_raw);
            v_cnt = warp_reduce_sum(v_cnt);
            if (lane == 0) {
                const float c_h = v_cnt;
                const float mean_bce = (c_h > 0.0f) ? v_raw / c_h : 0.0f;
                mean_bce_h_shared[h] = mean_bce;
                mean_bce_per_h_out[h] = mean_bce;

                const float ls = log_sigma_h[h];
                const float bw = (base_weights != nullptr) ? base_weights[h] : 1.0f;
                // w_h = bw / (2 * exp(2*ls)) = 0.5 * bw * exp(-2 * ls)
                const float w_h = 0.5f * bw * expf(-2.0f * ls);
                // Per-element grad prefactor = w_h / count_h.
                w_eff_shared[h] = (c_h > 0.0f) ? w_h / c_h : 0.0f;
                // ∂L/∂log_sigma_h = 1 − 2 * w_h * mean_bce_h.
                d_log_sigma_h[h] = 1.0f - 2.0f * w_h * mean_bce;
            }
        }
        int v_valid = (lane < BCS_N_WARPS) ? s_valid[lane] : 0;
        v_valid = warp_reduce_sum_int(v_valid);
        if (lane == 0) {
            valid_count_shared = v_valid;
            valid_count_out[0] = v_valid;
        }
    }
    __syncthreads();

    // Total loss — small serial reduce by thread 0 over 5 horizons.
    if (tid == 0) {
        float total_loss = 0.0f;
        #pragma unroll
        for (int h = 0; h < BCS_N_HORIZONS; ++h) {
            const float ls = log_sigma_h[h];
            const float bw = (base_weights != nullptr) ? base_weights[h] : 1.0f;
            const float w_h = 0.5f * bw * expf(-2.0f * ls);
            total_loss += w_h * mean_bce_h_shared[h] + ls;
        }
        loss_out[0] = total_loss;
    }
    __syncthreads();

    // Pass 2 — per-element grad. Coalesced strided access. w_eff_shared
    // is read-only at this point, broadcast freely across threads.
    for (int i = tid; i < total; i += BCS_BLOCK) {
        const float y = labels[i];
        if (isnan(y)) {
            continue;  // grad was already zeroed in pass 1
        }
        const int h = i % n_horizons;
        const float p = fminf(fmaxf(probs[i], 1e-6f), 1.0f - 1e-6f);
        grad_probs[i] = w_eff_shared[h] * (p - y) / (p * (1.0f - p));
    }
}