地震模板代码 - 第三部分

109702008

于 2024-09-02 05:08:44 发布

阅读量1.2k

点赞数 23

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

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

版权

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

257 篇文章 3 订阅

订阅专栏

ROCm

46 篇文章 1 订阅

订阅专栏

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

2024年8月12日，作者：Justin Chang 和 Ossian O’Reilly。

在前两篇博客文章中，我们开发了一个 HIP 内核，能够计算地震波传播中常用的高阶有限差分。经过优化后，z 方向的内核（在初始实现中表现最差的内核）在单个 MI250X GCD 上实现了近 1200 GB/s 的实际和有效内存带宽[1]。这种性能相当于达到可实现内存带宽的 85%，详见[2]。在地震模板博客系列的第三部分也是最后一部分中，我们优化了x方向和y方向的内核，以提高它们的性能。随后，我们将所有三个内核融合成一个单一的单片内核，遵循在本系列博客第一部分中简要提到的_一次通过方法_。

回顾第一部分的基线实现：

  template <int R>
  __launch_bounds__((BLOCK_DIM_X) * (BLOCK_DIM_Y) * (BLOCK_DIM_Z))
  __global__ void compute_fd_x_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;

  }

在 512 x 512 x 512 问题规模下，不同内存级别（L1、L2 和 HBM）的内存流量（前导维度 x 按 64 进行填充对齐）的具体数值为：

Kernel	L1	L2 read	L2 write	HBM
x-direction: Memory	6979 MB	604 MB	537 MB	1105 MB
x-direction: Memory / cube ratio	13.0	1.13	1.00	2.06
y-direction: Memory	5369 MB	1610 MB	537 MB	1080 MB
y-direction: Memory / cube ratio	10.0	3.00	1.00	2.01

尽管 HBM 和 L2 写入比率达到预期的水平（例如，在 HBM 级别上大约有 2 个立方体的移动，在 L2 级别上写入 1 个立方体），但 L1 缓存访问量却高得令人不快。这种通过 L1 的过多数据流动可能是导致有效和实际内存带宽都低于 1 TB/s[1] 的潜在原因，尽管对于简单的流式基准测试，可能达到 1.3 到 1.4 TB/s。接下来的两项优化探讨了减少这一流量的技术。同样，我们在 512 x 512 x 512 的立方体上并且在单个 MI250X GCD 上进行了所有实验，且`R=4`。确保内存对齐和偏移使得前导维度 x 可以被 64 整除。

优化1 - 向量化浮点数

在我们的有限差分拉普拉斯系列中（参见这里），我们介绍了循环展开的概念，它通过扩大块的大小并增加通过寄存器的数据重用来提高性能。我们将在这里应用类似的方法，通过利用向量化浮点数来实现这一点。该技术的核心思想是通过向量指令扩大x方向上的块大小。因此，线程块中的每个线程将负责计算一个、两个或四个输出元素，这取决于向量的大小。自然地，向量指令会增加每个线程的寄存器压力，可能会影响占用率。然而，向量指令的一个好处是它们利用了硬件可以针对全局内存加载和存储指令的每个通道请求高达128位的数据这一事实。由于指令更宽，它们减少了程序中的全局加载/存储指令总数以及一些与计算地址偏移相关的整数运算。如果占用率的降低比向量大小的因素小，并且没有寄存器溢出，这种优化将增加飞行中的内存指令数，这是饱和内存带宽的关键。
要在我们的示例中实现这一点，主导维度`nx`必须是向量大小的1、2或4的倍数（对应于`float`、`float2`或`float4`）。

首先，我们从引入一些预处理程序头文件开始：

// Vectorized floats
#ifndef VEC_EXP
#define VEC_EXP 0
#endif
#define VEC_LEN (1 << (VEC_EXP))
using vec = __attribute__((__vector_size__(VEC_LEN * sizeof(float)))) float;

