My goal is to speedup matrix multiplication using CUDA. Therefore I wrote a small python program (see code below) which compares the performance of just-in-time compilation using Numba and CUDA. The code will do a matrix multiplication using 3-dim matrices, the multiplication is done on the last two dimension, parallelization is done on the first dimension.
The code works, however I am wondering why the second Memcpy (D2H) is not done earlier (see screenshot). Can somebody please explain me that behavior? In general it is possible to have overlapping kernel execution and memcpy as visible in the first stream.
I am using a GTX 1050Ti (6.1 computing capability).
Python code
from numba import cuda
from numba import vectorize
from numba import guvectorize
from numba import cuda, float32, float64
from numba import njit
import numpy as np
import math
from time import time
dev = cuda.current_context().device
print(dev)
print('CUDA device [%s]' % dev.name.decode('utf-8'))
TPB = 8
@njit
def _mat_mul(X, Y):
'''Helper function to multiply two 3D matrices. Multiplication is done on last two axes.
'''
result = np.zeros((X.shape[0], X.shape[1], Y.shape[2]), dtype=np.float64)
for l in range(X.shape[0]):
# iterate through rows of X
for i in range(X.shape[1]):
# iterate through columns of Y
for j in range(Y.shape[2]):
# iterate through rows of Y
for k in range(Y.shape[1]):
result[l][i][j] += X[l][i][k] * Y[l][k][j]
return result
def matmul_njit(A, B):
# Matmul using njit
return _mat_mul(A, B)
@cuda.jit
def _fast_matmul_3d_streams(A, B, C, stream_size):
"""
Perform matrix multiplication of C = A * B
Each thread computes one element of the result matrix C
"""
# Define an array in the shared memory
# The size and type of the arrays must be known at compile time
sA = cuda.shared.array(shape=(TPB, TPB), dtype=float64)
sB = cuda.shared.array(shape=(TPB, TPB), dtype=float64)
n, x, y = cuda.grid(3) # get absolute position in grid
# Get thread ID inside block
tn = cuda.threadIdx.x
tx = cuda.threadIdx.y
ty = cuda.threadIdx.z
if x >= C.shape[1] or y >= C.shape[2] or n >= C.shape[0] or n >= stream_size: # Quit if (x, y) is outside of valid C boundary
return
# Each thread computes one element in the result matrix.
# The dot product is chunked into dot products of TPB-long vectors.
tmp = 0.0
for i in range(int(A.shape[2] / TPB)):
# Preload data into shared memory
sA[tx, ty] = A[n, x, ty + i * TPB]
sB[tx, ty] = B[n, tx + i * TPB, y]
# Wait until all threads finish preloading
cuda.syncthreads()
# Computes partial product on the shared memory
for j in range(TPB):
tmp += sA[tx, j] * sB[j, ty]
# Wait until all threads finish computing
cuda.syncthreads()
C[n, x, y] = tmp
def main():
# The data array
n = 512 * 512 * 8
print(n)
A = cuda.pinned_array((n, TPB, TPB), np.float32)
B = cuda.pinned_array((n, TPB, TPB), np.float32)
A_numba = np.zeros(shape=(n, TPB, TPB), dtype=np.float32)
B_numba = np.zeros(shape=(n, TPB, TPB), dtype=np.float32)
A[0, 0, 0] = 12.0
A[0, 1, 1] = 2.0
A[0, 1, 0] = 4.3
A[0, 1, 1] = 6.5
B[0, 1, 1] = 12.0
B[0, 1, 1] = 2.5
B[0, 0, 1] = 4.3
B[0, 1, 0] = 6.0
A[4, 0, 0] = 12.0
A[568, 1, 1] = 2.0
A[45, 1, 0] = 4.3
A[67, 1, 1] = 6.5
B[3, 1, 1] = 12.0
B[56923, 1, 1] = 2.5
B[10000, 0, 1] = 4.3
B[660, 1, 0] = 6.0
A[5, 0, 0] = 16.0
A[5, 1, 1] = 2.0
A[0, 1, 0] = 7.3
A[5, 1, 1] = 6.5
B[7, 1, 1] = 12.0
B[1, 1, 1] = 2.5
B[10, 0, 1] = 24.3
B[13, 1, 0] = 36.0
start = time()
# Cuda streams
n_streams = 32
stream_size = int(A.shape[0] / n_streams)
stream_list = []
for i in range(n_streams):
stream = cuda.stream()
stream_list.append(stream)
# Configure the blocks
threadsperblock = (1, TPB, TPB) # (32 x TPB x TPB) threads, maximum is 1024. Does not really matter
blockspergrid_n = int(math.ceil(stream_size / threadsperblock[0]))
blockspergrid_x = int(math.ceil(A.shape[1] / threadsperblock[1])) # such that number of threads matches the number of matrix elements
blockspergrid_y = int(math.ceil(B.shape[2] / threadsperblock[2])) # such that number of threads matches the number of matrix elements
blockspergrid = (blockspergrid_n, blockspergrid_x, blockspergrid_y)
# Result arrays
C_global_mem = cuda.device_array((stream_size, A.shape[1], B.shape[2]))
C = cuda.pinned_array((A.shape[0], A.shape[1], B.shape[2]))
C[:, :, :] = 0.0
# Streams
A_global_mem = []
B_global_mem = []
for i in range(n_streams):
offset = i * stream_size
# Copy data to device
A_global_mem.append(cuda.to_device(A[offset:offset + stream_size, :, :], stream=stream_list[i]))
B_global_mem.append(cuda.to_device(B[offset:offset + stream_size, :, :], stream=stream_list[i]))
for i in range(n_streams):
# Run kernel
_fast_matmul_3d_streams[blockspergrid, threadsperblock, stream_list[i]](A_global_mem[i], B_global_mem[i], C_global_mem, stream_size)
for i in range(n_streams):
# Copy to host
C[offset:offset + stream_size, :, :] = C_global_mem.copy_to_host(stream=stream_list[i])
print("GPU execution time using streams (including I/O): " + str(time() - start))
start = time()
C_res = matmul_njit(A_numba, B_numba)
print("Numba execution time: " + str(time() - start))
if __name__ == "__main__":
for _ in range(2):
main()