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

Question

我有一個使用以下內核的2D矩陣乘法程序：

__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矩陣乘法內核，但是我很難找到創建一個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產品，因為它沒有為超矩陣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）等的任何尺寸的數組。