用户将传递一个编译器标志`-DVEC_EXP`，其值为0（无打包数学）、1（`float2`）或2（`float4`）。我们的模板代码的基本设计是，普通的`float`数组仍然传递到HIP内核，但是在内核内部，我们通过引入新指针重新将`float *转换为新定义的vec`。接下来，我们为x方向的HIP内核引入一些额外的预处理程序头文件：

// x window stencil
#if VEC_LEN == 1
#define XREG RADIUS
#define XREG_VEC RADIUS
#elif VEC_LEN == 2
#if RADIUS > 2
#define XREG 4
#define XREG_VEC 2
#else
#define XREG 2
#define XREG_VEC 1
#endif
#else
#define XREG VEC_LEN
#define XREG_VEC 1
#endif
#define XREG_OFF (XREG-RADIUS)

向量化的浮点数需要从输入缓冲区`p_in`中处理相邻且对齐的值。例如，当`VEC_LEN == 2`即`float2`时，x方向内核中的每个线程计算2个网格点的模板。类似地，当`VEC_LEN == 4`即`float4`时，x方向内核中的每个线程计算4个网格点的模板。

考虑基线x方向内核中的for循环：

    // Compute the finite difference approximation
    float out = 0.0f;
    for (int r = 0; r <= 2 * R; ++r) {
        out += p_in[0] * d_dz<R>[r];
        p_in += 1;
    }

该代码示例需要重写以适应向量化的加载和存储指令：

    const vec *p_in_vec = reinterpret_cast<const vec*>(p_in);
    vec *p_out_vec = reinterpret_cast<vec*>(p_out);
    float x_reg[2 * XREG + VEC_LEN] = {0.0f};
    vec *x_reg_vec = reinterpret_cast<vec*>(x_reg);

    // Read x into registers
    for (int r = 0; r < 2 * XREG_VEC + 1; ++r)
        x_reg_vec[r] = p_in_vec[0 - XREG_VEC + r];
     
    // Compute the finite difference approximation
    vec out = {0.0f};
    for (int r = 0; r <= 2 * R; ++r) {
        for (int ii = 0; ii < VEC_LEN; ++ii)
            out[ii] += x_reg[XREG_OFF + r + ii] * d_dx<R>[r]; 
    }

进行了几项代码更改：
1. 将所有的`float *缓冲区重新转换为新定义的向量化vec *`
2. 类似于前一篇文章中介绍的滑动窗口概念，引入一个寄存器`x_reg`。当`R == 4`且`VEC_LEN == 1`时，该`x_reg`寄存器包含9个元素，但当`VEC_LEN == 2`和`VEC_LEN == 4`时，分别包含10个和12个元素。同样，`x_reg_vec`寄存器包含9个、5个和3个元素，当`VEC_LEN == 1`、`VEC_LEN == 2`和`VEC_LEN == 4`时
3. 将for循环分成两部分。第一个循环只是将向量化的浮点数加载到`x_reg_vec`中
4. 第二个循环是一个双嵌套循环，其中`x_reg`和`d_dx<R>被调整以计算每个VEC_LEN`网格点的有限差分。

下面是上述定义如何融入整体向量化的图示描述：

图1：本图说明了向量化加载和存储如何影响为一个线程计算的模板点数。每个线程执行相同的操作，需要从存储在寄存器中的数据中读取多个值，并计算多个输出的模板公式，由向量长度`VEC_LEN`确定。指向寄存器数组的箭头表示全局加载指令，指向寄存器数组外的箭头表示全局存储指令。因此，该实现依赖于L1缓存来高效重用许多请求的模板邻居。

另一个重大的变化是内核的启动配置：

#define BLOCK_DIM_X (64 * (4 / RADIUS))
#define BLOCK_DIM_Y RADIUS

以前，模板内核的启动线程块配置为256 x 1。现在，线程块配置根据有限差分近似的顺序选择。因此，如果`R=4`，则内核的线程块配置为64 x 4。

带有启动参数的完整代码如下：

