# Using Numba for GPUs

*Vasileios Karakasis, CSCS*

You may use Numba to generate GPU code from high-level Python functions. Both CUDA and ROC GPUs (NVIDIA and AMD, respectively) are supported, but we are going to focus only on the CUDA capabilities of Numba.

There is an almost one-to-one mapping between the Numba high-level Python abstractions and the different CUDA constructs. This practically means that, as a programmer, you need to take care explicitly of the host/device memory management and you need to be aware of the CUDA programming model.

This demo will teach you to write your first CUDA kernels using Numba. It will cover the basic CUDA programming principles, but it should be enough to kick start you in GPU programming. More specifically, we will cover the following topics:

- Writing a GPU kernel
- Moving data to/from the GPU.
- Spawning a GPU kernel
- Profiling a GPU kernel
- Optimizing memory accesses
- Making use of the shared memory

We will not cover CUDA streams.

## Verify that Numba sees the GPU and understands CUDA

First thing is to check if Numba can detect the GPU. You can achieve this by running the `numba` executable that comes with Numba's installation:

```python
!numba -s
```

    System info:
    --------------------------------------------------------------------------------
    __Time Stamp__
    2019-11-13 10:37:30.961347

    __Hardware Information__
    Machine                                       : x86_64
    CPU Name                                      : haswell
    Number of accessible CPU cores                : 24
    Listed accessible CPUs cores                  : 0-23
    CFS restrictions                              : Information not available
    CPU Features                                  :
    64bit aes avx avx2 bmi bmi2 cmov cx16 f16c fma fsgsbase invpcid lzcnt mmx movbe
    pclmul popcnt rdrnd sahf sse sse2 sse3 sse4.1 sse4.2 ssse3 xsave xsaveopt

    __OS Information__
    Platform                                      : Linux-4.12.14-150.17_5.0.85-cray_ari_c-x86_64-with-glibc2.9
    Release                                       : 4.12.14-150.17_5.0.85-cray_ari_c
    System Name                                   : Linux
    Version                                       : #1 SMP Thu Aug 22 18:29:02 UTC 2019 (24c42de)
    OS specific info                              :
    glibc info                                    : glibc 2.9

    __Python Information__
    Python Compiler                               : GCC 7.2.0
    Python Implementation                         : CPython
    Python Version                                : 3.6.2
    Python Locale                                 : en_US UTF-8

    __LLVM information__
    LLVM version                                  : 8.0.0

    __CUDA Information__
    Found 1 CUDA devices
    id 0    b'Tesla P100-PCIE-16GB'                              [SUPPORTED]
                          compute capability: 6.0
                               pci device id: 0
                                  pci bus id: 2
    Summary:
    	1/1 devices are supported
    CUDA driver version                           : 10010
    CUDA libraries:
    Finding cublas from Conda environment
    	named  libcublas.so.10.2.1.243
    	trying to open library...	ok
    Finding cusparse from Conda environment
    	named  libcusparse.so.10.3.0.243
    	trying to open library...	ok
    Finding cufft from Conda environment
    	named  libcufft.so.10.1.1.243
    	trying to open library...	ok
    Finding curand from Conda environment
    	named  libcurand.so.10.1.1.243
    	trying to open library...	ok
    Finding nvvm from Conda environment
    	named  libnvvm.so.3.3.0
    	trying to open library...	ok
    Finding libdevice from Conda environment
    	searching for compute_20...	ok
    	searching for compute_30...	ok
    	searching for compute_35...	ok
    	searching for compute_50...	ok

    __ROC Information__
    ROC available                                 : False
    Error initialising ROC due to                 : No ROC toolchains found.
    No HSA Agents found, encountered exception when searching:
    Error at driver init:
    NUMBA_HSA_DRIVER /opt/rocm/lib/libhsa-runtime64.so is not a valid file path.  Note it must be a filepath of the .so/.dll/.dylib or the driver:

    __SVML Information__
    SVML state, config.USING_SVML                 : False
    SVML library found and loaded                 : False
    llvmlite using SVML patched LLVM              : True
    SVML operational                              : False

    __Threading Layer Information__
    TBB Threading layer available                 : True
    OpenMP Threading layer available              : True
    Workqueue Threading layer available           : True

    __Numba Environment Variable Information__
    None set.

    __Conda Information__
    conda_build_version                           : not installed
    conda_env_version                             : 4.7.12
    platform                                      : linux-64
    python_version                                : 3.6.2.final.0
    root_writable                                 : False

    __Current Conda Env__
    _libgcc_mutex             0.1                        main
    asn1crypto                1.2.0                    py36_0
    backcall                  0.1.0                    py36_0    anaconda
    blas                      1.0                         mkl
    bokeh                     1.4.0                    py36_0
    ca-certificates           2019.10.16                    0
    cairo                     1.14.12              h8948797_3
    certifi                   2019.9.11                py36_0
    cffi                      1.13.1           py36h2e261b9_0    anaconda
    chardet                   3.0.4                 py36_1003
    click                     7.0                      py36_0
    cloudpickle               1.2.2                      py_0
    conda                     4.7.12                   py36_0    conda-forge
    conda-package-handling    1.6.0            py36h7b6447c_0
    cryptography              2.3.1            py36hc365091_0
    cudatoolkit               10.1.243             h6bb024c_0
    cudnn                     7.6.4                cuda10.1_0    anaconda
    cupy                      6.0.0            py36hc0ce245_0    anaconda
    cupy-cuda101              6.5.0                    pypi_0    pypi
    cycler                    0.10.0                     py_2    conda-forge
    cython                    0.29.13          py36he6710b0_0    anaconda
    cytoolz                   0.10.0           py36h7b6447c_0
    dask                      2.6.0                      py_0
    dask-core                 2.6.0                      py_0
    dbus                      1.13.2               h714fa37_1
    decorator                 4.4.1                      py_0    anaconda
    distributed               2.6.0                      py_0
    expat                     2.2.5             he1b5a44_1004    conda-forge
    fastrlock                 0.4              py36he6710b0_0    anaconda
    fontconfig                2.13.0               h9420a91_0
    freetype                  2.9.1                h8a8886c_1
    fribidi                   1.0.5             h516909a_1002    conda-forge
    fsspec                    0.5.2                      py_0
    gettext                   0.19.8.1          hc5be6a0_1002    conda-forge
    glib                      2.56.2            had28632_1001    conda-forge
    graphite2                 1.3.13            hf484d3e_1000    conda-forge
    graphviz                  2.40.1               h21bd128_2
    gst-plugins-base          1.14.0               hbbd80ab_1
    gstreamer                 1.14.0               hb453b48_1
    h5py                      2.9.0            py36h7918eee_0    anaconda
    harfbuzz                  1.9.0             he243708_1001    conda-forge
    hdf5                      1.10.4               hb1b8bf9_0    anaconda
    heapdict                  1.0.1                      py_0
    icu                       58.2              hf484d3e_1000    conda-forge
    idna                      2.8                      py36_0
    imageio                   2.6.1                    py36_0    conda-forge
    intel-openmp              2019.4                      243
    ipykernel                 5.1.3            py36h39e3cac_0    anaconda
    ipyparallel               6.2.4                    py36_0    anaconda
    ipython                   7.9.0            py36h39e3cac_0    anaconda
    ipython_genutils          0.2.0                    py36_0    anaconda
    jedi                      0.15.1                   py36_0    anaconda
    jinja2                    2.10.3                     py_0
    jpeg                      9b                   h024ee3a_2
    jupyter_client            5.3.4                    py36_0    anaconda
    jupyter_core              4.6.0                    py36_0    anaconda
    kiwisolver                1.1.0            py36hc9558a2_0    conda-forge
    libedit                   3.1                  heed3624_0
    libffi                    3.2.1                hd88cf55_4
    libgcc-ng                 8.2.0                hdf63c60_1
    libgfortran-ng            7.3.0                hdf63c60_0
    libiconv                  1.15              h516909a_1005    conda-forge
    libpng                    1.6.37               hbc83047_0
    libsodium                 1.0.16               h1bed415_0    anaconda
    libstdcxx-ng              8.2.0                hdf63c60_1
    libtiff                   4.0.10               h2733197_2
    libuuid                   1.0.3                h1bed415_2
    libxcb                    1.13              h14c3975_1002    conda-forge
    libxml2                   2.9.9                hea5a465_1
    line_profiler             2.1.2            py36h14c3975_0    anaconda
    llvmlite                  0.30.0           py36hd408876_0
    locket                    0.2.0                    py36_1
    markupsafe                1.1.1            py36h7b6447c_0
    matplotlib                3.1.1            py36h5429711_0
    memory_profiler           0.55.0                   py36_0    anaconda
    mkl                       2019.4                      243
    mkl-service               2.3.0            py36he904b0f_0
    mkl_fft                   1.0.15           py36ha843d7b_0
    mkl_random                1.1.0            py36hd6b4f25_0
    mpi4py                    3.0.3                    pypi_0    pypi
    msgpack-python            0.6.1            py36hfd86e86_1
    nccl                      1.3.5                cuda10.0_0    anaconda
    ncurses                   6.0                  h9df7e31_2
    networkx                  2.4                        py_0    conda-forge
    numba                     0.46.0           py36h962f231_0
    numexpr                   2.7.0            py36h9e4a6bb_0    anaconda
    numpy                     1.17.3           py36hd14ec0e_0
    numpy-base                1.17.3           py36hde5b4d6_0
    olefile                   0.46                     py36_0
    openssl                   1.0.2t               h7b6447c_1
    packaging                 19.2                       py_0
    pandas                    0.25.2           py36he6710b0_0
    pango                     1.42.4               h049681c_0
    parso                     0.5.1                      py_0    anaconda
    partd                     1.0.0                      py_0
    pcre                      8.43                 he1b5a44_0    conda-forge
    pexpect                   4.7.0                    py36_0    anaconda
    pickleshare               0.7.5                    py36_0    anaconda
    pillow                    6.2.1            py36h34e0f95_0
    pip                       19.3.1                   py36_0
    pixman                    0.38.0            h516909a_1003    conda-forge
    prompt_toolkit            2.0.10                     py_0    anaconda
    psutil                    5.6.3            py36h7b6447c_0
    pthread-stubs             0.4               h14c3975_1001    conda-forge
    ptyprocess                0.6.0                    py36_0    anaconda
    pycosat                   0.6.3            py36h14c3975_0
    pycparser                 2.19                     py36_0
    pygments                  2.4.2                      py_0    anaconda
    pyopenssl                 19.0.0                   py36_0
    pyparsing                 2.4.2                      py_0
    pyqt                      5.9.2            py36h05f1152_2
    pysocks                   1.7.1                    py36_0
    python                    3.6.2               hca45abc_19
    python-dateutil           2.8.0                    py36_0
    python-graphviz           0.13                       py_0    conda-forge
    pytz                      2019.3                     py_0
    pywavelets                1.1.1            py36hc1659b7_0    conda-forge
    pyyaml                    5.1.2            py36h7b6447c_0
    pyzmq                     18.1.0           py36he6710b0_0    anaconda
    qt                        5.9.6                h52aff34_0
    readline                  7.0                  ha6073c6_4
    requests                  2.22.0                   py36_0
    ruamel_yaml               0.15.46          py36h14c3975_0
    scikit-image              0.15.0           py36he6710b0_0
    scipy                     1.3.1            py36h7c811a0_0    anaconda
    setuptools                41.6.0                   py36_0
    sip                       4.19.8          py36hf484d3e_1000    conda-forge
    six                       1.12.0                   py36_0
    sortedcontainers          2.1.0                    py36_0
    sqlite                    3.23.1               he433501_0
    tbb                       2019.8               hfd86e86_0
    tblib                     1.4.0                      py_0
    tk                        8.6.8                hbc83047_0
    toolz                     0.10.0                     py_0
    tornado                   6.0.3            py36h7b6447c_0
    tqdm                      4.36.1                     py_0
    traitlets                 4.3.3                    py36_0    anaconda
    urllib3                   1.24.2                   py36_0
    wcwidth                   0.1.7                    py36_0    anaconda
    wheel                     0.33.6                   py36_0
    xorg-libxau               1.0.9                h14c3975_0    conda-forge
    xorg-libxdmcp             1.1.3                h516909a_0    conda-forge
    xz                        5.2.4                h14c3975_4
    yaml                      0.1.7                had09818_2
    zeromq                    4.3.1                he6710b0_3    anaconda
    zict                      1.0.0                      py_0
    zlib                      1.2.11               h7b6447c_3
    zstd                      1.3.7                h0b5b093_0
    --------------------------------------------------------------------------------
    If requested, please copy and paste the information between
    the dashed (----) lines, or from a given specific section as
    appropriate.

    =============================================================
    IMPORTANT: Please ensure that you are happy with sharing the
    contents of the information present, any information that you
    wish to keep private you should remove before sharing.
    =============================================================

