114 lines
4.2 KiB
Plaintext
114 lines
4.2 KiB
Plaintext
{
|
|
"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
|
|
}
|