Matrix Transpose: The Canonical CUDA Optimization Problem — From Naive to Bank-Conflict-Free

A naive matrix transpose achieves 150 GB/s on an A100. Add shared memory tiling to coalesce both reads and writes, and throughput jumps to 900 GB/s. Add one column of padding to eliminate bank conflicts, and throughput reaches 1,600 GB/s — 78% of peak HBM bandwidth. The progression from 150 to 1,600 GB/s represents three distinct optimizations, each targeting a different bottleneck: coalescing (10x), tiling (6x), and bank conflict elimination (1.8x). Matrix transpose is the canonical CUDA optimization problem because it forces you to confront every memory hierarchy issue simultaneously.

All measurements target NVIDIA Ampere (A100-80GB SXM, SM 8.0), 4096x4096 float matrix unless stated otherwise.

The Copy Kernel (Upper Bound)

The maximum achievable bandwidth is set by a simple copy kernel — both reads and writes are coalesced:

#include <cuda_runtime.h>

__global__ void copy_kernel(const float* __restrict__ src,
                            float* __restrict__ dst, int n) {
    int idx = threadIdx.x + blockIdx.x * blockDim.x;
    if (idx < n) {
        dst[idx] = src[idx];
    }
}

// For 4096x4096 float matrix (64 MB):
// A100 achieves ~1870 GB/s effective bandwidth (read + write)
// This is our target for transpose

Version 0: Naive Transpose (Uncoalesced Write)

// A is M x N (row-major), B is N x M (row-major)
// B[j][i] = A[i][j]
__global__ void transpose_naive(const float* __restrict__ A,
                                 float* __restrict__ B,
                                 int M, int N) {
    int col = blockIdx.x * blockDim.x + threadIdx.x;  // Column of A
    int row = blockIdx.y * blockDim.y + threadIdx.y;  // Row of A

    if (row < M && col < N) {
        // Read: A[row][col] = A[row * N + col]
        // Consecutive threads (varying threadIdx.x = varying col)
        // read A[row * N + col], A[row * N + col+1], ... -> COALESCED

        // Write: B[col][row] = B[col * M + row]
        // Consecutive threads write B[col * M + row], B[(col+1) * M + row], ...
        // These addresses differ by M (stride-M) -> NOT COALESCED

        B[col * M + row] = A[row * N + col];
    }
}

// Launch: dim3 block(32, 32); dim3 grid(N/32, M/32);
// Read bandwidth: ~1870 GB/s (coalesced)
// Write bandwidth: ~60 GB/s (stride-4096, every write to different cache line)
// Effective: ~120 GB/s

The read is perfectly coalesced: thread $(tx, ty)$ reads A[row * N + col] where col = blockIdx.x * 32 + tx. Consecutive tx values give consecutive addresses.

The write is strided: thread $(tx, ty)$ writes B[col * M + row] where col varies with tx. Consecutive tx values give addresses separated by M elements = 4096 floats = 16384 bytes. Each write goes to a different cache line.

Version 1: Coalesced Write (Uncoalesced Read)

We can swap which side is coalesced:

__global__ void transpose_coalesced_write(const float* __restrict__ A,
                                           float* __restrict__ B,
                                           int M, int N) {
    // Now threadIdx.x indexes the output column (row of B)
    int out_col = blockIdx.x * blockDim.x + threadIdx.x;  // Column of B = row of A
    int out_row = blockIdx.y * blockDim.y + threadIdx.y;  // Row of B = column of A

    if (out_row < N && out_col < M) {
        // Read: A[out_col][out_row] = A[out_col * N + out_row]
        // Consecutive threads (varying threadIdx.x = varying out_col)
        // read A[out_col * N + out_row], A[(out_col+1) * N + out_row], ...
        // Stride-N -> NOT COALESCED

        // Write: B[out_row][out_col] = B[out_row * M + out_col]
        // Consecutive threads write B[out_row * M + out_col], B[out_row * M + out_col+1]
        // -> COALESCED

        B[out_row * M + out_col] = A[out_col * N + out_row];
    }
}

Version 2: Tiled Transpose with Shared Memory

