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

Question

Based on my study, there are 2 different strategies to implement tiled version of convolution with cuda. I want to know more about this, and would like to see how they compare with each other, what is the advantage and disadvantage of each strategy, and how to choose. Below is the implementations of the two different strategies.

Strategy 1: the tile size matches with the output size, and needs multiple steps to load the input.

#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();
}

The second strategy is to set the block size to be the same with input tile. In calculating output, some threads are turned off.

#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;
    }
}

Any idea about how to choose between these? In addition, which version is used more often in practice, like in deep learning? Also if you have any comments on the code, please also let me know!

Answer 1

The general answer whenever it comes to the question of "which is faster?" is always: measure how fast each approach runs your application scenario to find out. In this case, I would say that the first approach would seem preferable most of the time (if you had to pick one of those two options for some reason). Unless you have some very tiny convolution kernels, the second approach would have lots of threads idle in the parts that do much of the actual work. Be sure to avoid bank conflicts within your tiles and think about the memory access patterns you get from your warps when moving data to and from global memory.

In the end, convolution is basically just computing sums over all possible combinations of kernel coefficients and input elements. Since the workload is essentially just repeatedly fetching these values in some order, convolution is almost necessarily going to be limited by bandwidth. Thus, doing convolution efficiently comes down to optimizing memory access and reducing bandwidth as much as possible.

[…] which version is used more often in practice, like in deep learning?

Neither. The naïve approach of throwing nested loops at it to brute-force convolution in the spatial domain is almost never an efficient way of computing convolutions. Convolution is such a fundamental operation for so many things that it has been studied extensively. There are literally hundreds, if not thousands of papers and books you could read on the subject. In deep learning, the problem of convolution has commonly been formulated in terms of general matrix multiplications (GEMMs) since this approach leads to rather nice memory access patterns and many efficient GEMM implementations are available for the GPU. But also FFT-based approaches as well as other algorithms are increasingly used depending on the application.