If this command does not work out of the box for your Numba installation, you may have to set the `CUDA_HOME` environment variable to point to your CUDA installation.

## Writing the first kernel

Here is a vector addition kernel that takes two vectors as input and writes the sum in a third vector:

```python
import numba.cuda as cuda

@cuda.jit
def _vecadd_cuda(z, x, y):
    i = cuda.grid(1)
    N = x.shape[0]
    if i >= N:
        return

    z[i] = x[i] + y[i]

```

Pretty simple, right? Let's explain each line of this code.

The `@cuda.jit` decorator will compile the following function into a CUDA kernel at runtime. We will see the cost of it later on.

A CUDA kernel in Numba is a Python function that does *not* return a value. This is in accordance with CUDA, where kernel functions are declared `void`. The arguments of the function can be either Numpy arrays or scalars of a Numba recognized type.

CUDA kernels specify the work to be done by a single GPU thread. Due to the massive hardware parallelism available on the GPU, a single GPU thread gets only a tiny portion of work, in this case, the sumation of a single element in the vectors. Since each element is independent from each other, we can safely create as many threads as the elements of the target vectors and spawn them (we will see how later).

Threads on the GPU as organized in groups, called CUDA blocks. There are constraints on the maximum number of threads a block can contain. For the P100 GPUs on Daint, this is 1024. This means that you will need multiple blocks to calculate the sum of large vectors. The blocks that comprise a kernel form the *grid*. Threads inside a block are numbered sequentially and the blocks inside the grid are numbered, too. The figure below shows the arrangement of threads for the vector addition example.

