Initial commit

2026-09-09 13:45:04 +02:00
commit 0f9b410151
5 changed files with 327 additions and 0 deletions
@@ -0,0 +1,113 @@
+{
+ "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
+}