LeetGPU算子优化记录【Vector Addition,Matrix Multiplication,Matrix Transpose】
LeetGPU 简介
LeetGPU 是一个专门面向 GPU 编程挑战的在线平台,可以把它理解成 “GPU 版的 LeetCode”。用户可以直接在网页里编写 GPU 程序并在真实硬件上执行,不需要本地有 NVIDIA 显卡,也不用配环境(CUDA Toolkit、驱动、编译器那一套)。
LeetGPU 目前有 50+ 道题目,范围从矩阵运算到内存优化、kernel fusion 都有。从最基础的 vector add(就是你刚才那题)一路到比较硬核的并行归约、卷积、attention 等。
除了原生 CUDA C++,LeetGPU 也支持 Triton、PyTorch、tinygrad 等其他 GPU 编程框架,可以用自己熟悉的工具解题。
提交后会在真机上跑测试用例,给出结果是否正确。需要订阅pro后,才能够查看性能指标(运行时间、吞吐量等),还能和其他人的解答比较排名。
LeetGPU官网:leetgpu.com
github仓库:https://github.com/AlphaGPU/leetgpu-challenges
算子实现与优化(测试平台:NVIDIA-A100-80GB)
1. Vector Addition
https://leetgpu.com/challenges/vector-addition
#include <cuda_runtime.h>
__global__ void vector_add(const float* A, const float* B, float* C, int N) {
int tid = blockIdx.x * blockDim.x + threadIdx.x;
if (tid < N) {
C[tid] = A[tid] + B[tid];
}
}
// A, B, C are device pointers (i.e. pointers to memory on the GPU)
extern "C" void solve(const float* A, const float* B, float* C, int N) {
int threadsPerBlock = 256;
int blocksPerGrid = (N + threadsPerBlock - 1) / threadsPerBlock;
vector_add<<<blocksPerGrid, threadsPerBlock>>>(A, B, C, N);
cudaDeviceSynchronize();
}

向量化访存:
#include <cuda_runtime.h>
__global__ void add_vectorized(const float *A, const float *B, float *C, int N) {
int vec_idx = blockIdx.x * blockDim.x + threadIdx.x;
int base_idx = vec_idx * 4;
if (base_idx + 3 < N) { // 确保剩余元素足够一个float4
// 类型转换:将float*转为float4*,直接读取4个连续元素
float4 a = reinterpret_cast<const float4*>(A)[vec_idx];
float4 b = reinterpret_cast<const float4*>(B)[vec_idx];
float4 c;
c.x = a.x + b.x;
c.y = a.y + b.y;
c.z = a.z + b.z;
c.w = a.w + b.w;
reinterpret_cast<float4*>(C)[vec_idx] = c;
}
else if (base_idx < N) {
for (int i = 0; base_idx + i < N; i++) {
C[base_idx + i] = A[base_idx + i] + B[base_idx + i];
}
}
}
// A, B, C are device pointers (i.e. pointers to memory on the GPU)
extern "C" void solve(const float* A, const float* B, float* C, int N) {
int threadsPerBlock = 256;
// 每个线程处理 4 个元素,总线程数按 ceil(N/4) 算
int totalThreads = (N + 3) / 4;
int blocksPerGrid = (totalThreads + threadsPerBlock - 1) / threadsPerBlock;
add_vectorized<<<blocksPerGrid, threadsPerBlock>>>(A, B, C, N);
cudaDeviceSynchronize();
}
但是实际上性能并没有提升
2. Matrix Multiplication
https://leetgpu.com/challenges/matrix-multiplication
朴素实现:
#include <cuda_runtime.h>
__global__ void matrix_multiplication_kernel(const float* A, const float* B, float* C, int M, int N,
int K) {
int col = blockIdx.x * blockDim.x + threadIdx.x;
int row = blockIdx.y * blockDim.y + threadIdx.y;
if (row < M && col < K) {
float sum = 0.0f;
for (int i = 0; i < N; ++i) {
sum += A[row * N + i] * B[i * K + col];
}
C[row * K + col] = sum;
}
}
// A, B, C are device pointers (i.e. pointers to memory on the GPU)
extern "C" void solve(const float* A, const float* B, float* C, int M, int N, int K) {
dim3 threadsPerBlock(16, 16);
dim3 blocksPerGrid((K + threadsPerBlock.x - 1) / threadsPerBlock.x,
(M + threadsPerBlock.y - 1) / threadsPerBlock.y);
matrix_multiplication_kernel<<<blocksPerGrid, threadsPerBlock>>>(A, B, C, M, N, K);
cudaDeviceSynchronize();
}

