地震微分方程代码

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

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

在高性能计算（HPC）领域，地震工作负载一直以来都依赖于结构网格上的高阶有限差分方法。这一趋势至今未变。尽管关于优化技术和策略的文献丰富，目前还没有一种“通用适用”的方法能在各种物理模型中都奏效。离散化和硬件特性的差异是影响传播算法整体性能的主要因素。GPU加速提供了一种具成本效益的方式来满足处理不断增长的地震数据量的巨大需求。与其他基于计算流的领域（例如气候和天气建模）相比，地震工作负载的一个显著特点是普遍使用高阶方法以获得更好的准确性。在地震勘探领域，高阶有限差分法是事实上的标准方法，因为它在数值精度和计算效率之间提供了一个有利的折中。高阶计算流的特点是更宽或更大的计算流半径。像许多计算流模式一样，地震代码中的计算流往往受限于内存或带宽。随着计算流半径的增加，只要能够实现高程度的数据重用，算术强度（浮点指令与内存加载存储指令的比率）也会提高。因此，高阶计算流代码往往会对内存子系统造成显著的压力。虽然没有“一刀切”的方法，但在广泛的地震模型实施中，有几种技术构成了基本构件。在这篇博客文章以及后续的文章中，我们将解释并演示如何在AMD Instinct加速器上实现这些技术，并充分利用硬件的性能来加速地震微分方程代码。

各向同性声波方程

地震模板代码主要涉及不同类型的波动方程的离散化。这些波动方程求解器构成了反射时间迁移（RTM）和全波形反演（FWI）求解器的基本构建块，旨在为油气勘探绘制地球地下结构。控制方程从求解声波方程到在各种介质中求解弹性波方程不等。可以说，研究得最透彻的方程是二阶形式的各向同性恒密度声波方程。该方程用于求解在以变化的声速 $c(x,y,z)$ 参数化的介质中的压力场 $p(x,y,z,t)$ 。

各向同性声波方程为：

