01b_fractal_numba.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "height = 2_000\n",
    "width = 3_000\n",
    "maxiterations = 20"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import numba\n",
    "import numba.cuda\n",
    "import math\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This time we can actually check the name of the device using the API:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "numba.cuda.get_current_device().name"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's make the data each time (we won't always use the output `fractal`)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def prepare(height, width):\n",
    "    x, y = np.ogrid[-1.5j : 1.5j : height * 1j, -2 : 2 : width * 1j]\n",
    "    c = x + y\n",
    "    fractal = np.zeros(c.shape, dtype=np.int32)\n",
    "    return c, fractal"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## NumPy\n",
    "\n",
    "Let's try a NumPy run (we will use `%%time` instead of `%%timeit`, since this takes several seconds to run so we don't need a precision measurement and don't want to waste time):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def fractal_numpy(c, maxiterations):\n",
    "    f = np.zeros_like(c, dtype=np.int32)\n",
    "    z = c.copy()\n",
    "\n",
    "    for i in range(1, maxiterations + 1):\n",
    "        z = z**2 + c  # Compute z\n",
    "        diverge = np.abs(z**2) > 2**2  # Divergence criteria\n",
    "\n",
    "        z[diverge] = 2  # Keep number size small\n",
    "        f[~diverge] = i  # Fill in non-diverged iteration number\n",
    "\n",
    "    return f"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "c, _ = prepare(height, width)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "_ = fractal_numpy(c, maxiterations)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Numba\n",
    "\n",
    "Let's do a quick check with Numba from the CPU course, just to see how fast we can get on single CPU:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@numba.vectorize([numba.int32(numba.complex128, numba.int32)])\n",
    "def fractal_numba_vectorize(cxy, maxiterations):\n",
    "    z = cxy\n",
    "    for i in range(maxiterations):\n",
    "        z = z**2 + cxy\n",
    "        if abs(z) > 2:\n",
    "            return i\n",
    "    return maxiterations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "c, _ = prepare(height, width)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit\n",
    "fractal_numba_vectorize(c, maxiterations)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Numba CUDA: vectorize, host memory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@numba.vectorize([numba.int32(numba.complex128, numba.int32)], target=\"cuda\")\n",
    "def fractal_cuda_vectorize(cxy, maxiterations):\n",
    "    z = cxy\n",
    "    for i in range(maxiterations):\n",
    "        z = z**2 + cxy\n",
    "        if abs(z) > 2:\n",
    "            return i\n",
    "    return maxiterations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "c, _ = prepare(height, width)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit\n",
    "fractal_cuda_vectorize(c, maxiterations);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Numba CUDA: vectorize, GPU memory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "c, _ = prepare(height, width)\n",
    "c = numba.cuda.to_device(c)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "c"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit\n",
    "fractal_cuda_vectorize(c, maxiterations)\n",
    "numba.cuda.synchronize()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that we did not copy the memory back to the CPU; it's still on the GPU."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Numba CUDA: vectorize, skip allocation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "c, f = prepare(height, width)\n",
    "c = numba.cuda.to_device(c)\n",
    "f = numba.cuda.to_device(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit\n",
    "fractal_cuda_vectorize(c, maxiterations, out=f)\n",
    "numba.cuda.synchronize()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Numba CUDA: custom kernel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@numba.cuda.jit\n",
    "def fractal_cuda_kernel(c_array, f, maxiterations):\n",
    "    x, y = numba.cuda.grid(2)\n",
    "    if x < c_array.shape[0] and y < c_array.shape[1]:\n",
    "        f[x, y] = 0\n",
    "        z = c_array[x, y]\n",
    "        for i in range(maxiterations):\n",
    "            z = z**2 + c_array[x, y]\n",
    "            if abs(z**2) > 4:\n",
    "                break\n",
    "            f[x, y] = i"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "c, f = prepare(height, width)\n",
    "c = numba.cuda.to_device(c)\n",
    "f = numba.cuda.to_device(f)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we have to specify a custom kernel launch, rather than having it automated."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "threadsperblock = (8, 8)\n",
    "blockspergrid = (\n",
    "    math.ceil(c.shape[0] / threadsperblock[0]),\n",
    "    math.ceil(c.shape[1] / threadsperblock[1]),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "blockspergrid"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit\n",
    "fractal_cuda_kernel[blockspergrid, threadsperblock](c, f, maxiterations)\n",
    "np.array(f, dtype=f.dtype)\n",
    "numba.cuda.synchronize()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can plot this, just in case we made a mistake (even though we ran timeit above, plotting is valid, since we are reusing the same preallocated memory location):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.imshow(np.array(f, dtype=f.dtype));"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "PyGPU Course 2019/12 [course/pygpu/default]",
   "language": "python",
   "name": "sys_pygpu201912"
  },
  "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.9.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}