cublasSgemmBatched使用说明

背景

在GPU进行计算的时候，很多时候都需要利用cublas的API, 常用的API有两个：cublasSgemm 和cublasSgemmBatched, 使用过MKL的可能觉得很熟悉，连参数都是一样的，但是这里有一比较坑的地方是，在mkl的矩阵乘法中我们可以设置使用行优先或者列优先，考虑到很多代码底层都是c/c++写的，所以平时矩阵都是按照行优先来写的，不过mkl是支持列优先的矩阵乘法，但是cublas只支持列优先，也不知道英伟达公司是怎么想的，做成兼容的就那么难？不管怎么样，反正既然别人是制定规则的，我们就必须按别人的游戏规则来玩。

cublasSgemm

cublasSgemm直接参考这里,建议大家必须把这个搞明白，不然下面的batch乘法更加会晕。

cublasSgemmBatched

很多时候我们不是简单的进行两个单独的矩阵乘法，而是将两个集合的矩阵进行相乘，例如下图，我们知道，如果利用之前的API.那么需要做一个循环，根据相关数据显示，cublasSgemmBatched的效果要远好于cublasSgemm，为此我们需要掌握该API的使用使用方式。

下面是该API的接口，接口说明可以在这里查看

cublasStatus_t cublasSgemmBatched(cublasHandle_t handle,
                                  cublasOperation_t transa,
                                  cublasOperation_t transb,
                                  int m, int n, int k,
                                  const float           *alpha,
                                  const float           *Aarray[], int lda,
                                  const float           *Barray[], int ldb,
                                  const float           *beta,
                                  float           *Carray[], int ldc,
                                  int batchCount)

接下来看看实际如何使用：假如我们有矩阵A和B如下，对应矩阵相乘得到C

事例代码如下，自己可以根据代码琢磨一下相关参数设置的原理

#include 
#include 
#include 
#include 

using namespace std;


__global__
void show(float* ptr, int size)
{
        for(int i =0; i<size; i++)
        printf("%f\n", ptr[i]);
}

int main()
{

        float* a = new float[16];
        for(int i=0; i<16; i++) a[i] = 1.0;

        float* b = new float[32];
        for(int i=0; i<32; i++) b[i] = i+1;

        float* c = new float[16];
        for(int i=0; i<16; i++) c[i] = 3.0;

        float* d_a, *d_b, *d_c;
        size_t size = sizeof(float) * 16;
        cudaMalloc(&d_a, size);
        cudaMalloc(&d_b, size*2);
        cudaMalloc(&d_c, size);

        cudaMemcpy(d_a, a, size, cudaMemcpyHostToDevice);
        cudaMemcpy(d_b, b, size*2, cudaMemcpyHostToDevice);
        cudaMemcpy(d_c, c, size, cudaMemcpyHostToDevice);


        cublasHandle_t handle;
        cublasStatus_t ret;
        ret = cublasCreate(&handle);
        float *a_array[8], *b_array[8];
        float *c_array[8];
        for (int i = 0; i < 2; ++i) {
            for (int j = 0; j < 4; ++j) {
                a_array[i*4+j] = d_a + i * 8 + j * 2;
                b_array[i*4+j] = d_b + i * 2 * 8  + j * 2;
                c_array[i*4+j] = d_c + i * 8 + j * 2;
            }
        }
        const float **d_Marray, **d_Narray;
        float **d_Parray;
        cudaMalloc((void**)&d_Marray, 8*sizeof(float *));
        cudaMalloc((void**)&d_Narray, 16*sizeof(float *));
        cudaMalloc((void**)&d_Parray, 8*sizeof(float *));
        cudaMemcpy(d_Marray, a_array, 8*sizeof(float *), cudaMemcpyHostToDevice);
        cudaMemcpy(d_Narray, b_array, 16*sizeof(float *), cudaMemcpyHostToDevice);
        cudaMemcpy(d_Parray, c_array, 8*sizeof(float *), cudaMemcpyHostToDevice);


        const float alpha  =  1.0f;
        const float beta  =  0.0f;
        int m = 2;
        int n = 1;
        int k = 2;
        int lda = 8;
        int ldb = 8;
        int ldc = 8;
        int batch = 8;

       ret = cublasSgemmBatched(handle,
                           CUBLAS_OP_N,
                           CUBLAS_OP_N,
                           m,n,k,
                           &alpha,
                           d_Narray,  ldb,
                           d_Marray,  lda,
                           &beta,
                           d_Parray,  ldc,
                           batch);
        cublasDestroy(handle);
        if (ret == CUBLAS_STATUS_SUCCESS)
        {
        printf("sgemm success  %d, line(%d)\n", ret, __LINE__);
        }

        show<<<1,1>>>(c_array[0], 16);
        cudaMemcpy(c, d_c, size, cudaMemcpyDeviceToHost);
        for(int i=0; i<16; i++) cout<<c[i]<<" "<<endl;


        return 0;
}

建议大家使用cublasGemmBatchedEx，因为很多gpu可是支持fp16, 如果后面想改用fp16的话不用来回折腾代码了，改改配置参数就行了