triton/master/_downloads/f191ee1e78dc52eb5f7cba88f71cef2f/01-vector-add.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "%matplotlib inline"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n# Vector Addition\nIn this tutorial, you will write a simple vector addition using Triton and learn about:\n\n- The basic programming model of Triton\n- The `triton.jit` decorator, which is used to define Triton kernels.\n- The best practices for validating and benchmarking your custom ops against native reference implementations\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Compute Kernel\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import torch\n\nimport triton\nimport triton.language as tl\n\n\n@triton.jit\ndef add_kernel(\n    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 'program's 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)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Let's also declare a helper function to (1) allocate the `z` tensor\nand (2) enqueue the above kernel with appropriate grid/block sizes.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "def add(x: torch.Tensor, y: torch.Tensor):\n    # We need to preallocate the output\n    output = torch.empty_like(x)\n    assert x.is_cuda and y.is_cuda and output.is_cuda\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 index 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": "markdown",
      "metadata": {},
      "source": [
        "We can now use the above function to compute the element-wise sum of two `torch.tensor` objects and test its correctness:\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "torch.manual_seed(0)\nsize = 98432\nx = torch.rand(size, device='cuda')\ny = torch.rand(size, device='cuda')\noutput_torch = x + y\noutput_triton = add(x, y)\nprint(output_torch)\nprint(output_triton)\nprint(\n    f'The maximum difference between torch and triton is '\n    f'{torch.max(torch.abs(output_torch - output_triton))}'\n)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Seems like we're good to go!\n\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Benchmark\nWe can now benchmark our custom op on vectors of increasing sizes to get a sense of how it does relative to PyTorch.\nTo make things easier, Triton has a set of built-in utilities that allow us to concisely plot the performance of your custom ops\nfor different problem sizes.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "@triton.testing.perf_report(\n    triton.testing.Benchmark(\n        x_names=['size'],  # argument names to use as an x-axis for the plot\n        x_vals=[\n            2 ** i for i in range(12, 28, 1)\n        ],  # different possible values for `x_name`\n        x_log=True,  # x axis is logarithmic\n        line_arg='provider',  # argument name whose value corresponds to a different line in the plot\n        line_vals=['triton', 'torch'],  # possible values for `line_arg`\n        line_names=['Triton', 'Torch'],  # label name for the lines\n        styles=[('blue', '-'), ('green', '-')],  # line styles\n        ylabel='GB/s',  # label name for the y-axis\n        plot_name='vector-add-performance',  # name for the plot. Used also as a file name for saving the plot.\n        args={},  # values for function arguments not in `x_names` and `y_name`\n    )\n)\ndef benchmark(size, provider):\n    x = torch.rand(size, device='cuda', dtype=torch.float32)\n    y = torch.rand(size, device='cuda', dtype=torch.float32)\n    if provider == 'torch':\n        ms, min_ms, max_ms = triton.testing.do_bench(lambda: x + y)\n    if provider == 'triton':\n        ms, min_ms, max_ms = triton.testing.do_bench(lambda: add(x, y))\n    gbps = lambda ms: 12 * size / ms * 1e-6\n    return gbps(ms), gbps(max_ms), gbps(min_ms)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "We can now run the decorated function above. Pass `print_data=True` to see the performance number, `show_plots=True` to plot them, and/or\n`save_path='/path/to/results/' to save them to disk along with raw CSV data\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "benchmark.run(print_data=True, show_plots=True)"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3",
      "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.8.10"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}