1 / 1

Tiled Transpose in Triton

Linear AlgebraGPU Programming

Medium

Implement a tiled transpose kernel in Triton. Given $A$ of shape $(M, N)$ , write $\texttt{out}$ of shape $(N, M)$ such that:

\texttt{out}[j, i] = A[i, j]

This is a 2D tile pattern with no reduction: each program loads a $\texttt{BLOCK\_M} \times \texttt{BLOCK\_N}$ tile of $A$ and stores its transpose into $\texttt{out}$ . The interesting part is that the load is row-contiguous in $A$ but the store is row-contiguous in a transposed coordinate system, so the per-axis strides on the two pointer expressions are swapped.

Loading visualization...

Examples

Input: A = [[1, 2, 3], [4, 5, 6]]

Output: [[1, 4], [2, 5], [3, 6]]

Input: A = [[1, 2], [3, 4], [5, 6]]

Output: [[1, 3, 5], [2, 4, 6]]

Hint 1

Use a 2D launch grid sized $(\lceil M / \text{BLOCK\_M} \rceil, \lceil N / \text{BLOCK\_N} \rceil)$ . Each program owns a tile of $A$ in the input coordinate system.

Hint 2

Build one set of row and column index vectors. Use them on the load with $A$ 's strides, and on the store with $\texttt{out}$ 's strides but with row and column roles swapped.

Hint 3

Both tile axes can be partial. Combine the row tail mask and the column tail mask with an AND for the 2D mask on the load and the store.

Requirements

Use one Triton program per input tile on a 2D grid
Swap the row and column strides between the load and the store to encode the transpose
Mask both the load and the store on the $M$ and $N$ tile tails
The $\texttt{solve}$ signature must not change

Constraints

$1 \le M, N \le 512$ .
Input and output are torch.float32 on cuda.
Tolerance: atol=1e-3. No accumulation, no precision loss.

Try Similar Problems

Autotuned Matrix Multiplication in TritonMedium GEMV: Matrix Vector Product in TritonMedium Fused Matmul + Bias + ReLU in TritonMedium Tiled Matrix Multiplication in TritonMedium Scaled Dot-Product Attention in TritonMedium Block-Pointer Matmul in TritonMedium Split-K Matmul in TritonHard Vector Addition in TritonEasy Fused Multiply-Add in TritonEasy Grouped Program-ID Matmul in TritonHard L2 Vector Norm in TritonEasy RMS Normalization as a CUDA KernelMedium

Tiled Transpose in Triton

Linear AlgebraGPU Programming

Medium

Implement a tiled transpose kernel in Triton. Given $A$ of shape $(M, N)$ , write $\texttt{out}$ of shape $(N, M)$ such that:

\texttt{out}[j, i] = A[i, j]

Loading visualization...

Examples

Input: A = [[1, 2, 3], [4, 5, 6]]

Output: [[1, 4], [2, 5], [3, 6]]

Input: A = [[1, 2], [3, 4], [5, 6]]

Output: [[1, 3, 5], [2, 4, 6]]

Hint 1

Use a 2D launch grid sized $(\lceil M / \text{BLOCK\_M} \rceil, \lceil N / \text{BLOCK\_N} \rceil)$ . Each program owns a tile of $A$ in the input coordinate system.

Hint 2

Build one set of row and column index vectors. Use them on the load with $A$ 's strides, and on the store with $\texttt{out}$ 's strides but with row and column roles swapped.

Hint 3

Both tile axes can be partial. Combine the row tail mask and the column tail mask with an AND for the 2D mask on the load and the store.

Requirements

Use one Triton program per input tile on a 2D grid
Swap the row and column strides between the load and the store to encode the transpose
Mask both the load and the store on the $M$ and $N$ tile tails
The $\texttt{solve}$ signature must not change

Constraints

$1 \le M, N \le 512$ .
Input and output are torch.float32 on cuda.
Tolerance: atol=1e-3. No accumulation, no precision loss.

Try Similar Problems

A =

Accepts: array

Examples

Requirements

Constraints

Kernel performance

Tiled Transpose in Triton

Examples

Requirements

Constraints

Kernel performance