Tiling + 共享内存:
#include <cuda_runtime.h>
#define TILE 16
__global__ void matrix_multiplication_kernel(const float* A, const float* B, float* C, int M, int N,
int K) {
__shared__ float As[TILE][TILE];
__shared__ float Bs[TILE][TILE];
int tx = threadIdx.x;
int ty = threadIdx.y;
int col = blockIdx.x * TILE + tx;
int row = blockIdx.y * TILE + ty;
float sum = 0.0f;
int num_tiles = (N + TILE - 1) / TILE;
for (int t = 0; t < num_tiles; ++t) {
int a_col = t * TILE + tx;
if (row < M && a_col < N) {
As[ty][tx] = A[row * N + a_col];
} else {
As[ty][tx] = 0.0f;
}
int b_row = t * TILE + ty;
if (b_row < N && col < K) {
Bs[ty][tx] = B[b_row * K + col];
} else {
Bs[ty][tx] = 0.0f;
}
__syncthreads();
#pragma unroll
for (int k = 0; k < TILE; ++k) {
sum += As[ty][k] * Bs[k][tx];
}
__syncthreads();
}
if (row < M && col < K) {
C[row * K + col] = sum;
}
}
// A, B, C are device pointers (i.e. pointers to memory on the GPU)
extern "C" void solve(const float* A, const float* B, float* C, int M, int N, int K) {
dim3 threadsPerBlock(16, 16);
dim3 blocksPerGrid((K + threadsPerBlock.x - 1) / threadsPerBlock.x,
(M + threadsPerBlock.y - 1) / threadsPerBlock.y);
matrix_multiplication_kernel<<<blocksPerGrid, threadsPerBlock>>>(A, B, C, M, N, K);
cudaDeviceSynchronize();
}

