cuBLAS¶

简介¶

cuBLAS 库可提供基本线性代数子程序 (BLAS) 的 GPU 加速实现。cuBLAS 利用针对 NVIDIA GPU 高度优化的插入式行业标准 BLAS API，加速 AI 和 HPC 应用。cuBLAS 库包含用于批量运算、跨多个 GPU 的执行以及混合精度和低精度执行的扩展程序。通过使用 cuBLAS，应用将能自动从定期性能提升及新的 GPU 体系架构中受益。cuBLAS 库包含在 NVIDIA HPC SDK 和 CUDA 工具包中。

cuBLAS使用说明¶

思源一号上的cuBLAS¶

先创建一个目录cuBLAStest并进入该目录：

mkdir cuBLAStest
cd cuBLAStest

在该目录下创建如下测试文件cuBLAStest.cu：

//矩阵乘法示例
#include <cstdio>
#include <cstdlib>
#include <vector>

#include <cublas_v2.h>
#include <cuda_runtime.h>

template <typename T>
void print_matrix(const int &m, const int &n, const T *A, const int &lda);
template <>
void print_matrix(const int &m, const int &n, const double *A, const int &lda)
{
    for (int i = 0; i < m; i++)
    {
        for (int j = 0; j < n; j++)
        {
            std::printf("%0.2f ", A[j * lda + i]);
        }
        std::printf("\n");
    }
}

using data_type = double;

int main(int argc, char *argv[])
{
    cublasHandle_t cublasH = NULL;
    cudaStream_t stream = NULL;

    const int m = 2;
    const int n = 2;
    const int k = 2;
    const int lda = 2;
    const int ldb = 2;
    const int ldc = 2;
    /*
    *   A = | 1.0 | 2.0 |
    *       | 3.0 | 4.0 |
    *
    *   B = | 5.0 | 6.0 |
    *       | 7.0 | 8.0 |
    */

    const std::vector<data_type> A = {1.0, 2.0, 3.0, 4.0};
    const std::vector<data_type> B = {5.0, 6.0, 7.0, 8.0};
    std::vector<data_type> C(m * n);
    const data_type alpha = 1.0;
    const data_type beta = 0.0;

    data_type *d_A = nullptr;
    data_type *d_B = nullptr;
    data_type *d_C = nullptr;

    cublasOperation_t transa = CUBLAS_OP_N;
    cublasOperation_t transb = CUBLAS_OP_N;

    printf("A\n");
    print_matrix(m, k, A.data(), lda);
    printf("=====\n");

    printf("B\n");
    print_matrix(k, n, B.data(), ldb);
    printf("=====\n");

    /* step 1: create cublas handle, bind a stream */
    (cublasCreate(&cublasH));

    (cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking));
    (cublasSetStream(cublasH, stream));

    /* step 2: copy data to device */
    (cudaMalloc(reinterpret_cast<void **>(&d_A), sizeof(data_type) * A.size()));
    (cudaMalloc(reinterpret_cast<void **>(&d_B), sizeof(data_type) * B.size()));
    (cudaMalloc(reinterpret_cast<void **>(&d_C), sizeof(data_type) * C.size()));

    (cudaMemcpyAsync(d_A, A.data(), sizeof(data_type) * A.size(), cudaMemcpyHostToDevice, stream));
    (cudaMemcpyAsync(d_B, B.data(), sizeof(data_type) * B.size(), cudaMemcpyHostToDevice, stream));

    /* step 3: compute */
    (cublasDgemm(cublasH, transa, transb, m, n, k, &alpha, d_A, lda, d_B, ldb, &beta, d_C, ldc));

    /* step 4: copy data to host */
    (cudaMemcpyAsync(C.data(), d_C, sizeof(data_type) * C.size(), cudaMemcpyDeviceToHost, stream));

    (cudaStreamSynchronize(stream));

    /*
    *   C = | 23.0 | 31.0 |
    *       | 34.0 | 46.0 |
    */

    printf("C\n");
    print_matrix(m, n, C.data(), ldc);
    printf("=====\n");

    /* free resources */
    (cudaFree(d_A));
    (cudaFree(d_B));
    (cudaFree(d_C));

    (cublasDestroy(cublasH));
    (cudaStreamDestroy(stream));
    (cudaDeviceReset());

    return EXIT_SUCCESS;
}

在该目录下创建如下作业提交脚本cuBLAStest.slurm:

#!/bin/bash

#SBATCH --job-name=cuBLAStest        # 作业名
#SBATCH --partition=a100             # a100 队列
#SBATCH -N 1
#SBATCH --ntasks-per-node=1
#SBATCH --cpus-per-task=1            # 1:1 的 GPU:CPU 配比
#SBATCH --gres=gpu:1                 # 1 块 GPU
#SBATCH --output=%j.out
#SBATCH --error=%j.err

module load cuda/11.5.0
module load gcc/11.2.0

nvcc cuBLAStest.cu -o cuBLAStest -lcublas
./cuBLAStest

使用如下命令提交作业：

sbatch cuBLAStest.slurm

作业完成后在.out文件中可看到如下结果：

A
1.00 3.00
2.00 4.00
=====
B
5.00 7.00
6.00 8.00
=====
C
23.00 31.00
34.00 46.00
=====

pi2.0上的cuBLAS¶

此步骤和上文完全相同；
此步骤和上文完全相同；
在该目录下创建如下作业提交脚本cuBLAStest.slurm:

#!/bin/bash

#SBATCH --job-name=cuBLAStest        # 作业名
#SBATCH --partition=dgx2             # dgx2 队列
#SBATCH -N 1
#SBATCH --ntasks-per-node=1
#SBATCH --cpus-per-task=1            # 1:1 的 GPU:CPU 配比
#SBATCH --gres=gpu:1                 # 1 块 GPU
#SBATCH --output=%j.out
#SBATCH --error=%j.err

module load cuda/11.6.2-gcc-8.3.0
module load gcc/8.3.0

nvcc cuBLAStest.cu -o cuBLAStest -lcublas
./cuBLAStest

使用如下命令提交作业：

sbatch cuBLAStest.slurm

作业完成后在.out文件中可看到如下结果：

A
1.00 3.00
2.00 4.00
=====
B
5.00 7.00
6.00 8.00
=====
C
23.00 31.00
34.00 46.00
=====

参考资料¶

最后更新: 2024 年 11 月 14 日