template <int R>
__launch_bounds__((BLOCK_DIM_X) * (BLOCK_DIM_Y))
__global__ void compute_fd_x_vec_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 + VEC_LEN * (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;

    // Shift pointers such that that p_in points to the first value in the stencil
    p_in += pos;
    p_out += pos;
    
    const vec *p_in_vec = reinterpret_cast<const vec*>(p_in);
    vec *p_out_vec = reinterpret_cast<vec*>(p_out);
    float x_reg[2 * XREG + VEC_LEN] = {0.0f};
    vec *x_reg_vec = reinterpret_cast<vec*>(x_reg);

    // Read x into registers
    for (int r = 0; r < 2 * XREG_VEC + 1; ++r)
        x_reg_vec[r] = p_in_vec[0 - XREG_VEC + r];
     
    // Compute the finite difference approximation
    vec out = {0.0f};
    for (int r = 0; r <= 2 * R; ++r) {
        for (int ii = 0; ii < VEC_LEN; ++ii)
            out[ii] += x_reg[XREG_OFF + r + ii] * d_dx<R>[r]; 
    }

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

}

template <int R>
void compute_fd_x_vec(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) {

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

    compute_fd_x_vec_kernel<R><<<grid, block>>>(p_out, p_in, d, line, slice, x0, x1, y0, y1, z0,
            z1);
    HIP_CHECK(hipGetLastError());
     
}

对y方向内核进行向量化的代码转换比较简单：

template <int R>
__launch_bounds__((BLOCK_DIM_X) * (BLOCK_DIM_Y))
__global__ void compute_fd_y_vec_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 + VEC_LEN * (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;
    size_t line_vec = line >> VEC_EXP;

    // Shift pointers such that that p_in points to the first value in the stencil
    p_in += pos - R * line;
    p_out += pos;
    
    const vec *p_in_vec = reinterpret_cast<const vec*>(p_in);
    vec *p_out_vec = reinterpret_cast<vec*>(p_out);
    
    // Compute the finite difference approximation
    vec out = {0.0f};
    for (int r = 0; r <= 2 * R; ++r) {
        out += p_in_vec[0] * d_dy<R>[r];
        p_in_vec += line_vec;
    }

    // Write the result
    p_out_vec[0] = out;
}

与x方向的核函数类似，每个线程在x方向上输出`VEC_LEN`个元素。然而，在加载数据时，是从向量化类型中加载的，同时在y方向上进行跨步。因此，步幅需要根据向量长度`VEC_LEN`进行修改。这个核函数使用与之前相同的线程块配置。

下面是内存带宽的数据：

Kernel	Effective memory bandwidth	Achieved memory bandwidth
x-direction: Baseline	990 GB/s	1018 GB/s
x-direction: `VEC_LEN=1`	980 GB/s	1007 GB/s
x-direction: `VEC_LEN=2`	1216 GB/s	1256 GB/s
x-direction: `VEC_LEN=4`	1240 GB/s	1280 GB/s
y-direction: Baseline	966 GB/s	970 GB/s
y-direction: `VEC_LEN=1`	967 GB/s	968 GB/s
y-direction: `VEC_LEN=2`	831 GB/s	834 GB/s
y-direction: `VEC_LEN=4`	1163 GB/s	1172 GB/s

几点观察结果。首先，基线和`VEC_LEN=1`的核函数表现相似，这符合预期，因为它们都操作`float`数组。其次，所有的核函数都从打包的FP32操作中受益。当`VEC_LEN=2`时，x方向核函数的有效和实际内存带宽都很高。y方向核函数的结果则有些混杂，具体取决于`VEC_LEN`的值，但当`VEC_LEN=4`时，核函数表现也很好。现在让我们更深入地探讨一下内存级别的流量：

Kernel	L1	L2 read	L2 write	HBM
x-direction: Baseline	6979 MB	604 MB	537 MB	1105 MB
x-direction: `VEC_LEN=1`	6979 MB	805 MB	537 MB	1105 MB
x-direction: `VEC_LEN=2`	5906 MB	672 MB	537 MB	1105 MB
x-direction: `VEC_LEN=4`	3221 MB	604 MB	537 MB	1105 MB
y-direction: Baseline	5369 MB	1610 MB	537 MB	1080 MB
y-direction: `VEC_LEN=1`	5369 MB	1610 MB	537 MB	1080 MB
y-direction: `VEC_LEN=2`	10737 MB	1611 MB	537 MB	1080 MB
y-direction: `VEC_LEN=4`	5369 MB	1758 MB	537 MB	1080 MB

