nvidia CUDA 高级编程使用cub库优化分布式计算下的原子操作

Posted 2023-01-22 从善若水

博主未授权任何人或组织机构转载博主任何原创文章，感谢各位对原创的支持！
博主链接

本人就职于国际知名终端厂商，负责modem芯片研发。
在5G早期负责终端数据业务层、核心网相关的开发工作，目前牵头6G算力网络技术标准研究。


博客内容主要围绕：
       5G/6G协议讲解
       算力网络讲解（云计算，边缘计算，端计算）
       高级C语言讲解
       Rust语言讲解

使用cub库优化分布式计算下的原子操作

       我们以 Jacobi 迭代的拉普拉斯方程求解器为例进行讲解。下面简单介绍以下算法，我们不需要深入了解这个算法的含义。

1. 拉普拉斯方程

       有限元/有限体积/有限差分应用中的一个常见主题是使用松弛方法求解椭圆偏微分方程。 也许最简单的椭圆偏微分方程是拉普拉斯方程：${▽}^{2}f=0$。其中， ${▽}^2=▽\ast ▽$ 是拉普拉斯算子（所有坐标方向的二阶导数之和）， 𝑓=𝑓(𝐫) 为标量场，是空间矢量坐标 𝐫 的函数。例如，拉普拉斯方程可以用来求解边缘被加热到固定温度的金属板上的温度平衡分布。

       在一维的 𝑓=𝑓(𝑥) 情况下，此方程为：$\frac{{\alpha }^2}{\alpha x^2}f=0$。

       假设我们想在给定固定边界条件 𝑓(0)=𝑇left 和 𝑓(𝐿)=𝑇right 下，在域 𝑥=[0,𝐿] 上求解这个方程。也就是说，我们想知道作为 𝑥 的函数，温度在 𝑥 取值域内的分布情况。一种常见的方法是将空间离散为 𝑁 个点的集合，这些点位于 0,𝐿/(𝑁−1),2𝐿/(𝑁−1),…,(𝑁−2)𝐿/(𝑁−1),𝐿 。最左侧和最右侧的点将分别保持在固定温度 𝑇left 和 𝑇right ，而内部的 𝑁−2 点的温度是我们需要求解的未知数。这些点之间的距离是 Δ𝑥=𝐿/(𝑁−1) ，我们将这些点存储在一个长度为 𝑁 的数组中。对于（零索引的）数组中的每个索引 𝑖 ，其坐标位置为 𝑖𝐿/(𝑁−1)=𝑖Δ𝑥 。

       在离散空间域中，索引 𝑖 处的场导数是附近点的函数。例如，一阶导数的一个简单的离散形式是：$\frac{{\alpha }^2}{\alpha x^2}{f}_i=(f_{i+1}-f_{i-1})/(2\Delta x)$。

而二阶导数的简单的离散形式是：$\frac{{\alpha }^2}{\alpha x^2}{f}_i=(f_{i+1}-2f_i+f_{i-1})/(\Delta x^2)$

如果我们让这个表达式等于 0 来满足拉普拉斯方程，我们得到：$f_{i+1}-2f_i+f_{i-1}=0$

通过对 𝑓𝑖 的求解，我们得到：$f_i=(f_{i+1}+f_{i-1})/2$

2. Jacobi 迭代求解

       尽管 𝑓𝑖+1 和 𝑓𝑖−1 也在变化（边界点 𝑖==0 和 𝑖==𝑁−1 除外），我们只需在这个解之上为 𝑓_𝑖 迭代 多次，直到解达到充分均衡。也就是说，如果在每次迭代中我们都采用 𝑓 的旧解来计算两个相邻点的平均值，并以此作为新解中的每个点，我们最终将求解出 𝑓 的平衡分布。

以（串行）伪代码描述这种方法：

while (error > tolerance):
    l2_norm = 0
    for i = 1, N-2:
        f[i] = 0.5 * (f_old[i-1] + f_old[i+1])
        l2_norm += (f[i] - f_old[i]) * (f[i] - f_old[i])
    error = sqrt(l2_norm / N)
    swap(f_old, f)

3. 单 GPU 的 CUDA 实现

cuda代码实现如下：

#include <iostream>
#include <limits>
#include <cstdio>

inline void CUDA_CHECK (cudaError_t err) 
    if (err != cudaSuccess) 
        fprintf(stderr, "CUDA error: %s\\n", cudaGetErrorString(err));
        exit(-1);
    


#define NUM_POINTS 4194304
#define TOLERANCE  0.0001
#define MAX_ITERS  1000

__global__ void jacobi (const float* f_old, float* f, float* l2_norm, int N)

    int idx = threadIdx.x + blockIdx.x * blockDim.x;

    if (idx > 0 && idx < N - 1) 
        f[idx] = 0.5f * (f_old[idx+1] + f_old[idx-1]);

        float l2 = (f[idx] - f_old[idx]) * (f[idx] - f_old[idx]);

        atomicAdd(l2_norm, l2);
    


__global__ void initialize (float* f, float T_left, float T_right, int N)

    int idx = threadIdx.x + blockIdx.x * blockDim.x;
    
    if (idx == 0) 
        f[idx] = T_left;
    
    else if (idx == N - 1) 
        f[idx] = T_right;
    
    else if (idx < N - 1) 
        f[idx] = 0.0f;
    


