CUDA 加速的基础线性代数库cuBLAS

原创于 2025-05-22 22:59:05 发布 · 445 阅读

10 ·

CC 4.0 BY-SA版权

文章标签：

#线性代数

NVIDIA 专栏收录该内容

14 篇文章

订阅专栏

cuBLAS 是 NVIDIA 提供的 GPU 加速版 BLAS (Basic Linear Algebra Subprograms) 实现，专门针对 NVIDIA GPU 进行了高度优化。

cuBLAS 核心特性

完整 BLAS 功能支持：
- Level 1：向量-向量运算 (如点积、axpy)
- Level 2：矩阵-向量运算 (如矩阵向量乘法)
- Level 3：矩阵-矩阵运算 (如矩阵乘法 GEMM)
数据类型支持：
- 单精度 (S)、双精度 (D)
- 复数 (C/Z)
- 半精度 (H) 和混合精度计算
- Tensor Core 加速 (支持 FP16/FP32 混合精度)
高性能特性：
- 异步执行
- 流并行
- 批处理操作

cuBLAS API 结构

cuBLAS 提供两种 API 风格：

传统 API (legacy API)：
- 较老的函数命名风格 (如 cublasSgemm)
- 需要显式管理设备指针
新版 API (自 CUDA 6.0)：
- 更现代的接口设计
- 支持指针模式自动管理
- 更完善的错误处理

基本使用模式

1. 初始化和资源管理

cpp

#include <cublas_v2.h>

cublasHandle_t handle;
cublasCreate(&handle); // 创建句柄

// 设置数学模式 (如允许Tensor Core使用)
cublasSetMathMode(handle, CUBLAS_TENSOR_OP_MATH);

// ...执行运算...

cublasDestroy(handle); // 销毁句柄

2. 矩阵乘法示例 (SGEMM)

cpp

// 矩阵尺寸
int m = 1024, n = 1024, k = 1024;

// 主机内存分配和初始化
float *h_A = (float*)malloc(m*k*sizeof(float));
float *h_B = (float*)malloc(k*n*sizeof(float));
float *h_C = (float*)malloc(m*n*sizeof(float));

// 设备内存分配
float *d_A, *d_B, *d_C;
cudaMalloc(&d_A, m*k*sizeof(float));
cudaMalloc(&d_B, k*n*sizeof(float));
cudaMalloc(&d_C, m*n*sizeof(float));

// 数据传输到设备
cublasSetVector(m*k, sizeof(float), h_A, 1, d_A, 1);
cublasSetVector(k*n, sizeof(float), h_B, 1, d_B, 1);

// 执行矩阵乘法 (C = α*A*B + β*C)
const float alpha = 1.0f, beta = 0.0f;
cublasSgemm(handle, 
           CUBLAS_OP_N, CUBLAS_OP_N, 
           m, n, k, 
           &alpha, 
           d_A, m, 
           d_B, k, 
           &beta, 
           d_C, m);

// 结果取回主机
cublasGetVector(m*n, sizeof(float), d_C, 1, h_C, 1);

// 释放资源
cudaFree(d_A); cudaFree(d_B); cudaFree(d_C);
free(h_A); free(h_B); free(h_C);

关键函数分类

Level 1 函数 (向量运算)

cublas<>axpy：向量加权加法 (y = αx + y)
cublas<>dot：向量点积
cublas<>scal：向量缩放 (x = αx)
cublas<>nrm2：向量范数计算

Level 2 函数 (矩阵-向量运算)

cublas<>gemv：通用矩阵-向量乘法
cublas<>symv：对称矩阵-向量乘法
cublas<>trmv：三角矩阵-向量乘法

Level 3 函数 (矩阵-矩阵运算)

cublas<>gemm：通用矩阵乘法 (核心函数)
cublas<>symm：对称矩阵乘法
cublas<>trmm：三角矩阵乘法
cublas<>syrk：对称秩k更新

高级特性

1. 批处理 GEMM

cpp

// 批处理矩阵乘法 (多个小矩阵同时计算)
cublasSgemmStridedBatched(
    handle,
    CUBLAS_OP_N, CUBLAS_OP_N,
    m, n, k,
    &alpha,
    d_A, m, m*k,
    d_B, k, k*n,
    &beta,
    d_C, m, m*n,
    batchCount);

2. Tensor Core 加速

cpp

// 设置使用Tensor Core
cublasSetMathMode(handle, CUBLAS_TENSOR_OP_MATH);

// 使用FP16输入和FP32累加
cublasGemmEx(
    handle,
    CUBLAS_OP_N, CUBLAS_OP_N,
    m, n, k,
    &alpha,
    d_A, CUDA_R_16F, m,
    d_B, CUDA_R_16F, k,
    &beta,
    d_C, CUDA_R_16F, m,
    CUDA_R_32F, CUBLAS_GEMM_DEFAULT_TENSOR_OP);

3. 异步执行

cpp

cudaStream_t stream;
cudaStreamCreate(&stream);
cublasSetStream(handle, stream);

// 异步GEMM调用
cublasSgemm(handle, ...);

// 可以在此插入其他操作
cudaDeviceSynchronize(); // 等待完成