Triton Autotuned Vector Add
Problem Statement
Add two vectors but let Triton pick the best BLOCK_SIZE using @triton.autotune.
Background
@triton.autotune benchmarks a list of triton.Configs (keyed on input sizes) and caches the fastest. The grid becomes a lambda over meta so it can read the chosen BLOCK_SIZE.
Your Task
Decorate add_kernel with @triton.autotune over at least two block sizes and implement run(n=4096) that launches with a meta-aware grid.
How it is tested
Your solution must define a top-level function run(...) that allocates inputs on the GPU, launches your Triton kernel, and returns a boolean from torch.allclose(triton_out, torch_reference, ...). The grader prints run(...); the expected output is True.
Example:
n = 4096
True
- The
runfunction is called with the inputn = 4096, which represents the size of the vectors to be added. - The
add_kernelfunction is decorated with@triton.autotuneto benchmark different block sizes and determine the fastest configuration for the given input size. - The Triton kernel is launched with the chosen block size, and the vector addition is performed on the GPU, producing the
triton_outresult. - The
triton_outresult is compared to thetorch_referenceresult usingtorch.allclose, which checks for element-wise equality within a tolerance, yielding the outputTrueif the results match.
Constraints:
- Use @triton.autotune with configs of differing BLOCK_SIZE, key=['n']
- grid = lambda meta: (triton.cdiv(n, meta['BLOCK_SIZE']),)
- BLOCK_SIZE is supplied by the config, not passed explicitly
Background Knowledge
The problem involves using Triton, a Python-based programming language and framework for writing high-performance GPU code. Triton allows users to define custom kernels, which are small programs that run on the GPU. In this problem, we need to define a kernel that adds two vectors. The Triton Autotune feature is used to automatically find the optimal block size for the kernel, which can significantly impact performance. The block size determines how many threads are executed in parallel on the GPU.
To understand this problem, it's essential to have a basic knowledge of GPU programming and parallel computing. In GPU programming, the GPU is divided into a grid of blocks, and each block is further divided into threads. The number of threads in a block is determined by the block size. A larger block size can lead to better performance, but it also increases the memory usage and can lead to slower performance if the block size is too large. Triton Autotune helps to find the optimal block size by benchmarking different block sizes and caching the fastest one.
The problem also involves using PyTorch, a popular deep learning framework, to allocate inputs on the GPU and launch the Triton kernel. The torch.allclose function is used to compare the output of the Triton kernel with the reference output. The goal is to define a run function that launches the kernel and returns a boolean indicating whether the output is correct.
Algorithm/Approach
The general approach to solve this problem involves defining a Triton kernel that adds two vectors, decorating the kernel with @triton.autotune to find the optimal block size, and launching the kernel using a meta-aware grid. The meta-aware grid allows the kernel to adapt to the optimal block size found by Triton Autotune.
Step-by-Step Strategy
To solve this problem, follow these steps:
- Define a Triton kernel that adds two vectors.
- Decorate the kernel with @triton.autotune and specify a list of block sizes to benchmark.
- Define a run function that allocates inputs on the GPU, launches the Triton kernel, and returns a boolean indicating whether the output is correct.
- Use torch.allclose to compare the output of the Triton kernel with the reference output.
Common Pitfalls
When implementing the solution, watch out for the following common pitfalls:
- Incorrectly defining the Triton kernel or the run function.
- Failing to decorate the kernel with @triton.autotune or specifying an incorrect list of block sizes.
- Incorrectly allocating inputs on the GPU or launching the kernel.
- Failing to use torch.allclose to compare the output of the Triton kernel with the reference output.
Time & Space Complexity
The time complexity of the solution depends on the block size and the number of threads launched. In general, the time complexity is O(n), where n is the length of the input vectors. The space complexity is also O(n), as we need to allocate memory for the input and output vectors on the GPU. However, the actual time and space complexity may vary depending on the specific implementation and the hardware used.