如何在 CUDA 中有效地计算矩阵的所有列和行的总和？

Question

我想计算 CUDA 中矩阵的所有列的总和和所有行的总和。一种方法是使用 BLAS 中的 SGEMV 子例程，将矩阵乘以 1 的向量。

然而，这会导致对矩阵进行两次扫描，假设它比 L1 缓存大得多：一次用于行，另一次用于列。另外，我打算进一步修改其他运算符的代码，所以这就是我编写自己的内核的原因。

到目前为止，我的方法是将矩阵分解为大小为 32 x 32 的子矩阵。每个线程块将这样一个子矩阵加载到共享内存中，计算子矩阵的行和列的总和，并将它们自动添加到适当的输出（下面的 row 和 col）。这样，矩阵数据只需要从VRAM中读取一次。

为简单起见，代码假设矩阵为n x n，n % 32 == 0，线程块为32 x 32

__global__ void sum_cols_and_rows(size_t n, const float* matrix, float* col, float* row)
{   
    __shared__ float sh[32][32];

    size_t x = blockDim.x * blockIdx.x + threadIdx.x;
    size_t y = blockDim.y * blockIdx.y + threadIdx.y;

    float sum = matrix[x + n * y];
    sh[threadIdx.x][threadIdx.y] = sum;

    for(unsigned w = 16; w >= 1; w /= 2)
        sum += __shfl_down(sum, w);
    const size_t laneID = threadIdx.x & 0x1f; // 32-1
    if(laneID == 0)
        atomicAdd(row + y, sum);
    __syncthreads();

    sum = sh[threadIdx.y][threadIdx.x]; // swapped indexes
    for(unsigned w = 16; w >= 1; w /= 2)
        sum += __shfl_down(sum, w);
    if(laneID == 0)
        atomicAdd(col + blockDim.x * blockIdx.x + threadIdx.y, sum);
}

// launch :
sum_cols_and_rows<<<dim3(n/32, n/32), dim3(32, 32), 32*32*sizeof(float)>>>(n, matrix, col, row);

然而，表现却相当令人失望。我在 GTX 980 上看到理论 224GB/s 内存带宽的大约 20%，即使在大型矩阵上也是如此，例如 16384x16384。

有什么方法可以使这种方法成为理论带宽限制吗？

Answer 1

在您的解决方案中，矩阵的每个 NxN 块都由单独的 NxN 线程块处理。实际上，每个单独的线程只做很少的工作，因此开销在实际计算中占主导地位。您可以通过让线程块处理多个矩阵块来改进它。

但是有一个更简单的解决方案，每个矩阵块只使用 N 个线程，其中一个线程对整个列求和。

实现与此类似：

__global__ void sum_cols_and_rows(size_t n, const float* matrix, float* col, float* row)
{   
    size_t laneID = threadIdx.x & 31;

    size_t x = blockDim.x * blockIdx.x + threadIdx.x;
    size_t y = N_ITERATIONS * blockIdx.y;

    size_t idx = y * n + x;

    float vertical = 0;

    for(int i = 0; i < N_ITERATIONS; i++) {
        float v = matrix[idx];
        vertical += v;
        for(unsigned w = 16; w >= 1; w /= 2)
            v += __shfl_down(v, w);
        if(laneID == 0)
            atomicAdd(&row[y], v);
        y++;
        idx += n;
    }

    atomicAdd(&col[x], vertical);
}

这里的可调参数是每个线程组的扭曲数和每个矩阵块中的行数 (N_ITERATIONS)。更大的值可能会以并行度为代价减少开销。

另一个可以尝试的想法是 vectorized loading - 以下之一：

float2 v2 = reinterpret_cast<float2*>(matrix)[idx];
float v = v2.x + v2.y;

float4 v4 = reinterpret_cast<float4*>(matrix)[idx];
float v = (v4.x + v4.y) + (v4.z + v4.w);

如何在 CUDA 中有效地计算矩阵的所有列和行的总和？

How to calculate the sum of all columns and rows of a matrix efficiently in CUDA?

cuda

gpgpu

nvidia