Triton Clamp Kernel

easyPROElementwise & Fused OpsTriton Programming

Problem Statement

Implement a clamp kernel: out = min(max(x, lo), hi) with runtime scalars lo and hi.

Background

Compose tl.maximum and tl.minimum. Clamping bounds activations or gradients.

Your Task

Implement clamp_kernel and run(n=1024, lo=-0.5, hi=0.5) comparing to torch.clamp.

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, lo = -0.5, hi = 0.5

Output

True

Reasoning

The input values are $n = 1024$ , $lo = -0.5$ , and $hi = 0.5$ , which represent the number of elements to clamp and the lower and upper bounds, respectively.
The clamp_kernel function is applied to an array of $n$ elements, effectively computing $out = \min(\max(x, lo), hi)$ for each element $x$ .
The result from the clamp_kernel function is then compared to the result from torch.clamp using torch.allclose, which checks if the two arrays are element-wise equal within a tolerance.
The comparison yields True if the results are close enough, indicating that the custom clamp_kernel implementation matches the PyTorch reference implementation.

Constraints

out = tl.minimum(tl.maximum(x, lo), hi)
lo, hi are runtime scalars

Learn this in

Your First Kernel: Vector AdditionGPU Basics with Triton

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

Concept primerPREVIEW

Background Knowledge

The problem involves implementing a clamp kernel, which is a fundamental operation in deep learning and computer vision. The clamp kernel applies a simple transformation to an input tensor x, where each element is clamped to a range defined by lo and hi. This operation is essential in various applications, such as activation functions and gradient clipping. The clamp kernel can be composed using two basic operations: tl.maximum and tl.minimum, which are used to compute the maximum and minimum values between two tensors.

In the context of Triton Programming, the problem requires implementing a custom kernel that can be executed on a GPU. Triton is a programming language and framework for writing high-performance, GPU-accelerated code. To solve this problem, you need to understand how to define and launch a Triton kernel, as well as how to use the tl library to perform element-wise operations. Additionally, you should be familiar with PyTorch, as the problem involves comparing the output of your custom kernel with the result of torch.clamp.

The mathematical formulation of the clamp kernel is straightforward: out = min(max(x, lo), hi). This equation can be broken down into two steps: first, compute the maximum value between x and lo, and then compute the minimum value between the result and hi. This process can be expressed using the following equation: $y = \text{max}(x, lo)$ , followed by $out = \text{min}(y, hi)$ . By composing these two operations, you can implement the clamp kernel.

Algorithm/Approach

The general approach to solving this problem involves the following steps:

Define a Triton kernel that takes an input tensor x, and runtime scalars lo and hi.
Use the tl library to compute the maximum value between x and lo, and store the result in a temporary tensor y.
Compute the minimum value between y and hi, and store the result in the output tensor out.
Launch the Triton kernel on the GPU and allocate the input and output tensors.
Compare the output of your custom kernel with the result of torch.clamp using torch.allclose.

PRO

Continue the full explanation

You're reading the free preview. Unlock the complete walkthrough, the code editor, test runner and reference solution with Premium.

PRO

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.

solution.py

Python 3.13

GPU · T4

Test results