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 <cuda_runtime.h>
#include <cublas_v2.h>
#include <iostream>
#include <stdio.h>

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的話不用來回折騰代碼了，改改配置參數就行了