{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "869892b8",
   "metadata": {},
   "outputs": [],
   "source": [
    "import triton\n",
    "import triton.language as tl\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "\n",
    "torch.set_default_device(\"cuda\")\n",
    "torch.set_default_dtype(torch.bfloat16)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1a4b0ef4",
   "metadata": {},
   "outputs": [],
   "source": [
    "@triton.jit\n",
    "def add_kernel(x_ptr,  # *Pointer* to first input vector.\n",
    "               y_ptr,  # *Pointer* to second input vector.\n",
    "               output_ptr,  # *Pointer* to output vector.\n",
    "               n_elements,  # Size of the vector.\n",
    "               BLOCK_SIZE: tl.constexpr,  # Number of elements each program should process.\n",
    "               # NOTE: `constexpr` so it can be used as a shape value.\n",
    "               ):\n",
    "    # There are multiple 'programs' processing different data. We identify which program\n",
    "    # we are here:\n",
    "    pid = tl.program_id(axis=0)  # We use a 1D launch grid so axis is 0.\n",
    "    # This program will process inputs that are offset from the initial data.\n",
    "    # For instance, if you had a vector of length 256 and block_size of 64, the programs\n",
    "    # would each access the elements [0:64, 64:128, 128:192, 192:256].\n",
    "    # Note that offsets is a list of pointers:\n",
    "    block_start = pid * BLOCK_SIZE\n",
    "    offsets = block_start + tl.arange(0, BLOCK_SIZE)\n",
    "    # Create a mask to guard memory operations against out-of-bounds accesses.\n",
    "    mask = offsets < n_elements\n",
    "    # Load x and y from DRAM, masking out any extra elements in case the input is not a\n",
    "    # multiple of the block size.\n",
    "    x = tl.load(x_ptr + offsets, mask=mask)\n",
    "    y = tl.load(y_ptr + offsets, mask=mask)\n",
    "    output = x + y\n",
    "    # Write x + y back to DRAM.\n",
    "    tl.store(output_ptr + offsets, output, mask=mask)\n",
    "\n",
    "def add(x: torch.Tensor, y: torch.Tensor):\n",
    "    # We need to preallocate the output.\n",
    "    output = torch.empty_like(x)\n",
    "    n_elements = output.numel()\n",
    "    # The SPMD launch grid denotes the number of kernel instances that run in parallel.\n",
    "    # It is analogous to CUDA launch grids. It can be either Tuple[int], or Callable(metaparameters) -> Tuple[int].\n",
    "    # In this case, we use a 1D grid where the size is the number of blocks:\n",
    "    grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE']), )\n",
    "    # NOTE:\n",
    "    #  - Each torch.tensor object is implicitly converted into a pointer to its first element.\n",
    "    #  - `triton.jit`'ed functions can be indexed with a launch grid to obtain a callable GPU kernel.\n",
    "    #  - Don't forget to pass meta-parameters as keywords arguments.\n",
    "    add_kernel[grid](x, y, output, n_elements, BLOCK_SIZE=1024)\n",
    "    # We return a handle to z but, since `torch.cuda.synchronize()` hasn't been called, the kernel is still\n",
    "    # running asynchronously at this point.\n",
    "    return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7bcab582",
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.manual_seed(0)\n",
    "size = 98432\n",
    "x = torch.rand(size)\n",
    "y = torch.rand(size)\n",
    "output_torch = x + y\n",
    "output_triton = add(x, y)\n",
    "print(output_torch)\n",
    "print(output_triton)\n",
    "print(f'The maximum difference between torch and triton is '\n",
    "      f'{torch.max(torch.abs(output_torch - output_triton))}')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pyvenv",
   "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.13.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}