<jupyter_start><jupyter_text>Understanding the memory performance of the GPU*<NAME>, CSCS*In this section we are going to investigate a crucial aspect of the memory locality on the GPUs. It should be perceived in a slightly different way than on the CPUs. To demonstrate this, we will use BLAS matrix-vector kernel using all the tricks we have learned so far. The threads of the GPU operate row-wise in the input matrix, each one taking care of a single row to compute:<jupyter_code>import numba
import numba.cuda as cuda
import numpy as np
import time


class time_region:
    def __init__(self, time_offset=0):
        self._time_off = time_offset

    def __enter__(self):
        self._t_start = time.time()
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        self._t_end = time.time()

    def elapsed_time(self):
        return self._time_off + (self._t_end - self._t_start)


class time_region_cuda:
    def __init__(self, time_offset=0, cuda_stream=0):
        self._t_start = cuda.event(timing=True)
        self._t_end = cuda.event(timing=True)
        self._time_off = time_offset
        self._cuda_stream = cuda_stream

    def __enter__(self):
        self._t_start.record(self._cuda_stream)
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        self._t_end.record(self._cuda_stream)
        self._t_end.synchronize()

    def elapsed_time(self):
        return self._time_off + 1.e-3*cuda.event_elapsed_time(self._t_start,
                                                              self._t_end)


@cuda.jit('void(float64, Array(float64, 2, "C"), Array(float64, 1, "C"), '
          'float64, Array(float64, 1, "C"))')
def _gemv_cuda(alpha, A, x, beta, y):
    i = cuda.grid(1)
    N, M = A.shape
    if i >= N:
        return

    prod = 0.0
    for j in range(M):
        prod += A[i, j]*x[j]

    y[i] = alpha*prod + beta*y[i]


def gemv_gpu(alpha, A, x, beta, y):
    # Works only for square matrices
    N = A.shape[0]
    with time_region_cuda() as t_xfer:
        d_A = cuda.to_device(A)
        d_x = cuda.to_device(x)
        d_y = cuda.to_device(y)
        y_ret = cuda.pinned_array(N)
        
    block_size = 128
    num_blocks = N // block_size
    if N % block_size:
        num_blocks += 1

    with time_region_cuda() as t_kernel:
        _gemv_cuda[num_blocks, block_size](alpha, d_A, d_x, beta, d_y)

    with time_region_cuda(t_xfer.elapsed_time()) as t_xfer:
        d_y.copy_to_host(y_ret)

    print(f'  CUDA transfer overheads: {t_xfer.elapsed_time()}')
    print(f'  CUDA kernel time: {t_kernel.elapsed_time()}')
    print(f'  Consumed memory bandwidth: {1e-9*8*N*(N+2)/t_kernel.elapsed_time()} GB/s')
    return y_ret

N = 1024*16
A = np.random.rand(N, N)
x = np.random.rand(N)
y_orig = np.ones(N)
alpha = 0.2
beta = 1

with time_region() as t_gpu:
    y = gemv_gpu(alpha, A, x, beta, y_orig)

with time_region() as t_ref:
    y_ref = alpha*(A @ x) + beta*y_orig
    
    
cuda.profile_stop()

print(f'Total time (GPU): {t_gpu.elapsed_time()} s')
print(f'Total time (CPU): {t_ref.elapsed_time()} s')

assert np.allclose(y, y_ref)<jupyter_output>CUDA transfer overheads: 0.2124068799316883
  CUDA kernel time: 0.015253503799438477
  Consumed memory bandwidth: 140.8034390157011 GB/s
