Dot Product

原始题目：LeetGPU - Dot Product

题目描述

编写一个 GPU 程序，计算两个 32 位浮点数向量的点积。点积是两个向量对应元素的乘积之和：

$$A \cdot B = \sum_{i=0}^{n-1} A_i \cdot B_i = A_0 B_0 + A_1 B_1 + \dots + A_{n-1} B_{n-1}$$

实现要求

• 只允许使用 GPU 原生功能（不允许使用外部库）。
• solve 函数签名必须保持不变。
• 最终结果必须存储在 output 变量中。

示例

示例 1

Input:  A = [1.0, 2.0, 3.0, 4.0], B = [5.0, 6.0, 7.0, 8.0]
Output: 70.0 (1*5 + 2*6 + 3*7 + 4*8)

示例 2

Input:  A = [0.5, 1.5, 2.5], B = [2.0, 3.0, 4.0]
Output: 15.5 (0.5*2 + 1.5*3 + 2.5*4)

约束条件

• $A$ 和 $B$ 长度相同。
• $1 \le N \le 100{,}000{,}000$。
• 性能测试在 $N = 5$ 的规模下进行。

解题思路

点积 = 逐元素乘法 + 规约求和。第一步是完全并行的逐元素乘法，第二步使用并行规约。标准实现先做乘法得到乘积向量，再做分块规约 + warp shuffle + 全局 atomicAdd。当 $N$ 极大时，可以用多级规约树。

代码实现

CUDA

#include <cuda_runtime.h>
__global__ void dot_kernel(const float* A, const float* B, 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)
        sum += A[i] * B[i];
    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]);
}
extern "C" void solve(const float* A, const float* B, float* output, int N) {
    dot_kernel<<<min((N+255)/256, 1024), 256>>>(A, B, output, N);
    cudaDeviceSynchronize();
}

Triton

import triton, triton.language as tl
@triton.jit
def dot_kernel(A_ptr, B_ptr, output_ptr, N: tl.constexpr, BLOCK: tl.constexpr):
    idx = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)
    mask = idx < N
    a = tl.load(A_ptr + idx, mask=mask, other=0.0)
    b = tl.load(B_ptr + idx, mask=mask, other=0.0)
    tl.atomic_add(output_ptr, tl.sum(a * b, axis=0))