$p_{tt} = L(p, c) + s(x, y, z, t), \quad L = c^2 \nabla^2 p, \quad s(x, y, z, t) = \delta(x - x') \delta(y - y') \delta(z - z') g(t)$

并且满足齐次初始条件 $p(x, y, z, 0) = 0$ 。上式中， $\nabla^2 p$ 是压力场的拉普拉斯算子，而 $g(t)$ 是一个源项。该源项用于扰动初始压力场，以便从预定位置（ $x', y', z'$ ）产生波动。随着时间的推移，从源点出发的波会在整个区域传播。

理论上讲，波可以无限传播或继续传播直到消失。然而，由于计算资源的限制，这在模拟特定区域时是不可行的。因此，我们将模型截断到一个感兴趣区域及其边界区域。在计算网格的边界处吸收所有入射能量可以模拟出实际的无限介质。

关于边界条件的注释

强制执行边界条件极为重要。在实际操作中，理想的边界条件是吸收型的，模拟无边界介质。计算域的截断引入了一个人工边界，该边界会产生虚假反射，这可能会污染计算域内部的解。因此，开发了许多吸收边界方法来处理这个问题，它们在准确性和计算成本之间有所折衷；有些方法依赖于物理衰减（如随机边界），有些方法依赖于人工衰减（如海绵边界和完美匹配层（PML））。在本文中，我们关注内部离散化，并将边界处理放在一边，因为只要体积与表面积的比率较大，边界处理的计算影响通常比内部计算要小。

离散化概述

常见的离散化声波方程的方法是通过行方法分两步进行：
1. 时间积分器在时间上是显式的，并使用二阶跃蛙中心差分离散 $p_{tt}$ 项，而右侧在当前时间步 $t_n$ 评价。
2. 空间部分通常使用高阶有限差分离散，通常是第八阶。

稍加滥用符号，得到的时间离散化变为：

$p^{n+1} - 2p^n + p^{n-1} = \Delta t s(p^n, t^n)$

在上述公式中，可以理解为 $p^n = p(x,y,z,t_n)$ 且 $t_n = \Delta t n$ 是第 $n$ 个时间步。

主循环

下节展示了声波方程求解器的主要结构。

for (int n = 0; n < nsteps; ++n) {
     // Time step: t = n\Delta t 
     float t = n * dt;
     // Solves  p^{n+1} = -p^{n-1} + 2p^n + \Delta t L(p^n,t^n)
     solve_acoustic_wave_equation(p_next, p_curr, p_prev, c_squared, dt);
     // Treat snapshots
     // Per each n cycles; store compressed or uncompressed snapshots
     // Adds the discretization of the source term: s(x,y,z,t) evaluated at the grid index: si, sj, sk
     apply_source(p_next, si,sj,sk, t);
     // Cycle pointers to advance the pressure field for each time step
     swap(p_prev , p_curr);   
}

需要注意的是，该实现仅在三个时间点存储压力场。更详细地说，`p_next` 对应 $p^{n+1}$ ，`p_curr` 对应 $p^n$ ，而 p_prev 对应 $p^{n-1}$ 。指针交换技术是一种在推进波场时避免数据移动的简单方法。

空间离散化

在单向等方声波方程中，大部分复杂性来源于对空间部分的高阶离散化。对于该方程，有必要离散化拉普拉斯算子 $\nabla^2$ ，表示在笛卡尔坐标系中，即 $\nabla^2 = \partial ^2_{xx} + \partial ^2_{yy} + \partial ^2_{zz}$ 。

多次计算和单次计算方法

对于空间部分的离散化，有两种方法：多次计算方法和单次计算方法。这两种方法处于光谱的相对两端。在多次计算方法中，空间导数通过一次处理一个方向的方式在多个过程计算。例如，拉普拉斯算子可以通过依次处理x方向、y方向和z方向来离散化。这是一种“三次计算”方法。以下是该技术的伪代码示例，仅演示了拉普拉斯算子：

// split `p_out = \partial_x^2 p_in + \partial_y^2 p_in + \partial_z^2 p_in` into:
// p_out = \partial_x^2 p_in
compute_full_fd_d2<x>(p_out, p_in, ...);
// p_out += \partial_y^2 p_in
compute_full_fd_d2<y>(p_out, p_in, ...);
// p_out += \partial_z^2 p_in
compute_full_fd_d2<z>(p_out, p_in, ...);

在此示例中，可以理解为，函数 compute_full_fd_d2<x>(...) 负责在 x 方向处理所有网格点的有限差分离散化。同样地，在 y 方向和 z 方向上，每一步会逐步更新输出场数组。

这种“三次计算”方法的缺点是需要多次访问主内存，同一个数据要多次读取以执行计算。在这种情况下，`p_in` 需要读取三次，而 p_out 需要读取两次并写入三次。由于地震模板通常受内存带宽限制，这些额外的内存访问代价不菲。当性能十分重要时，建议实现“单次计算”方法。然而，多次计算方法的主要优点是实现简化，因为它一次只处理一个网格方向。如果例如在边界处有限差分离散化发生变化，这一特点尤为明显。在这种情况下，可能需要在边界附近修改计算。结果是每个网格方向会有三种不同的计算。然而，如果在单次过程中处理所有网格方向，就需要进行 3^3 = 27 种不同的计算。

光谱的另一端是“单次计算方法”。顾名思义，单次计算方法使用单个函数一次性计算出整个拉普拉斯算子。

// Compute the entire Laplacian: p_out = \partial_x^2 p_in + \partial_y^2 p_in + \partial_z^2 p_in
compute_full_fd_laplacian(p_out, p_in, ...);

这种技术在我们的博客中用于离散化空间算子，链接为这里。虽然这种方法最小化了访问主内存的次数，但与多次计算方法相比，通常需要更多的 GPU 硬件资源，即寄存器和共享内存。由于单次计算方法将更多的工作集中在一个内核中，这给了编译器更多优化代码的机会，同时手工优化的机会也大大增加。

出于我们的目的，我们将展示一些流行技术，重点讨论多次计算方法。我们选择这种方法是因为它可以简化接下来的展示。然而，需要注意的是，为了获得硬件可以提供的最佳性能，这些技术在大多数情况下都应扩展为单次计算方法。

高阶有限差分近似

下面是应用于栅格点 ( $x_i$ , $y_j$ , $z_k$ ) 的拉普拉斯算子的二阶导数项的一些高阶有限差分近似，对于每个方向的统一栅格间距分别为 $h_x$ 、 $h_y$ 和 $h_z$ :

$\frac{\partial^2 p}{\partial x^2} \approx \sum_{r=-R}^{R} \frac{d_r}{h_x^2} p(x_i + r h_x, y_j, z_k, t),$ $\frac{\partial^2 p}{\partial y^2} \approx \sum_{r=-R}^{R} \frac{d_r}{h_y^2} p(x_i, y_j + r h_y, z_k, t),$ $\frac{\partial^2 p}{\partial z^2} \approx \sum_{r=-R}^{R} \frac{d_r}{h_z^2} p(x_i, y_j, z_k + r h_z, t).$

在上述近似中，模板的半径为 𝑅，模板的宽度为 $2R + 1$ 个栅格点宽度。如果模板是对称的，即 $d_r = d_{-r}$ ，则可以构造出 $2R$ 阶精度的差分近似。对于 $R = 1$ ，对称的二阶精度系数为： $d_{-1} = d_1 = 1$ 和 $d_0 = -2$ 。更高阶的系数可以在[这里]查阅。对于地震应用，𝑅 通常至少为 4（第 8 阶）或更多。

以下示例代码在给定方向上对 nx * ny * nz 的内部栅格点应用高阶有限差分算子。

template <int stride>
__inline__ void compute_d2(float *p_out, const float *p_in, const float *d, const int R, int pos) {
    /* Computes a central high-order finite difference approximation in a given direction
     p_out: Array to write the result to
     p_in: Array to apply the approximation to
     d: Array of length 2 * R + 1 that contains the finite difference approximations, including scaling by grid spacing. 
     R: Stencil radius
     pos: The current grid point to apply the approximation to
     stride: The stride to use for the stencil. This parameter controls the direction in which the stencil is applied.
     */
    float out = 0.0f;
    for (int r = -R; r < R; ++r) 
        out += d[r + R] * p_in[pos + r * stride]; 
    }
    p_out[pos] += out;
}

