3_Writing_CUDA_Kernels.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.7.1"
    },
    "colab": {
      "name": "3 - Writing CUDA Kernels.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "include_colab_link": true
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/geoffwoollard/gpu-speedups-mbptechtalk2020/blob/master/3_Writing_CUDA_Kernels.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_rHG1prajec3",
        "colab_type": "text"
      },
      "source": [
        "# MBP Tech Talk 2020 :: Writing CUDA Kernels\n",
        "\n",
        "## The CUDA Programming Model\n",
        "\n",
        "Ufuncs (and generalized ufuncs mentioned in the bonus notebook at the end of the tutorial) are the easiest way in Numba to use the GPU, and present an abstraction that requires minimal understanding of the CUDA programming model.  However, not all functions can be written as ufuncs.  Many problems require greater flexibility, in which case you want to write a *CUDA kernel*, the topic of this notebook. \n",
        "\n",
        "Fully explaining the CUDA programming model is beyond the scope of this tutorial.  I highly recommend that everyone writing CUDA kernels with Numba take the time to read Chapters 1 and 2 of the CUDA C Programming Guide. If the learning barrier is too steep, there are various pedagogical videos out there explaining these ideas and walking through similar examples:\n",
        "\n",
        " * Introduction: http://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#introduction\n",
        " * Programming Model: http://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#programming-model\n",
        "\n",
        "The programming model chapter gets a little in to C specifics, but familiarity with CUDA C can help write better CUDA kernels in Python.\n",
        "\n",
        "For the purposes of this tutorial, the most important thing is to understand this diagram:\n",
        "![Thread Hierarchy](http://docs.nvidia.com/cuda/cuda-c-programming-guide/graphics/grid-of-thread-blocks.png \"Thread Hierarchy (from CUDA C Programming Guide)\")\n",
        "\n",
        "We will be writing a *kernel* that decribes the execution of a single thread in this hierarchy.  The CUDA compiler and driver will execute our kernel across a *thread grid* that is divided into *blocks* of threads.  Threads within the same block can exchange data very easily during the execution of a kernel, whereas threads in different blocks should generally not communicate with each other (with a few exceptions).\n",
        "\n",
        "Deciding the best size for the CUDA thread grid is a complex problem (and depends on both the algorithm and the specific GPU compute capability), but here are some very rough heuristics that we follow:\n",
        "\n",
        "  * the size of a block should be a multiple of 32 threads, with typical block sizes between 128 and 512 threads per block.\n",
        "  * the size of the grid should ensure the full GPU is utilized where possible.  Launching a grid where the number of blocks is 2x-4x the number of \"multiprocessors\" on the GPU is a good starting place.  Something in the range of 20 - 100 blocks is usually a good starting point.\n",
        "  * The CUDA kernel launch overhead does depend on the number of blocks, so we find it best not to launch a grid where the number of threads equals the number of input elements when the input size is very big.  We'll show a pattern for dealing with large inputs below.\n",
        "\n",
        "Each thread distinguishes itself from the other threads using its unique thread (`threadIdx`) and block (`blockIdx`) index values, which can be multidimensional if launched that way."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2FjvHQ5kjec6",
        "colab_type": "text"
      },
      "source": [
        "## A First Example\n",
        "\n",
        "This all will sound a little overwhelming at first, so let's start with a concrete example.  Let's write our addition function for 1D NumPy arrays.  CUDA kernels are compiled using the `numba.cuda.jit` decorator (not to be confused with the `numba.jit` decorator for the CPU). `jit` stands for *just in time*.  `jit` decorates a kernel function (written in Python, not CUDA C!), so some of the learning material out there for writing CUDA C kernels will help you understand what is happening.:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "eFjjSvPnjec8",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "from numba import cuda\n",
        "\n",
        "@cuda.jit\n",
        "def add_kernel(x, y, out):\n",
        "    tx = cuda.threadIdx.x # this is the unique thread ID within a 1D block\n",
        "    ty = cuda.blockIdx.x  # Similarly, this is the unique block ID within the 1D grid\n",
        "\n",
        "    block_size = cuda.blockDim.x  # number of threads per block\n",
        "    grid_size = cuda.gridDim.x    # number of blocks in the grid\n",
        "    \n",
        "    start = tx + ty * block_size\n",
        "    stride = block_size * grid_size\n",
        "\n",
        "    # assuming x and y inputs are same length\n",
        "    for i in range(start, x.shape[0], stride):\n",
        "        out[i] = x[i] + y[i]"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "57UjuGuBjedB",
        "colab_type": "text"
      },
      "source": [
        "That's a lot more typing than our ufunc example, and it is much more limited: only works on 1D arrays, doesn't verify input sizes match, etc.  Most of the function is spent figuring out how to turn the block and grid indices and dimensions into unique offsets into the input arrays.  The pattern of computing a starting index and a stride is a common way to ensure that your grid size is independent of the input size.  The striding will maximize bandwidth by ensuring that threads with consecuitive indices are accessing consecutive memory locations as much as possible.  Thread indices beyond the length of the input (`x.shape[0]`, since `x` is a NumPy array) automatically skip over the for loop.\n",
        "\n",
        "Also note that we did not need to specify a type signature for the CUDA kernel.  Unlike `@vectorize`, Numba can infer the type signature from the inputs automatically, and much more reliably.\n",
        "\n",
        "Let's call the function now on some data:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "scrolled": false,
        "id": "sWCep3hwjedC",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import numpy as np\n",
        "\n",
        "n = 100000\n",
        "x = np.arange(n).astype(np.float32)\n",
        "y = 2 * x\n",
        "out = np.empty_like(x)\n",
        "\n",
        "threads_per_block = 128\n",
        "blocks_per_grid = 30\n",
        "\n",
        "add_kernel[blocks_per_grid, threads_per_block](x, y, out)\n",
        "print(out[:10])"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EiMGDAsNjedG",
        "colab_type": "text"
      },
      "source": [
        "The unusual syntax for calling the kernel function is designed to mimic the CUDA Runtime API in C, where the above call would look like:\n",
        "```\n",
        "add_kernel<<<blocks_per_grid, threads_per_block>>>(x, y, out)\n",
        "```\n",
        "The arguments within the square brackets define the size and shape of the thread grid, and the arguments with parentheses correspond to the kernel function arguments.\n",
        "\n",
        "Note that, unlike the ufunc, the arguments are passed to the kernel as full NumPy arrays.  The kernel can access any element in the array it wants, regardless of its position in the thread grid.  This is why CUDA kernels are significantly more powerful that ufuncs.  (But with great power, comes a greater amount of typing...)\n",
        "\n",
        "Numba includes [several helper functions](http://numba.pydata.org/numba-doc/dev/cuda/kernels.html#absolute-positions) to simplify the thread offset calculations above.  You can write the function much more simply as:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "mnM5SWUZjedI",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "@cuda.jit\n",
        "def add_kernel(x, y, out):\n",
        "    start = cuda.grid(1)      # 1 = one dimensional thread grid, returns a single value\n",
        "    stride = cuda.gridsize(1) # ditto\n",
        "\n",
        "    # assuming x and y inputs are same length\n",
        "    for i in range(start, x.shape[0], stride):\n",
        "        out[i] = x[i] + y[i]"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dR7pbeZejedL",
        "colab_type": "text"
      },
      "source": [
        "As before, using NumPy arrays forces Numba to allocate GPU memory, copy the arguments to the GPU, run the kernel, then copy the argument arrays back to the host.  This not very efficient, so you will often want to allocate device arrays:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "vYSgr1HzjedM",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "x_device = cuda.to_device(x)\n",
        "y_device = cuda.to_device(y)\n",
        "out_device = cuda.device_array_like(x)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "OBuA73l4jedQ",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "%timeit add_kernel[blocks_per_grid, threads_per_block](x, y, out)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "0y3sklaXjedU",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "%timeit add_kernel[blocks_per_grid, threads_per_block](x_device, y_device, out_device); out_device.copy_to_host()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8wxxp9nrjedW",
        "colab_type": "text"
      },
      "source": [
        "## Kernel Synchronization\n",
        "\n",
        "*One extremely important caveat should be mentioned here*: CUDA kernel execution is designed to be asynchronous with respect to the host program.  This means that the kernel launch (`add_kernel[blocks_per_grid, threads_per_block](x_device, y_device, out_device)`) returns immediately, allowing the CPU to continue executing while the GPU works in the background.  Only host<->device memory copies or an explicit synchronization call will force the CPU to wait until previously queued CUDA kernels are complete.\n",
        "\n",
        "When you pass host NumPy arrays to a CUDA kernel, Numba has to synchronize on your behalf, but if you pass device arrays, processing will continue.  If you launch multiple kernels in sequence without any synchronization in between, they will be queued up to run sequentially by the driver, which is usually what you want.  If you want to run multiple kernels on the GPU in parallel (sometimes a good idea, but beware of race conditions!), take a look at [CUDA streams](http://numba.pydata.org/numba-doc/dev/cuda-reference/host.html?highlight=synchronize#stream-management).\n",
        "\n",
        "Here's some sample timings (using `%time`, which only runs the statement once to ensure our measurement isn't affected by the finite depth of the CUDA kernel queue):"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "WGqJ5YCbjedX",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# CPU input/output arrays, implied synchronization for memory copies\n",
        "%time add_kernel[blocks_per_grid, threads_per_block](x, y, out)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "4sRFOB13jeda",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# GPU input/output arrays, no synchronization (but force sync before and after)\n",
        "cuda.synchronize()\n",
        "%time add_kernel[blocks_per_grid, threads_per_block](x_device, y_device, out_device)\n",
        "cuda.synchronize()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "90Ch7Sw6jede",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# GPU input/output arrays, include explicit synchronization in timing\n",
        "cuda.synchronize()\n",
        "%time add_kernel[blocks_per_grid, threads_per_block](x_device, y_device, out_device); cuda.synchronize()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "q51708Lxjedh",
        "colab_type": "text"
      },
      "source": [
        "**Always be sure to synchronize with the GPU when benchmarking CUDA kernels!**"
      ]
    }
  ]
}