cudf/notebooks_numba_cuDF_integration.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Objective\n",
    "\n",
    "In my previous tutorial, I showed how to use the `apply_rows` and `apply_chunks` methods in cuDF to implement customized data transformations. Under the hood, they are all using [Numba library](https://numba.pydata.org/) to compile the normal python code into GPU kernels. Numba is an excellent python library that accelerates the numerical computations. Most importantly, Numba has direct CUDA programming support. For detailed information, refer to the [Numba CUDA documentation](https://numba.pydata.org/numba-doc/dev/cuda/index.html). As we know, the underlying data structure of cuDF is a GPU version of Apache Arrow. We can directly pass the GPU array around without the copying operation. Once we have the nice Numba library and standard GPU array, the sky is the limit. In this tutorial, I will show how to use Numba CUDA to accelerate cuDF data transformation and how to step by step accelerate it using CUDA programming tricks. \n",
    "\n",
    "The following experiments are performed at DGX V100 node."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## A simple example\n",
    "As usual, I am going to start with a simple example of doubling the numbers in an array:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<class 'numba.cuda.cudadrv.devicearray.DeviceNDArray'>\n"
     ]
    }
   ],
   "source": [
    "import cudf\n",
    "import numpy as np\n",
    "from numba import cuda\n",
    " \n",
    "array_len = 1000\n",
    "number_of_threads = 128\n",
    "number_of_blocks = (array_len + (number_of_threads - 1)) // number_of_threads\n",
    "df = cudf.DataFrame()\n",
    "df['in'] = np.arange(array_len, dtype=np.float64)\n",
    " \n",
    " \n",
    "@cuda.jit\n",
    "def double_kernel(result, array_len):\n",
    "    \"\"\"\n",
    "    double each element of the array\n",
    "    \"\"\"\n",
    "    i = cuda.grid(1)\n",
    "    if i < array_len:\n",
    "        result[i] = result[i] * 2.0\n",
    " \n",
    " \n",
    "before = df['in'].sum()\n",
    "gpu_array = df['in'].to_gpu_array()\n",
    "print(type(gpu_array))\n",
    "double_kernel[(number_of_blocks,), (number_of_threads,)](gpu_array, array_len)\n",
    "after = df['in'].sum()\n",
    "assert(np.isclose(before * 2.0, after))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "From the output of this code, it shows the underlying GPU array is of type `numba.cuda.cudadrv.devicearray.DeviceNDArray`. We can directly pass it to the kernel function that is compiled by the `cuda.jit`. Because we passed in the reference, the effect of number transformation will automatically show up in the original cuDF DataFrame. Note we have to manually enter the block size and grid size, which gives us the maximum of GPU programming control. The `cuda.grid` is a convenient method to compute the absolute position for the threads. It is equivalent to the normal `block_id * block_dim + thread_id` formula."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Practical example\n",
    "\n",
    "### Baseline\n",
    "\n",
    "We will work on the moving average problem as the last time. Because we have the full control of the grid and block size allocation, the vanilla moving average implementation code is much simpler compared to the `apply_chunks` implementation. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "%reset -s -f"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Numba with compile time 2.067620038986206\n",
      "Numba without compile time 1.9229750633239746\n",
      "pandas time 5.2932703495025635\n"
     ]
    }
   ],
   "source": [
    "import cudf\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from numba import cuda\n",
    "import numba\n",
    "import time\n",
    " \n",
    "array_len = int(5e8)\n",
    "average_window = 3000\n",
    "number_of_threads = 128\n",
    "number_of_blocks = (array_len + (number_of_threads - 1)) // number_of_threads\n",
    "df = cudf.DataFrame()\n",
    "df['in'] = np.arange(array_len, dtype=np.float64)\n",
    "df['out'] = np.arange(array_len, dtype=np.float64)\n",
    " \n",
    " \n",
    "@cuda.jit\n",
    "def kernel1(in_arr, out_arr, average_length, arr_len):\n",
    "    s = numba.cuda.local.array(1, numba.float64)\n",
    "    s[0] = 0.0\n",
    "    i = cuda.grid(1)\n",
    "    if i < arr_len:\n",
    "        if i < average_length-1:\n",
    "            out_arr[i] = np.inf\n",
    "        else:\n",
    "            for j in range(0, average_length):\n",
    "                s[0] += in_arr[i-j]\n",
    "            out_arr[i] = s[0] / np.float64(average_length)\n",
    " \n",
    " \n",
    "gpu_in = df['in'].to_gpu_array()\n",
    "gpu_out = df['out'].to_gpu_array()\n",
    "start = time.time()\n",
    "kernel1[(number_of_blocks,), (number_of_threads,)](gpu_in, gpu_out,\n",
    "                                                   average_window, array_len)\n",
    "cuda.synchronize()\n",
    "end = time.time()\n",
    "print('Numba with compile time', end-start)\n",
    " \n",
    "start = time.time()\n",
    "kernel1[(number_of_blocks,), (number_of_threads,)](gpu_in, gpu_out,\n",
    "                                                   average_window, array_len)\n",
    "cuda.synchronize()\n",
    "end = time.time()\n",
    "print('Numba without compile time', end-start)\n",
    " \n",
    "pdf = pd.DataFrame()\n",
    "pdf['in'] = np.arange(array_len, dtype=np.float64)\n",
    "start = time.time()\n",
    "pdf['out'] = pdf.rolling(average_window).mean()\n",
    "end = time.time()\n",
    "print('pandas time', end-start)\n",
    " \n",
    "assert(np.isclose(pdf.out.values[average_window:].mean(),\n",
    "       df.out.to_array()[average_window:].mean()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note, in order to compare the computation time accurately, I launch the kernel twice. The first time kernel launching will include the kernel compilation time. In this example, it takes 1.9s for the kernel to run without compilation. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Use shared memory\n",
    "\n",
    "In the baseline code, each thread is reading the numbers from the global memory. When doing the moving average, the same number is read multiple times by different threads. GPU global memory IO, in this case, is the speed bottleneck. To mitigate it, we load the data into shared memory for each of the computation blocks. Then the threads are doing summation from the numbers in the cache. To do the moving average for the elements at the beginning of the array, we make sure to load the `average_window` more data in the shared_memory. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "%reset -s -f"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Numba with compile time 1.3115026950836182\n",
      "Numba without compile time 1.085998773574829\n",
      "pandas time 5.594487428665161\n"
     ]
    }
   ],
   "source": [
    "import cudf\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from numba import cuda\n",
    "import numba\n",
    "import time\n",
    " \n",
    "array_len = int(5e8)\n",
    "average_window = 3000\n",
    "number_of_threads = 128\n",
    "number_of_blocks = (array_len + (number_of_threads - 1)) // number_of_threads\n",
    "shared_buffer_size = number_of_threads + average_window - 1\n",
    "df = cudf.DataFrame()\n",
    "df['in'] = np.arange(array_len, dtype=np.float64)\n",
    "df['out'] = np.arange(array_len, dtype=np.float64)\n",
    " \n",
    " \n",
    "@cuda.jit\n",
    "def kernel1(in_arr, out_arr, average_length, arr_len):\n",
    "    block_size = cuda.blockDim.x\n",
    "    shared = cuda.shared.array(shape=(shared_buffer_size),\n",
    "                               dtype=numba.float64)\n",
    "    i = cuda.grid(1)\n",
    "    tx = cuda.threadIdx.x\n",
    "    # Block id in a 1D grid\n",
    "    bid = cuda.blockIdx.x\n",
    "    starting_id = bid * block_size\n",
    " \n",
    "    shared[tx + average_length - 1] = in_arr[i]\n",
    "    cuda.syncthreads()\n",
    "    for j in range(0, average_length - 1, block_size):\n",
    "        if (tx + j) < average_length - 1:\n",
    "            shared[tx + j] = in_arr[starting_id -\n",
    "                                                 average_length + 1 +\n",
    "                                                 tx + j]\n",
    "    cuda.syncthreads()\n",
    " \n",
    "    s = numba.cuda.local.array(1, numba.float64)\n",
    "    s[0] = 0.0\n",
    "    if i < arr_len:\n",
    "        if i < average_length-1:\n",
    "            out_arr[i] = np.inf\n",
    "        else:\n",
    "            for j in range(0, average_length):\n",
    "                s[0] += shared[tx + average_length - 1 - j]\n",
    "            out_arr[i] = s[0] / np.float64(average_length)\n",
    " \n",
    " \n",
    "gpu_in = df['in'].to_gpu_array()\n",
    "gpu_out = df['out'].to_gpu_array()\n",
    "start = time.time()\n",
    "kernel1[(number_of_blocks,), (number_of_threads,)](gpu_in, gpu_out,\n",
    "                                                   average_window, array_len)\n",
    "cuda.synchronize()\n",
    "end = time.time()\n",
    " \n",
    "print('Numba with compile time', end-start)\n",
    " \n",
    "start = time.time()\n",
    "kernel1[(number_of_blocks,), (number_of_threads,)](gpu_in, gpu_out,\n",
    "                                                   average_window, array_len)\n",
    "cuda.synchronize()\n",
    "end = time.time()\n",
    "print('Numba without compile time', end-start)\n",
    " \n",
    "pdf = pd.DataFrame()\n",
    "pdf['in'] = np.arange(array_len, dtype=np.float64)\n",
    "start = time.time()\n",
    "pdf['out'] = pdf.rolling(average_window).mean()\n",
    "end = time.time()\n",
    "print('pandas time', end-start)\n",
    " \n",
    "assert(np.isclose(pdf.out.values[average_window:].mean(),\n",
    "       df.out.to_array()[average_window:].mean()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Running this, the computation time is reduced to 1.09s without kernel compilation time. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Reduced redundant summations\n",
    "\n",
    "Each thread in the above code is doing one moving average in a for-loop. It is easy to see that there are a lot of redundant summation operations done by different threads. To reduce the redundancy, the following code is changed to let each thread to compute a consecutive number of moving averages. The later moving average step is able to reuse the sum of the previous steps. This eliminated `thread_tile` number of for-loops.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "%reset -s -f"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Numba with compile time 0.6331000328063965\n",
      "Numba without compile time 0.30219364166259766\n",
      "pandas time 6.03054666519165\n"
     ]
    }
   ],
   "source": [
    "import cudf\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from numba import cuda\n",
    "import numba\n",
    "import time\n",
    " \n",
    "array_len = int(5e8)\n",
    "average_window = 3000\n",
    "number_of_threads = 64\n",
    "thread_tile = 48\n",
    "number_of_blocks = (array_len + (number_of_threads * thread_tile - 1)) // (number_of_threads * thread_tile)\n",
    "shared_buffer_size = number_of_threads * thread_tile + average_window - 1\n",
    "df = cudf.DataFrame()\n",
    "df['in'] = np.arange(array_len, dtype=np.float64)\n",
    "df['out'] = np.arange(array_len, dtype=np.float64)\n",
    " \n",
    " \n",
    "@cuda.jit\n",
    "def kernel1(in_arr, out_arr, average_length, arr_len):\n",
    "    block_size = cuda.blockDim.x\n",
    "    shared = cuda.shared.array(shape=(shared_buffer_size),\n",
    "                               dtype=numba.float64)\n",
    "    tx = cuda.threadIdx.x\n",
    "    # Block id in a 1D grid\n",
    "    bid = cuda.blockIdx.x\n",
    "    starting_id = bid * block_size * thread_tile\n",
    " \n",
    "    for j in range(thread_tile):\n",
    "        shared[tx + j * block_size + average_length - 1] = in_arr[starting_id\n",
    "                                                                   + tx +\n",
    "                                                                   j * block_size]\n",
    "        cuda.syncthreads()\n",
    "    for j in range(0, average_length - 1, block_size):\n",
    "        if (tx + j) < average_length - 1:\n",
    "            shared[tx + j] = in_arr[starting_id -\n",
    "                                                 average_length + 1 +\n",
    "                                                 tx + j]\n",
    "    cuda.syncthreads()\n",
    " \n",
    "    s = numba.cuda.local.array(1, numba.float64)\n",
    "    first = False\n",
    "    s[0] = 0.0\n",
    "    for k in range(thread_tile):\n",
    "        i = starting_id + tx * thread_tile + k\n",
    "        if i < arr_len:\n",
    "            if i < average_length-1:\n",
    "                out_arr[i] = np.inf\n",
    "            else:\n",
    "                if not first:\n",
    "                    for j in range(0, average_length):\n",
    "                        s[0] += shared[tx * thread_tile + k + average_length - 1 - j]\n",
    "                    s[0] = s[0] / np.float64(average_length)\n",
    "                    out_arr[i] = s[0]\n",
    "                    first = True\n",
    "                else:\n",
    "                    s[0] = s[0] + (shared[tx * thread_tile + k + average_length - 1]\n",
    "                                   - shared[tx * thread_tile + k + average_length - 1 - average_length])  / np.float64(average_length)\n",
    " \n",
    "                    out_arr[i] = s[0]\n",
    " \n",
    " \n",
    "gpu_in = df['in'].to_gpu_array()\n",
    "gpu_out = df['out'].to_gpu_array()\n",
    "start = time.time()\n",
    "kernel1[(number_of_blocks,), (number_of_threads,)](gpu_in, gpu_out,\n",
    "                                                   average_window, array_len)\n",
    "cuda.synchronize()\n",
    "end = time.time()\n",
    "print('Numba with compile time', end-start)\n",
    " \n",
    "start = time.time()\n",
    "kernel1[(number_of_blocks,), (number_of_threads,)](gpu_in, gpu_out,\n",
    "                                                   average_window, array_len)\n",
    "cuda.synchronize()\n",
    "end = time.time()\n",
    "print('Numba without compile time', end-start)\n",
    " \n",
    "pdf = pd.DataFrame()\n",
    "pdf['in'] = np.arange(array_len, dtype=np.float64)\n",
    "start = time.time()\n",
    "pdf['out'] = pdf.rolling(average_window).mean()\n",
    "end = time.time()\n",
    "print('pandas time', end-start)\n",
    " \n",
    "assert(np.isclose(pdf.out.values[average_window:].mean(),\n",
    "       df.out.to_array()[average_window:].mean()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "After this change, the computation time is reduced to 0.3s without kernel compilation time, we achieved a total of 6x speedup compared with the baseline."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Conclusion\n",
    "\n",
    "In this tutorial, we take advantage of the CUDA programming model in the Numba library to do a moving average computation. We show by using a few CUDA programming tricks, that we can achieve a **6x** speed up in moving average computations for long arrays.\n",
    "\n",
    "cuDF is a powerful tool for data scientists to use that provides both a high-level API that covers most of the use cases as well as access to its lower-level components. Those components, including gpu_array and Numba integration, make the cuDF library very flexible for processing data in a customized way. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.6.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}