Matrix Addition

原始题目：LeetGPU - Matrix Addition

题目描述

编写一个 GPU 程序，对两个 $N \times N$ 的 32 位浮点数矩阵进行逐元素加法。程序应接收两个维度相同的输入矩阵，输出一个包含逐元素之和的矩阵：

$$C[i][j] = A[i][j] + B[i][j]$$

实现要求

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

示例

示例 1

Input:  A = [[1.0, 2.0], [3.0, 4.0]]
        B = [[5.0, 6.0], [7.0, 8.0]]
Output: C = [[6.0, 8.0], [10.0, 12.0]]

示例 2

Input:  A = [[1.5, 2.5, 3.5], [4.5, 5.5, 6.5], [7.5, 8.5, 9.5]]
        B = [[0.5, 0.5, 0.5], [0.5, 0.5, 0.5], [0.5, 0.5, 0.5]]
Output: C = [[2.0, 3.0, 4.0], [5.0, 6.0, 7.0], [8.0, 9.0, 10.0]]

约束条件

• 输入矩阵 $A$ 和 $B$ 维度相同。
• $1 \le N \le 4{,}096$。
• 所有元素均为 32 位浮点数。
• 性能测试在 $N = 4{,}096$ 的规模下进行。

解题思路

矩阵加法与向量加法类似，属于典型的逐元素映射（Element-wise Map）操作。可以将二维索引 $(i, j)$ 映射到一维线程索引，每个线程独立完成一个元素的加法。计算强度极低（每个元素 1 FLOP，3 次 4 字节内存访问），是典型的内存带宽受限内核。

代码实现

CUDA

#include <cuda_runtime.h>

// 解法一：基础逐元素映射
__global__ void matrix_add(const float* A, const float* B, float* C, int N2) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < N2) {
        C[i] = A[i] + B[i];
    }
}

// 解法二：Grid-Stride Loop（解耦数据规模与网格大小）
__global__ void matrix_add_grid_stride(const float* A, const float* B, float* C, int N2) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int stride = gridDim.x * blockDim.x;
    for (int i = idx; i < N2; i += stride) {
        C[i] = A[i] + B[i];
    }
}

// 解法三：float4 向量化加载（一次搬运 16 字节 = 4 个 float）
// 要求内存 16 字节对齐（cudaMalloc 默认满足）
__global__ void matrix_add_float4(const float* A, const float* B, float* C, int N2) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int stride = gridDim.x * blockDim.x;
    int N4 = N2 / 4;

    const float4* A4 = (const float4*)A;
    const float4* B4 = (const float4*)B;
    float4* C4 = (float4*)C;

    for (int i = idx; i < N4; i += stride) {
        float4 a = A4[i];
        float4 b = B4[i];
        float4 c;
        c.x = a.x + b.x;
        c.y = a.y + b.y;
        c.z = a.z + b.z;
        c.w = a.w + b.w;
        C4[i] = c;
    }

    // 处理尾部（不足 4 的余数）
    for (int i = N4 * 4 + idx; i < N2; i += stride) {
        C[i] = A[i] + B[i];
    }
}

extern "C" void solve(const float* A, const float* B, float* C, int N) {
    int N2 = N * N;
    int threadsPerBlock = 256;
    int blocksPerGrid = (N2 + threadsPerBlock - 1) / threadsPerBlock;
    matrix_add<<<blocksPerGrid, threadsPerBlock>>>(A, B, C, N2);
    cudaDeviceSynchronize();
}

Triton

import triton
import triton.language as tl
import torch


@triton.jit
def matrix_add_kernel(
    A_ptr, B_ptr, C_ptr,
    N: tl.constexpr,
    BLOCK_SIZE: tl.constexpr,
):
    idx = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    mask = idx < N * N
    a = tl.load(A_ptr + idx, mask=mask)
    b = tl.load(B_ptr + idx, mask=mask)
    tl.store(C_ptr + idx, a + b, mask=mask)


def solve(A: torch.Tensor, B: torch.Tensor, C: torch.Tensor):
    """Triton 解法：编译器自动处理向量化宽度和网格调度"""
    N = A.shape[0]
    grid = lambda meta: (triton.cdiv(N * N, meta['BLOCK_SIZE']),)
    matrix_add_kernel[grid](A, B, C, N, BLOCK_SIZE=1024)