有限差分法 - 拉普拉斯算子 (Part 1)

Finite difference method - Laplacian part 1 — ROCm Blogs (amd.com)

2022年11月14日, Justin Chang, Rajat Arora, Thomas Gibson, Sean Miller, Ossian O’Reilly撰写。

有限差分法是一种在计算物理中常用的网格离散化方法，广泛应用于从地球物理（天气和石油&天然气）和电磁学（半导体和天体物理学）到气体动力学（气流和等离子体）的各种应用中。Stencil代码的特点是需要访问网格点的本地邻域（即Stencil），以便在单个网格点上计算一个值，这意味着算法的性能与内存访问模式密切相关。由于算法结构的规则性、高度暴露的并行性和GPU高内存带宽的高效利用，GPU在加速Stencil代码方面表现出色。在这篇博客文章中，我们将使用AMD的异构接口可移植性（HIP）API来开发一个GPU加速的Stencil代码，这将使我们可以显式控制内存访问模式。在这个开发过程中，我们将讨论以下概念：

1. 使用有限差分法离散化拉普拉斯算子
2. 提供一个简单的HIP实现版本
3. 指导读者选择优化策略以提高Stencil代码的性能
4. 演示如何使用`rocprof`获取内核的性能指标

Laplacian

拉普拉斯算子 ∇² 是一种微分算子，可以通过偏微分方程描述各种物理现象。例如，它可以用来描述声波的传播（波动方程）、热流（热方程），以及静电和重力势（泊松方程）。它在计算机视觉和网格生成中也有应用。更多信息请参见 [维基百科]。

作为拉普拉斯算子可以描述的物理现象的一个例子，下面的动画展示了一个二维热方程的数值模拟，显示了热量从热区（凸起的字母）向冷区（平面）的传递过程。

初始条件

模拟

../../../_images/heat_eq.gif

在笛卡尔坐标系中，拉普拉斯算子表示为一个标量场 u(x, y, z) 的梯度的散度：

$\nabla \cdot \nabla u = \nabla^2 u = \frac{\partial^2 u}{\partial x^2} + \frac{\partial^2 u}{\partial y^2} + \frac{\partial^2 u}{\partial z^2}$

其中，u 是空间坐标 x、y 和 z 的适当光滑函数。

离散化和主机实现

为了离散化拉普拉斯算子，我们引入了一个有 $n_{x}$ × $n_{y}$ × $n_{z}$ 个等间距网格点的网格，并在每个网格方向上应用二阶有限差分近似。例如，x方向的导数可以近似为：

$\frac{\partial^2 u}{\partial x^2} \approx \frac{u(x-h_x,y,z) - 2u(x,y,z) + u(x+h_x,y,z)}{h_x^2}$

其中 $h_x$ 是 x 方向的网格间距（即 x 方向上相邻网格点之间的距离）。将其扩展到三维空间需要在所有空间方向上应用这个模板：

$\nabla^2 u(x,y,z) \approx \frac{u(x-h_x,y,z) - 2u(x,y,z) + u(x+h_x,y,z)}{h_x^2} + \frac{u(x,y-h_y,z) - 2u(x,y,z) + u(x,y+h_y,z)}{h_y^2} + \frac{u(x,y,z-h_z) - 2u(x,y,z) + u(x,y,z+h_z)}{h_z^2}$

下面的代码示例演示了如何在主机代码中实现拉普拉斯模板：

template <typename T>
void laplacian_host(T *h_f, const T *h_u, T hx, T hy, T, hz, int nx, int ny,
int nz) {
#define u(i, j, k) h_u[(i) + (j) * nx  + (k) * nx * ny]
#define f(i, j, k) h_f[(i) + (j) * nx  + (k) * nx * ny]// Skip boundary pointsfor (k = 1; k < nz - 1; ++k)for (j = 1; j < ny - 1; ++j)for (i = 1; i < nx - 1; ++i) {// Apply stencil approximation of Laplacian to all interior pointsu_xx = ( u(i + 1, j, k) - 2 * u(i, j, k) + u(i - 1, j, k) ) / (hx * hx);u_yy = ( u(i, j + 1, k) - 2 * u(i, j, k) + u(i, j - 1, k) ) / (hy * hy);u_zz = ( u(i, j, k + 1) - 2 * u(i, j, k) + u(i, j, k - 1) ) / (hz * hz);f(i, j, k) = u_xx + u_yy + u_zz;}
#undef u
#undef f
}

