triton_fused_vector_addition_kernel_with_pytorch_benchmark.py

1import torch
2
3import triton
4import triton.language as tl
5
6
7@triton.jit
8def add_kernel(x_ptr,  # Pointer to first input vector.
9               y_ptr,  # Pointer to second input vector.
10               output_ptr,  # Pointer to output vector.
11               n_elements,  # Size of the vector.
12               BLOCK_SIZE: tl.constexpr,  # Number of elements each program should process.
13               # NOTE: `constexpr` so it can be used as a shape value.
14               ):
15    # There are multiple 'programs' processing different data. We identify which program
16    # we are here:
17    pid = tl.program_id(axis=0)  # We use a 1D launch grid so axis is 0.
18    # This program will process inputs that are offset from the initial data.
19    # For instance, if you had a vector of length 256 and block_size of 64, the programs
20    # would each access the elements [0:64, 64:128, 128:192, 192:256].
21    # Note that offsets are a list of pointers:
22    block_start = pid * BLOCK_SIZE
23    offsets = block_start + tl.arange(0, BLOCK_SIZE)
24    # Create a mask to guard memory operations against out-of-bounds accesses.
25    mask = offsets < n_elements
26    # Load x and y from DRAM, masking out any extra elements in case the input is not a
27    # multiple of the block size.
28    x = tl.load(x_ptr + offsets, mask=mask)
29    y = tl.load(y_ptr + offsets, mask=mask)
30    output = x + y
31    # Write x + y back to DRAM.
32    tl.store(output_ptr + offsets, output, mask=mask)
33
34
35def add(x: torch.Tensor, y: torch.Tensor):
36    # We need to preallocate the output.
37    output = torch.empty_like(x)
38    assert x.is_cuda and y.is_cuda and output.is_cuda
39    n_elements = output.numel()
40    # The L2 cache is shared among all next-generation GPUs.
41    # The grid is a 1D grid where each program computes a block of elements.
42    grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE']), )
43    # NOTE:
44    #  - Each torch.tensor object is implicitly converted into a pointer to its first element.
45    #  - `triton.jit`'ed functions can be indexed with a launch grid to launch a kernel.
46    #  - The keyword arguments are used to initialize the `tl.constexpr` objects.
47    add_kernel[grid](x, y, output, n_elements, BLOCK_SIZE=1024)
48    # We return a handle to the z tensor, but this is actually void since the kernel writes
49    # to z in place.
50    return output
51
52
53# %%
54# We can now use the above function to compute the element-wise sum of two 1D tensors and test its correctness:
55
56torch.manual_seed(0)
57size = 98432
58x = torch.rand(size, device='cuda')
59y = torch.rand(size, device='cuda')
60output_torch = x + y
61output_triton = add(x, y)
62print(output_torch)
63print(output_triton)
64print(f'The maximum difference between torch and triton is '
65      f'{torch.max(torch.abs(output_torch - output_triton))}')
66
67# %%
68# Benchmark
69# ---------
70#
71# We can now benchmark our custom op on vectors of increasing sizes to get a sense of how it performs relative to PyTorch.
72# To make things easier, Triton has a set of built-in utilities that allow us to construct the closing plots in a few lines of code.
73
74@triton.testing.perf_report(
75    triton.testing.Benchmark(
76        x_names=['size'],  # Argument names to use as an x-axis for the plot.
77        x_vals=[2**i for i in range(12, 28, 1)],  # Different values for `x_name`.
78        line_arg='provider',  # Argument name whose value corresponds to a different line in the plot.
79        line_vals=['triton', 'torch'],  # Possible values for `line_arg`.
80        line_names=['Triton', 'PyTorch'],  # Label name for the lines.
81        styles=[('blue', '-'), ('green', '-')],  # Line styles.
82        ylabel='GB/s',  # Label name for the y-axis.
83        plot_name='vector-add-performance',  # Name for the plot. Used also as a file name for saving the plot.
84        args={},  # Values for function arguments not in `x_names` and `y_name`.
85    ))
86def benchmark(size, provider):
87    x = torch.rand(size, device='cuda', dtype=torch.float32)
88    y = torch.rand(size, device='cuda', dtype=torch.float32)
89    quantiles = [0.5, 0.2, 0.8]
90    if provider == 'torch':
91        ms, min_ms, max_ms = triton.testing.do_bench(lambda: x + y, quantiles=quantiles)
92    if provider == 'triton':
93        ms, min_ms, max_ms = triton.testing.do_bench(lambda: add(x, y), quantiles=quantiles)
94    gbps = lambda ms: 3 * x.numel() * x.element_size() * 1e-9 / (ms * 1e-3)
95    return gbps(ms), gbps(max_ms), gbps(min_ms)
96
97
98# %%
99# Run the benchmark
100# -----------------
101#
102# Finally, we can run the benchmark and check the results.
103# The `show_plots` argument will display the plot if it's set to True.
104benchmark.run(print_data=True, show_plots=True)