地震模板代码 - 第二部分

109702008

于 2024-09-02 05:07:19 发布

阅读量972

点赞数 20

分类专栏： # ROCm 人工智能文章标签：人工智能学习

本文链接：https://blog.csdn.net/eidolon_foot/article/details/141748748

版权

人工智能同时被 2 个专栏收录

257 篇文章 3 订阅

订阅专栏

ROCm

46 篇文章 1 订阅

订阅专栏

Seismic stencil codes - part 2 — ROCm Blogs (amd.com)

发布于 2024年8月12日，作者为 [Justin Chang](Justin Chang — ROCm Blogs) 和 [Ossian O'Reilly](Ossian O’Reilly — ROCm Blogs)。

在上一篇文章中，我们提到，z方向上的模板计算内核由于低有效带宽而表现不佳。这种低性能是由于大量数据需要移动到全局内存造成的。为了解决这一性能瓶颈，我们将考虑两种优化策略：一种是通过增加内存占用来提升性能，另一种是将压力从内存子系统转移到GPU的向量寄存器上。为了简化实验，所有实验都使用三维立方网格，尺寸 nx, ny, nz = 512, 512, 512，模板半径 R=4。

对齐和未对齐的内存访问

第一个优化方法是对齐内存访问，使每个线程块中的波段写入一整套缓存行。为了对齐内存访问，需要填充领先维度并偏移基指针。领先维度应是缓存行大小的倍数。否则，缓存行的组织方式会根据访问的计算网格的某一行或某一层发生变化。下面的图示展示了在未对齐情况下（左）和通过正确的偏移和填充后的对齐情况下（右），访问缓存行数量如何变化。

图1：二维数组的未对齐和对齐的缓存行访问情况。（左）未对齐的访问；（右）对齐的访问。每个浅蓝色阴影框表示一个缓存行。

在这里，我们通过填充使得前导维度总是能被64整除。

int align = 64;
// Define mx = nx + 2 * R and pad it by some amount such that mx is divisible
// by 'align'
int mx = nx + 2 * R;
int mx = ((mx - 1) /align + 1) * align;

自然，另一种选择是限制网格大小，使其尽可能是波长的倍数。这种策略可以带来更进一步的性能提升，因为它确保波内的所有线程都被充分利用，即在边界附近没有空闲的线程。

为了在不改变内核本身的情况下确保内存访问对齐，我们偏移了基指针。

int offset = align - R;
hipMalloc(&d_p_in, (mx * (ny + 2 * R) * (nz + 2 * R) + offset) * sizeof(float));
hipMalloc(&d_p_out, (mx * (ny + 2 * R) * (nz + 2 * R) + offset) * sizeof(float));
d_p_in += offset;
d_p_out += offset;

这个偏移假设输入参数指定了基线内核在x方向上的起始索引为 x0=R。回顾基线内核的实现：

  template <int R>
  __launch_bounds__((BLOCK_DIM_X) * (BLOCK_DIM_Y) * (BLOCK_DIM_Z))
  __global__ void compute_fd_x_gpu_kernel(float *__restrict__ p_out, const float *__restrict__ p_in,
          int line, int slice, int x0, int x1, int y0, int y1, int z0, int z1) {

      const int i = x0 + threadIdx.x + blockIdx.x * blockDim.x;
      const int j = y0 + threadIdx.y + blockIdx.y * blockDim.y;
      const int k = z0 + threadIdx.z + blockIdx.z * blockDim.z;

      if (i >= x1 || j >= y1 || k >= z1) return;

      size_t pos = i + line * j + slice * k;
      int stride = 1; // x-direction

      // Shift pointers such that that p_in points to the first value in the stencil
      p_in += pos - R * stride;
      p_out += pos;

      // Compute the finite difference approximation
      float out = 0.0f;
      for (int r = 0; r <= 2 * R; ++r)
          out += p_in[r * stride] * d_dx<R>[r]; // x-direction

      // Write the result
      p_out[0] = out;

  }

基准性能

接下来，我们展示对齐与非对齐内存访问对性能的影响。在实验中，除了是否对齐内存访问之外，所有设置都保持相同。当内存访问对齐时，数组会进行填充，使得主维度可以被64整除，并且基指针会在半径`R=4`的情况下偏移60个元素。下表显示了在AMD Instinct MI200系列GPU上，内存访问的正确对齐对性能有显著影响。

Kernel	Effective memory bandwidth	Speedup
x-direction: Unaligned	956 GB/s	1.00
x-direction: Aligned	990 GB/s	1.04
y-direction: Unaligned	940 GB/s	1.00
y-direction: Aligned	966 GB/s	1.03
z-direction: Unaligned	217 GB/s	1.00
z-direction: Aligned	297 GB/s	1.37

需要提到的是，对齐版本由于数组填充会使用略多的内存。对于一个512x512x512的立方体，这种内存消耗的增加约为10%。另一方面，有效的内存带宽提高了最多1.37倍。

在L1和L2缓存级别，对齐的内存访问模式减少了请求次数，这是因为在非对齐的情况下，波请求的数据可能会分布在完整和部分的缓存行上。缓存以其缓存行粒度操作，这意味着任何部分缓存行请求仍然需要通过分配和移动完整缓存行来服务。由于邻近的线程块也需要访问这些部分缓存行，如果它们被缓存至L2，它们会被重用。因此，当切换对齐与非对齐内存访问时，从主存/全局内存（对于MI250来说称为高带宽内存或HBM）移动的数据量不会有显著差异。您可以通过收集以下性能计数器来观察这些行为：

Performance counter	Description
`TCP_TOTAL_CACHE_ACCESSES_sum`	The total number of cache line accessed in the L1 cache.
`TCP_TCC_READ_REQ_SUM`	The total number of read requests that missed in L1 and went to the L2 cache.
`TCP_TCC_WRITE_REQ_SUM`	The total number of write requests that missed in L1 and went to the L2 cache.

要将请求数量转换为移动的数据量（以字节为单位），需要将它们乘以64 B - 即MI250的L1缓存行大小。

Kernel	L1	L2 read	L2 write	HBM
x-direction: Unaligned	7247 MB	570 MB	671 MB	1088 MB
x-direction: Aligned	6979 MB	604 MB	537 MB	1105 MB
y-direction: Unaligned	10737 MB	2009 MB	671 MB	1097 MB
y-direction: Aligned	5369 MB	1610 MB	537 MB	1080 MB
z-direction: Unaligned	10737 MB	6040 MB	671 MB	4447 MB
z-direction: Aligned	5369 MB	4832 MB	537 MB	4268 MB

为了进一步了解整个内存子系统中的数据移动量，将表格中的每条数据除以立方体大小 (512 * 512 * 512) * sizeof(float)，这大约是536 MB。请参阅新的比率：

Kernel	L1	L2 read	L2 write	HBM
x-direction: Unaligned	13.5	1.06	1.25	2.03
x-direction: Aligned	13.0	1.13	1.00	2.06
y-direction: Unaligned	20.0	3.75	1.25	2.05
y-direction: Aligned	10.0	3.00	1.00	2.01
z-direction: Unaligned	20.0	11.3	1.25	8.30
z-direction: Aligned	10.0	9.00	1.00	7.96

这个数字表示通过内存层次结构每一级移动的立方体数量。对于x方向和y方向的内核，大约有两个立方体在L2缓存和HBM之间移动。这种情况是理想的，因为它对应于读写一次立方体。然而，我们注意到从L1和L2缓存中最多处理了20个立方体的数据。理想情况下，我们希望这些立方体尽可能少地加载和存储，而对齐内存访问减少了这些操作。接下来的子部分将分析这种现象在每个方向上的表现。

X方向访问

在HBM数据传输量方面，内存对齐优化对比非对齐的差异较小，这表明这种优化主要影响的是通过缓存移动的缓存行数量。对于x方向核算，即使数据是否缓存行对齐，大多数邻近访问的结果仍然会命中相同的缓存行（如图2所示）。因此，在内存子系统的各个层级中，数据传输量的差异较小。

图2: 由于半径为R的halo区域，在x方向的线程块边界处需要加载部分缓存行，无论数据是否对齐缓存行。

图3: 在非对齐情况下，当加载邻近数值时，访问的缓存行数量可以与对齐情况下相同，或者少一条缓存行，这取决于缓存行边界的位置。

Y方向访问

对比之下，对于y方向核算，每个邻近访问都会命中不同的缓存行。在没有对齐的情况下，每个线程块访问的缓存行数量远大于对齐情况（如图3和图4所示）。

图3: 该图显示了一组线程沿着线程块底部边缘在y方向上计算stencil需要访问的缓存行（只显示了底部一半的邻居）。y方向的每个邻居访问命中不同的缓存行。由于数据对齐，所有的缓存行是共线的，这在使用缓存计算y方向stencil时，最小化了所需的访问次数。

图4: 与图3相同，但数据是未对齐的。在这种情况下，每行邻居的y方向stencil计算访问了额外的缓存行。

Z方向访问

尽管对齐的Z方向核高度饱和了主内存带宽，但该核的性能仍远远没有达到其极限。由于理想情况下该核只需加载和写入两个数组，最佳的读写比应接近于1。X和Y方向的内核已经接近这个最佳限度，但Z方向的内核远未达到：

Kernel	HBM Read / Write
x-direction: Unaligned	1.01
x-direction: Aligned	1.07
y-direction: Unaligned	1.02
y-direction: Aligned	1.02
z-direction: Unaligned	7.05
z-direction: Aligned	7.02

注意： 正如下述链接中的[有限差分拉普拉斯博客系列第4部分](Seismic stencil codes - part 2 — ROCm Blogs)中所述，一个512 x 512 x 512立方体的XY平面可以适应单个 MI250X GCD的L2缓存，因此有所重复使用，这解释了为什么Z方向的读写比是7.0而不是9.0。在未包含于这篇博客文章的实验中，一个1024 x 1024 x 1024的立方体确实具有9.0的读写比。

图5：由于数据重用不足，在Z方向的栅格运算中存在大量的HBM数据移动。线程块是二维的，设备上所有当前运行的线程块都处于立方体的狭窄带状区域内（图中的活动计算集）。线程块将其上下的halo数据加载到L2缓存中。L2缓存不足以同时容纳多个数据切片。当一个线程块在另一个线程块下方计算时，该数据很有可能已从缓存中逐出。因此，数据需要再次从HBM重新加载。

这种Z方向上的过多读操作源于栅格计算中几乎没有数据重用（图5）。回顾一下，栅格计算需要加载8个邻近值。理想情况下，这些值应从缓存或寄存器中加载，而不是从主内存中加载。由于当前核实现、线程块配置和问题大小，这些值未能在L2缓存中保持以供重复使用。在当前的线程块配置中，线程以二维模式排列。因此，在垂直方向上没有波动的线程块，也没有共享数据的垂直方向的邻近波动线程。因此，任何线程块加载到L2缓存中的数据极不可能被重复使用，因为这需要一个线程块在其上方或下方的切片上工作。由于这些原因，L2缓存中几乎没有可被其他线程块重复使用的数据（因此上面提到的7.0读写比）。虽然可以通过选择其他线程块大小在一定程度上缓解这种情况，但我们将通过更改算法来解决这一问题。
由于Z方向栅格计算是表现最差的内核，下一步的优化将专注于提高其性能。

滑动窗口

z方向的内核展示了实现高阶有限差分时遇到的一个关键挑战，即如何在进行模版计算时有效地重用数据。为了达到出色的性能，避免昂贵的全局内存读操作是至关重要的。幸运的是，高阶有限差分产生的模版具有重复的模式，可以使用滑动窗口技术来处理（图6）。

图6：滑动窗口技术用于模版计算的关键思想是将输入数组数据的局部体积保存在寄存器中，并完全使用这些寄存器计算z方向的模版。要计算下一个z值的模版，只需将一个新的数据切片加载到寄存器中，并重复使用之前在寄存器中的数据。
滑动窗口算法可以分为四个步骤实现：
1. 将模版中除了最后一个值的所有值加载到一个寄存器数组中。
2. 将模版中的最后一个值加载到寄存器数组中。
3. 使用寄存器数组中的所有值计算有限差分近似。
4. 通过向前移动一步更新滑动窗口，将寄存器数组中的第一个值覆盖掉。

对于滑动窗口更新的每个网格点，重复步骤2到4。下面的图像描述了假设R=4和输入数组`p[i]`初始化为`p[i] = i * i`时上述步骤的图示：

图7：步骤1 - 初始化窗口，加载前8个值

图8：步骤2和3 - 加载窗口的最后一个值，计算有限差分并存储结果。

滑动窗口将N次更新的主内存加载指令从(2R+1)*N减少到2R+N，代价是需要2R+1个寄存器。

图9：步骤4 - 更新滑动窗口

图10：用新移动的滑动窗口重复步骤2到4

滑动窗口将N次更新的主内存加载指令从(2R+1)*N减少到2R+N，代价是需要2R+1个寄存器。

这是我们之前看到的高阶模版函数的滑动窗口实现：

template <int stride, int R>
__inline__ void compute_fd_z_sliding_window(float *p_out, const float *p_in, const float *d,
        int begin, int end, int pos, int stride) {
    // Sliding window
    float w[2 * R + 1];

    // 1. Prime the sliding window
    for (int r = 0; r < 2 * R; r++)
        w[r] = p_in[pos + (begin - R + r) * stride];

    // Apply the sliding window along the given grid direction determined by `stride`
    for (int i = begin; i < end; ++i) {

        // 2. Load the next value into the sliding window at its last position
        w[2 * R] = p_in[pos + (i + R) * stride];

        // 3. Compute the finite difference approximation using the sliding window
        float out = 0.0f;
        for (int r = 0; r <= 2 * R; r++)
            out += w[r] * d[r];
        p_out[pos + i * stride] = out;

        // 4. Update the sliding window by shifting it forward one step
        for (int r = 0; r < 2 * R; r++)
            w[r] = w[r+1];

    }
}

这个函数可以如下使用来计算z方向的滑动窗口：

// Apply approximation in the x-direction for all interior grid points
for (int j = R; j < ny + R; ++j) {
    for (int i = R; i < nx + R; ++i) {
        const uint64_t pos = i + line * j;          
        compute_fd_z_sliding_window<R>(p_out, p_in, d, R, nz + R, pos, slice);
    }
}

GPU上的滑动窗口

基于GPU的滑动窗口实现与主机上的实现有很多相似之处。以下是实现方式：

template <int R>
__launch_bounds__((BLOCK_DIM_X) * (BLOCK_DIM_Y))
__global__ void compute_fd_z_window_kernel(float *p_out, const float *p_in, const float *d, int line, int
        slice, int x0, int x1, int y0, int y1, int z0, int z1, int dimz) {

    const int i = x0 + threadIdx.x + blockIdx.x * blockDim.x;
    const int j = y0 + threadIdx.y + blockIdx.y * blockDim.y;

    if (i >= x1 || j >= y1) return;

    // Determine the k indices covered by this sliding window
    // The extent cannot exceed the z1
    const int kbegin = z0 + blockIdx.z * dimz;
    const int kend = kbegin + dimz > z1 ? z1 : kbegin + dimz;

    // Sliding window
    float w[2 * R + 1];

    size_t pos = i + line * j + slice * kbegin;

    // Shift pointers such that that p_in points to the first value in the sliding window
    p_in += pos - R * slice;
    p_out += pos;

    // 1. Prime the sliding window
    for (int r = 0; r < 2 * R; ++r) {
        w[r] = p_in[0];
          p_in += slice;
      }

   // Apply the sliding window along the given grid direction
   for (int k = kbegin; k < kend; ++k) {

       // 2. Load the next value into the sliding window at its last position
       w[2 * R] = p_in[0];

       // 3. Compute the finite difference approximation using the sliding window
       float out = 0.0f;
       for (int r = 0; r <= 2 * R; ++r)
           out += w[r] * d_dz<R>[r];
       p_out[0] = out;

       // 4. Update the sliding window by shifting it forward one step
       for (int r = 0; r < 2 * R; ++r)
           w[r] = w[r+1];

       // Increment pointers
       p_in += slice;
       p_out += slice;
   }
}

在此实现中，线程块配置必须为2D，线程布局在XY平面上。每个线程块内的每个线程负责计算一个从 kbegin 开始到 kend-1 结束的滑动窗口。这些起始和终止点取决于线程块在z方向上应更新的网格点数。在此实现中，除非线程块位于边界且工作量不能被网格尺寸整除，否则每个线程块滑动相同数量的网格点数 𝑛𝑤。

如前所述，虽然线程块维度是2D的，但线程块覆盖的是一个3D线程块网格。内核是通过以下方式启动的：

    dim3 block(BLOCK_DIM_X, BLOCK_DIM_Y);
    dim3 grid;
    grid.x = ceil(x1 - x0, block.x);
    grid.y = ceil(y1 - y0, block.y);
    grid.z = ceil(z1 - z0, nw);

    compute_fd_z_window_kernel<R><<<grid, block>>>(p_out, p_in, d, line, slice, x0, x1,
            y0, y1, z0, z1, nw);

性能考虑

控制线程块滑动网格点数的参数𝑛𝑤可能有重大性能影响。如果𝑛𝑤=1，滑动窗口简化为仅更新每个线程一个网格点的基线内核。在这种情况下，我们知道加载指令的数量将爆炸至(2𝑅+1)𝑛𝑧。另一方面，这种选择暴露了𝑛𝑧量的并行性。在相反的极限下，如果𝑛𝑤=𝑛𝑧，加载指令的数量为2𝑅+𝑛𝑧，但在z方向上没有暴露并行性。在一般情况下，加载指令的数量是(2𝑅+𝑛𝑤)𝑛𝑧/𝑛𝑤，而𝑛𝑧/𝑛𝑤 是暴露的并行性。因此，根据问题的维度，在减少加载指令数量和暴露足够多的并行性以饱和硬件资源之间可能存在权衡。如果我们假设并行性减少是可以忽略的，那么速度提升的上限是基线内核和最佳情况加载和存储次数的比率：𝑆≤(1+(2𝑅+1))/(1+1)=𝑅+1。对于𝑅=4，速度提升最多是5倍。

图11：滑动窗口内核的有效内存带宽高达1200 GB/s[^1]，这大约是基线内核的4倍速度提升。这种速度提升来自于将读/写比率减少了大约7倍。
在上图中，左侧子面板显示了有效和实际内存带宽作为窗口维度𝑛𝑤的函数，使用的是半对数比例图。对于有效带宽来说，有一个收益递减点。这一点出现在当增加𝑛𝑤时HBM读写比几乎没有变化的时候。右侧子面板显示了比较的HBM读写比，预测结果和实际结果。回顾我们之前的讨论，由于xy平面的尺寸，在L2缓存中存在一些重用，因此实际值为7.0而预测值为9.0。不管怎样，两条曲线都接近1。
由于在增加并行性和减少读写比率之间的权衡，最高速度提升大约是4倍，而不是5倍。在我们的实验中，暴露并行性的减少逐渐降低了实际带宽。z方向内核的实际和有效内存带宽现在接近1200 GB/s[1]。

结论

缓存行对齐是一项需要考虑的重要优化，对于模板计算来说，因为它可以减少缓存行访问的次数。对于任何利用缓存的模板计算来说，使用鼓励模板邻近值高重用率的线程块配置是很重要的。对于z方向的模板，如果问题规模足够大且线程块是二维的，那么可能几乎没有数据重用。这一问题出现的原因在于，缓存不够大，无法容纳为了缓存z方向模板邻近值所需的数据切片数量。因此，高带宽内存（HBM）的数据移动量会被放大，性能会急剧下降。解决这一问题的一个有效方法是在AMD GPU上使用滑动线程块技术。当滑动线程块技术与对齐优化结合使用时，相较于原始的、基线情况下，可以实现近5.5倍的速度提升。另外，进一步减少全局内存请求次数的另一个方法是使用用户可编程缓存，即所谓的本地数据共享（LDS）。下一篇博客将重点讨论这一主题，涉及x方向、y方向内核，以及将所有三个内核结合成一个单一的整体内核。

附带代码示例
如果您有任何问题或意见，请通过GitHub与我们联系Discussions

[1]
测试使用ROCm版本6.1.0-82进行。基准测试结果并非经过验证的性能数据，仅用于展示代码修改后的相对性能提升。实际性能结果受多种因素影响，包括系统配置和环境设置，结果的可重复性不作保证。