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<T>(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

最后编辑于：2020.09.03 10:48:56