Categorical Cross Entropy Loss

原始题目：LeetGPU - Categorical Cross Entropy Loss

题目描述

编写一个 GPU 程序，计算一批预测结果的分类交叉熵损失。给定预测 logits 矩阵 $Z$（$N \times C$）和真实类别标签向量 true_labels（长度为 $N$），计算批量平均交叉熵损失。

单个样本 $j$ 的损失（logits 为 $z_j = [z_{j1}, \dots, z_{jC}]$，真实标签为 $y_j$）使用数值稳定公式计算：

$$Loss_j = \log\!\left(\sum_{k=1}^{C} e^{z_{jk}}\right) - z_{j, y_j}$$

最终输出存储在 loss 变量中，为 $N$ 个样本的平均损失：

$$L = \frac{1}{N}\sum_{j=1}^{N} Loss_j$$

实现要求

• 不允许使用外部库。
• solve 函数签名必须保持不变。
• 最终结果（平均损失）必须存储在 loss 中。

示例

示例 1

Input:  N = 2, C = 3
        logits = [[1.0, 2.0, 0.5], [0.1, 3.0, 1.5]]
        true_labels = [1, 1]
Output: loss = 0.3548926

示例 2

Input:  N = 3, C = 4
        logits = [[-0.5, 1.5, 0.0, 1.0], [2.0, -1.0, 0.5, 0.5], [0.0, 0.0, 0.0, 0.0]]
        true_labels = [3, 0, 1]
Output: loss = 0.98820376

约束条件

• $1 \le N \le 10{,}000$，$2 \le C \le 1{,}000$。
• $-10.0 \le logits[i, j] \le 10.0$。
• $0 \le true\_labels[i] \le C$。
• 性能测试在 $N = 10{,}000$ 的规模下进行。

解题思路

交叉熵损失 = log-softmax + 负对数似然。与 Softmax 类似，需要先找每行的最大值（"max trick"），再计算 log-sum-exp 减去真实标签对应的 logit。核心是每行独立计算，可以用一个线程块处理一行或一个 warp 处理一行。当 $C$ 较大时，warp-level reduction 非常高效。

代码实现

CUDA

#include <cuda_runtime.h>
#include <math.h>
__global__ void ce_kernel(const float* logits, const int* labels, float* loss, int N, int C) {
    __shared__ float sdata[256];
    int tid = threadIdx.x; float sum=0.0f;
    for (int n=blockIdx.x*blockDim.x+tid; n<N; n+=gridDim.x*blockDim.x) {
        float mx=logits[n*C]; for(int c=1;c<C;c++)mx=fmaxf(mx,logits[n*C+c]);
        float lse=0.0f; for(int c=0;c<C;c++)lse+=expf(logits[n*C+c]-mx);
        lse=logf(lse)+mx; sum += lse - logits[n*C+labels[n]];
    }
    sdata[tid]=sum; __syncthreads();
    for(int s=blockDim.x/2;s>0;s>>=1){if(tid<s)sdata[tid]+=sdata[tid+s];__syncthreads();}
    if(tid==0)atomicAdd(loss, sdata[0]/N);
}
extern "C" void solve(const float* logits, const int* labels, float* loss, int N, int C) {
    ce_kernel<<<min(N,1024),256>>>(logits, labels, loss, N, C);
    cudaDeviceSynchronize();
}

Triton

import triton, triton.language as tl
@triton.jit
def ce_kernel(logits_ptr, labels_ptr, loss_ptr, N: tl.constexpr, C: tl.constexpr, BLOCK: tl.constexpr):
    nid = tl.program_id(0)*BLOCK + tl.arange(0, BLOCK)
    mask = nid < N
    cid = tl.arange(0, C)
    logits = tl.load(logits_ptr + nid[:,None]*C + cid[None,:], mask=mask[:,None], other=-float('inf'))
    mx = tl.max(logits, axis=1); lse = tl.log(tl.sum(tl.exp(logits-mx[:,None]), axis=1)) + mx
    label = tl.load(labels_ptr + nid, mask=mask)
    loss_i = lse - tl.load(logits_ptr + nid*C + label, mask=mask)
    tl.atomic_add(loss_ptr, tl.sum(loss_i, axis=0) / N)