主机代码使用 C 预处理宏更容易地将数学符号 $u(x,y,z)$ 转换为代码：u(i, j, k) - 你很快就会学习到更多关于这种转换的内容。函数 laplacian_host 接受输入 h_u[] 和输出 h_f[]，并期望它们至少包含 nx * ny * nz 个元素（每个方向上的网格点数）。为了防止越界访问的内存访问错误，循环边界排除了每个网格方向上的第一个和最后一个元素。最内层的循环对 h_u[] 应用拉普拉斯模板，并将结果存储在 h_f[] 中。为了区分主机数组和设备数组，我们分别使用前缀 h_（主机）和 d_（设备）。当我们访问设备代码时，这个区分非常重要。

数据布局

3D数据的布局方式是在内存中 i 方向上的网格点是连续存储的，而 k 方向上的网格点是通过 nx * ny 间隔的。这种映射通过宏来实现：

#define u(i, j, k) h_u[(i) + (j) * nx  + (k) * nx * ny]

括号围绕 i，j 和 k，确保表达式如 u(i, j + 1, k) 等的正确扩展。

下图展示了在网格点 pos 处应用拉普拉斯算子的模板：

图1：三维空间中的有限差分模板。

在存储布局中，由于步长存储布局，相邻 j 和 k 方向的网格点在内存中实际上相距很远。例如，考虑 nx = ny = nz = 5 的情况：

图2：用于中心差分模板的步长内存空间。这里我们可以看到 h_u[] 数组（红色）中的哪些条目被用来更新输出数组 h_f[]。注意，这些条目不是连续的。

本文的目标之一是展示如何在 AMD GPU 上高效地处理这种内存访问模式。

HIP 实现

首先，考虑一个大立方体，其中 nx = ny = nz = 512。我们先使用 double 精度为输入数组 d_u 和输出数组 d_f 分配内存：

// 参数
using precision = double;
size_t nx = 512, ny = 512, nz = 512; // Cube dimensions// 网格点间距
precision hx = 1.0 / (nx - 1), hy = 1.0 / (ny - 1), hz = 1.0 / (nz - 1);// 输入和输出数组
precision *d_u, *d_f;// 在设备上分配内存
size_t numbytes = nx * ny * nz * sizeof(precision);
hipMalloc((void**)&d_u, numbytes);
hipMalloc((void**)&d_f, numbytes);

d_u 输入数组使用一个二次测试函数初始化，以简化验证正确性的任务（为了简洁未展示）。

下面的代码片段展示了拉普拉斯的初始 HIP 实现：

template <typename T>
__global__ void laplacian_kernel(T * f, const T * u, int nx, int ny, int nz, T invhx2, T invhy2, T invhz2, T invhxyz2) {int i = threadIdx.x + blockIdx.x * blockDim.x;int j = threadIdx.y + blockIdx.y * blockDim.y;int k = threadIdx.z + blockIdx.z * blockDim.z;// 如果该线程在边界上，则退出if (i == 0 || i >= nx - 1 ||j == 0 || j >= ny - 1 ||k == 0 || k >= nz - 1)return;const int slice = nx * ny;size_t pos = i + nx * j + slice * k;// 计算模板操作的结果f[pos] = u[pos] * invhxyz2+ (u[pos - 1]     + u[pos + 1]) * invhx2+ (u[pos - nx]    + u[pos + nx]) * invhy2+ (u[pos - slice] + u[pos + slice]) * invhz2;
}template <typename T>
void laplacian(T *d_f, T *d_u, int nx, int ny, int nz, int BLK_X, int BLK_Y, int BLK_Z, T hx, T hy, T hz) {dim3 block(BLK_X, BLK_Y, BLK_Z);dim3 grid((nx - 1) / block.x + 1, (ny - 1) / block.y + 1, (nz - 1) / block.z + 1);T invhx2 = (T)1./hx/hx;T invhy2 = (T)1./hy/hy;T invhz2 = (T)1./hz/hz;T invhxyz2 = -2. * (invhx2 + invhy2 + invhz2);laplacian_kernel<<<grid, block>>>(d_f, d_u, nx, ny, nz, invhx2, invhy2, invhz2, invhxyz2);
}

laplacian 主函数负责启动设备内核 laplacian_kernel，线程块大小为 BLK_X = 256, BLK_Y = 1, BLK_Z = 1。当前线程块大小 BLK_X * BLK_Y * BLK_Z = 256 是推荐的默认选择，因为它适配硬件工作调度，但其他选择可能表现更好。一般来说，为了性能选择这个参数可能高度依赖于问题规模、内核实现以及设备的硬件特性。我们将在未来的帖子中更详细地讨论这个主题。尽管这个内核实现支持 3D 线程块大小，GPU 并没有 2D 和 3D 的概念。定义 3D 组的可能性是为程序员提供的方便功能。

