GPU 编程实战——GPU 上的线性代数基础

408 阅读11分钟

使用 cuBLAS 进行稠密向量与矩阵操作

cuBLAS 简介

cuBLAS 是 NVIDIA 提供的 GPU 加速线性代数库,涵盖从简单的向量加法到大规模矩阵乘法等高度优化的基础运算。该库由 NVIDIA 针对 CUDA 硬件精心调优,利用了专用指令、内存布局和批量并行技术,能够充分发挥 GPU 性能。cuBLAS 遵循经典的 BLAS(Basic Linear Algebra Subprograms)API,科学与工程领域用户对此接口十分熟悉,但其在 GPU 上的高性能实现使其成为机器学习、仿真、信号处理等领域稠密线性代数计算的首选。

尽管我们可以自己用 CUDA kernel 编写向量和矩阵运算(正如前面章节所练习的),但这些手写 kernel 很难达到 cuBLAS 的性能,因为:

  • cuBLAS 使用了分块(tiling)、寄存器阻塞(register blocking)、内存合并访问(coalescing)和操作融合等高级技术;
  • 库会根据实际 GPU 架构动态调整并优化;
  • cuBLAS 提供了从最基础的向量加法到高级的 LU 分解等全系列例程,不必我们自行实现和调试。

在 CuPy 中使用 cuBLAS

CuPy 提供了对 cuBLAS 的无缝 Python 接口,让我们像调用 NumPy 函数一样调用高性能 GPU 例程。

稠密向量运算

以下示例演示向量加法与点积:

import cupy as cp
import numpy as np
import time

N = 1_000_000
a = cp.random.rand(N, dtype=cp.float32)
b = cp.random.rand(N, dtype=cp.float32)

# 向量加法(逐元素相加,底层调用 cuBLAS)
start = time.time()
c = a + b
cp.cuda.Stream.null.synchronize()
elapsed_add = time.time() - start

# 点积(内积,cuBLAS 加速)
start = time.time()
dot_result = cp.dot(a, b)
cp.cuda.Stream.null.synchronize()
elapsed_dot = time.time() - start

print(f"向量加法(cuBLAS)耗时:{elapsed_add:.5f} 秒")
print(f"点积(cuBLAS)耗时:{elapsed_dot:.5f} 秒")

稠密矩阵-矩阵乘法

矩阵乘法是 cuBLAS 最具代表性的应用。下面演示 2048×2048 的矩阵相乘:

M, K, N_ = 2048, 2048, 2048
A = cp.random.rand(M, K, dtype=cp.float32)
B = cp.random.rand(K, N_, dtype=cp.float32)

start = time.time()
C = cp.matmul(A, B)  # cuBLAS 驱动的 GEMM
cp.cuda.Stream.null.synchronize()
elapsed_mm = time.time() - start

print(f"矩阵乘法(cuBLAS)耗时:{elapsed_mm:.5f} 秒")

手写 Kernel 对比

为了直观对比,我们实现一个最简朴的矩阵乘法 kernel,不做分块或其他优化:

from cupy import RawKernel

matmul_kernel = RawKernel(r'''
extern "C" __global__
void matmul(const float* A, const float* B, float* C,
            int M, int K, int N) {
    int row = blockDim.y * blockIdx.y + threadIdx.y;
    int col = blockDim.x * blockIdx.x + threadIdx.x;
    if (row < M && col < N) {
        float sum = 0.0f;
        for (int i = 0; i < K; ++i)
            sum += A[row * K + i] * B[i * N + col];
        C[row * N + col] = sum;
    }
}
''', 'matmul')

