Triton Write Index Kernel
Problem Statement
Write a kernel that fills an output tensor with its own global indices: out[i] = i. This is the Triton equivalent of torch.arange.
Background
The global index of each lane is exactly the offsets value you compute from program_id and tl.arange. Storing offsets (cast to float) demonstrates that you understand the program/block addressing scheme.
Your Task
Implement iota_kernel and run(n=1024) returning whether out equals torch.arange(n).
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 output tensor. - The
iota_kernelfunction is launched, which fills the output tensor with its own global indices: out[i]=i. - The resulting output tensor
triton_outis compared to the reference tensortorch.arange(n)usingtorch.allclose. - The comparison returns
Trueif the two tensors are equal within a certain tolerance, indicating that the kernel correctly filled the output tensor with its global indices.
Constraints:
- out[i] = i (as float)
- Derive i from program_id and tl.arange
- Mask the tail block
Background Knowledge
The problem involves working with the Triton programming language, which is used for writing high-performance kernels. To understand this problem, it's essential to have a grasp of parallel computing concepts, particularly those related to GPU acceleration. In parallel computing, a kernel is a small program that runs on a GPU, performing a specific task. The program_id and tl.arange are used to compute the offsets, which represent the global index of each lane.
In the context of Triton, lanes refer to the individual threads that execute a kernel. Each lane has a unique program_id, which is used to identify the lane and compute its global index. The tl.arange function is used to generate a sequence of numbers, similar to the torch.arange function in PyTorch. Understanding how to work with these concepts is crucial to solving this problem.
The problem also involves comparing the output of the Triton kernel with the output of a PyTorch reference implementation using torch.allclose. This function checks if two tensors are element-wise equal within a certain tolerance. To solve this problem, you'll need to have a good understanding of how to work with tensors, GPU memory allocation, and kernel launches.
Algorithm/Approach
The general approach to solving this problem involves:
- Computing the global index of each lane using the program_id and tl.arange
- Storing the global index in the output tensor
- Launching the kernel and allocating memory on the GPU
- Comparing the output of the Triton kernel with the PyTorch reference implementation
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.