进一步优化:
- 每线程计算多个输出(Thread Coarsening)
- 向量化访存
- 用双缓冲实现软件预取
- 异步拷贝
(claude opus 4.7生成)
#include <cuda_runtime.h>
#include <cuda/pipeline>
#include <cooperative_groups.h>
#include <cooperative_groups/memcpy_async.h>
namespace cg = cooperative_groups;
#define BM 128
#define BN 128
#define BK 16
#define TM 8
#define TN 8
#define NUM_THREADS 256
#define STAGES 2
__global__ void matmul_kernel(const float* __restrict__ A,
const float* __restrict__ B,
float* __restrict__ C,
int M, int N, int K,
bool use_vec_A, bool use_vec_B) {
__shared__ float As[STAGES][BM][BK];
__shared__ float Bs[STAGES][BK][BN];
const int tid = threadIdx.y * blockDim.x + threadIdx.x;
const int block_row = blockIdx.y * BM;
const int block_col = blockIdx.x * BN;
const int thread_row = (tid / (BN / TN)) * TM;
const int thread_col = (tid % (BN / TN)) * TN;
float sum[TM][TN] = {0.0f};
float a_reg[TM];
float b_reg[TN];
constexpr int A_LOADS_PER_THREAD = (BM * BK) / (NUM_THREADS * 4); // 2
constexpr int B_LOADS_PER_THREAD = (BK * BN) / (NUM_THREADS * 4); // 2
cuda::pipeline<cuda::thread_scope_thread> pipe = cuda::make_pipeline();
auto load_A = [&](int t, int stage) {
#pragma unroll
for (int i = 0; i < A_LOADS_PER_THREAD; ++i) {
int float4_idx = tid + i * NUM_THREADS;
int a_row = float4_idx / (BK / 4);
int a_col = (float4_idx % (BK / 4)) * 4;
int g_row = block_row + a_row;
int g_col = t * BK + a_col;
float4* dst = reinterpret_cast<float4*>(&As[stage][a_row][a_col]);
if (g_row >= M) {
*dst = make_float4(0.f, 0.f, 0.f, 0.f);
continue;
}
// 向量化路径:要求 N 是 4 的倍数(起始对齐)且 4 个元素都在界内
if (use_vec_A && g_col + 4 <= N) {
const float4* src =
reinterpret_cast<const float4*>(&A[g_row * N + g_col]);
cuda::memcpy_async(dst, src, sizeof(float4), pipe);
} else {
// 标量路径:逐个加载,越界填 0
float tmp[4] = {0.f, 0.f, 0.f, 0.f};
#pragma unroll
for (int k = 0; k < 4; ++k) {
if (g_col + k < N) tmp[k] = A[g_row * N + g_col + k];
}
*dst = make_float4(tmp[0], tmp[1], tmp[2], tmp[3]);
}
}
};
auto load_B = [&](int t, int stage) {
#pragma unroll
for (int i = 0; i < B_LOADS_PER_THREAD; ++i) {
int float4_idx = tid + i * NUM_THREADS;
int b_row = float4_idx / (BN / 4);
int b_col = (float4_idx % (BN / 4)) * 4;
int g_row = t * BK + b_row;
int g_col = block_col + b_col;
float4* dst = reinterpret_cast<float4*>(&Bs[stage][b_row][b_col]);
if (g_row >= N) {
*dst = make_float4(0.f, 0.f, 0.f, 0.f);
continue;
}
// 向量化路径:要求 K 是 4 的倍数
if (use_vec_B && g_col + 4 <= K) {
const float4* src =
reinterpret_cast<const float4*>(&B[g_row * K + g_col]);
cuda::memcpy_async(dst, src, sizeof(float4), pipe);
} else {
float tmp[4] = {0.f, 0.f, 0.f, 0.f};
#pragma unroll
for (int k = 0; k < 4; ++k) {
if (g_col + k < K) tmp[k] = B[g_row * K + g_col + k];
}
*dst = make_float4(tmp[0], tmp[1], tmp[2], tmp[3]);
}
}
};
const int num_tiles = (N + BK - 1) / BK;
pipe.producer_acquire();
load_A(0, 0);
load_B(0, 0);
pipe.producer_commit();
for (int t = 0; t < num_tiles; ++t) {
int cur_stage = t % STAGES;
int next_stage = (t + 1) % STAGES;
if (t + 1 < num_tiles) {
pipe.producer_acquire();
load_A(t + 1, next_stage);
load_B(t + 1, next_stage);
pipe.producer_commit();
}
pipe.consumer_wait();
__syncthreads();
#pragma unroll
for (int k = 0; k < BK; ++k) {
#pragma unroll
for (int i = 0; i < TM; ++i)
a_reg[i] = As[cur_stage][thread_row + i][k];
#pragma unroll
for (int j = 0; j < TN; ++j)
b_reg[j] = Bs[cur_stage][k][thread_col + j];
#pragma unroll
for (int i = 0; i < TM; ++i)
#pragma unroll
for (int j = 0; j < TN; ++j)
sum[i][j] += a_reg[i] * b_reg[j];
}
pipe.consumer_release();
__syncthreads();
}
// 写回(一律用标量,安全)
#pragma unroll
for (int i = 0; i < TM; ++i) {
int g_row = block_row + thread_row + i;
if (g_row >= M) continue;
#pragma unroll
for (int j = 0; j < TN; ++j) {
int g_col = block_col + thread_col + j;
if (g_col < K) {
C[g_row * K + g_col] = sum[i][j];
}
}
}
}
extern "C" void solve(const float* A, const float* B, float* C, int M, int N, int K) {
// 运行时决定是否启用向量化加载:要求 N、K 是 4 的倍数
// (另外还应检查指针本身 16 字节对齐,通常 cudaMalloc 保证这一点)
bool use_vec_A = (N % 4 == 0);
bool use_vec_B = (K % 4 == 0);
dim3 threadsPerBlock(NUM_THREADS, 1, 1);
dim3 blocksPerGrid((K + BN - 1) / BN,
(M + BM - 1) / BM);
matmul_kernel<<<blocksPerGrid, threadsPerBlock>>>(
A, B, C, M, N, K, use_vec_A, use_vec_B);
cudaDeviceSynchronize();
}