In order to obtain the i-th element of the vector given a block size $B$, you would have to calculate the following:

\begin{equation}
i = i_{b}B + i_{t}
\end{equation}

where $i_{b}$ is the block index and $i_{t}$ is the thread index. CUDA provides this information and Numba makes it available, so that the above statement would be written as follows:

```python
i = cuda.blockIdx.x*cuda.blockDim.x + threadIdx.x
```

Blocks and grids can be three dimensional, thus the `.x` attribute in all of these variables (the other dimensions are accessed through the `.y` and `.z` attributes).

Since obtaining the absolute position of a thread in a CUDA is quite common operation, Numba offers the convenience function `grid()`, which essentially does automatically the above calculation:

```python
i = cuda.grid(1)
```

The argument passed to the `grid()` function is the number of dimensions of the grid.

The next three statements in the code simply check that we don't overrun the arrays in case that their dimension is not a multiple of the block size. In this case, some threads of the last block will remain idle:

```python
    N = x.shape[0]
    if i >= N:
        return
```

> The threads inside a block are run in batches of 32 at once, called *warps*. All the threads of the warp execute the same instruction. If the program control flow diverges for some of the threads of the warp, e.g., due to an `if` condition, the branches will be executed sequentially by disabling the non participating threads. This condition is called *warp divergence* and may incur a performance penalty.

