13_CuPyAndLegate/13_CuPyAndLegate.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "e6597bb3-15b6-4639-82ed-841386591567",
   "metadata": {},
   "source": [
    "![NCAR UCAR Logo](../NCAR_CISL_NSF_banner.jpeg)\n",
    "# Python GPU Session: CuPy and Legate\n",
    "\n",
    "By: Brett Neuman [bneuman@ucar.edu](mailto:bneuman@ucar.edu), Consulting Services Group, CISL & NCAR\n",
    "\n",
    "Date: July 28th 2022"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "300697eb-0d7a-4e09-80fd-869a2d502971",
   "metadata": {},
   "source": [
    "Head to the [NCAR JupyterHub portal](https://jupyterhub.hpc.ucar.edu/stable) and __start a JupyterHub session on a Casper batch node__ (using 1 CPU, 1 GPU) and open the notebook in `12_CuPyAndLegate.ipynb`. Be sure to clone (if needed) and update/pull the NCAR GPU_workshop directory.\n",
    "\n",
    "```shell\n",
    "# Use the JupyterHub GitHub GUI on the left panel or the below shell commands\n",
    "git clone git@github.com:NCAR/GPU_workshop.git\n",
    "git pull\n",
    "```\n",
    "\n",
    "# Workshop Etiquette\n",
    "* Please mute yourself and turn off video during the session.\n",
    "* Questions may be submitted in the chat and will be answered when appropriate. You may also raise your hand, unmute, and ask questions during Q&A at the end of the presentation.\n",
    "* By participating, you are agreeing to [UCAR’s Code of Conduct](https://www.ucar.edu/who-we-are/ethics-integrity/codes-conduct/participants)\n",
    "* Recordings & other material will be archived & shared publicly.\n",
    "* Feel free to follow up with the GPU workshop team via Slack or submit support requests to [support.ucar.edu](https://support.ucar.edu)\n",
    "    * Office Hours: Asynchronous support via [Slack](https://ncargpuusers.slack.com) or schedule a time with an organizer"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "553aa25c-d12f-42f1-bb23-1eb92d9fe3f2",
   "metadata": {},
   "source": [
    "## Notebook Setup\n",
    "This notebook will require running within a JupyterHub PBS Batch (interactive) session.  Much of the Python code will fail if a GPU is not detected on the node.  Select the PBS Batch option when launching from JupyterHub and set the `PROJECT` code to a currently active project, ie `UCIS0004` for the GPU workshop, and `QUEUE` to the appropriate routing queue depending on if during a live workshop session (`gpuworkshop`), during weekday 8am to 5:30pm MT (`gpudev`), or all other times (`casper`). Due to limited shared GPU resources, please use `GPU_TYPE=gp100` during the workshop. Otherwise, set `GPU_TYPE=v100` (required for `gpudev`) for independent work. See [Casper queue documentation](https://arc.ucar.edu/knowledge_base/72581396#StartingCasperjobswithPBS-Concurrentresourcelimits) for more info.\n",
    "\n",
    "If running this notebook outside of the NCAR computing environment, a compatible GPU is required on a node or your local machine to execute the code in the CuPy sections of this notebook.\n",
    "\n",
    "### Changing Notebook Kernel\n",
    "\n",
    "You'll need to use a JupyterHub kernel that has the CuPy module.  There are several ways to do this but you can navigate to the `Kernel` dropdown and select `Change Kernel`.  Select **GPU Workshop** kernel from the dropdown to run the CuPy examples in this notebook."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "40669051-77a8-44fc-9f81-a96645fbf270",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Python Virtual Environment Setup\n",
    "\n",
    "You can set up your own virtual environment for running the Python code outside of this notebook.  This will also be useful if you would like to create your own virtual environment for GPU programming experimentation with Python.\n",
    "\n",
    "See [Python virtual environment documentation](https://kb.ucar.edu/display/RC/Using+conda+environments+for+Python+access) for using the NCAR Python Library (npl) or for setting up your own virtual environment.  Here is a bash example for creating your own virtual environment."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f106e693-527a-49c0-a072-f5ca95b74db6",
   "metadata": {
    "tags": []
   },
   "source": [
    "```bash\n",
    "# Conda used for virtual environment in NCAR clusters\n",
    "module load conda/latest\n",
    "# Create your own virtual environment named 'pygpu'\n",
    "mamba create -n pygpu python==3.7.9 numpy scipy matplotlib pandas geocat-comp xarray dask wrf-python cupy \n",
    "# Use the virtual environment (don't forget to deactivate \"conda deactivate\")\n",
    "conda activate pygpu\n",
    "\n",
    "# You can also create a virtual environment from existing .yml file\n",
    "conda env create --file envs/environment.yml\n",
    "```\n",
    "\n",
    "----"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8dfb6890-a078-4cfa-9182-d1e475c9706d",
   "metadata": {},
   "source": [
    "# Python GPU Itinerary\n",
    "\n",
    "1) Python GPU Programming Overview\n",
    "2) CuPy\n",
    "    - Overview and Setup\n",
    "    - Basic Drop-in Replacement Example\n",
    "    - Computational Fluid Dynamics Example\n",
    "    - Additional CuPy Functionality\n",
    "    - GeoCAT Example   \n",
    "3) Legate\n",
    "    - Overview, Philosophy, and Setup\n",
    "    - Basic Drop-in Replacement Examples\n",
    "\n",
    "----"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "88cb4cf8-273d-4b89-a745-1f6378afd4f9",
   "metadata": {},
   "source": [
    "## GPU Acceleration Overview\n",
    "\n",
    "The Central Processing Unit (CPU) is a low latency focused, general purpose processor that handles sequential and complex tasks well.  The Graphics Processing Unit (GPU) is a throughput focused set of computational units that excel at parallel processing, particularly on single instruction, multiple data (SIMD) type problems.  The decision to port your code to run on a GPU is problem independent.\n",
    "\n",
    "|  |  |\n",
    "| :--- | :--- | \n",
    "| **CPU** | **GPU** |\n",
    "| Several cores | Many cores |\n",
    "| Low latency | High throughput |\n",
    "| Good for serial processing | Good for parallel processing |\n",
    "| Can do a handful of operations at once| Can do thousands of operations at once |\n",
    "\n",
    "It is important to be able to identify problem types that are likely to benefit from porting code to the GPU.  Python offers APIs that make quick drop-in replacement of an entire code attempting but this can still take significant time, especially verifying correctness of your changes.\n",
    "\n",
    "*Source:* https://blogs.nvidia.com/blog/2009/12/16/whats-the-difference-between-a-cpu-and-a-gpu/\n",
    "\n",
    "The NCAR GPU Workshop series contains more detailed information on GPU programming topics."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f3b610dc-8da4-44b7-a8ec-e70df2b23526",
   "metadata": {},
   "source": [
    "## Interpreted vs. Compiled Languages for GPU Programming\n",
    "\n",
    "Python is a partially interpreted language, meaning that the code requires an additional step to run as machine code. This additional step will need to be considered when writing performant code.  One trade off for Python's platform independence, dynamic typing, and modularity is the additional step to interpret the bytecode generated by Python and then running machine code on your device.\n",
    "\n",
    "|  |  |\n",
    "| :--- | :--- | \n",
    "| **Interpreted** | **Compiled** |\n",
    "| Ease of use | Performance |\n",
    "| Parsed, Interpreted, and Executed each time run | Compilation overhead occurs once |\n",
    "| Lenient datatype enforcement | Strict datatype enforcement |\n",
    "| Platform independent | Platform dependent |\n",
    "\n",
    "Fortran and C/C++ examples of languages that are translated by running the source code through a compiler. This results in very efficient code that can be executed any number of times. The overhead for the translation is incurred just once, when the source is compiled.\n",
    "\n",
    "Interpreted languages must be parsed, interpreted, and executed each time the program is run, thereby greatly adding to the cost of running the program. For this reason, interpreted programs are usually less efficient than compiled programs.  __The trade-off is machine resources for programmer time__.\n",
    "\n",
    "Keeping this in mind, we can very generally say that it would make sense to use a compiled language for the intensive parts of an application, whereas interpreted could be use for more rapid prototyping, experimentation, or low computational workload sections of an application.\n",
    "\n",
    "*Source:* https://www.ibm.com/docs/en/zos-basic-skills?topic=zos-compiled-versus-interpreted-languages\n",
    "\n",
    "Some of the considerations needed for compiled languages that are potentially hidden from Python GPU implementations:\n",
    "  - Memory movement\n",
    "  - Garbage collection\n",
    "  - Instruction flow\n",
    "  - Loop control\n",
    "  - Parallelization\n",
    "\n",
    "Comparing a Python GPU implementation to the CUDA / C implementation:\n",
    "\n",
    "![gpu_py_c.png](img/gpu_py_c.png)\n",
    "*Source: NCAR SIParCS 2022 Haniye Kashgarani (Implementation of GeoCat on GPUs)*"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "900a3ab1-b677-4d48-8198-002572bf6c68",
   "metadata": {},
   "source": [
    "### Compiled Python\n",
    "\n",
    "Compiled Python allows calls to compiled subsections of Python code.  Common compilation options include Cython and Numba which uses the Just-In-Time (JIT) compiler.  Compiled sections of code often see a performance boost from running the code in a compiled language and avoids the slowdowns from interpreted drawbacks covered in the previous section.  Compiled Python provides opportunities for improved performance at the cost of programming ease.\n",
    "\n",
    "We'll provide an example of a custom kernel that uses Compiled Python near the end of the notebook. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1452652f-019c-4d04-b19d-743cf9583968",
   "metadata": {},
   "source": [
    "## Python Scientific Computing Tools for the GPU\n",
    "\n",
    "Many of the tools to enable performant Python code on the GPU attempt to be drop-in replacements.  Well known modules Numpy, Scipy, and Pandas are primary targets for GPU porting efforts.\n",
    "\n",
    "You might be familiar with GPU accelerated Python ecosystems such as TensorFlow, PyTorch, or RAPIDS.  These tools attempt to provide a suite of APIs to cover specific areas of research, most commonly for data science.  Our focus today won't be on these broader scope suites but on individual APIs that allow developers to port their code with drop-in replacements for common scientific computing modules.\n",
    "\n",
    "Drop-in replacement is used to describe a new API or module that can be imported but retains the function syntax of another module.  The user can replace the module name and expect the function calls to remain the same.\n",
    "\n",
    "One consideration introduced with Python GPU programming is portability.  Different tools used for Python GPU programming need to be designed for specific device families.\n",
    "\n",
    "Below is a chart showing scientific computing GPU APIs, device compatibility, and additional compute functionality for known Python GPU programming tool suites: \n",
    "\n",
    "![Python GPU Chart](img/pygpuchart.png)\n",
    "  \n",
    "    (* : Options to use compiled Python)\n",
    "    \n",
    "----"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f2bf9c1a-0179-40eb-9db1-259771d4c3c5",
   "metadata": {},
   "source": [
    "# CuPy\n",
    "\n",
    "![CuPy Logo](img/cupy_logo.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3546a7f5-e7d5-4eeb-a27a-1076d46fab77",
   "metadata": {},
   "source": [
    "CuPy is a NumPy/SciPy-compatible array library for GPU-accelerated computing with Python. CuPy acts as a drop-in replacement to run existing NumPy/SciPy code on NVIDIA CUDA or AMD ROCm platforms.\n",
    "\n",
    "CuPy provides ndarray, sparse matrices, and the associated routines for GPU devices, all having the same API as NumPy and SciPy.\n",
    "\n",
    "CuPy's stated goal is to __provide Python users GPU acceleration capabilities, without the in-depth knowledge of underlying GPU technologies__. The CuPy team focuses on providing:\n",
    "\n",
    "  - *A complete NumPy and SciPy API coverage to become a full drop-in replacement plus advanced CUDA features to maximize performance.*\n",
    "\n",
    "  - *A quality, mature, and fundamental package for all projects needing acceleration at any scale.*\n",
    "\n",
    "[CuPy Homepage](https://docs.cupy.dev/en/stable/index.html#)\n",
    "\n",
    "[CuPy Github](https://github.com/cupy/cupy)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "80f55cac-6dc5-4ca8-9a51-34ed4f30df13",
   "metadata": {},
   "source": [
    "## Setup\n",
    "\n",
    "### Requirements\n",
    "\n",
    "NVIDIA CUDA GPU with Compute Capability 3.0+.\n",
    "\n",
    "CUDA Toolkit: v10.2 / v11.0 / v11.1 / v11.2 / v11.3 / v11.4 / v11.5 / v11.6 / v11.7\n",
    "\n",
    "Python: \n",
    "  - Versions: v3.7.0+ / v3.8.0+ / v3.9.0+ / v3.10.0+\n",
    "  - Modules: Numpy, SciPy (optional)\n",
    "  \n",
    "CuPy's [installation page](https://docs.cupy.dev/en/stable/install.html) has step-by-step details to install using pip, conda, or from source.  Setting up your virtual environment to include CuPy is an easy way to maintain package dependencies and share your code with others. \n",
    "\n",
    "### Importing\n",
    "\n",
    "At the most basic setup, import the CuPy module in your Python file:\n",
    "\n",
    "```python\n",
    "import cupy as cp\n",
    "```\n",
    "\n",
    "Then use `cp` instead of Numpy's standard `np` to create arrays and perform computations.  \n",
    "\n",
    "Another option to determine if CuPy is available is to check during import.  This method checks if CuPy is available and avoids errors on import. \n",
    "\n",
    "```python\n",
    "def _import_cupy():\n",
    "    \"\"\"imports the cupy package and checks if not installed\"\"\"\n",
    "    try:\n",
    "        import cupy as cp\n",
    "        return cp\n",
    "    except ImportError as e:\n",
    "        print(f\"CuPy is not installed for GPU computation!\")\n",
    "        pass  # module doesn't exist, deal with it. \n",
    "```\n",
    "*Source: Haniye Kashgarani SIParCS GeoCAT GPU Project*\n",
    "\n",
    "We'll go through a couple of examples below and then provide resources to see which Numpy functions have been implemented in the CuPy API."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9b23a921-ac6c-4fef-aec0-476b84cb3f16",
   "metadata": {},
   "source": [
    "## CuPy Functions Walkthrough"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "71235ea5-c0fb-4afb-a30e-72275ad0b7f0",
   "metadata": {},
   "source": [
    "### Import NumPy and CuPy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "860567c3-18b2-444c-93a9-96ce620ea08f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import cupy as cp\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "be9ca22a-8e41-46b1-a43d-6cf93e4fd677",
   "metadata": {},
   "source": [
    "### Creating Arrays on the CPU and GPU\n",
    "\n",
    "Let's start by using some basic NumPy functions to create arrays and how that same procedure is done in CuPy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8aba4551-3b83-4f2b-9fc6-541b0ae071e6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# NumPy data (host / cpu)\n",
    "x_cpu = np.linspace(0, 2, 5)\n",
    "print(\"On the CPU: \", x_cpu)\n",
    "\n",
    "# CuPy data\n",
    "x_gpu = cp.linspace(2, 4, 5)\n",
    "print(\"On the GPU: \", x_gpu)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "10fa58ed-0911-4dea-80fb-d4061b95e2d3",
   "metadata": {},
   "source": [
    "### Moving Data between Host and Device"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b233a422-513d-40ad-b271-facd356cb5b6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Move data to GPU\n",
    "x_gpu = cp.asarray(x_cpu)\n",
    "\n",
    "# Move data back to host\n",
    "x_cpu = cp.asnumpy(x_gpu)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4fd163f6-8efb-4796-9cf2-af406323a551",
   "metadata": {},
   "source": [
    "### Device Information\n",
    "\n",
    "Combined with the `nvidia-smi` command you can identify then use devices with CuPy.  This would become more important if your code utilizes multiple GPUs.  The default behavior runs code on Device 0.  Here is an example of forcing an operation to run on Device 0:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3d117f27-d2a1-4bfd-846f-a439ebebc9b1",
   "metadata": {},
   "outputs": [],
   "source": [
    "with cp.cuda.Device(0):\n",
    "    x_on_gpu0 = cp.array([1, 2, 3, 4, 5])\n",
    "x_on_gpu0 = cp.array([1, 2, 3, 4, 5])\n",
    "print(\"X on GPU0: \", x_on_gpu0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f45b1dd6-16c2-4a21-b4a1-c65658f8eb86",
   "metadata": {},
   "source": [
    "### CuPy Implemented Functions\n",
    "\n",
    "CuPy has equivalents for many of the commonly used NumPy functions, but not all.  Here is a short list of the NumPy function with it's CuPy equivalent.  You can see almost all of CuPy's functions will use the same function call as its NumPy equivalent."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5e32819b-41aa-4a59-9837-69a8c5479958",
   "metadata": {},
   "source": [
    "| | |\n",
    "| :--- | :--- |\n",
    "| **NumPy** | **CuPy** |\n",
    "| numpy.identity        | cupy.identity       |\n",
    "| numpy.matmul        | cupy.matmul       |\n",
    "| numpy.nan_to_num        | cupy.nan_to_num       |\n",
    "| numpy.zeros        | cupy.zeros       |\n",
    "| numpy.ones        | cupy.ones       |\n",
    "| numpy.shape        | cupy.shape       |\n",
    "| numpy.reshape        | cupy.reshape       |\n",
    "| numpy.tensordot        | cupy.tensordot       |\n",
    "| numpy.transpose        | cupy.transpose       |\n",
    "| numpy.fft.fft        | cupy.fft.fft       |\n",
    "\n",
    "CuPy also has equivalent functions used within SciPy but fewer are implemented than NumPy.  A full list of CuPy's Numpy and SciPy equivalent functions are provided on the link below.\n",
    "\n",
    "[Complete Comparison of NumPy and SciPy to CuPy functions](https://docs.cupy.dev/en/stable/reference/comparison.html)\n",
    "\n",
    "[CuPy API Reference](https://docs.cupy.dev/en/stable/reference/index.html)\n",
    "\n",
    "----"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bdaacdc2-eed5-4e22-8da7-65ff450715ff",
   "metadata": {
    "tags": []
   },
   "source": [
    "## NumPy to CuPy Example\n",
    "\n",
    "Here we want to show an example of easily porting a short script to run on the GPU.  We'll import some time series data and extract simple statistics on both the CPU and the GPU.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cea2b704-7da4-4cb3-8c4c-17a3e2a0bb38",
   "metadata": {},
   "source": [
    "### The Gravity Recovery and Climate Experiment (GRACE) Dataset\n",
    "\n",
    "Refers to a pair of NASA satellites that has flown in low-Earth orbit since 2002.  We've collected a time series data representation of mean annual amplitude of total water storage in a 200km square area near Colorado.\n",
    "\n",
    "Import matplotlib for visualization of the time series data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1fffa6cd-3d30-49ab-b425-95fa889c27b7",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1711f8ad-c78a-46cb-a7ef-3ca53e25efd6",
   "metadata": {},
   "source": [
    "### CPU Implementation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e096d000-5e6f-45e2-8e53-562e33401418",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Pull in the time series data to CPU\n",
    "data_cpu = np.loadtxt(\"code/gracets/data/grace_raw.o\")\n",
    "\n",
    "# Separate dates and amplitude\n",
    "date_cpu = data_cpu[:, 0]\n",
    "amp_cpu = data_cpu[:, 1]\n",
    "\n",
    "# Get min, max, and mean from amplitude\n",
    "y_min_cpu = np.min(amp_cpu)\n",
    "y_max_cpu = np.max(amp_cpu)\n",
    "y_mean_cpu = np.round(np.mean(amp_cpu), 2)\n",
    "print(\"Min:\", y_min_cpu, \"\\nMax:\", y_max_cpu, \"\\nMean:\", y_mean_cpu)\n",
    "\n",
    "# Plot the time series\n",
    "ts = plt.plot(date_cpu, amp_cpu)\n",
    "plt.xlabel(\"Year\")\n",
    "plt.ylabel(\"GRACE Annual Amplitude (mm)\")\n",
    "plt.title(\"GRACE Time Series on CPU\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8b291172-3cf6-46e5-944c-1b8362fa7e1a",
   "metadata": {},
   "source": [
    "### GPU Implementation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "95e10e41-f294-4224-989a-7f3053cd7166",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Pull in the time series data to GPU\n",
    "data_gpu = cp.loadtxt(\"code/gracets/data/grace_raw.o\")\n",
    "\n",
    "# Separate dates and amplitude\n",
    "date_gpu = data_gpu[:, 0]\n",
    "amp_gpu = data_gpu[:, 1]\n",
    "\n",
    "# Get min, max, and mean from amplitude\n",
    "y_min_gpu = cp.min(amp_gpu)\n",
    "y_max_gpu = cp.max(amp_gpu)\n",
    "y_mean_gpu = cp.round(cp.mean(amp_gpu), 2)\n",
    "print(\"Min:\", y_min_gpu, \"\\nMax:\", y_max_gpu, \"\\nMean:\", y_mean_gpu)\n",
    "\n",
    "# Conversion back to Numpy to plot the data on the CPU\n",
    "ts = plt.plot(cp.asnumpy(date_gpu), cp.asnumpy(amp_gpu))\n",
    "plt.xlabel(\"Year\")\n",
    "plt.ylabel(\"GRACE Annual Amplitude (mm)\")\n",
    "plt.title(\"GRACE Time Series on GPU\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c6a905df-0193-41f2-8805-4ec60b3ad552",
   "metadata": {},
   "source": [
    "*GRACE Source:* Swenson, Sean & National Center for Atmospheric Research Staff (Eds). Last modified 08 Oct 2013. \"The Climate Data Guide: GRACE: Gravity Recovery and Climate Experiment: Surface mass, total water storage, and derived variables.\" Retrieved from https://climatedataguide.ucar.edu/climate-data/grace-gravity-recovery-and-climate-experiment-surface-mass-total-water-storage-and."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3c63e73a-e0f0-4f9a-9920-7f6f327ca233",
   "metadata": {},
   "source": [
    "## Developing for both CPU and GPU\n",
    "\n",
    "Once you begin porting code to the GPU you will have to consider how to handle creating data on either the CPU or GPU.  CuPy has the ability to identify array types to write device agnostic code.  Below is a code sample using the community standard for reading an array and setting the package to use either NumPy or CuPy based on the type.\n",
    "\n",
    "```python\n",
    "xp = cp.get_array_module(x)\n",
    "xp.linspace(0, 2, 5)\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0481d9f7-e3fe-45a1-bfee-39199f696dcb",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def addone(x):\n",
    "    xp = cp.get_array_module(x)\n",
    "    print(\"Using:\", xp.__name__)\n",
    "    return x+1\n",
    "\n",
    "# Calls and Output\n",
    "print(addone(x_cpu))\n",
    "print(addone(x_gpu))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c570b5eb-1650-49f9-addb-adf9a160f4f6",
   "metadata": {},
   "source": [
    "There are other possibilities to consider when writing CPU and GPU compatible code:\n",
    "  - Separate development branches for GPU code\n",
    "  - Implementing GPU flags within code to use the GPU portion of your code\n",
    "  - Allocating all NumPy or Cupy datatypes with a function or script\n",
    " \n",
    "You may also need to consider new unit or regression tests to verify your GPU code does not inject any unexpected behavior both on the CPU and GPU side."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "812c70fa-6efa-4d77-bab5-1c69438437a9",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Computational Fluid Dynamics (CFD) Example\n",
    "\n",
    "Let's look at an example where you might expect performance improvements by running it on the GPU.  The [12 Step's to Navier-Stokes](https://lorenabarba.com/blog/cfd-python-12-steps-to-navier-stokes/) program was designed by Lorena Barba at George Washington University to help students learn and implement a Navier-Stokes finite-difference solver using Python.  Our example looks at the cavity flow, which solves the system of 2D Navier-Stokes differential equations.\n",
    "\n",
    "Below are the computationally intensive discretized equations used to solve for `u`, `v`, and `p` in the Cavity Flow problem:\n",
    "\n",
    "![Discretized Steps](img/uvp_cfd.png)\n",
    "\n",
    "More detail on this specific step can be found on the website [12 Steps to Navier Stokes: Step 11 Cavity Flow](https://nbviewer.org/github/barbagroup/CFDPython/blob/master/lessons/14_Step_11.ipynb).\n",
    "\n",
    "We'll walk through steps to implement the GPU version of this code and also compare the performance results between the CPU and GPU."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3744872d-e13d-4e9a-a87e-81f1b6048a15",
   "metadata": {},
   "source": [
    "### File Locations and Descriptions\n",
    "  \n",
    "In the `code` folder there are several Python files used to create these results.  Feel free to view and modify these files to experiment further.\n",
    "\n",
    "|Folder |File |Description |\n",
    "|:--- |:--- |:--- |\n",
    "|Numpy   |                             |                                         |\n",
    "|        | [step11_numpy.py](code/cfd/numpy/step11_numpy.py)                | cpu numpy using float       |\n",
    "|        | [step11_numpy_double.py](code/cfd/numpy/step11_numpy_double.py)         | cpu numpy using double              |\n",
    "|Cupy    |                             |                                         |\n",
    "|        | [step11_cupy.py](code/cfd/numpy/step11_cupy.py)              | cpu/gpu agnostic with gpu flag using float   |\n",
    "|        | [step11_cupy_double.py](code/cfd/numpy/step11_cupy_double.py)       | cpu/gpu agnostic with gpu flag using double  |\n",
    "|        | [step11_nsight.py](code/cfd/numpy/step11_nsight.py)              | cpu/gpu with cupyx profiler nvtx flags for nsight  |"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "62e4187b-a68f-42a4-bb4d-a9668649f7e1",
   "metadata": {},
   "source": [
    "### CPU Implementation\n",
    "\n",
    "Below we initialize the data used in our application with NumPy.  This will run on the CPU and use whatever default NumPy optimizations:\n",
    "\n",
    "```python\n",
    "    def _init_params(self, n):\n",
    "        # NumPy setup\n",
    "        self.x = np.linspace(0, 2, n)\n",
    "        self.y = np.linspace(0, 2, n)\n",
    "        self.X, self.Y = np.meshgrid(self.x, self.y)\n",
    "        self.u = np.zeros((n, n))\n",
    "        self.v = np.zeros((n, n))\n",
    "        self.p = np.zeros((n, n))         \n",
    "        self.un = np.empty_like(self.u)\n",
    "        self.vn = np.empty_like(self.v)\n",
    "        self.pn = np.empty_like(self.p)\n",
    "        self.b = np.zeros((n, n)) \n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7ce9acc6-b8da-4a1d-af25-6e34488f1c25",
   "metadata": {},
   "source": [
    "It should be noted that even the NumPy CPU implementation hides the computational heavy lifting unless you look for it.  Here is an example of uncertainity in how many threads NumPy will use by default:\n",
    "\n",
    "https://stackoverflow.com/questions/30791550/limit-number-of-threads-in-numpy"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "856356e4-9ffe-45c6-8374-7299e443d5ce",
   "metadata": {},
   "source": [
    "### Basic GPU Port\n",
    "\n",
    "Here we'll use the device agnostic method and a flag to determine how to allocate the data, either with Numpy on the host (CPU) or CuPy on the device (GPU).  Drop-in, device agnostic replacement:\n",
    "  \n",
    "```python\n",
    "\n",
    "    class CavityFlow(object):\n",
    "        # ...\n",
    "        \n",
    "        def __init__(self, dims, timesteps, use_gpu=False):\n",
    "            # ...\n",
    "\n",
    "        def _init_params(self, n):\n",
    "            # CuPy setup\n",
    "            if self.use_gpu:\n",
    "                xp = cp\n",
    "            else:\n",
    "                xp = np\n",
    "\n",
    "            self.x = xp.linspace(0, 2, n)\n",
    "            self.y = xp.linspace(0, 2, n)\n",
    "            self.X, self.Y = xp.meshgrid(self.x, self.y)\n",
    "            self.u = xp.zeros((n, n))\n",
    "            self.v = xp.zeros((n, n))\n",
    "            self.p = xp.zeros((n, n))         \n",
    "            self.un = xp.empty_like(self.u)\n",
    "            self.vn = xp.empty_like(self.v)\n",
    "            self.pn = xp.empty_like(self.p)\n",
    "            self.b = xp.zeros((n, n))\n",
    "        \n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8ec8bc27-d307-4652-9494-44af7f92733a",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Import additional modules"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6becb6a9-74be-4662-a001-3e0e10ea15ff",
   "metadata": {},
   "outputs": [],
   "source": [
    "from matplotlib import pyplot, cm\n",
    "import time"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "51a0709f-a95a-4117-81d3-1d12c3836bde",
   "metadata": {},
   "source": [
    "### Class and Functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3fcf2964-fdda-4e99-9efe-025ca75ce83a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Cavity Flow Class\n",
    "class CavityFlow(object):\n",
    "\n",
    "    def __init__(self, dims, timesteps, use_gpu, plot_result):\n",
    "        # Parameter Initialization\n",
    "        self.n = dims\n",
    "        self.nt = timesteps\n",
    "        self.use_gpu = use_gpu\n",
    "        self.plot_result = plot_result\n",
    "        self.nit = 50\n",
    "        self.c = 1\n",
    "        self.dx = 2 / (self.n-1)\n",
    "        self.dy = 2 / (self.n-1)\n",
    "        self.rho = 1\n",
    "        self.nu = 0.1\n",
    "        self.dt = 0.0001\n",
    "\n",
    "        self._init_params(self.n)\n",
    "\n",
    "    def _init_params(self, n):\n",
    "        # CuPy setup\n",
    "        if self.use_gpu:\n",
    "            xp = cp\n",
    "        else:\n",
    "            xp = np\n",
    "\n",
    "        self.x = xp.linspace(0, 2, n, dtype=xp.double)\n",
    "        self.y = xp.linspace(0, 2, n, dtype=xp.double)\n",
    "        self.X, self.Y = xp.meshgrid(self.x, self.y)\n",
    "        self.u = xp.zeros((n, n), dtype=xp.double)\n",
    "        self.v = xp.zeros((n, n), dtype=xp.double)\n",
    "        self.p = xp.zeros((n, n), dtype=xp.double)         \n",
    "        self.un = xp.empty_like(self.u, dtype=xp.double)\n",
    "        self.vn = xp.empty_like(self.v, dtype=xp.double)\n",
    "        self.pn = xp.empty_like(self.p, dtype=xp.double)\n",
    "        self.b = xp.zeros((n, n), dtype=xp.double)\n",
    "\n",
    "    def plot(self):\n",
    "        fig = pyplot.figure(figsize=(11,7), dpi=100)\n",
    "        # plotting the pressure field as a contour\n",
    "        pyplot.contourf(self.X, self.Y, self.p, alpha=0.5, cmap=cm.viridis)  \n",
    "        pyplot.colorbar()\n",
    "        # plotting the pressure field outlines\n",
    "        pyplot.contour(self.X, self.Y, self.p, cmap=cm.viridis)  \n",
    "        # plotting velocity field\n",
    "        pyplot.quiver(self.X[::2, ::2], self.Y[::2, ::2], self.u[::2, ::2], self.v[::2, ::2]) \n",
    "        pyplot.xlabel('X')\n",
    "        pyplot.ylabel('Y')\n",
    "        pyplot.title('Cavity Flow')\n",
    "\n",
    "    def _build_up_b(self):\n",
    "        self.b[1:-1, 1:-1] = (self.rho * (1 / self.dt * \n",
    "                             ((self.u[1:-1, 2:] - self.u[1:-1, 0:-2]) / \n",
    "                             (2 * self.dx) + (self.v[2:, 1:-1] - self.v[0:-2, 1:-1]) / (2 * self.dy)) -\n",
    "                             ((self.u[1:-1, 2:] - self.u[1:-1, 0:-2]) / (2 * self.dx))**2 -\n",
    "                             2 * ((self.u[2:, 1:-1] - self.u[0:-2, 1:-1]) / (2 * self.dy) *\n",
    "                             (self.v[1:-1, 2:] - self.v[1:-1, 0:-2]) / (2 * self.dx))-\n",
    "                             ((self.v[2:, 1:-1] - self.v[0:-2, 1:-1]) / (2 * self.dy))**2))\n",
    "\n",
    "    def _pressure_poisson(self):\n",
    "        self.pn = cp.copy(self.p)\n",
    "        for q in range(self.nit):\n",
    "            self.pn = self.p.copy()\n",
    "            self.p[1:-1, 1:-1] = (((self.pn[1:-1, 2:] + self.pn[1:-1, 0:-2]) * self.dy**2 + \n",
    "                                 (self.pn[2:, 1:-1] + self.pn[0:-2, 1:-1]) * self.dx**2) /\n",
    "                                 (2 * (self.dx**2 + self.dy**2)) -\n",
    "                                 self.dx**2 * self.dy**2 / (2 * (self.dx**2 + self.dy**2)) * \n",
    "                                 self.b[1:-1,1:-1])\n",
    "\n",
    "            # dp/dx = 0 at x = 2\n",
    "            self.p[:, -1] = self.p[:, -2]\n",
    "            # dp/dy = 0 at y = 0\n",
    "            self.p[0, :] = self.p[1, :]\n",
    "            # dp/dx = 0 at x = 0\n",
    "            self.p[:, 0] = self.p[:, 1]\n",
    "            # p = 0 at y = 2 \n",
    "            self.p[-1, :] = 0\n",
    "\n",
    "    def compute(self):\n",
    "        for n in range(self.nt):\n",
    "            un = cp.copy(self.u)\n",
    "            vn = cp.copy(self.v)\n",
    "            self._build_up_b()\n",
    "            self._pressure_poisson()\n",
    "            self.u[1:-1, 1:-1] = (un[1:-1, 1:-1]-\n",
    "                             un[1:-1, 1:-1] * self.dt / self.dx *\n",
    "                            (un[1:-1, 1:-1] - un[1:-1, 0:-2]) -\n",
    "                             vn[1:-1, 1:-1] * self.dt / self.dy *\n",
    "                            (un[1:-1, 1:-1] - un[0:-2, 1:-1]) -\n",
    "                             self.dt / (2 * self.rho * self.dx) * (self.p[1:-1, 2:] - self.p[1:-1, 0:-2]) +\n",
    "                             self.nu * (self.dt / self.dx**2 *\n",
    "                            (un[1:-1, 2:] - 2 * un[1:-1, 1:-1] + un[1:-1, 0:-2]) +\n",
    "                             self.dt / self.dy**2 *\n",
    "                            (un[2:, 1:-1] - 2 * un[1:-1, 1:-1] + un[0:-2, 1:-1])))\n",
    "\n",
    "            self.v[1:-1, 1:-1] = (vn[1:-1, 1:-1] -\n",
    "                             un[1:-1, 1:-1] * self.dt / self.dx *\n",
    "                            (vn[1:-1, 1:-1] - vn[1:-1, 0:-2]) -\n",
    "                             vn[1:-1, 1:-1] * self.dt / self.dy *\n",
    "                            (vn[1:-1, 1:-1] - vn[0:-2, 1:-1]) -\n",
    "                             self.dt / (2 * self.rho * self.dy) * (self.p[2:, 1:-1] - self.p[0:-2, 1:-1]) +\n",
    "                             self.nu * (self.dt / self.dx**2 *\n",
    "                            (vn[1:-1, 2:] - 2 * vn[1:-1, 1:-1] + vn[1:-1, 0:-2]) +\n",
    "                             self.dt / self.dy**2 *\n",
    "                            (vn[2:, 1:-1] - 2 * vn[1:-1, 1:-1] + vn[0:-2, 1:-1])))\n",
    "\n",
    "            self.u[0, :]  = 0\n",
    "            self.u[:, 0]  = 0\n",
    "            self.u[:, -1] = 0\n",
    "            # Set velocity on cavity lid equal to 1\n",
    "            self.u[-1, :] = 1\n",
    "            self.v[0, :]  = 0\n",
    "            self.v[-1, :] = 0\n",
    "            self.v[:, 0]  = 0\n",
    "            self.v[:, -1] = 0\n",
    "        if (self.plot_result):\n",
    "            self.plot()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7283eede-b12e-4762-a837-d841a546ef44",
   "metadata": {},
   "source": [
    "Set up timing for compute section:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f77b2bf9-4233-4fb2-84e9-bcf73c015c8e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def launch_test(n, ts, use_gpu=False, plot_result=False):\n",
    "    start_datamv = time.perf_counter()\n",
    "    flow = CavityFlow(n, ts, use_gpu, plot_result)\n",
    "    end_datamv = time.perf_counter()\n",
    "    if (use_gpu):\n",
    "        start_gpu = time.perf_counter()\n",
    "        flow.compute()\n",
    "        end_gpu = time.perf_counter()\n",
    "        t_gpu = end_gpu - start_gpu\n",
    "        total_time = (end_datamv - start_datamv) + t_gpu\n",
    "        print(\"--- Cavity Flow Performance Test ---\")\n",
    "        print(\"Dimension: \", n, \"\\nTimesteps: \", ts, \"\\nGPU Run\")\n",
    "        print(\"Computation Time (s): \", t_gpu, \"\\nTotal Time (s): \", total_time, \"\\n\")\n",
    "    else:\n",
    "        start_cpu = time.perf_counter()\n",
    "        flow.compute()\n",
    "        end_cpu = time.perf_counter()\n",
    "        t_cpu = end_cpu - start_cpu\n",
    "        total_time = (end_datamv - start_datamv) + t_cpu\n",
    "        print(\"--- Cavity Flow Performance Test ---\")\n",
    "        print(\"Dimension: \", n, \"\\nTimesteps: \", ts, \"\\nCPU Run\")\n",
    "        print(\"Computation Time (s): \", t_cpu, \"\\nTotal Time (s): \", total_time, \"\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c1aa11b5-8f49-4c27-b475-6a0c3d294b1b",
   "metadata": {},
   "source": [
    "CuPy provides a basic profiling tool that allows you to repeat a function call and take the average runtime for results.  Below is an example of benchmarking the entire `launch_test` function:\n",
    "\n",
    "```python  \n",
    "    benchmark(launch_test, (n, ts, use_gpu), n_repeat=20, n_warmup=5))\n",
    "    \n",
    "    launch_test         :    \n",
    "        CPU: 2633938.903 us   +/-57936.881 (min:2431833.191 / max:2880125.065) us                     \n",
    "      GPU-0: 6149643.408 us   +/-57938.509 (min:6057958.008 / max:6283191.406) us\n",
    "``` "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "50ac668e-57b2-466d-81c2-ed1cb05e0df4",
   "metadata": {},
   "source": [
    "### Launching Tests\n",
    "\n",
    "Start the test with input for gridsize (n) and timesteps (ts).  Feel free to re-run the code with your own inputs for gridsize and timesteps to see performance results.  The simulation is sensitive to the ratio of dx/dt, so be mindful that very large grid sizes could cause the simulation to go out of bounds.  Lowering the dt value is a good way to prevent out of bounds cases with larger grid sizes.\n",
    "\n",
    "The cavity flow code can be parallelized over the grid size but not over the timesteps.  Modifying timesteps in it's current state will see increase in runtime but not parallelization. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f27974aa-414d-44a7-bd6d-b4ec8e37f16c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 41x41 Grid, 100 Timesteps\n",
    "n = 41\n",
    "ts = 100\n",
    "launch_test(n, ts, use_gpu=False, plot_result=True)\n",
    "launch_test(n, ts, use_gpu=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bbc88876-6990-473f-bc15-9680c4ebeb77",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 256x256 Grid, 100 Timesteps\n",
    "n = 256\n",
    "ts = 100\n",
    "launch_test(n, ts, use_gpu=False, plot_result=True)\n",
    "launch_test(n, ts, use_gpu=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b9e1e05f-d679-430d-b87e-6c21689e19af",
   "metadata": {},
   "source": [
    "### Plotting Results\n",
    "\n",
    "Identifying the approximate grid size that it is beneficial to run the code on the GPU:\n",
    "    \n",
    "![step11_results](img/step11_results1.png)\n",
    "\n",
    "Comparing additional grid sizes from 32x32 to 4096x4096, doubling grid size at each point:\n",
    "\n",
    "![step11_gridsizes](img/perf_logscale.png)\n",
    "\n",
    "You'll notice the time difference at the 4096x4096 scale where the CPU wall-time executed in 2000 seconds, over 30 minutes, where the GPU wall-time executed in just over 20 seconds."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5a0a3598-2b9a-463a-be26-a793aae652b9",
   "metadata": {},
   "source": [
    "### Advanced Profiling\n",
    "\n",
    "CuPy includes Nvidia Tools Extension (NVTX) profiling capabilities that can be imported from the `cupyx.profiler` module.  This allows you to flag and label a function when running the Nsight Systems profiler for your GPU application.\n",
    "\n",
    "Nsight Systems benchmarking and setting time ranges in the Python code: \n",
    "\n",
    "```python\n",
    "    from cupyx.profiler import time_range, benchmark \n",
    "\n",
    "    @time_range()\n",
    "    def compute(self):\n",
    "        ...\n",
    "```\n",
    "\n",
    "The above code example would flag the `compute` function of our CFD Python program to be profiled.  Any operations will be traced within the generated report.  This can add overhead to your runtime.\n",
    "\n",
    "To run the Nsys profiler with your Python GPU application, run the following within your terminal window:\n",
    "        \n",
    "```sh\n",
    "    ### Run\n",
    "    nsys profile -o step11_profile python step11_cupy_double.py -t cuda,cublas,nvtx -f true\n",
    "```\n",
    "\n",
    "The results for a profiled run of 128 gridsize (n) and 500 timesteps (ts) is shown below.  The arguments genereate a profile report called `step11_profile.qdrep` from the `step11_cupy_double.py` Python file, and will trace Cuda, CuBLAS, and NVTX events when profiling the application.  A detailed overview of the profiler arguments can be found in a previous workshop session.\n",
    "\n",
    "Nsight Systems generated a report with file extension `.qdrep`.  You can also view this generated report [*step11_profile.qdrep*](code/cfd/cupy/step11_profile.qdrep) in the `code/cfd/cupy` folder if you have previously installed Nsight Systems.  Right click on the file to download it to your local machine.\n",
    "\n",
    "Nsight Systems GUI viewing the profiled application report:\n",
    "\n",
    "![nsys_step11.png](img/nsys_step11.png)\n",
    "\n",
    "If you've used NVTX previously, the NVTX API can be used to annotate sections of your code without CuPy's imported benchmark and timing utilities from the cupyx.profiler.  NVIDIA provides a tutorial on using NVTX to profile your Python code in more detail:\n",
    "\n",
    "[Nvidia Tools Extension API for Profiling Python and C code](https://developer.nvidia.com/blog/nvidia-tools-extension-api-nvtx-annotation-tool-for-profiling-code-in-python-and-c-c/)\n",
    "\n",
    "A previous GPU Workshop session covers Nsight Systems and Compute for profiling:\n",
    "\n",
    "[NCAR GPU Workshop - Session 10: Hands on Nsight](https://github.com/NCAR/GPU_workshop/tree/workshop/10_HandsOnNsight)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cf09471d-6df5-4bd4-b78d-3ab5ffefc371",
   "metadata": {},
   "source": [
    "## Additional CuPy Topics\n",
    "\n",
    "### Custom kernels\n",
    "\n",
    "Custom kernels in CuPy use C language syntax to write compiled code within Python.  There are code optimization opportunities by writing your own kernels and leveraging concepts of compiled languages and CUDA architecture. Here is a basic example of defining datatypes and calling a custom kernel within Python."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "50389d84-49e6-44a9-9bd8-679457752280",
   "metadata": {},
   "outputs": [],
   "source": [
    "x = cp.arange(6, dtype='f').reshape(2, 3)\n",
    "y = cp.arange(3, dtype='f')\n",
    "\n",
    "kernel = cp.ElementwiseKernel(\n",
    "    'float32 x, float32 y', 'float32 z',\n",
    "    '''\n",
    "    if (x - 2 > y) {\n",
    "      z = x * y;\n",
    "    } else {\n",
    "      z = x + y;\n",
    "    }\n",
    "    ''', 'my_kernel')\n",
    "\n",
    "kernel(x, y)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9dccc6c4-b167-4d52-b2aa-39abf1050bb1",
   "metadata": {},
   "source": [
    "In addition to custom kernels, low-level CUDA features are available:\n",
    "\n",
    "  - __Stream and Event__: CUDA stream and per-thread default stream\n",
    "\n",
    "  - __Memory Pool__: Customizable memory allocator with a built-in memory pool\n",
    "\n",
    "  - __Host API Binding__: Directly call CUDA libraries, such as NCCL, cuDNN, cuTENSOR, and cuSPARSELt APIs from Python\n",
    "\n",
    "References to get started with custom kernels:\n",
    "\n",
    "[CuPy User-Defined Kernels](https://docs.cupy.dev/en/stable/user_guide/kernel.html)\n",
    "\n",
    "[Nvidia Introduction to Cuda Python](https://nvidia.github.io/cuda-python/overview.html)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a8299fdd-b9c1-4e5f-b990-2ad029e10645",
   "metadata": {},
   "source": [
    "### Distributed Communication\n",
    "\n",
    "The `cupyx.distributed` API can be used to perform communication between different processes in CuPy.  It provides access to the NCCL operations to communication between GPUs using the following arguments and parameters. \n",
    "\n",
    "\n",
    "```python\n",
    "class cupyx.distributed.NCCLBackend(n_devices, rank, host='127.0.0.1', port=13333, use_mpi=False)\n",
    "\n",
    "\"\"\"\n",
    "Interface that uses NVIDIA’s NCCL to perform communications.\n",
    "\n",
    "Parameters\n",
    "    n_devices (int) – Total number of devices that will be used in the distributed execution.\n",
    "\n",
    "    rank (int) – Unique id of the GPU that the communicator is associated to its value needs to be 0 <= rank < n_devices.\n",
    "\n",
    "    host (str, optional) – host address for the process rendezvous on initialization. Defaults to “127.0.0.1”.\n",
    "\n",
    "    port (int, optional) – port used for the process rendezvous on initialization. Defaults to 13333.\n",
    "\n",
    "    use_mpi (bool, optional) – switch between MPI and use the included TCP server for initialization & synchronization. Defaults to False.\n",
    "\"\"\"\n",
    "```\n",
    "\n",
    "In the next session we'll look at using Dask for distributed Multi-GPU programming."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3ce0e0bc-0d9f-466e-80b3-ceb830ff6552",
   "metadata": {
    "tags": []
   },
   "source": [
    "## GeoCAT Example\n",
    "\n",
    "NCAR's GeoCat team has started exploring the potential of porting computational meteorological code and other scripts to run on GPUs.  GeoCAT is working towards porting the NCAR Command Line scripts to Python and also ways to extend the existing code while porting.\n",
    "\n",
    "Below is a more formalized version to convert datatypes to run on the GPU using CuPy.  You'll notice the first function checks to see if CuPy can be imported before attempting to use the module for conversion.  This example also shows how you might add Xarray support for your code by converting the NumPy arrays into Xarray compatible datatypes."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "79131aef-a792-47b0-b5fc-f4e60ef6e3c0",
   "metadata": {},
   "source": [
    "```python\n",
    "def _import_cupy():\n",
    "    \"\"\"imports the cupy and checks if not installed.\"\"\"\n",
    "    try:\n",
    "        import cupy as cp\n",
    "        return cp\n",
    "    except ImportError as e:\n",
    "        print(f\"Cupy is not installed for GPU computation!\")\n",
    "        pass  # module doesn't exist, deal with it. \"\"\"\n",
    "    \n",
    "def _convert_to_gpu_array(inputs):\n",
    "    xp = _import_cupy()\n",
    "    inputs_gpu = []\n",
    "    in_types = [type(item) for item in inputs]\n",
    "    if xp.ndarray in in_types:\n",
    "        #if the inputs are already cupy arrays\n",
    "        return inputs\n",
    "    elif np.ndarray in in_types:\n",
    "        #convert numpy to cupy\n",
    "        for item in inputs:\n",
    "            inputs_gpu.append(xp.asarray(item))\n",
    "    elif xr.DataArray in in_types:\n",
    "        #convert xarray\n",
    "        in_types = [type(item.data) for item in inputs]\n",
    "        if np.ndarray in in_types:\n",
    "            #xarray with type(item.data) = numpy.ndarray\n",
    "            for item in inputs:\n",
    "                inputs_gpu.append(xr.DataArray(xp.asarray(item.data)))\n",
    "        elif da.Array in in_types:\n",
    "            #xarray with type(item.data) = dask array\n",
    "            for item in inputs:\n",
    "                inputs_gpu.append(xr.DataArray(item.data.map_blocks(\n",
    "                    xp.asarray)))\n",
    "        else:\n",
    "            inputs_gpu = inputs\n",
    "    elif da.Array in in_types:\n",
    "        #convert dask array\n",
    "        for item in inputs:\n",
    "            inputs_gpu.append(item.map_blocks(xp.asarray))\n",
    "    else:\n",
    "        return inputs\n",
    "    return inputs_gpu\n",
    "```\n",
    "*Source: NCAR SIParCS 2022 Haniye Kashgarani (Implementation of GeoCAT on GPUs)*"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5b9710bf-d100-4b34-b58d-47f8efce6efb",
   "metadata": {},
   "source": [
    "This code is in the process of being pushed to production with the goal of integration directly into the main branch using the device agnostic approach shown earlier in the notebook.  GeoCAT is Open Development and you can contribute via their GitHub page.  Instructions for contributing are included on their website and via the link below.\n",
    "\n",
    "[GeoCat Homepage](https://geocat.ucar.edu/)\n",
    "\n",
    "[Contribute to GeoCat](https://geocat.ucar.edu/pages/contributing.html)\n",
    "\n",
    "----"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "48dc40aa-f960-4700-ad7d-c9f050b27ee2",
   "metadata": {},
   "source": [
    "# Legate\n",
    "\n",
    "Legate is another option for drop-in replacement for accelerated Python GPU code.  Numpy and Pandas are supported for drop-in replacement via `Legate.cuNumerics` and `Legate.Pandas`, respectively. \n",
    "\n",
    "Legate's approach to Python GPU acceleration:\n",
    "\n",
    "  1) The programming model must be identical to programming a single sequential CPU on their laptop or desktop. All concerns relating to parallelism, data distribution, and synchronization must be implicit. \n",
    "  2) Software must be compositional and not just interoperable (i.e. functionally correct). Libraries developed in the Legate ecosystem must be able to exchange partitioned and distributed data without requiring \"shuffles\" or unnecessary blocking synchronization.\n",
    " \n",
    "These two core concepts of Legate reduce the effort of the developer and shifts the memory management responsibilities onto the backend of the Legate Core API.\n",
    "\n",
    "## Setup\n",
    "\n",
    "Legate uses a driver script to launch your Python GPU programs.  The legate driver script is in the bin directory of the Legate installation directory. Note that the default python interpreter (python) will not work with programs that use Legate libraries, you need to use this custom driver script.\n",
    "\n",
    "For example, to run your script in the default configuration (4 CPUs cores and 4 GB of memory) run:\n",
    "\n",
    "``` Python\n",
    "installdir/bin/legate my_python_program.py [other args]\n",
    "```\n",
    "\n",
    "Documentation on downloading and installing Legate can be found here:\n",
    "\n",
    "https://github.com/nv-legate/legate.core"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c290d968-a6fc-46d7-bf28-6e0df1a65b54",
   "metadata": {},
   "source": [
    "## CuNumerics"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "63c86435-9478-4378-bced-336278027eca",
   "metadata": {},
   "source": [
    "### NumPy Drop-in Replacement\n",
    "\n",
    "After installing the Legate Core library, the next step is to install a Legate application library such as cuNumeric. The installation process for a Legate application library will require you to provide a pointer to the location of your Legate Core library installation as this will be used to configure the installation of the Legate application library. After you finish installing any Legate application libraries, you can then simply replace their import statements with the equivalent ones from any Legate application libraries you have installed. For example, you can change this:\n",
    "\n",
    "```python\n",
    "import numpy as np\n",
    "```\n",
    "\n",
    "to this:\n",
    "\n",
    "```python\n",
    "import cunumeric as np\n",
    "```\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "df8a9f0d-6cf3-41e8-ab89-07146d42a1a5",
   "metadata": {},
   "source": [
    "## Legate.Pandas\n",
    "\n",
    "### Pandas Drop-in Replacement\n",
    "\n",
    "Legate Pandas is a distributed and accelerated drop-in replacement of Pandas. Legate Pandas enables high-performance, scalable execution of dataframe programs on multi-GPU systems by combining the Legion runtime with GPU accelerated dataframe kernels in cuDF. Legate Pandas targets dataframe programs with data processing requirements that cannot be fulfilled by a single GPU.  Legate.Pandas provides an alternative to a CuDF+Dask implementation.\n",
    "\n",
    "The process to use Legate.Pandas is the same as CuNumerics, using a simple drop-in:\n",
    "\n",
    "```python\n",
    "import pandas as pd\n",
    "```\n",
    "\n",
    "with this:\n",
    "\n",
    "```python\n",
    "import legate.pandas as pd\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ba6c30ad-ae31-4178-895b-feefc090ea7c",
   "metadata": {},
   "source": [
    "## Legate Resources\n",
    "\n",
    "[Legate Main Resource](https://github.com/nv-legate/)\n",
    "\n",
    "[Legate CuNumeric](https://github.com/nv-legate/cunumeric)\n",
    "\n",
    "[Legate Pandas](https://github.com/nv-legate/legate.pandas)\n",
    "\n",
    "[NVidia Scaling the PyData Ecosystem with Legate](https://www.nvidia.com/en-us/on-demand/session/gtcfall21-a31168/)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "29c9fa25-d281-4b3a-a591-bf27b776545b",
   "metadata": {},
   "source": [
    "# Resources\n",
    "\n",
    "[CuPy Website and Documentation](https://docs.cupy.dev/en/stable/index.html)\n",
    "\n",
    "[12 Steps to CFD GitHub](https://github.com/barbagroup/CFDPython)\n",
    "\n",
    "[12 Steps to CFD Main page](https://lorenabarba.com/blog/cfd-python-12-steps-to-navier-stokes/)\n",
    "\n",
    "[GRACE Data NCAR](https://climatedataguide.ucar.edu/climate-data/grace-gravity-recovery-and-climate-experiment-surface-mass-total-water-storage-and)\n",
    "\n",
    "[GRACE Data CU](http://geoid.colorado.edu/grace/index.html)\n",
    "\n",
    "[Python Intrepreter and How Python Works](https://towardsdatascience.com/how-does-python-work-6f21fd197888)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "GPU Workshop",
   "language": "python",
   "name": "pygpu"
  },
  "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.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}