对于x方向的核函数，打包FP32操作和数据移动减少了L1缓存的流量。当`VEC_LEN=2`时，L1流量几乎翻倍，这表明存在合并不足的情况。y方向的核函数，即使在`VEC_LEN=4`的情况下，也稍微落后于相应的x方向核函数，但经历了高水平的L1流量。现在让我们考虑针对y方向的第二个优化。

优化二 - 本地数据共享（LDS）

本地数据共享（LDS）是一个位于计算单元（CU）上的快速、用户可编程的缓存，可用于高效地在线程块中的所有线程之间共享数据。如在[寄存器压力](Register pressure in AMD CDNA™2 GPUs — ROCm Blogs)一文中讨论的那样，LDS是线程块中所有线程共有的几种内存资源之一。在MI200 GPU的每个CU上都有64KiB的LDS容量。由于每个CU上的所有活动线程块共享LDS资源，如果每个线程块需要过多的LDS，则可能会导致占用率下降。然而，即使占用率下降，内核的性能也可能会提高，因为LDS指令可以取代需要通过缓存甚至HBM（高带宽内存）迁移数据的全局内存指令。

尽管HBM流量已经达到了预期水平，利用LDS可以缓解这个特定内核的L1缓存压力。以下是应用了向量化加载和存储指令以及LDS的修改代码：

template <int R>
__launch_bounds__((BLOCK_DIM_X) * (BLOCK_DIM_Y))
__global__ void compute_fd_y_vec_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 + VEC_LEN * (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;

    size_t pos = i + line * j + slice * k;
    size_t spos = threadIdx.x + (y0 + threadIdx.y) * BLOCK_DIM_X;
    size_t line_vec = line >> VEC_EXP;

    // Shift pointers such that that p_in points to the first value in the stencil
    p_in += pos;
    p_out += pos;
    
    const vec *p_in_vec = reinterpret_cast<const vec*>(p_in);
    vec *p_out_vec = reinterpret_cast<vec*>(p_out);

    const int lds_y = BLOCK_DIM_Y + 2 * R;
    __shared__ vec smem[BLOCK_DIM_X * lds_y];

    // Read y into LDS
    smem[spos - (BLOCK_DIM_X * R)          ] = p_in_vec[0 - R * line_vec];
    smem[spos                              ] = p_in_vec[0];
    smem[spos + (BLOCK_DIM_X * BLOCK_DIM_Y)] = p_in_vec[0 + line_vec * BLOCK_DIM_Y];
    __syncthreads();
    
    if (i >= x1 || j >= y1 || k >= z1) return;

    // Compute the finite difference approximation
    vec out = {0.0f};
    for (int r = 0; r <= 2 * R; ++r) {
        out += smem[spos + (r - R) * BLOCK_DIM_X] * d_dy<R>[r];
    }

    // Write the result
    p_out_vec[0] = out;
}

上面的代码示例引入了以下几点：
1. 一个新寄存器`spos`，跟踪LDS内存中的线程索引。
2. 一个静态的LDS内存分配，`smem`，用于保存大小为`BLOCK_DIM_X * (BLOCK_DIM_Y + 2 * R)`的`vec`类型。
3. 以循环方式将`p_in_vec`读入`smem`。
4. 使用加载到LDS数组`smem`中的数据计算和存储有限差分模板。

之前我们看到`VEC_LEN=4`表现最好，因此我们选择此VEC_LEN进行LDS实现。将`-Rpass-analysis=kernel-resource-usage` 参数添加到编译标志中，可以快速检查内核资源和添加LDS的影响：

