lecture_018/FusedKernels.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "62bca490-d730-4aad-9cd9-1a7551ea2eae",
   "metadata": {},
   "source": [
    "# Fused Kernels"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8643efce-dc8f-43a2-b28c-27a7c41f66f3",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "## Abstract\n",
    "\n",
    "With focus on performance to get the most out of hardware, fusing of kernels has been a popular technique. At times, researchers/practitioners will re-write their code in native cuda or cpu kernels to get optimal performance, but projects such as torch.compile aim to make this simpler. Talk will focus on generating fused kernels and how to leverage torch.compile to be able to do that. We will shift a bit from all LLM talk and look into recommendation algorithms. In the process, we will work on creating fused kernels (triton and cuda) with the help of `torch.compile`. \n",
    "\n",
    "\n",
    "## About Me\n",
    "\n",
    "Software Engineer @ Meta building Rec Sys components for modelling lifecycle such as preprocessing, training, eval, inference, etc using Pytorch. \n",
    "\n",
    "If you want to get in touch with me:\n",
    "\n",
    "- Discord (@sk4301)\n",
    "- Linkedin (https://www.linkedin.com/in/sharma-k/)\n",
    "- Twitter (@kapil_sh_)\n",
    "- Github (https://github.com/kapilsh)\n",
    "\n",
    "## How is the talk structured\n",
    "\n",
    "- Dive into DLRM (open source deep rec sys model)\n",
    "- Build it from scrarch\n",
    "- Go through some some profiler traces\n",
    "- torch.compile\n",
    "- Writing fused kernels\n",
    "- Case Study: LoRA"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "43570ae8-5020-480c-84cb-4004722a51f4",
   "metadata": {},
   "source": [
    "# Setup"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "be61622a-2bcb-4f93-910b-639e10d0b69f",
   "metadata": {},
   "source": [
    "## Code and other artifacts\n",
    "\n",
    "- Lecture code: https://github.com/kapilsh/cuda-mode-lecture\n",
    "- How to open chrome trace: chrome://tracing\n",
    "- DLRM Blog Post: https://ai.meta.com/blog/dlrm-an-advanced-open-source-deep-learning-recommendation-model/\n",
    "- DLRM Paper: https://arxiv.org/pdf/1906.00091\n",
    "- DLRM github repo: https://github.com/facebookresearch/dlrm\n",
    "- Criteo Dataset: https://ailab.criteo.com/download-criteo-1tb-click-logs-dataset/\n",
    "- [Pytorch Profiler with Tensorboard](https://pytorch.org/tutorials/intermediate/tensorboard_profiler_tutorial.html?source=post_page-----2cb7e0fef30e--------------------------------)\n",
    "- [TORCH_LOGS with torch.compile](https://pytorch.org/tutorials/recipes/torch_logs.html#beta-using-torch-logs-python-api-with-torch-compile)\n",
    "- LoRA Paper: https://arxiv.org/abs/2106.09685\n",
    "- LoRA from scratch: https://lightning.ai/lightning-ai/studios/code-lora-from-scratch\n",
    "- Netron: https://netron.app/\n",
    "- GPUs go brrr https://horace.io/brrr_intro.html\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d80a7390-000f-46b8-99fb-a7896ffda565",
   "metadata": {},
   "source": [
    "## DLRM (Deep Learning Recommendation Model)\n",
    "\n",
    "### MODEL ARCHITECTURE \n",
    "\n",
    "- Takes a bunch of dense and sparse features\n",
    "- Dense features feed to a dense MLP layer\n",
    "- Sparse features feed into embedding layers\n",
    "- Interaction layers between dense NN layers and spares embeddings combine dense and sparse outputs\n",
    "- Interaction \"features\" output prediction click/no-click as output from an MLP \n",
    "\n",
    "![DLRM Model](./data/dlrm_model.png)\n",
    "\n",
    "### System Constrants\n",
    "\n",
    "![System Constraints](./data/66324023_2056206621174067_2937830378620059648_n.gif)\n",
    "\n",
    "- D Dense Features MLP output\n",
    "- S Sparse Features\n",
    "- E Embedding Dimension on average\n",
    "\n",
    "Interaction: O((D * S * E) ^ 2)\n",
    "\n",
    "### Criteo Dataset\n",
    "\n",
    "- Training dataset with 24 days of ad display and click data (positive: clicked and negatives: non-clicked)\n",
    "- 13 features taking integer values (mostly count features)\n",
    "- 26 anonymized categorical features\n",
    "- Corresponding Kaggle competition: https://www.kaggle.com/c/criteo-display-ad-challenge"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d76eb45b-1ee6-4257-98a6-5185f57987a6",
   "metadata": {},
   "source": [
    "# Exploring DLRM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "a8b063e1-6f04-4665-ac70-8508641d37b6",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import time\n",
    "from dataclasses import dataclass\n",
    "from typing import Mapping, List, Dict, Union\n",
    "\n",
    "import click\n",
    "import torch\n",
    "import torch._dynamo\n",
    "from loguru import logger\n",
    "from torch import nn, Tensor\n",
    "from torch.utils.data import DataLoader\n",
    "\n",
    "from criteo_dataset import CriteoParquetDataset\n",
    "from model import DenseArch, read_metadata, SparseArch, DenseSparseInteractionLayer, PredictionLayer, Parameters, DLRM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "1069bd19-a050-4c9b-8d1f-051e3bd7e319",
   "metadata": {},
   "outputs": [],
   "source": [
    "file_path = \"./data/sample_criteo_data.parquet\"\n",
    "metadata_path = \"./data/sample_criteo_metadata.json\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "42b72f15-8b72-4345-8d2f-a71ad4b8f01e",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024-05-12 06:47:32.120 | INFO     | __main__:<module>:1 - Reading the parquet file ./data/sample_criteo_data.parquet...\n",
      "2024-05-12 06:47:32.121 | INFO     | __main__:<module>:2 - Reading the metadata file ./data/sample_criteo_metadata.json...\n",
      "2024-05-12 06:47:33.233 | INFO     | __main__:<module>:7 - Labels size: torch.Size([2])\n",
      "2024-05-12 06:47:33.234 | INFO     | __main__:<module>:8 - Dense size: torch.Size([2, 13])\n",
      "2024-05-12 06:47:33.234 | INFO     | __main__:<module>:9 - Sparse size: torch.Size([2, 26])\n"
     ]
    }
   ],
   "source": [
    "logger.info(\"Reading the parquet file {}...\".format(file_path))\n",
    "logger.info(\"Reading the metadata file {}...\".format(metadata_path))\n",
    "\n",
    "dataset = CriteoParquetDataset(file_path)\n",
    "data_loader = DataLoader(dataset, batch_size=2, shuffle=False)\n",
    "labels, dense, sparse = next(iter(data_loader))\n",
    "logger.info(\"Labels size: {}\".format(labels.size()))\n",
    "logger.info(\"Dense size: {}\".format(dense.size()))\n",
    "logger.info(\"Sparse size: {}\".format(sparse.size()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "884bc4d3-a489-45a1-b269-141db29a3b7c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[5.0000e+00, 1.1000e+02, 0.0000e+00, 1.6000e+01, 0.0000e+00, 1.0000e+00,\n",
       "         0.0000e+00, 1.4000e+01, 7.0000e+00, 1.0000e+00, 0.0000e+00, 3.0600e+02,\n",
       "         0.0000e+00],\n",
       "        [3.2000e+01, 3.0000e+00, 5.0000e+00, 0.0000e+00, 1.0000e+00, 0.0000e+00,\n",
       "         0.0000e+00, 6.1000e+01, 5.0000e+00, 0.0000e+00, 1.0000e+00, 3.1570e+03,\n",
       "         5.0000e+00]])"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dense"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "c896c0f1-3773-4af5-bcc4-1aa1b17fb409",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[1651969401, 3793706328, 2951365679, 2489089999,  951068488, 1875733963,\n",
       "          897624609,  679512323, 1189011366,  771915201,  209470001, 2509774111,\n",
       "           12976055, 3192841527, 2316006604, 1289502458, 3523761834, 3088518074,\n",
       "         2501034507, 3280875304,  351689309,  632402057, 3619814411, 2091868316,\n",
       "          809724924, 3977271069],\n",
       "        [3857972621, 2695561126, 1873417685, 3666490401, 1020698403, 1875733963,\n",
       "         2870406529, 1128426537,  502653268, 2112471209, 1716706404, 2582335015,\n",
       "           12976055, 3192841527, 4089183897, 1289502458, 3523761834, 2716538129,\n",
       "         2501034507, 4273985635, 2737978529, 3370249814,  391309800, 1966410890,\n",
       "         2568167914, 3075991895]])"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sparse"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "1e0476a3-0af3-4de7-985e-76a2f13d8da6",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024-05-12 06:47:33.257 | INFO     | __main__:<module>:7 - Dense out size: torch.Size([2, 16])\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "tensor([[ -19.3222,   16.1915,  -14.7218,  -48.3280,   -7.9050,   49.6840,\n",
       "          -32.6965,   20.7256,  -46.5937,  -18.3307,  -17.1108,   25.4100,\n",
       "          -16.2373,  -18.1572,   -2.6305,   19.7319],\n",
       "        [-205.2248,  222.4405,  -45.1590, -361.0247,  -57.2848,  468.3490,\n",
       "         -252.3489,    6.0886, -417.9690, -245.4523, -240.5372,  289.3233,\n",
       "         -145.1369, -100.0480,  -92.8139,  117.6506]],\n",
       "       grad_fn=<AddmmBackward0>)"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dense_mlp_out_size = 16\n",
    "num_dense_features = dense.size()[1]\n",
    "dense_arch = DenseArch(dense_feature_count=num_dense_features,\n",
    "                       dense_hidden_layers_sizes=[32],\n",
    "                       output_size=dense_mlp_out_size)\n",
    "dense_out = dense_arch(dense)\n",
    "logger.info(\"Dense out size: {}\".format(dense_out.size()))\n",
    "dense_out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "eca531ef-68ee-4773-821a-69518e3fed07",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024-05-12 06:47:35.113 | INFO     | __main__:<module>:11 - Sparse out size: torch.Size([2, 16])\n",
      "2024-05-12 06:47:35.113 | INFO     | __main__:<module>:11 - Sparse out size: torch.Size([2, 16])\n",
      "2024-05-12 06:47:35.114 | INFO     | __main__:<module>:11 - Sparse out size: torch.Size([2, 16])\n",
      "2024-05-12 06:47:35.114 | INFO     | __main__:<module>:11 - Sparse out size: torch.Size([2, 16])\n",
      "2024-05-12 06:47:35.114 | INFO     | __main__:<module>:11 - Sparse out size: torch.Size([2, 16])\n",
      "2024-05-12 06:47:35.114 | INFO     | __main__:<module>:11 - Sparse out size: torch.Size([2, 16])\n",
      "2024-05-12 06:47:35.115 | INFO     | __main__:<module>:11 - Sparse out size: torch.Size([2, 16])\n",
      "2024-05-12 06:47:35.115 | INFO     | __main__:<module>:11 - Sparse out size: torch.Size([2, 16])\n",
      "2024-05-12 06:47:35.115 | INFO     | __main__:<module>:11 - Sparse out size: torch.Size([2, 16])\n",
      "2024-05-12 06:47:35.116 | INFO     | __main__:<module>:11 - Sparse out size: torch.Size([2, 16])\n",
      "2024-05-12 06:47:35.116 | INFO     | __main__:<module>:11 - Sparse out size: torch.Size([2, 16])\n",
      "2024-05-12 06:47:35.116 | INFO     | __main__:<module>:11 - Sparse out size: torch.Size([2, 16])\n",
      "2024-05-12 06:47:35.116 | INFO     | __main__:<module>:11 - Sparse out size: torch.Size([2, 16])\n",
      "2024-05-12 06:47:35.117 | INFO     | __main__:<module>:11 - Sparse out size: torch.Size([2, 16])\n",
      "2024-05-12 06:47:35.117 | INFO     | __main__:<module>:11 - Sparse out size: torch.Size([2, 16])\n",
      "2024-05-12 06:47:35.117 | INFO     | __main__:<module>:11 - Sparse out size: torch.Size([2, 16])\n",
      "2024-05-12 06:47:35.118 | INFO     | __main__:<module>:11 - Sparse out size: torch.Size([2, 16])\n",
      "2024-05-12 06:47:35.118 | INFO     | __main__:<module>:11 - Sparse out size: torch.Size([2, 16])\n",
      "2024-05-12 06:47:35.118 | INFO     | __main__:<module>:11 - Sparse out size: torch.Size([2, 16])\n",
      "2024-05-12 06:47:35.119 | INFO     | __main__:<module>:11 - Sparse out size: torch.Size([2, 16])\n",
      "2024-05-12 06:47:35.120 | INFO     | __main__:<module>:11 - Sparse out size: torch.Size([2, 16])\n",
      "2024-05-12 06:47:35.120 | INFO     | __main__:<module>:11 - Sparse out size: torch.Size([2, 16])\n",
      "2024-05-12 06:47:35.120 | INFO     | __main__:<module>:11 - Sparse out size: torch.Size([2, 16])\n",
      "2024-05-12 06:47:35.121 | INFO     | __main__:<module>:11 - Sparse out size: torch.Size([2, 16])\n",
      "2024-05-12 06:47:35.121 | INFO     | __main__:<module>:11 - Sparse out size: torch.Size([2, 16])\n",
      "2024-05-12 06:47:35.121 | INFO     | __main__:<module>:11 - Sparse out size: torch.Size([2, 16])\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "tensor([[-0.6661,  0.9471,  0.2120, -0.9125,  0.6718, -0.6560, -1.9422,  0.4658,\n",
       "         -0.2885, -0.0596,  1.5758, -0.3784,  0.2411, -1.4830, -0.0242,  1.3030],\n",
       "        [ 0.4879,  1.3181, -1.8291,  1.3051, -0.3013,  0.0863, -0.5329,  0.4685,\n",
       "         -0.7085,  2.1623, -1.8609, -0.9163,  0.2337, -0.6242, -1.0802, -0.6486]],\n",
       "       grad_fn=<EmbeddingBackward0>)"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "metadata = read_metadata(metadata_path)\n",
    "embedding_size = 16\n",
    "embedding_sizes = {fn: embedding_size for fn in metadata.keys()}\n",
    "sparse_mlp_out_size = 16\n",
    "sparse_arch = SparseArch(metadata=metadata,\n",
    "                         embedding_sizes=embedding_sizes)\n",
    "# compiled model hangs on running with inputs\n",
    "# sparse_arch_optim = torch.compile(sparse_arch)\n",
    "sparse_out = sparse_arch(sparse)\n",
    "for v in sparse_out:\n",
    "    logger.info(\"Sparse out size: {}\".format(v.size()))\n",
    "sparse_out[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "303ed5c6-d3f4-47b6-8ec4-02b1fd889444",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024-05-12 06:47:35.127 | INFO     | __main__:<module>:3 - Dense sparse interaction out size: torch.Size([2, 186624])\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "tensor([[ 3.7335e+02, -3.1286e+02,  2.8446e+02,  ..., -2.1380e-01,\n",
       "          4.0430e-02,  3.9008e-01],\n",
       "        [ 4.2117e+04, -4.5650e+04,  9.2678e+03,  ...,  1.3297e-01,\n",
       "          2.5754e-01,  3.5446e-02]], grad_fn=<ViewBackward0>)"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dense_sparse_interaction_layer = DenseSparseInteractionLayer()\n",
    "ds_out = dense_sparse_interaction_layer(dense_out, sparse_out)\n",
    "logger.info(\"Dense sparse interaction out size: {}\".format(ds_out.size()))\n",
    "ds_out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "965bb04a-8025-4398-8924-7e2d7b1e857e",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024-05-12 06:47:35.152 | INFO     | __main__:<module>:5 - Prediction out size: torch.Size([2, 1])\n",
      "2024-05-12 06:47:35.153 | INFO     | __main__:<module>:6 - Prediction out value: tensor([[0.4819],\n",
      "        [0.0000]], grad_fn=<SigmoidBackward0>)\n"
     ]
    }
   ],
   "source": [
    "prediction_layer = PredictionLayer(dense_out_size=dense_mlp_out_size,\n",
    "                                   sparse_out_sizes=[sparse_mlp_out_size] * len(metadata),\n",
    "                                   hidden_sizes=[16])\n",
    "pred_out = prediction_layer(ds_out)\n",
    "logger.info(\"Prediction out size: {}\".format(pred_out.size()))\n",
    "logger.info(\"Prediction out value: {}\".format(pred_out))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "66fe15e0-15a1-4cc2-8406-a46ebb489c42",
   "metadata": {},
   "source": [
    "# Model Graph"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "e291be4b-6e30-4ba1-aedb-5b54cba4feeb",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "## Model Graph\n",
    "\n",
    "![Model Graph](./data/model_graph.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "03e1ac03-4806-4a78-ae31-31dd469e932d",
   "metadata": {},
   "source": [
    "# Profiling"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "031a5bea-c03b-466b-a507-9f3b4497bc7a",
   "metadata": {},
   "source": [
    "### Initial Setup: Simple 2 layered MLP used for each triangle\n",
    "\n",
    "### Baseline\n",
    "\n",
    "> python model_train.py --config ./model_hyperparameters_small.json\n",
    "\n",
    "### Initial Distribution - Naive Implementation of index_hash\n",
    "\n",
    "```\n",
    "...\n",
    "# mapping loaded as is from the metadata - so a python list\n",
    "self.mapping = [metadata[f\"SPARSE_{i}\"][\"tokenizer_values\"] for i in range(self.num_sparse_features)]\n",
    "...\n",
    "# in forward\n",
    "tokenizers = torch.tensor(tokenizer_values).reshape(1, -1)\n",
    "if input_tensor.is_cuda:\n",
    "   tokenizers = tokenizers.cuda()\n",
    "...\n",
    "```\n",
    "\n",
    "<img src=\"./perf_screenshots/pytorch_profile_initial_index_hash.png\" width=\"800\"/>\n",
    "*Pytorch Profiler trace (initial)*\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f9e58625-5f27-4183-975d-49badcd370d8",
   "metadata": {},
   "source": [
    "### Using tensorboard for high level info\n",
    "\n",
    "> tensorboard --logdir tb_logs --bind_all\n",
    "\n",
    "\n",
    "<img src=\"./perf_screenshots/summary_initial_index_hash.png\" width=\"800\"/>\n",
    "*Initial distribution of ops - summary from tensorboard*\n",
    "\n",
    "---\n",
    "\n",
    "### Tensor.item() takes a lot of running time\n",
    "\n",
    "- What's going on - what is _local_scalar_dense and why is item() taking so long?\n",
    "    - https://discuss.pytorch.org/t/tensor-item-takes-a-lot-of-running-time/16683\n",
    "    - https://discuss.pytorch.org/t/calling-loss-item-is-very-slow/99774 \n",
    "\n",
    "> CUDA_LAUNCH_BLOCKING=1 python model_train.py\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fc3cf0ac-4705-4010-8798-4ebc831ec413",
   "metadata": {},
   "source": [
    "### After passing `CUDA_LAUNCH_BLOCKING=1`\n",
    "\n",
    "> CUDA_LAUNCH_BLOCKING=1 python model_train.py --config ./model_hyperparameters_small.json\n",
    "\n",
    "<img src=\"./perf_screenshots/summary_initial_cuda_launch_blocking.png\" width=\"800\"/>\n",
    "\n",
    "*New distribution of ops after `CUDA_LAUNCH_BLOCKING=1` - summary from tensorboard*\n",
    "\n",
    "---\n",
    "\n",
    "model_hyperparameters_initial.1714869603606159866.pt.trace.json\n",
    "\n",
    "![Initial Index Hash Profile](./perf_screenshots/index_hash_profile_1.png)\n",
    "*Profile initial index hash implementation*\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "52c5ae46-4c45-49a6-b962-2153a6c39e85",
   "metadata": {},
   "source": [
    "\n",
    "### Improvements\n",
    "\n",
    "```\n",
    "# in ctor - Put metadata needed for model directly on the gpu\n",
    "self.mapping = [torch.tensor(metadata[f\"SPARSE_{i}\"][\"tokenizer_values\"], device=device) for i in\n",
    "                range(self.num_sparse_features)]\n",
    "\n",
    "# in forward - dont use reshape if you can avoid\n",
    "tokenizers = tokenizer_values.view(1, -1)\n",
    "```\n",
    "\n",
    "Ref Trace: model_hyperparameters_initial.1714870277384855181.pt.trace.json\n",
    "\n",
    "\n",
    "<img src=\"./perf_screenshots/improve_index_hash.png\" width=\"800\"/>\n",
    "\n",
    "*Profile after improvements*\n",
    "\n",
    "\n",
    "### What's next\n",
    "\n",
    "- Index hash seems pretty expensive\n",
    "- Can we improve/simplify the hash function/tokenization\n",
    "- Let's just calculate the modulus hash based on cardinality\n",
    "    - Maybe not representative of data if distribution is non uniform across categories (but that's fine for now) \n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a65ea844-886a-450d-b513-a84a36e597b3",
   "metadata": {},
   "source": [
    "### Using Modulus Hash\n",
    "\n",
    "```\n",
    "def modulus_hash(tensor: torch.Tensor, cardinality: torch.Tensor):\n",
    "    return (tensor + 1) % cardinality\n",
    "```\n",
    "\n",
    "\n",
    "<img src=\"./perf_screenshots/optimized_modulus_hash.png\" width=\"800\"/>\n",
    "\n",
    "*Pytorch Profiler trace for optimized modulus hash*\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "34436457-a6c9-4489-bc64-b487266c2ca4",
   "metadata": {},
   "source": [
    "# torch.compile"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4cbdc60a-5d38-4fd8-8498-3f9a7f283220",
   "metadata": {},
   "source": [
    "##  torch.compile DLRM\n",
    "\n",
    "> TORCH_COMPILE_DEBUG_DIR=/home/ksharma/logs TORCH_LOGS=recompiles,+dynamo,inductor,guards,graph_breaks python model.py\n",
    "\n",
    "> CUDA_LAUNCH_BLOCKING=1 python model_train.py\n",
    "\n",
    "- GPU utilization goes up\n",
    "- memory footprint goes down\n",
    "\n",
    "## Memory Footprint\n",
    "\n",
    "### Pre `torch.compile`\n",
    "\n",
    "<img src=\"./perf_screenshots/pre_torch_compile_initial.png\" width=\"800\"/>\n",
    "\n",
    "\n",
    "### Post `torch.compile`\n",
    "\n",
    "<img src=\"./perf_screenshots/post_torch_compile_initial.png\" width=\"800\"/>\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "51084a89-1bae-4fb5-ab33-47a5f40157ae",
   "metadata": {},
   "source": [
    "### Chrome Trace after `torch.compile`\n",
    "<img src=\"./perf_screenshots/pytorch_profile_torch_compile.png\" width=\"800\"/>\n",
    "\n",
    "*Pytorch Profile Trace after `torch.compile`"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7048a179-9077-4e60-9199-977392ced657",
   "metadata": {},
   "source": [
    "### Let's look deeper into what's going on\n",
    "\n",
    "<img src=\"./perf_screenshots/torch_compile_triton_kernels.png\" width=\"800\"/>\n",
    "\n",
    "*Custom triton kernel scheduled on the cuda stream*\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2c5571bf-9d9d-45cf-bed0-8cb9fd460629",
   "metadata": {},
   "source": [
    "### Increase complexity\n",
    "\n",
    "Source: https://ai.meta.com/blog/dlrm-an-advanced-open-source-deep-learning-recommendation-model/\n",
    "\n",
    "```shell\n",
    "python dlrm_s_pytorch.py --arch-sparse-feature-size=16 --arch-mlp-bot=\"13-512-256-64-16\" --arch-mlp-top=\"512-256-1\" --data-generation=dataset --data-set=kaggle --processed-data-file=./input/kaggle_processed.npz --loss-function=bce --round-targets=True --learning-rate=0.1 --mini-batch-size=128 --print-freq=1024 --print-time\n",
    "```\n",
    "\n",
    "### Let's change the model architecture\n",
    "\n",
    "- --arch-mlp-bot=\"13-512-256-64-16\"\n",
    "- --arch-mlp-top=\"512-256-1\"\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d2746386-b28c-49e8-a9f1-b013b52d397d",
   "metadata": {},
   "source": [
    "### Eager view\n",
    "\n",
    "<img src=\"./perf_screenshots/full_model_eager_view.png\" width=\"800\"/>\n",
    "\n",
    "*Full Eager Model - Pytorch Profiler trace*\n",
    "\n",
    "- Sparse Arch is now not the biggest piece of the pie\n",
    "- PredictionLayer is the highest\n",
    "    - Top MLP and sigmoid"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8db1dfe8-05a1-4c5f-a708-adba0ddd70fb",
   "metadata": {},
   "source": [
    "### `torch.compile` view\n",
    "\n",
    "<img src=\"./perf_screenshots/full_model_torch_compiled.png\" width=\"800\"/>\n",
    "\n",
    "*Full `torch.compile` Model - Pytorch Profiler trace*"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "19abe6a3-74fa-47b8-b439-1b4af946ffa0",
   "metadata": {},
   "source": [
    "# torch.compile -> triton code generation \n",
    "\n",
    "How do we see what is going on with the triton kernels"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d0b529a0-b65a-48aa-8ec9-dd58064aabc9",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "## Generate triton code\n",
    "\n",
    "> `TORCH_LOGS=output_code CUDA_LAUNCH_BLOCKING=1 python model_train.py --config ./model_hyperparameters_main.json --use_torch_compile`\n",
    "\n",
    "## Inspect\n",
    "\n",
    "- Prints generated code for you\n",
    "- Should see `... torch._inductor.graph.__output_code: [INFO] Output code written to: ...`\n",
    "- Shows source nodes from where the code was generated\n",
    "- Fused kernels:\n",
    "    - fused_relu\n",
    "    - fused_cat\n",
    "    - fused_embedding\n",
    "    - fused_sigmoid_squeeze\n",
    "- Reinterpret_tensor: https://github.com/pytorch/pytorch/blob/ca98c2a932132e49559bf777c02798633d585e66/torch/csrc/inductor/inductor_ops.cpp#L54\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "87661f85-04e4-4649-acb4-4a19061d8f05",
   "metadata": {},
   "source": [
    "## Write our own\n",
    "\n",
    "### Kernel\n",
    "\n",
    "```python\n",
    "@triton.jit\n",
    "def pointwise_add_relu_fusion_512(in_out_ptr0, in_ptr0, XBLOCK : tl.constexpr):\n",
    "    # Number of elements in in_out_ptr0 (B X N)\n",
    "    xnumel = 65536\n",
    "    # This program will process inputs that are offset from the initial data.\n",
    "    # For instance, if you had a strided tensor of 65536 i.e. 128 X 512 and XBLOCK = 512\n",
    "    # the programs will each access the elements [0:512, 512:1024, ...].\n",
    "    # i.e. offsets is a list of pointers:\n",
    "    # Question: Can you see how torch.compile is allocating blocks here? \n",
    "    # below we will call this N = 512\n",
    "    xoffset = tl.program_id(0) * XBLOCK\n",
    "    # block threads\n",
    "    xindex = xoffset + tl.arange(0, XBLOCK)[:]\n",
    "    # masks to guard against overflow\n",
    "    xmask = xindex < xnumel\n",
    "    # xindex will have elements from 0:N, N:2N where N = dense @ weights\n",
    "    x2 = xindex\n",
    "    # bias i.e. 1D tensor with only N elements\n",
    "    # mod will give the us the right \n",
    "    x0 = xindex % 512\n",
    "    # load the N elements\n",
    "    tmp0 = tl.load(in_out_ptr0 + (x2), xmask)\n",
    "    # load the 1D tensor\n",
    "    tmp1 = tl.load(in_ptr0 + (x0), xmask, eviction_policy='evict_last')\n",
    "    # result = bias + dense @ weights\n",
    "    tmp2 = tmp0 + tmp1\n",
    "    # relu: can also use tl.maximum\n",
    "    tmp3 = triton_helpers.maximum(0, tmp2) \n",
    "    # output moved over\n",
    "    tl.store(in_out_ptr0 + (x2), tmp3, None)\n",
    "```\n",
    "\n",
    "### Test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "d0f9d6ad-c572-44da-944a-8e11a875d477",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[1., 1., 1.],\n",
      "        [1., 1., 1.],\n",
      "        [1., 1., 1.]], device='cuda:0')\n",
      "tensor([1., 1., 1.], device='cuda:0')\n",
      "tensor([[2., 2., 2.],\n",
      "        [2., 2., 2.],\n",
      "        [2., 2., 2.]], device='cuda:0')\n",
      "tensor([[2., 2., 2.,  ..., 2., 2., 2.],\n",
      "        [2., 2., 2.,  ..., 2., 2., 2.],\n",
      "        [2., 2., 2.,  ..., 2., 2., 2.],\n",
      "        ...,\n",
      "        [2., 2., 2.,  ..., 2., 2., 2.],\n",
      "        [2., 2., 2.,  ..., 2., 2., 2.],\n",
      "        [2., 2., 2.,  ..., 2., 2., 2.]], device='cuda:0')\n"
     ]
    }
   ],
   "source": [
    "import triton\n",
    "import torch\n",
    "import triton.language as tl\n",
    "from torch._inductor import triton_helpers\n",
    "from torch._inductor.triton_heuristics import grid\n",
    "\n",
    "@triton.jit\n",
    "def pointwise_add_relu_fusion_512(in_out_ptr0, in_ptr0, XBLOCK : tl.constexpr):\n",
    "    xnumel = 65536\n",
    "    xoffset = tl.program_id(0) * XBLOCK\n",
    "    xindex = xoffset + tl.arange(0, XBLOCK)[:]\n",
    "    xmask = xindex < xnumel\n",
    "    # dense @ weights\n",
    "    x2 = xindex\n",
    "    # bias\n",
    "    x0 = xindex % 512\n",
    "    tmp0 = tl.load(in_out_ptr0 + (x2), xmask)\n",
    "    tmp1 = tl.load(in_ptr0 + (x0), xmask, eviction_policy='evict_last')\n",
    "    # bias + dense @ weights\n",
    "    tmp2 = tmp0 + tmp1\n",
    "    tmp3 = triton_helpers.maximum(0, tmp2)\n",
    "    tl.store(in_out_ptr0 + (x2), tmp3, None)\n",
    "\n",
    "\n",
    "torch.cuda.set_device(0)  # no-op to ensure context\n",
    "X = torch.ones(size=(128, 512), device='cuda')\n",
    "print(X[:3, :3])\n",
    "Y = torch.ones(size=(512,), device='cuda')\n",
    "print(Y[:3])\n",
    "eager_result = torch.maximum(X + Y, torch.tensor(0., device='cuda'))\n",
    "print(eager_result[:3, :3])\n",
    "pointwise_add_relu_fusion_512[grid(65536)](X, Y, 512)\n",
    "print(X)\n",
    "torch.testing.assert_close(X, eager_result, rtol=1e-4, atol=1e-4)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "32a42047-da22-4496-ad93-0b890d5851d0",
   "metadata": {},
   "source": [
    "## Cuda Kernel\n",
    "\n",
    "### Ask ChatGPT to generate the kernel for us\n",
    "\n",
    "![Chat GPT Input](./data/chatgpt_input.png)\n",
    "\n",
    "### ChatGPT output (without any changes)\n",
    "\n",
    "```cpp\n",
    "#include <cuda_fp16.h>\n",
    "\n",
    "__global__ void pointwise_add_relu_fusion_512(float* in_out_ptr0, const float* in_ptr0, const int XBLOCK) {\n",
    "    const int xnumel = 65536;\n",
    "    const int N = 512; // Value of N from the Triton kernel\n",
    "    const int tid = threadIdx.x;\n",
    "    const int xoffset = blockIdx.x * XBLOCK;\n",
    "    const int xindex = xoffset + tid;\n",
    "    const bool xmask = xindex < xnumel;\n",
    "    \n",
    "    if (xmask) {\n",
    "        int x2 = xindex;\n",
    "        int x0 = xindex % N;\n",
    "        \n",
    "        float tmp0 = in_out_ptr0[x2];\n",
    "        float tmp1 = in_ptr0[x0];\n",
    "        float tmp2 = tmp0 + tmp1;\n",
    "        float tmp3 = max(0.0f, tmp2); // ReLU operation\n",
    "        \n",
    "        in_out_ptr0[x2] = tmp3;\n",
    "    }\n",
    "}\n",
    "```\n",
    "\n",
    "### Let's run the generated CUDA kernel\n",
    "\n",
    "> NOTE: To run torch native, you can download it as below or add conda environment to $CMAKE_PREFIX_PATH\n",
    "\n",
    "```\n",
    "wget https://download.pytorch.org/libtorch/cu121/libtorch-cxx11-abi-shared-with-deps-2.2.1%2Bcu121.zip # Download torch native lib\n",
    "```\n",
    "\n",
    "### Build the cmake project"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "1f371d14-f084-4826-b750-f35ce40afe9b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "-- CMake version: 3.22.1\n",
      "-- Caffe2: CUDA detected: 12.1\n",
      "-- Caffe2: CUDA nvcc is: /home/ksharma/anaconda3/envs/cuda-learn/bin/nvcc\n",
      "-- Caffe2: CUDA toolkit directory: /home/ksharma/anaconda3/envs/cuda-learn\n",
      "-- Caffe2: Header version is: 12.1\n",
      "-- /home/ksharma/anaconda3/envs/cuda-learn/lib/libnvrtc.so shorthash is c993a6f1\n",
      "-- USE_CUDNN is set to 0. Compiling without cuDNN support\n",
      "-- USE_CUSPARSELT is set to 0. Compiling without cuSPARSELt support\n",
      "-- Autodetected CUDA architecture(s):  7.5\n",
      "-- Added CUDA NVCC flags for: -gencode;arch=compute_75,code=sm_75\n",
      "-- Configuring done\n",
      "-- Generating done\n",
      "-- Build files have been written to: /home/ksharma/dev/git/cuda-mode-lecture/kernels/cmake-build-debug\n",
      "ninja: no work to do.\n"
     ]
    }
   ],
   "source": [
    "! mkdir -p kernels/cmake-build-debug && cd kernels/cmake-build-debug && cmake .. -G Ninja && ninja"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "8ef7bde3-d0b8-435b-898b-d5baa009105d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Tensor x:\n",
      "-0.9247 -0.4253 -2.6438  0.1452 -0.1209 -0.5797 -0.6229 -0.3284 -1.0745 -0.3631\n",
      "-1.6711  2.2655  0.3117 -0.1842  1.2866  1.1820 -0.1271  1.2169  1.4353  1.0605\n",
      "-0.4941 -1.4244 -0.7244 -1.2973  0.0697 -0.0074  1.8969  0.6878 -0.0779 -0.8373\n",
      " 1.3506 -0.2879 -0.5965 -0.3283 -0.9086 -0.8059 -0.7407 -0.0504  0.5435  1.5150\n",
      " 0.0141  0.4532  1.6349  0.7124 -0.1806  1.0252 -1.4622 -0.7554 -0.1836  0.3824\n",
      " 0.3918 -0.0830  0.8971 -1.1123  0.1116  0.4863 -0.5499 -0.3231 -0.5469  0.9049\n",
      " 0.2837  0.1210  0.4730 -1.0823 -0.0334 -0.9734  0.9559 -1.1795 -1.0064  0.1160\n",
      " 0.6852 -0.4124 -0.6738 -0.5404  0.6898 -1.5517  0.3805 -0.0436  0.3597 -0.5043\n",
      "[ CUDAFloatType{8,10} ]\n",
      "Tensor y:\n",
      " 0.1808\n",
      "-0.5523\n",
      " 0.9238\n",
      "-0.7350\n",
      " 1.3800\n",
      " 0.8676\n",
      " 0.1297\n",
      "-0.9406\n",
      " 0.8109\n",
      " 0.8821\n",
      "[ CUDAFloatType{10} ]\n",
      "Expected:\n",
      " 0.0000  0.0000  0.0000  0.0000  1.2591  0.2879  0.0000  0.0000  0.0000  0.5189\n",
      " 0.0000  1.7132  1.2355  0.0000  2.6666  2.0496  0.0026  0.2763  2.2462  1.9425\n",
      " 0.0000  0.0000  0.1994  0.0000  1.4497  0.8602  2.0266  0.0000  0.7330  0.0448\n",
      " 1.5315  0.0000  0.3273  0.0000  0.4714  0.0617  0.0000  0.0000  1.3544  2.3971\n",
      " 0.1949  0.0000  2.5587  0.0000  1.1994  1.8929  0.0000  0.0000  0.6273  1.2644\n",
      " 0.5726  0.0000  1.8209  0.0000  1.4916  1.3539  0.0000  0.0000  0.2640  1.7869\n",
      " 0.4645  0.0000  1.3968  0.0000  1.3465  0.0000  1.0856  0.0000  0.0000  0.9980\n",
      " 0.8660  0.0000  0.2500  0.0000  2.0698  0.0000  0.5102  0.0000  1.1706  0.3778\n",
      "[ CUDAFloatType{8,10} ]\n",
      "Result:\n",
      " 0.0000  0.0000  0.0000  0.0000  1.2591  0.2879  0.0000  0.0000  0.0000  0.5189\n",
      " 0.0000  1.7132  1.2355  0.0000  2.6666  2.0496  0.0026  0.2763  2.2462  1.9425\n",
      " 0.0000  0.0000  0.1994  0.0000  1.4497  0.8602  2.0266  0.0000  0.7330  0.0448\n",
      " 1.5315  0.0000  0.3273  0.0000  0.4714  0.0617  0.0000  0.0000  1.3544  2.3971\n",
      " 0.1949  0.0000  2.5587  0.0000  1.1994  1.8929  0.0000  0.0000  0.6273  1.2644\n",
      " 0.5726  0.0000  1.8209  0.0000  1.4916  1.3539  0.0000  0.0000  0.2640  1.7869\n",
      " 0.4645  0.0000  1.3968  0.0000  1.3465  0.0000  1.0856  0.0000  0.0000  0.9980\n",
      " 0.8660  0.0000  0.2500  0.0000  2.0698  0.0000  0.5102  0.0000  1.1706  0.3778\n",
      "[ CUDAFloatType{8,10} ]\n",
      "All Match: true\n"
     ]
    }
   ],
   "source": [
    "!./kernels/cmake-build-debug/dlrm_kernels_test"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0b9afaf0-576a-4a57-b9ac-91ea779d75f4",
   "metadata": {},
   "source": [
    "### (OR) Run it locally with pytorch utils"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "81ca18d8-86c9-48aa-801f-b48a14bb9613",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "__global__ void add_relu_fusion_kernel(float* in_out_ptr0, const float* in_ptr0, const int xnumel ,const int XBLOCK) {\n",
      "    const int tid = threadIdx.x;\n",
      "    const int xoffset = blockIdx.x * XBLOCK;\n",
      "    const int xindex = xoffset + tid;\n",
      "    const bool xmask = xindex < xnumel;\n",
      "\n",
      "    if (xmask) {\n",
      "        int x2 = xindex;\n",
      "        int x0 = xindex % XBLOCK;\n",
      "        float tmp0 = in_out_ptr0[x2];\n",
      "        float tmp1 = in_ptr0[x0];\n",
      "        float tmp2 = tmp0 + tmp1;\n",
      "        float tmp3 = max(0.0f, tmp2); // ReLU operation\n",
      "\n",
      "        in_out_ptr0[x2] = tmp3;\n",
      "    }\n",
      "}\n",
      "\n",
      "torch::Tensor add_relu_fusion(torch::Tensor in_out, const torch::Tensor& in) {\n",
      "    auto sizes = in_out.sizes();\n",
      "    auto XBLOCK = sizes[1];\n",
      "    auto numel = in_out.numel();\n",
      "    dim3 threadsPerBlock(XBLOCK);\n",
      "    dim3 numBlocks((numel + XBLOCK - 1) / XBLOCK);\n",
      "    add_relu_fusion_kernel<<<numBlocks, threadsPerBlock>>>(in_out.data_ptr<float>(), in.data_ptr<float>(), numel, XBLOCK);\n",
      "    cudaDeviceSynchronize();\n",
      "    return std::move(in_out);\n",
      "}\n",
      "----\n",
      "\n",
      "torch::Tensor add_relu_fusion(torch::Tensor in_out, const torch::Tensor& in);\n"
     ]
    }
   ],
   "source": [
    "cuda_code_file = \"./kernels/src/pointwise_add_relu_fused.cu\"\n",
    "header_code_file = \"./kernels/src/pointwise_add_relu_fused.cuh\"\n",
    "\n",
    "with open(cuda_code_file) as f:\n",
    "    cuda_code = \"\".join([f for f in f.readlines() if not f.startswith(\"#include\")])\n",
    "    print(cuda_code)\n",
    "\n",
    "print(\"----\")\n",
    "\n",
    "with open(header_code_file) as f:\n",
    "    header_code = \"\".join([f for f in f.readlines() if not f.startswith(\"#include\")])\n",
    "    print(header_code)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "ed787f6f-23eb-4622-bdb2-4ae99f5ceabb",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "mkdir: cannot create directory ‘./build’: File exists\n"
     ]
    }
   ],
   "source": [
    "!mkdir ./build"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "87c2998a-f14a-4d8f-9eaa-554169c33e39",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Detected CUDA files, patching ldflags\n",
      "Emitting ninja build file ./build/build.ninja...\n",
      "Building extension module kernel_extension...\n",
      "Allowing ninja to set a default number of workers... (overridable by setting the environment variable MAX_JOBS=N)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[1/3] c++ -MMD -MF main.o.d -DTORCH_EXTENSION_NAME=kernel_extension -DTORCH_API_INCLUDE_EXTENSION_H -DPYBIND11_COMPILER_TYPE=\\\"_gcc\\\" -DPYBIND11_STDLIB=\\\"_libstdcpp\\\" -DPYBIND11_BUILD_ABI=\\\"_cxxabi1011\\\" -isystem /home/ksharma/anaconda3/envs/cuda-learn/lib/python3.11/site-packages/torch/include -isystem /home/ksharma/anaconda3/envs/cuda-learn/lib/python3.11/site-packages/torch/include/torch/csrc/api/include -isystem /home/ksharma/anaconda3/envs/cuda-learn/lib/python3.11/site-packages/torch/include/TH -isystem /home/ksharma/anaconda3/envs/cuda-learn/lib/python3.11/site-packages/torch/include/THC -isystem /home/ksharma/anaconda3/envs/cuda-learn/include -isystem /home/ksharma/anaconda3/envs/cuda-learn/include/python3.11 -D_GLIBCXX_USE_CXX11_ABI=0 -fPIC -std=c++17 -c /home/ksharma/dev/git/cuda-mode-lecture/build/main.cpp -o main.o \n",
      "[2/3] /home/ksharma/anaconda3/envs/cuda-learn/bin/nvcc --generate-dependencies-with-compile --dependency-output cuda.cuda.o.d -DTORCH_EXTENSION_NAME=kernel_extension -DTORCH_API_INCLUDE_EXTENSION_H -DPYBIND11_COMPILER_TYPE=\\\"_gcc\\\" -DPYBIND11_STDLIB=\\\"_libstdcpp\\\" -DPYBIND11_BUILD_ABI=\\\"_cxxabi1011\\\" -isystem /home/ksharma/anaconda3/envs/cuda-learn/lib/python3.11/site-packages/torch/include -isystem /home/ksharma/anaconda3/envs/cuda-learn/lib/python3.11/site-packages/torch/include/torch/csrc/api/include -isystem /home/ksharma/anaconda3/envs/cuda-learn/lib/python3.11/site-packages/torch/include/TH -isystem /home/ksharma/anaconda3/envs/cuda-learn/lib/python3.11/site-packages/torch/include/THC -isystem /home/ksharma/anaconda3/envs/cuda-learn/include -isystem /home/ksharma/anaconda3/envs/cuda-learn/include/python3.11 -D_GLIBCXX_USE_CXX11_ABI=0 -D__CUDA_NO_HALF_OPERATORS__ -D__CUDA_NO_HALF_CONVERSIONS__ -D__CUDA_NO_BFLOAT16_CONVERSIONS__ -D__CUDA_NO_HALF2_OPERATORS__ --expt-relaxed-constexpr -gencode=arch=compute_75,code=compute_75 -gencode=arch=compute_75,code=sm_75 --compiler-options '-fPIC' -O2 -std=c++17 -c /home/ksharma/dev/git/cuda-mode-lecture/build/cuda.cu -o cuda.cuda.o \n",
      "[3/3] c++ main.o cuda.cuda.o -shared -L/home/ksharma/anaconda3/envs/cuda-learn/lib/python3.11/site-packages/torch/lib -lc10 -lc10_cuda -ltorch_cpu -ltorch_cuda -ltorch -ltorch_python -L/home/ksharma/anaconda3/envs/cuda-learn/lib -lcudart -o kernel_extension.so\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading extension module kernel_extension...\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "from torch.utils.cpp_extension import load_inline\n",
    "\n",
    "cuda_extension = load_inline(\n",
    "    name='kernel_extension',\n",
    "    cpp_sources=header_code,\n",
    "    cuda_sources=cuda_code,\n",
    "    functions=[\"add_relu_fusion\"],\n",
    "    with_cuda=True,\n",
    "    verbose=True,\n",
    "    extra_cuda_cflags=[\"-O2\"],\n",
    "    build_directory='./build',\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "c4674f5f-5132-4703-ae1d-af28ef8f8b4b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['__doc__',\n",
       " '__file__',\n",
       " '__loader__',\n",
       " '__name__',\n",
       " '__package__',\n",
       " '__spec__',\n",
       " 'add_relu_fusion']"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dir(cuda_extension)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "2420a7ce-8ae6-4b09-8315-a58a8d9a6cdf",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[2., 2., 2.,  ..., 2., 2., 2.],\n",
      "        [2., 2., 2.,  ..., 2., 2., 2.],\n",
      "        [2., 2., 2.,  ..., 2., 2., 2.],\n",
      "        ...,\n",
      "        [2., 2., 2.,  ..., 2., 2., 2.],\n",
      "        [2., 2., 2.,  ..., 2., 2., 2.],\n",
      "        [2., 2., 2.,  ..., 2., 2., 2.]], device='cuda:0')\n"
     ]
    }
   ],
   "source": [
    "torch.cuda.set_device(0)  # no-op to ensure context\n",
    "X = torch.ones(size=(128, 512), device='cuda')\n",
    "Y = torch.ones(size=(512,), device='cuda')\n",
    "cuda_extension.add_relu_fusion(X, Y)\n",
    "print(X)\n",
    "torch.testing.assert_close(X, eager_result, rtol=1e-4, atol=1e-4)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "12f1cf41c8096476",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "# `torch.compile` is your triton learning companion\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fb8294d1f06f54d9",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "## Example: LoRA Fused Kernels\n",
    "\n",
    "### LoRA (LOW-RANK ADAPTATION)\n",
    "\n",
    "<img src=\"./data/lora.png\" width=\"400\"/>\n",
    "\n",
    "Source: https://arxiv.org/pdf/2106.09685\n",
    "\n",
    "- Simple low rank reparametization of weight matrices\n",
    "- Singular value decomposition"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ff5da8bc2255f63c",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "## Fused Kernels\n",
    "\n",
    "> `TORCH_LOGS=output_code CUDA_LAUNCH_BLOCKING=1 python lora_on_simple_mlp.py `"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "56bdb71cea7e209b",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "from lora_on_simple_mlp import *\n",
    "from kernels.triton_fused_add_mul_relu import * "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3d8c54ec754688b4",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "### [TRITON] Fused Mul Add Relu"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "3c4f5f53c03abe97",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2.2.0\n",
      "Input tensor([[-0.9247, -0.4253, -2.6438,  0.1452, -0.1209, -0.5797, -0.6229, -0.3284],\n",
      "        [-1.0745, -0.3631, -1.6711,  2.2655,  0.3117, -0.1842,  1.2866,  1.1820],\n",
      "        [-0.1271,  1.2169,  1.4353,  1.0605, -0.4941, -1.4244, -0.7244, -1.2973],\n",
      "        [ 0.0697, -0.0074,  1.8969,  0.6878, -0.0779, -0.8373,  1.3506, -0.2879],\n",
      "        [-0.5965, -0.3283, -0.9086, -0.8059, -0.7407, -0.0504,  0.5435,  1.5150],\n",
      "        [ 0.0141,  0.4532,  1.6349,  0.7124, -0.1806,  1.0252, -1.4622, -0.7554],\n",
      "        [-0.1836,  0.3824,  0.3918, -0.0830,  0.8971, -1.1123,  0.1116,  0.4863],\n",
      "        [-0.5499, -0.3231, -0.5469,  0.9049,  0.2837,  0.1210,  0.4730, -1.0823]],\n",
      "       device='cuda:0')\n",
      "Expected Output tensor([[1.7098e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 6.7049e-01, 0.0000e+00,\n",
      "         3.4424e-01, 1.1223e-02],\n",
      "        [1.8750e+00, 0.0000e+00, 0.0000e+00, 2.2105e+00, 6.6808e-01, 0.0000e+00,\n",
      "         1.8445e+00, 1.9246e+00],\n",
      "        [2.2614e+00, 1.3964e+00, 5.1802e-01, 2.0114e+00, 0.0000e+00, 0.0000e+00,\n",
      "         2.7715e-01, 0.0000e+00],\n",
      "        [2.4750e+00, 0.0000e+00, 1.9079e+00, 5.4869e-01, 3.7923e-01, 0.0000e+00,\n",
      "         3.1791e+00, 6.6843e-01],\n",
      "        [1.8234e+00, 0.0000e+00, 1.1147e-01, 0.0000e+00, 0.0000e+00, 0.0000e+00,\n",
      "         2.4250e+00, 3.2593e+00],\n",
      "        [2.2903e+00, 0.0000e+00, 1.1721e+00, 6.9331e-01, 1.0583e+00, 6.7518e-01,\n",
      "         0.0000e+00, 2.6185e-01],\n",
      "        [1.5804e+00, 0.0000e+00, 9.0740e-01, 1.6670e-01, 5.5230e-02, 0.0000e+00,\n",
      "         1.5325e+00, 3.5984e-01],\n",
      "        [1.5554e+00, 0.0000e+00, 0.0000e+00, 9.4380e-01, 2.9795e-03, 7.4125e-02,\n",
      "         1.6286e+00, 0.0000e+00]], device='cuda:0')\n"
     ]
    }
   ],
   "source": [
    "print(triton.__version__)\n",
    "in_out_tensor, in_tensor, bias = get_inputs(add_manual_seed=True)\n",
    "expected_output = torch.maximum(in_out_tensor + 0.5 * in_tensor + bias, torch.tensor(0., device='cuda'))\n",
    "print(\"Input\", in_out_tensor)\n",
    "print(\"Expected Output\", expected_output)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "cdd2e0fbb38ab43c",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Output 1 tensor([[1.7098e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 6.7049e-01, 0.0000e+00,\n",
      "         3.4424e-01, 1.1223e-02],\n",
      "        [1.8750e+00, 0.0000e+00, 0.0000e+00, 2.2105e+00, 6.6808e-01, 0.0000e+00,\n",
      "         1.8445e+00, 1.9246e+00],\n",
      "        [2.2614e+00, 1.3964e+00, 5.1802e-01, 2.0114e+00, 0.0000e+00, 0.0000e+00,\n",
      "         2.7715e-01, 0.0000e+00],\n",
      "        [2.4750e+00, 0.0000e+00, 1.9079e+00, 5.4869e-01, 3.7923e-01, 0.0000e+00,\n",
      "         3.1791e+00, 6.6843e-01],\n",
      "        [1.8234e+00, 0.0000e+00, 1.1147e-01, 0.0000e+00, 0.0000e+00, 0.0000e+00,\n",
      "         2.4250e+00, 3.2593e+00],\n",
      "        [2.2903e+00, 0.0000e+00, 1.1721e+00, 6.9331e-01, 1.0583e+00, 6.7518e-01,\n",
      "         0.0000e+00, 2.6185e-01],\n",
      "        [1.5804e+00, 0.0000e+00, 9.0740e-01, 1.6670e-01, 5.5230e-02, 0.0000e+00,\n",
      "         1.5325e+00, 3.5984e-01],\n",
      "        [1.5554e+00, 0.0000e+00, 0.0000e+00, 9.4380e-01, 2.9795e-03, 7.4125e-02,\n",
      "         1.6286e+00, 0.0000e+00]], device='cuda:0')\n"
     ]
    }
   ],
   "source": [
    "BLOCK_SIZE = 8\n",
    "grid = lambda meta: (triton.cdiv(in_out_tensor.numel(), meta['BLOCK_SIZE']),)\n",
    "fused_add_mul_relu[grid](in_out_tensor, \n",
    "                         bias, \n",
    "                         in_tensor, \n",
    "                         in_out_tensor.numel(), \n",
    "                         BLOCK_SIZE=BLOCK_SIZE)\n",
    "print(\"Output 1\", in_out_tensor)\n",
    "torch.testing.assert_close(in_out_tensor, expected_output, rtol=1e-4, atol=1e-4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "64731e8e8a09fe44",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Output 2 tensor([[1.7098e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 6.7049e-01, 0.0000e+00,\n",
      "         3.4424e-01, 1.1223e-02],\n",
      "        [1.8750e+00, 0.0000e+00, 0.0000e+00, 2.2105e+00, 6.6808e-01, 0.0000e+00,\n",
      "         1.8445e+00, 1.9246e+00],\n",
      "        [2.2614e+00, 1.3964e+00, 5.1802e-01, 2.0114e+00, 0.0000e+00, 0.0000e+00,\n",
      "         2.7715e-01, 0.0000e+00],\n",
      "        [2.4750e+00, 0.0000e+00, 1.9079e+00, 5.4869e-01, 3.7923e-01, 0.0000e+00,\n",
      "         3.1791e+00, 6.6843e-01],\n",
      "        [1.8234e+00, 0.0000e+00, 1.1147e-01, 0.0000e+00, 0.0000e+00, 0.0000e+00,\n",
      "         2.4250e+00, 3.2593e+00],\n",
      "        [2.2903e+00, 0.0000e+00, 1.1721e+00, 6.9331e-01, 1.0583e+00, 6.7518e-01,\n",
      "         0.0000e+00, 2.6185e-01],\n",
      "        [1.5804e+00, 0.0000e+00, 9.0740e-01, 1.6670e-01, 5.5230e-02, 0.0000e+00,\n",
      "         1.5325e+00, 3.5984e-01],\n",
      "        [1.5554e+00, 0.0000e+00, 0.0000e+00, 9.4380e-01, 2.9795e-03, 7.4125e-02,\n",
      "         1.6286e+00, 0.0000e+00]], device='cuda:0')\n"
     ]
    }
   ],
   "source": [
    "in_out_tensor, in_tensor, bias = get_inputs(add_manual_seed=True)\n",
    "num_weights = bias.numel()\n",
    "fused_add_mul_relu_cleaner[grid](in_out_tensor, \n",
    "                                 bias, \n",
    "                                 in_tensor, \n",
    "                                 num_weights, \n",
    "                                 in_out_tensor.numel(), \n",
    "                                 multiplier=0.5,\n",
    "                                 BLOCK_SIZE=BLOCK_SIZE)\n",
    "print(\"Output 2\", in_out_tensor)\n",
    "torch.testing.assert_close(in_out_tensor, expected_output, rtol=1e-4, atol=1e-4)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3718837bd93be3c6",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "### [CUDA] Fused Mul Add Relu"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "b1d9236ebca65c91",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "__global__ void fused_add_mul_relu_kernel(float *dense_in_out_ptr,\n",
      "                                          const float *scalar_ptr,\n",
      "                                          const float *dense_ptr,\n",
      "                                          const int num_weights,\n",
      "                                          const int xnumel,\n",
      "                                          const double multiplier) {\n",
      "    int index = blockIdx.x * blockDim.x + threadIdx.x;\n",
      "    if (index < xnumel) {\n",
      "        int scalar_index = index % num_weights;\n",
      "        float tmp0 = dense_in_out_ptr[index];\n",
      "        float tmp1 = scalar_ptr[scalar_index];\n",
      "        float tmp3 = dense_ptr[index];\n",
      "        float ma_result = max(0.0f, multiplier * tmp3 + tmp0 + tmp1);\n",
      "        dense_in_out_ptr[index] = ma_result;\n",
      "    }\n",
      "}\n",
      "\n",
      "torch::Tensor fused_add_mul_relu(torch::Tensor in_out,\n",
      "                                 const torch::Tensor &bias,\n",
      "                                 const torch::Tensor &in,\n",
      "                                 const double multiplier) {\n",
      "    auto numel = in_out.numel();\n",
      "    auto sizes = in_out.sizes();\n",
      "    const int XBLOCK = sizes[0];\n",
      "    dim3 threadsPerBlock(sizes[1]);\n",
      "    dim3 numBlocks((numel + XBLOCK - 1) / XBLOCK);\n",
      "    fused_add_mul_relu_kernel<<<numBlocks, threadsPerBlock>>>(\n",
      "            in_out.data_ptr<float>(),\n",
      "            bias.data_ptr<float>(),\n",
      "            in.data_ptr<float>(),\n",
      "            sizes[1],\n",
      "            numel,\n",
      "            multiplier);\n",
      "    cudaDeviceSynchronize();\n",
      "    return std::move(in_out);\n",
      "}\n",
      "----\n",
      "\n",
      "torch::Tensor fused_add_mul_relu(torch::Tensor in_out,\n",
      "                                 const torch::Tensor &bias,\n",
      "                                 const torch::Tensor &in,\n",
      "                                 const double multiplier);\n"
     ]
    }
   ],
   "source": [
    "cuda_code_file = \"./kernels/src/fused_kernels_lora_on_mlp.cu\"\n",
    "header_code_file = \"./kernels/src/fused_kernels_lora_on_mlp.cuh\"\n",
    "\n",
    "with open(cuda_code_file) as f:\n",
    "    cuda_code = \"\".join([f for f in f.readlines() if not f.startswith(\"#include\")])\n",
    "    print(cuda_code)\n",
    "\n",
    "print(\"----\")\n",
    "\n",
    "with open(header_code_file) as f:\n",
    "    header_code = \"\".join([f for f in f.readlines() if not f.startswith(\"#include\")])\n",
    "    print(header_code)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "a0f47ff1e55ef9c3",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Tensor x:\n",
      "-0.9247 -0.4253 -2.6438  0.1452 -0.1209 -0.5797 -0.6229 -0.3284 -1.0745 -0.3631\n",
      "-1.6711  2.2655  0.3117 -0.1842  1.2866  1.1820 -0.1271  1.2169  1.4353  1.0605\n",
      "-0.4941 -1.4244 -0.7244 -1.2973  0.0697 -0.0074  1.8969  0.6878 -0.0779 -0.8373\n",
      " 1.3506 -0.2879 -0.5965 -0.3283 -0.9086 -0.8059 -0.7407 -0.0504  0.5435  1.5150\n",
      " 0.0141  0.4532  1.6349  0.7124 -0.1806  1.0252 -1.4622 -0.7554 -0.1836  0.3824\n",
      " 0.3918 -0.0830  0.8971 -1.1123  0.1116  0.4863 -0.5499 -0.3231 -0.5469  0.9049\n",
      " 0.2837  0.1210  0.4730 -1.0823 -0.0334 -0.9734  0.9559 -1.1795 -1.0064  0.1160\n",
      " 0.6852 -0.4124 -0.6738 -0.5404  0.6898 -1.5517  0.3805 -0.0436  0.3597 -0.5043\n",
      "[ CUDAFloatType{8,10} ]\n",
      "Tensor bias:\n",
      " 0.1808\n",
      "-0.5523\n",
      " 0.9238\n",
      "-0.7350\n",
      " 1.3800\n",
      " 0.8676\n",
      " 0.1297\n",
      "-0.9406\n",
      " 0.8109\n",
      " 0.8821\n",
      "[ CUDAFloatType{10} ]\n",
      "Tensor y:\n",
      " 2.5441 -0.7163 -0.4934  0.1267  0.1014 -0.4035  0.9023  0.8099 -0.6884  0.1372\n",
      " 1.0377  0.0925 -0.3752 -0.0908  2.0639 -1.8164 -0.2719  0.2811 -1.0399  0.7765\n",
      " 0.8814  0.0444 -1.4870  1.1334  1.3268 -1.2616  0.9501 -0.6558  0.9098 -0.6290\n",
      "-0.6587  2.0811  1.4151 -0.3091 -0.2055  2.0562 -0.0490 -0.6361 -0.5359 -0.1310\n",
      "-0.2945  1.2275  1.0549  0.3576  1.6378 -0.2310  0.7883 -0.0807 -0.3924  1.2673\n",
      " 1.0420 -0.4945 -1.1637  1.5740  0.7116  0.6104  1.2852 -0.6533  1.1171 -1.0067\n",
      " 1.2912  1.6028  0.1332  1.0703 -1.1161 -0.8396 -3.6680  0.8189  0.1255 -0.7691\n",
      " 0.1552 -0.8782 -0.4734  0.9690 -1.9985  0.1030  0.8580  0.7625 -1.2587 -0.8183\n",
      "[ CUDAFloatType{8,10} ]\n",
      "Expected:\n",
      " 5.1076  0.0000  0.0000  0.0000  1.4922  0.0000  1.5820  0.5938  0.0000  0.8346\n",
      " 0.8966  1.9261  0.3726  0.0000  7.4136  0.0000  0.0000  0.9228  0.0000  3.7285\n",
      " 1.7140  0.0000  0.0000  0.5746  4.5014  0.0000  4.2118  0.0000  2.8254  0.0000\n",
      " 0.0165  3.9464  3.5820  0.0000  0.0000  4.7910  0.0000  0.0000  0.1218  2.0957\n",
      " 0.0000  2.7243  4.9850  0.7998  4.9664  1.3615  0.4806  0.0000  0.0000  4.1793\n",
      " 2.9691  0.0000  0.0000  1.7729  3.1283  2.7577  2.5356  0.0000  2.8334  0.0000\n",
      " 3.4343  3.2553  1.7032  0.6444  0.0000  0.0000  0.0000  0.0000  0.0932  0.0000\n",
      " 1.2229  0.0000  0.0000  0.9534  0.0000  0.0000  2.4837  0.7694  0.0000  0.0000\n",
      "[ CUDAFloatType{8,10} ]\n",
      "Result:\n",
      " 5.1076  0.0000  0.0000  0.0000  1.4922  0.0000  1.5820  0.5938  0.0000  0.8346\n",
      " 0.8966  1.9261  0.3726  0.0000  7.4136  0.0000  0.0000  0.9228  0.0000  3.7285\n",
      " 1.7140  0.0000  0.0000  0.5746  4.5014  0.0000  4.2118  0.0000  2.8254  0.0000\n",
      " 0.0165  3.9464  3.5820  0.0000  0.0000  4.7910  0.0000  0.0000  0.1218  2.0957\n",
      " 0.0000  2.7243  4.9850  0.7998  4.9664  1.3615  0.4806  0.0000  0.0000  4.1793\n",
      " 2.9691  0.0000  0.0000  1.7729  3.1283  2.7577  2.5356  0.0000  2.8334  0.0000\n",
      " 3.4343  3.2553  1.7032  0.6444  0.0000  0.0000  0.0000  0.0000  0.0932  0.0000\n",
      " 1.2229  0.0000  0.0000  0.9534  0.0000  0.0000  2.4837  0.7694  0.0000  0.0000\n",
      "[ CUDAFloatType{8,10} ]\n",
      "All Match: true\n"
     ]
    }
   ],
   "source": [
    "! ./kernels/cmake-build-debug/fused_kernels_lora_test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "b9624b8b71cef929",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "The input conditions for extension module kernel_extension have changed. Bumping to version 1 and re-building as kernel_extension_v1...\n",
      "Detected CUDA files, patching ldflags\n",
      "Emitting ninja build file ./build/build.ninja...\n",
      "Building extension module kernel_extension_v1...\n",
      "Allowing ninja to set a default number of workers... (overridable by setting the environment variable MAX_JOBS=N)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[1/3] c++ -MMD -MF main.o.d -DTORCH_EXTENSION_NAME=kernel_extension_v1 -DTORCH_API_INCLUDE_EXTENSION_H -DPYBIND11_COMPILER_TYPE=\\\"_gcc\\\" -DPYBIND11_STDLIB=\\\"_libstdcpp\\\" -DPYBIND11_BUILD_ABI=\\\"_cxxabi1011\\\" -isystem /home/ksharma/anaconda3/envs/cuda-learn/lib/python3.11/site-packages/torch/include -isystem /home/ksharma/anaconda3/envs/cuda-learn/lib/python3.11/site-packages/torch/include/torch/csrc/api/include -isystem /home/ksharma/anaconda3/envs/cuda-learn/lib/python3.11/site-packages/torch/include/TH -isystem /home/ksharma/anaconda3/envs/cuda-learn/lib/python3.11/site-packages/torch/include/THC -isystem /home/ksharma/anaconda3/envs/cuda-learn/include -isystem /home/ksharma/anaconda3/envs/cuda-learn/include/python3.11 -D_GLIBCXX_USE_CXX11_ABI=0 -fPIC -std=c++17 -c /home/ksharma/dev/git/cuda-mode-lecture/build/main.cpp -o main.o \n",
      "[2/3] /home/ksharma/anaconda3/envs/cuda-learn/bin/nvcc --generate-dependencies-with-compile --dependency-output cuda.cuda.o.d -DTORCH_EXTENSION_NAME=kernel_extension_v1 -DTORCH_API_INCLUDE_EXTENSION_H -DPYBIND11_COMPILER_TYPE=\\\"_gcc\\\" -DPYBIND11_STDLIB=\\\"_libstdcpp\\\" -DPYBIND11_BUILD_ABI=\\\"_cxxabi1011\\\" -isystem /home/ksharma/anaconda3/envs/cuda-learn/lib/python3.11/site-packages/torch/include -isystem /home/ksharma/anaconda3/envs/cuda-learn/lib/python3.11/site-packages/torch/include/torch/csrc/api/include -isystem /home/ksharma/anaconda3/envs/cuda-learn/lib/python3.11/site-packages/torch/include/TH -isystem /home/ksharma/anaconda3/envs/cuda-learn/lib/python3.11/site-packages/torch/include/THC -isystem /home/ksharma/anaconda3/envs/cuda-learn/include -isystem /home/ksharma/anaconda3/envs/cuda-learn/include/python3.11 -D_GLIBCXX_USE_CXX11_ABI=0 -D__CUDA_NO_HALF_OPERATORS__ -D__CUDA_NO_HALF_CONVERSIONS__ -D__CUDA_NO_BFLOAT16_CONVERSIONS__ -D__CUDA_NO_HALF2_OPERATORS__ --expt-relaxed-constexpr -gencode=arch=compute_75,code=compute_75 -gencode=arch=compute_75,code=sm_75 --compiler-options '-fPIC' -O2 -std=c++17 -c /home/ksharma/dev/git/cuda-mode-lecture/build/cuda.cu -o cuda.cuda.o \n",
      "[3/3] c++ main.o cuda.cuda.o -shared -L/home/ksharma/anaconda3/envs/cuda-learn/lib/python3.11/site-packages/torch/lib -lc10 -lc10_cuda -ltorch_cpu -ltorch_cuda -ltorch -ltorch_python -L/home/ksharma/anaconda3/envs/cuda-learn/lib -lcudart -o kernel_extension_v1.so\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading extension module kernel_extension_v1...\n"
     ]
    }
   ],
   "source": [
    "cuda_extension = load_inline(\n",
    "    name='kernel_extension',\n",
    "    cpp_sources=header_code,\n",
    "    cuda_sources=cuda_code,\n",
    "    functions=[\"fused_add_mul_relu\"],\n",
    "    with_cuda=True,\n",
    "    verbose=True,\n",
    "    extra_cuda_cflags=[\"-O2\"],\n",
    "    build_directory='./build',\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "b9a2f5d64f138d09",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "in_out_tensor, in_tensor, bias = get_inputs(add_manual_seed=True)\n",
    "num_weights = bias.numel()\n",
    "result = cuda_extension.fused_add_mul_relu(in_out_tensor, bias, in_tensor, 0.5)\n",
    "torch.testing.assert_close(result, expected_output, rtol=1e-4, atol=1e-4)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6c04b5b64738982b",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "## Combine the kernels\n",
    "\n",
    "### Fused add mul sigmoid/relu/etc\n",
    "\n",
    "```\n",
    "@triton.jit\n",
    "def fused_add_mul_activation_kernel(x_ptr, bias_ptr, in_ptr,\n",
    "                                    num_weights: tl.constexpr,\n",
    "                                    xnumel: tl.constexpr,\n",
    "                                    multiplier: tl.constexpr,\n",
    "                                    activation: tl.constexpr,\n",
    "                                    BLOCK_SIZE: tl.constexpr):\n",
    "    xoffset = tl.program_id(0) * BLOCK_SIZE\n",
    "    index = xoffset + tl.arange(0, BLOCK_SIZE)[:]\n",
    "    mask = index < xnumel\n",
    "    bias_index = index % num_weights\n",
    "    tmp0 = tl.load(x_ptr + index, mask)\n",
    "    tmp1 = tl.load(bias_ptr + bias_index, mask, eviction_policy='evict_last')\n",
    "    tmp3 = tl.load(in_ptr + index, mask)\n",
    "    activ_input = multiplier * tmp3 + tmp0 + tmp1\n",
    "    if activation == \"sigmoid\":\n",
    "        ma_result = tl.sigmoid(activ_input)\n",
    "    elif activation == \"relu\":\n",
    "        ma_result = tl.maximum(0, activ_input)\n",
    "    # elif ...\n",
    "\n",
    "    tl.store(x_ptr + index, ma_result, mask)\n",
    "```\n",
    "\n",
    "### Let's check the perf of this kernel wrt torch.script, eager torch\n",
    "\n",
    "### Can we write the whole thing as triton/cuda kernel? Let's look MLP without LoRA layers\n",
    "\n",
    "\n",
    "#### BLOCKS: May be I could do something like this? INCORRECT VERSION\n",
    "\n",
    "- Block b x n\n",
    "- Block n x w1\n",
    "- Block 1 x w1\n",
    "- Block w1 x w2\n",
    "- Block 1 x w2\n",
    "\n",
    "<img src=\"./data/fused_lora_on_mlp.png\" width=\"800\"/>\n",
    "\n",
    "#### Pseudo code:\n",
    "\n",
    "- Step 1: `step1 = tl.dot(block_bn, block_nw1)`\n",
    "- Step 2: `block_1w1 = index_bn % W1.size()[1]`\n",
    "- Step 3: `step3 = step1 + block_1w1`\n",
    "- Step 4: `step4 = Relu(step3)`\n",
    "- Step 5: `block_w1w2 = ...`\n",
    "- Step 6: `step6 = tl.dot(step4, block_w1w2)`\n",
    "- Step 7: `block_1w2 = index_w1w2 % W2.size()[1]`\n",
    "- Step 8: `step8 = step6 + block_1w2`\n",
    "- Step 9: `result = sigmoid(step8)`\n",
    "\n",
    "\n",
    "#### ======================\n",
    "#### What's wrong with this?\n",
    "#### ======================\n",
    "\n",
    "### Can we do this?\n",
    "\n",
    "<img src=\"./data/fused_partial_mlp.png\" width=\"800\"/>\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ab4fbb17-3077-42a1-aa84-1430d8681d5c",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.11.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}