2.0 Numba.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "05db90c8-a261-49ec-a4a2-c97cb8013ec7",
   "metadata": {},
   "source": [
    "# Numba\n",
    "\n",
    "[Numba](http://numba.pydata.org/) is an open source JIT compiler that translates a subset of Python and NumPy code into fast machine code. \n",
    "\n",
    "Numba supports CUDA GPU programming by directly compiling a restricted subset of Python code into CUDA kernels and device functions following the CUDA execution model. Kernels written in Numba appear to have direct access to NumPy arrays. NumPy arrays are transferred between the CPU and the GPU automatically."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1df11e1b-234c-4008-bf0c-bbdb897e0f07",
   "metadata": {},
   "source": [
    "## What is a kernel?\n",
    "\n",
    "A kernel is similar to a function, it is a block of code which takes some inputs and is executed by a processor.\n",
    "\n",
    "The difference between a function and a kernel is:\n",
    "- A kernel cannot return anything, it must instead modify memory\n",
    "- A kernel must specify its thread hierarchy (threads and blocks)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "412bb30f-31d3-41b5-bec7-9d7dd72472fc",
   "metadata": {},
   "source": [
    "## What are grids, threads and blocks (and warps)?\n",
    "\n",
    "[Threads and blocks](https://en.wikipedia.org/wiki/Thread_block_(CUDA_programming)) are how you instruct you GPU to process some code in parallel. Our GPU is a parallel processor, so we need to specify how many times we want our kernel to be executed.\n",
    "\n",
    "Threads have the benefit of havign some shared cache memory between them, but there are a limited number of cores on each GPU so we need to break our work down into blocks which will be run after another on the GPU.\n",
    "\n",
    "<figure>\n",
    "\n",
    "![CPU GPU Comparison](images/threads-blocks-warps.png)\n",
    "\n",
    "<figcaption style=\"text-align: center;\"> \n",
    "    \n",
    "Image source <a href=\"https://docs.nvidia.com/cuda/cuda-c-programming-guide/\">https://docs.nvidia.com/cuda/cuda-c-programming-guide/</a>\n",
    "    \n",
    "</figcaption>\n",
    "</figure>\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fad6a21f-ce9c-49a2-9b6a-41fbd29bf6d1",
   "metadata": {},
   "source": [
    "### What??\n",
    "\n",
    "Don't worry too much about this now. Just take away the idea that **we need to specify the number of times we want our kernel to be called**, and that is given as two numbers which are multiplied together to give your overall grid size.\n",
    "\n",
    "Rules of thumb for threads per block:\n",
    "- Should be a round multiple of the warp size (32)\n",
    "- A good place to start is 128-512 but benchmarking is required to determine the optimal value.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d36336db-8514-4b21-8072-62f085779f93",
   "metadata": {},
   "source": [
    "## Hello world"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e8dded05-6f39-4ea3-8ac6-45247ebc35d1",
   "metadata": {},
   "source": [
    "Let's dig in with some code and hopefully things will become more clear.\n",
    "\n",
    "To start off let's write a simple CPU based Python function which we will call repeatedly within a [list comprehension](https://docs.python.org/3/tutorial/datastructures.html#list-comprehensions). From a Python perspective list comprehensions can be a good jumping off point for parallel computing because they feel somewhat parallel already."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f4a31223-019a-403b-bbd9-c3daa2f1042f",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "data = range(10)\n",
    "\n",
    "def foo(i):\n",
    "    return i\n",
    "    \n",
    "[foo(i) for i in data]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9c420044-2d0a-4ca6-8643-82f96390ff2a",
   "metadata": {},
   "source": [
    "Here we have our `foo` function return its index value and use a for loop to iterate over our data which is generated by `range`.\n",
    "\n",
    "Next we will step by step convert this over to a CUDA kernel and run it on our GPU with numba CUDA."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "81a1d262-f928-4473-b6b1-450bcb9fe60b",
   "metadata": {},
   "source": [
    "First we need to remember that our kernel cannot return anything. Instead we will use an output list to store the values we would return."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ea8d3d9c-72fe-4ff7-b9f9-44e23dcf258e",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "data = range(10)\n",
    "output = []\n",
    "\n",
    "def foo(i):\n",
    "    output.append(i)\n",
    "    \n",
    "[foo(i) for i in data]\n",
    "\n",
    "output"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9d161843-ee8d-4cb0-bf02-481c61b8b4e0",
   "metadata": {},
   "source": [
    "Our next challenge is that our output array on our GPU must have a fixed length. We can't start off with an empty array and keep appending things. So let's use numpy to create a `ndarray` with the same length as our input data. We will also convert our input list to a `numpy` array too as that's what we can move to our GPU."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "099b300d-0f4d-4f05-bbc7-b6b0e888f226",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d6832e81-5a23-4af2-bdd1-c562a21e6d33",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "data = np.asarray(range(10))\n",
    "output = np.zeros(len(data))\n",
    "\n",
    "def foo(i):\n",
    "    output[i] = i\n",
    "    \n",
    "[foo(i) for i in data]\n",
    "\n",
    "output"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "15a42ff5-cc99-4889-acc5-443819664fdd",
   "metadata": {},
   "source": [
    "Now that our pure Python function behaves like a kernel let's use numba to convert it into one."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e60e9835-253f-4786-be3a-780f08df89da",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from numba import cuda"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "edd640a5-ceaf-4fc1-9d87-9f9bda5d1d87",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "data = np.asarray(range(10))\n",
    "output = np.zeros(len(data))\n",
    "\n",
    "@cuda.jit\n",
    "def foo(input_array, output_array):\n",
    "    i = cuda.grid(1)\n",
    "    output_array[i] = input_array[i]\n",
    "    \n",
    "foo[1, len(data)](data, output)\n",
    "\n",
    "output"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "758f7db8-9673-4b04-a7b3-2b8ef0e21907",
   "metadata": {},
   "source": [
    "**Woo the above code ran on our GPU!**\n",
    "\n",
    "Now let's unpack this a bit.\n",
    "\n",
    "To convert our CPU function into a GPU kernel we need to add the `@cuda.jit` decorator. This tells numba to compile our code down to CUDA compatible byte code at runtime.\n",
    "\n",
    "Next we changed our kernels inputs to `input_array` and `output_array`. This is because our kernel needs a reference to both arrays in order to interact with them. (More on this later).\n",
    "\n",
    "But what about `i`? Instead of passing our function the index each time we call it we can rely on a nice CUDA function called `cuda.grid` which allows our kernel to get it's own thread index while it is running.\n",
    "\n",
    "Lastly we make a funny looking function call `foo[blocks, threads](input, output)`. In order to run our kernel on the GPU in parallel we need to specify how many times we want it to run. Kernel functions are configured using square brackets and passing the block size and thread size. With our array only being `10` elements long we specify a blocksize of `1` and a thread size of `10` which means our kernel will be executed `10` times. Then we pass our arguments as normal."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f1d83f9b-c4c7-4b18-b83f-21c85cb0a5de",
   "metadata": {},
   "source": [
    "## Something a little bigger\n",
    "\n",
    "Now that we've run our first CUDA kernel with Numba let's try something a little bigger.\n",
    "\n",
    "This time we are going to take a large array and double every number in it. We will do it first in pure Python on the CPU and then in a CUDA kernel on the GPU.\n",
    "\n",
    "Let's start with a large `30m` long array of random numbers. And an output array of equal length."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0d06dc1c-b36e-48df-9c7f-0bf109c1ceae",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "random_array = np.random.random((30_000_000))\n",
    "random_array"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2791c50f-a80b-4ccc-86a1-5955c508dc78",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "output = np.zeros_like(random_array)\n",
    "output"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1d90ac1c-932a-4345-9e70-8a2e24f15c10",
   "metadata": {},
   "source": [
    "Then in Python let's iterate over this array and double each item into the output array. We can time the cell to see how long this takes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "195b97db-2efc-4e64-8f95-3ade92c52a07",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "def foo(i):\n",
    "    output[i] = random_array[i] * 2\n",
    "    \n",
    "[foo(i) for i in range(len(random_array))]\n",
    "\n",
    "output"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d881d061-b924-48b5-97d9-8ce5bc862e69",
   "metadata": {},
   "source": [
    "For me this takes around 10 seconds for the CPU to do this calculation."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "966920e3-653d-458c-9f9b-1e740e1db024",
   "metadata": {},
   "source": [
    "Next let's write a CUDA kernel which does exactly the same thing. The only difference to the previous example is that we set out thread size to a fixes value of `128` and then calculate how many blocks we need to cover the whole array."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c98f0967-ad98-420f-9ccb-1eae9e242bc7",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import math"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9723dc8a-9627-442f-becc-3d9fc5423096",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "output = np.zeros_like(random_array)\n",
    "\n",
    "threads = 128\n",
    "blocks = math.ceil(random_array.shape[0] / threads)\n",
    "\n",
    "@cuda.jit\n",
    "def foo(input_array, output_array):\n",
    "    i = cuda.grid(1)\n",
    "    output_array[i] = input_array[i] * 2\n",
    "    \n",
    "foo[blocks, threads](random_array, output)\n",
    "output"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "abec8daa-6d42-46cb-b95e-9f896de3a69d",
   "metadata": {},
   "source": [
    "Hooray this is now orders of magnitude faster, only taking a few hundred milliseconds."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f4043550-55eb-4a74-bdb7-32e97f58314e",
   "metadata": {},
   "source": [
    "The savvy among you may be wondering though, numpy is already a C based optimised library and we are comparing our GPU kernel with some pure Python code. What if we did it in numpy?\n",
    "\n",
    "Well you would be right, for this example numpy is still faster than our GPU."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ac1d6108-8d8f-4318-b7d8-f08cbf5c8b4d",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%time \n",
    "\n",
    "random_array * 2"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8b5a9d0e-775a-4871-8416-f81e232e728e",
   "metadata": {},
   "source": [
    "But the reason this is the case is because of memory management."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "383ba959-6e68-4f89-a887-8c139e847789",
   "metadata": {},
   "source": [
    "## Memory management\n",
    "\n",
    "Earlier we discussed that the CPU and GPU are effectively two separate computers. Each of these computers has it's own memory. \n",
    "\n",
    "All of the data we've worked with so far has been created with numpy on the CPU. So in order for `numba` to work with this data on the GPU it has been quietly copying data back and forth for us.\n",
    "\n",
    "This data movement comes at a penalty.\n",
    "\n",
    "We also have the option of being in control of the data ourselves. We can explicitly move our numpy arrays to the GPU ahead of time with `cuda.to_device`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a96de9eb-2adc-4179-89ce-5be1bdc2fb6b",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "gpu_random_array = cuda.to_device(random_array)\n",
    "gpu_output = cuda.to_device(np.zeros_like(random_array))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8dccdba5-b280-448e-9d26-98e0120b283d",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "gpu_random_array"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0d552a86-08da-4324-bfc4-1c726aca528d",
   "metadata": {},
   "source": [
    "Now if we run our kernel again and pass it our GPU memory arrays we see it does in fact out perform numpy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "651b7b29-4bad-4015-af25-1040f15b1946",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%timeit -n 100\n",
    "foo[blocks, threads](gpu_random_array, gpu_output)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "39df78d6-6ea7-4252-9bcd-203d0920a81b",
   "metadata": {},
   "source": [
    "However our output result is also still on the GPU. We explicitly copied it there, so we need to explicitly copy it back with `copy_to_host()`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6d1d93c1-f54d-4c4f-a789-237ff74409d9",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "gpu_output.copy_to_host()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b6dc8ac8-433c-46e0-a978-065364089e18",
   "metadata": {},
   "source": [
    "Both of these data movement operations take time. But the calculation we are performing here is trivial. As the calculation within our kernel gets more complex the percentage of time spent copying data becomes smaller and smaller.\n",
    "\n",
    "Memory management is useful in other places too. We may wish to write code where we write many kernels and chain them together. It would be inefficient to copy that data to the GPU and back again between each kernel call.\n",
    "\n",
    "```python\n",
    "# move array to GPU\n",
    "foo[blocks, threads](data, output)\n",
    "# move data back to CPU\n",
    "\n",
    "# move array to GPU\n",
    "bar[blocks, threads](data, output)\n",
    "# move data back to CPU\n",
    "\n",
    "# move array to GPU\n",
    "baz[blocks, threads](data, output)\n",
    "# move data back to CPU\n",
    "```\n",
    "\n",
    "So by explicitly putting our data there we can cut down on this time and be more in control of our computation.\n",
    "\n",
    "```python\n",
    "# move array to GPU\n",
    "data = cuda.to_device(data)\n",
    "output = cuda.to_device(output)\n",
    "\n",
    "foo[blocks, threads](data, output)\n",
    "bar[blocks, threads](data, output)\n",
    "baz[blocks, threads](data, output)\n",
    "\n",
    "# move data back to CPU\n",
    "data = data.copy_to_host()\n",
    "output = output.copy_to_host()\n",
    "```"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}