cupy.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c615dde9-1407-41e5-b635-b1d71a047358",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import cupy as cp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "909858c9-0945-4cb3-805b-63e8fd20edaf",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "# The jupyter notebook is launched from your $HOME directory.\n",
    "# Change the working directory to the workshop directory\n",
    "# which was created in your username directory under /scratch/vp91\n",
    "os.chdir(os.path.expandvars(\"/scratch/vp91/$USER/\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "41ee8aff-e582-4902-bb70-6353264b89d6",
   "metadata": {},
   "source": [
    " #####  ``nvidia-smi``  is a command we can use to get the statistics GPU devices available in a Node"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8b0a69c5-6b65-41a4-bfb9-3c23fc054f64",
   "metadata": {},
   "outputs": [],
   "source": [
    "!nvidia-smi"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5ad13132-8e7f-4215-98b8-d6696307f69e",
   "metadata": {},
   "source": [
    "We can do the same using CuPy functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2cc753ae-405c-4f3d-bd3e-4b62482d4447",
   "metadata": {},
   "outputs": [],
   "source": [
    "ngpus = cp.cuda.runtime.getDeviceCount()\n",
    "print('#GPU = ' + str(ngpus))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b29dda8c-0043-4997-b005-f27484ba4dcc",
   "metadata": {},
   "outputs": [],
   "source": [
    "for g in range(ngpus):\n",
    "    print(cp.cuda.runtime.getDeviceProperties(g))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "125340c4-c876-41ed-b77c-36af46c03e00",
   "metadata": {},
   "source": [
    "##### CuPy has a concept of a current device, which is the default GPU device on which on which all operation of related to CuPy takes place. Unless specifically mentioned, all operation taskes place in this default device. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "47933ba3-b1ae-4ff1-9fb6-c82e0094e685",
   "metadata": {},
   "outputs": [],
   "source": [
    "print('Current Device = ' + str(cp.cuda.runtime.getDevice()))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "05a4e8a2-3966-45f2-b44d-52605f2ed55f",
   "metadata": {},
   "source": [
    "# NDARRAY"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "da4a8b68-dc05-4eca-9948-42e3b82817c0",
   "metadata": {},
   "source": [
    "##### CuPy implements a subset of Numpy interface. \n",
    "\n",
    "``cupy.ndarray `` is akin to the ``numpy.ndarray ``. It is an array object that represents a multidimensional, homogeneous array of fixed-size items. This is the core of CuPy\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "33220008-eab5-4a93-b8b0-d82029d6f864",
   "metadata": {},
   "source": [
    "##### A call to the ``numpy.array()`` allocates the data in the main memory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d714e196-c6f2-4e92-b27f-f8b682ee92e1",
   "metadata": {},
   "outputs": [],
   "source": [
    "x_cpu = np.array([1, 2, 3]) # allocate an ndarray in the main memory"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "92f07393-5192-4c7e-a2a1-0cf52a097397",
   "metadata": {},
   "source": [
    "##### While a call to the ``cupy.array()`` allocates the data in the GPU memory. If no device is specified the memory gets allocated in the ``current`` device."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "03e15ea4-e83d-4a1e-afe8-4a3f261aef39",
   "metadata": {},
   "outputs": [],
   "source": [
    "x_gpu = cp.array([1, 2, 3])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "83697d4c-b915-472c-8192-081c47b08cd2",
   "metadata": {},
   "source": [
    "##### We can use the  device context manage to switch between the devices"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "320e01cb-2a24-4d4e-96a6-5bef6e09e8d9",
   "metadata": {},
   "outputs": [],
   "source": [
    "with cp.cuda.Device(1):  # alloacte an ndarray in the gpu memory of the 1st device\n",
    "    x_on_gpu1 = cp.array([1, 2, 3, 4, 5])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "504742fe-0d8d-4135-95bc-a225b7dbe02c",
   "metadata": {},
   "outputs": [],
   "source": [
    "print('location of x_gpu = ' + str(x_gpu.device) )\n",
    "print('location of x_gpu1 = ' + str(x_on_gpu1.device) )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f30f79a9-8826-4f29-ac18-bc0b584d6cea",
   "metadata": {},
   "source": [
    "##### CuPy always assumes that the operations are performed on the currently active device. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dbf1b207-2c36-4406-b123-1006cb7ffa5e",
   "metadata": {},
   "source": [
    "# Data Transfers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4a7b3a12-599d-4de2-aa8a-5c2adb76328b",
   "metadata": {},
   "outputs": [],
   "source": [
    "x_cpu = np.array([1, 2, 3])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c68ab494-608b-4dc0-8e65-a51d8e2a39c8",
   "metadata": {},
   "source": [
    "In a normal CUDA workflow we have to allocate the memory on the GPU and the move the data to the GPU memory. In CuPy this is not required, the memory allocation and data movement can be done in a single operation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "085ffb9c-2af0-426b-999c-d04649258360",
   "metadata": {},
   "outputs": [],
   "source": [
    "x_gpu_0 = cp.asarray(x_cpu)  # move the ndarray from host mem to GPU0 memeory."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "22c8b38e-9be0-481d-9ea8-4fec4f3a4635",
   "metadata": {},
   "source": [
    "In the past any communication between two GPUs had to go throgh the PCIe card. But now NVIDIA offeres a technology called NVLink. NVLink is a direct GPU-to-GPU interconnect that scales multi-GPU input/output (IO) within a node. This makes GPU-to-GPU transfer (D2D tranfer) much faster than GPU-to-Host (D2H transfer) or Host-to-GPU transfer (H2D transfer). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5d447bec-d662-4c4b-98ee-e49964c207b0",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "with cp.cuda.Device(1):\n",
    "    x_gpu_1 = cp.asarray(x_gpu_0)  # move the ndarray to GPU0 to GPU1."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "448cf1e0-b0cf-4a6f-a4f7-f5eaf45e28fa",
   "metadata": {},
   "source": [
    "##### There are two ways to fetch the data from the GPU to host ``cupy.ndarray.get()`` or ``cupy.asnumpy``"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ae639333-8f9d-4020-a2a3-7406cfde7e97",
   "metadata": {},
   "outputs": [],
   "source": [
    "with cp.cuda.Device(0):\n",
    "    x_cpu = cp.asnumpy(x_gpu_0)  # move the array from GPU 0 back to the host."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8252e144-2683-4dbb-9b2b-a348e26cd1e8",
   "metadata": {},
   "outputs": [],
   "source": [
    "type(x_cpu)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "db9435e4-d446-47f8-9de0-2d1aadbb7b9a",
   "metadata": {},
   "outputs": [],
   "source": [
    "with cp.cuda.Device(1):\n",
    "    x_cpu = x_gpu_1.get()  # move the array from GPU 1 back to the host."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2c0823ee-0006-4f13-b753-43ccf75f5248",
   "metadata": {},
   "outputs": [],
   "source": [
    "type(x_cpu)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8ce05330-ed0f-43d1-8024-fe9d5a51f802",
   "metadata": {},
   "source": [
    "# Device Agnostic Codes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ec37d380-1f8d-42fb-964d-8232dd080851",
   "metadata": {},
   "outputs": [],
   "source": [
    "x_cpu = np.array([1, 2, 3])\n",
    "x_gpu = cp.array([1, 2, 3, 4, 5])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3e990496-c772-46d1-9751-01245469e565",
   "metadata": {},
   "source": [
    " As cupy mimicks numpy we can build device agnostics codes. That is, we can make function calls to a data, without the knowledge of where they reside. The ``cupy.get_array_module()`` function in CuPy returns a reference to cupy if any of its arguments resides on a GPU and numpy otherwise."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2d251de1-87ba-47f7-814b-d519215a1c09",
   "metadata": {},
   "outputs": [],
   "source": [
    "def log_array(x):\n",
    "    xp = cp.get_array_module(x)  # cupy ndarray array reference is returned if x is in GPU\n",
    "                                 # numpy ndarray array reference is returned if x is in host memory\n",
    "    xp.log1p(xp.exp(-abs(x))) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "298c31de-76cc-40ec-b3a0-cf2ba2cc662c",
   "metadata": {},
   "outputs": [],
   "source": [
    "log_array(x_cpu) # the function works with numpy ndarray"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fe2745cc-8ee0-4451-9436-92bdf9305088",
   "metadata": {},
   "outputs": [],
   "source": [
    "log_array(x_gpu) # the function works with cupy ndarray"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4687ae64-2b0f-4c16-8cc7-4086bb55358b",
   "metadata": {},
   "source": [
    "# Explicit Data Transfer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1076bd6b-a80f-4ad2-8117-e8f43f0cda53",
   "metadata": {},
   "outputs": [],
   "source": [
    "x_cpu = np.array([1, 2, 3])\n",
    "y_cpu = np.array([4, 5, 6])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "07cfb002-45f2-406a-859c-c9f80b332afe",
   "metadata": {},
   "outputs": [],
   "source": [
    "z_cpu = x_cpu + y_cpu"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ed466b47-31d6-4b3e-9f75-b15101aa8a55",
   "metadata": {},
   "source": [
    "While CuPy automates most data transfers, we may need to move things explicitly between host and GPU, depdending on the application we are implementing.  CuPy provides two methods ``cupy.asnumpy()`` and ``cupy.asarray()`` to do this. \n",
    "\n",
    "The ``cupy.asnumpy() `` method returns a NumPy array (array on the host), whereas ``cupy.asarray()`` method returns a CuPy array (array on the current device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9188782f-f6c1-4631-b8da-ba0d832cbc2f",
   "metadata": {},
   "outputs": [],
   "source": [
    "x_gpu = cp.asarray(x_cpu) #copy host data to GPU mem"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6f181971-f084-4591-9998-1b77a54bf2e7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# x_gpu is located in device memory\n",
    "# y_cpu is located in host memory\n",
    "# z_cpu is located in host memory\n",
    "\n",
    "z_cpu = x_gpu + y_cpu # generats an error\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f0549e6b-d27c-47c3-a25e-eff41ab7b2b4",
   "metadata": {},
   "source": [
    "We cant have an opeartion where on memory is located on on memory domain (Host) and the other is located on another memory domain (GPU). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "640b5bfe-2c58-43e9-a2b5-1b41c37daceb",
   "metadata": {},
   "outputs": [],
   "source": [
    "x_gpu = cp.asarray(x_cpu) #copy host data to GPU mem\n",
    "\n",
    "# x_gpu is located in device memory\n",
    "# y_cpu is located in host memory\n",
    "# z_cpu is located in host memory\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "231ccdbc-356a-42ee-9a7c-96b3a8106bea",
   "metadata": {},
   "source": [
    "#### So either both the data we operate on must be in host memory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7aad6bcf-a3d8-45e1-89fe-3ea148c06315",
   "metadata": {},
   "outputs": [],
   "source": [
    "# move x_gpu to host memory\n",
    "x_cpu = cp.asnumpy(x_gpu) \n",
    "z_cpu = x_cpu + y_cpu\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d0f2787c-5b27-4f79-b16b-2b881a105bff",
   "metadata": {},
   "source": [
    "#### .... or on the GPU memory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "db2410b7-7beb-4eda-a1d0-0b4ac8969335",
   "metadata": {},
   "outputs": [],
   "source": [
    "# or move y_cpu to device memory\n",
    "y_gpu = cp.asarray(y_cpu)\n",
    "z_gpu = x_gpu + y_gpu"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "285502dd-9fbf-4cb9-9c3c-38f46d4605d5",
   "metadata": {},
   "source": [
    "# User defined kernels"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f4df3017-42cc-421b-a8af-37fd005f9160",
   "metadata": {},
   "source": [
    "Kernels are functions that will be run on the GPU. CuPy provides easy ways to define three types of CUDA kernels: elementwise kernels, reduction kernels and raw kernels"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b3e61d3b-a7dc-49a2-8c3d-f9c89050777d",
   "metadata": {},
   "source": [
    "## Elementwise Kernels\n",
    "\n",
    "An element wise kerenel has 4 parts :\n",
    "1. an input argument list:\n",
    "2. an output argument list\n",
    "3. loop body : defines the operation on each element of the argument\n",
    "4. kernel name\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "008fccda-c646-435a-9237-1e8aae2f5737",
   "metadata": {},
   "outputs": [],
   "source": [
    "element_diff = cp.ElementwiseKernel('float32 x, float32 y', \n",
    "                                    'float32 z', \n",
    "                                    'z = (x - y)', \n",
    "                                    'element_diff')\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ecd0810a-0e64-483c-8934-efc1aeddfdbd",
   "metadata": {},
   "outputs": [],
   "source": [
    "x = cp.array([1, 2, 3], dtype=np.float32)\n",
    "y = cp.array([4, 5, 6], dtype=np.float32)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "429363d4-4edd-4918-912e-9ee2519498d8",
   "metadata": {},
   "outputs": [],
   "source": [
    "z = element_diff(x, y)\n",
    "print(z)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0c2ce202-6f63-4a8a-95b2-11912293b140",
   "metadata": {},
   "outputs": [],
   "source": [
    "z = cp.empty((1, 3), dtype=np.float32)\n",
    "element_diff(x, y, z)\n",
    "print(z)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "07f25bb6-39f2-4ab7-8054-f0d2cdc5d3e5",
   "metadata": {},
   "source": [
    "#### What happens when the ndarray is of different dimension?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f7d4c01a-1ec2-4923-a37a-6abc594e3351",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "e74d40ce-0283-4402-b5d5-820c9d0ee6ca",
   "metadata": {},
   "source": [
    "#### Type can be inferred from the arguments"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "134bcc42-6018-4679-9200-3dddacf06105",
   "metadata": {},
   "outputs": [],
   "source": [
    "element_diff = cp.ElementwiseKernel('T x, T y', \n",
    "                                    'T z', 'z = (x - y)', \n",
    "                                    'element_diff')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "49ec617b-0993-47bf-a916-de3a6f002b52",
   "metadata": {},
   "outputs": [],
   "source": [
    "x = cp.array([1, 2, 3], dtype=np.float32)\n",
    "y = cp.array([4, 5, 6], dtype=np.float32)\n",
    "\n",
    "z = element_diff(x, y)\n",
    "print(z)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3ea09a6b-8c11-42a9-915a-5f9909589d97",
   "metadata": {},
   "outputs": [],
   "source": [
    "x = cp.array([1, 2, 3], dtype=np.int32)\n",
    "y = cp.array([4, 5, 6], dtype=np.int32)\n",
    "\n",
    "z = element_diff(x, y)\n",
    "print(z)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2bd07ce5-bd80-4afc-b1ee-b68769eaea18",
   "metadata": {},
   "source": [
    "We can also write a kernel with manual indexing for some arguments by telling ElementwiseKernel class to use manual indexing. This is done by adding the ``raw`` keyword preceding the type specifier"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "719cbd6f-7a39-4df3-b07e-b0ac61e0f702",
   "metadata": {},
   "outputs": [],
   "source": [
    "# i indicates the index within the loop\n",
    "# _ind.size() indicates total number of elements to apply the elementwise operation\n",
    "\n",
    "element_sum = cp.ElementwiseKernel('T x, raw T y', \n",
    "                                   'T z', 'z = x + y[_ind.size() - i - 1]', \n",
    "                                   'element_sum')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "365d7488-56f0-46c7-b035-af65a8583953",
   "metadata": {},
   "outputs": [],
   "source": [
    "x = cp.array([1, 2, 3], dtype=np.float32)\n",
    "y = cp.array([4, 5, 6], dtype=np.float32)\n",
    "\n",
    "z = element_sum(x, y)\n",
    "print(z)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9c5de77e-1d2c-400e-988d-7524ff7d06fd",
   "metadata": {},
   "source": [
    "## Reduction kernels\n",
    "\n",
    "A reduction kernel has 4 parts :\n",
    "1. Identity value: This value is used for the initial value of reduction.\n",
    "2. Mapping expression: It is used for the pre-processing of each element to be reduced.\n",
    "3. Reduction expression: It is an operator to reduce the multiple mapped values. The special variables a and b are used for its operands.\n",
    "4. Post mapping expression: It is used to transform the resulting reduced values. The special variable a is used as its input. Output should be written to the output parameter."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ababfaf8-739b-43f5-aa57-b6847c4196f3",
   "metadata": {},
   "outputs": [],
   "source": [
    "reduction_kernel = cp.ReductionKernel(\n",
    "    'T x', # input param\n",
    "    'T y',  # output param\n",
    "    'x',  # map\n",
    "    'a + b',  # reduce\n",
    "    'y = a',  # post-reduction map\n",
    "    '0',  # identity value\n",
    "    'reduction_kernel'  # kernel name\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8edb46dc-07d8-42d1-a5aa-84cda4825b99",
   "metadata": {},
   "outputs": [],
   "source": [
    "x = cp.arange(10, dtype=np.float32).reshape(1, 10)\n",
    "print(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e3322e6c-ffcc-4d13-9753-3cae7f59069e",
   "metadata": {},
   "outputs": [],
   "source": [
    "reduction_kernel(x, axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "785106b9-1726-4c70-853a-27813ed18334",
   "metadata": {},
   "outputs": [],
   "source": [
    "x = cp.arange(10, dtype=np.float32).reshape(2, 5)\n",
    "print(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e2ff105a-e917-4d30-8eb6-e394c327b038",
   "metadata": {},
   "outputs": [],
   "source": [
    "reduction_kernel(x, axis=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7b4bd42e-f493-4259-9333-07672dad4504",
   "metadata": {},
   "source": [
    "#### TODO: Change the above function to find the sum of squares of each element in the array\n",
    "That is, If the array is [2, 3, 4, 5], find 2^2 + 3^2 + 4^2 + 5^2 "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1bd0bb9f-dc2a-44a7-9009-1772cf081707",
   "metadata": {},
   "source": [
    "## Raw kernels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "788113f2-d6f3-401b-9d09-686e10b41baa",
   "metadata": {},
   "outputs": [],
   "source": [
    "add_kernel = cp.RawKernel(r'''\n",
    "    extern \"C\" __global__\n",
    "    void my_add(const float* x1, const float* x2, float* y) \n",
    "    {\n",
    "        int tid = blockDim.x * blockIdx.x + threadIdx.x;\n",
    "        y[tid] = x1[tid] + x2[tid];\n",
    "    }''', 'my_add')\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3542798a-059c-4017-813c-351097c5eec8",
   "metadata": {},
   "outputs": [],
   "source": [
    "x = cp.arange(25, dtype=cp.float32).reshape(5, 5)\n",
    "y = cp.arange(25, dtype=cp.float32).reshape(5, 5)\n",
    "z = cp.zeros((5, 5), dtype=cp.float32)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0e6e0a85-55f2-44f8-93a6-7009be4bd549",
   "metadata": {},
   "outputs": [],
   "source": [
    "# When calling a raw kernel ypu have to specify  \n",
    "# how threads are grouped (grids and blocks)\n",
    "# https://developer.nvidia.com/blog/cuda-refresher-cuda-programming-model/\n",
    "\n",
    "add_kernel((5,), (5,), (x, y, z))\n",
    "print(z)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2bdc0b17-acd8-43b5-b756-1ce98d89bbc7",
   "metadata": {},
   "source": [
    "# CUDA Events"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "38edac10-35f5-4f5f-995d-d47b48dd11d6",
   "metadata": {},
   "outputs": [],
   "source": [
    "a_cp = cp.arange(10)\n",
    "b_cp = cp.arange(10)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a18ecd42-6c36-42bb-94c1-16490c34e883",
   "metadata": {},
   "outputs": [],
   "source": [
    "e1 = cp.cuda.Event() # create an event"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6fee2020-d251-41d9-874b-cad7b962de26",
   "metadata": {},
   "outputs": [],
   "source": [
    "e1.record() # Records an event on the stream."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ddcf9ab0-e98b-4a19-8f3d-1b5f823b0d6c",
   "metadata": {},
   "outputs": [],
   "source": [
    "a_cp = b_cp * a_cp + 8"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e0b79f10-7518-4109-aefc-0ba97d7082e5",
   "metadata": {},
   "outputs": [],
   "source": [
    "e2 = cp.cuda.get_current_stream().record() # create and record the event"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9ef7ad37-319c-49f5-83e3-4879f028a117",
   "metadata": {},
   "outputs": [],
   "source": [
    "s = cp.cuda.Stream.null\n",
    "# make sure the stream wait for the second event\n",
    "s.wait_event(e2)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a3287f69-ce45-4613-935e-f0356a8463ea",
   "metadata": {},
   "outputs": [],
   "source": [
    "with s:\n",
    "    # all actions in a stream happens seuentially\n",
    "    # as the stream is waiting for the second event to complete\n",
    "    # we can be assured that all the operations before it also has been complete.\n",
    "    a_np = cp.asnumpy(a_cp)\n",
    "\n",
    "# Waits for the stream that track an event to complete that event\n",
    "e2.synchronize()\n",
    "t = cp.cuda.get_elapsed_time(e1, e2)\n",
    "\n",
    "print(t)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "399d4fa3-1951-4f4f-89a2-d628443cfa92",
   "metadata": {},
   "outputs": [],
   "source": [
    "# https://developer.nvidia.com/blog/how-implement-performance-metrics-cuda-cc/"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "edd43f62-e1c6-45b7-904b-2f09fda65dfd",
   "metadata": {},
   "source": [
    "# CUDA Streams"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a6828bb2-6c68-432f-a865-2c788ba8ef38",
   "metadata": {},
   "outputs": [],
   "source": [
    "a_np = np.arange(10)\n",
    "s = cp.cuda.Stream()\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9af27361-1ec3-40b7-b0dd-c15693b3a49a",
   "metadata": {},
   "outputs": [],
   "source": [
    "with s:\n",
    "   a_cp = cp.asarray(a_np)  # H2D transfer on stream s\n",
    "   b_cp = cp.sum(a_cp)      # kernel launched on stream s \n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8d498daa-242b-40a7-9462-a280ebc2078a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# or we can use 'use()'\n",
    "# if we use 'use() any subsequent CUDA operation will completed\n",
    "# using the stream we specify, until we make a change \n",
    "s.use()\n",
    "\n",
    "b_np = cp.asnumpy(b_cp)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bd831863-29dd-4573-abfa-afae159804c9",
   "metadata": {},
   "outputs": [],
   "source": [
    "assert s == cp.cuda.get_current_stream()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b2c640f1-951e-4c42-a8e2-bbe89905a4fb",
   "metadata": {},
   "outputs": [],
   "source": [
    "# go back to the default stream\n",
    "cp.cuda.Stream.null.use()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9cd6da6d-fc04-4933-bc8f-8194eb98b9aa",
   "metadata": {},
   "outputs": [],
   "source": [
    "assert s == cp.cuda.get_current_stream() # run fails is assert condition is false\n",
    "                                         # generates an error \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4c1062ca-ff32-4021-8925-7360f5aafe0a",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "47e91bf2-7486-4aea-8b8c-2ac7cc11f21c",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}