将C外积函数转为CUDA核：三维网格转二维数组索引求解

原C函数逻辑拆解 #

原函数本质是标准N×N矩阵乘法：C = A × B，其中每个元素C[r][c]是A的第r行与B的第c列的点积，计算公式为：
$$C[r][c] = \sum_{cc=0}^{N-1} A[r][cc] \times B[cc][c]$$
矩阵采用行优先一维存储，即元素(row, col)的一维索引为 row * N + col。

你的实现问题 #

1. 线程配置错误：原函数需要计算N²个C元素，每个元素需要N次乘加，你设置的3D线程网格会导致重复计算（比如N=3时，27个线程每个又做3次循环，共81次计算，远多于实际需要的27次乘加）。
2. 索引计算完全错误：你误用了3D存储索引，但原函数的矩阵都是2D的一维存储，混淆了维度映射关系，导致访问的矩阵元素完全不符合逻辑。

正确的索引计算思路 #

核心原则 #

• 明确矩阵的行优先存储规则：对于N×N矩阵，元素(row, col)的一维索引为 row * N + col。
• 根据线程负责的计算单元，将线程坐标映射到矩阵的行、列或求和维度：
  • 方式1（推荐，无原子操作，效率更高）：每个线程负责计算一个C[r][c]元素，内部完成点积循环。
  • 方式2（3D线程演示）：每个线程负责一次乘加操作，通过原子操作累加到C[r][c]。

方式1：每个线程负责一个C元素（推荐实现） #

线程配置示例 #

以N=3为例，可设置2D块和网格：

dim3 threads_per_block(3, 3); // 每个块3×3线程
dim3 blocks_per_grid(1, 1);    // 1×1个块，刚好覆盖3×3矩阵

对于任意N，推荐块大小设为(16,16)（CUDA硬件的warp大小为32，16×16的块能充分利用硬件资源），网格大小计算为：

dim3 threads_per_block(16, 16);
dim3 blocks_per_grid((N + 15) / 16, (N + 15) / 16);

核函数实现 #

__global__ void OuterProductKernel(float* A, float* B, float* C, int N)
{
    // 计算当前线程对应的C矩阵的行r和列c
    int r = blockIdx.y * blockDim.y + threadIdx.y;
    int c = blockIdx.x * blockDim.x + threadIdx.x;

    // 边界检查：超出矩阵范围则直接返回
    if (r >= N || c >= N)
        return;

    // 初始化求和变量
    float sum = 0.0f;
    // 计算A的第r行与B的第c列的点积
    for (int cc = 0; cc < N; cc++)
    {
        // A[r][cc]的一维索引：行r × 列数N + 列cc
        float a_val = A[r * N + cc];
        // B[cc][c]的一维索引：行cc × 列数N + 列c
        float b_val = B[cc * N + c];
        sum += a_val * b_val;
    }
    // 将结果写入C[r][c]
    C[r * N + c] = sum;
}

方式2：3D线程负责单次乘加（演示3D网格使用） #

如果要充分利用3D网格维度，可让每个线程对应(r, c, cc)三元组，负责一次乘加并原子累加到C[r][c]：

线程配置示例（N=3） #

dim3 threads_per_block(3, 3, 3);
dim3 blocks_per_grid(1, 1, 1);

核函数实现 #

__global__ void OuterProductKernel3D(float* A, float* B, float* C, int N)
{
    // 计算当前线程对应的r（C的行）、c（C的列）、cc（求和维度）
    int r = blockIdx.x * blockDim.x + threadIdx.x;
    int c = blockIdx.y * blockDim.y + threadIdx.y;
    int cc = blockIdx.z * blockDim.z + threadIdx.z;

    // 边界检查
    if (r < N && c < N && cc < N)
    {
        // 计算单次乘积
        float product = A[r * N + cc] * B[cc * N + c];
        // 原子累加：多个线程会写入同一个C[r][c]，必须用原子操作避免竞争
        atomicAdd(&C[r * N + c], product);
    }
}

注意事项 #

• 方式2需要使用atomicAdd，会带来一定性能开销，因此实际场景中优先选择方式1。
• 调用核函数前，需确保C数组已在设备端清零（可使用cudaMemset或在核函数中初始化）。

内容的提问来源于stack exchange，提问作者user366312