HIP 实现适用于任何能够适应 GPU 内存的问题大小。对无法被线程块大小均匀分割的问题大小，我们引入了附加的工作组来处理余下的部分。这些附加的线程块在余数较小的情况下可能具有低线程利用率，即许多线程映射到计算域外的区域没有工作要做。这种低线程利用率会对性能产生负面影响。就本文而言，我们将重点关注能被线程块大小均匀分割的问题大小。要调度的块数由 (nx - 1) / block.x + 1, ... 的上取整操作决定。正如之前解释的那样，它将组数向上取整以保证覆盖整个计算域。

在内核中的退出条件也需要注意。我们在线程位于或超出边界时返回（退出），因为我们只计算内部点的有限差分。内核退出策略的实际工作方式是，由于波前（wavefront）以锁步方式执行指令，任何在或超过边界的线程将在内核的其余部分被屏蔽掉。被屏蔽的线程在执行指令时将处于空闲状态。让我们回到块大小不能均匀分割问题大小的性能问题。如果我们选择 nx=258，在 x 方向上会得到两个线程块，但第二线程块中只有一个线程有工作要做。

主要计算公式如下所示，

    // 计算模板操作的结果f[pos] = u[pos] * invhxyz2+ (u[pos - 1]     + u[pos + 1]) * invhx2+ (u[pos - nx]    + u[pos + nx]) * invhy2+ (u[pos - slice] + u[pos + slice]) * invhz2;

在这段代码中，多次访问数组 u。对于每次访问，编译器会生成一个全局加载指令，该指令将从全局内存中加载数据，除非这些数据已经存在于缓存中。正如我们将看到的，这些全局加载指令可能对性能产生显著影响。

性能优化建议
为了优化性能，可以考虑以下几点：
1. 内存访问模式：
- 共线：确保内存访问共线，例如多个线程访问连续的内存位置，这样可以充分利用内存带宽。
- 缓存友好：尽可能让数据在缓存中有一个高的命中率，通过适当地调整数据结构和访问模式，可以减少全局内存访问延迟。
2. 线程块大小选择：
- 根据硬件特性和具体问题规模选择最优的线程块大小，一般来说，线程块大小会选择为系统支持的最大值，以充分利用 GPU 计算能力。
3. 边界处理：
- 可以将边界处理与内核计算分离，将边界线程的计算单独提取出来，这样可以避免在内核中频繁检查边界条件。
优化后的模板实现
优化后的模板可能会类似于以下代码：
template <typename T>
__global__ void laplacian_kernel_optimized(T * f, const T * u, int nx, int ny, int nz, T invhx2, T invhy2, T invhz2, T invhxyz2) {int i = threadIdx.x + blockIdx.x * blockDim.x;int j = threadIdx.y + blockIdx.y * blockDim.y;int k = threadIdx.z + blockIdx.z * blockDim.z;// 使用共享内存减少全局内存访问extern __shared__ T shared_u[];int lx = threadIdx.x, ly = threadIdx.y, lz = threadIdx.z;int local_index = lx + blockDim.x * ly + blockDim.x * blockDim.y * lz;size_t global_index = i + nx * j + nx * ny * k;// 只复制非边界部分到共享内存if(i < nx && j < ny && k < nz)shared_u[local_index] = u[global_index];__syncthreads();// 边界不计算梯度if (i == 0 || i >= nx - 1 || j == 0 || j >= ny - 1 || k == 0 || k >= nz - 1)return;// 计算模板操作的结果f[global_index] = shared_u[local_index] * invhxyz2+ (shared_u[local_index - 1]     + shared_u[local_index + 1]) * invhx2+ (shared_u[local_index - nx]    + shared_u[local_index + nx]) * invhy2+ (shared_u[local_index - nx * ny] + shared_u[local_index + nx * ny]) * invhz2;
}template <typename T>
void laplacian_optimized(T *d_f, T *d_u, int nx, int ny, int nz, int BLK_X, int BLK_Y, int BLK_Z, T hx, T hy, T hz) {dim3 block(BLK_X, BLK_Y, BLK_Z);dim3 grid((nx - 1) / block.x + 1, (ny - 1) / block.y + 1, (nz - 1) / block.z + 1);T invhx2 = (T)1./hx/hx;T invhy2 = (T)1./hy/hy;T invhz2 = (T)1./hz/hz;T invhxyz2 = -2. * (invhx2 + invhy2 + invhz2);size_t shared_memory_size = (BLK_X + 2) * (BLK_Y + 2) * (BLK_Z + 2) * sizeof(T);laplacian_kernel_optimized<<<grid, block, shared_memory_size>>>(d_f, d_u, nx, ny, nz, invhx2, invhy2, invhz2, invhxyz2);
}
在这个优化后的实现中，我们引入了共享内存来减少全局内存访问，同时通过适当地选择线程块大小和处理边界条件的方式来提高性能。这个版本的代码可以根据具体的硬件特性和问题规模做进一步调整。