C_manual = cp.zeros((M, N_), dtype=cp.float32)
block    = (16, 16)
grid     = (N_ // block[0], M // block[1])

start = time.time()
matmul_kernel(grid, block,
    (A, B, C_manual, np.int32(M),
     np.int32(K), np.int32(N_))
)
cp.cuda.Stream.null.synchronize()
elapsed_manual = time.time() - start

print(f"矩阵乘法(手写 kernel)耗时:{elapsed_manual:.5f} 秒")

性能对比与验证

# 验证两者结果差异
max_diff = cp.max(cp.abs(C - C_manual))
print("cuBLAS 与手写 kernel 结果最大差异:", max_diff)

# 计算加速比
print(f"加速比(手写 / cuBLAS):{elapsed_manual / elapsed_mm:.2f}×")

通常会发现:

  • cuBLAS 相对于天真的手写 Kernel 有 5–10 倍甚至更高的速度提升;
  • cuBLAS 在数值精度上也十分可靠;
  • 随着问题规模增大,特别是大规模矩阵乘法或批量操作时,性能差距进一步拉大。

CuPy 与 cuBLAS 的紧密集成,让我们能够像写 NumPy 代码一样简单地调用高性能 GPU 线性代数例程,从而轻松扩展科学、工程和机器学习工作流到现代 GPU。

简单矩阵–向量乘法

矩阵–向量乘法是线性代数中的基本操作,广泛应用于科学计算、数据分析和机器学习,构成神经网络推理、迭代求解等流程的核心。要在现代硬件上构建可扩展且高性能的管道,高效实现此操作至关重要。

使用 PyCUDA 手写内核

假设我们有一个大小为 M×NM\times N 的矩阵 AA 和长度为 NN 的向量 xx,目标是计算结果向量 y=A×xy = A \times x。步骤如下。

准备数据

import numpy as np
import pycuda.autoinit
import pycuda.driver as drv
import pycuda.gpuarray as gpuarray
from pycuda.compiler import SourceModule
import time

M, N = 2048, 2048
A_host = np.random.rand(M, N).astype(np.float32)
x_host = np.random.rand(N).astype(np.float32)

A_gpu = gpuarray.to_gpu(A_host)
x_gpu = gpuarray.to_gpu(x_host)
y_gpu = gpuarray.zeros(M, dtype=np.float32)

PyCUDA 内核实现

kernel_code = """
__global__ void matvec(const float *A, const float *x, float *y, int M, int N)
{
    int row = blockIdx.x * blockDim.x + threadIdx.x;
    if (row < M) {
        float sum = 0.0f;
        for (int col = 0; col < N; ++col)
            sum += A[row * N + col] * x[col];
        y[row] = sum;
    }
}
"""
mod    = SourceModule(kernel_code)
matvec = mod.get_function("matvec")

启动内核并计时

threads_per_block = 256
blocks_per_grid   = (M + threads_per_block - 1) // threads_per_block

start = time.time()
matvec(
    A_gpu, x_gpu, y_gpu, np.int32(M), np.int32(N),
    block=(threads_per_block, 1, 1),
    grid =(blocks_per_grid,   1   )
)
drv.Context.synchronize()
elapsed_pycuda = time.time() - start

y_result_pycuda = y_gpu.get()

使用 cuBLAS(通过 CuPy)

CuPy 提供了直接且高度优化的矩阵–向量乘法接口,调用了 cuBLAS 的 gemv 例程。我们可以用 cp.dot@ 操作符:

import cupy as cp

A_cupy = cp.array(A_host)
x_cupy = cp.array(x_host)

start = time.time()
y_cupy = cp.dot(A_cupy, x_cupy)
cp.cuda.Stream.null.synchronize()
elapsed_cublas = time.time() - start

y_result_cublas = cp.asnumpy(y_cupy)

结果比较与验证

import numpy as np

max_diff = np.max(np.abs(y_result_pycuda - y_result_cublas))
print("最大差异:", max_diff)
print(f"PyCUDA 内核耗时: {elapsed_pycuda:.5f} 秒")
print(f"cuBLAS(CuPy)耗时:{elapsed_cublas:.5f} 秒")
print(f"加速比(PyCUDA / cuBLAS):{elapsed_pycuda / elapsed_cublas:.2f}×")

通常情况下,除非有极为特殊的需求或非标准数据布局,否则在 GPU 上进行稠密矩阵–向量(及矩阵–矩阵)操作时,使用 cuBLAS 是最优选。它在速度、稳定性和简易性上都具有显著优势,是科学计算和机器学习领域 GPU 加速的基础。

批量 GEMM(通用矩阵乘法)

为什么需要批量 GEMM?

在许多科学、工程和机器学习应用中,我们经常不止对一对矩阵进行乘法,而是一次性对数百或数千对矩阵进行操作。例如神经网络的小批量推理、3D 图形变换和物理系统模拟等场景。若将每对矩阵乘法依次循环调用,会导致 GPU 空闲周期增多,吞吐量远未饱和。

cuBLAS 提供了批量 GEMM 接口(batched GEMM),允许我们将一批矩阵乘法任务一次性提交给 GPU,库内部并行处理所有批次,与顺序执行相比能显著提升性能和总吞吐量。

使用 cuBLAS(CuPy)执行批量 GEMM

准备多组矩阵

以下示例生成 1000 对 32×3232\times32 矩阵:

import cupy as cp
import numpy as np
import time

batch_size = 1000
M = K = N = 32
A = cp.random.rand(batch_size, M, K).astype(cp.float32)
B = cp.random.rand(batch_size, K, N).astype(cp.float32)

批量矩阵乘法

在 CuPy 中,对具有批次维度的数组调用 cp.matmul@ 运算符时,即会调用 cuBLAS 的 batched GEMM:

start = time.time()
C_batched = cp.matmul(A, B)  # 输出形状 (batch_size, M, N)
cp.cuda.Stream.null.synchronize()
elapsed_batched = time.time() - start

顺序矩阵乘法

为了对比吞吐,这里逐对矩阵相乘:

C_sequential = cp.empty_like(C_batched)
start = time.time()
for i in range(batch_size):
    C_sequential[i] = cp.matmul(A[i], B[i])
cp.cuda.Stream.null.synchronize()
elapsed_sequential = time.time() - start

性能对比与验证

max_diff = cp.max(cp.abs(C_batched - C_sequential))
print("批量与顺序结果最大差异:", float(max_diff))
print(f"批量 GEMM 耗时:      {elapsed_batched:.5f} 秒")
print(f"顺序 GEMM 耗时:      {elapsed_sequential:.5f} 秒")
print(f"加速比:{elapsed_sequential / elapsed_batched:.2f}×")

通常,批量 GEMM 比顺序执行快一个数量级甚至更多,因为库能在内部并行化运算,重叠内存传输与计算,实现 GPU 资源的极致利用。该方法在从小规模实验到生产级大批量工作流的多种批次大小和矩阵形状下均表现出优异的可扩展性和性能。

使用 CuPy 的 dotmatmul

CuPy 的 dotmatmul 函数概述

CuPy 提供了熟悉的高层线性代数接口:cupy.dotcupy.matmul。它们可用于对向量和矩阵进行简洁、富有表现力且高效的数学运算。两者底层均调用了 cuBLAS,在 NVIDIA GPU 上具有高度优化性能。

  • cupy.dot

    • 对于一维向量,计算标量(内积)。
    • 对于二维矩阵与向量,执行矩阵–向量乘法。
    • 对于两个二维矩阵,执行矩阵–矩阵乘法。
  • cupy.matmul

    • 用于矩阵乘法,支持广播(broadcasting)和对更高维数组的批量(batched)操作。

使用这些函数无需手写循环或管理内核启动,我们可以专注于科学计算或数据分析,同时充分享受 GPU 加速带来的性能优势。

高层向量与矩阵运算示例

向量点积

import cupy as cp
import numpy as np

N = 500_000
a = cp.random.rand(N, dtype=cp.float32)
b = cp.random.rand(N, dtype=cp.float32)

# 使用 CuPy 计算点积
dot_result = cp.dot(a, b)
print("点积结果:", float(dot_result))

矩阵–向量乘法

M, N = 1024, 1024
A = cp.random.rand(M, N, dtype=cp.float32)
x = cp.random.rand(N, dtype=cp.float32)

# 矩阵–向量乘法
y = cp.dot(A, x)
print("矩阵–向量结果形状:", y.shape)

矩阵–矩阵乘法

对二维矩阵,dotmatmul 效果相同:

K = 512
B = cp.random.rand(N, K, dtype=cp.float32)

# 使用 dot
C1 = cp.dot(A, B)
print("矩阵–矩阵结果形状(dot):", C1.shape)

# 使用 matmul
C2 = cp.matmul(A, B)
print("矩阵–矩阵结果形状(matmul):", C2.shape)

# 使用 @ 运算符(等同于 matmul)
C3 = A @ B
print("矩阵–矩阵结果形状(@):", C3.shape)

批量矩阵乘法

对带批次维度的更高维数组,matmul 自动调用 cuBLAS 的批量 GEMM:

batch    = 100
A_batch = cp.random.rand(batch, M, N, dtype=cp.float32)
B_batch = cp.random.rand(batch, N, K, dtype=cp.float32)

C_batch = cp.matmul(A_batch, B_batch)
print("批量矩阵–矩阵乘法结果形状:", C_batch.shape)  # (batch, M, K)

cp.dot 处理向量内积、矩阵–向量与 2D 矩阵–矩阵乘法;cp.matmul(及 @)则对 2D 及更高维批量矩阵运算提供统一支持。两者均调用高度优化的 GPU 代码,使我们能够用简洁、可读的高层 API 完成标准线性代数任务,而无需手写任何底层优化或内核开发。

评估数值精度

在本章中,我们使用 GPU 加速了线性代数运算,既有自定义内核,也有通过 CuPy 调用的高效 cuBLAS 例程。虽然性能提升往往十分显著,但同样重要的是确保 GPU 计算结果与传统 CPU 工作流相比具有一致性和可靠性。GPU 与 CPU 在浮点运算时可能会因指令执行顺序、精度差异、并行化程度以及融合操作的使用而产生细微差别。

在科学或工程计算中,验证数值精度至关重要。本节我们将使用标准误差度量,比较 GPU 与 CPU 的输出。

1. 准备相同输入

以矩阵–矩阵乘法为例:

import numpy as np
import cupy as cp

M, N, K = 512, 512, 512
A_cpu = np.random.rand(M, N).astype(np.float32)
B_cpu = np.random.rand(N, K).astype(np.float32)

A_gpu = cp.array(A_cpu)
B_gpu = cp.array(B_cpu)

2. 分别计算 CPU 与 GPU 结果

# CPU(NumPy)
C_cpu = np.dot(A_cpu, B_cpu)

# GPU(CuPy,调用 cuBLAS)
C_gpu       = cp.dot(A_gpu, B_gpu)
C_gpu_cpu   = cp.asnumpy(C_gpu)  # 拷回主机,便于比较

3. 计算误差度量

常用的误差度量包括:

  • 最大绝对误差(L∞ 范数):对应元素最大绝对差。
  • 均方根误差(RMSE,L₂ 范数):差值平方的均值后开方。
  • 相对误差:RMSE 与参考值范数之比。
abs_diff       = np.abs(C_cpu - C_gpu_cpu)
max_error      = np.max(abs_diff)
rmse           = np.sqrt(np.mean((C_cpu - C_gpu_cpu) ** 2))
relative_error = rmse / np.linalg.norm(C_cpu)

print(f"最大绝对误差:{max_error:e}")
print(f"均方根误差:{rmse:e}")
print(f"相对误差(L2 范数):{relative_error:e}")

对于对精度要求更高的场景,可考虑使用 float64(双精度),但需注意许多消费级 GPU 在双精度上的性能可能明显低于单精度。此精度验证流程是科学计算工作流中的关键环节,有助于确保 GPU 加速方案的可靠性。

总结

本章重点介绍了如何利用 GPU 强大的计算能力来完成核心线性代数任务,并结合低级与高级两种方法构建了稳固的理解。

  1. cuBLAS 介绍

    • 学习了 NVIDIA 提供的高度优化线性代数库 cuBLAS,及其如何在单精度向量和矩阵运算上比自定义内核取得更显著的加速效果。
  2. 稠密向量与矩阵操作

    • 通过 CuPy 无缝调用 cuBLAS 执行向量加法、矩阵–向量乘法和大规模矩阵–矩阵乘法,持续关注性能对比。
  3. 批量 GEMM

    • 演示了一次性对多对矩阵执行 GEMM(批量矩阵乘法)如何显著提升 GPU 吞吐,相比循环顺序执行效率更高。
  4. 高层接口 dotmatmul

    • 使用 cupy.dotcupy.matmul(及 @)在不同维度上完成内积、矩阵乘法和批量乘法,代码简洁同时调用 cuBLAS 获取高性能。
  5. 数值精度评估

    • 通过最大绝对误差、均方根误差和相对误差,将 GPU 计算结果与 CPU 参考值进行严格对比,确保加速方案的可靠性。

通过本章内容,您已掌握在 GPU 上高效完成稠密线性代数运算的完整流程,从库调用到自定义内核,从性能优化到结果验证,为科研、工程和机器学习等领域的高性能计算打下了坚实基础。