remark: ./compute_fd_y_vec.hpp:13:0:     SGPRs: 44 [-Rpass-analysis=kernel-resource-usage]
remark: ./compute_fd_y_vec.hpp:13:0:     VGPRs: 39 [-Rpass-analysis=kernel-resource-usage]
remark: ./compute_fd_y_vec.hpp:13:0:     AGPRs: 0 [-Rpass-analysis=kernel-resource-usage]
remark: ./compute_fd_y_vec.hpp:13:0:     ScratchSize [bytes/lane]: 0 [-Rpass-analysis=kernel-resource-usage]
remark: ./compute_fd_y_vec.hpp:13:0:     Occupancy [waves/SIMD]: 8 [-Rpass-analysis=kernel-resource-usage]
remark: ./compute_fd_y_vec.hpp:13:0:     SGPRs Spill: 0 [-Rpass-analysis=kernel-resource-usage]
remark: ./compute_fd_y_vec.hpp:13:0:     VGPRs Spill: 0 [-Rpass-analysis=kernel-resource-usage]
remark: ./compute_fd_y_vec.hpp:13:0:     LDS Size [bytes/block]: 0 [-Rpass-analysis=kernel-resource-usage]

...

remark: ./compute_fd_y_lds_vec.hpp:13:0: Function Name: _Z27compute_fd_y_lds_vec_kernelILi4EEvPfPKfS2_iiiiiiii [-Rpass-analysis=kernel-resource-usage]
remark: ./compute_fd_y_lds_vec.hpp:13:0:     SGPRs: 24 [-Rpass-analysis=kernel-resource-usage]
remark: ./compute_fd_y_lds_vec.hpp:13:0:     VGPRs: 21 [-Rpass-analysis=kernel-resource-usage]
remark: ./compute_fd_y_lds_vec.hpp:13:0:     AGPRs: 0 [-Rpass-analysis=kernel-resource-usage]
remark: ./compute_fd_y_lds_vec.hpp:13:0:     ScratchSize [bytes/lane]: 0 [-Rpass-analysis=kernel-resource-usage]
remark: ./compute_fd_y_lds_vec.hpp:13:0:     Occupancy [waves/SIMD]: 5 [-Rpass-analysis=kernel-resource-usage]
remark: ./compute_fd_y_lds_vec.hpp:13:0:     SGPRs Spill: 0 [-Rpass-analysis=kernel-resource-usage]
remark: ./compute_fd_y_lds_vec.hpp:13:0:     VGPRs Spill: 0 [-Rpass-analysis=kernel-resource-usage]
remark: ./compute_fd_y_lds_vec.hpp:13:0:     LDS Size [bytes/block]: 12288 [-Rpass-analysis=kernel-resource-usage]

具有`VEC_LEN=4`的LDS实现对于线程块大小为256（每SIMD波一组）的线程块需要12288字节。尽管降低了SGPR（标量通用寄存器）和VGPR（向量通用寄存器）的使用量，LDS的使用导致占用率下降到5：`12.288 KB x 5 = 61.44 KB < 64 KiB`。下表显示了性能情况。

Kernel	Effective memory bandwidth	Achieved memory bandwidth
y-direction: `VEC_LEN=4`	1163 GB/s	1172 GB/s
y-direction: `VEC_LEN=4` + LDS	1272 GB/s	1289 GB/s

尽管使用LDS的内核占用率略有降低，但相比不使用LDS的内核，有效和实际的内存带宽显著提高。事实上，这种性能几乎接近前面讨论的BabelStream性能。接下来，我们调查了内存流量：

Kernel	L1	L2 read	L2 write	HBM	LDS Instructions / wave
y-direction: `VEC_LEN=4`	5369 MB	1758 MB	537 MB	1080 MB	0.0
y-direction: `VEC_LEN=4` + LDS	2147 MB	1611 MB	537 MB	1080 MB	12.0

L1流量降低了一倍以上，L2读取流量略有减少。根据`rocprof`，每个波的LDS指令数量从0增加到12 - 这对应于3条从HBM读取指令和9条计算并将有限差分存储到`out`寄存器的指令。

所有三个方向（x、y和z）上的内核在单个MI250X GCD上获得了可接受的性能。下一部分探讨内核融合，将所有三个单独方向模板内核合并为一个单体内核，以获得进一步的加速。

