📘

Triton ReLU Kernel

EasyElementwise & Fused Ops 📚 Triton Programming

Problem Statement

Implement the ReLU activation out = max(x, 0) as a Triton kernel.

Background

Use tl.maximum to clamp negatives to zero. ReLU is the most common activation and a perfect single-op kernel.

Your Task

Implement relu_kernel and run(n=1024) comparing to torch.relu(x).

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:

Input:

n = 1024

Output:

True

Reasoning:

The input value n = 1024 is used to allocate an input tensor x of size n on the GPU.
The relu_kernel function is launched, applying the ReLU activation function to each element of x using tl.maximum(x, 0), effectively setting all negative values to zero: $out = \max(x, 0)$ .
The output of the relu_kernel function is compared to the output of torch.relu(x) using torch.allclose(triton_out, torch_reference, ...).
The comparison returns True if the two outputs are close enough, indicating that the relu_kernel function has been correctly implemented.

Constraints:

Use tl.maximum(x, 0.0)
Mask the tail block

More from Triton Programming

Triton Single-Query Attention KernelHard Triton Row Cumulative SumHard Triton Stable LogSumExp KernelHard Triton Seeded Fused DropoutHard Triton Autotuned Vector AddMedium

Background Knowledge

The ReLU (Rectified Linear Unit) activation function is a widely used element-wise operation in deep learning models. It maps all negative values to zero and all non-negative values to themselves, which can be expressed as $out = \max(x, 0)$ . This function is essential in introducing non-linearity into neural networks, allowing them to learn more complex relationships between inputs and outputs. The ReLU function is defined as:

\text{ReLU}(x) = \max(x, 0) = \begin{cases} x, & \text{if } x \geq 0 \\ 0, & \text{if } x < 0 \end{cases}

In the context of Triton Programming, a kernel refers to a small program that runs on the GPU, performing a specific computation. Triton provides a Python-based interface for defining and launching these kernels, allowing for high-performance computations on large datasets. The tl.maximum function in Triton can be used to implement the ReLU activation by clamping negative values to zero.

Understanding the basics of GPU programming and parallel computing is also crucial for working with Triton. GPUs have thousands of cores that can perform computations simultaneously, making them ideal for element-wise operations like ReLU. However, this also means that the code must be designed to take advantage of this parallelism, using techniques like data parallelism and thread blocking to minimize overhead and maximize performance.

Algorithm/Approach

The general approach to solving this problem involves defining a Triton kernel that applies the ReLU activation function to an input tensor. This can be achieved by using the tl.maximum function to clamp negative values to zero. The kernel will need to be launched on the GPU, and the output will need to be compared to a reference implementation using torch.relu to verify correctness.

Step-by-Step Strategy

To implement the solution, follow these steps:

Define a Triton kernel function relu_kernel that takes an input tensor x and applies the ReLU activation using tl.maximum.
Allocate input data on the GPU using Triton's memory management functions.
Launch the relu_kernel on the GPU, passing in the input data.
Allocate output memory on the GPU to store the result of the kernel.
Compare the output of the Triton kernel to the result of torch.relu using torch.allclose.
Define a top-level function run that encapsulates the entire process and returns a boolean indicating whether the results match.

Common Pitfalls

When implementing the solution, watch out for the following common pitfalls:

Forgetting to allocate memory on the GPU for the input and output data.
Incorrectly launching the kernel or specifying the wrong block and grid dimensions.
Failing to synchronize the GPU after launching the kernel, leading to incorrect results.
Not using torch.allclose with the correct tolerance when comparing the results.

Time & Space Complexity

The time complexity of the ReLU activation function is $O(n)$ , where $n$ is the number of elements in the input tensor, since it involves a single element-wise operation. The space complexity is also $O(n)$ , as the output tensor has the same size as the input tensor. However, the actual performance of the implementation will depend on the efficiency of the Triton kernel and the underlying GPU architecture.

solution.py

Python 3.13

GPU · T4

Test Results

0/0

Run code to see test results.