The solution: use shared memory as a staging buffer. Read a tile from the input (coalesced), write it to shared memory, read it from shared memory in transposed order, write to output (coalesced):

#define TILE_DIM 32
#define BLOCK_ROWS 8  // Each thread processes TILE_DIM/BLOCK_ROWS = 4 elements

__global__ void transpose_tiled(const float* __restrict__ A,
                                 float* __restrict__ B,
                                 int M, int N) {
    __shared__ float tile[TILE_DIM][TILE_DIM];

    // Input coordinates
    int x = blockIdx.x * TILE_DIM + threadIdx.x;  // Column of A
    int y = blockIdx.y * TILE_DIM + threadIdx.y;  // Row of A

    // Load TILE_DIM x TILE_DIM tile from A into shared memory
    // Block is TILE_DIM x BLOCK_ROWS = 32 x 8 = 256 threads
    // Each thread loads TILE_DIM/BLOCK_ROWS = 4 elements (one per row)
    for (int j = 0; j < TILE_DIM; j += BLOCK_ROWS) {
        if ((y + j) < M && x < N) {
            // A[(y+j) * N + x]: consecutive threadIdx.x -> consecutive x -> COALESCED
            tile[threadIdx.y + j][threadIdx.x] = A[(y + j) * N + x];
        }
    }

    __syncthreads();

    // Output coordinates (transposed block position)
    x = blockIdx.y * TILE_DIM + threadIdx.x;  // Note: blockIdx.y, not blockIdx.x
    y = blockIdx.x * TILE_DIM + threadIdx.y;  // Note: blockIdx.x, not blockIdx.y

    // Write transposed tile to B
    for (int j = 0; j < TILE_DIM; j += BLOCK_ROWS) {
        if ((y + j) < N && x < M) {
            // Read from shared memory: tile[threadIdx.x][threadIdx.y + j]
            // (transposed indexing)
            // Write to B: B[(y+j) * M + x]: consecutive threadIdx.x -> COALESCED
            B[(y + j) * M + x] = tile[threadIdx.x][threadIdx.y + j];
        }
    }
}

The key insight: both the global read and global write are now coalesced. The transpose happens in shared memory, which has much higher bandwidth and lower latency than global memory.

Visualizing the Tile Flow

A (input, row-major):                  tile (shared memory):
[a00 a01 a02 ... a031]                [a00 a01 a02 ... a031]
[a10 a11 a12 ... a131]    READ -->    [a10 a11 a12 ... a131]
[...]                     (coalesced)  [...]
[a310 a311 ... a3131]                  [a310 a311 ... a3131]

tile[tx][ty+j] (transposed read):     B (output, row-major):
[a00 a10 a20 ... a310]                [a00 a10 a20 ... a310]
[a01 a11 a21 ... a311]    WRITE -->   [a01 a11 a21 ... a311]
[...]                     (coalesced)  [...]
[a031 a131 ... a3131]                  [a031 a131 ... a3131]

Version 3: Bank-Conflict-Free Tiled Transpose

Version 2 has a subtle problem: shared memory bank conflicts during the transposed read.

When writing to shared memory: tile[threadIdx.y + j][threadIdx.x] — consecutive threads (varying threadIdx.x) access consecutive columns of the same row. Each column maps to a different bank (32 columns, 32 banks). No conflicts.

When reading from shared memory: tile[threadIdx.x][threadIdx.y + j] — consecutive threads (varying threadIdx.x) access consecutive ROWS of the same column. All 32 threads in a warp access the same column (threadIdx.y + j is fixed within one iteration), but different rows. Row 0 is bank 0, row 1 is bank 1, …, row 31 is bank 31 (since each row is 32 floats = 128 bytes = 32 banks). Wait — that is actually conflict-free!

Actually, the conflict arises from the tile width being exactly 32. tile[threadIdx.x][col] has address: base + threadIdx.x * 32 + col. For a fixed col, consecutive threadIdx.x values differ by 32 floats = 128 bytes. Since bank = (address / 4) % 32 = threadIdx.x * 32 + col) % 32 = col (constant). ALL 32 threads hit the SAME bank.