int main() 
    // 为“旧”数据的网格数据和临时缓冲区
    // 分配空间。
    float* f_old;
    float* f;

    CUDA_CHECK(cudaMalloc(&f_old, NUM_POINTS * sizeof(float)));
    CUDA_CHECK(cudaMalloc(&f, NUM_POINTS * sizeof(float)));

    // 在主机和设备上为 L2 范数分配内存。
    float* l2_norm = (float*) malloc(sizeof(float));
    float* d_l2_norm;
    CUDA_CHECK(cudaMalloc(&d_l2_norm, sizeof(float)));

    // 将误差初始化为较大的数字。
    float error = std::numeric_limits<float>::max();

    // 初始化数据。
    int threads_per_block = 256;
    int blocks = (NUM_POINTS + threads_per_block - 1) / threads_per_block;

    float T_left = 5.0f;
    float T_right = 10.0f;
    initialize<<<blocks, threads_per_block>>>(f_old, T_left, T_right, NUM_POINTS);
    initialize<<<blocks, threads_per_block>>>(f, T_left, T_right, NUM_POINTS);
    CUDA_CHECK(cudaDeviceSynchronize());

    // 现在进行迭代，直到误差足够小为止。
    // 限制迭代次数，将其用作一种安全机制。
    int num_iters = 0;

    while (error > TOLERANCE && num_iters < MAX_ITERS) 
        // 将范数数据初始化为零
        CUDA_CHECK(cudaMemset(d_l2_norm, 0, sizeof(float)));

        // 启动核函数进行计算
        jacobi<<<blocks, threads_per_block>>>(f_old, f, d_l2_norm, NUM_POINTS);
        CUDA_CHECK(cudaDeviceSynchronize());

        // 交换 f_old 和 f
        std::swap(f_old, f);

        // 更新主机范数；通过按区域数进行规范化并取平方根来计算误差
        CUDA_CHECK(cudaMemcpy(l2_norm, d_l2_norm, sizeof(float), cudaMemcpyDeviceToHost));

        if (*l2_norm == 0.0f) 
            // 处理第一次迭代
            error = 1.0f;
        
        else 
            error = std::sqrt(*l2_norm / NUM_POINTS);
        

        if (num_iters % 10 == 0) 
            std::cout << "Iteration = " << num_iters << " error = " << error << std::endl;
        

        ++num_iters;
    

    if (error <= TOLERANCE && num_iters < MAX_ITERS) 
        std::cout << "Success!\\n";
    
    else 
        std::cout << "Failure!\\n";
    

    // 清理
    free(l2_norm);
    CUDA_CHECK(cudaFree(d_l2_norm));
    CUDA_CHECK(cudaFree(f_old));
    CUDA_CHECK(cudaFree(f));

    return 0;

代码编译和执行命令：

nvcc -x cu -arch=sm_70 -o jacobi jacobi.cpp
./jacobi

程序执行结果：

Iteration = 0 error = 0.00272958
Iteration = 10 error = 0.00034546
Iteration = 20 error = 0.000210903
Iteration = 30 error = 0.000157015
Iteration = 40 error = 0.000127122
Iteration = 50 error = 0.00010783
Success!

4. 利用 NVSHMEM 在多 GPU 上实现

       对于多个 GPU，一个非常简单的分配策略是将域划分为 𝑀 个块(其中 𝑀 为 GPU 的数量)。PE 0 将获得 [0,𝑁/𝑀−1] 内的点，PE 1 将获得 [𝑁/𝑀,2𝑁/𝑀−1] 内的点，以此类推。在这种方法中，PE 之间的通信需要发生在 PE 之间的边界点上。例如，PE0 上点 𝑖=𝑁/𝑀−1 的更新为：$f[N/M-1]=(f[N/M]+f[N/M-2])/2$。

       但此 PE 不包含 𝑖=𝑁/𝑀 的数据点，它属于 PE1。我们需要从远端 PE 获取这个数据点。为此，我们可以使用 nvshmem_float_g() API 来获取远端 PE 上的标量。

float r = nvshmem_float_g(source, target_pe);

结果如下所示。请注意，对于 PE0，位置 N / M对应于 PE1 的索引 0。

f_left = f[N / M - 2]
f_right = nvshmem_float_g(&f[0], 1)
f[N / M - 1] = (f_right + f_left) / 2

代码实现如下：

#include <iostream>
#include <limits>
#include <cstdio>

#include <nvshmem.h>
#include <nvshmemx.h>

inline void CUDA_CHECK (cudaError_t err) 
    if (err != cudaSuccess) 
        fprintf(stderr, "CUDA error: %s\\n", cudaGetErrorString(err));
        exit(-1);
    


#define NUM_POINTS 4194304
以上是关于nvidia CUDA 高级编程使用cub库优化分布式计算下的原子操作的主要内容，如果未能解决你的问题，请参考以下文章
 
 在 nvidia nsight 中安装 CUB
 nvidia CUDA 高级编程NVSHMEM 直方图——复制式方法
 nvidia CUDA 高级编程NVSHMEM 直方图——复制式方法
 nvidia CUDA 高级编程NVSHMEM 直方图——复制式方法
 nvidia CUDA 高级编程NVSHMEM 直方图——分布式方法
 nvidia CUDA 高级编程NVSHMEM 直方图——分布式方法