SPH 粒子领域搜索

问题描述

在光滑粒子流体动力学（Smoothed Particle Hydrodynamics，简称 SPH）方法中，“邻居搜索”（Neighbor Search）是一个至关重要的计算任务。它直接影响到模拟的效率与精度，尤其在涉及上百万粒子的三维问题中，计算瓶颈往往出现在如何高效查找每个粒子周围的邻居上。

由于 SPH 是一种无网格方法，天然适合处理自由表面、大变形和断裂等复杂物理现象，因此被广泛应用于天体物理、流体力学、固体力学等领域。我主要从事小行星撞击问题的数值模拟研究，并开发了一套针对该问题的 SPH 代码。虽然不同领域在实现细节上存在差异，但“邻居搜索”作为核心模块，其数学模型相对简单，却几乎在所有 SPH 实现中都不可或缺。

本文旨在系统介绍 SPH 中邻居搜索的常见算法、性能优化方法，以及在 CUDA 等并行计算平台上的实现，作为我个人学习与研究的记录，同时希望对同领域的研究者提供参考。

首先给出 SPH 粒子邻居搜索的抽象数学问题：设在三维空间中存在 $ N $ 个粒子，每个粒子的位置已知，记为 $ \mathbf{x}_i $，每个粒子有一个核长度（smoothing length）$ h_i $。我们需要找到每个粒子 $ i $ 的邻居粒子集合 $ j $，使得满足以下条件：

\[ \|\mathbf{x}_i - \mathbf{x}_j\| < f(h_i, h_j) \]

其中，函数 $ f(h_i, h_j) $ 用于定义粒子间的交互距离，其常见定义包括：

$ f(h_i, h_j) = \eta \cdot \frac{1}{2}(h_i + h_j) $
$ f(h_i, h_j) = \eta \cdot \min(h_i, h_j) $
$ f(h_i, h_j) = \eta \cdot \max(h_i, h_j) $

这里，$ \eta $ 是一个无量纲系数，称为核支持半径因子（kernel support radius factor），通常取值在 $ [1.2, 2.5] $ 之间，用于控制粒子的影响范围（本文取2）。

在邻居搜索中，常见的方法包括：

暴力搜索（Brute Force）：简单可靠，但时间复杂度高（$O(N^2)$），在大规模问题中效率低下；
链表搜索（Linked-Cell/Grid-based）：通过空间划分显著减少搜索粒子数，是实际中常用的高效方法；
树搜索（如 Octree 或 KD-tree）：适合不均匀粒子分布，尤其适用于天体物理模拟中的自适应精度问题，良好的树代码具有非常高的鲁棒性。如果要考虑自引力，那么树搜索是效率和精度都不错的选择。

本文将以暴力搜索作为结果基线，重点介绍如何高效实现链表搜索和树搜索，并比较它们在实际 SPH 模拟中的性能表现。

1. 链表搜索（Linked-Cell/Grid-based）

在光滑长度为空间常量的情况下，也即所有粒子的$ h_i $都相等，应用链表搜索法非常有效。Monaghan 和 Gingold(1983)提出，可以通过对粒子的空间区域划分网格，记录每个网格内的粒子编号。这样在搜索粒子的邻居时，只需要遍历当前粒子网格的邻居网格内的粒子即可。此方法在传统SPH代码中使用非常多，如Monaghan(1985)，Rhoades(1992)，Simpson(1995)等。

在实现链表算法时，要在问题域上铺设一临时网格。网格单元的空间大小应选取与支持域的空间大小一致。若光滑函数支持域的计量尺度为 $ \eta h $，则网格单元的尺度也必须设置为 $ \eta h $。那么，对于给定的粒子i，其相邻粒子只能在同一网格单元内，或者在紧密相邻的单元内。所以，当 $ \eta = 2 $ 时，在一维、二维和三维空间里的搜索范围分别是在 3,9,27 个单元内。链表搜索法将每个粒子都分布在网格单元内，并通过简单的存储规则将每个网格内的所有粒子连接起来。若每个单元内的平均粒子数量足够小，则链表搜索法的复杂度阶数为 $ O(N) $。

