Mean Squared Error

原始题目：LeetGPU - Mean Squared Error

题目描述

编写一个 GPU 程序，计算预测值与目标值之间的均方误差（MSE）。给定两个等长的数组 predictions 和 targets，计算：

$$MSE = \frac{1}{N} \sum_{i=1}^{N} (predictions_i - targets_i)^2$$

其中 $N$ 为每个数组的元素数量。

实现要求

• 不允许使用外部库。
• solve 函数签名必须保持不变。
• 最终结果必须存储在 mse 变量中。

示例

示例 1

Input:  predictions = [1.0, 2.0, 3.0, 4.0]
        targets     = [1.5, 2.5, 3.5, 4.5]
Output: mse = 0.25

示例 2

Input:  predictions = [10.0, 20.0, 30.0]
        targets     = [12.0, 18.0, 33.0]
Output: mse = 5.67

约束条件

• $1 \le N \le 100{,}000{,}000$。
• $-1000.0 \le predictions[i], targets[i] \le 1000.0$。
• 性能测试在 $N = 50{,}000{,}000$ 的规模下进行。

解题思路

MSE 计算分三步：逐元素差平方（map）、求和（reduce）、除以 $N$（标量运算）。前两步都在 GPU 上，每个线程先计算 $(pred_i - targ_i)^2$，再用分块规约 + warp shuffle 求和。最后一步除以 $N$ 只需一个线程执行。与纯规约相比，多了一步逐元素运算，但整体仍然是内存带宽受限的。

代码实现

CUDA

#include <cuda_runtime.h>
__global__ void mse_kernel(const float* pred, const float* target, float* output, int N) {
    __shared__ float sdata[256];
    int tid = threadIdx.x; float sum = 0.0f;
    for (int i=blockIdx.x*blockDim.x+tid; i<N; i+=gridDim.x*blockDim.x) {
        float d = pred[i]-target[i]; sum += d*d;
    }
    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(output, sdata[0]/N);
}
extern "C" void solve(const float* pred, const float* target, float* output, int N) {
    mse_kernel<<<min((N+255)/256,1024),256>>>(pred,target,output,N);
    cudaDeviceSynchronize();
}

Triton

import triton, triton.language as tl
@triton.jit
def mse_kernel(pred_ptr, target_ptr, output_ptr, N: tl.constexpr, BLOCK: tl.constexpr):
    idx = tl.program_id(0)*BLOCK + tl.arange(0, BLOCK)
    mask = idx < N
    p = tl.load(pred_ptr + idx, mask=mask); t = tl.load(target_ptr + idx, mask=mask)
    tl.atomic_add(output_ptr, tl.sum((p-t)*(p-t), axis=0) / N)