3. Matrix Transpose
https://leetgpu.com/challenges/matrix-transpose
朴素实现:
#include <cuda_runtime.h>
__global__ void matrix_transpose_kernel(const float* input, float* output, int rows, int cols) {
int col = blockIdx.x * blockDim.x + threadIdx.x;
int row = blockIdx.y * blockDim.y + threadIdx.y;
if(row < rows && col < cols) {
output[col * rows + row] = input[row * cols + col];
}
}
// input, output are device pointers (i.e. pointers to memory on the GPU)
extern "C" void solve(const float* input, float* output, int rows, int cols) {
dim3 threadsPerBlock(16, 16);
dim3 blocksPerGrid((cols + threadsPerBlock.x - 1) / threadsPerBlock.x,
(rows + threadsPerBlock.y - 1) / threadsPerBlock.y);
matrix_transpose_kernel<<<blocksPerGrid, threadsPerBlock>>>(input, output, rows, cols);
cudaDeviceSynchronize();
}

Tiling + 共享内存:
#include <cuda_runtime.h>
#define TILE_DIM 16
__global__ void matrix_transpose_kernel(const float* input, float* output, int rows, int cols) {
// +1 避免 bank conflict
__shared__ float tile[TILE_DIM][TILE_DIM + 1];
int x = blockIdx.x * TILE_DIM + threadIdx.x; // 列
int y = blockIdx.y * TILE_DIM + threadIdx.y; // 行
// 合并读取:从 input 读入 tile
if (y < rows && x < cols) {
tile[threadIdx.y][threadIdx.x] = input[y * cols + x];
}
__syncthreads();
// 交换 block 坐标,使得写入也是合并的
x = blockIdx.y * TILE_DIM + threadIdx.x; // 输出的列(对应原来的行)
y = blockIdx.x * TILE_DIM + threadIdx.y; // 输出的行(对应原来的列)
if (y < cols && x < rows) {
output[y * rows + x] = tile[threadIdx.x][threadIdx.y];
}
}
extern "C" void solve(const float* input, float* output, int rows, int cols) {
dim3 threadsPerBlock(16, 16);
dim3 blocksPerGrid((cols + threadsPerBlock.x - 1) / threadsPerBlock.x,
(rows + threadsPerBlock.y - 1) / threadsPerBlock.y);
matrix_transpose_kernel<<<blocksPerGrid, threadsPerBlock>>>(input, output, rows, cols);
cudaDeviceSynchronize();
}
为什么非要交换 block,而不能只交换 thread? 因为如果只交换 threadIdx,warp 内相邻线程(相邻 threadIdx.x)写入的地址就会相隔 rows 个元素,这就是非合并访问,性能很差。
交换 blockIdx 的同时保持 threadIdx.x 走行方向,warp 内相邻线程仍然写入连续地址——这就是为什么要“交换 block 坐标而不是在 shared memory 里换”。

