{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook shows how to compute forward Layer Norm (training) + clamped ReLU (with bitmask generation), then compute the backward equivalent (DReLU + DLN) using the bitmask. "
   ]
  },
  {
   "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 and Setup\n",
    "This notebook requires an NVIDIA GPU. If `nvidia-smi` fails, go to Runtime -> Change runtime type -> Hardware accelerator and confirm a GPU is selected."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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": null,
   "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",
    "Create a cudnn handle, which is a per device handle used to initialize cudnn context."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import cudnn\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import sys\n",
    "\n",
    "torch.manual_seed(1)\n",
    "handle = cudnn.create_handle()\n",
    "\n",
    "print(\"Running with cudnn backend version:\", cudnn.backend_version())\n",
    "\n",
    "assert torch.cuda.is_available()\n",
    "\n",
    "assert (\n",
    "    cudnn.backend_version() >= 91100\n",
    "), \"LayerNorm with relu bitmask generation is only supported cuDNN version 9.11.0 or above\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### LayerNorm Relu Bitmask Training\n",
    "Problem Sizes:\n",
    "- Batch Size: 4\n",
    "- Sequence Size: 1024\n",
    "- Embedding Dimension: 768"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch, seq_size, embedding_dim = 4, 1024, 768\n",
    "\n",
    "input_type = torch.float32\n",
    "\n",
    "# Epsilon is a small number to prevent division by 0.\n",
    "epsilon_value = 1e-3\n",
    "\n",
    "# Set clamped ReLU limits\n",
    "lower_clip_val = 0\n",
    "upper_clip_val = 6"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create input tensor GPU buffers. We use PyTorch to allocate GPU tensors so we can reuse them easily when we calculate reference outputs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Allocate input tensor memory, initialize them to random numbers\n",
    "x_gpu = torch.randn(\n",
    "    batch * seq_size,\n",
    "    embedding_dim,\n",
    "    1,\n",
    "    1,\n",
    "    dtype=input_type,\n",
    "    requires_grad=True,\n",
    "    device=\"cuda\",\n",
    ").to(memory_format=torch.channels_last)\n",
    "\n",
    "scale_gpu = torch.randn(\n",
    "    1, embedding_dim, 1, 1, dtype=input_type, requires_grad=True, device=\"cuda\"\n",
    ").to(memory_format=torch.channels_last)\n",
    "\n",
    "bias_gpu = torch.randn(\n",
    "    1, embedding_dim, 1, 1, dtype=input_type, requires_grad=True, device=\"cuda\"\n",
    ").to(memory_format=torch.channels_last)\n",
    "\n",
    "# Allocate output tensor memory.\n",
    "out_gpu = torch.empty_like(x_gpu)\n",
    "mean_gpu = torch.empty(batch * seq_size, dtype=torch.float32, device=\"cuda\")\n",
    "inv_var_gpu = torch.empty(batch * seq_size, dtype=torch.float32, device=\"cuda\")\n",
    "\n",
    "# CuDNN stores boolean bitmask values as bit-packed int8_t.\n",
    "mask_gpu = torch.empty(\n",
    "    ((batch * seq_size), embedding_dim // 8, 1, 1), dtype=torch.uint8, device=\"cuda\"\n",
    ")\n",
    "\n",
    "# Epsilon, lower clip, and upper clip must be a scalar value on the cpu.\n",
    "epsilon_cpu = torch.full(\n",
    "    (1, 1, 1, 1), epsilon_value, dtype=torch.float32, requires_grad=False, device=\"cpu\"\n",
    ")\n",
    "lower_clip_cpu = torch.full(\n",
    "    (1, 1, 1, 1), lower_clip_val, dtype=torch.float32, requires_grad=False, device=\"cpu\"\n",
    ")\n",
    "upper_clip_cpu = torch.full(\n",
    "    (1, 1, 1, 1), upper_clip_val, dtype=torch.float32, requires_grad=False, device=\"cpu\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Compute reference ouputs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create the reference computation outputs here before the cuDNN computation, in order to use .empty_like() to create our output buffers\n",
    "x_ref = x_gpu.clone().float()\n",
    "\n",
    "normalized_x = torch.nn.functional.layer_norm(\n",
    "    x_ref,\n",
    "    [embedding_dim, 1, 1],\n",
    "    weight=scale_gpu.squeeze(0),\n",
    "    bias=bias_gpu.squeeze(0),\n",
    "    eps=epsilon_value,\n",
    ")\n",
    "\n",
    "out_expected = torch.clamp(normalized_x, min=lower_clip_val, max=upper_clip_val)\n",
    "\n",
    "mask_expected = (lower_clip_val < out_expected) & (out_expected < upper_clip_val)\n",
    "\n",
    "mean_expected = x_gpu.to(torch.float32).mean(dim=(1, 2, 3))\n",
    "\n",
    "inv_var_expected = torch.rsqrt(\n",
    "    torch.var(x_gpu.to(torch.float32), dim=(1, 2, 3)) + epsilon_value\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Create cuDNN graph and tensors\n",
    "\n",
    "Here we assign UIDs for tensors. UIDs are a unique identifier that will allow us to provide a mapping from tensors created from cuDNN graph api calls, such as `graph.tensor_like()`, to the underlying device memory that will be used to store these tensors. Virtual tensors don't require explicit memory allocated for them, but non-vritual tensors like inputs or outputs will need to have UIDs assigned to them. \n",
    "\n",
    "Alternatively, one can use handles directly in the mapping, however using UIDs can be more convinient for caching of cuDNN graphs.\n",
    "\n",
    "For each of our inputs {X, Scale, Bias, Epsilon} and our outputs {Out, Mean, Inverse Variance}, we allocate a UID. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create the cuDNN graph.\n",
    "graph = cudnn.pygraph(\n",
    "    handle=handle,\n",
    "    io_data_type=cudnn.data_type.FLOAT,\n",
    "    intermediate_data_type=cudnn.data_type.FLOAT,\n",
    "    compute_data_type=cudnn.data_type.FLOAT,\n",
    ")\n",
    "\n",
    "# Create tensor handles with the graph API, assign UIDs.\n",
    "x = graph.tensor_like(x_gpu.detach()).set_name(\"X\")\n",
    "scale = graph.tensor_like(scale_gpu.detach()).set_name(\"scale\")\n",
    "bias = graph.tensor_like(bias_gpu.detach()).set_name(\"bias\")\n",
    "epsilon = graph.tensor_like(epsilon_cpu).set_name(\"epsilon\")\n",
    "lower_clip = graph.tensor_like(lower_clip_cpu).set_name(\"lower_clip\")\n",
    "upper_clip = graph.tensor_like(upper_clip_cpu).set_name(\"upper_clip\")\n",
    "\n",
    "# Add a layernorm operation\n",
    "(norm_out, mean, inv_var) = graph.layernorm(\n",
    "    name=\"layernorm\",\n",
    "    input=x,\n",
    "    norm_forward_phase=cudnn.norm_forward_phase.TRAINING,\n",
    "    scale=scale,\n",
    "    bias=bias,\n",
    "    epsilon=epsilon,\n",
    ")\n",
    "\n",
    "# Add a relu operation\n",
    "out = graph.relu(\n",
    "    name=\"relu\", input=norm_out, lower_clip=lower_clip_val, upper_clip=upper_clip_val\n",
    ")\n",
    "\n",
    "# Add logical operations for generating bitmask\n",
    "lower_clip_mask = graph.cmp_gt(\n",
    "    name=\"cmp_gt_lower_clip\", input=out, comparison=lower_clip\n",
    ")\n",
    "lower_clip_mask.set_name(\"lower_clip\").set_data_type(cudnn.data_type.BOOLEAN)\n",
    "upper_clip_mask = graph.cmp_lt(\n",
    "    name=\"cmp_lt_upper_clip\", input=out, comparison=upper_clip\n",
    ")\n",
    "upper_clip_mask.set_name(\"upper_clip\").set_data_type(cudnn.data_type.BOOLEAN)\n",
    "bitmask = graph.logical_and(name=\"and_bitmask\", a=lower_clip_mask, b=upper_clip_mask)\n",
    "bitmask.set_name(\"upper_clip\").set_data_type(cudnn.data_type.BOOLEAN)\n",
    "\n",
    "# Enable all outputs, by default outputs are disabled\n",
    "out.set_name(\"output\").set_output(True)\n",
    "mean.set_name(\"mean\").set_output(True).set_data_type(cudnn.data_type.FLOAT)\n",
    "inv_var.set_name(\"inv_var\").set_output(True).set_data_type(cudnn.data_type.FLOAT)\n",
    "bitmask.set_name(\"relu_bitmask\").set_output(True)\n",
    "\n",
    "print(graph)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Build the graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Build the graph\n",
    "graph.build([cudnn.heur_mode.A])\n",
    "\n",
    "# To run this block more than once, we need to re-run the previous block to get a new graph.\n",
    "# The same instance of a graph can not be built twice."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Execute the graph"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "After validating and building a cuDNN graph,  we can now execute it. To do this, we have to provide input and output buffers. We do this by using the previously allocated UIDs to associate between tensor handles generated from the graph API, and their underlying memory. \n",
    "\n",
    "The desired input values need to be stored in these buffers before the `graph.execute` call. Because we have done a reference computation, we can simply reuse the buffers we have allocated via PyTorch.\n",
    "\n",
    "Note that the EPISLON UID expects a cpu buffer, "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "variant_pack = {\n",
    "    x: x_gpu,\n",
    "    scale: scale_gpu,\n",
    "    bias: bias_gpu,\n",
    "    epsilon: epsilon_cpu,\n",
    "    out: out_gpu,\n",
    "    mean: mean_gpu,\n",
    "    inv_var: inv_var_gpu,\n",
    "    lower_clip: lower_clip_cpu,\n",
    "    upper_clip: upper_clip_cpu,\n",
    "    bitmask: mask_gpu,\n",
    "}\n",
    "workspace = torch.empty(graph.get_workspace_size(), device=\"cuda\", dtype=torch.uint8)\n",
    "graph.execute(variant_pack, workspace)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Test cuDNN's output against PyTorch's and check correctness"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def unpack_cudnn_bitmask(bitmask_tensor, N, C, H=1, W=1):\n",
    "    \"\"\"\n",
    "    Helper function to unpack a cuDNN bitmask tensor of shape [N, C//8, H, W] and dtype=torch.uint8 (stored as packed bits)\n",
    "    into a boolean tensor of shape [N, C, H, W] for assert testing.\n",
    "    \"\"\"\n",
    "    bitmask_flat = bitmask_tensor.view(N, C // 8, H * W)\n",
    "    unpacked = torch.zeros(\n",
    "        (N, C, H * W), dtype=torch.bool, device=bitmask_tensor.device\n",
    "    )\n",
    "\n",
    "    for bit in range(8):\n",
    "        bit_values = (bitmask_flat >> bit) & 1\n",
    "        unpacked[:, bit::8, :] = bit_values\n",
    "\n",
    "    unpacked = unpacked.view(N, C, H, W)\n",
    "    return unpacked"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.cuda.synchronize()\n",
    "\n",
    "# compare to reference output\n",
    "torch.testing.assert_close(out_gpu, out_expected, rtol=5e-3, atol=5e-3)\n",
    "torch.testing.assert_close(inv_var_gpu, inv_var_expected, rtol=5e-3, atol=5e-3)\n",
    "torch.testing.assert_close(mean_gpu, mean_expected, rtol=5e-3, atol=5e-3)\n",
    "\n",
    "# Unpack the bitmask tensor and compare to reference output\n",
    "unpacked_mask = unpack_cudnn_bitmask(mask_gpu, batch * seq_size, embedding_dim, 1, 1)\n",
    "torch.testing.assert_close(unpacked_mask, mask_expected, atol=1e-3, rtol=1e-3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "### LayerNorm Relu Bitmask Backward Pass\n",
    "\n",
    "Compute references values for backward graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Reference backward operation using PyTorch\n",
    "target = torch.randn_like(out_expected)\n",
    "criterion = torch.nn.MSELoss()\n",
    "loss = criterion(out_expected, target)\n",
    "\n",
    "out_expected.retain_grad()\n",
    "x_gpu.retain_grad()\n",
    "scale_gpu.retain_grad()\n",
    "bias_gpu.retain_grad()\n",
    "\n",
    "loss.backward()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Create cuDNN graph and tensors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bwd_graph = cudnn.pygraph(\n",
    "    handle=handle,\n",
    "    intermediate_data_type=cudnn.data_type.FLOAT,\n",
    "    compute_data_type=cudnn.data_type.FLOAT,\n",
    ")\n",
    "\n",
    "# Create tensors associated with the backwards graph. DO NOT reuse tensor handles from the forward graph.\n",
    "d_out = bwd_graph.tensor(\n",
    "    name=\"d_out\", dim=x_gpu.size(), stride=x_gpu.stride(), data_type=x_gpu.dtype\n",
    ")\n",
    "\n",
    "x_bwd = bwd_graph.tensor_like(x, name=\"x\")\n",
    "scale_bwd = bwd_graph.tensor_like(scale, name=\"scale\")\n",
    "mean_bwd = bwd_graph.tensor_like(mean, name=\"mean\")\n",
    "inv_var_bwd = bwd_graph.tensor_like(inv_var, name=\"inv_var\")\n",
    "bitmask_bwd = bwd_graph.tensor(\n",
    "    name=\"bitmask\",\n",
    "    dim=(batch * seq_size, embedding_dim, 1, 1),\n",
    "    stride=(embedding_dim, 1, 1, 1),\n",
    "    data_type=cudnn.data_type.BOOLEAN,\n",
    ")\n",
    "\n",
    "# Add a pointwise mul operation for dRelu using the bitmask\n",
    "drelu_dY = bwd_graph.mul(name=\"drelu_bitmask_mul\", a=d_out, b=bitmask_bwd)\n",
    "drelu_dY.set_name(\"dRelu(dY)\")\n",
    "print(\"drelu_x_bwd:\", drelu_dY.get_dim())\n",
    "\n",
    "# Add the adaptive layernorm backward operation\n",
    "(d_x, d_scale, d_bias) = bwd_graph.layernorm_backward(\n",
    "    name=\"DLN\",\n",
    "    grad=drelu_dY,\n",
    "    input=x_bwd,\n",
    "    scale=scale_bwd,\n",
    "    mean=mean_bwd,\n",
    "    inv_variance=inv_var_bwd,\n",
    ")\n",
    "\n",
    "# Enable outputs.\n",
    "d_x.set_output(True).set_data_type(x_gpu.dtype)\n",
    "d_scale.set_output(True).set_data_type(x_gpu.dtype)\n",
    "d_bias.set_output(True).set_data_type(x_gpu.dtype)\n",
    "\n",
    "print(bwd_graph)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Build the graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Build the bwd_graph\n",
    "bwd_graph.build([cudnn.heur_mode.A])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Execute the graph "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create output buffers for gradients\n",
    "d_x_gpu = torch.empty_like(x_gpu)\n",
    "d_scale_gpu = torch.empty_like(scale_gpu)\n",
    "d_bias_gpu = torch.empty_like(bias_gpu)\n",
    "\n",
    "workspace = torch.empty(\n",
    "    bwd_graph.get_workspace_size(), device=\"cuda\", dtype=torch.uint8\n",
    ")\n",
    "\n",
    "# For the inputs of the backwards graph (x_bwd, d_out, scale_bwd, mean_bwd, inv_var_bwd), we use the outputs of the forwards graph. For d_out we use pytorches autograd .grad functionality.\n",
    "variant_pack = {\n",
    "    x_bwd: x_gpu.detach(),\n",
    "    scale_bwd: scale_gpu.detach(),\n",
    "    d_out: out_expected.grad,\n",
    "    mean_bwd: mean_gpu.detach(),\n",
    "    inv_var_bwd: inv_var_gpu.detach(),\n",
    "    d_x: d_x_gpu,\n",
    "    d_scale: d_scale_gpu,\n",
    "    d_bias: d_bias_gpu,\n",
    "    bitmask_bwd: mask_gpu.detach(),\n",
    "}\n",
    "bwd_graph.execute(variant_pack, workspace)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Test cuDNN's output against PyTorch's and check correctness"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.cuda.synchronize()\n",
    "\n",
    "# compare to reference output\n",
    "torch.testing.assert_close(x_gpu.grad, d_x_gpu, atol=2e-4, rtol=2e-4)\n",
    "torch.testing.assert_close(scale_gpu.grad, d_scale_gpu, atol=2e-4, rtol=2e-4)\n",
    "torch.testing.assert_close(bias_gpu.grad, d_bias_gpu, atol=2e-4, rtol=2e-4)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Perform Cleanup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cudnn.destroy_handle(handle)"
   ]
  }
 ],
 "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": 4
}
