Triton Fused Add + ReLU Kernel
Problem Statement
Fuse a residual add and a ReLU into one kernel: out = max(x + y, 0).
Background
Fusing avoids materializing the intermediate x + y to global memory — one read of each input, one write of the output.
Your Task
Implement add_relu_kernel and run(n=1024) comparing to torch.relu(x + y).
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 = 1024
True
- The input
n = 1024determines the size of the input tensorsxandy. - The
add_relu_kernelfunction is launched, which computesout = max(x + y, 0)for each element in the tensors, fusing the addition and ReLU operations into a single kernel. - The result
triton_outis compared to the reference outputtorch_referencecomputed usingtorch.relu(x + y). - The
runfunction returnsTrueiftriton_outis close totorch_referenceaccording totorch.allclose, indicating that the custom kernel produces the same result as the PyTorch reference implementation.
Constraints:
- Compute x + y then tl.maximum(., 0.0) without a second kernel
- Mask the tail block
Background Knowledge
The problem involves fusing two common operations in deep learning: element-wise addition and the ReLU (Rectified Linear Unit) activation function. The ReLU function is defined as f(x)=max(x,0), which outputs x if x is positive and 0 otherwise. In the context of this problem, we are tasked with implementing a kernel that performs the operation out = max(x + y, 0), where x and y are input tensors. This operation is a fundamental component of many neural network architectures.
Fusing operations, as mentioned in the problem statement, is a technique used to improve performance by reducing the number of memory accesses. In this case, fusing the addition and ReLU operations avoids the need to store the intermediate result of x + y in memory, which can lead to significant performance gains, especially for large tensors. This is particularly important in the context of GPU (Graphics Processing Unit) computing, where memory bandwidth is often a bottleneck.
The problem also involves using the Triton programming language, which is a Python-based language for writing high-performance GPU kernels. Triton provides a simple and intuitive API for launching GPU kernels, making it an ideal choice for this type of problem. To solve this problem, you will need to have a basic understanding of GPU programming, as well as familiarity with the Triton language and its ecosystem.
Algorithm/Approach
The general approach to solving this problem involves defining a Triton kernel that performs the fused addition and ReLU operation. This will involve using Triton's API to launch a GPU kernel that takes the input tensors x and y as input and produces the output tensor out. The kernel will need to be implemented using Triton's programming model, which involves defining a kernel function that is executed in parallel across the GPU.
Continue the full explanation
You're reading the free preview. Unlock the complete walkthrough, the code editor, test runner and reference solution with Premium.
Editor locked
The code editor is locked for Pro problems. It is only available for free problems. Please upgrade to gain access to the code editor for all problems.