链表搜索法存在的问题是，当光滑长度可变时，尤其是模拟分辨率变化的问题时，网格空间就不能适应每一个粒子，此时若再应用链表搜索法，则搜索效率会很低。除此之外，该方法在CUDA上实现时，需要对显存分配进行小心处理，不然很容易占用超大显存（主要在存储网格粒子编号时）。下面首先就光滑长度为空间常量的情况下进行代码说明。

光滑长度为空间常量的链表搜索

我们先来看链表搜索算法的基础版本。实现该算法需要两个步骤：

记录每个网格单元中包含哪些粒子，并记录每个粒子所属的单元；
遍历所有粒子，对每个粒子的单元及其邻接单元进行扫描，查找可能的邻居粒子。

为了提高效率，本文只展示核函数的编写，且不在每次调用中反复申请或释放内存。

第一个核函数较为简单，其用于构建粒子与网格单元之间的映射关系。

核函数声明与粒子网格索引计算

下面是用于建立粒子与网格单元之间映射关系的 CUDA 核函数 particleLoop，以及配套的粒子网格索引计算函数 particleGridIndex。本版本假设所有粒子的核长度 $h$ 是一个常量，记为 h[0]。

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
#define DIM 3                    // 空间维度：可设为 1, 2 或 3
#define MAX_PARTICLES_PER_GRID  200  // 每个网格单元最多可容纳的粒子数
#define eta 2  // 核长度比例因子

// CUDA 核函数：构建每个粒子所属网格，以及每个网格中的粒子列表
__global__ void particleLoop(
    const double *x,              // 粒子 x 坐标，长度为 numParticles
    const double minx,            // x 方向最小坐标
    const int nx,                 // x 方向网格数 = ceil((maxx - minx) / (η * h))
#if DIM > 1
    const double *y,              // 粒子 y 坐标
    const double miny,
    const int ny,
#endif
#if DIM > 2 
    const double *z,              // 粒子 z 坐标
    const double minz,
    const int nz,
#endif
    const double *h,              // 每个粒子的核长度（此版本为常量）
    int *gridParticlesList,       // 网格中粒子索引列表，大小为 numGrids * MAX_PARTICLES_PER_GRID 无须初始化
    int *gridWritingPointer,      // 为避免数据竞争，使用一个gridWritingPointer来记录写入位置，长度为numGrids，需要初始化为0
    int *particleGridList,        // 每个粒子所属网格索引，长度为 numParticles
    int numParticles,             // 粒子总数
    int numGrids                  // 网格总数
) {
    int tid = threadIdx.x + blockIdx.x * blockDim.x;
    if (tid >= numParticles) return;

    // 计算粒子在网格中的索引坐标
    int ix = int((x[tid] - minx) / (h[0])*eta);  // η = 2
#if DIM > 1
    int iy = int((y[tid] - miny) / (h[0])*eta);
#else
    int iy = 0;
#endif
#if DIM > 2
    int iz = int((z[tid] - minz) / (h[0])*eta);
#else
    int iz = 0;
#endif

    // 获取粒子所在网格的线性索引
    int gridIndex;
#if DIM == 1
    gridIndex = ix;
#elif DIM == 2
    gridIndex = ix + iy * nx;
#else  // DIM == 3
    gridIndex = ix + iy * nx + iz * nx * ny;
#endif
    particleGridList[tid] = gridIndex;
    
    // 获取当前网格的gridWritingPointer位置
    int writingPointer = atomicAdd(&gridWritingPointer[gridIndex], 1);  // atomicAdd返回写入位置
    gridParticlesList[gridIndex * MAX_PARTICLES_PER_GRID + writingPointer] = tid;
}

使用上述函数时要注意，当网格个数很多时，容易超过int的表达范围。此外，前置条件是需要知道最小和最大的粒子坐标，长数组求最大最小值可以使用scan算法，以后有空也会介绍此类算法。

至此，我们已经建立了空间粒子与网格之间的联系，现在可以遍历粒子，获取他们的邻居粒子。我目前想到两种遍历方法，1. 按粒子顺序遍历（下面称A1） 2. 按网格顺序遍历（下面称A1）。