计算绑定 vs 内存绑定

在分析我们的HIP内核性能时，理解其是计算绑定还是内存绑定是有帮助的。计算绑定的内核受计算单元吞吐量限制，而内存绑定的内核则受内存带宽或内存访问延迟限制。要对我们的内核进行分类，我们可以通过以下公式来研究其算术强度（AI）：

$\text{AI} = \frac{\text{FLOPs}}{\text{BYTES}}$

其中，FLOPs指的是浮点运算的总数，而BYTES指的是主存储器的流量。如果AI $\le\ \frac{\text{PF}}{\text{BW}}$，那么算法被认为是内存绑定的，其中 PF代表峰值浮点运算速率，BW代表峰值内存带宽。相反，如果AI $\ge\ \frac{\text{PF}}{\text{BW}}$，算法就被认为是计算绑定的。请注意，我们忽略了整型运算，只专注于浮点运算。这种对浮点运算和内存流量的分析通常用于创建屋顶线模型。读者若想了解更多关于屋顶线模型的信息，可以参考[维基百科上的条目]及其附带的参考资料。

在上述的 laplacian_kernel 中，每个线程对 d_u 元素执行10次浮点运算，并将计算结果存储到一个 d_f 元素中。然而，我们将假设理想条件下的无限缓存，这意味着每个线程只需获取一个 d_u 元素，并且其他线程将该元素用于模板计算。每个线程的 FLOPs 和 BYTES 分别等于 10 和 2 * sizeof(double) = 16，所以我们的AI是 0.63。像这样的模板代码通常有较低的 FLOPs 和 BYTES 比率，当我们取消无限缓存假设时，这个比率会更低。

在这篇博文的剩余部分，我们考虑一个来自[MI250X GPU] 的单个图形计算单元（GCD），其中 PF=24 TFLOPs/秒，BW=1.6 TBytes/秒的高带宽内存（HBM）带宽。PF/BW 比率是我们估计的 AI 的 24 倍之多，明显将我们的HIP内核归类为内存绑定而不是计算绑定。因此，我们的优化将重点放在解决内存瓶颈上。

性能评估

为了评估性能，我们可以将有效内存带宽作为我们的衡量标准（FOM）。从广义上讲，有效内存带宽是考虑整个内存子系统（从具有高延迟的芯片外高带宽内存（HBM）到低延迟寄存器）中必须移动的最少数据量后的平均数据传输率。有效内存带宽的定义是：

effective_memory_bandwidth = (theoretical_fetch_size + theoretical_write_size) / average_kernel_execution_time;

为了测量平均内核执行时间，我们使用函数`hipEventElapsedTime()`（有关完整实现细节，请参阅附带的`laplacian.cpp`）。计算`nx * ny * nz`立方体的理论读取和写入大小（以字节为单位）的公式是：