该示例代码使用单精度浮点，因为这种精度在实际应用中是最广泛使用的选择之一。实现高阶有限差分方法时，方便的是填充数组以留出光晕区域的空间。光晕区域可以由邻居或用于施加边界条件的边界值占据。

// Interior grid size
int nx, ny, nz;
nx=ny=nz=10;
// Padded grid size
int mx = nx + 2 * R;
int my = ny + 2 * R;
int mz = nz + 2 * R;
uint64_t m = mx * my * mz;
float *p_out = new float[m];
float *p_in = new float[m];

由于输入和输出数组被合并到一维中，变量 pos 指定如何通过以下公式线性访问栅格点 ( $x_i$ , $y_j$ , $z_k$ ) 的栅格值：

pos = i + line * j + slice * k;

如果数组的主要维度沿 x 方向排列，则为：

int line = nx + 2 * R;
int slice = line * (ny + 2 * R);

下面的代码展示了高阶有限差分方法的应用。

const float h = 0.1;
const int R = 1;
const float d[2*R + 1] = {1.0 / (h * h), -2.0 / (h * h), 1.0 / (h * h)};

// Define each direction to apply the finite difference approximation in
const int x = 1;
const int y = line;
const int z = stride;

const int line = mx;
const int slice = mx * my;
const uint64_t n = mx * my * mz;

// zero-initialize the memory
memset(p_out, 0, n * sizeof(float));    
    
// Apply approximation in the x-direction for all interior grid points
for (int k = R; k < nz + R; ++k) {
    for (int j = R; j < ny + R; ++j) {
        for (int i = R; i < nx + R; ++i) {
            const uint64_t pos = i + line * j + slice * k;          
            compute_d2<x>(p_out, p_in, d, R, pos, line, slice);
        }
    }
}

通过更改 stride 参数，代码可以很容易地变为特定方向的有限差分近似。

// Apply approximation in the x-direction
compute_d2<x>(p_out, p_in, d, R, pos, line, slice);
// Apply approximation in the y-direction
compute_d2<y>(p_out, p_in, d, R, pos, line, slice);
// Apply approximation in the z-direction
compute_d2<z>(p_out, p_in, d, R, pos, line, slice);

基线内核

现在是时候开始研究在 GPU 上实现高阶有限差分近似了。作为起点，让我们考虑一个简单的内核，用于在 x 方向上应用有限差分方法。这个内核将用于将来的所有性能比较。

  // Table containing finite difference coefficients
  template <int R> __constant__ float d_dx[2 * R + 1];
  template <int R> __constant__ float d_dy[2 * R + 1];
  template <int R> __constant__ float d_dz[2 * R + 1];

  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;

  }

