Triton Row Sum Reduction
Problem Statement
Given a 2D tensor x of shape (M, N), compute the sum along each row, producing a length-M vector. Assume each row fits in one block.
Background
Launch one program per row (grid = (M,)). Load the whole row with BLOCK_SIZE = triton.next_power_of_2(N), mask columns >= N (load other=0.0), then reduce with tl.sum(row, axis=0).
Your Task
Implement row_sum_kernel and run(M=64, N=300) comparing to x.sum(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 are
M = 64andN = 300, representing the shape of the 2D tensorx. - We launch one program per row, resulting in a total of
Mprograms, each loading a row of lengthNinto a block of sizeBLOCK_SIZE = triton.next_power_of_2(N). - Within each block, we mask columns
>= Nby loadingother=0.0and then reduce the row usingtl.sum(row, axis=0), effectively calculating the sum of each row. - The resulting row sums from the Triton kernel are compared to the reference solution
x.sum(dim=1)usingtorch.allclose, yieldingTrueif the outputs match within a certain tolerance.
Constraints:
- One program per row, grid = (M,)
- BLOCK_SIZE = triton.next_power_of_2(N), pass as constexpr
- Use tl.sum(row, axis=0); mask invalid columns with other=0.0
Background Knowledge
The problem involves computing the sum along each row of a 2D tensor using Triton, a programming language for writing high-performance GPU code. To tackle this problem, it's essential to understand the basics of parallel computing and GPU programming. In parallel computing, tasks are divided into smaller sub-tasks that can be executed simultaneously, leveraging multiple processing units. GPU programming, in particular, involves writing code that can be executed on a Graphics Processing Unit (GPU), which has many more processing units than a Central Processing Unit (CPU).
In the context of this problem, we're dealing with a 2D tensor x of shape (M, N), where each row needs to be reduced to a single value by summing its elements. The reduction operation is a common pattern in parallel computing, where an array of values is reduced to a single value by applying a binary operation (in this case, addition). To achieve this, we'll be using Triton's programming model, which involves launching a program per row and using blocks to manage memory and computation.
The problem also mentions using triton.next_power_of_2(N) to determine the block size, which is a common technique in GPU programming to ensure efficient memory access and minimize memory fragmentation. Additionally, we'll be using masking to handle rows with fewer than BLOCK_SIZE elements, which is a common pattern in parallel computing to handle irregularly-shaped data.
Algorithm/Approach
The general approach to solving this problem involves:
- Launching a program per row, using a grid of size (M,)
- Loading each row into a block of size BLOCK_SIZE, which is a power of 2
- Masking columns that are out of bounds (i.e., >= N)
- Reducing each row using a sum operation, which can be implemented using Triton's tl.sum function
This approach leverages the parallel computing capabilities of the GPU to compute the sum along each row simultaneously, while also handling irregularly-shaped data using masking.
Step-by-Step Strategy
To implement the solution, follow these steps:
- Define the row_sum_kernel function, which will be launched per row
- Determine the block size using triton.next_power_of_2(N)
- Load each row into a block, masking columns that are out of bounds
- Reduce each row using a sum operation, storing the result in a register
- Define the run function, which will allocate inputs on the GPU, launch the row_sum_kernel, and compare the result to the reference implementation using torch.allclose
Common Pitfalls
When implementing the solution, watch out for:
- Incorrectly determining the block size, which can lead to inefficient memory access
- Failing to mask columns that are out of bounds, which can result in incorrect results
- Incorrectly implementing the sum operation, which can lead to incorrect results
Time & Space Complexity
The expected time complexity of the solution is O(M×N), where M is the number of rows and N is the number of columns. The space complexity is O(M), as we need to store the sum along each row. Note that the actual time and space complexity may vary depending on the specific implementation and the characteristics of the input data.