theoretical_fetch_size = ((nx * ny * nz - 8 - 4 * (nx - 2) - 4 * (ny - 2) - 4 * (nz - 2) ) * sizeof(double);
theoretical_write_size = ((nx - 2) * (ny - 2) * (nz - 2)) * sizeof(double);

由于模板的形状，理论读取大小包括读取立方体中`d_u`的所有网格点，除了8个角和12个边。在理论写入大小中，只有`d_f`中的内部网格点会被写回内存。这些计算忽略了内存访问的粒度，同样假设存在无限缓存。因此，它们给出了实际数据移动的下限。

以下是单个MI250X GCD上的性能数据：

$ ./laplacian_dp_kernel1
Kernel: 1
Precision: double
nx,ny,nz = 512, 512, 512
block sizes = 256, 1, 1
Laplacian kernel took: 2.64172 ms, effective memory bandwidth: 808.148 GB/s

对于一个优化后的内核和足够大的工作负载，有能力饱和设备时，有效内存带宽应尽可能接近理论峰值HBM带宽。实现这一目标需要满足两个目标：

1. 将过量的内存流量减少到接近理论极限，
2. 饱和设备的HBM带宽。

鉴于MI250X每个GCD的理论带宽为`1638.4 GB/s`，基线内核实现了理论峰值的50%左右[1]。是什么瓶颈阻碍了进一步的性能？是由于数据移动过多还是内存带宽饱和度不理想？例如，如果计算模板时没有良好的数据重用，可能会造成过量的数据移动。也许一些邻近值被缓存带入，但被其他值逐出，导致缓存未命中和额外的主内存访问。当内核生存期内没有足够的内存请求在进行时，内存带宽不能得到饱和。通常，有很多原因可能导致这个问题。也许每个波请求的数据很少，并且由于资源限制，无法同时在每个计算单元（CU）上运行足够多的波。另一种可能是某些波被阻塞，等待数据依赖，无法发出内存指令。要了解瓶颈并解决它，我们需要收集性能计数器数据。

为此，我们使用AMD ROCm™ profiler（`rocprof`）来捕获三个重要指标：

FETCH_SIZE : The total kilobytes fetched from the video memory. This ismeasured with all extra fetches and any cache or memory effectstaken into account.
WRITE_SIZE : The total kilobytes written to the video memory. This is measuredwith all extra fetches and any cache or memory effects takeninto account.
L2CacheHit : The percentage of fetch, write, atomic, and other instructions thathit the data in L2 cache. Value range: 0% (no hit) to 100% (optimal).

请注意，`FETCH_SIZE`和`WRITE_SIZE`的总和可能并不完全代表所有内存流量。尽管如此，这些指标无疑可以为HIP内核的设计提供更深入的洞见。可以从`rocprof`中查询更多的指标（参见`rocprof --list-basic`和`rocprof --list-derived`以了解选项）。`rocprof`允许我们在ascii文本文件中组织我们的请求。例如，在附带的工作中，我们使用`rocprof_input.txt`：

pmc: FETCH_SIZE
pmc: WRITE_SIZE
pmc: L2CacheHit
kernel: laplacian_kernel

这将捕获所列内核的性能计数器（pmc）`FETCH_SIZE`、`WRITE_SIZE`和`L2CacheHit`。我们可以使用以下命令对可执行文件运行此命令：

rocprof -i rocprof_input.txt -o rocprof_output.csv laplacian_dp_kernel

这将生成一个包含每次内核启动请求的指标的`rocprof_output.csv`文件。有关`rocprof`的更多信息，我们将所有感兴趣的读者转至ROCm™ Profiler 文档。

FETCH_SIZE和`WRITE_SIZE`越接近`theoretical_fetch_size`和`theoretical_write_size`，我们的HIP内核性能可能就越好。`L2CacheHit`越接近100%也是表明内核优化的另一个指标。下表比较了上述三个`rocprof`指标：

	FETCH_SIZE (GB)	WRITE_SIZE (GB)	Fetch efficiency (%)	Write efficiency (%)	L2CacheHit (%)
Theoretical	1.074	1.061	-	-	-
Kernel 1	2.014	1.064	53.3	99.7	65.0

65%的`L2CacheHit`率本身并不能提供太多洞见。虽然`WRITE_SIZE`指标与其理论估计值相匹配，但`FETCH_SIZE`几乎翻倍了。这一观察告诉我们，大约50%的模板被重用。如果我们能够将模板重用提高到100%，`FETCH_SIZE`将减少超过0.9 GB，从而使总内存流量减少31%。假设HBM带宽饱和保持不变，这种减少可能会转化为内核执行时间的约1.44倍的加速。因此，有效内存带宽的更现实目标将是： 808.148 GB/s * 1.44 = 1165 GB/s, 这大约是每个GCD理论峰值HBM带宽的71%[1]。

优化有限差分内核有几种不同的方法。在下一篇文章中，我们将解释和讨论一些可以应用于减少过量内存流量并达到新的有效内存带宽目标的优化。

结论

这篇文章是使用HIP开发和优化Laplacian内核的第一部分。在开始任何优化工作之前，收集性能测量数据以建立基准并识别性能瓶颈是非常重要的。在这项工作中，我们确定了基准内核由于其算术强度而受内存带宽的限制。我们还发现，与我们的估计相比，内核从全局内存空间加载的数据量显著增加。这些条件告诉我们，减少多余的内存加载将对性能产生影响。在该系列的第二部分中，我们将介绍优化技术，旨在减少全局内存数据移动，这些技术不仅对模板内核有价值，对使用规则、非连续内存访问模式的其他带宽受限内核同样适用。

伴随的代码示例

如有任何问题或评论，请在GitHub上联系我们进行讨论： Discussions

[1](1,2)

测试使用的为ROCm版本5.3.0-63。基准测试结果并非验证的性能数据，仅用于展示代码修改后的相对性能改进。实际性能结果取决于多个因素，包括系统配置和环境设置，不能保证结果的可重复性。