The fix: add one column of padding to make the stride 33 instead of 32:

#define TILE_DIM 32
#define BLOCK_ROWS 8
#define PAD 1

__global__ void transpose_tiled_padded(const float* __restrict__ A,
                                        float* __restrict__ B,
                                        int M, int N) {
    // PAD = 1: stride is 33 instead of 32
    // Bank for tile[row][col] = (row * 33 + col) % 32
    // Consecutive rows with same col: banks differ by 33 % 32 = 1
    // -> all 32 threads hit 32 DIFFERENT banks = no conflict
    __shared__ float tile[TILE_DIM][TILE_DIM + PAD];

    int x = blockIdx.x * TILE_DIM + threadIdx.x;
    int y = blockIdx.y * TILE_DIM + threadIdx.y;

    // Load (coalesced global read, no bank conflict on write)
    for (int j = 0; j < TILE_DIM; j += BLOCK_ROWS) {
        if ((y + j) < M && x < N) {
            tile[threadIdx.y + j][threadIdx.x] = A[(y + j) * N + x];
        }
    }

    __syncthreads();

    // Write (no bank conflict on transposed read, coalesced global write)
    x = blockIdx.y * TILE_DIM + threadIdx.x;
    y = blockIdx.x * TILE_DIM + threadIdx.y;

    for (int j = 0; j < TILE_DIM; j += BLOCK_ROWS) {
        if ((y + j) < N && x < M) {
            B[(y + j) * M + x] = tile[threadIdx.x][threadIdx.y + j];
        }
    }
}

Version 4: Vectorized Loads

#define TILE_DIM 32
#define BLOCK_ROWS 8

__global__ void transpose_vectorized(const float* __restrict__ A,
                                      float* __restrict__ B,
                                      int M, int N) {
    __shared__ float tile[TILE_DIM][TILE_DIM + 1];

    int x_base = blockIdx.x * TILE_DIM;
    int y_base = blockIdx.y * TILE_DIM;

    // Vectorized load: each thread loads float4 (4 consecutive floats)
    // Block: 32 x 8 threads. Each thread loads 4 rows worth of 1 float.
    // Alternative: 8 x 8 threads, each loads float4 across columns
    // Let's use 8 threads in x, float4 per thread -> covers 32 columns
    // 32 threads in y -> covers 32 rows with BLOCK_ROWS=8 -> 4 iterations

    int tx = threadIdx.x;
    int ty = threadIdx.y;

    // Load using standard scalar pattern (compatible with all sizes)
    for (int j = 0; j < TILE_DIM; j += BLOCK_ROWS) {
        int row = y_base + ty + j;
        int col = x_base + tx;
        if (row < M && col < N) {
            tile[ty + j][tx] = A[row * N + col];
        }
    }

    __syncthreads();

    // Write transposed
    for (int j = 0; j < TILE_DIM; j += BLOCK_ROWS) {
        int row = x_base + ty + j;  // Transposed
        int col = y_base + tx;
        if (row < N && col < M) {
            B[row * M + col] = tile[tx][ty + j];
        }
    }
}

For a truly vectorized version, restructure the block layout:

// float4 transpose: each thread handles 4 elements
// Block: 8 x 32 threads (256 total)
// Each thread loads one float4 = 4 consecutive floats per row
// 32 rows covered by threadIdx.y with 4 loads each (32/8 = 4 iterations)

