3.2 Numba lab 2 (solution).ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "f95cdd35-30d2-405c-b92e-9b98d3848b19",
   "metadata": {},
   "source": [
    "# Black and white\n",
    "\n",
    "Now that we've had a look at multi-dimensional indexing why don't you try and use two-dimensional indexing to make our image black and white.\n",
    "\n",
    "Instead of operating over all pixels channel by channel we want to just operate over all pixels and average the channels out."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e67a1a0e-d8c5-481c-82c6-588a6f939bf6",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import matplotlib as mpl\n",
    "import matplotlib.pyplot as plt\n",
    "from numba import cuda\n",
    "import numpy as np\n",
    "import math\n",
    "\n",
    "plt.rcParams[\"figure.figsize\"] = (30,4)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a5424e22-3b39-4a01-9f2d-cd3e70693d8e",
   "metadata": {},
   "source": [
    "**1. Load our image with matplotlib.**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4658f0e0-732b-4741-860e-bacebf3520b7",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "!wget https://raw.githubusercontent.com/jacobtomlinson/gpu-python-tutorial/main/images/numba.png\n",
    "im = plt.imread(\"numba.png\")\n",
    "plt.imshow(im)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ae2f31f2-5d10-480a-87e8-224ea9c77634",
   "metadata": {},
   "source": [
    "**2. Move our image to the GPU and create an output array of the same size.**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "38a8b64b-d552-476c-ac97-8f0b4a3f9c11",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "gpu_im = cuda.to_device(im)\n",
    "gpu_output = cuda.to_device(np.zeros_like(gpu_im))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4e78dc11-5467-4a51-8e83-c5f5169a54a0",
   "metadata": {},
   "source": [
    "**3. Set our two-dimensional thead size and block size.** _Hint: Our `threadsperblock` should still multiply to `128`._"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c3b5749e-d2fc-4252-af7b-ed2440b9fafd",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "threadsperblock = (16, 16)\n",
    "blockspergrid_x = math.ceil(gpu_im.shape[0] / threadsperblock[0])\n",
    "blockspergrid_y = math.ceil(gpu_im.shape[1] / threadsperblock[1])\n",
    "blockspergrid = (blockspergrid_x, blockspergrid_y)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0d60af57-2d91-46a9-a9c4-6976712c7498",
   "metadata": {},
   "source": [
    "**4. Write our kernel.**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f27721dc-9e01-4461-80bb-681f08f8d81f",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "@cuda.jit\n",
    "def black_white(im, output):\n",
    "    # With our two-dimensional grid we can get our index position in two dimensions\n",
    "    x, y = cuda.grid(2)\n",
    "\n",
    "    # Because our grid is slightly larger than our image anything outside the image should be ignored\n",
    "    if x < im.shape[0] and y < im.shape[1]:\n",
    "\n",
    "        # Calculate the average across the RGB channels\n",
    "        average = (im[x, y, 0] + im[x, y, 1] + im[x, y, 2]) / 3\n",
    "\n",
    "        # Set all output RGB channels to the average\n",
    "        output[x, y, 0] = average\n",
    "        output[x, y, 1] = average\n",
    "        output[x, y, 2] = average\n",
    "\n",
    "        # Pass the alpha channel through\n",
    "        output[x, y, 3] = im[x, y, 3]\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8d4cb160-9fd9-4158-8916-50941396ac6d",
   "metadata": {},
   "source": [
    "**5. Run the kernel.**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b8951dce-73ef-41b7-97e3-97be08fe6ad4",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "black_white[blockspergrid, threadsperblock](gpu_im, gpu_output)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "02bc9ea1-4020-473a-9ed8-9b71e3d60822",
   "metadata": {},
   "source": [
    "**6. Move the data back from the GPU and plot it.**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "168c8971-e236-40eb-aea3-a673077fe066",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "plt.imshow(gpu_output.copy_to_host())\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "339118e0-c024-4830-95e6-58e163c25739",
   "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.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}