Total time (GPU): 0.23821115493774414 s
Total time (CPU): 0.04474353790283203 s<jupyter_text>Exercise> Increase the array size and record the CUDA kernel performance time. How much faster is it compared to the CPU version?As you might have noticed already, the GPU kernels is only about 2-3x faster than the CPU version. Not as good it has been with the vector addition kernel. Is this expected, is this how it should be? Let's look into the kernel in more detail.The kernel needs to read the whole matrix $A$ and the vectors $x$ and $y$, i.e., $8(N^2 + 2N)$ bytes need to be transferred to/from main memory in total. At the same time, the kernel performs $2N^2 + 3N$ floating point operations in total. This leads to an arithmetic intensity or flop:byte ratio equals to $\frac{2N(N+\frac{3}{2})}{8N(N+2)} \approx 0.25$. This ratio is much higher than that for the vector addition kernel, but it is very low to make the kernel compute bound.> Given the nominal peak double precision performance (5.3 Tflop/s) and the nominal peak memory bandwidth of the P100 GPUs (732 GB/s), a kernel would need a flop:byte ratio of at least 7.24, so as to be compute bound.So, theoretically, we should be approaching the effective memory bandwidth limit of the device, but we only achieve 1/4 of it. The CPU kernel on the other hand seems to hit its memory bandwidth limit:<jupyter_code>print(f'CPU memory bandwidth consumed: {1e-9*8*N*(N+2)/t_ref.elapsed_time()}')<jupyter_output>CPU memory bandwidth consumed: 47.884390070318666<jupyter_text>What is going on? The following figure shows the memory layout of the $A$ matrix, the $x$ and $y$ vectors and how the CUDA threads are arranged for the computation.![Memory layout and thread arrangement for the threads](figs/matvec-memlayout.png)The matrix is stored in the default, row-major, order and each thread is assigned a row of the matrix.> Assigning a column to each thread would require a reduction operations across all the threads on the device. Global reduction on GPUs is not straightforward to achieve, because there can't be any sort of global synchronization. You would have to write multiple kernels to achieve a global reduction. For reduction operations, you should use the `@cuda.reduce` decorator (see [here](http://numba.pydata.org/numba-doc/latest/cuda/reduction.html) for more information).In this arrangement subsequent threads access non-contiguous memory, but each thread accesses memory sequentially. This would be almost ideal on the CPU: perfect read access on both the $A$ matrix and the $x$ vector. What is "wrong" with the GPUs?> For the CPUs the ideal is to assign a chunk of rows to each thread, because the above arrangement would lead to false sharing.Recall from our discussion on CUDA blocks in the beginning how threads are executed on the GPU. Threads are executed in batches of 32, called warps, all of them executing the same instruction. In our example, that means that each thread in the warp will access memory allocated for the $A$ matrix in $8N$ bytes strides.The global memory on the GPU is organized in 256-byte memory segments and can be accessed in transactions of 32, 64 or 128 bytes. If all the threads of a warp access the same memory segments, a maximum of two memory transactions will be performed in order to fetch the values required by all the threads of the warp. This is called *memory coalescing* in CUDA's terminology. In our example, a warp needs to fetch $32\times 8 = 256$ bytes, i.e., 2 memory transactions of 128 bytes each ideally. However, due to the row-major layout of the matrix $A$, each thread does a separate 32 byte transaction, generating $32\times 32 = 1024$ bytes memory traffic per warp, which is four times more than the ideal! How much of the memory bandwidth did our code utilize? Exercise> Change the memory layout of matrix $A$ to column major (or Fortran in NumPy's nomenclature) and measure again the performance. How close is it to the effective memory bandwidth limit? Hints: (a) `numpy.random.rand()` does not accept an `order` argument. Use the `asarray()` function to create matrix $A$ from the result of `numpy.random.rand()` and with the desired order. (b) Make sure to adjust order in the JIT function signature. Taking advantage of the shared memoryOn every GPU multiprocessor chip (SM), there is an on-chip memory called *shared memory*. This can either act as L1 data cache memory (default mode) or it can serve as programmable scratchpad at the disposal of the programmer. This memory is shared among all the threads of a CUDA block. We are going to modify our matrix-vector kernel to make use of it. If you inspect closer the algorithm, you will see that although there is no temporal locality in accesses to the matrix $A$, the the vector $x$ is reused $N$ times. For this reason, we are going to cache the vector $x$ manually into the shared memory. The following figure shows the principle in more detail:![Matrix-vector multiplication using shared memory](figs/matvec-shared.png)We essentially process the matrix with a sliding window. The threads of a block (still 1D) undertake a double role: first, they fetch the $x$ vector elements into the shared memory of the block and, second, they perform the multiplication and the reduction to a local register variable. As soon as they process the full row, the store this value back to vector $y$. Here is the implementation in Numba:<jupyter_code>import numba
import numba.cuda as cuda
import numpy as np
import time


class time_region:
    def __init__(self, time_offset=0):
        self._time_off = time_offset

    def __enter__(self):
        self._t_start = time.time()
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        self._t_end = time.time()

    def elapsed_time(self):
        return self._time_off + (self._t_end - self._t_start)


class time_region_cuda:
    def __init__(self, time_offset=0, cuda_stream=0):
        self._t_start = cuda.event(timing=True)
        self._t_end = cuda.event(timing=True)
        self._time_off = time_offset
        self._cuda_stream = cuda_stream

    def __enter__(self):
        self._t_start.record(self._cuda_stream)
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        self._t_end.record(self._cuda_stream)
        self._t_end.synchronize()

    def elapsed_time(self):
        return self._time_off + 1.e-3*cuda.event_elapsed_time(self._t_start,
                                                              self._t_end)
BLOCK_SIZE = 128
    
@cuda.jit('void(float64, Array(float64, 2, "F"), Array(float64, 1, "F"), '
          'float64, Array(float64, 1, "F"))')
def _gemv_cuda_shared(alpha, A, x, beta, y):
    i = cuda.grid(1)
    N, M = A.shape
    if i >= N:
        return

    lx = cuda.shared.array(shape=BLOCK_SIZE, dtype=numba.float64)
    bsize = cuda.blockDim.x
    tid = cuda.threadIdx.x
    num_blocks = cuda.gridDim.x

    prod = 0.0
    for b in range(num_blocks):
        lx[tid] = x[tid + b*bsize]
        cuda.syncthreads()

        for j in range(BLOCK_SIZE):
            prod += A[i, j + b*bsize]*lx[j]

        cuda.syncthreads()

    y[i] = alpha*prod + beta*y[i]

def gemv_gpu(alpha, A, x, beta, y):
    # Works only for square matrices
    N = A.shape[0]
    with time_region_cuda() as t_xfer:
        d_A = cuda.to_device(A)
        d_x = cuda.to_device(x)
        d_y = cuda.to_device(y)
        y_ret = cuda.pinned_array(N)
        
    num_blocks = N // BLOCK_SIZE
    if N % BLOCK_SIZE:
        num_blocks += 1

    with time_region_cuda() as t_kernel:
        _gemv_cuda_shared[num_blocks, BLOCK_SIZE](alpha, d_A, d_x, beta, d_y)

    with time_region_cuda(t_xfer.elapsed_time()) as t_xfer:
        d_y.copy_to_host(y_ret)

    print(f'  CUDA transfer overheads: {t_xfer.elapsed_time()}')
    print(f'  CUDA kernel time: {t_kernel.elapsed_time()}')
    print(f'  Consumed memory bandwidth: {1e-9*8*N*(N+2)/t_kernel.elapsed_time()} GB/s')
    return y_ret

N = 1024*16
A = np.asarray(np.random.rand(N, N), order='F')
x = np.random.rand(N)
y_orig = np.ones(N)
alpha = 0.2
beta = 1

with time_region() as t_gpu:
    y = gemv_gpu(alpha, A, x, beta, y_orig)

with time_region() as t_ref:
    y_ref = alpha*(A @ x) + beta*y_orig
    
    
cuda.profile_stop()

print(f'Total time (GPU): {t_gpu.elapsed_time()} s')
print(f'Total time (CPU): {t_ref.elapsed_time()} s')

assert np.allclose(y, y_ref)<jupyter_output>CUDA transfer overheads: 0.21740512585639954
  CUDA kernel time: 0.005159103870391846
  Consumed memory bandwidth: 416.3021032249296 GB/s
Total time (GPU): 0.22838807106018066 s
Total time (CPU): 0.0459444522857666 s
{'fraction_non_alphanumeric': 0.09007199062447681, 'fraction_numerical': 0.0285451197053407, 'mean_word_length': 3.5100037750094377, 'pattern_counts': {'":': 0, '<': 11, '<?xml version=': 0, '>': 18, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '3', 'dir': 'jupyter-structured-clean-dedup', 'id': '264016', 'license': 'BSD-3-Clause', 'n_tokens_mistral': 4086, 'n_tokens_neox': 3777, 'n_words': 1436, 'path': 'numba-cuda/3.memory-optimization.ipynb', 'repo_name': 'edoardob90/PythonHPC', 'size': '11959'}