【CUDA】CUBLAS

【CUDA】CUBLAS

在深入了解之前，提前运行预热（warmup）和基准测试（benchmark runs） 是获得准确执行时间的关键。如果不进行预热运行，cuBLAS 的首次运行会有较大的开销（大约 45 毫秒），会导致结果偏差。基准测试能够更准确地计算出平均执行时间。

---

cuBLAS：CUDA 基本线性代数子程序库

• 简介：cuBLAS 是 NVIDIA 的 GPU 加速线性代数运算库，广泛应用于 人工智能（AI） 和 高性能计算（HPC） 领域。
• 功能：提供了行业标准的 BLAS 和 GEMM（矩阵乘法）API，并支持高度优化的融合操作（fusion），充分发挥 NVIDIA GPU 的性能。
• 使用提示：在使用 GEMM 操作时，需要特别注意矩阵的存储布局（行优先或列优先），例如参考这里。

---

cuBLASLt：轻量级扩展

• 简介：cuBLASLt 是 cuBLAS 的轻量级扩展，提供了更灵活的 API，专注于提升特定工作负载（如深度学习模型）的性能。
• 特点：
  • 如果单个内核无法处理问题，cuBLASLt 会将问题分解为多个子问题，并在每个子问题上运行内核。
  • 支持混合精度计算，如 fp16、fp8 和 int8，可显著提升深度学习的推理速度。

---

cuBLASXt：支持多 GPU 扩展

• 简介：cuBLASXt 是 cuBLAS 的扩展版，主要针对超大规模计算，支持多 GPU 运行。
• 特点：
  • 多 GPU 支持：能够将 BLAS 操作分布到多块 GPU 上，适合处理需要扩展 GPU 计算的大型数据集。
  • 线程安全：支持多线程并发执行，在不同的 GPU 上同时运行多个 BLAS 操作。
  • 适用场景：特别适用于超出单块 GPU 显存限制的大规模线性代数问题。
• 缺点：由于需要在 主板 DRAM 和 GPU VRAM 之间频繁传输数据，会造成内存带宽瓶颈，导致计算速度较慢。

---

cuBLASDx：设备端扩展（未在课程中使用）

• 简介：cuBLASDx 是一个设备端 API 扩展，用于直接在 CUDA 内核中执行 BLAS 计算。
• 特点：
  • 通过融合（fusion）数值操作，进一步降低延迟，提升应用程序性能。
  • 注意：cuBLASDx 并不包含在 CUDA Toolkit 中，需要单独下载。

---

CUTLASS：CUDA 模板线性代数子程序库

• 简介：cuBLAS 及其变体主要在 主机（host）端 运行，而 CUTLASS 提供了模板库，允许开发者实现高度自定义和优化的线性代数操作。
• 特点：
  • 支持在 CUDA 中轻松融合矩阵运算。
  • 对于深度学习，矩阵乘法是最重要的操作，而 cuBLAS 无法轻松实现复杂操作的融合。
• 补充说明：
  • CUTLASS 并未用于实现 Flash Attention，后者是通过高度优化的 CUDA 内核实现的（详见论文）。

---

总结与应用场景

库	特点与适用场景
cuBLAS	高性能线性代数运算，适用于 AI 和 HPC。
cuBLASLt	灵活 API 和混合精度支持，专注深度学习工作负载。
cuBLASXt	多 GPU 支持，适合超大规模计算，但受限于内存带宽瓶颈。
cuBLASDx	设备端 API，可在 CUDA 内核中实现 BLAS 操作，进一步优化延迟。
CUTLASS	提供模板库，支持复杂运算融合，深度学习中高效矩阵运算的选择。

通过根据工作负载的需求选择合适的库，可以充分利用 NVIDIA GPU 的计算能力。

cuBLAS示例

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

#define M 3
#define K 4
#define N 2

#define CHECK_CUDA(call)                                                                               \
    {
                                                                                                       \
        cudaError_t err = call;                                                                        \
        if (err != cudaSuccess) {
                                                                           \
            fprintf(stderr, "CUDA error in %s:%d: %s\n", __FILE__, __LINE__, cudaGetErrorString(err)); \
            exit(EXIT_FAILURE);                                                                        \
        }                                                                                              \
    }

#define CHECK_CUBLAS(call)                                                              \
    {
                                                                                        \
        cublasStatus_t status = call;                                                   \
        if (status != CUBLAS_STATUS_SUCCESS) {
                                               \
            fprintf(stderr, "cuBLAS error in %s:%d: %d\n", __FILE__, __LINE__, status); \
            exit(EXIT_FAILURE);                                                         \
        }                                                                               \
    }

