具有共享内存的 cuda 平铺 3d 卷积实现

Question

根据我的研究，使用 cuda 实现平铺版卷积有两种不同的策略。我想更多地了解这方面，想看看他们之间的比较，每种策略的优缺点是什么，以及如何选择。下面是两种不同策略的实现。

策略一：瓦片大小与输出大小相匹配，需要多步加载输入。

#define MASK_WIDTH 3
#define MASK_RADIUS 1

#define TILE_WIDTH 8

#define SHAREDMEM_DIM (TILE_WIDTH + (MASK_RADIUS * 2))

__constant__ float deviceMask[MASK_WIDTH * MASK_WIDTH * MASK_WIDTH];

__global__ void conv3d(float *inputArray, 
                   float *outputArray, 
                   const int z_size,
                   const int y_size, 
                   const int x_size) {
    __shared__ float subTile[SHAREDMEM_DIM][SHAREDMEM_DIM][SHAREDMEM_DIM];

    int bx = blockIdx.x, tx = threadIdx.x;
    int by = blockIdx.y, ty = threadIdx.y;
    int bz = blockIdx.z, tz = threadIdx.z;

    int destination = (tz * TILE_WIDTH * TILE_WIDTH) + (ty * TILE_WIDTH) + tx;
    int destTmp = destination;
    int dX = destTmp % SHAREDMEM_DIM;
    destTmp = destTmp / SHAREDMEM_DIM;
    int dY = destTmp % SHAREDMEM_DIM;
    destTmp = destTmp / SHAREDMEM_DIM;
    int dZ = destTmp;

    int inputZ = dZ + (bz * TILE_WIDTH) - MASK_RADIUS;
    int inputY = dY + (by * TILE_WIDTH) - MASK_RADIUS;
    int inputX = dX + (bx * TILE_WIDTH) - MASK_RADIUS;
    int input = (inputZ * y_size * x_size) + (inputY * x_size) + inputX;

    if(   inputZ >= 0 && inputZ < z_size 
       && inputY >= 0 && inputY < y_size 
       && inputX >= 0 && inputX < x_size){
           subTile[dZ][dY][dX] = inputArray[input];
    }
    else{
        subTile[dZ][dY][dX] = 0;
    }

    destination = TILE_WIDTH * TILE_WIDTH * TILE_WIDTH 
            + (tz * TILE_WIDTH * TILE_WIDTH) + (ty * TILE_WIDTH) + tx;
    destTmp = destination;
    dX = destTmp % SHAREDMEM_DIM;
    destTmp = destTmp / SHAREDMEM_DIM;
    dY = destTmp % SHAREDMEM_DIM;
    destTmp = destTmp / SHAREDMEM_DIM;
    dZ = destTmp;

    inputZ = dZ + (bz * TILE_WIDTH) - MASK_RADIUS;
    inputY = dY + (by * TILE_WIDTH) - MASK_RADIUS;
    inputX = dX + (bx * TILE_WIDTH) - MASK_RADIUS;
    input = (inputZ * y_size * x_size) + (inputY * x_size) + inputX;

    if(dZ < SHAREDMEM_DIM){
        if(   inputZ >= 0 && inputZ < z_size 
           && inputY >= 0 && inputY < y_size 
           && inputX >= 0 && inputX < x_size ) {
                subTile[dZ][dY][dX] = inputArray[input];
           }
        else{
            subTile[dZ][dY][dX] = 0;
        }
    }

    __syncthreads();  

    float sum = 0;
    int z, y, x;
    for(z = 0; z < MASK_WIDTH; z++){
        for(y = 0; y < MASK_WIDTH; y++){
            for(x = 0; x < MASK_WIDTH; x++){
                sum += subTile[tz + z][ty + y][tx + x] 
                   * deviceMask[x + (y * MASK_WIDTH) + (z * MASK_WIDTH * MASK_WIDTH)];
            }
        }
    }
    z = tz + (bz * TILE_WIDTH);
    y = ty + (by * TILE_WIDTH);
    x = tx + (bx * TILE_WIDTH);
    if(z < z_size && y < y_size && x < x_size){
        outputArray[x + (y * x_size) + (z * y_size * x_size)] = sum;
    }

    __syncthreads();
}

