3 - Parallel features.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Numba's Parallel Features\n",
    "=============\n",
    "\n",
    "Numba supports a variety of different ways to use threads to do parallel computation.  Let's start with the easiest method, and parallelize a ufunc.\n",
    "\n",
    "@vectorize\n",
    "----------\n",
    "\n",
    "NumPy is built on Universal Functions (\"ufuncs\") as a key concept, these functions perform their operations element wise on the input and handle all the broadcasting over arrays etc. What if you want to define your own ufunc? Meet Numba's `@vectorize`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from numba import vectorize, jit\n",
    "\n",
    "@vectorize\n",
    "def ReynoldsNo(rho, u, L, mu):\n",
    "    return (rho * u * L) / mu\n",
    "\n",
    "rho = 1000 # density of pure water\n",
    "speed = 2 # human swimming\n",
    "length = 1.8 # human height\n",
    "dyn_visc = 1e-3 # dynamic viscosity of water at about 20 C\n",
    "\n",
    "ReynoldsNo(rho, speed, length, dyn_visc)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "How fast do you need to swim to match the Reynold number of a blue whale?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "u = np.arange(0, 1000, 10)\n",
    "re = ReynoldsNo(rho, u, length, dyn_visc)\n",
    "\n",
    "blue_whale = 4e8\n",
    "u[np.argmin(np.abs(re - blue_whale))]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "By default all `@vectorize` ufuncs are \"Dynamic ufuncs\", which means they will recompile themselves on encountering a new type. Express \"length\" in the complex domain just as an example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ReynoldsNo(rho, speed, 2j, dyn_visc)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Check the type signatures cf. NumPy:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ReynoldsNo.types"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Switching the target backend...\n",
    "\n",
    "The vectorize decorator was just targetting a single core of a CPU, this is the default target. Other targets of interest include:\n",
    " * \"cpu\" - serial CPU execution\n",
    " * \"parallel\" - parallel CPU execution\n",
    " * \"CUDA\" - NVIDIA CUDA\n",
    " \n",
    "Let's try out the CPU targets with the classic form of solving:\n",
    "\n",
    "$ a x^2 + b x + c = 0 $\n",
    "\n",
    "Using the formula for returning the one of the solutions:\n",
    "\n",
    "$ x = \\frac{-b - \\sqrt{b^2 - 4ac}}{2a} $"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# some inputs\n",
    "from math import sqrt\n",
    "n = 100000\n",
    "a = np.random.random((n,))\n",
    "b = 4 + np.random.random((n,))\n",
    "c = np.random.random((n,))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<h3><span style=\"color:blue\"> Task 1a: Implement the above equation as a `@vectorize` function</span></h3>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# <Add the decorator here>\n",
    "def my_solver(a, b, c):\n",
    "    # You write this!\n",
    "\n",
    "# check it's correct\n",
    "assert my_solver(3, 10, 3) == -3"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<h3><span style=\"color:blue\"> Task 1b: Measure the performance of your implementation</span></h3>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%timeit my_solver(a, b, c)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Some targets require signatures still (this is a hang over from older code paths), a signature is a way to spell out the the compiler the permitted types of the inputs (and outputs). In the case above, there's three \"float64\" scalar inputs and a \"float64\" scalar output. If you ever aren't sure of the type of something, use `numba.typeof()` and it'll tell you. The types can be spelled as strings or as instances of `numba.types` classes. It's these classes that you'll see in the type inference system/error message feedback.\n",
    "\n",
    "<h3><span style=\"color:blue\"> Task 1c: Use your function body from above in the stub below</span></h3>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sig = \"float64(float64, float64, float64,)\"\n",
    "@vectorize([sig,], target=\"parallel\")\n",
    "def my_solver_parallel(a, b, c):\n",
    "    # <copy your solver function body here>\n",
    "    return x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<h3><span style=\"color:blue\"> Task 1d: Check the performance of the function on the \"parallel\" target</span></h3>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%timeit my_solver_parallel(a, b, c)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "@guvectorize\n",
    "----------\n",
    "\n",
    "Numba extends the ufunc concept to \"Generalised ufuncs\" these are indeed a generalisation on \"ufuncs\" in that they allow arrays as inputs but still handle all the broadcasting etc. Things to note:\n",
    " * You have to spell out the type signature(s)\n",
    " * You have to supply an input/output layout (this is so the compiler knows how to generate the inner loop nest).\n",
    " * The function has a `void` return type, the layout specification is used to determine which arguments are inputs/outputs.\n",
    " \n",
    "Using matrix multiplication solely as a pedagogical example (naturally, you'd use `np.dot`!):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from numba import guvectorize\n",
    "\n",
    "@guvectorize([\"float64[:,:], float64[:,:], float64[:,:]\",], '(m, k),(k, n)->(m, n)')\n",
    "def matmul(x, y, res):\n",
    "    m, k = x.shape\n",
    "    k2, n = y.shape\n",
    "    out_m, out_n = res.shape\n",
    "    assert k == k2\n",
    "    assert m == out_m\n",
    "    assert n == out_n\n",
    "    for p in range(m):\n",
    "        for q in range(n):\n",
    "            res[p, q] = 0\n",
    "            for r in range(k):\n",
    "                res[p, q] += x[p, r] * y[r, q]\n",
    "\n",
    "\n",
    "m, n, k = 3, 4, 5\n",
    "a = np.arange(m * k).reshape(m, k)\n",
    "b = np.arange(k * n).reshape(k, n)\n",
    "out = np.zeros((m, n))\n",
    "\n",
    "expected = np.dot(a, b)\n",
    "\n",
    "matmul(a, b, out)\n",
    "np.testing.assert_allclose(out, expected)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It should be noted that the layout specification is the layout of the data needed for the kernel (the function body) to operate, **not** the layout of the arrays supplied as input. The `@guvectorize` decorator will do all the broadcasting necessary to use lower dimensioned kernels on higher dimensioned data. For example, using the `matmul` function above on some 3D input, note how the kernel is broadcast across the two inner most dimensions (as determined by the layout specification)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "outer, m, n, k = 7, 3, 4, 5\n",
    "a = np.arange(outer * m * k).reshape(outer, m, k)\n",
    "b = np.arange(outer * k * n).reshape(outer, k, n)\n",
    "out = np.zeros((outer, m, n))\n",
    "expected = np.zeros((outer, m, n))\n",
    "\n",
    "# np.dot has different broadcasting rules to the kernel above\n",
    "for x in range(outer):\n",
    "    expected[x] = np.dot(a[x], b[x])\n",
    "\n",
    "matmul(a, b, out)\n",
    "\n",
    "np.testing.assert_allclose(out, expected)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `@guvectorize` decorator supports the same targets as the `@vectorize` decorator."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "@jit parallelism\n",
    "-------------\n",
    "We've seen that Numba has a parallel target accessible, this is also available for the `@jit` family and is switched on via the kwarg `parallel=True`. Revisit `awkward_sine` from the first notebook..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n = 10000\n",
    "\n",
    "def awkward_sine(n):\n",
    "    a = np.ones(n)\n",
    "    return np.imag(np.exp(1j * a))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Check the performance of the NumPy based function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%timeit awkward_sine(n)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "JIT compile the function and check the performance of the JIT compiled version:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "jit_awkward_sine = jit()(awkward_sine)\n",
    "\n",
    "%timeit jit_awkward_sine(n)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "JIT compile the function with `parallel=True` set and check the performance of this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "parallel_jit_awkward_sine = jit(parallel=True)(awkward_sine)\n",
    "\n",
    "%timeit parallel_jit_awkward_sine(n)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Numba has some parallel diagnostic output available, it's accessible via the `.parallel_diagnostics()` method on the dispatcher [docs](http://numba.pydata.org/numba-doc/latest/user/parallel.html#diagnostics). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "parallel_jit_awkward_sine.parallel_diagnostics(level=4)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The above is implicit parallelism, array analysis is being done to work out what can be executed safely in parallel. Explicit parallelism is also possible using a `numba.prange` driven loop. The `prange` function is a like `range` but will cause the loop body to execute in parallel if the `parallel=True` target is in use (else it just behaves like normal `range`).\n",
    "\n",
    "\n",
    "<h3><span style=\"color:blue\"> Task 2a: Compute the 2-norm of this vector using a `prange` loop with an accumulator</span></h3>\n",
    "\n",
    "For reference the 2-norm of vector $x$ is:\n",
    "\n",
    "$|x|_2 = \\sqrt{\\sum_{i=0}^{n} {x_i}^2}$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from numba import prange\n",
    "\n",
    "@jit(parallel=True)\n",
    "def two_norm(x):\n",
    "    acc = 0.0\n",
    "    for i in prange(x.shape[0]):\n",
    "        # You write the loop body\n",
    "\n",
    "    return np.sqrt(acc)\n",
    "\n",
    "x = np.arange(1000000.)\n",
    "\n",
    "np.testing.assert_allclose(two_norm(x), np.linalg.norm(x))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<h3><span style=\"color:blue\"> Task 2b: Investigate the performance</span></h3>\n",
    "\n",
    "Some `timeit` helpers, first the python function, second the `np.linalg.norm` implementation from NumPy and last, the parallel implementation above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%timeit two_norm.py_func(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%timeit np.linalg.norm(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%timeit two_norm(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<h3><span style=\"color:blue\"> Task 2c: Investigate how performance changes with input size</span></h3>\n",
    "\n",
    "Launching a parallel calculation in Numba has some fixed overhead.   Go back and change the size of the array `x` (the argument to `np.arange` selects how many elements x has) and see how the relative speed of NumPy vs. parallel Numba changes.  Try arrays of size:\n",
    "* 100\n",
    "* 50000000 (don't run `two_norm.py_func(x)` for this case or you will wait a while)\n",
    "\n",
    "Was NumPy or parallel Numba faster in each case?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}