CUDA中的3D元素矩阵乘法？

Question

我有一个使用以下内核的二维矩阵乘法程序：

__global__ void multKernel(int *a, int *b, int *c, int N)
{
    int column  = threadIdx.x + blockDim.x * blockIdx.x;
    int row     = threadIdx.y + blockDim.y * blockIdx.y;

    int index = row * N + column;

    if(column < N && row < N)
    {
        c[index] = a[index] * b[index];
    }
}

现在，我想创建一个 3D 矩阵乘法内核，但我很难找到如何创建一个的示例（另外，我不擅长阅读数学公式，这是我需要改进的地方） .

我知道 GPU 示例将涉及使用

threadIdx.z

等等，但我对如何做有点迷茫。

谁能指出一些公式或示例代码的正确方向？甚至提供一个基本的例子？我认为我有一个应该工作的 CPU 示例。

void matrixMult3D(int *a, int *b, int *c, int *z, int N)
{
    int index;

    for(int column = 0; column < N; column++)
    {
        for(int row = 0; row < N; row++)
        {
            for (int z = 0; z < N; z++)
            {
                index = row * N + column + z;
                c[index] = a[index] * b[index] * z[index];
            }
        }
    }
}

我至少在正确的轨道上吗？

Answer 1

因为您实际上在做的只是一个元素方面的产品（我不敢称它为 Hadamard Product，因为它不是为超矩阵 AFAIK 定义的），您不需要做任何与最简单的不同的事情内核代码的 1D 版本。像这样的：

template<int ndim>
__global__ void multKernel(int *a, int *b, int *c, int *z, int N)
{
    int idx  = threadIdx.x + blockDim.x * blockIdx.x;
    int stride = blockDim.x * gridDim.x;

    int idxmax = 1;
    #pragma unroll
    for(int i=0; i < ndim; i++) {
        idxmax *= N;
    }
    for(; idx < idxmax; idx+=stride) {
       c[index] = a[index] * b[index] * z[index];
    }
}

[免责声明：在浏览器中编写的代码，从未编译或运行。使用风险自负]

适用于维度为 N (ndim=1)、N*N (ndim=2)、N*N*N (ndim=3) 等的任何维度的数组。