Triton Masked Copy Kernel
Problem Statement
Copy a 1D tensor into an output buffer using a Triton kernel, correctly handling a length that is not a multiple of BLOCK_SIZE.
Background
The final block usually has lanes past the end of the tensor. mask = offsets < n ensures those lanes are not loaded or stored, preventing out-of-bounds memory access.
Your Task
Implement copy_kernel and run(n=1000) (deliberately non-power-of-two) that returns whether the copy is exact.
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 = 1000
True
- The input
n = 1000is used to allocate a 1D tensor of length 1000 on the GPU. - The
copy_kernelfunction is launched with this tensor, using aBLOCK_SIZEthat does not divide 1000 evenly, resulting in a final block with lanes past the end of the tensor. - To prevent out-of-bounds memory access, a mask is applied using
mask = offsets < n, ensuring that only valid lanes are loaded and stored. - The copied tensor is then compared to a PyTorch reference implementation using
torch.allclose, which checks for element-wise equality within a tolerance, resulting in the outputTrueif the copy is exact.
Constraints:
- Must use a mask so the tail block does not read/write OOB
- out must equal x exactly
Background Knowledge
The problem involves using Triton, a Python-based programming language and framework for writing high-performance GPU code. Triton is designed to be more accessible and easier to use than traditional CUDA or CUDA Python, while still providing low-level control over GPU operations. The Triton kernel is the core component of a Triton program, responsible for executing computations on the GPU. In this problem, we need to implement a Triton kernel that copies a 1D tensor into an output buffer, handling cases where the length of the tensor is not a multiple of the block size.
The concept of block size is crucial in GPU programming. A block is a group of threads that can cooperate with each other, sharing memory and synchronizing their execution. The block size determines how many threads are executed together, which can significantly impact performance. However, when the length of the input data is not a multiple of the block size, some threads in the last block will be idle or may access out-of-bounds memory locations. To prevent this, we need to use a mask to selectively enable or disable threads based on their position within the block.
In this problem, we are also required to use PyTorch, a popular deep learning framework, to allocate inputs on the GPU and launch the Triton kernel. We will need to use PyTorch's torch.allclose function to compare the output of the Triton kernel with a reference implementation, ensuring that the copy is exact.
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.