Finally, the `z[i] = x[i] + y[i]` computes the actual sum.

## Preparing the data for the GPU

All Numba CUDA kernels operate on data residing on the GPU. That means that the `x`, `y` and `z` arrays must be  transferred from the host (CPU) to the device (accelerator) before calling the CUDA kernel.

Let's create first two vectors on the host:

```python
import numpy as np

N = 1024
x = np.random.rand(N)
y = np.random.rand(N)
```

Let's also compute our reference result:

```python
z_ref = x + y
```

Now it's time to transfer the vectors to the GPU:

```python
d_x = cuda.to_device(x)
d_y = cuda.to_device(y)
```

The `d_x` and `d_y` are Numpy *array-like* objects that have their data mapped in the GPU memory:

```python
d_x
```

    <numba.cuda.cudadrv.devicearray.DeviceNDArray at 0x2aaae185d240>

For the `z` array, we don't need to create it on the host and copy it, since it is essentially an output only array. We can simply allocate it directly on the GPU and copy it out when the kernel finishes.

```python
d_z = cuda.device_array_like(x)
```

This will simply allocate an array like `x`, i.e., with the same element type, shape and order, on the GPU.

Data is ready, now it's time to call the kernel!

## Invoking a CUDA kernel

When invoking a CUDA kernel you have to specify the block size to use and the corresponding grid size. Picking the right size of block is not always straightforward, but usually values between 64 and 256 are good enough.