#undef PRINT_MATRIX
#define PRINT_MATRIX(mat, rows, cols)                                       \
    for (int i = 0; i < rows; i++) {
                                             \
        for (int j = 0; j < cols; j++) printf("%8.3f ", mat[i * cols + j]); \
        printf("\n");                                                       \
    }                                                                       \
    printf("\n");

void CpuMatmul(float* A, float* B, float* C) {
   
    for (int i = 0; i < M; ++i) {
   
        for (int j = 0; j < N; ++j) {
   
            float sum = 0;
            for (int k = 0; k < K; ++k) {
   
                sum += A[i * K + k] * B[k * N + j];
            }
            C[i * N + j] = sum;
        }
    }
}

int main() {
   
    float A[M * K] = {
   1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f, 9.0f, 10.0f, 11.0f, 12.0f};
    float B[K * N] = {
   1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f};
    float C_cpu[M * N], C_cublas_s[M * N], C_cublas_h[M * N];

    // CPU matmul
    CpuMatmul(A, B, C_cpu);

    cublasHandle_t handle;
    CHECK_CUBLAS(cublasCreate(&handle));

    float *d_A, *d_B, *d_C;
    CHECK_CUDA(cudaMalloc(&d_A, M * K * sizeof(float)));
    CHECK_CUDA(cudaMalloc(&d_B, K * N * sizeof(float)));
    CHECK_CUDA(cudaMalloc(&d_C, M * N * sizeof(float)));

    CHECK_CUDA(cudaMemcpy(d_A, A, M * K * sizeof(float), cudaMemcpyHostToDevice));
    CHECK_CUDA(cudaMemcpy(d_B, B, K * N * sizeof(float), cudaMemcpyHostToDevice));

    // cuBLAS SGEMM(单精度)
    float alpha = 1.0f, beta = 0.0f;
    /*cublasSgemm ：
    矩阵乘法公式 C = alpha x op(A) @ op(B) + beta x C
    cublasStatus_t cublasSgemm(
        cublasHandle_t handle,      // cuBLAS上下文
        // CUBLAS_OP_N：不转置（默认）。 CUBLAS_OP_T：转置。CUBLAS_OP_C：共轭转置。
        cublasOperation_t transa,  //指定矩阵A是否转置,
        cublasOperation_t transb,  // 指定矩阵B是否转置
        int m,                     // 矩阵C的行数
        int n,                     // 矩阵C的列数
        int k,                     // A的列数或转置后行数；B的行数或转置后列数
        const float *alpha,        // 标量alpha
        const float *A,            // 矩阵A
        // 主维度（leading dimension），表示存储矩阵时每列或每行的跨度,如果矩阵是列主序，则 lda 是矩阵 A 的行数。
        int lda,                   // A的主维（leading dimension）
        const float *B,            // 矩阵B
        int ldb,                   // B的主维
        const float *beta,         // 标量beta
        float *C,                  // 输出矩阵C
        int ldc                    // C的主维
    );
    */
    // cuBLAS 默认使用列主序存储矩阵。如果输入矩阵是行主序存储，则需要手动调整主维度，或使用技巧重新解释。
    /*
    通过技巧将问题转换为计算 (B^T * A^T)^T，从而直接获得期望结果：
    矩阵 A 和 B 在内存中以行主序存储，等价于将其转置解释为列主序存储。
    调用 cublasSgemm 时，将矩阵按未转置（CUBLAS_OP_N）处理，令 cuBLAS 按列主序解释输入。
    调整矩阵的传递顺序和参数，实际计算 (B^T * A^T)^T，最终结果矩阵在内存中直接符合行主序的期望。
    */
    CHECK_CUBLAS(cublasSgemm(handle, CUBLAS_OP_N, CUBLAS_OP_N, N, M, K, &alpha, d_B, N, d_A, K, &beta, d_C, N));
    CHECK_CUDA(cudaMemcpy(C_cublas_s, d_C, M * N * sizeof(float), cudaMemcpyDeviceToHost));

    // cuBLAS HGEMM
    half *d_A_h, *d_B_h, *d_C_h;
    CHECK_CUDA(cudaMalloc(&d_A_h, M * K * sizeof(half)));
    CHECK_CUDA(cudaMalloc(&d_B_h, K * N * sizeof(half)));
    CHECK_CUDA(cudaMalloc(&d_C_h, M * N * sizeof(half)));

    // Convert to half percision to CPU
    half A_h[M * K], B_h[K * N];
    for (int i = 0; i < M * K; ++i) {
   
        A_h[i] = __float2half(A[i]);
    }
    for (int i = 0; i < K * N; ++i) {
   
        B_h[i] = __float2half(B[i]);
    }

    CHECK_CUDA(cudaMemcpy(d_A_h, A_h, M * K * sizeof(half), cudaMemcpyHostToDevice));
    CHECK_CUDA(cudaMemcpy(d_B_h, B_h, K