notebooks/dpep_scalability_journey_example-pairwise_distance.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "9a5db442-1e0d-42d6-8d1c-c50808f2fabf",
   "metadata": {},
   "source": [
    "# Pairwise Distance: Open-Source XPU Programming with Data Parallel Extensions for Python* language"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "ecfe8042-d9dc-4d74-adad-a4b6103113d6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# SPDX-FileCopyrightText: 2020 - 2023 Intel Corporation\n",
    "#\n",
    "# SPDX-License-Identifier: Apache-2.0"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fe38f7a4-ad5f-4b9a-bac4-6aa023f2ba06",
   "metadata": {},
   "source": [
    "In this notebook, learn how to compute pairwise distance, a very popular distance metric applicable in many applications such as machine learning, geospatial, biology, HPC, and so on. \n",
    "\n",
    "This tutorial will cover how Data Parallel Extension for NumPy* and Data Parallel Extension for Numba* can be used to compute pairwise distance in NumPy* code efficiently and scale performance through open-source heterogenous computing and DPC++/SYCL compilation."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "88a670af-8d6c-4df0-9c53-26ff63e56760",
   "metadata": {},
   "source": [
    "## Simple NumPy Script"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c3ffc2f5-4d3a-49d9-86db-46697d94e532",
   "metadata": {},
   "source": [
    "Here is a simple NumPy* script running pairwise distance on CPU - this is a normal NumPy script without any Data Parallel Extensions for Python* language (dpep) products included. \n We are using Intel® Optimizations for NumPy* for out-of-box acceleration powered by Intel® oneAPI Math Kernel Library optimizations. This can be downloaded via Anaconda using the command `conda install -c intel numpy` or as part of [Intel® Distribution for Python*](https://www.intel.com/content/www/us/en/developer/tools/oneapi/distribution-for-python.html). \n",
    "\n",
    "We will also record the time it takes to compute pariwise distance with this script for performance purposes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "b7a89df0-7e71-4fcb-8e4e-358ef6790ca0",
   "metadata": {},
   "outputs": [],
   "source": [
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "161d3106-e5e2-459e-9699-2d607e563f9b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Time to Compute Pairwise Distance with Stock NumPy on CPU: 2.764047861099243\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "import time\n",
    "\n",
    "#compute pairwise distance\n",
    "def pairwise_distance(data, distance):\n",
    "    data_sqr = np.sum(np.square(data), dtype = np.float32, axis=1)\n",
    "    np.dot(data, data.T, distance)\n",
    "    distance *= -2\n",
    "    np.add(distance, data_sqr.reshape(data_sqr.size, 1), distance)\n",
    "    np.add(distance, data_sqr, distance)\n",
    "    np.sqrt(distance, distance)\n",
    "\n",
    "data = np.random.ranf((10*1024,3)).astype(np.float32)\n",
    "distance = np.empty(shape=(data.shape[0], data.shape[0]), dtype=np.float32)\n",
    "\n",
    "\n",
    "#time pairwise distance calculations\n",
    "start = time.time()\n",
    "pairwise_distance(data, distance)\n",
    "print(\"Time to Compute Pairwise Distance with Stock NumPy on CPU:\", time.time() - start)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8126c45f-a863-42f5-b589-7b4e006eca96",
   "metadata": {},
   "source": [
    "## Data Parallel Extension for NumPy* Running on CPU for Additional Speed-Up"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3dd0c2fa-590c-4f0a-95b6-7e32561934de",
   "metadata": {},
   "source": [
    "Now let's convert this script into a Data Parallel Extension for NumPy* (dpnp) Script that can be run on any SYCL-supported device (CPU, GPU, etc) by a simple import statement change. \n",
    "\n",
    "Instead of `import numpy as np`, let's replace this statement with `import dpnp as np`. That is the only change required!\n",
    "\n",
    "The script will run on a Intel GPU SYCL-device by default, if available, so we will specify `\"cpu\"` as the device for the arrays in order to run on a Intel® Xeon 4th Generation CPU. Even though it is running on CPU like NumPy, it has SYCL and oneAPI optimizations underneath it to give it more performance."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "72373f34-d925-4326-b3f3-8d4451cc3586",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Time to Compute Pairwise Distance with DPNP on CPU: 0.2046055793762207\n"
     ]
    }
   ],
   "source": [
    "#using dpnp to allow heterogenous computing of NumPy!\n",
    "import dpnp as np \n",
    "#import numpy as np\n",
    "\n",
    "import time\n",
    "\n",
    "#compute pairwise distance\n",
    "def pairwise_distance(data, distance):\n",
    "    data_sqr = np.sum(np.square(data), dtype = np.float32, axis=1) #specify cpu as the device\n",
    "    np.dot(data, data.T, distance)\n",
    "    distance *= -2\n",
    "    np.add(distance, data_sqr.reshape(data_sqr.size, 1), distance)\n",
    "    np.add(distance, data_sqr, distance)\n",
    "    np.sqrt(distance, distance)\n",
    "    \n",
    "#specify CPU device type to run on CPU\n",
    "data = np.random.ranf((10*1024,3), device = \"cpu\").astype(np.float32)\n",
    "distance = np.empty(shape=(data.shape[0], data.shape[0]), dtype=np.float32, device = \"cpu\")\n",
    "\n",
    "\n",
    "#do compilation of kernel first run, we will calculate performance on subsequent runs\n",
    "pairwise_distance(data, distance)\n",
    "#time pairwise distance calculations\n",
    "start = time.time()\n",
    "pairwise_distance(data, distance)\n",
    "print(\"Time to Compute Pairwise Distance with DPNP on CPU:\", time.time() - start)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6264484d-2cc1-482d-949c-8a12b994fc93",
   "metadata": {},
   "source": [
    "We already see some speedup just by using DPNP on CPU!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b1c25cfa-0d13-4e04-8053-d98376414b09",
   "metadata": {},
   "source": [
    "## Data Parallel Extension for NumPy* Running on Intel GPU"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7ce6099a-be10-493c-925a-839591c2c4aa",
   "metadata": {},
   "source": [
    "Now instead of running this script on Intel CPU, let's run it on a Intel® Data Center GPU Max Series GPU. Since dpnp will use a Intel GPUS SYCL-device by device if it's available, we can remove the device parameter from our calls."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "e4843868-d044-4a05-9cca-1327486c2a62",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Time to Compute Pairwise Distance with DPNP on GPU: 0.012220144271850586\n"
     ]
    }
   ],
   "source": [
    "#using dpnp to allow heterogenous computing of NumPy!\n",
    "import dpnp as np \n",
    "#import numpy as np\n",
    "\n",
    "import time\n",
    "\n",
    "#compute pairwise distance\n",
    "def pairwise_distance(data, distance):\n",
    "    data_sqr = np.sum(np.square(data), dtype = np.float32, axis=1) #specify cpu as the device\n",
    "    np.dot(data, data.T, distance)\n",
    "    distance *= -2\n",
    "    np.add(distance, data_sqr.reshape(data_sqr.size, 1), distance)\n",
    "    np.add(distance, data_sqr, distance)\n",
    "    np.sqrt(distance, distance)\n",
    "\n",
    "#no device specification necessary, by default it will run on Intel GPU SYCL-device if available\n",
    "data = np.random.ranf((10*1024,3)).astype(np.float32)\n",
    "distance = np.empty(shape=(data.shape[0], data.shape[0]), dtype=np.float32)\n",
    "\n",
    "\n",
    "#do compilation of kernel first run, we will calculate performance on subsequent runs\n",
    "pairwise_distance(data, distance)\n",
    "#time pairwise distance calculations\n",
    "start = time.time()\n",
    "pairwise_distance(data, distance)\n",
    "print(\"Time to Compute Pairwise Distance with DPNP on GPU:\", time.time() - start)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fdaa94a8-330d-42a8-9e05-45c7fc02250d",
   "metadata": {},
   "source": [
    "Look at that great speedup! All done by a line of code to allow heterogenous computing. \n",
    "\n",
    "You can specify the device that you would like to execute on if you have a preferred device. See more in the [dpnp documentation](https://intelpython.github.io/dpnp/)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e0eb30c5-a79b-4e30-aa02-4337dc92d5df",
   "metadata": {},
   "source": [
    "## Data Parallel Extension for Numba* Enhanced DPNP Script for Further Performance and Compilation on Intel GPU"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "19540dfe-149b-4199-a803-0418aee2b24d",
   "metadata": {},
   "source": [
    "## Numba Performance Gains Made Easy on any SYCL device"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "40b86fda-a9b3-4683-94d9-84df3a556889",
   "metadata": {},
   "source": [
    "With Data Parallel Extension for Numba* (numba-dpex), you can now run Python code with the Numba-JIT compiler to take advantage of these optimizations on any SYCL device, including Intel GPU! \n",
    "\n",
    "Let's try it out with the `@dpjit` function decorator, similar to the Numba `@jit` decorator in allowing JIT compilation but on any SYCL device rather than just CPU."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "fc1ac4a2-3250-44cb-92cf-a484aae39da0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Time to Compute Pairwise Distance with DPNP and Numba-DPEX on GPU: 0.007500171661376953\n"
     ]
    }
   ],
   "source": [
    "import dpnp as np\n",
    "from numba import prange\n",
    "from numba_dpex import dpjit\n",
    "import time\n",
    "\n",
    "@dpjit #using Numba JIT compiler optimizations, now on any SYCL device\n",
    "def pairwise_distance(data, distance):\n",
    "    float0 = data.dtype.type(0)\n",
    "\n",
    "    #prange used for parallel loops for further opt\n",
    "    for i in prange(data.shape[0]):\n",
    "        for j in range(data.shape[0]):\n",
    "            d = float0\n",
    "            for k in range(data.shape[1]):\n",
    "                d += (data[i, k] - data[j, k])**2\n",
    "\n",
    "            distance[j, i] = np.sqrt(d)\n",
    "\n",
    "data = np.random.ranf((10*1024,3)).astype(np.float32)\n",
    "distance = np.empty(shape=(data.shape[0], data.shape[0]), dtype=np.float32)\n",
    "\n",
    "#do compilation first run, we will calculate performance on subsequent runs\n",
    "#wrapper consistent with previous np and dpnp scripts\n",
    "pairwise_distance(data, distance)\n",
    "\n",
    "#time pairwise distance calculations\n",
    "start = time.time()\n",
    "pairwise_distance(data, distance)\n",
    "print(\"Time to Compute Pairwise Distance with DPNP and Numba-DPEX on GPU:\", time.time() - start)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4ae08f4f-c272-4ecc-95f6-c1a6a068237d",
   "metadata": {},
   "source": [
    "Wow, even more speedup than in our previous DPNP script from compilation optimizations and SYCL!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "484b0db2-bf77-4c13-bb24-fef643159b67",
   "metadata": {},
   "source": [
    "### For advanced programmers: SYCL Programming Made Simple by Data Parallel Extension for Numba*"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5956296a-36a0-4e61-9400-9559b0a5138b",
   "metadata": {},
   "source": [
    "For more advanced programmers, Python developers may want to take advantage of Numba-JIT optimizations and write efficient kernel functions for even more speed-up with  DPNP scripts. \n",
    "\n",
    "Now you can write SYCL kernels directly with the `@dpex.kernel` function decorator to use DPC++ SYCL run-time and compilation for near-native speed. No DPC++ SYCL language coding required!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "6c124ae6-135d-4543-98c0-e6e68ff6bdec",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Time to Compute Pairwise Distance with DPNP and Numba-DPEX written SYCL Kernel on GPU: 0.005505800247192383\n"
     ]
    }
   ],
   "source": [
    "import dpnp as np\n",
    "import numba_dpex as dpex\n",
    "import time\n",
    "\n",
    "# kernel programming similar to Numba's cuda.jit sub-module\n",
    "# API modeled after SYCL lanaguage and uses the DPC++ run-time\n",
    "@dpex.kernel\n",
    "def _pairwise_distance_kernel(data, distance):\n",
    "    i = dpex.get_global_id(0)\n",
    "    j = dpex.get_global_id(1)\n",
    "\n",
    "    data_dims = data.shape[1]\n",
    "\n",
    "    d = data.dtype.type(0.0)\n",
    "    for k in range(data_dims):\n",
    "        tmp = data[i, k] - data[j, k]\n",
    "        d += tmp * tmp\n",
    "        \n",
    "    distance[j, i] = np.sqrt(d)\n",
    "\n",
    "# for this example, let's put the kernel we just wrote in a wrapper to call it for consistency with our previous scipts\n",
    "# you can call the kernel directly as well\n",
    "def pairwise_distance(data, distance):\n",
    "    _pairwise_distance_kernel[dpex.Range(data.shape[0], data.shape[0])](data, distance)\n",
    "\n",
    "\n",
    "data = np.random.ranf((10*1024,3)).astype(np.float32)\n",
    "distance = np.empty(shape=(data.shape[0], data.shape[0]), dtype=np.float32)\n",
    "\n",
    "\n",
    "#do compilation first run, we will calculate performance on subsequent runs\n",
    "#wrapper consistent with previous np and dpnp scripts\n",
    "pairwise_distance(data, distance)\n",
    "\n",
    "# alternatively, call kernel directly without wrapper:\n",
    "#_pairwise_distance_kernel[dpex.Range(data.shape[0], data.shape[0])](data, distance)\n",
    "\n",
    "#time pairwise distance calculations\n",
    "start = time.time()\n",
    "pairwise_distance(data, distance)\n",
    "print(\"Time to Compute Pairwise Distance with DPNP and Numba-DPEX written SYCL Kernel on GPU:\", time.time() - start)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "762c8862-6d34-4aa4-bbe3-947902e00f4e",
   "metadata": {},
   "source": [
    "## For more information on Data Parallel Extensions for Python* language, please visit:\n",
    "\n",
    "- [Data Parallel Extensions for Python* language (dpep) Website](https://intelpython.github.io/DPEP/main/)\n",
    "- [Data Parallel Extension for NumPy* (dpnp) Documentation](https://intelpython.github.io/dpnp/)\n",
    "- [Data Parallel Extension for Numba* (numba-dpex) Documentation](https://intelpython.github.io/numba-dpex/latest/index.html)\n",
    "- [Data Parallel Control (dpctl) Documentation](https://intelpython.github.io/dpctl/latest/index.html)\n",
    "\n",
    "For support: [Intel Machine Learning and Data Analytics Support Forum](https://community.intel.com/t5/Intel-Distribution-for-Python/bd-p/distribution-python)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "318dac21-5744-48e9-8238-cef5b78185da",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "language": "python",
   "name": "base"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}