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();
}

性能有小幅度的提升
在这里插入图片描述

Logo

汇聚全球AI编程工具,助力开发者即刻编程。

更多推荐