gtc2019-numba.git/4 - Writing CUDA Kernels.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# GTC 2019 Numba Tutorial Notebook 4: 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.  We highly recommend that everyone writing CUDA kernels with Numba take the time to read Chapters 1 and 2 of the CUDA C Programming Guide:\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": {},
   "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):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "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]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "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",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "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])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "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",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "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]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "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",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x_device = cuda.to_device(x)\n",
    "y_device = cuda.to_device(y)\n",
    "out_device = cuda.device_array_like(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%timeit add_kernel[blocks_per_grid, threads_per_block](x, y, out)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%timeit add_kernel[blocks_per_grid, threads_per_block](x_device, y_device, out_device); out_device.copy_to_host()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "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",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# CPU input/output arrays, implied synchronization for memory copies\n",
    "%time add_kernel[blocks_per_grid, threads_per_block](x, y, out)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "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()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "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()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Always be sure to synchronize with the GPU when benchmarking CUDA kernels!**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Atomic Operations and Avoiding Race Conditions\n",
    "\n",
    "CUDA, like many general purpose parallel execution frameworks, makes it possible to have race condtions in your code.  A race condition in CUDA arises when threads read or write a memory location that might be modified by another independent thread.  Generally speaking, you need to worry about:\n",
    "\n",
    " * read-after-write hazards: One thread is reading a memory location at the same time another thread might be writing to it.\n",
    " * write-after-write hazards: Two threads are writing to the same memory location, and only one write will be visible when the kernel is complete.\n",
    " \n",
    "A common strategy to avoid both of these hazards is to organize your CUDA kernel algorithm such that each thread has exclusive responsibility for unique subsets of output array elements, and/or to never use the same array for both input and output in a single kernel call.  (Iterative algorithms can use a double-buffering strategy if needed, and switch input and output arrays on each iteration.)\n",
    "\n",
    "However, there are many cases where different threads need to combine results.  Consider something very simple, like: \"every thread increments a global counter.\"  Implementing this in your kernel requires each thread to:\n",
    "\n",
    "1. Read the current value of a global counter.\n",
    "2. Compute `counter + 1`.\n",
    "3. Write that value back to global memory.\n",
    "\n",
    "However, there is no guarantee that another thread has not changed the global counter between steps 1 and 3.  To resolve this problem, CUDA provides \"atomic operations\" which will read, modify and update a memory location in one, indivisible step.  Numba supports several of these functions, [described here](http://numba.pydata.org/numba-doc/dev/cuda/intrinsics.html#supported-atomic-operations).\n",
    "\n",
    "Let's make our thread counter kernel:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@cuda.jit\n",
    "def thread_counter_race_condition(global_counter):\n",
    "    global_counter[0] += 1  # This is bad\n",
    "    \n",
    "@cuda.jit\n",
    "def thread_counter_safe(global_counter):\n",
    "    cuda.atomic.add(global_counter, 0, 1)  # Safely add 1 to offset 0 in global_counter array"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# This gets the wrong answer\n",
    "global_counter = cuda.to_device(np.array([0], dtype=np.int32))\n",
    "thread_counter_race_condition[64, 64](global_counter)\n",
    "\n",
    "print('Should be %d:' % (64*64), global_counter.copy_to_host())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# This works correctly\n",
    "global_counter = cuda.to_device(np.array([0], dtype=np.int32))\n",
    "thread_counter_safe[64, 64](global_counter)\n",
    "\n",
    "print('Should be %d:' % (64*64), global_counter.copy_to_host())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "## Exercise\n",
    "\n",
    "For this exercise, create a histogramming kernel.  This will take an array of input data, a range and a number of bins, and count how many of the input data elements land in each bin.  Below is an example CPU implementation of histogramming:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def cpu_histogram(x, xmin, xmax, histogram_out):\n",
    "    '''Increment bin counts in histogram_out, given histogram range [xmin, xmax).'''\n",
    "    # Note that we don't have to pass in nbins explicitly, because the size of histogram_out determines it\n",
    "    nbins = histogram_out.shape[0]\n",
    "    bin_width = (xmax - xmin) / nbins\n",
    "    \n",
    "    # This is a very slow way to do this with NumPy, but looks similar to what you will do on the GPU\n",
    "    for element in x:\n",
    "        bin_number = np.int32((element - xmin)/bin_width)\n",
    "        if bin_number >= 0 and bin_number < histogram_out.shape[0]:\n",
    "            # only increment if in range\n",
    "            histogram_out[bin_number] += 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = np.random.normal(size=10000, loc=0, scale=1).astype(np.float32)\n",
    "xmin = np.float32(-4.0)\n",
    "xmax = np.float32(4.0)\n",
    "histogram_out = np.zeros(shape=10, dtype=np.int32)\n",
    "\n",
    "cpu_histogram(x, xmin, xmax, histogram_out)\n",
    "\n",
    "histogram_out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@cuda.jit\n",
    "def cuda_histogram(x, xmin, xmax, histogram_out):\n",
    "    '''Increment bin counts in histogram_out, given histogram range [xmin, xmax).'''\n",
    "    \n",
    "    pass  # Replace this with your implementation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "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"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}