__global__ void transpose_float4(const float* __restrict__ A,
                                  float* __restrict__ B,
                                  int M, int N) {
    __shared__ float tile[32][33];  // Padded

    int bx = blockIdx.x * 32;
    int by = blockIdx.y * 32;
    int tx = threadIdx.x;  // 0-7: each handles 4 columns via float4
    int ty = threadIdx.y;  // 0-31: row within tile

    // Load: each thread loads 4 floats from one row
    // threadIdx.x * 4 gives the starting column within the tile
    int row = by + ty;
    int col = bx + tx * 4;

    if (row < M && col + 3 < N) {
        float4 val = *reinterpret_cast<const float4*>(&A[row * N + col]);
        tile[ty][tx * 4 + 0] = val.x;
        tile[ty][tx * 4 + 1] = val.y;
        tile[ty][tx * 4 + 2] = val.z;
        tile[ty][tx * 4 + 3] = val.w;
    }

    __syncthreads();

    // Write transposed: read columns from tile, write as float4 to rows of B
    // Now row of B = column of A, col of B = row of A
    int out_row = bx + ty;
    int out_col = by + tx * 4;

    if (out_row < N && out_col + 3 < M) {
        float4 val;
        val.x = tile[tx * 4 + 0][ty];
        val.y = tile[tx * 4 + 1][ty];
        val.z = tile[tx * 4 + 2][ty];
        val.w = tile[tx * 4 + 3][ty];
        *reinterpret_cast<float4*>(&B[out_row * M + out_col]) = val;
    }
}

Version 5: Diagonal Block Ordering

For square matrices, the standard block ordering can cause partition camping — multiple blocks simultaneously access the same DRAM partitions. Diagonal block ordering distributes memory accesses more evenly:

#define TILE_DIM 32
#define BLOCK_ROWS 8

__global__ void transpose_diagonal(const float* __restrict__ A,
                                    float* __restrict__ B,
                                    int M, int N) {
    __shared__ float tile[TILE_DIM][TILE_DIM + 1];

    // Remap block coordinates along the diagonal
    // Standard: blockIdx.x, blockIdx.y
    // Diagonal: remap so that blocks accessing the same DRAM partition
    //           are spread across different SMs
    int bx, by;
    if (N == M) {
        // Square matrix: diagonal mapping
        bx = blockIdx.x;
        by = (blockIdx.x + blockIdx.y) % gridDim.y;
    } else {
        bx = blockIdx.x;
        by = blockIdx.y;
    }

    int x = bx * TILE_DIM + threadIdx.x;
    int y = by * TILE_DIM + threadIdx.y;

    for (int j = 0; j < TILE_DIM; j += BLOCK_ROWS) {
        if ((y + j) < M && x < N) {
            tile[threadIdx.y + j][threadIdx.x] = A[(y + j) * N + x];
        }
    }
    __syncthreads();

    x = by * TILE_DIM + threadIdx.x;
    y = bx * TILE_DIM + threadIdx.y;

    for (int j = 0; j < TILE_DIM; j += BLOCK_ROWS) {
        if ((y + j) < N && x < M) {
            B[(y + j) * M + x] = tile[threadIdx.x][threadIdx.y + j];
        }
    }
}

All Versions Compared

Non-Square Matrix Transpose

For non-square matrices ($M \ne N$), the grid dimensions differ and diagonal ordering is less effective:

// Non-square: A is M x N, B is N x M
// Grid: dim3((N + TILE_DIM - 1) / TILE_DIM, (M + TILE_DIM - 1) / TILE_DIM)
// The same tiled+padded kernel works without modification

void launch_transpose(const float* d_A, float* d_B, int M, int N) {
    dim3 block(TILE_DIM, BLOCK_ROWS);
    dim3 grid((N + TILE_DIM - 1) / TILE_DIM, (M + TILE_DIM - 1) / TILE_DIM);

    transpose_tiled_padded<<<grid, block>>>(d_A, d_B, M, N);
}

Batched Transpose

For batched operations (e.g., transposing attention heads in a transformer):

// Batched transpose: transpose B matrices of size M x N
__global__ void transpose_batched(const float* __restrict__ A,
                                   float* __restrict__ B,
                                   int M, int N, int batch_count) {
    __shared__ float tile[TILE_DIM][TILE_DIM + 1];

    int batch = blockIdx.z;
    if (batch >= batch_count) return;

    int matrix_offset = batch * M * N;
    const float* A_batch = A + matrix_offset;
    float* B_batch = B + matrix_offset;

    int x = blockIdx.x * TILE_DIM + threadIdx.x;
    int y = blockIdx.y * TILE_DIM + threadIdx.y;

    for (int j = 0; j < TILE_DIM; j += BLOCK_ROWS) {
        if ((y + j) < M && x < N) {
            tile[threadIdx.y + j][threadIdx.x] = A_batch[(y + j) * N + x];
        }
    }
    __syncthreads();

    x = blockIdx.y * TILE_DIM + threadIdx.x;
    y = blockIdx.x * TILE_DIM + threadIdx.y;

    for (int j = 0; j < TILE_DIM; j += BLOCK_ROWS) {
        if ((y + j) < N && x < M) {
            B_batch[(y + j) * M + x] = tile[threadIdx.x][threadIdx.y + j];
        }
    }
}