内核融合

现在所有三个方向的内核都已经高度优化，下一步任务是将它们合并为一个单一的内核，以进一步减少数据移动。进行内核融合的一个挑战是，它通常会增加对可用硬件资源（如寄存器和LDS）的压力。因此，可能需要减少块大小或占用率。在这方面，我们将看到我们可以保持之前的块大小。首先，通过应用向量化和LDS，我们将x方向和y方向的内核合并：

template <int R>
__launch_bounds__((BLOCK_DIM_X) * (BLOCK_DIM_Y))
__global__ void compute_fd_xy_lds_vec_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) {

    const int sj = y0 + threadIdx.y;
    const int lds_y = BLOCK_DIM_Y + 2*R;
    const int i = x0 + VEC_LEN*(threadIdx.x + blockIdx.x * blockDim.x);
    const int j = sj + blockIdx.y * blockDim.y;
    const int k = z0 + threadIdx.z + blockIdx.z * blockDim.z;
    size_t pos = i + line * j + slice * k;
    size_t spos = threadIdx.x + sj * BLOCK_DIM_X;
    size_t line_vec = line >> VEC_EXP;

    p_in += pos;
    p_out += pos;
    float x_reg[2 * XREG + VEC_LEN] = {0.0f};
    __shared__ vec smem[BLOCK_DIM_X * lds_y];

    // Recast as vectorized floats
    const vec *p_in_vec = reinterpret_cast<const vec*>(p_in);
    vec *p_out_vec = reinterpret_cast<vec*>(p_out);
    vec *x_reg_vec = reinterpret_cast<vec*>(x_reg);

    // Read x into registers
    for (int r = 0; r < 2*XREG_VEC+1; ++r)
        x_reg_vec[r] = p_in_vec[0 - XREG_VEC + r];

    // Read y into LDS
    smem[spos] = x_reg_vec[XREG_VEC];
    smem[spos - (BLOCK_DIM_X * R)] = p_in_vec[0 - R*line_vec];
    smem[spos + BLOCK_DIM_X * BLOCK_DIM_Y] = p_in_vec[0 + line_vec * BLOCK_DIM_Y];
    __syncthreads();

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

    // Compute the finite difference approximation in the xy-direction
    vec out = {0.0f};
    for (int r = 0; r <= 2 * R; ++r) {
        out += smem[spos + (r - R) * BLOCK_DIM_X] * d_dy<R>[r];
        for (int ii = 0; ii < VEC_LEN; ++ii)
            out[ii] += x_reg[XREG_OFF + r + ii] * d_dx<R>[r];
    }

    __builtin_nontemporal_store(out,&p_out_vec[0]);

}

另外请注意，我们通过内置函数 __builtin_nontemporal_store 引入了非临时存储，参见 [Finite Difference Laplacian Part 3](Finite difference method - Laplacian part 3 — ROCm Blogs) 获取更多详细信息。下面是与单方向内核相比的性能数据：

Kernel	Effective memory bandwidth	Achieved memory bandwidth
x-direction: `VEC_LEN=4`	1240 GB/s	1280 GB/s
y-direction: `VEC_LEN=4` + LDS	1272 GB/s	1289 GB/s
xy-direction: `VEC_LEN=4` + LDS	1252 GB/s	1292 GB/s

