将 3D 网格转换为 2D 数组索引

Question

假设我想将以下 C 例程翻译成 CUDA 内核。

而且，我想使用网格中的所有维度来运行内核。

如何计算矩阵的行和列的索引？

void OuterProduct(float* A, float* B, float** C, int N)
{
    for(int r=0 ; r<N ; r++)
    {
        for(int c=0 ; c<N ; c++)
        {
            for(int cc=0 ; cc<N ; cc++)
            {
                (*C)[r * N + c] += A[r * N + cc] * B[cc * N + c];
            }
        }
    }
}

以下是我的理解：

假设上面的例程是要将两个 3x3 矩阵相乘。因此，计算次数为 3x3x3 = 27。因此，我们需要 27 个线程来完成乘法运算。

假设我们将在每个块中使用一个线程。所以，我们需要 27 个区块。

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

typedef float I;  
__global__ void OuterProductKernel(I* A, I* B, I* C, int N)
{
    int dimx = N;
    int dimy = N;
    int dimz = N;

    int r = blockIdx.x * blockDim.x + threadIdx.x;
    int c = blockIdx.y * blockDim.y + threadIdx.y;
    int d = blockIdx.z * blockDim.z + threadIdx.z;

    if (r < N && c < N && d < N) 
    {
        int loc_c = d * dimx * dimy + c * dimx + r;
 
        for (int cc=0; cc<N; cc++) 
        {
            int loc_a = (cc * dimx * dimy) + (c * dimx) + r;
            int loc_b = (d * dimx * dimy) + (cc * dimx) + r;
                    C[loc_c] += A[loc_a]*B[loc_b];
        }
    }
}

这是正确的吗？我认为不是。

你能告诉我计算

loc_a

、

loc_b

和

loc_c

的正确理由吗？

Answer 1

CUDA（作为 C 和 C++）使用 行优先顺序，所以代码像

int loc_c = d * dimx * dimy + c * dimx + r;

应该改写为

int loc_c = d * dimx * dimy + r * dimx + c;

与其他“位置”相同：loc_a 和 loc_b。

还有：

1. 确保 C 数组归零，你永远不会在代码中这样做
2. 如果能看到调用代码就好了

但是如果你的目标是在大多数情况下的性能，你将从缓存、内存对齐和别名中获得更多。

参见，例如：CUDA 矩阵-矩阵乘法

Answer 2

这里在很多层面上都有很多混淆——数组索引、CUDA 执行模型、数学运算本身。

从基础开始：矩阵乘法或两个矩阵 A 和 B 之间的点积的元素运算基本上是

C[x,y] = dot(A[x,:], B[:,y]) for all [x,y] in [0...N,0...N]

在哪里

dot(A[x,:], B[:,y]) = A[x,0]*B[0,y] + A[x,1]*B[1,y] + ......

（你称之为“外积”，但这是完全不同的 Hadamard 积）

如果您想使用三维网格，那么第三维需要沿着点积的“内部”进行维数。假设我们在两个数组中都有行主要排序，让我们定义一个小辅助函数：

__device__ size_t rowmajoridx(size_t row, size_t col, size_t lda)
{
    return col * lda + row;
}

然后内核就变成了：

template<typename T>  
__global__ void DotProductKernel(T* A, T* B, T* C, int N)
{
    int dimx = N;
    int dimy = N;
    int dimz = N;

    int r = blockIdx.x * blockDim.x + threadIdx.x;
    int c = blockIdx.y * blockDim.y + threadIdx.y;
    int d = blockIdx.z * blockDim.z + threadIdx.z;

    if (r < N && c < N && d < N) 
    {
        C[rowmajidx(r,c,N)] += A[rowmajoridx(r,d,N)] * B[rowmajoridx(d,c,N)); 
    }
}

稍微抽象一下，可以清楚地看到三个线程索引中只有两个曾经用于索引任何数组 r 和 c 用于 C，r 和 d 用于 A，d 和 c 用于 B。它们是，毕竟只有 N x N 大小。

但是请注意，这个内核实际上并不能正常工作。因为有很多线程为

C

中的每个输出条目做出贡献，所以你有很多方式的内存竞争。

C

需要在内核运行之前归零。要修复内存竞争，您需要使用原子内存事务，它比标准内存写入慢很多数量级，并且并非所有硬件上的每种类型都支持。在那种情况下，内核变成这样的：

template<typename T>  
__global__ void DotProductKernel(T* A, T* B, T* C, int N)
{
    int dimx = N;
    int dimy = N;
    int dimz = N;

    int r = blockIdx.x * blockDim.x + threadIdx.x;
    int c = blockIdx.y * blockDim.y + threadIdx.y;
    int d = blockIdx.z * blockDim.z + threadIdx.z;

    if (r < N && c < N && d < N) 
    {
        T ans = A[rowmajoridx(r,d,N)] * B[rowmajoridx(d,c,N));
        atomicAdd(C + rowmajidx(r,c,N), ans);
    }  
}

[注意所有这些代码都是在浏览器中编写的，从未编译从未测试过，使用风险自负]