void launch_batched_transpose(const float* d_A, float* d_B,
                               int M, int N, int batch) {
    dim3 block(TILE_DIM, BLOCK_ROWS);
    dim3 grid((N + TILE_DIM - 1) / TILE_DIM,
              (M + TILE_DIM - 1) / TILE_DIM,
              batch);

    transpose_batched<<<grid, block>>>(d_A, d_B, M, N, batch);
}

In-Place Transpose (Square Matrices Only)

In-place transpose is possible only for square matrices. Each off-diagonal element pair is swapped:

__global__ void transpose_inplace(float* A, int N) {
    __shared__ float tile1[TILE_DIM][TILE_DIM + 1];
    __shared__ float tile2[TILE_DIM][TILE_DIM + 1];

    int bx = blockIdx.x;
    int by = blockIdx.y;

    // Only process upper triangle of block grid (including diagonal)
    if (bx > by) return;

    int x1 = bx * TILE_DIM + threadIdx.x;
    int y1 = by * TILE_DIM + threadIdx.y;
    int x2 = by * TILE_DIM + threadIdx.x;
    int y2 = bx * TILE_DIM + threadIdx.y;

    // Load both tiles
    for (int j = 0; j < TILE_DIM; j += BLOCK_ROWS) {
        if ((y1 + j) < N && x1 < N)
            tile1[threadIdx.y + j][threadIdx.x] = A[(y1 + j) * N + x1];
        if ((y2 + j) < N && x2 < N)
            tile2[threadIdx.y + j][threadIdx.x] = A[(y2 + j) * N + x2];
    }
    __syncthreads();

    // Write transposed
    // tile1 -> position of tile2, tile2 -> position of tile1
    for (int j = 0; j < TILE_DIM; j += BLOCK_ROWS) {
        if ((y2 + j) < N && x2 < N)
            A[(y2 + j) * N + x2] = tile1[threadIdx.x][threadIdx.y + j];
        if ((y1 + j) < N && x1 < N)
            A[(y1 + j) * N + x1] = tile2[threadIdx.x][threadIdx.y + j];
    }
}

Using cuBLAS for Production Transpose

For production code, cuBLAS provides a highly optimized transpose:

#include <cublas_v2.h>

void cublas_transpose(const float* d_A, float* d_B, int M, int N) {
    cublasHandle_t handle;
    cublasCreate(&handle);

    float alpha = 1.0f, beta = 0.0f;

    // cublasSgeam: C = alpha * op(A) + beta * op(B)
    // With beta = 0, this is just C = alpha * transpose(A)
    cublasSgeam(handle,
                CUBLAS_OP_T,  // Transpose A
                CUBLAS_OP_N,  // Don't transpose B (unused since beta=0)
                N, M,         // Output dimensions (N rows, M cols)
                &alpha,
                d_A, M,       // A is M x N in column-major = N x M in row-major
                &beta,
                d_B, N,       // B unused
                d_B, N);      // C output

    cublasDestroy(handle);
}

Summary

Matrix transpose is the canonical CUDA optimization problem because it demonstrates every major memory system concept. The naive version achieves 6% of copy bandwidth due to uncoalesced memory access. Tiling with shared memory fixes coalescing on both sides, reaching 72%. Adding one column of padding (TILE_DIM + 1) eliminates shared memory bank conflicts, jumping to 91%. Vectorized loads and diagonal block ordering add incremental gains to 94%. The tiled + padded version is the minimum viable optimization for production code. The takeaway generalizes: whenever a memory access pattern has inherent conflicts between the read and write sides, use shared memory as a transpose buffer with padding to avoid bank conflicts.