Triton Clamp Kernel
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:
n = 1024, lo = -0.5, hi = 0.5
True
- 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_kernelfunction 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_kernelfunction is then compared to the result fromtorch.clampusingtorch.allclose, which checks if the two arrays are element-wise equal within a tolerance. - The comparison yields
Trueif the results are close enough, indicating that the customclamp_kernelimplementation matches the PyTorch reference implementation.
Constraints:
- out = tl.minimum(tl.maximum(x, lo), hi)
- lo, hi are runtime scalars
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=max(x,lo), followed by out=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.
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.