cublasSgemm Случайное использование с jcuda

Question

Я пытался использовать функцию cublasSgemmBatched() в jcuda для умножения матриц, и я не уверен, как правильно обрабатывать передачу указателей и векторы пакетных матриц. Я буду очень благодарен, если кто-то знает, как изменить мой код для правильного решения этой проблемы. В этом примере массив C остается неизменным после cublasGetVector.

public static void SsmmBatchJCublas(int m, int n, int k, float A[], float B[]){

    // Create a CUBLAS handle
    cublasHandle handle = new cublasHandle();
    cublasCreate(handle);

    // Allocate memory on the device
    Pointer d_A = new Pointer();
    Pointer d_B = new Pointer();
    Pointer d_C = new Pointer();


    cudaMalloc(d_A, m*k * Sizeof.FLOAT);
    cudaMalloc(d_B, n*k * Sizeof.FLOAT);
    cudaMalloc(d_C, m*n * Sizeof.FLOAT);

    float[] C = new float[m*n];
    // Copy the memory from the host to the device
    cublasSetVector(m*k, Sizeof.FLOAT, Pointer.to(A), 1, d_A, 1);
    cublasSetVector(n*k, Sizeof.FLOAT, Pointer.to(B), 1, d_B, 1);
    cublasSetVector(m*n, Sizeof.FLOAT, Pointer.to(C), 1, d_C, 1);

    Pointer[] Aarray = new Pointer[]{d_A};
    Pointer AarrayPtr = Pointer.to(Aarray);
    Pointer[] Barray = new Pointer[]{d_B};
    Pointer BarrayPtr = Pointer.to(Barray);
    Pointer[] Carray = new Pointer[]{d_C};
    Pointer CarrayPtr = Pointer.to(Carray);

    // Execute sgemm
    Pointer pAlpha = Pointer.to(new float[]{1});
    Pointer pBeta = Pointer.to(new float[]{0});


    cublasSgemmBatched(handle, CUBLAS_OP_N, CUBLAS_OP_N, m, n, k, pAlpha, AarrayPtr, Aarray.length, BarrayPtr, Barray.length, pBeta, CarrayPtr, Carray.length, Aarray.length);
    // Copy the result from the device to the host
    cublasGetVector(m*n, Sizeof.FLOAT, d_C, 1, Pointer.to(C), 1);

    // Clean up
    cudaFree(d_A);
    cudaFree(d_B);
    cudaFree(d_C);
    cublasDestroy(handle);
}

Answer 1

Я спросил на официальном форуме jcuda и быстро получил ответ здесь.