#include "cuda_runtime.h"
#include "cublas_v2.h"
#include <time.h>
#include <iostream>

using namespace std;


// cuBLAS實(shí)現(xiàn)矩陣乘法
int **matMult_cuBLAS(int **A, int **B, int rowSizeA, int colSizeA, int colSizeB, cublasHandle_t cuHandle){
    // 結(jié)果矩陣
    int** C = new int*[rowSizeA];
    for(int i = 0; i < rowSizeA; i++){
        C[i] = new int[colSizeB];
    }
    for (int i = 0; i < rowSizeA; i++){
        for (int j = 0; j < colSizeB; j++){
            C[i][j] = 0;
        }
    }

    // 在內(nèi)存中為將要計(jì)算的矩陣開辟空間
    float *h_A = (float*)malloc (rowSizeA * colSizeA * sizeof(float));
    float *h_B = (float*)malloc (colSizeA * colSizeB * sizeof(float));
    float *h_C = (float*)malloc (rowSizeA * colSizeB * sizeof(float));

    // 初始化計(jì)算矩陣h_A和h_B
    for (int i = 0; i < rowSizeA; i++) {
        for (int j = 0; j < colSizeA; j++) {
            h_A[i * colSizeA + j] = (float)A[i][j];
        }
    }
    for (int i = 0; i < colSizeA; i++) {
        for (int j = 0; j < colSizeB; j++) {
            h_B[i * colSizeB + j] = (float)B[i][j];
        }
    }

    // 在顯存中為將要計(jì)算矩陣與結(jié)果矩陣開辟空間
    float *d_A, *d_B, *d_C;
    cudaMalloc (
        (void**)&d_A,    // 指向開辟的空間的指針
        rowSizeA * colSizeA * sizeof(float)    //　需要開辟空間的字節(jié)數(shù)
    );
    cudaMalloc (
        (void**)&d_B,    
        colSizeA * colSizeB * sizeof(float)    
    );
    cudaMalloc (
        (void**)&d_C,
        rowSizeA * colSizeB * sizeof(float)    
    );

    // 將矩陣數(shù)據(jù)傳遞進(jìn)顯存中已經(jīng)開辟好了的空間
    cublasSetVector (
        rowSizeA * colSizeA,    // 要存入顯存的元素個數(shù)
        sizeof(float),    // 每個元素大小
        h_A,    // 主機(jī)端起始地址
        1,    // 連續(xù)元素之間的存儲間隔
        d_A,    // GPU 端起始地址
        1    // 連續(xù)元素之間的存儲間隔
    );
    cublasSetVector (colSizeA * colSizeB, sizeof(float), h_B, 1, d_B, 1);

    // 傳遞進(jìn)矩陣相乘函數(shù)中的參數(shù)，具體含義請參考函數(shù)手冊．
    float a=1; float b=0;
    // 矩陣相乘．該函數(shù)必然將數(shù)組解析成列優(yōu)先數(shù)組
    cublasSgemm (
        cuHandle,    // blas 庫對象 
        CUBLAS_OP_T,    // 矩陣 A 屬性參數(shù)
        CUBLAS_OP_T,    // 矩陣 B 屬性參數(shù)
        rowSizeA,    // A, C 的行數(shù) 
        colSizeB,    // B, C 的列數(shù)
        colSizeA,    // A 的列數(shù)和 B 的行數(shù)
        &a,    // 運(yùn)算式的 \alpha 值
        d_A,    // A 在顯存中的地址
        colSizeA,    // lda
        d_B,    // B 在顯存中的地址
        colSizeB,    // ldb
        &b,    // 運(yùn)算式的 \beta 值
        d_C,    // C 在顯存中的地址(結(jié)果矩陣)
        rowSizeA    // ldc
    );
    
    // 從 顯存 中取出運(yùn)算結(jié)果至 內(nèi)存中去
    cublasGetVector (
        rowSizeA * colSizeB,    //  要取出元素的個數(shù)
        sizeof(float),    // 每個元素大小
        d_C,    // GPU 端起始地址
        1,    // 連續(xù)元素之間的存儲間隔
        h_C,    // 主機(jī)端起始地址
        1    // 連續(xù)元素之間的存儲間隔
    );

    for (int i = 0; i < rowSizeA; i++) {
        for (int j = 0; j < colSizeB; j++) {
            C[i][j] = (int)h_C[j * rowSizeA + i];
        }
    }
    
    // 清理掉使用過的內(nèi)存
    free (h_A); free (h_B); free (h_C); cudaFree (d_A);
    cudaFree (d_B); cudaFree (d_C);

    return C;
}