xy方向内核本身与x方向和y方向内核相当。通过这种优化，在数据移动上最多可节省2倍。接下来的挑战是将这个融合的内核与使用滑动窗口技术的z方向内核合并。

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

    const size_t i = (x0 + VEC_LEN * (threadIdx.x + blockIdx.x * blockDim.x));
    const size_t j = y0 + threadIdx.y + blockIdx.y * blockDim.y;

    if (i >= x1) return;

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

    size_t pos = i + line * j + slice * kbegin;
    size_t slice_vec = slice >> VEC_EXP;
    size_t line_vec = line >> VEC_EXP;

    // Shift pointers such that that p_in points to the first value in the sliding window
    p_in += pos - R * slice;
    p_out += pos;
    const vec *p_in_vec = reinterpret_cast<const vec*>(p_in);
    vec *p_out_vec = reinterpret_cast<vec*>(p_out);

     // LDS for y direction
    const int lds_y = BLOCK_DIM_Y + 2*R;
    const int sj = y0 + threadIdx.y;
    size_t spos = threadIdx.x + sj * BLOCK_DIM_X;
    __shared__ vec smem[BLOCK_DIM_X * lds_y];

    // z direction sliding window
    vec w[2 * R + 1];

    // solution register
    vec out[R+1];

    // x direction stencil
    float x_reg[2 * XREG + VEC_LEN];
    vec *x_reg_vec = reinterpret_cast<vec*>(x_reg);

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

        // 2. Load x into registers
        for (int r2 = 0; r2 < 2*XREG_VEC + 1; ++r2)
            x_reg_vec[r2] = p_in_vec[0 - XREG_VEC + r2];

        // 3. Load y into LDS
        __syncthreads();
        {
            smem[spos - (BLOCK_DIM_X * R)] = p_in_vec[0 - R * line_vec];
            smem[spos] = x_reg_vec[XREG_VEC];
            smem[spos + (BLOCK_DIM_X * BLOCK_DIM_Y)] = p_in_vec[0 + line_vec * BLOCK_DIM_Y];
        }
        __syncthreads();

        // 4. Compute xy stencils
        out[r-R] = {0.0f};
        for (int r2 = 0; r2 <= 2 * R; ++r2) {
            out[r-R] += smem[spos + (r2 - R) * BLOCK_DIM_X] * d_dy<R>[r2]; // y-direction
            for (int ii = 0; ii < VEC_LEN; ++ii)
                out[r-R][ii] += x_reg[XREG_OFF + r2 + ii] * d_dx<R>[r2]; // x-direction
        }

        // Prime the z sliding window
        w[r] = x_reg_vec[XREG_VEC];
        p_in_vec += slice_vec;
    }

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

        // 2. Load x into registers
        for (int r2 = 0; r2 < 2*XREG_VEC+1; ++r2)
            x_reg_vec[r2] = p_in_vec[0 - XREG_VEC + r2]; // x - R

        // 3. Load y into LDS
        __syncthreads();
        {
            smem[spos - (BLOCK_DIM_X * R)] = p_in_vec[0 - R * line_vec]; // y - R
            smem[spos] = x_reg_vec[XREG_VEC];
            smem[spos + (BLOCK_DIM_X * BLOCK_DIM_Y)] = p_in_vec[0 + line_vec * BLOCK_DIM_Y]; // y + R
        }
        __syncthreads();

        // 4. Compute xyz stencils
        w[2*R] = x_reg_vec[XREG_VEC];
        out[R] = {0.0f};
        for (int r = 0; r <= 2 * R; ++r) {
            out[0] += w[r] * d_dz<R>[r]; // z-direction
            out[R] += smem[spos + (r - R) * BLOCK_DIM_X] * d_dy<R>[r]; // y-direction
            for (int ii = 0; ii < VEC_LEN; ++ii)
                out[R][ii] += x_reg[XREG_OFF + r + ii] * d_dx<R>[r]; // x-direction
        }

        // 5. Write only if within y boundary
        if (j < y1)
            __builtin_nontemporal_store(out[0],&p_out_vec[0]);

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

        // Increment pointers
        p_in_vec += slice_vec;
        p_out_vec += slice_vec;
    }
}

对于这个最终的单块内核，我们总结了以下几个关键点：
1. 只有在线程超过 x1 时退出。
2. 将启动步骤分成两部分。
3. 第一部分将z方向模板加载到滑动窗口寄存器 w 中。
4. 第二部分继续加载z方向模板，但也计算x和y方向的模板，并将结果存储到临时输出寄存器 out 中。
5. 在滑动窗口阶段，z方向模板计算并与之前存储的xy方向模板结合。
6. 在LDS加载操作之前引入额外的 __syncthreads() 以避免数据冲突。
7. 只有在线程未超过 y1 时才写入结果。

