当前位置：首页 > news >正文

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

news 来源：原创 2025/5/23 6:04:26

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(); // 等待完成