07-coding-fast-and-fast-code.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<br><br><br><br><br>\n",
    "\n",
    "# Coding fast and fast code\n",
    "\n",
    "<br><br><br><br><br>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<center><img src=\"img/commute-by-plane.png\" width=\"65%\"></center>\n",
    "\n",
    "_(Stolen from Jake Vanderplas.)_"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<br><br><br><br><br>\n",
    "\n",
    "**Extended analogy:** drive your car to the airport, then take an airplane.\n",
    "\n",
    "<br><br><br><br><br>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<br><br>\n",
    "\n",
    "#### Typical analysis workflow\n",
    "\n",
    "   1. Start with small dataset; plot some things, compute some things, get a sense of the data.\n",
    "   2. Collect one-liners, notebook cells, scripts into an executable.\n",
    "   3. Run the executable on ever-larger datasets.\n",
    "   4. Run it in different ways to quantify systematic errors.\n",
    "   5. Realize that it wasn't written efficiently enough to finish analysis before deadline.\n",
    "   6. Rewrite it in a faster language (e.g. Python → C++).\n",
    "   7. Wonder why the new code doesn't reproduce the old results.\n",
    "\n",
    "The problem is **step 6**: frameworks should allow you to improve (cleanliness, execution time) in _small steps,_ not _major rewrites._\n",
    "\n",
    "<br><br>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Example problem: computing a fractal (calendar/mousepad/T-shirt...).\n",
    "import time, numpy\n",
    "\n",
    "def run_python(height, width, maxiterations=20):\n",
    "    y, x = numpy.ogrid[-1:0:height*1j, -1.5:0:width*1j]\n",
    "    c = x + y*1j\n",
    "    fractal = numpy.full(c.shape, maxiterations, dtype=numpy.int32)\n",
    "    for h in range(height):\n",
    "        for w in range(width):                  # for each pixel (h, w)...\n",
    "            z = c[h, w]\n",
    "            for i in range(maxiterations):      # iterate at most 20 times\n",
    "                z = z**2 + c[h, w]              # applying z → z² + c\n",
    "                if abs(z) > 2:                  # if it diverges (|z| > 2)\n",
    "                    fractal[h, w] = i           # color the plane with the iteration number\n",
    "                    break                       # we're done, no need to keep iterating\n",
    "    return fractal\n",
    "\n",
    "starttime = time.time()\n",
    "fractal = run_python(800, 1200)\n",
    "print(\"{0} ns per pixel\".format(1e9 * (time.time() - starttime) / (800 * 1200)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import matplotlib.pyplot\n",
    "fig, ax = matplotlib.pyplot.subplots(figsize=(10, 5.5))\n",
    "ax.imshow(fractal);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Numpy operates array-at-a-time, you have to completely rethink the \"iterate until converged\" part.\n",
    "\n",
    "def run_numpy(height, width, maxiterations=20):\n",
    "    y, x = numpy.ogrid[-1:0:height*1j, -1.5:0:width*1j]\n",
    "    c = x + y*1j\n",
    "    fractal = numpy.full(c.shape, maxiterations, dtype=numpy.int32)\n",
    "    z = c\n",
    "    for i in range(maxiterations):\n",
    "        z = z**2 + c                                            # applying z → z² + c\n",
    "        diverged = numpy.absolute(z) > 2                        # |z| > 2 is \"divergence\"\n",
    "        diverging_now = diverged & (fractal == maxiterations)   # some are already done\n",
    "        fractal[diverging_now] = i                              # just set the new ones\n",
    "        z[diverged] = 2                                         # clamp diverged at 2\n",
    "    return fractal\n",
    "\n",
    "starttime = time.time()\n",
    "fractal = run_numpy(1600, 2400)\n",
    "print(\"{0} ns per pixel\".format(1e9 * (time.time() - starttime) / (1600 * 2400)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# But once it's in a Numpy-friendly form, it's ready for GPUs.\n",
    "import cupy\n",
    "\n",
    "def run_cupy(height, width, maxiterations=20):\n",
    "    y, x = cupy.ogrid[-1:0:height*1j, -1.5:0:width*1j]               # <--- changed this line\n",
    "    c = x + y*1j\n",
    "    fractal = cupy.full(c.shape, maxiterations, dtype=numpy.int32)   # <--- changed this line\n",
    "    z = c\n",
    "    for i in range(maxiterations):\n",
    "        z = z**2 + c                                            # applying z → z² + c\n",
    "        diverged = numpy.absolute(z) > 2                        # |z| > 2 is \"divergence\"\n",
    "        diverging_now = diverged & (fractal == maxiterations)   # some are already done\n",
    "        fractal[diverging_now] = i                              # just set the new ones\n",
    "        z[diverged] = 2                                         # clamp diverged at 2\n",
    "    return fractal\n",
    "\n",
    "starttime = time.time()\n",
    "fractal = run_cupy(3200, 4800)\n",
    "print(\"{0} ns per pixel\".format(1e9 * (time.time() - starttime) / (3200 * 4800)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Dask can keep all of your processors busy, but it might be making a more complex graph than you want.\n",
    "import dask.array\n",
    "\n",
    "def run_dask(height, width, maxiterations=20, hchunks=3, vchunks=4):\n",
    "    chunked = lambda a: dask.array.from_array(a, chunks=(height // hchunks, width // vchunks))\n",
    "    y, x = numpy.ogrid[-1:0:height*1j, -1.5:0:width*1j]\n",
    "    c = chunked(x + y*1j)\n",
    "    fractal = chunked(numpy.full(c.shape, maxiterations, dtype=numpy.int32))\n",
    "    z = c\n",
    "    for i in range(maxiterations):\n",
    "        z = z**2 + c                                            # applying z → z² + c\n",
    "        diverged = numpy.absolute(z) > 2                        # |z| > 2 is \"divergence\"\n",
    "        diverging_now = diverged & (fractal == maxiterations)   # some are already done\n",
    "        fractal[diverging_now] = i                              # just set the new ones\n",
    "        z[diverged] = 2                                         # clamp diverged at 2\n",
    "    return fractal\n",
    "\n",
    "starttime = time.time()\n",
    "fractal = run_dask(1600, 2400).compute()\n",
    "print(\"{0} ns per pixel\".format(1e9 * (time.time() - starttime) / (1600 * 2400)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "run_dask(1600, 2400, maxiterations=1, hchunks=3, vchunks=4).visualize()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Numba, on the other hand, lets you run the original Python code, but faster.\n",
    "import numba\n",
    "\n",
    "@numba.jit\n",
    "def run_numba(height, width, maxiterations=20):\n",
    "    y, x = numpy.ogrid[-1:0:height*1j, -1.5:0:width*1j]\n",
    "    c = x + y*1j\n",
    "    fractal = numpy.full(c.shape, maxiterations, dtype=numpy.int32)\n",
    "    for h in range(height):\n",
    "        for w in range(width):                  # for each pixel (h, w)...\n",
    "            z = c[h, w]\n",
    "            for i in range(maxiterations):      # iterate at most 20 times\n",
    "                z = z**2 + c[h, w]              # applying z → z² + c\n",
    "                if abs(z) > 2:                  # if it diverges (|z| > 2)\n",
    "                    fractal[h, w] = i           # color the plane with the iteration number\n",
    "                    break                       # we're done, no need to keep iterating\n",
    "    return fractal\n",
    "\n",
    "starttime = time.time()\n",
    "fractal = run_numba(6400, 9600)\n",
    "print(\"{0} ns per pixel\".format(1e9 * (time.time() - starttime) / (6400 * 9600)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Numba also has an OpenMP-like, \"parallel for\" (prange) to parallelize without complex dependencies.\n",
    "import numba\n",
    "\n",
    "@numba.jit(parallel=True)\n",
    "def run_numba_parallel(height, width, maxiterations=20):\n",
    "    y, x = numpy.ogrid[-1:0:height*1j, -1.5:0:width*1j]\n",
    "    c = x + y*1j\n",
    "    fractal = numpy.full(c.shape, maxiterations, dtype=numpy.int32)\n",
    "    for h in numba.prange(height):\n",
    "        for w in numba.prange(width):           # for each pixel (h, w)...\n",
    "            z = c[h, w]\n",
    "            for i in range(maxiterations):      # iterate at most 20 times\n",
    "                z = z**2 + c[h, w]              # applying z → z² + c\n",
    "                if abs(z) > 2:                  # if it diverges (|z| > 2)\n",
    "                    fractal[h, w] = i           # color the plane with the iteration number\n",
    "                    break                       # we're done, no need to keep iterating\n",
    "    return fractal\n",
    "\n",
    "starttime = time.time()\n",
    "fractal = run_numba_parallel(6400, 9600)\n",
    "print(\"{0} ns per pixel\".format(1e9 * (time.time() - starttime) / (6400 * 9600)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Numba can compile to GPUs, too, but you have to identify the kernel and set block sizes.\n",
    "import numba.cuda, math\n",
    "def run_numba_cuda(height, width, maxiterations=20):\n",
    "    fractal = numpy.empty((height, width), dtype=numpy.int32)\n",
    "    griddim = (math.ceil(height / 32), math.ceil(width / 32))\n",
    "    blockdim = (32, 32)\n",
    "    inner_loop_cuda[griddim, blockdim](height, width, maxiterations, fractal)\n",
    "    return fractal\n",
    "@numba.cuda.jit\n",
    "def inner_loop_cuda(height, width, maxiterations, fractal):\n",
    "    x, y = numba.cuda.grid(2)                   # 2-dimensional CUDA grid\n",
    "    z = c = -1.5 + y*1.0/(height + 1) + -1j + x*1j*1.5/(width + 1)\n",
    "    fractal[x, y] = maxiterations\n",
    "    for i in range(maxiterations):      # iterate at most 20 times\n",
    "        z = z**2 + c                    # applying z → z² + c\n",
    "        if abs(z) > 2:                  # if it diverges (|z| > 2)\n",
    "            fractal[x, y] = i           # color the plane with the iteration number\n",
    "            break                       # we're done, no need to keep iterating\n",
    "starttime = time.time()\n",
    "fractal = run_numba_cuda(12800, 19200)\n",
    "print(\"{0} ns per pixel\".format(1e9 * (time.time() - starttime) / (12800 * 19200)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext Cython\n",
    "\n",
    "# Cython is a language that mixes C++ and Python. Can it help?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%cython --cplus -c-O3\n",
    "import cython, numpy       # load Python interface to Numpy\n",
    "cimport numpy              # load C++ interface to Numpy (types end in _t)\n",
    "\n",
    "@cython.boundscheck(False) # turn off bounds-checking\n",
    "@cython.wraparound(False)  # turn off negative index wrapping (e.g. -1 for last element)\n",
    "def run_cython(int height, int width, int maxiterations=20):\n",
    "    y, x = numpy.ogrid[-1:0:height*1j, -1.5:0:width*1j]\n",
    "    c = x + y*1j\n",
    "    fractal = numpy.zeros(c.shape, dtype=numpy.int32) + maxiterations\n",
    "    cdef numpy.ndarray[numpy.complex128_t, ndim=2, mode=\"c\"] c_array = c\n",
    "    cdef numpy.ndarray[numpy.int32_t,      ndim=2, mode=\"c\"] fractal_array = fractal\n",
    "    cdef numpy.complex128_t z\n",
    "    for h in range(height):\n",
    "        for w in range(width):\n",
    "            z = c_array[h, w]\n",
    "            for i in range(maxiterations):\n",
    "                z = z**2 + c_array[h, w]\n",
    "                if abs(z) > 2:\n",
    "                    fractal_array[h, w] = i\n",
    "                    break\n",
    "    return fractal"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Even with all the work we had to do, Cython's performance is underwhelming.\n",
    "\n",
    "starttime = time.time()\n",
    "fractal = run_cython(800, 1200)\n",
    "print(\"{0} ns per pixel\".format(1e9 * (time.time() - starttime) / (800 * 1200)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile run_pybind11.cpp\n",
    "#include <complex>\n",
    "#include <pybind11/pybind11.h>    // pybind11: another (better) way to compile C++ for Python.\n",
    "#include <pybind11/numpy.h>\n",
    "namespace py = pybind11;\n",
    "void run(int height, int width, int maxiterations,\n",
    "         py::array_t<std::complex<double>> np_c, py::array_t<int> np_fractal) {\n",
    "    std::complex<double>* c = static_cast<std::complex<double>*>(np_c.request().ptr);\n",
    "    int* fractal = static_cast<int*>(np_fractal.request().ptr);\n",
    "    for (int h = 0;  h < height;  h++) {\n",
    "        for (int w = 0;  w < width;  w++) {\n",
    "            std::complex<double> ci = c[h + height*w];\n",
    "            std::complex<double> z = ci;\n",
    "            for (int i = 0;  i < maxiterations;  i++) {\n",
    "                z = z*z + ci;\n",
    "                if (std::abs(z) > 2) {\n",
    "                    fractal[h + height*w] = i;\n",
    "                    break;\n",
    "}}}}}\n",
    "PYBIND11_MODULE(run_pybind11, m) {\n",
    "    m.def(\"run\", &run, \"the inner loop\");\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash\n",
    "\n",
    "# Compile it as a Python extension module.\n",
    "\n",
    "c++ -Wall -shared -std=c++11 -fPIC -O3                 \\\n",
    "    `python -m pybind11 --includes` run_pybind11.cpp   \\\n",
    "    -o run_pybind11`python3-config --extension-suffix`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import run_pybind11\n",
    "\n",
    "starttime = time.time()\n",
    "height, width, maxiterations = 3200, 4800, 20\n",
    "y, x = numpy.ogrid[-1:0:height*1j, -1.5:0:width*1j]\n",
    "c = x + y*1j\n",
    "fractal = numpy.full(c.shape, maxiterations, dtype=numpy.int32)\n",
    "run_pybind11.run(height, width, maxiterations, c, fractal)\n",
    "print(\"{0} ns per pixel\".format(1e9 * (time.time() - starttime) / (3200 * 4800)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile run_pybind11_fastmath.cpp\n",
    "#include <complex>\n",
    "#include <pybind11/pybind11.h>    // We'll be compiling with -ffast-math (so we need a new name)!\n",
    "#include <pybind11/numpy.h>\n",
    "namespace py = pybind11;\n",
    "void run(int height, int width, int maxiterations,\n",
    "         py::array_t<std::complex<double>> np_c, py::array_t<int> np_fractal) {\n",
    "    std::complex<double>* c = static_cast<std::complex<double>*>(np_c.request().ptr);\n",
    "    int* fractal = static_cast<int*>(np_fractal.request().ptr);\n",
    "    for (int h = 0;  h < height;  h++) {\n",
    "        for (int w = 0;  w < width;  w++) {\n",
    "            std::complex<double> ci = c[h + height*w];\n",
    "            std::complex<double> z = ci;\n",
    "            for (int i = 0;  i < maxiterations;  i++) {\n",
    "                z = z*z + ci;\n",
    "                if (std::abs(z) > 2) {\n",
    "                    fractal[h + height*w] = i;\n",
    "                    break;\n",
    "}}}}}\n",
    "PYBIND11_MODULE(run_pybind11_fastmath, m) {\n",
    "    m.def(\"run\", &run, \"the inner loop\");\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash\n",
    "\n",
    "# This time, compile with -ffast-math: trades floating point accuracy for speed.\n",
    "\n",
    "c++ -Wall -shared -std=c++11 -fPIC -O3 -ffast-math              \\\n",
    "    `python -m pybind11 --includes` run_pybind11_fastmath.cpp   \\\n",
    "    -o run_pybind11_fastmath`python3-config --extension-suffix`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import run_pybind11_fastmath\n",
    "\n",
    "starttime = time.time()\n",
    "height, width, maxiterations = 3200, 4800, 20\n",
    "y, x = numpy.ogrid[-1:0:height*1j, -1.5:0:width*1j]\n",
    "c = x + y*1j\n",
    "fractal = numpy.full(c.shape, maxiterations, dtype=numpy.int32)\n",
    "run_pybind11_fastmath.run(height, width, maxiterations, c, fractal)\n",
    "print(\"{0} ns per pixel\".format(1e9 * (time.time() - starttime) / (3200 * 4800)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Summary\n",
    "\n",
    "The times measured above are collected here, in order of my subjective estimate of setup difficulty.\n",
    "\n",
    "| Method            | Setup       | time (ns/px) | speedup | Cores    |\n",
    "|:------------------|:-----------:|-------------:|--------:|:--------:|\n",
    "| Python            | automatic   | 5588.5       | 1×      | 1        |\n",
    "| Numba             | automatic   | 111.4        | 50×     | 1        |\n",
    "| Numba-parallel    | easy        | 33.89        | 165×    | all (12) |\n",
    "| Numpy             | medium      | 359.6        | 15×     | 1        |\n",
    "| CuPy              | medium      | 72.8         | 77×     | GPU      |\n",
    "| Dask              | medium      | 214.8        | 26×     | all (12) |\n",
    "| Numba-CUDA        | difficult   | 6.95         | 800×    | GPU      |\n",
    "| pybind11          | challenging | 165.4        | 34×     | 1        |\n",
    "| pybind11-fastmath | challenging | 63.1         | 90×     | 1        |\n",
    "| Cython            | maddening   | 1501.8       | 3.7×    | 1        |\n",
    "\n",
    "Pure Python is \"automatic\" because it's what you start with, and the Numpy-based methods are \"medium\" because you have to rethink your problem. CUDA kernel management is another layer of difficulty, and translating into another language (C++ or Cython) is the biggest hurdle.\n",
    "\n",
    "The number of cores qualifies the speedup, and GPU methods rely on special hardware. Perfect scaling was never observed; independent, pinned, single-core processes could possibly do better."
   ]
  }
 ],
 "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.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}