cublas 中 gemm 的用法

最近在看 llama2 的推理过程的源码,其中的矩阵乘法操作用到了 cublas 库。对其中的参数有较大的疑惑,所以打算梳理一下。

有一个很好的解释 cublas 中 gemm 的文章,链接如下有关CUBLAS中的矩阵乘法函数 - 爨爨爨好 - 博客园 (cnblogs.com)

具体的详细操作,请看给出的文章,我只是对其进行一个总结。

cublas 中矩阵乘法主要就是 14 个参数,该函数实际上是用于计算 C=\alphaAB+\betaC, 其中 A 大小默认为 m×k,B 大小默认为 k×n,C 大小默认为 m×n,\alpha 和 \beta 为标量。

cublasHandle_t handle  为调用 cuBLAS 库时的句柄

cublasOperation_t transa 是否转置矩阵 A

cublasOperation_t transb 是否转置矩阵 B

int m 矩阵A的行数,等于矩阵C的列数。

int n 矩阵B的列数,等于矩阵C的行数。

int k 矩阵A的列数,等于矩阵B的行数。

m, n, k 这三个字母的作用就是确定输入矩阵 A, B, C 中元素的个数(因为float * 这种类型中并未蕴含元素个数的信息)。因为要运行的函数是矩阵乘法,那么只需要这三个数字就能覆盖A,B,C中元素的个数。

const float *alpha 标量 α 的指针,可以是主机指针或设备指针,只需要计算矩阵乘法时命 α = 1.0f

const float *A  矩阵(数组)A 的指针,必须是设备指针

const float *B 矩阵(数组)B 的指针,必须是设备指针

float *C 矩阵(数组)C 的指针,必须是设备指针

const float *beta 标量 β 的指针,可以是主机指针或设备指针,只需要计算矩阵乘法时命 β = 0.0f

int lda 矩阵 A 的主维(leading dimension)(lda 要大于等于 m)

int ldb 矩阵 B 的主维 (ldb 要等于等于 k)

int ldc 矩阵 C 的主维 (ldc 要大于等于 m)

重点说下这个主维的作用,cublas 中是矩阵是以列优先的,而C++中是以行优先的,MATLAB 也是以列优先。

cublas 这个函数中有一个自己的 A1, B1, C1。输入的 A, B 都换先转换到 A1, B1, 然后 A1和B1做矩阵乘法得到C1,然后再将C1转换到 C中。

以C++进行举例,

大小为 m×k 的矩阵 A(行主序)会首先转换到 lda × k 的矩阵 A1,用0补充可能多余的空位。

大小为 k×n 的矩阵 B(行主序)会首先转换到 ldb × n 的矩阵 B1,用0补充可能多余的空位。

然后用 A1 乘以部分 B1 (前k行) 得到的大小为 lda×n C1(前ldc行)。

然后将列主序的C1转换到行主序的大小为 m×n的 C。

### 使用 cuBLAS 实现矩阵相乘 为了使用 cuBLAS 进行高效的矩阵相乘操作,需要初始化 CUDA 和 cuBLAS 库,并调用相应的 API 函数完成计算。下面是一个完整的例子说明如何通过 cuBLAS 完成两个浮点数矩阵之间的乘法。 #### 初始化环境与资源分配 在执行任何 cuBLAS 调用之前,先要创建一个上下文句柄用于后续的操作: ```cpp #include <cuda_runtime.h> #include <cublas_v2.h> // 错误处理宏定义简化错误检测过程 #define cudaErrorCheck(call) \ do { \ cudaError_t err = call; \ if (err != cudaSuccess) { \ fprintf(stderr, "CUDA Error: %s\n", cudaGetErrorString(err)); \ exit(EXIT_FAILURE); \ } \ } while(0) int main() { cublasHandle_t handle; cublasStatus_t status; // 创建 cuBLAS 句柄 status = cublasCreate(&handle); if (status != CUBLAS_STATUS_SUCCESS){ printf("Failed to create cuBLAS handle.\n"); return -1; } } ``` 此部分代码负责设置必要的运行时环境以及异常情况下的报错机制[^1]。 #### 数据准备 接着准备好待运算的数据集并将它们传输到 GPU 设备上: ```cpp const int m = 2; // 行数 const int n = 3; // 列数 const int k = 4; // 中间维度大小 float h_A[m*k] = {...}; // 主机端输入矩阵A数据填充... float h_B[k*n] = {...}; // 主机端输入矩阵B数据填充... float *d_A, *d_B, *d_C; size_t size_A = m * k * sizeof(float); size_t size_B = k * n * sizeof(float); size_t size_C = m * n * sizeof(float); // 分配设备内存空间给三个矩阵变量 cudaErrorCheck(cudaMalloc((void**)&d_A, size_A)); cudaErrorCheck(cudaMalloc((void**)&d_B, size_B)); cudaErrorCheck(cudaMalloc((void**)&d_C, size_C)); // 将主机上的初始值复制到对应的设备指针位置处 cudaErrorCheck(cudaMemcpy(d_A, h_A, size_A, cudaMemcpyHostToDevice)); cudaErrorCheck(cudaMemcpy(d_B, h_B, size_B, cudaMemcpyHostToDevice)); ``` 这里假设 `h_A` 和 `h_B` 已经被适当赋值为所需的测试数值。 #### 执行矩阵乘法 现在可以调用 cuBLAS 提供的 gemm 接口来进行实际的矩阵乘法运算了: ```cpp // 设置 alpha=1.0f,beta=0.0f 参数控制线性组合系数 const float alf = 1.f; const float bet = 0.f; // 调用 cuBLASGEMM 方法实现 C := α*A*B + β*C status = cublasSgemm(handle, CUBLAS_OP_N /* op(A)=A */, CUBLAS_OP_N /* op(B)=B */, n /* Number of rows in matrix op(B) and C*/, m /* Number of columns in matrix op(A) and C*/, k /* Number of columns/rows in matrices A/B */, &alf, d_B, n /* Leading dimension of B */, d_A, k /* Leading dimension of A */, &bet, d_C, n /* Leading dimension of C */ ); if(status != CUBLAS_STATUS_SUCCESS){ printf("GEMM failed!\n"); return -1; } // 把结果从设备拷贝回主机侧以便验证正确性 float h_C[m*n]; cudaErrorCheck(cudaMemcpy(h_C, d_C, size_C, cudaMemcpyDeviceToHost)); ``` 这段程序片段展示了怎样配置参数并通过一次函数调用来高效地完成大规模矩阵间的乘积运算。 #### 清理工作 最后不要忘记释放所占用的所有硬件资源: ```cpp // 解除对 cuBLAS 上下文对象的引用并销毁它 cublasDestroy(handle); // 释放先前申请过的显存区域 cudaFree(d_A); cudaFree(d_B); cudaFree(d_C); ``` 上述流程构成了基于 NVIDIA cuBLAS 库的一个典型应用案例,能够显著提升涉及大量密集型代数变换任务的工作效率。
评论 1
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值