GEMM 通用矩阵乘

CUDA9.0中GEMM接口不支持数据按列存储(即横向排布),当只有C横向排布时会报第15个值错误,当A、B横向排布时,虽然不会报错,但后续操作会访存越界。

接口如下

/** cublasStatus_t cublasSgemmStridedBatched(cublasHandle_t handle,
 *                                           cublasOperation_t transa, cublasOperation_t transb,
 *                                           int m, int n, int k,
 *                                           const float *alpha, 
 *                                           const float *A, int lda, long long int strideA,
 *                                           const float *B, int ldb, long long int strideB,
 *                                           const float *beta,
 *                                           float *C, int ldc, long long int strideC,
 *                                           int batchCount)
 */

其中,A,B,C内部进行运算的小矩阵分别是MxK, KxN, MxN大小,TRANSA, TRANSB表示是否使用对应矩阵的转置,ALPHA, BETA为对应的系数。而LDA, LDB, LDC表示对应矩阵的leading dimension,即第一维度的大小。LDA表示一个batch中一个矩阵行的长度, 因为矩阵在内存中是连续存放的,而这个leading dimension的量用来定义元素换行后的位置,即A[i, j] = A + i*lda + j

C++矩阵默认行优先存储,BLAS库默认列优先存储,所以两个矩阵要反过来输入(无需指定trans=CUBLAS_OP_T),transa就是对应于第一个矩阵是否转置。不过如果反着输入,且两个trans都为CUBLAS_OP_N, 则我们令MxK为第一个输入中小矩阵形状, KxN为第二个输入中小矩阵形状, 然后将N,M,K对应于接口的m,n,k进行输入。如果其中trans为T,则m,n,k用来指定其小矩阵转置后的形状,表示实际参与运算的矩阵大小,对应于逻辑上的AxB(而不是反着输入的B和A),但trans的顺序是对应反着输入的两个矩阵的。

// 以BERT中Q*K'为例
// 对于K, 大矩阵是lda * m(ldQKV * S), 小矩阵转置后是k * m(headSize * S), 转置只是转置小矩阵
// 对于Q, 大矩阵是ldb * n(ldQKV * S), 小矩阵转置后是k * n(headSize x S)
CUBLAS_CHECK(cublasGemmStridedBatched(cublas, CUBLAS_OP_T, CUBLAS_OP_N, S, S, headSize, 1.f, kptr, ldQKV, strideQKV, 
             qptr, ldQKV, strideQKV, 0.f, qkptr, S, omatSize, numMats));

参考https://blog.csdn.net/feng__shuai/article/details/107091684

你可能感兴趣的:(GEMM 通用矩阵乘)