fields_intro_to_GPUs_solutions.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/carlnotsagan/LSST-DSFP-Session15-Materials/blob/main/fields_intro_to_GPUs_solutions.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "id": "quHFe1dTJVan"
      },
      "outputs": [],
      "source": [
        "from __future__ import print_function, division, absolute_import "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yuAyY7U9JVbK"
      },
      "source": [
        "# Introduction to GPUs (in Python):\n",
        "\n",
        "By Carl Fields (Los Alamos National Lab)\n",
        "\n",
        "*This exercise was designed heavily based on tutorials from the [GTC 2018 Conference](https://github.com/ContinuumIO/gtc2018-numba).* \n",
        "\n",
        "---"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NHPCMLpBJVbd"
      },
      "source": [
        "## Problem 0) Creating Our HPC Environment\n",
        "**Before** beginning, we want to prepare a new environment for the purpose of this exercise.\n",
        "\n",
        "Be sure to activate our HPC environment from last time:\n",
        "\n",
        "```linux\n",
        "$ conda activate hpc\n",
        "```\n",
        "---\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Mb4EgY5YJVbb"
      },
      "source": [
        "## Learning Objectives\n",
        "\n",
        "- Using Numba decorators to speed up algorithms targeting GPUs\n",
        "- Creating complex alogrithms utilizing Numba decorators targeting GPUs\n",
        "- Memory management with _CUDA_ kernels\n",
        "- Exploring General _CUDA_ kernels\n",
        "---"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hMU9pYvQJVbq"
      },
      "source": [
        "We can check that our installation worked by trying to import the required libraries:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "id": "sPild-HfJVbs"
      },
      "outputs": [],
      "source": [
        "import numpy as np\n",
        "import math\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "%matplotlib notebook"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dRApbp_zJVbu"
      },
      "source": [
        "## Problem 1) Exploring a basic algorithm targeting the GPU\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "w1GFBEnNJVb8"
      },
      "source": [
        "We want to begin by compaaring a native numpy function which will leverage the CPU and compare that to a Ufunc that will target the CPU. \n",
        "\n",
        "Lets start by considering the addition of two numbers."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "T4UTxhKsJVb_"
      },
      "source": [
        "**Problem 1a** Run the following cell to compute the addition of two arrays:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "DqcpLS9BJVcM",
        "outputId": "3522b579-b690-468a-ecde-b8325898be7a"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "array([11, 22, 33, 44])"
            ]
          },
          "metadata": {},
          "execution_count": 3
        }
      ],
      "source": [
        "import numpy as np\n",
        "\n",
        "a = np.array([1, 2, 3, 4])\n",
        "b = np.array([10, 20, 30, 40])\n",
        "\n",
        "np.add(a, b)"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "More information on Numpy Ufuncs can be found [here](https://docs.scipy.org/doc/numpy/reference/ufuncs.html)."
      ],
      "metadata": {
        "id": "P_KDgS8TLIIE"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Next, we want to explore using Numba to create Ufuncs that target the GPU."
      ],
      "metadata": {
        "id": "HigqTW-4LTRB"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "**Problem 1b** Use the `vectorize` decorator in Numba to write a function that adds to arrays. Use the `int64` data types and `target='cuda'`. Note: An overview including some common terminology for CUDA programming can be found [here](https://numba.readthedocs.io/en/stable/cuda/overview.html)."
      ],
      "metadata": {
        "id": "ZL_tQqRILk8_"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from numba import vectorize\n",
        "\n",
        "@vectorize(['int64(int64, int64)'], target='cuda')\n",
        "def add_ufunc(x, y):\n",
        "    return x + y"
      ],
      "metadata": {
        "id": "_iSVUWE-MeFD"
      },
      "execution_count": 4,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Before running we need to be sure we have the hardware we would like to target!\n",
        "\n",
        "First, you'll need to enable GPUs for the notebook:\n",
        "\n",
        "- Navigate to Edit→Notebook Settings\n",
        "- select GPU from the Hardware Accelerator drop-down\n",
        "\n",
        "Check that the GPU was found:"
      ],
      "metadata": {
        "id": "byv3W6puNqAC"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "%tensorflow_version 2.x\n",
        "import tensorflow as tf\n",
        "device_name = tf.test.gpu_device_name()\n",
        "if device_name != '/device:GPU:0':\n",
        "  raise SystemError('GPU device not found')\n",
        "print('Found GPU at: {}'.format(device_name))"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "B240Pi_UN34I",
        "outputId": "3909ef1e-7c9e-4c14-d825-86268096c814"
      },
      "execution_count": 5,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Found GPU at: /device:GPU:0\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "More information can be found about the available types of devices:"
      ],
      "metadata": {
        "id": "mu9dqxDfOruh"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from tensorflow.python.client import device_lib\n",
        "device_lib.list_local_devices()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Hk8mU8y7OXUh",
        "outputId": "feadab89-6871-4655-e3d9-3580ed14279f"
      },
      "execution_count": 6,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "[name: \"/device:CPU:0\"\n",
              " device_type: \"CPU\"\n",
              " memory_limit: 268435456\n",
              " locality {\n",
              " }\n",
              " incarnation: 9588685152280568412\n",
              " xla_global_id: -1, name: \"/device:GPU:0\"\n",
              " device_type: \"GPU\"\n",
              " memory_limit: 14444920832\n",
              " locality {\n",
              "   bus_id: 1\n",
              "   links {\n",
              "   }\n",
              " }\n",
              " incarnation: 1214201806784238051\n",
              " physical_device_desc: \"device: 0, name: Tesla T4, pci bus id: 0000:00:04.0, compute capability: 7.5\"\n",
              " xla_global_id: 416903419]"
            ]
          },
          "metadata": {},
          "execution_count": 6
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "**Problem 1c** Run your Ufunc which utilizes the GPU and compare the numerical result to Numpy."
      ],
      "metadata": {
        "id": "s39l-vQ2MpBz"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "print('a+b:\\n', add_ufunc(a, b))\n",
        "assert np.allclose(np.add(a, b),add_ufunc(a, b))"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "yfoLLhlkMoH8",
        "outputId": "8f42c5ad-249d-49a3-ec81-a6b25397fd42"
      },
      "execution_count": 7,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "a+b:\n",
            " [11 22 33 44]\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Now, lets use our favorite `timeit` magic command to see how much we benefitted from targeting the GPU.\n"
      ],
      "metadata": {
        "id": "ISBs9a8IPbN-"
      }
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zHdNjXQKJVcP"
      },
      "source": [
        "**Problem 1d** Use `timeit` to compare the execution time of the default Numpy Ufunc and our new Numba Ufunc"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {
        "scrolled": true,
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Ty9hOhdFJVcc",
        "outputId": "17a964d0-bdca-4a38-bf65-269fbabc15e9"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "The slowest run took 79.32 times longer than the fastest. This could mean that an intermediate result is being cached.\n",
            "1000000 loops, best of 5: 506 ns per loop\n"
          ]
        }
      ],
      "source": [
        "%timeit np.add(a,b)   # NumPy on CPU"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "%timeit add_ufunc(a,b) # Numba on GPU"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "PHyk6SEnP2LG",
        "outputId": "4b63f923-02d0-4c1d-804b-1dd4f4756f68"
      },
      "execution_count": 9,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "1000 loops, best of 5: 1.19 ms per loop\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "The GPU result is... *slower*?\n",
        "\n",
        "**Problem 1e** Discuss in this situation why our GPU result may be slower and how we can modify the problem to benfit fromm the GPU."
      ],
      "metadata": {
        "id": "it4I4hqrQQfw"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "**Answer** \n",
        "\n",
        "Some points to consider: \n",
        "\n",
        " - **Our inputs are too small**: the GPU achieves performance through parallelism, operating on thousands of values at once. Our test inputs have only 4 and 16 integers, respectively. We need a much larger array to even keep the GPU busy.\n",
        "\n",
        "\n",
        "- **Our calculation is too simple**: Sending a calculation to the GPU involves quite a bit of overhead compared to calling a function on the CPU. If our calculation does not involve enough math operations (often called \"arithmetic intensity\"), then the GPU will spend most of its time waiting for data to move around.\n",
        "\n",
        "\n",
        "- **We copy the data to and from the GPU**: While including the copy time can be realistic for a single function, often we want to run several GPU operations in sequence. In those cases, it makes sense to send data to the GPU and keep it there until all of our processing is complete.\n",
        "\n",
        "\n",
        "- **Our data types are larger than necessary**: Our example uses int64 when we probably don't need it. "
      ],
      "metadata": {
        "id": "8j-yapHRQiVU"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Problem 2) Exploring a more complex, data-intensive algorithm targeting the GPU"
      ],
      "metadata": {
        "id": "q8Pip2rbRoEj"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "We will now consider a more complex example where we can take some of the necessary steps to make the problem more efficient to run on GPUs. A few of these steps include: \n",
        "\n",
        "1. Using native math module functions described [here](https://docs.python.org/3/library/math.html).\n",
        "2. Using a less precise datatype than necessary. Consider using `float32` instead. \n",
        "3. Solving a more complex algorithm - one with more math operations in this case than the addition of two arrays. \n",
        "4. Precompute constant values when possible.\n"
      ],
      "metadata": {
        "id": "2iy_vIiXRzA7"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "**Problem 2a** Take the above steps to define the Ufunc for a Gaussian PDF using Numba `vectorize` again targeting `cuda`: \n",
        "\n",
        " $f(x) = \\frac{1}{\\sigma \\sqrt{2\\pi}} e^{-\\frac{1}{2} \\left ( \\frac{x-\\mu}{\\sigma} \\right )^{2}}$.\n",
        "\n",
        " Information on solving Normal distributions are discussed [here](https://en.wikipedia.org/wiki/Normal_distribution).[link text](https://)"
      ],
      "metadata": {
        "id": "olNN1JTwTFNi"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import math  # Note that for the CUDA target, we need to use the scalar functions from the math module, not NumPy\n",
        "\n",
        "SQRT_2PI = np.float32((2*math.pi)**0.5)  # Precompute this constant as a float32.  Numba will inline it at compile time.\n",
        "\n",
        "@vectorize(['float32(float32, float32, float32)'], target='cuda')\n",
        "def gaussian_pdf(x, mean, sigma):\n",
        "    '''Compute the value of a Gaussian probability density function at x with given mean and sigma.'''\n",
        "    return math.exp(-0.5 * ((x - mean) / sigma)**2) / (sigma * SQRT_2PI)"
      ],
      "metadata": {
        "id": "VYyEI4LrRx5L"
      },
      "execution_count": 10,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "**Problem 2b** Evaluate our Ufunc a million times, set $\\mu=0$ and $\\sigma=1$. Use `np.random.uniform` to create our `x` array for a bound of [-3,3].: "
      ],
      "metadata": {
        "id": "p8uM5sPSUR0R"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Evaluate the Gaussian a million times!\n",
        "x = np.random.uniform(-3, 3, size=1000000).astype(np.float32)\n",
        "mean = np.float32(0.0)\n",
        "sigma = np.float32(1.0)\n",
        "\n",
        "# Quick test\n",
        "gaussian_pdf(x[0], 0.0, 1.0)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "vxNO39mtUSfx",
        "outputId": "0fb4d2b3-1f95-4bfb-fbf8-cd4d88a9bbda"
      },
      "execution_count": 11,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "array([0.30788687], dtype=float32)"
            ]
          },
          "metadata": {},
          "execution_count": 11
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "**Problem 2c** Perform the same calculation using scipys native `norm` function (details [here](https://docs.scipy.org/doc/scipy/reference/generated/scipy.stats.norm.html)). Time the results using `timeit`."
      ],
      "metadata": {
        "id": "MuyzndFdU5sZ"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import scipy.stats # for definition of gaussian distribution\n",
        "norm_pdf = scipy.stats.norm\n",
        "%timeit norm_pdf.pdf(x, loc=mean, scale=sigma)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "JuTYq6gjUz2I",
        "outputId": "7bd21fc3-0fde-43e4-f771-42d6d7b9d8fc"
      },
      "execution_count": 12,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "10 loops, best of 5: 36.2 ms per loop\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "**Problem 2d** Time the result for our GPU Ufunc for comparison:"
      ],
      "metadata": {
        "id": "B7JZNpAwVVpE"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "%timeit gaussian_pdf(x, mean, sigma)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "0J0hlysUVWN6",
        "outputId": "46747317-a0e3-49e7-93cf-9218a3263e57"
      },
      "execution_count": 13,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "The slowest run took 21.22 times longer than the fastest. This could mean that an intermediate result is being cached.\n",
            "100 loops, best of 5: 4.87 ms per loop\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "**Problem 2e** Thats a big improvement! \n",
        "\n",
        "Discuss some of the overhead costs still associated with this approach.\n",
        "\n",
        "**Answer**: Copying data to and from the GPU."
      ],
      "metadata": {
        "id": "-0FrZT5VVlgG"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Problem 3) Memory Management "
      ],
      "metadata": {
        "id": "_iOIEzs_Wgsj"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Some key terms before proceeding:\n",
        "\n",
        "- host: the CPU\n",
        "\n",
        "- device: the GPU\n",
        "\n",
        "- host memory: the system main memory\n",
        "\n",
        "- device memory: onboard memory on a GPU card\n",
        "\n",
        "- kernels: a GPU function launched by the host and executed on the device\n",
        "\n",
        "- device function: a GPU function executed on the device which can only be called from the device (i.e. from a kernel or another device function)\n",
        "\n",
        "In this problem we are going to explore explicit memory management as a way to further increase the speed of our algorithms. "
      ],
      "metadata": {
        "id": "rsQax-fPaWwL"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "**Problem 3a** Use the `vectorize` decorator in Numba to write a function that adds two arrays. \n",
        "\n",
        "Use the `float32` data types and `target='cuda'`. \n",
        "\n",
        "This will serve as the basis of our timing analysis before optimization. "
      ],
      "metadata": {
        "id": "Ep5xoWkQdY_H"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "@vectorize(['float32(float32, float32)'], target='cuda')\n",
        "def add_ufunc(x, y):\n",
        "    return x + y"
      ],
      "metadata": {
        "id": "HM5LkRnaaRx8"
      },
      "execution_count": 14,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "**Problem 3b** Construct an array `x` of a range of 100000 values. Create an array `y=2*x`."
      ],
      "metadata": {
        "id": "WkqOJZFIaVcr"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "n = 100000\n",
        "x = np.arange(n).astype(np.float32)\n",
        "y = 2 * x"
      ],
      "metadata": {
        "id": "7U3C1Qg1aYAN"
      },
      "execution_count": 15,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "**Problem 3c** Use these arrays to run and time (using `timeit` our baseline CUDA function)"
      ],
      "metadata": {
        "id": "NLngR4PzacsD"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "%timeit add_ufunc(x, y)  # Baseline performance with host arrays"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "oi63pz1vaddC",
        "outputId": "cf784c7d-25ca-4445-879c-9f5243878f1e"
      },
      "execution_count": 16,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "The slowest run took 50.25 times longer than the fastest. This could mean that an intermediate result is being cached.\n",
            "1000 loops, best of 5: 1.72 ms per loop\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Next, we want to utilize native Numba tools to copy the arrays we created _to device_ (to the GPU - recall our teriminology cheat sheet from earlier). Details about Memory Management using Numba CUDA is found [here](https://numba.readthedocs.io/en/stable/cuda/memory.html).\n",
        "\n",
        "\n",
        "**Problem 3d** Use Numba to copy our `x` and `y` arrays to device."
      ],
      "metadata": {
        "id": "M1QeQABkboKX"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from numba import cuda\n",
        "\n",
        "x_device = cuda.to_device(x)\n",
        "y_device = cuda.to_device(y)\n",
        "\n",
        "print(x_device)\n",
        "print(x_device.shape)\n",
        "print(x_device.dtype)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "U83bH98vbose",
        "outputId": "7ed9d8bb-bc7d-4325-e39d-8a7c242d8550"
      },
      "execution_count": 17,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "<numba.cuda.cudadrv.devicearray.DeviceNDArray object at 0x7fe0da2f7510>\n",
            "(100000,)\n",
            "float32\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "**Problem 3e** Call our `add_ufunc` routine using the _on device_ arrays and measure the execution speed. "
      ],
      "metadata": {
        "id": "4JNu74Slb02c"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "%timeit add_ufunc(x_device, y_device)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "qxvVIc95b1U0",
        "outputId": "d09da5bc-bf04-4a8d-acff-6c8c4191e4ac"
      },
      "execution_count": 18,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "The slowest run took 4.08 times longer than the fastest. This could mean that an intermediate result is being cached.\n",
            "1000 loops, best of 5: 709 µs per loop\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "**Problem 3f** What is the main reason for the speedup observed here?\n",
        "\n",
        "**Answer:** We removed a key step from the process in our original call of `add_ufunc` as done in `Problem 3c`. We copied the array to device outside of the function call, thus removing the step from the calculation.\n",
        "\n"
      ],
      "metadata": {
        "id": "Sh0rmDmhb7ou"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "With our current optimization, there is still an array created in our call to `add_ufunc` that then needs to be copied _back to host_ (the CPU). We can create the output array ahead of time and set it in our call using _device arrays_.\n",
        "\n",
        "**Problem 3g** Create an output array on device using Numba."
      ],
      "metadata": {
        "id": "NeQUEL5hg44R"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "out_device = cuda.device_array(shape=(n,), dtype=np.float32)  # does not initialize the contents, like np.empty()"
      ],
      "metadata": {
        "id": "6HdWfCEhb8JK"
      },
      "execution_count": 19,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "**Problem 3h** Finally, call our function passing all input and output device arrays we have created and time the execution. "
      ],
      "metadata": {
        "id": "NZAiN8G_b_ZC"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "%timeit add_ufunc(x_device, y_device, out=out_device)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "h8Gxray7b_yt",
        "outputId": "ea2537d5-0093-4e1d-a8cd-eefcd6afa800"
      },
      "execution_count": 20,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "The slowest run took 6.77 times longer than the fastest. This could mean that an intermediate result is being cached.\n",
            "1000 loops, best of 5: 621 µs per loop\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "**Problem 3i** Perform the last step by hand, copy the output of the device array to host and print the first 10 elements to verify the result. "
      ],
      "metadata": {
        "id": "H0uRuCiUcEbA"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "out_host = out_device.copy_to_host()\n",
        "print(out_host[:10])"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "flKXjWiacE6L",
        "outputId": "2b55bb13-ed72-4f18-acac-e5b63efb21ab"
      },
      "execution_count": 21,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "[ 0.  3.  6.  9. 12. 15. 18. 21. 24. 27.]\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Problem 4) Writing CUDA Kernels\n",
        "\n",
        "Some definitions:\n",
        "\n",
        "A **kernel** function is a GPU function that is meant to be called from CPU code (*). \n",
        "\n",
        "It gives it two fundamental characteristics:\n",
        "\n",
        "- kernels cannot explicitly return a value; all result data must be written to an array passed to the function (if computing a scalar, you will probably pass a one-element array);\n",
        "\n",
        "- kernels explicitly declare their thread hierarchy when called: i.e. the number of thread blocks and the number of threads per block (note that while a kernel is compiled once, it can be called multiple times with different block sizes or grid sizes).\n"
      ],
      "metadata": {
        "id": "H-WXt2f9YymF"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "**Problem 4a** Write a _CUDA_ kernel version of our add function.\n",
        "\n",
        "Details on the API is available [here]().."
      ],
      "metadata": {
        "id": "895tb_hqMFpu"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from numba import cuda\n",
        "\n",
        "@cuda.jit\n",
        "def add_kernel(x, y, out):\n",
        "    tx = cuda.threadIdx.x # this is the unique thread ID within a 1D block\n",
        "    ty = cuda.blockIdx.x  # Similarly, this is the unique block ID within the 1D grid\n",
        "\n",
        "    block_size = cuda.blockDim.x  # number of threads per block\n",
        "    grid_size = cuda.gridDim.x    # number of blocks in the grid\n",
        "    \n",
        "    start = tx + ty * block_size\n",
        "    stride = block_size * grid_size\n",
        "\n",
        "    # assuming x and y inputs are same length\n",
        "    for i in range(start, x.shape[0], stride):\n",
        "        out[i] = x[i] + y[i]"
      ],
      "metadata": {
        "id": "OtaWbDwzjDFo"
      },
      "execution_count": 22,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "**Problem 4b** Call our new _CUDA_ kernal:"
      ],
      "metadata": {
        "id": "lxGRcSZbMrag"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import numpy as np\n",
        "\n",
        "n = 5000000\n",
        "x = np.arange(n).astype(np.float32)\n",
        "y = 2 * x\n",
        "out = np.empty_like(x)\n",
        "\n",
        "threads_per_block = 128\n",
        "blocks_per_grid = 30\n",
        "\n",
        "add_kernel[blocks_per_grid, threads_per_block](x, y, out)\n",
        "print(out[:10])"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "hDubS9k3OCa0",
        "outputId": "baf0fd25-580e-4783-dea9-e356888c89e6"
      },
      "execution_count": 23,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "[ 0.  3.  6.  9. 12. 15. 18. 21. 24. 27.]\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "**Problem 4c** Simplify our add kernel using `cuda` grid tools:"
      ],
      "metadata": {
        "id": "hOWfCNaeOMaV"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "@cuda.jit\n",
        "def add_kernel(x, y, out):\n",
        "    start = cuda.grid(1)      # 1 = one dimensional thread grid, returns a single value\n",
        "    stride = cuda.gridsize(1) # ditto\n",
        "\n",
        "    # assuming x and y inputs are same length\n",
        "    for i in range(start, x.shape[0], stride):\n",
        "        out[i] = x[i] + y[i]"
      ],
      "metadata": {
        "id": "_XjXGE85ONCe"
      },
      "execution_count": 24,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "**Problem 4d** As before, create device arrays instead of the implicit copy that takes place when passing NumPy arrays to our GPU kernel:"
      ],
      "metadata": {
        "id": "wQm0ZJ0vOhoO"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "x_device = cuda.to_device(x)\n",
        "y_device = cuda.to_device(y)\n",
        "out_device = cuda.device_array_like(x)"
      ],
      "metadata": {
        "id": "AEmjl39IOicc"
      },
      "execution_count": 25,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "**Problem 4e** Perform timing for the kernel with and without the use of device arrays.  "
      ],
      "metadata": {
        "id": "AVzLZpWIOy-U"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "%timeit add_kernel[blocks_per_grid, threads_per_block](x, y, out)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "det7GaQ-OyJn",
        "outputId": "482080dc-55c9-419f-f817-3d9d74b1e0f3"
      },
      "execution_count": 26,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "The slowest run took 4.29 times longer than the fastest. This could mean that an intermediate result is being cached.\n",
            "10 loops, best of 5: 29 ms per loop\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "%timeit add_kernel[blocks_per_grid, threads_per_block](x_device, y_device, out_device); out_device.copy_to_host()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Ur7DOM6uQcC0",
        "outputId": "e62aae38-6526-4473-9e15-42141b52f99d"
      },
      "execution_count": 27,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "100 loops, best of 5: 5.21 ms per loop\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Some comments about Kernel synchronization:\n",
        "\n",
        "\n",
        "- One extremely important caveat should be mentioned here: CUDA kernel execution is designed to be asynchronous with respect to the host program. This means that the kernel launch `add_kernel[blocks_per_grid, threads_per_block](x_device, y_device, out_device)` returns immediately, allowing the CPU to continue executing while the GPU works in the background. Only host<->device memory copies or an explicit synchronization call will force the CPU to wait until previously queued CUDA kernels are complete.\n",
        "\n",
        "When you pass host NumPy arrays to a CUDA kernel, Numba has to synchronize on your behalf, but if you pass device arrays, processing will continue. If you launch multiple kernels in sequence without any synchronization in between, they will be queued up to run sequentially by the driver, which is usually what you want. If you want to run multiple kernels on the GPU in parallel (sometimes a good idea, but beware of race conditions!), take a look at CUDA streams."
      ],
      "metadata": {
        "id": "EtyLO1xpQpwr"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "**Problem 4f** Perform timing for the kernel with and without the use of device arrays and with explicit versus implicit synchronization for memory copies."
      ],
      "metadata": {
        "id": "ahWqEf_ERvIY"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# CPU input/output arrays, implied synchronization for memory copies\n",
        "%time add_kernel[blocks_per_grid, threads_per_block](x, y, out)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "4WqqEPSaQqZ-",
        "outputId": "0cf7e9df-8fab-42a4-f075-aeb742940fbe"
      },
      "execution_count": 28,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "CPU times: user 30.2 ms, sys: 0 ns, total: 30.2 ms\n",
            "Wall time: 32.2 ms\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# GPU input/output arrays, no synchronization (but force sync before and after)\n",
        "cuda.synchronize()\n",
        "%time add_kernel[blocks_per_grid, threads_per_block](x_device, y_device, out_device)\n",
        "cuda.synchronize()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "3yj7k7vlQvqT",
        "outputId": "d83407c1-7f62-47f4-8429-9a6006089e88"
      },
      "execution_count": 29,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "CPU times: user 0 ns, sys: 573 µs, total: 573 µs\n",
            "Wall time: 583 µs\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# GPU input/output arrays, include explicit synchronization in timing\n",
        "cuda.synchronize()\n",
        "%time add_kernel[blocks_per_grid, threads_per_block](x_device, y_device, out_device); cuda.synchronize()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "6mh2FFO-Q0JL",
        "outputId": "a57bdf11-9c11-4da0-ede0-c7dfb8c66e89"
      },
      "execution_count": 30,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "CPU times: user 2.43 ms, sys: 0 ns, total: 2.43 ms\n",
            "Wall time: 1.68 ms\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "**Question:** What is observed? When should we be mindful to sync?\n",
        "\n",
        "**Answer:** Always be sure to synchronize with the GPU when benchmarking CUDA kernels!"
      ],
      "metadata": {
        "id": "Pa6u-kBCR83k"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "\n",
        "## Problem 5) Matrix multiplication (Optional/Challenge)"
      ],
      "metadata": {
        "id": "ds3xczZIdmdO"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "**Problem 5a**: Rewrite the following implication of matrix multiplication using shared memory. \n",
        "\n",
        "Details on shared memory usage in _CUDA_ is described [here](https://numba.readthedocs.io/en/stable/cuda/memory.html#cuda-shared-memory)."
      ],
      "metadata": {
        "id": "rkhOQBarScY9"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "@cuda.jit\n",
        "def matmul(A, B, C):\n",
        "    \"\"\"Perform square matrix multiplication of C = A * B.\"\"\"\n",
        "    i, j = cuda.grid(2)\n",
        "    if i < C.shape[0] and j < C.shape[1]:\n",
        "        tmp = 0.\n",
        "        for k in range(A.shape[1]):\n",
        "            tmp += A[i, k] * B[k, j]\n",
        "        C[i, j] = tmp"
      ],
      "metadata": {
        "id": "r1TF10tOO03r"
      },
      "execution_count": 31,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        ""
      ],
      "metadata": {
        "id": "YPbvl1c4aqEX"
      },
      "execution_count": null,
      "outputs": []
    }
  ],
  "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.5"
    },
    "colab": {
      "name": "fields_intro_to_GPUs_solutions.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "toc_visible": true,
      "include_colab_link": true
    },
    "accelerator": "GPU",
    "gpuClass": "standard"
  },
  "nbformat": 4,
  "nbformat_minor": 0
}