> The block size has a direct effect on the *occupancy* of each GPU SM, i.e., how much the actual hardware threads of the SM are utilized, and as a result to the occupancy of the whole GPU. Performance-wise, though, it does not have such a big impact. Nvidia provides a nice tool for calculating the occupancy of the GPU, that you can find [here](https://docs.nvidia.com/cuda/cuda-occupancy-calculator/CUDA_Occupancy_Calculator.xls).

For our kernel, we will select a block size of 128 threads.

```python
block_size = 128
```

Having decided the block size we need to set up the grid:

```python
num_blocks = N // block_size
if N % block_size:
    num_blocks += 1
```

```python
num_blocks
```

    8

Now we are ready to call the kernel!

```python
_vecadd_cuda[num_blocks, block_size](d_z, d_x, d_y)
```

Both the block and the grid could have been two-dimensional (not in this example), in which case you would just define them as tuples.

## Copying back the results

As mentioned before, kernels operate on GPU data only. We need a way to transfer back to the host the result, which is the `d_z` array. Here's how:

```python
res = d_z.copy_to_host()
```

Let's validate the result though to make sure that everything has worked fine:

```python
assert np.allclose(z_ref, res)
```

## Putting it all together

For completeness and easy reference, here is the whole vector addition example:

```python
import numba.cuda as cuda
import numpy as np

@cuda.jit
def _vecadd_cuda(z, x, y):
    '''
    The CUDA kernel

    There's no `cache=True` option like @numba.jit. Every time we're calling the kernel, Numba is going to recompile it
    '''
    i = cuda.grid(1)
    N = x.shape[0]
    if i >= N:
        return

    z[i] = x[i] + y[i]

# Set up the host vectors
N = 1000 * 1000 * 100
x = np.random.rand(N)
y = np.random.rand(N)

# Copy and allocate data on the device
d_x = cuda.to_device(x)
d_y = cuda.to_device(y)
d_z = cuda.device_array_like(x)

# Set up the kernel invocation
block_size = 128
num_blocks = N // block_size
if N % block_size:
    num_blocks += 1

%timeit -n 1 -r 5 ref = x + y

# Call the kernel
%timeit -n 1 -r 5 _vecadd_cuda[num_blocks, block_size](d_z, d_x, d_y)

# Copy back the result to the host
res = d_z.copy_to_host()

# Validate the result
assert np.allclose(x + y, res)
```

    49.5 ms ± 67.7 µs per loop (mean ± std. dev. of 5 runs, 1 loop each)
    The slowest run took 438.93 times longer than the fastest. This could mean that an intermediate result is being cached.
    22.8 ms ± 45.1 ms per loop (mean ± std. dev. of 5 runs, 1 loop each)

### Exercise

> 1. Make the array sufficiently large and time the Numpy version of the sum `x + y`.
> 2. Now time the call to the CUDA kernel with `%timeit -n1 -r1`. What do you see?
> 3. Try timing the CUDA kernels with `%timeit -n1 -r2`. What is happening?

```python

```