上述内核可以通过修改计算步长来计算 y 方向和 z 方向上的模板，其中 int stride = line; 和 int stride = slice; 分别对应 y 方向和 z 方向。该内核使用 x0,y0,z0 和 x1,y1,z1 来定义感兴趣的区域以应用模板。以这种方式定义边界很方便，因为它可以明确控制模板的应用位置。如果需要重叠通信和计算，可以多次调用相同的内核，参数对应于不同的 Halo 切片。为了限制所有内部点的注意范围，边界需要设置为 x0=R, x1=nx+R, y0=R, y1=ny+R, z0=R, z1=nz+R。在内核中，索引 i、`j`、`k` 被 x0、`y0`、`z0` 所偏移。此偏移节省了三个比较指令，因为它消除了对下边界的检查，同时避免了一些线程在下边界外空闲。另一方面，x 方向的偏移可能会由于内存访问未对齐而对性能产生负面影响。

初步性能评估

让我们快速评估这些基线核函数的性能，使用两种不同的评价标准（FOM）。首先，像在[Laplacian 系列](Seismic stencil codes - part 1 — ROCm Blogs)中一样，我们从模板算法设计了解到，每个网格点只需要加载一次，并且可以被邻近的线程重复使用。为此，我们将使用_有效内存带宽_，其定义如下：

effective_memory_bandwidth = (theoretical_fetch_size + theoretical_write_size) / average_kernel_execution_time;

理论上的提取和写入大小（以字节为单位）对于一个`nx * ny * nz`的立方体是：

theoretical_fetch_size = (nx * ny * nz + 2 * R * PLANE) * sizeof(float);
theoretical_write_size = (nx * ny * nz) * sizeof(float);

其中`PLANE`分别为x方向、y方向和z方向的`ny * nz`、`nx * nz`和`ny * nz`。理想情况下，我们希望这个FOM尽可能接近可实现的内存带宽。在一个单独的MI250X GCD上，以下是对于不同的`R`值，512 x 512 x 512立方体的有效内存带宽[1]：

Radius ®	x-direction	y-direction	z-direction
1	971 GB/s	948 GB/s	915 GB/s
2	995 GB/s	982 GB/s	753 GB/s
3	977 GB/s	965 GB/s	394 GB/s
4	956 GB/s	940 GB/s	217 GB/s

较大的模板半径和较大的内存跨度会恶化带宽。所有这些数值，包括较小的模板半径和跨度，都明显低于单个MI250X GCD可实现的带宽，根据[BabelStream 案例研究](https://www.olcf.ornl.gov/wp-content/uploads/2-16-23-node_performance.pdf)，大约为1.3 TB/s到1.4 TB/s。然而，单凭这些FOM不足以深入了解我们的GPU实现对硬件的利用程度。

因此，我们提供了可通过两种不同方法实现的_实测内存带宽_。第一种方法是直接使用`rocprof`并将`FETCH_SIZE`和`WRITE_SIZE`指标之和除以平均核函数执行时间。第二种方法是使用[omniperf](https://amdresearch.github.io/omniperf/)并取报告的`L2-EA Rd BW`和`L2-EA Wr BW`之和。理想情况下，我们也希望这个数值尽可能接近可实现的内存带宽。以下是相应的实测内存带宽数值[1]：

Radius	x-direction	y-direction	z-direction
1	976 GB/s	953 GB/s	931 GB/s
2	1003 GB/s	995 GB/s	953 GB/s
3	986 GB/s	981 GB/s	925 GB/s
4	967 GB/s	959 GB/s	891 GB/s

x方向和y方向的实测内存带宽数值和报告的有效内存带宽数值紧密对齐。当有效和实测内存带宽数值一致时，这表明移动的数据量（即分子）是一样的。两个带宽度量使用相同的核函数执行时间，所以如果两个带宽度量都较低，这表明核函数的部分是计算密集型而未完全覆盖内存传输和/或存在潜在的延迟问题。

然而，z方向的实测内存带宽显示更高，这表明虽然硬件利用率可能比x和y方向的实现更好，但核函数从全局内存提取和/或写入的次数远远超出理想情况。

结论

这就结束了开发地震模板代码的第一部分，这些代码通常依赖于使用高阶有限差分方法对波动方程进行离散化。我们首先从三通道方法开始，并展示了不同模板半径和方向的初始性能数据。没有任何实现方案在有效或达到的内存带宽方面表现出令人满意的性能。在接下来的系列文章中，我们将深入探讨可能的优化措施，以提升内存带宽数据。
[伴随的代码示例]

如果您有任何问题或评论，请通过GitHub上的[讨论](amd/amd-lab-notes · Discussions · GitHub)联系我们。

[1]
(1, 2)
测试使用ROCm版本6.1.0-82进行。基准测试结果并非验证的性能数据，仅用于展示代码修改后的相对性能提升。实际性能结果取决于多个因素，包括系统配置和环境设置，结果的可重复性无法保证。