{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Matrix multiplication operation with fused bias using cudnn FE\n",
    "This notebook shows how a matmul operation with fused bias can be done using cudnn."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/NVIDIA/cudnn-frontend/blob/main/samples/python/01_matmul_bias.ipynb)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prerequisites for running on Colab\n",
    "This notebook requires an NVIDIA GPU H100 or newer. If `nvidia-smi` fails, go to Runtime -> Change runtime type -> Hardware accelerator and confirm a GPU is selected."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "# get_ipython().system('nvidia-smi')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If running on Colab, you will need to install the cudnn python interface."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "# get_ipython().system('pip install nvidia-cudnn-cu12')\n",
    "# get_ipython().system('pip install nvidia-cudnn-frontend')\n",
    "# get_ipython().system('pip3 install --pre torch --index-url https://download.pytorch.org/whl/nightly/cu128')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### General Setup\n",
    "We are going to call the cudnn through torch in this example. In general any dlpack tensor should work.\n",
    "cudnn handle is a per device handle used to initialize cudnn context.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "import cudnn\n",
    "import torch\n",
    "import sys\n",
    "\n",
    "handle = cudnn.create_handle()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Create input tensors and calculate reference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch, m, n, k = 16, 128, 128, 512\n",
    "\n",
    "input_type = torch.float16\n",
    "\n",
    "# input tensors\n",
    "a = torch.randn(batch, m, k, dtype=input_type, device=\"cuda\")\n",
    "b = torch.randn(batch, k, n, dtype=input_type, device=\"cuda\")\n",
    "B = torch.randn(1, m, n, dtype=torch.float16, device=\"cuda\")\n",
    "\n",
    "# reference output\n",
    "c_ref = torch.matmul(a, b) + B\n",
    "\n",
    "# place holder for cudnn output\n",
    "c = torch.randn_like(c_ref, device=\"cuda\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Define the hash of the given graph in terms of tensors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "def matmul_cache_key(handle, a, b, bias):\n",
    "    \"\"\"Custom key function for matmul + bias\"\"\"\n",
    "    return (\n",
    "        tuple(a.shape),\n",
    "        tuple(b.shape),\n",
    "        tuple(a.stride()),\n",
    "        tuple(b.stride()),\n",
    "        a.dtype,\n",
    "        b.dtype,\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Create cudnn matmul + bias fused graph."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "@cudnn.jit(heur_modes=[cudnn.heur_mode.A, cudnn.heur_mode.B])\n",
    "@cudnn.graph_cache(key_fn=matmul_cache_key)\n",
    "def create_matmul_bias_graph(handle, a, b, bias):\n",
    "    with cudnn.graph(handle) as (g, _):\n",
    "        a_cudnn = g.tensor_like(a)\n",
    "        b_cudnn = g.tensor_like(b)\n",
    "        bias_cudnn = g.tensor_like(bias)\n",
    "        c_cudnn = g.matmul(name=\"matmul\", A=a_cudnn, B=b_cudnn)\n",
    "        out = g.bias(name=\"bias\", input=c_cudnn, bias=bias_cudnn)\n",
    "        out.set_output(True).set_data_type(cudnn.data_type.HALF)\n",
    "\n",
    "    return g, [a_cudnn, b_cudnn, bias_cudnn, out]  # Return raw graph and tensors"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Build the graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "g, uids = create_matmul_bias_graph(handle, a, b, B)\n",
    "\n",
    "a_uid, b_uid, bias_uid, out_uid = uids"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Execute the code"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "variant_pack = {\n",
    "    a_uid: a,\n",
    "    b_uid: b,\n",
    "    bias_uid: B,\n",
    "    out_uid: c,\n",
    "}\n",
    "\n",
    "workspace = torch.empty(g.get_workspace_size(), device=\"cuda\", dtype=torch.uint8)\n",
    "g.execute(variant_pack, workspace)\n",
    "torch.cuda.synchronize()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.testing.assert_close(c, c_ref, rtol=5e-3, atol=5e-3)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "fe_0",
   "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.9.21"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
