notebooks/matmul/02_epilogs.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "fb0742d2-785b-4def-9939-9309b0f5c3e7",
   "metadata": {},
   "source": [
    "Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES\n",
    "\n",
    "SPDX-License-Identifier: BSD-3-Clause"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1388882b-0c36-4ac2-ac0b-9e909f543605",
   "metadata": {},
   "source": [
    "# Epilogs"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "292c2024-bd73-48de-8320-62c0ec0df645",
   "metadata": {},
   "source": [
    "In this tutorial, we will demonstrate the use of cuBLAS *epilogs*. An epilog is a simple computation executed after performing the matrix multiplication. Epilogs are usually much faster than the same computation executed manually on the result."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "73306371-66d8-499d-a022-a747323a654a",
   "metadata": {},
   "outputs": [],
   "source": [
    "from nvmath.linalg.advanced import MatmulEpilog as Epilog"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "58b562ec-c5fb-4779-80e8-d9508013b3ea",
   "metadata": {},
   "source": [
    "Let us begin by listing the available epilogs:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "c95a8dd3-d57a-4e68-94f0-837037196159",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'DEFAULT, RELU, RELU_AUX, BIAS, RELU_BIAS, RELU_AUX_BIAS, DRELU, DRELU_BGRAD, GELU, GELU_AUX, GELU_BIAS, GELU_AUX_BIAS, DGELU, DGELU_BGRAD, BGRADA, BGRADB'"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\", \".join(e.name for e in Epilog)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "953b9ef8-f212-44fc-81ef-fea6bd582c08",
   "metadata": {},
   "source": [
    "Epilogs can be specified by `epilog=` keyword argument of `matmul` (or `Matmul.plan`, if you are using stateful API)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c2aeb960-de6b-42ed-9899-088c8bcddd71",
   "metadata": {},
   "source": [
    "Now, let us describe in detail few of the epilogs available."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e5d6dc1e-370c-4af1-90ad-2240889a09cf",
   "metadata": {},
   "source": [
    "## RELU"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0a0361f1-0cad-4834-8815-50bd10c23874",
   "metadata": {},
   "source": [
    "`RELU` epilog executes ReLU (Rectified Linear Unit) on the result. It replaces all negative elements with zeros:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "c6251e5e-2681-466e-8cd7-238440ed831f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "A @ B:\n",
      "[[ 19. -22.]\n",
      " [-13.  14.]]\n",
      "\n",
      "relu(A @ B):\n",
      "[[19.  0.]\n",
      " [ 0. 14.]]\n"
     ]
    }
   ],
   "source": [
    "import cupy\n",
    "from nvmath.linalg.advanced import matmul\n",
    "\n",
    "a = cupy.asarray([[1, 2], [3, -4]], dtype=cupy.float32)\n",
    "b = cupy.asarray([[5, -6], [7, -8]], dtype=cupy.float32)\n",
    "\n",
    "print(\"A @ B:\")\n",
    "print(matmul(a, b))\n",
    "print()\n",
    "print(\"relu(A @ B):\")\n",
    "print(matmul(a, b, epilog=Epilog.RELU))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "54df0799-41d7-4347-b467-eb95c342eb23",
   "metadata": {},
   "source": [
    "## BIAS"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cc229c70-1c5e-4d4c-aa8c-77e0f95d1432",
   "metadata": {},
   "source": [
    "`BIAS` epilog adds 1-dimensional bias to the result, broadcasting it along the columns of the result. This epilog requires an extra input (the bias to add), which we can provide with `epilog_inputs` argument:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "e9ee3bd7-4ece-47f0-9a33-ca5d839e1945",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "a @ b:\n",
      "[[5. 5. 5. 5.]\n",
      " [5. 5. 5. 5.]\n",
      " [5. 5. 5. 5.]]\n",
      "\n",
      "a @ b + bias:\n",
      "[[ 5.  5.  5.  5.]\n",
      " [15. 15. 15. 15.]\n",
      " [25. 25. 25. 25.]]\n"
     ]
    }
   ],
   "source": [
    "a = cupy.asarray(cupy.ones((3, 5)), dtype=cupy.float32)\n",
    "b = cupy.asarray(cupy.ones((5, 4)), dtype=cupy.float32)\n",
    "bias = cupy.asarray([0, 10, 20], dtype=cupy.float32)\n",
    "\n",
    "print(\"a @ b:\")\n",
    "print(matmul(a, b))\n",
    "print()\n",
    "print(\"a @ b + bias:\")\n",
    "print(matmul(a, b, epilog=Epilog.BIAS, epilog_inputs={\"bias\": bias}))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0b900950-83fa-44a4-bef8-23b1d104eba0",
   "metadata": {},
   "source": [
    "## RELU_BIAS"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fdc619cf-2388-4220-87de-64ee978fe4e7",
   "metadata": {},
   "source": [
    "Also, there is a `RELU_BIAS` epilog which first adds the bias and then applies ReLU:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "92ca664f-63bd-4f9e-80fa-5fb4ab03f8d0",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "a @ b:\n",
      "[[ 19. -22.]\n",
      " [-13.  14.]]\n",
      "\n",
      "a @ b + bias:\n",
      "[[  9. -32.]\n",
      " [ -3.  24.]]\n",
      "\n",
      "relu(a @ b + bias):\n",
      "[[ 9.  0.]\n",
      " [ 0. 24.]]\n"
     ]
    }
   ],
   "source": [
    "a = cupy.asarray([[1, 2], [3, -4]], dtype=cupy.float32)\n",
    "b = cupy.asarray([[5, -6], [7, -8]], dtype=cupy.float32)\n",
    "bias = cupy.asarray([-10, 10], dtype=cupy.float32)\n",
    "\n",
    "# No epilog\n",
    "print(\"a @ b:\")\n",
    "print(matmul(a, b))\n",
    "print()\n",
    "\n",
    "# BIAS epilog, for reference only\n",
    "print(\"a @ b + bias:\")\n",
    "print(matmul(a, b, epilog=Epilog.BIAS, epilog_inputs={\"bias\": bias}))\n",
    "print()\n",
    "\n",
    "# RELU_BIAS epilog\n",
    "print(\"relu(a @ b + bias):\")\n",
    "print(matmul(a, b, epilog=Epilog.RELU_BIAS, epilog_inputs={\"bias\": bias}))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "78e1ad0d-0a59-4d24-bf6b-cab9a75e7521",
   "metadata": {},
   "source": [
    "## Gradient computations (`_AUX` epilogs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "39305501-4412-4147-8e2a-3204895e6c2c",
   "metadata": {},
   "source": [
    "As you may have noticed, the functions available in epilogs are commonly used in neural networks training. During backpropagation in a neural network, we need to compute the gradient of the transformations applied. For this reason, we need to store an auxiliary information about their inputs.\n",
    "\n",
    "For example, to compute the gradient of ReLU, we need to know which elements of the input matrix were negative.\n",
    "To get this information, we can use `RELU_AUX` epilog, which returns an auxiliary output indicating the sign of elements before ReLU.\n",
    "In case of ReLU, this auxiliary output is a bitmask for space optimization reasons. You don't need to unpack and interpret this bitmask\n",
    "in any way - there is a dedicated `DRELU` epilog which can handle this for you\n",
    "\n",
    "Let us show a simple example with `RELU_AUX`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "2a6e1270-7e4d-46c6-942e-453b80b9f5b8",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "a @ b:\n",
      "[[ 19. -22.]\n",
      " [-13.  14.]]\n",
      "\n",
      "relu(a @ b):\n",
      "[[19.  0.]\n",
      " [ 0. 14.]]\n",
      "\n",
      "aux is a <class 'dict'> with the following keys: dict_keys(['relu_aux'])\n",
      "aux[relu_aux] is <class 'cupy.ndarray'> with dtype int8\n"
     ]
    }
   ],
   "source": [
    "a = cupy.asarray([[1, 2], [3, -4]], dtype=cupy.float32)\n",
    "b = cupy.asarray([[5, -6], [7, -8]], dtype=cupy.float32)\n",
    "\n",
    "print(\"a @ b:\")\n",
    "print(matmul(a, b))\n",
    "print()\n",
    "\n",
    "print(\"relu(a @ b):\")\n",
    "result, aux = matmul(a, b, epilog=Epilog.RELU_AUX)  # Note that a tuple is returned!\n",
    "print(result)\n",
    "print()\n",
    "\n",
    "print(\"aux is a\", type(aux), \"with the following keys:\", aux.keys())\n",
    "print(\"aux[relu_aux] is\", type(aux[\"relu_aux\"]), \"with dtype\", aux[\"relu_aux\"].dtype)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e1b24e08-858d-4303-b540-3d1307f325da",
   "metadata": {},
   "source": [
    "## Learning more\n",
    "\n",
    "We will show a practical use case for the epilogs in the next tutorial, in which we will implement a simple digit recognition neural network using nvmath-python `matmul` and its epilogs.\n",
    "\n",
    "To learn more about the available epilogs, you can visit [cuBLAS documentation on epilogs](https://docs.nvidia.com/cuda/cublas/#cublasltepilogue-t)."
   ]
  }
 ],
 "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.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}