Triton Square Kernel

easyPROTriton FundamentalsTriton Programming

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

Input

n = 1024

Output

True

Reasoning

The input n = 1024 is the length of a 1-D vector $x$ of random standard-normal floats that lives on the GPU, plus an output vector of the same length waiting to be filled.
Each output element is the square of the matching input: $\text{out}_i = x_i \cdot x_i = x_i^2$ for $i = 0, \dots, n-1$ . For example $x_i = -1.5$ gives $2.25$ and $x_i = 0.3$ gives $0.09$ .
Notice that every output is non-negative whatever the sign of $x_i$ , and each position depends only on its own input, so all $n$ positions can be computed independently and in parallel.
The work is split into fixed-size blocks of positions. When $n$ is not a multiple of the block size (for instance $n = 333$ ), the last block reaches past index $n-1$ ; those positions must be neither read nor written, or memory outside the vectors gets touched.
The result True means every computed value matches the PyTorch reference $x_i^2$ to within an absolute tolerance of $10^{-5}$ . A single skipped or wrong element would make the check fail.

Constraints

out = x * x
Mask the tail block

Learn this in

Introduction to TritonGPU 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 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.

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