Triton Square Kernel
Problem Statement
Implement an elementwise square kernel: out = x * x.
Background
A warm-up for elementwise math in registers. Compute the product of the loaded value with itself and store it.
Your Task
Implement square_kernel and run(n=1024).
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 value
n = 1024determines the size of the input arrayx. - The
square_kernelfunction computes the elementwise square ofx, resulting in an output arrayoutwhere each element is xi⋅xi=xi2. - The
runfunction allocates the input arrayxon the GPU, launches thesquare_kernel, and computes a reference solution using PyTorch. - The final output is
Trueif the Triton outputtriton_outis close to the PyTorch reference solutiontorch_reference, as determined bytorch.allclose.
Constraints:
- out = x * x
- Mask the tail block
Background Knowledge
The problem is based on the Triton programming language, which is a Python-based programming language and framework for writing high-performance GPU code. The Triton language is designed to provide a simple and easy-to-use interface for programming GPUs, while also providing low-level control over memory management and parallelism. In this problem, we are asked to implement an elementwise square kernel, which is a fundamental operation in many numerical computations. The kernel takes an input array x and produces an output array out where each element is the square of the corresponding element in x, i.e., out = x * x.
The problem also involves using PyTorch, a popular deep learning framework, to allocate inputs on the GPU and launch the Triton kernel. PyTorch provides a dynamic computation graph and automatic differentiation, making it a popular choice for deep learning research and development. In this problem, we will use PyTorch to allocate inputs on the GPU and compare the output of our Triton kernel with a reference implementation using PyTorch's torch module.
To solve this problem, we need to understand the basics of GPU programming, including memory management, parallelism, and synchronization. We also need to be familiar with the Triton language and its syntax, as well as PyTorch and its API. Additionally, we need to understand the concept of elementwise operations, which are operations that are applied to each element of an array independently.
Algorithm/Approach
The general approach to solving this problem is to use the Triton language to define a kernel function that takes an input array x and produces an output array out where each element is the square of the corresponding element in x. We will use PyTorch to allocate inputs on the GPU and launch the Triton kernel. We will also use PyTorch's torch.allclose function to compare the output of our Triton kernel with a reference implementation using PyTorch's torch module.
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.