numba_jit_parallel.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "d22be65d-8a71-4c72-872a-2184d1a013f7",
   "metadata": {},
   "source": [
    "# Just-in-time compilation and parallelisation with Numba"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6b3ba435-7b48-4dd2-8316-71410026deca",
   "metadata": {},
   "source": [
    "## Python interpreter overheads"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2e8568f5-428f-41c8-99fa-a34b0bc384af",
   "metadata": {},
   "source": [
    "Python being an interpreted language, executes each line of source code one at a time in a sequential manner. An interpreter is a program that directly executes the instructions in the source code without requiring them to be compiled into an executable first. Interpreters tend to be more flexible than compilers, but are less efficient when running programs. The efficiency penalty arises because the interpreting process needs to be done every time by the python interpreter when the program is run. \n",
    "\n",
    "Due to these overheads of the Python interpreter at runtime, running a pure python program is generally slow which may not be suitable for large-scale numerical computing demands. Each mathematical operation usually has a considerable overhead arising from the Python interpreter, making especially time-critical operations inefficient."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "78e29b2b-db63-4754-9cc1-0432fa9b4a8d",
   "metadata": {},
   "source": [
    "## Compilation to machine code"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "879bee30-e5ee-4dd2-9196-c0c11fd0a19c",
   "metadata": {},
   "source": [
    "One approach to avoid the slowdown due to the interpreting program is to compile (i.e. translate) the user written code into machine-level instructions that can be directly executed by the CPU thereby circumventing the need for an interpreting process. In the domain of scientific computing, this paradigm is usually adopted to execute programs written in languages such as C, C++, Fortran. \n",
    "\n",
    "The process of translating the user code to machine code typically involves several steps and is performed by a compiler toolchain (a set of related programs) installed on the machine. Since compilers can view the entire source code upfront, they can perform a number of analyses and optimizations when generating machine instructions that speeds up execution than just interpreting each line individually.\n",
    "\n",
    "However, compiled languages are not usually flexible e.g. they typically do not have features like dynamic typing wherein the data type of each variable need not be known apriori. In general, it is accepted that languages that traditionally use interpreter runtimes are easier for new programmers to learn and prototype to quickly evaluate scientific concepts.  "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9d7b95c6-3524-4973-bacc-ca8d103e9535",
   "metadata": {},
   "source": [
    "## Just-in-time (JIT) compilation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "131e31ee-1e65-49d0-9e42-6c8324731c75",
   "metadata": {},
   "source": [
    "Since interpreters and compilers have complementary strengths and weaknesses, one possibility is to implement facilities to somehow combine aspects of both. A \"Just-in-time\" compilation is a process wherein only certain critical sub-sections of code are identified to be compiled at runtime. The rest of the user code shall continue to be evaluated by the python interpreter as usual. \n",
    "\n",
    "The critical sections of code that need to be run several times (e.g. compulational bottlenecks such as slow functions identified with the help of a profiler) can be instructed to be compiled into machine code."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e374cc22-48db-4690-82f5-462f777dda5f",
   "metadata": {},
   "source": [
    "## Parallel Computing"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8c2121f6-a546-4d6c-90f0-af76c6a39b22",
   "metadata": {},
   "source": [
    "Modern computers are highly parallel systems. We already covered SIMD parallelisation, where within each CPU cores there exist hardware vector units that allow the parallel execution of certain floating point operations. \n",
    "\n",
    "CPUs themselves consists of multiple CPU cores, and they can be leveraged i.e. the task to solve can be distributed across these different cores to further achieve speedup. However, it is worth remembering that the serial part of the code (i.e. code that needs to be run sequentially due to its fundamental nature of operations such as in time-stepping for Ordinary and Partial Differential Equations) shall dominate and the speed-up shall most likely be far below the number of CPU cores across which the program code was distributed across. \n",
    "\n",
    "We can also leverage General Purpose Graphics Processing Units (GPGPUs) as accelerator devices that contain highly parallel floating point units. If we scale to larger compute clusters then there is also a further level of parallelism between the individual hardware nodes."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "56ed2dd9-d785-46a8-90cf-9b15c763a02e",
   "metadata": {},
   "source": [
    "## The `numba` library"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6cfb67db-bdf2-4819-9143-a1a64cf9a82c",
   "metadata": {},
   "source": [
    "[`Numba`](https://numba.pydata.org/) is a third party accelerator library for Python that allows to just-in-time compile Python functions into fast, direct machine code that need not access the Python interpreter. Moreover, it allows to explore parallelisations that cover a lot of use-cases for parallel computing on a single machine using shared-memory parallelism. In addition to all this, Numba has features to directly cross-compile code for use on NVidia GPU accelerators.\n",
    "\n",
    "Depending on the mathematical operations within the function being targeted for compilation, its performance can be close to hand-optimized C code. It is cross-platform and bundles the `llvmlite` library (using the LLVM compiler toolchain) for providing cross-platform high-quality compilation of python functions.  "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2001118f-a4e0-4ffc-be9c-606d55ff1615",
   "metadata": {},
   "source": [
    "## Exercise: Accelerate large-scale vector addition using `numba`"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6f848070-4af7-41bb-90e8-6890c5eeeac1",
   "metadata": {},
   "source": [
    "Although `numba` provides an extensive set of facilities and features to tackle various compilation aspects, easy gains can usually be made by making use of the parallelisati. This is done by decorating critical functions of interest with the `@njit(parallel=True)` decorator. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a887acfd-beef-40de-82f1-6b2bb8e8a4b1",
   "metadata": {},
   "source": [
    "Let us accelerate the addition of two numpy arrays and see the difference in performance afforded by parallelisation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "f51941b1-3211-4910-a060-2419c0adce55",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import numba\n",
    "import time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "bae62aca-8aa3-40cf-9c36-7e014873a0ea",
   "metadata": {},
   "outputs": [],
   "source": [
    "n = 1000000                       # 1 million array elements\n",
    "rng = np.random.default_rng()     # Set up a random number generator\n",
    "a = rng.random(n)                 # n uniformly distributed floating point numbers in [0, 1]\n",
    "b = rng.standard_normal(n)        # n Gaussian distributed floating point numbers\n",
    "c = np.empty(n, dtype='float64')  # Placeholder for result vector"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "28d8658b-19fd-4284-8665-6b892aaee91e",
   "metadata": {},
   "source": [
    "First, we benchmark the numpy vector addition"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "e0225494-c16b-40c1-bbcb-f766937a20f4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Numpy vectorised operation on a single CPU core took 9.108473 ms\n"
     ]
    }
   ],
   "source": [
    "t1 = time.perf_counter_ns()\n",
    "c = a + b\n",
    "t2 = time.perf_counter_ns()\n",
    "print(f'Numpy vectorised operation on a single CPU core took {(t2-t1)/1e6} ms')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2e9e5d1a-891f-402a-9c59-b9c530819e35",
   "metadata": {},
   "source": [
    "Next, we refactor the vector addition operation into a function, and mark that function for jit compilation by decorating it suitably."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "1c549fc8-0184-4ed7-b296-f1124d5399eb",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Running the function (including compilation time) with parallel CPU cores took 819.015109 ms\n"
     ]
    }
   ],
   "source": [
    "@numba.njit(parallel=True)\n",
    "def numba_fun(in1, in2, out):\n",
    "    out = in1 + in2\n",
    "        \n",
    "# numba.set_num_threads(8)\n",
    "t1 = time.perf_counter_ns()\n",
    "numba_fun(a,b,c)\n",
    "t2 = time.perf_counter_ns()\n",
    "print(f'Running the function (including compilation time) with parallel CPU cores took {(t2-t1)/1e6} ms')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5dff9cb9-d967-43c2-8eb3-401fbeac1aaf",
   "metadata": {},
   "source": [
    "However, the performance seems to be rather poor.  This is because the compilation process incurs significant overhead. However, once compiled, evaluating the function is significantly faster. Hence, JIT compilation is most beneficial in cases where a time-consuming function is repeatedly evaluated (as in our `jacobi` function in the CFD code exercise). The speed-up of compiled function evaluation can be validated by repeatedly calling the function a few times."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "64f82c9a-4725-4415-b16a-7f07a7eaae43",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Evaluating the pre-compiled function with parallel CPU cores took 12.687 us\n",
      "Evaluating the pre-compiled function with parallel CPU cores took 3.233 us\n",
      "Evaluating the pre-compiled function with parallel CPU cores took 1.68 us\n",
      "Evaluating the pre-compiled function with parallel CPU cores took 1.43 us\n",
      "Evaluating the pre-compiled function with parallel CPU cores took 1.52 us\n"
     ]
    }
   ],
   "source": [
    "for i in range(5):\n",
    "    t1 = time.perf_counter_ns()\n",
    "    numba_fun(a,b,c)\n",
    "    t2 = time.perf_counter_ns()\n",
    "    print(f'Evaluating the pre-compiled function with parallel CPU cores took {(t2-t1)/1e3} us')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "817e5f5b-c99b-4054-9c0a-317729a176c1",
   "metadata": {},
   "source": [
    "As seen from the timing results, a significant speed-up was achieved by distributing the compiled function over multiple cores. The number of cores to use for parallelisation can optionally be set using `numba.set_num_threads(n)` before the `@njit` decorator line. Otherwise, numba shall use all available CPU cores, which may not be the most efficient when considering the scaling characteristics of the problem at hand."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1c22c04f-9cb7-404f-b617-364a54792164",
   "metadata": {},
   "source": [
    "## Exercise: Refactor `jacobi` function of CFD code to use `numba` JIT and autoparallelisation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "257cde62-c75d-4734-971c-5dae04aea872",
   "metadata": {},
   "source": [
    "In the `cfd_numba` directory, we have provided a version of the CFD code that uses numba to accelerate the computations. Upon inspecting the code, we see that the `numba` version of the CFD code differs mainly in the fact that the `jacobi` function is annotated to be just-in-time compiled and distributed across cores for parallel computation of results. Not all methods in the `numpy` library are compatible with `numba` though. In particular, we notice through experimentation that the `np.copyto` method cannot be used here, and hence the code provided reverts to the simpler `np.copy()` for assigning the stream function's updated values. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "07b08786-e173-49d9-8237-98106752449a",
   "metadata": {},
   "source": [
    "Navigate to the `cfd_numba` directory, and run the program with the same problem size as before, i.e. a 128 x 128 grid (scaling factor of 4) and 10000 Jacobi iterations:\n",
    "\n",
    "```bash\n",
    "prompt:/path/to/cfd_numba>python cfd.py 4 10000\n",
    "```\n",
    "\n",
    "The timing output of the program should unfortunatel result in a slowdown over the single-core `numpy` version e.g.\n",
    "\n",
    "```\n",
    "2D CFD Simulation\n",
    "=================\n",
    "Scale factor = 4\n",
    "Iterations   = 10000\n",
    "\n",
    "Initialisation took 0.00021s\n",
    "\n",
    "Grid size = 128 x 128\n",
    "\n",
    "Starting main Jacobi loop...\n",
    "\n",
    "...finished\n",
    "\n",
    "Calculation took 5.27128s\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "688aa874-8fe5-4a58-9ee9-93d06f50f63e",
   "metadata": {},
   "source": [
    "However, one key observation is that, we can now refine the grid size and/or crank up the number of Jacobi iterations to profit from the amortised cost of repeatedly invoking a pre-compiled function. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6c9e8d36-faf7-404c-a9ba-ccf4f76f0125",
   "metadata": {},
   "source": [
    "## Exercise: Compare the `numba` and `numpy` versions of CFD code on a 512 x 512 grid (scale factor of 16) and with 10000 Jacobi iterations"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d47ff01a-bd8a-4d99-b164-62aae5491d23",
   "metadata": {},
   "source": [
    "Sample timing output from `numba` parallel version:"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f2fb16d1-b8bb-48ce-9126-81f4c8cb043e",
   "metadata": {},
   "source": [
    "```\n",
    "2D CFD Simulation\n",
    "=================\n",
    "Scale factor = 16\n",
    "Iterations   = 10000\n",
    "\n",
    "Initialisation took 0.00142s\n",
    "\n",
    "Grid size = 512 x 512\n",
    "\n",
    "Starting main Jacobi loop...\n",
    "\n",
    "...finished\n",
    "\n",
    "Calculation took 23.12356s\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "423578ed-2357-4db3-9e60-86233633f2b7",
   "metadata": {},
   "source": [
    "Sample timing output from `numpy` serial version:"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cc430b26-905b-4095-afc4-f1177e55f54f",
   "metadata": {},
   "source": [
    "```\n",
    "2D CFD Simulation\n",
    "=================\n",
    "Scale factor = 16\n",
    "Iterations   = 10000\n",
    "\n",
    "Initialisation took 0.00133s\n",
    "\n",
    "Grid size = 512 x 512\n",
    "\n",
    "Starting main Jacobi loop...\n",
    "\n",
    "...finished\n",
    "\n",
    "Calculation took 49.91595s\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "79674fb8-62b6-44c2-8167-1c3589c78a44",
   "metadata": {},
   "source": [
    "We get a speed up of approx 2. Experiment to see if there are further improvements in speed-up using a finer grid and/or by increasing the number of Jacobi iterations"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e22c0e8f-2176-4c69-94be-a91c100ab364",
   "metadata": {},
   "source": [
    "## Concluding remarks"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d757c83e-a7a1-4ae9-a2fa-878011ec0cd0",
   "metadata": {},
   "source": [
    "Just-in-time compilation is not a guaranteed panacea for accelerating scientific codes. The potential speed-up that can be achieved depends on a number of factors such as the nature of floating point operations used, the library methods used within the function to be compiled etc. This requires further investigation on the part of the user to determine a reasonable trade-off between pace of development vs runtime performance.\n",
    "\n",
    "There exist various other approaches to accelerate scientific python codes. This includes distributing computations over many computational nodes using the MPI protocol and `mpi4py` library, performing distributed task/data parallelism/pipelining using the `dask` library, exploring concurrency and multithreading using `joblib`, offloading computationally intensive functions as kernels to dedicated accelerator hardware such as GPUs and FPGAs using libraries such as `dace` etc.     "
   ]
  }
 ],
 "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.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}