Triton Row Cumulative Sum
Problem Statement
Compute the inclusive prefix sum (cumulative sum) along each row of a (M, N) tensor using tl.cumsum. Each row fits in one block.
Background
tl.cumsum(x, axis=0) returns the running sum along the block axis. Pad invalid lanes with 0 so they don't affect earlier prefixes, and mask them out on store.
Your Task
Implement cumsum_kernel and run(M=64, N=300) comparing to torch.cumsum(x, dim=1).
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:
M = 64, N = 300
True
- The input values
M = 64andN = 300define the size of the tensor, with 64 rows and 300 columns. - The
cumsum_kernelfunction computes the cumulative sum along each row usingtl.cumsum, padding invalid lanes with 0 and masking them out on store. - The result from
cumsum_kernelis compared to the result fromtorch.cumsum(x, dim=1)usingtorch.allclose, which checks if the two tensors are element-wise equal within a tolerance. - Since the implementation of
cumsum_kernelis correct and matches the result fromtorch.cumsum, the output ofrun(M=64, N=300)isTrue, indicating that the two results are identical.
Constraints:
- Use tl.cumsum(row, axis=0)
- Pad invalid lanes with other=0.0 and mask the store
- One program per row
Background Knowledge
The problem involves computing the cumulative sum along each row of a 2D tensor using Triton, a Python-based programming language and framework for writing high-performance GPU code. The cumulative sum, also known as the prefix sum, is a fundamental operation in linear algebra and is used extensively in various fields, including machine learning, data analysis, and scientific computing. In this context, the cumulative sum is computed along the rows of a (M, N) tensor, where M is the number of rows and N is the number of columns.
The Triton framework provides a tl.cumsum function that computes the running sum along the block axis. However, this function needs to be used in conjunction with other Triton functions and programming constructs to achieve the desired result. The problem also involves comparing the result with the torch.cumsum function from the PyTorch library, which provides a reference implementation for computing the cumulative sum. Understanding the basics of GPU programming, parallel computing, and linear algebra is essential for tackling this problem.
The problem requires a good understanding of memory management, data parallelism, and synchronization techniques to ensure that the computation is performed efficiently and correctly on the GPU. The use of masks and padding is also crucial to handle invalid lanes and ensure that the computation is accurate. The problem is designed to test the learner's ability to write efficient and correct GPU code using the Triton framework.
Algorithm/Approach
The general approach to solving this problem involves using the tl.cumsum function in conjunction with other Triton functions and programming constructs to compute the cumulative sum along each row of the input tensor. The algorithm can be broken down into several steps, including:
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.