CUDA optimization for a vector tensor product using a custom kernel or CUBLAS

Viewed 175

I have two vectors a and b. Each vector contains the coordinates of a 3d points x, y, z vector3f.

struct Vector3f
{ 
    float x;
    float y;
    float z;
}

vector a has a size of n = 5000 points and vector b has a size of m = 4000. I need to do a tensor vector product between them like on the right side of the picture. the resulted vector should have a length size of 5000 * 4000 and contain float point where results are stored at c. enter image description here

__global__ void tensor3dProdcutClassic(const int n, const int m, const Vector3f *a, const Vector3f *b, float *c) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    // int j = blockIdy.y * blockDim.y + threadIdx.y;

    //check if  the idx is out of range
    if (i < n) {
        for (int j = 0; j < m; j++) {
            int idx = j + m * i;
            c[idx] = a[i].x * b[j].x + a[i].y * b[j].y + a[i].z * b[j].z;
        } 
    }
} 

dim3 blockSize(32, 1, 1);
dim3 gridSize((n + blockSize.x - 1) / blockSize.x, 1, 1);

tensor3dProdcutClassic<<<gridSize, blockSize>>>(n, m, x, y, out);

I get high execution time on Volta arch which is a lot.
My question is how can I optimize the kernel to reduce time which is mainly because of the for loop inside the kernel. I know here that all global reads and writes are not coalesced.

1 Answers

You can make the kernel move through both a and b simultaneously, like this:

__global__ void tensor3dProdcutClassic(const int n, const int m, const Vector3f *a, const Vector3f *b, float *c)
{
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    int j = blockIdy.y * blockDim.y + threadIdx.y;

    if (i < n && j < m)
    {
        int idx = j + m * i;
        c[idx] = a[i].x * b[j].x + a[i].y * b[j].y + a[i].z * b[j].z;
    }
}

dim3 blockSize(32, 32);
dim3 gridSize((int)ceil(n / 32.0), (int)ceil(m / 32.0));

tensor3dProdcutClassic<<<gridSize, blockSize>>>(n, m, x, y, out);

Update
I tried to modify the code to use a single array with and without shared memory, the code without shared memory with always faster 3 or 4 times.

With shared memory:

#define BLOCK_SIZE 32
void tensor3dProdcut(const int n, const int m, const float* a, const float* b, float* c)
{
    float* d_a;
    size_t size = (uint64_t)n * 3 * sizeof(float);
    cudaMalloc(&d_a, size);
    cudaMemcpy(d_a, a, size, cudaMemcpyHostToDevice);
    float* d_b;
    size = (uint64_t)m * 3 * sizeof(float);
    cudaMalloc(&d_b, size);
    cudaMemcpy(d_b, b, size, cudaMemcpyHostToDevice);
    float* d_c;
    size = (uint64_t)n * m * sizeof(float);
    cudaMalloc(&d_c, size);
    dim3 dimBlock(BLOCK_SIZE, BLOCK_SIZE);
    dim3 dimGrid((int)ceil((double)n / BLOCK_SIZE), (int)ceil((double)m / BLOCK_SIZE));
    tensor3dProdcutKernel<<<dimGrid, dimBlock>>>(d_a, d_b, d_c, n, m);
    cudaMemcpy(c, d_c, size, cudaMemcpyDeviceToHost);
    cudaFree(d_a);
    cudaFree(d_b);
    cudaFree(d_c);
}

__global__ void tensor3dProdcutKernel(float* a, float* b, float* c, int n, int m)
{
    int i, blockRow, blockCol, row, col;
    float Cvalue;
    blockRow = blockIdx.x;
    blockCol = blockIdx.y;
    row = threadIdx.x;
    col = threadIdx.y;
    if (blockRow * BLOCK_SIZE + row >= n || blockCol * BLOCK_SIZE + col >= m)
        return;
    __shared__ double as[BLOCK_SIZE][3];
    __shared__ double bs[BLOCK_SIZE][3];
    for (i = 0; i < 3; i++)
    {
        as[row][i] = a[(BLOCK_SIZE * blockRow + row) * 3 + i];
        bs[col][i] = b[(BLOCK_SIZE * blockCol + col) * 3 + i];
    }
    __syncthreads();
    Cvalue = 0;
    for (i = 0; i < 3; i++)
        Cvalue += as[row][i] * bs[col][i];
    c[(BLOCK_SIZE * blockRow + row) * m + BLOCK_SIZE * blockCol + col] = Cvalue;
}

Without shared memory:

__global__ void tensor3dProdcutKernel(float* a, float* b, float* c, int n, int m)
{
    int i, blockRow, blockCol, row, col;
    float Cvalue;
    blockRow = blockIdx.x;
    blockCol = blockIdx.y;
    row = threadIdx.x;
    col = threadIdx.y;
    if (blockRow * BLOCK_SIZE + row >= n || blockCol * BLOCK_SIZE + col >= m)
        return;
    Cvalue = 0;
    for (i = 0; i < 3; i++)
        Cvalue += a[(BLOCK_SIZE * blockRow + row) * 3 + i] * b[(BLOCK_SIZE * blockCol + col) * 3 + i];
    c[(BLOCK_SIZE * blockRow + row) * m + BLOCK_SIZE * blockCol + col] = Cvalue;
}
Related