下面是各个 nw 值的结果：

图1：完全组合的xyz内核在实现的和实际的内存带宽上均超过1000 GB/s。随着窗口大小的增加，读写比下降并接近1。
在 nw = 100 和 nw = 200 附近，两个内存带宽FOM都刚刚超过1000 GB/s。还应注意的是，随着 nw 增加，读写比略高于1，并且比在 [Part 2](Seismic stencil codes - part 3 — ROCm Blogs) 中计算的预测值更差。另一个有趣的观察是，与z方向滑动窗口内核不同，当 nw 过小或过大时，实现的和内存带宽都相对较低。
即使在最佳范围内的数值也比单独的方向内核观察到的1200 GB/s要差。然而，执行融合的xyz内核仍然比顺序执行三个单独的优化内核或xy方向和z方向滑动窗口内核要快。

我们可以用这些公式来近似速度提升：

$T_{xyz} = \frac{N}{1000 GB/s}$ , $T_{xy} = \frac{N}{1252 GB/s}$ , $T_{x} = \frac{N}{1240 GB/s}$ , $T_{y} = \frac{N}{1272 GB/s}$ , $T_{z} = \frac{N}{1200 GB/s}$

其中𝑇表示内核执行时间，𝑁表示立方体的大小。通过以下公式近似地计算可以实现的速度提升：

$\frac{T_{x}+T_{y}+T_{z}}{T_{xyz}} = \frac{1000}{1240} + \frac{1000}{1272} + \frac{1000}{1200} = 2.42$

同样，如果我们只希望合并xy方向内核，速度提升为：

$\frac{T_{xy} + T_{z}}{T_{xyz}} = \frac{1000}{1252} + \frac{1000}{1200} = 1.63$

即使是通过这种初步尝试将三个内核合并为单一的单次通过方法，它仍然提供高达1.63倍至2.42倍的速度提升。

总结和下一步计划

在这最后的三篇博客文章中，我们介绍了计算地震模板的基本HIP实现，采用高阶有限差分法进行计算。使用与[有限差分拉普拉斯系列](Seismic stencil codes - part 3 — ROCm Blogs)相同的性能指标，我们开发了近似有效内存带宽的方法，并将其与实际内存带宽进行比较。这些文章集中在以下优化策略上：

1. 对齐内存：此优化使用填充内存分配，使得主维度是GPU缓存行大小的整数倍。
2. 滑动窗口：此技术在寄存器中保留输入数组数据的局部体积，并完全使用这些寄存器计算z方向的模板。此“窗口”有助于消除内核在逐一处理xy平面时冗余的全局内存获取。
3. 向量化：将`float`设备缓冲区重新定义为`float2`或`float4`缓冲区。这项技术增加了在途字节数并减少了全局加载/存储指令的数量。
4. 本地数据共享 (LDS)：利用LDS存储网格数据并计算y方向的模板，减轻了对全局内存（缓存和HBM）的压力。虽然将所有网格数据存储在LDS中而非寄存器中可能有助于减轻寄存器压力，但为每个线程块分配过多的LDS会降低占用率。

所有四种优化都显著提升了单方向内核的性能。此外，结合上述优化的x和y内核的性能几乎与x方向和y方向内核相同。最终的优化将所有三个模板内核融合为一个。最终结果是一个*单次通过*方法，比*三次通过*方法快了近2.5倍。

在下一篇博客文章中，我们将把这一工作扩展到更高阶的有限差分模板，即超出`R=4`，并研究不同网格大小和硬件架构下的性能。如果您有任何问题，请随时在[Github Discussions](ROCm/rocm-blogs · Discussions · GitHub)上与我们联系。

伴随的代码示例

[1](1, 2)测试使用ROCm版本6.1.0-82进行。基准测试结果不是经过验证的性能数据，提供的仅是为了展示代码修改相对性能提升的参考。实际性能结果取决于包括系统配置和环境设置在内的多种因素，结果的再现性无法保证。
[2]巴别流案例研究