// 構(gòu)造一個隨機(jī)二維數(shù)組（矩陣）
int** uniformMat(int rowSize, int colSize, int minValue, int maxValue) {
    int** mat = new int* [rowSize];
    for (int i = 0; i < rowSize; i++)
        mat[i] = new int[colSize];


    // srand(1024);
    srand((unsigned)time(NULL));  //隨機(jī)數(shù)種子采用系統(tǒng)時鐘
    for (int i = 0; i < rowSize; i++) {
        for (int j = 0; j < colSize; j++) {
            mat[i][j] = (int)(rand() % (maxValue - minValue + 1)) + minValue;
        }
    }

    return mat;
}

int main(void) 
{   
    // 創(chuàng)建并初始化 CUBLAS 庫對象
    // 若是CUBLAS對象在主函數(shù)中初始化，cuBLAS方法在其他函數(shù)中調(diào)用，需要將cuHandle傳入該函數(shù)，并在該函數(shù)內(nèi)創(chuàng)建status對象
    cublasHandle_t cuHandle;
    cublasStatus_t status = cublasCreate(&cuHandle);
    if (status != CUBLAS_STATUS_SUCCESS)
    {
        if (status == CUBLAS_STATUS_NOT_INITIALIZED) {
            cout << "CUBLAS 對象實(shí)例化出錯" << endl;
        }
        getchar ();
        return EXIT_FAILURE;
    }

    // 矩陣大小定義
    int rowSizeA = 3; // 矩陣A的行數(shù)
    int colSizeA = 4; // 矩陣A的列數(shù)和矩陣B的行數(shù)
    int colSizeB = 2; // 矩陣B的列數(shù)

    // 構(gòu)造一個3行4列的矩陣A，矩陣元素在(0,4)內(nèi)隨機(jī)選取
    int **A = uniformMat(rowSizeA, colSizeA, 0, 4);
    // 構(gòu)造一個4行2列的矩陣B，矩陣元素在(5,9)內(nèi)隨機(jī)選取
    int **B = uniformMat(colSizeA, colSizeB, 5, 9);

    // 輸出矩陣A和B
    cout << "矩陣 A :" << endl;
    for (int i = 0; i < rowSizeA; i++) {
        for (int j = 0; j < colSizeA; j++) {
            cout << A[i][j] << " ";
        }
        cout << endl;
    }
    cout << endl;

    cout << "矩陣 B :" << endl;
    for (int i = 0; i < colSizeA; i++) {
        for (int j = 0; j < colSizeB; j++) {
            cout << B[i][j] << " ";
        }
        cout << endl;
    }
    cout << endl;

    // 使用cuBLAS進(jìn)行矩陣乘法運(yùn)算：C = A * B
    int **C = matMult_cuBLAS(A, B, rowSizeA, colSizeA, colSizeB, cuHandle);

    // 輸出矩陣C，即運(yùn)算結(jié)果
    cout << "矩陣 C :" << endl;
    for (int i = 0; i < rowSizeA; i++) {
        for (int j = 0; j < colSizeB; j++) {
            cout << C[i][j] << " ";
        }
        cout << endl;
    }
    cout << endl;

    // 釋放 CUBLAS 庫對象
    cublasDestroy (cuHandle);
    return 0;
}