victory的博客

长安一片月，万户捣衣声

0%

GPU | FMA

发表于 分类于 ， ， ，

什么是FMA？

GPU能够加速AI模型训练和推理速度的原因是GPU拥有众多的CUDA core、Tensor Core。

AI模型的核心计算是矩阵乘加运算，Tensor Core实现了不同浮点精度的矩阵乘加运算（FMA），从而加速了AI模型的训练、推理过程。

下面是使用FMA进行矩阵乘加运算D = A * B + C的CUDA示例代码：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
#include 
#include 

// 假设矩阵A, B, C和D都是NxN大小的方阵
#define N 1024

__global__ void matrixFMA(float *A, float *B, float *C, float *D, int n) {
    int row = blockIdx.y * blockDim.y + threadIdx.y;
    int col = blockIdx.x * blockDim.x + threadIdx.x;

    if (row < n && col < n) {
        float sum = 0.0f;
        for (int k = 0; k < n; ++k) {
            // 使用FMA操作计算每个元素
            sum = __fmaf_rn(A[row * n + k], B[k * n + col], sum);
        }
        D[row * n + col] = sum + C[row * n + col];
    }
}

int main() {
    float *A, *B, *C, *D;
    float *d_A, *d_B, *d_C, *d_D;

    // 分配主机内存
    A = new float[N * N];
    B = new float[N * N];
    C = new float[N * N];
    D = new float[N * N];

    // 初始化输入数据
    for (int i = 0; i < N * N; i++) {
        A[i] = static_cast(i);
        B[i] = static_cast(2 * i);
        C[i] = static_cast(3 * i);
    }

    // 分配设备内存
    cudaMalloc(&d_A, N * N * sizeof(float));
    cudaMalloc(&d_B, N * N * sizeof(float));
    cudaMalloc(&d_C, N * N * sizeof(float));
    cudaMalloc(&d_D, N * N * sizeof(float));

    // 将输入数据复制到设备
    cudaMemcpy(d_A, A, N * N * sizeof(float), cudaMemcpyHostToDevice);
    cudaMemcpy(d_B, B, N * N * sizeof(float), cudaMemcpyHostToDevice);
    cudaMemcpy(d_C, C, N * N * sizeof(float), cudaMemcpyHostToDevice);

    // 配置并启动内核
    int threadsPerBlock = 16;
    int blocksPerGrid = (N + threadsPerBlock - 1) / threadsPerBlock;
    dim3 dimGrid(blocksPerGrid, blocksPerGrid, 1);
    dim3 dimBlock(threadsPerBlock, threadsPerBlock, 1);
    matrixFMA<<>>(d_A, d_B, d_C, d_D, N);

    // 将结果从设备复制回主机
    cudaMemcpy(D, d_D, N * N * sizeof(float), cudaMemcpyDeviceToHost);

    // 验证结果（可选）
    // ...

    // 释放内存
    cudaFree(d_A);
    cudaFree(d_B);
    cudaFree(d_C);
    cudaFree(d_D);
    delete[] A;
    delete[] B;
    delete[] C;
    delete[] D;

    return 0;
}

  • FMA与Tensor Core

    Tensor Core 采用融合乘法加法（FMA）的方式来高效地处理计算任务。每个 Tensor Core 每周期能执行 4x4x4 GEMM，64 个浮点乘法累加（FMA）运算。