搜索

按粒子顺序遍历(A1)

首先来看第一个遍历方法：按粒子顺序遍历(A1)。事实上

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
#define MAX_NEIGHBORS 200  // 每个粒子最多的邻居数

__global__ void neighborSearchByParticles(
    const double *x, const double *y, const double *z,
    const double *h,
    const int *gridParticlesList,    // 网格中粒子索引
    const int *gridWritingPointer,   // 每个网格中的粒子数量（已由atomicAdd更新）
    const int *particleGridList,     // 每个粒子所在网格索引
    const double minx, const double miny, const double minz,
    const int nx, const int ny, const int nz,
    int *neighborList,               // 输出：每个粒子的邻居列表，大小为 numParticles * MAX_NEIGHBORS
    int *neighborCount,              // 输出：每个粒子的邻居数
    int numParticles
) {
    int tid = threadIdx.x + blockIdx.x * blockDim.x;
    if (tid >= numParticles) return;

    double xi = x[tid], yi = y[tid], zi = z[tid];
    double hi = h[0];
    double hi2 = hi * hi;
    int gridIndex = particleGridList[tid];

    // 计算该粒子所在网格的坐标索引
    int ix = gridIndex % nx;
    int iy = (gridIndex / nx) % ny;
    int iz = gridIndex / (nx * ny);

    int count = 0;

    // 遍历该粒子所在网格及其周围 3x3x3 网格
    for (int dx = -1; dx <= 1; ++dx) {
        int nix = ix + dx;
        if (nix < 0 || nix >= nx) continue;
        for (int dy = -1; dy <= 1; ++dy) {
            int niy = iy + dy;
            if (niy < 0 || niy >= ny) continue;
            for (int dz = -1; dz <= 1; ++dz) {
                int niz = iz + dz;
                if (niz < 0 || niz >= nz) continue;

                // 相邻网格索引
                int neighborGrid = nix + niy * nx + niz * nx * ny;
                int npg = gridWritingPointer[neighborGrid];

                for (int k = 0; k < npg; ++k) {
                    int j = gridParticlesList[neighborGrid * MAX_PARTICLES_PER_GRID + k];
                    if (j == tid) continue;

                    double dx = x[j] - xi;
                    double dy = y[j] - yi;
                    double dz = z[j] - zi;
                    double dist2 = dx*dx + dy*dy + dz*dz;

                    if (dist2 < hi2 && count < MAX_NEIGHBORS) {
                        neighborList[tid * MAX_NEIGHBORS + count] = j;
                        count++;
                    }
                }
            }
        }
    }

    neighborCount[tid] = count;
}

事实上，particleGridList 数组是专门为 A1 方案 设计的；而若采用 A2 方案，则无需该数组。

A1 的缺点在于：当粒子在数组中呈随机排布时，线程束（warp）之间对粒子属性和网格数据的访问也将是随机的。在 CUDA 编程中，访存模式对性能影响极大，因此可以预期 A1 的效率不会非常理想。

在一个失眠的深夜，我思考了两个问题：

如何优化 A1 的访存模式？
当粒子的核尺度 $ h $ 相差悬殊（例如在模拟月球遭受小行星撞击时，$h_{min} ≈ 1$，$h_{max} ≈ 500$），如何在链表结构中高效支持如此大的跨度？

关于这两个问题，我各自有一些设想和初步实现，接下来将逐一介绍，并进行测试和分析。

A2：按网格顺序遍历粒子

（此处补充 A2 的实现简介和性能特点。）更多关于树搜索的内容可以参见[树搜索]更多关于树搜索的内容可以参见树搜索.

缺陷

网格搜索我个人感觉效率是比树搜索好的，尤其是当粒子的光滑长度都一致或者差不多的时候。缺点在于，如果粒子在空间中分散的非常广泛，比如撞击引起的dust喷发。这会导致网格数量激增，显存很快就会用光，导致计算失败。因此网格搜索还是适合模拟空间分布稳定的问题，比如溃坝。