进一步优化:
- 向量化访存
- 对齐时用 float4
- swizzle(异或变换索引)
#include <cuda_runtime.h>
#define TILE_DIM 32
#define BLOCK_ROWS 8
// swizzle: 把 (row, col) 的 col 异或上 row,保证同一 warp 内访问的地址
// 落在不同的 bank 里。由于异或是自逆的,读写用同一个函数即可。
__device__ __forceinline__ int swz(int row, int col) {
return col ^ row;
}
__global__ void matrix_transpose_kernel(const float* __restrict__ input,
float* __restrict__ output,
int rows, int cols) {
// 零 padding 的 shared memory,靠 swizzle 避免 bank conflict
__shared__ float tile[TILE_DIM][TILE_DIM];
const int tx = threadIdx.x;
const int ty = threadIdx.y;
// ========== 读取阶段 ==========
// 每个 block 负责 32×32 tile;32×8 个线程,每个线程处理 4 行
const int x_in = blockIdx.x * TILE_DIM + tx; // 输入列
const int y_in = blockIdx.y * TILE_DIM + ty; // 输入行起点
// 判断当前 tile 是否完全落在矩阵内(避免边界检查的开销)
const bool full_tile_in =
(blockIdx.x * TILE_DIM + TILE_DIM <= cols) &&
(blockIdx.y * TILE_DIM + TILE_DIM <= rows);
// 判断是否可以用 float4:tile 完整 + 起始地址 16 字节对齐 + cols 是 4 的倍数
const bool can_vec4 =
full_tile_in &&
((cols & 3) == 0) &&
(((blockIdx.x * TILE_DIM) & 3) == 0);
if (can_vec4 && (tx & 3) == 0) {
// 向量化路径:tx 每 4 个线程中只有 tx%4==0 负责加载
// 每次加载 float4 = 4 个连续列
#pragma unroll
for (int j = 0; j < TILE_DIM; j += BLOCK_ROWS) {
const float4 v = *reinterpret_cast<const float4*>(
&input[(y_in + j) * cols + x_in]);
tile[ty + j][swz(ty + j, tx + 0)] = v.x;
tile[ty + j][swz(ty + j, tx + 1)] = v.y;
tile[ty + j][swz(ty + j, tx + 2)] = v.z;
tile[ty + j][swz(ty + j, tx + 3)] = v.w;
}
} else if (full_tile_in) {
// 标量路径(tile 完整但不能向量化)
#pragma unroll
for (int j = 0; j < TILE_DIM; j += BLOCK_ROWS) {
tile[ty + j][swz(ty + j, tx)] = input[(y_in + j) * cols + x_in];
}
} else {
// 边界路径:tile 越界,需要逐元素检查
#pragma unroll
for (int j = 0; j < TILE_DIM; j += BLOCK_ROWS) {
const int y = y_in + j;
if (x_in < cols && y < rows) {
tile[ty + j][swz(ty + j, tx)] = input[y * cols + x_in];
}
}
}
__syncthreads();
// ========== 写回阶段 ==========
// 交换 blockIdx.x 和 blockIdx.y,保证写入仍然是合并访问
const int x_out = blockIdx.y * TILE_DIM + tx; // 输出列 (= 原行)
const int y_out = blockIdx.x * TILE_DIM + ty; // 输出行 (= 原列)
const bool full_tile_out =
(blockIdx.y * TILE_DIM + TILE_DIM <= rows) &&
(blockIdx.x * TILE_DIM + TILE_DIM <= cols);
if (full_tile_out) {
// 读取 shared memory 时做转置:tile[col][row] -> output[row][col]
// 这里原始坐标是 (row = tx, col = ty + j),转置后写入
#pragma unroll
for (int j = 0; j < TILE_DIM; j += BLOCK_ROWS) {
// 注意:写入 global 的是转置后的数据
// 原 tile 中的 (row=tx, col=ty+j) 对应输出的 (y_out+j, x_out)
output[(y_out + j) * rows + x_out] = tile[tx][swz(tx, ty + j)];
}
} else {
#pragma unroll
for (int j = 0; j < TILE_DIM; j += BLOCK_ROWS) {
const int y = y_out + j;
if (x_out < rows && y < cols) {
output[y * rows + x_out] = tile[tx][swz(tx, ty + j)];
}
}
}
}
extern "C" void solve(const float* input, float* output, int rows, int cols) {
dim3 threadsPerBlock(TILE_DIM, BLOCK_ROWS); // 32 × 8 = 256 线程
dim3 blocksPerGrid((cols + TILE_DIM - 1) / TILE_DIM,
(rows + TILE_DIM - 1) / TILE_DIM);
matrix_transpose_kernel<<<blocksPerGrid, threadsPerBlock>>>(
input, output, rows, cols);
cudaDeviceSynchronize();
}
性能有小幅度的提升
更多推荐



所有评论(0)