第二种策略是将块大小设置为与输入图块相同。在计算输出时，关闭了一些线程。

#define TILE_X 14 
#define TILE_Y 6 
#define TILE_Z 6 
#define MASK_WIDTH 3
#define MASK_SIZE MASK_WIDTH * MASK_WIDTH * MASK_WIDTH
__constant__ float mask[MASK_WIDTH][MASK_WIDTH][MASK_WIDTH];
__global__ void conv3d(float *input, float *output, const int z_size, const int y_size, const int x_size) {
    __shared__ float inputTile [TILE_Z+MASK_WIDTH-1][TILE_Y+MASK_WIDTH-1][TILE_X+MASK_WIDTH-1];
    int tx = threadIdx.x; int ty = threadIdx.y; int tz = threadIdx.z;
    int bx = blockIdx.x; int by = blockIdx.y; int bz = blockIdx.z;

    int x_o = bx * TILE_X + tx
    int y_o = by * TILE_Y + ty;
    int z_o = bz * TILE_Z + tz;

    int x_i = x_o - MASK_WIDTH/2;
    int y_i = y_o - MASK_WIDTH/2;
    int z_i = z_o - MASK_WIDTH/2;
    if (x_i >= 0 && y_i >= 0 && z_i >= 0 && x_i < x_size && y_i < y_size && z_i < z_size)
        inputTile[tz][ty][tx] = input[(z_i * y_size + y_i) * x_size + x_i];
    else
        inputTile[tz][ty][tx] = 0.0;
    __syncthreads();
    float acc = 0.0;
    if(tz < TILE_Z && ty < TILE_Y && tx < TILE_X) {
        for(int z_mask = 0; z_mask < Z_MASK_WIDTH; z_mask++) {
            for(int y_mask = 0; y_mask < Y_MASK_WIDTH; y_mask++) {
                for(int x_mask = 0; x_mask < X_MASK_WIDTH; x_mask++) {
                    acc += mask[z_mask][y_mask][x_mask] * 
                           inputTile[tz+z_mask][ty+y_mask][tx+x_mask];
                }
             }
         }
    if(z_o < z_size && y_o < y_size && x_o < x_size)
        output[(z_o * y_size + y_o) * x_size + x_o] = acc;
    }
}

知道如何在这些之间进行选择吗？此外，哪个版本在实践中更常用，比如在深度学习中？另外，如果您对代码有任何 cmet，也请告诉我！

Answer 1

当涉及到“哪个更快？”这个问题时的一般答案总是：测量每种方法运行您的应用程序场景的速度以找出答案。在这种情况下，我会说第一种方法在大多数情况下似乎更可取（如果您出于某种原因必须选择这两个选项中的一个）。除非您有一些非常小的卷积核，否则第二种方法将在执行大部分实际工作的部分中有大量线程空闲。在将数据移入和移出全局内存时，请务必避免磁贴内的存储库冲突，并考虑从 warp 中获得的内存访问模式。

最后，卷积基本上只是计算内核系数和输入元素的所有可能组合的总和。由于工作负载本质上只是以某种顺序重复获取这些值，因此卷积几乎必然会受到带宽的限制。因此，有效地进行卷积归结为优化内存访问并尽可能减少带宽。

[...] 哪个版本在实践中更常用，比如在深度学习中？

两者都不是。将嵌套循环扔给它以在空间域中进行蛮力卷积的幼稚方法几乎从来都不是计算卷积的有效方法。卷积对于许多事物来说都是如此基本的运算，因此它已被广泛研究。实际上，您可以阅读有关该主题的数百甚至数千种论文和书籍。在深度学习中，卷积问题有commonly been formulated in terms of general matrix multiplications (GEMMs)，因为这种方法导致了相当好的内存访问模式，并且许多高效的 GEMM 实现可用于 GPU。但也越来越多地使用基于 FFT 的方法以及 other algorithms，具体取决于应用程序。