hands-on/cuda/03_Kernel_Tuner_Advanced.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "fxycY0vBU8oQ"
   },
   "source": [
    "# Kernel Tuner Tutorial\n",
    "\n",
    "## Advanced Hands-on\n",
    "\n",
    "In this hands-on we will look at few of the features of Kernel Tuner that have been recently introduced to you: **search optimization strategies** and **custom observers**.\n",
    "\n",
    "But first, if you have not done it already, it is time to install and import `kernel_tuner` and its dependencies."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "95uRehlTY67Z"
   },
   "outputs": [],
   "source": [
    "%pip install kernel_tuner\n",
    "\n",
    "import numpy as np\n",
    "import kernel_tuner as kt\n",
    "import collections"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "GmCh0nrscYZL"
   },
   "source": [
    "To work with these features we will use a matrix multiplication kernel.\n",
    "\n",
    "Matrix multiplication is one of the most well-known and widely-used linear algebra operations, and is frequently used to demonstrate the high-performance computing capabilities of GPUs. As such, matrix multiplication presents a familiar starting point for many GPU programmers. More information about matrix multiplication can be found on [Wikipedia](https://en.wikipedia.org/wiki/Matrix_multiplication)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "dkekhvw_cZu0"
   },
   "outputs": [],
   "source": [
    "%%writefile matmul.cu\n",
    "\n",
    "#define WIDTH 512\n",
    "\n",
    "__global__ void matmul_kernel(float *C, float *A, float *B) {\n",
    "\n",
    "    __shared__ float sA[block_size_y*tile_size_y][block_size_x];\n",
    "    __shared__ float sB[block_size_y*tile_size_y][block_size_x * tile_size_x];\n",
    "\n",
    "    int tx = threadIdx.x;\n",
    "    int ty = threadIdx.y;\n",
    "    int x = blockIdx.x * block_size_x * tile_size_x + threadIdx.x;\n",
    "    int y = blockIdx.y * block_size_y * tile_size_y + threadIdx.y;\n",
    "    int k, kb;\n",
    "\n",
    "    float sum[tile_size_y][tile_size_x];\n",
    "    #pragma unroll\n",
    "    for (int i = 0; i < tile_size_y; i++) {\n",
    "        #pragma unroll\n",
    "        for (int j = 0; j < tile_size_x; j++) {\n",
    "            sum[i][j] = 0.0f;\n",
    "        }\n",
    "    }\n",
    "\n",
    "    for (k = 0; k < WIDTH; k += block_size_x) {\n",
    "\n",
    "        __syncthreads();\n",
    "        #pragma unroll\n",
    "        for (int i = 0; i < tile_size_y; i++) {\n",
    "            sA[ty + block_size_y * i][tx] = A[(y+i*block_size_y) * WIDTH + k + tx];\n",
    "\n",
    "            #pragma unroll\n",
    "            for (int j = 0; j < tile_size_x; j++) {\n",
    "                sB[ty + block_size_y * i][tx + j * block_size_x] = B[(k + ty + block_size_y * i) * WIDTH + x + j * block_size_x];\n",
    "            }\n",
    "        }\n",
    "        __syncthreads();\n",
    "\n",
    "        //compute\n",
    "        #pragma unroll\n",
    "        for (kb = 0; kb < block_size_x; kb++) {\n",
    "\n",
    "            #pragma unroll\n",
    "            for (int i = 0; i < tile_size_y; i++) {\n",
    "            #pragma unroll\n",
    "                for (int j = 0; j < tile_size_x; j++) {\n",
    "                    sum[i][j] += sA[ty + block_size_y * i][kb] * sB[kb][tx + j * block_size_x];\n",
    "                }\n",
    "            }\n",
    "\n",
    "        }\n",
    "\n",
    "    }\n",
    "\n",
    "    //store result\n",
    "    #pragma unroll\n",
    "    for (int i = 0; i < tile_size_y; i++) {\n",
    "        #pragma unroll\n",
    "        for (int j = 0; j < tile_size_x; j++) {\n",
    "            C[y * WIDTH + x + block_size_y * i * WIDTH + j * block_size_x] = sum[i][j];\n",
    "        }\n",
    "    }\n",
    "\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "NGfOQTE4idrv"
   },
   "source": [
    "We now allocate memory, define tunable parameters and constraints, and tune the kernel."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "FgSLwkmGinyg"
   },
   "outputs": [],
   "source": [
    "# matrix width needs to match the value in the kernel source\n",
    "problem_size = (512, 512)\n",
    "\n",
    "A = np.random.randn(*problem_size).astype(np.float32)\n",
    "B = np.random.randn(*problem_size).astype(np.float32)\n",
    "C = np.zeros_like(A)\n",
    "\n",
    "args = [C, A, B]\n",
    "\n",
    "tune_params = collections.OrderedDict()\n",
    "tune_params[\"block_size_x\"] = [2**i for i in range(0, 11)]\n",
    "tune_params[\"block_size_y\"] = [2**i for i in range(0, 11)]\n",
    "tune_params[\"tile_size_x\"] = [2**i for i in range(0, 6)]\n",
    "tune_params[\"tile_size_y\"] = [2**i for i in range(0, 6)]\n",
    "\n",
    "restrict = [\"block_size_x == block_size_y * tile_size_y\"]\n",
    "\n",
    "grid_div_x = [\"block_size_x\", \"tile_size_x\"]\n",
    "grid_div_y = [\"block_size_y\", \"tile_size_y\"]\n",
    "\n",
    "answer = [np.matmul(A,B), None, None]\n",
    "\n",
    "metrics = collections.OrderedDict()\n",
    "metrics[\"GFLOP/s\"] = lambda p : (2 * 512**3 / 1e9) / (p[\"time\"] / 1e3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "W6ZqN7GP6Gwe"
   },
   "outputs": [],
   "source": [
    "results, env = kt.tune_kernel(\"matmul_kernel\", \"matmul.cu\",\n",
    "                             problem_size, args, tune_params,\n",
    "                             grid_div_y=grid_div_y, grid_div_x=grid_div_x,\n",
    "                             answer=answer, atol=1e-4,\n",
    "                             restrictions=restrict, verbose=True, iterations=32, metrics=metrics, lang=\"cupy\", cache=\"matmul_cache.json\")\n",
    "print(f\"Number of configurations: {len(results)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "eK2IdErE_hL_"
   },
   "source": [
    "We can also visualize the tuning results using [KTdashboard](https://github.com/KernelTuner/dashboard)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "5uIYKswZKVe6"
   },
   "outputs": [],
   "source": [
    "%pip install git+https://github.com/KernelTuner/dashboard\n",
    "import panel as pn\n",
    "pn.extension(comms='colab')\n",
    "import ktdashboard.ktdashboard as ktd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "2cxXYJ2KKv2U"
   },
   "outputs": [],
   "source": [
    "ktd.KTdashboard(\"matmul_cache.json\").notebook()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "WfHCkXDM16NQ"
   },
   "source": [
    "There are times when the amount of possible configurations of tunable parameters is too high, or other time constraints do not allow to perform a full search. In those cases, it could be beneficial to use one of Kernel Tuner **search optimization strategies**.\n",
    "\n",
    "You can experiment with them in the next block. Try different strategies, and compare the optimum found with the overall optimum found previously. You can also time the tuning process to see the differences there.\n",
    "\n",
    "The strategies and how to enable them are described in Kernel Tuner's [API](https://KernelTuner.github.io/kernel_tuner/stable/user-api.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "fE8v8_Bg3XlS"
   },
   "outputs": [],
   "source": [
    "# experiment with enabling a search optimization strategy\n",
    "strategy = \"\"\n",
    "\n",
    "# tell the strategy to compile and benchmark at most 40 kernel configurations\n",
    "strategy_options = dict(max_fevals=40)\n",
    "\n",
    "results_opt, env_opt = kt.tune_kernel(\"matmul_kernel\", \"matmul.cu\",\n",
    "                                      problem_size, args, tune_params,\n",
    "                                      grid_div_y=grid_div_y, grid_div_x=grid_div_x,\n",
    "                                      answer=answer, atol=1e-4,\n",
    "                                      restrictions=restrict, verbose=True, iterations=32,\n",
    "                                      metrics=metrics, lang=\"cupy\",\n",
    "                                      strategy=strategy, strategy_options=strategy_options)\n",
    "print(f\"Number of configurations: {len(results_opt)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "uzVlmRtD6cBI"
   },
   "source": [
    "Next we are going to add a **custom observer** to the kernel. One possibility is to add an observer to compute the number of registers used by the kernel, and add this value to the metrics.\n",
    "\n",
    "In order to create a new observer it is necessary to extend the class `BenchmarkObserver` provided by Kernel Tuner in the `kt.observers` package. In case you want to access the number of registers used by a kernel instance, this is available inside your observer class in `self.dev.func.num_regs`.\n",
    "\n",
    "As usual, how to add observers is described in Kernel Tuner's [API](https://KernelTuner.github.io/kernel_tuner/stable/user-api.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "A9p8PPXoAZYU"
   },
   "outputs": [],
   "source": [
    "observers = []\n",
    "\n",
    "# add a custom observer\n",
    "from kernel_tuner.observers import BenchmarkObserver\n",
    "\n",
    "# define your own observer class that extends BenchmarkObserver\n",
    "#class CustomObserver(BenchmarkObserver)\n",
    "\n",
    "# implement the get_results method of this class\n",
    "# ...\n",
    "\n",
    "# create an instance of your custom observer\n",
    "custom_observer = None\n",
    "\n",
    "# append it to the list of observers by uncommenting the line below\n",
    "#observers.append(custom_observer)\n",
    "\n",
    "# add a metric so that our observed number of registers appears in the console output\n",
    "#metrics[\"regs\"] = lambda p:p[\"num_regs\"]\n",
    "\n",
    "\n",
    "# add an NVMLObserver\n",
    "from kernel_tuner.nvml import NVMLObserver\n",
    "nvml_observer = NVMLObserver([\"nvml_energy\", \"temperature\"])\n",
    "\n",
    "observers.append(nvml_observer)\n",
    "\n",
    "# add metrics to enable console output for observed quantities\n",
    "metrics[\"GFLOPS/W\"] = lambda p : (2 * 512**3 / 1e9) / (p[\"nvml_energy\"])\n",
    "metrics[\"T\"] = lambda p:p[\"temperature\"]\n",
    "\n",
    "# call tune_kernel to tune using our new Observers and additional metrics\n",
    "results, env = kt.tune_kernel(\"matmul_kernel\", \"matmul.cu\",\n",
    "                             problem_size, args, tune_params,\n",
    "                             observers=observers,\n",
    "                             grid_div_y=grid_div_y, grid_div_x=grid_div_x,\n",
    "                             answer=answer, atol=1e-4,\n",
    "                             restrictions=restrict, verbose=True, iterations=32, metrics=metrics, lang=\"cupy\")\n",
    "print(f\"Number of configurations: {len(results)}\")"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "03-Kernel_Tuner-Advanced.ipynb",
   "provenance": []
  },
  "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.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}