让 Julia 达到 C++ 的速度（2019）

# 让 Julia 达到 C++ 的速度 · FLOW Lab 来源：https://flow.byu.edu/posts/julia-c++ 2019年4月18日，Eduardo Alvarez 视频链接 ## 介绍 传言称 Julia 能够达到与 C++ 和 FORTRAN 等编译语言相同的[计算性能](https://julialang.org/benchmarks/)。在过去两年使用 Julia 编程的过程中，我深深爱上了它类似 Python 的语法、多重分发机制，以及类似 MATLAB 的线性代数便利性，同时还能使用显式类型声明等编译语言特性来预防 bug。总的来说，Julia 的设计理念兼具解释型语言的灵活性，而其即时（JIT）编译机制使其实际上成为一门编译型语言。 Julia 的高级语法让来自任何背景的程序员都能轻松上手，然而实现高性能在某种程度上是一门艺术。在这篇文章中，我总结了在为高性能计算优化 Julia 代码过程中学到的一些经验。我将通过一个来自空气动力学领域的实际计算应用——[涡粒子法](https://scholarsarchive.byu.edu/facpub/2116/)\(^{[1,\,2]}\)——来展示代码优化的过程。 ## 问题定义 在[涡粒子法](https://scholarsarchive.byu.edu/facpub/2116/)中，我们关注的是计算由 \(N\) 个涡粒子组成的粒子场在任意位置 \(\mathbf{x}\) 处所诱导的速度 \(\mathbf{u}\) 和速度雅可比矩阵 \(\mathbf{J}\)。计算公式如下： \[{\bf u}\left( {\bf x} \right) = \sum\limits_p^N g_\sigma\left( {\bf x}-{\bf x}_p \right) {\bf K}\left( {\bf x}-{\bf x}_p \right) \times \boldsymbol\Gamma_p\] \[\frac{\partial {\bf u}}{\partial x_j}\left( {\bf x} \right) = \sum\limits_p^N \left[ \left( \frac{1}{\sigma } \frac{\Delta x_j}{\Vert \Delta \mathbf{x} \Vert} \frac{\partial g}{\partial r} \left( \frac{\Vert \Delta\mathbf{x} \Vert}{\sigma} \right) - 3 g_\sigma\left( \Delta{\bf x} \right) \frac{\Delta x_j}{\Vert \Delta\mathbf{x} \Vert^2} \right) {\bf K}\left( \Delta\mathbf{x} \right) \times \boldsymbol\Gamma_p - \frac{g_\sigma\left( \Delta{\bf x} \right) }{4\pi \Vert \Delta{\bf x} \Vert^3} \delta_{ij} \times \boldsymbol\Gamma_p \right] ,\] 其中 \({\bf K}\) 为奇异牛顿核 \({\bf K}\left( {\bf x}\right)=-\frac{ {\bf x} }{4\pi \Vert{\bf x}\Vert^3}\)，\(g_\sigma\) 为平滑半径为 \(\sigma\) 的正则化函数，\(\mathbf{x}_p\) 和 \(\boldsymbol{\Gamma}_p\) 分别为第 \(p\) 个粒子的位置和矢量强度。 此外，该方法的控制方程需要计算粒子集合对自身所诱导的速度 \(\mathbf{u}\) 和雅可比矩阵 \(\mathbf{J}\)，从而引出众所周知的计算复杂度为 \(\mathcal{O}(N^2)\) 的 [N 体问题](https://en.wikipedia.org/wiki/N-body_problem)。 选用 Winckelmans 正则化核\(^{[1]}\)： \[g(r) = r^3 \frac{r^2 + 5/2}{\left( r^2 + 1 \right)^{5/2}}\] \[\frac{\partial g}{\partial r} (r) = \frac{15}{2} \frac{r^2}{\left( r^2 + 1 \right)^{7/2}} ,\] 上述方程在 C++ 中的实现如下： ```cpp // 粒子间相互作用 void P2P(Particle * P, const int numParticles) { real_t r, ros, aux, g_sgm, dgdr; vec3 dX; for (int i=0; i ``` **输入 \[1\]：** ```julia """ 这是一个 Python 风格的实现， 没有任何为提升性能而做的调整 """ function P2P_pythonic(particles, g, dgdr) for Pi in particles for Pj in particles dX = Pi.X - Pj.X r = norm(dX) if r != 0 # g_σ 和 ∂gσ∂r gsgm = g(r / Pj.sigma) dgsgmdr = dgdr(r / Pj.sigma) # K × Γp crss = cross(-const4 * (dX/r^3), Pj.Gamma) # U = ∑g_σ(x-xp) * K(x-xp) × Γp Pi.U[:] += gsgm * crss # ∂u∂xj(x) = ∑[ ∂gσ∂xj(x−xp) * K(x−xp)×Γp + gσ(x−xp) * ∂K∂xj(x−xp)×Γp ] for j in 1:3 Pi.J[:, j] += ( dX[j] / (Pj.sigma*r) * (dgsgmdr*crss) - gsgm * (3*dX[j]/r^2) * crss - gsgm * (const4/r^3) * cross([i==j for i in 1:3], Pj.Gamma) ) end end end end end ``` **输入 \[2\]：** ```julia # 生成一个 6x6x6 的粒子网格 n = 6 # 每轴粒子数 lambda = 2.0 # 粒子间距与 sigma 之比 # 创建一个简单的粒子 struct（无具体类型） mutable struct ParticleAmbiguous X # 位置 Gamma # 矢量环量 sigma # 平滑半径 U # 粒子处的速度 J # 粒子处的雅可比矩阵 J[i,j]=dUi/dxj end particles = generate_particles(ParticleAmbiguous, n, lambda) ``` **输入 \[3\]：** ```julia args = (particles, g_wnk, dgdr_wnk) @benchmark P2P_pythonic(args...); ``` ``` BenchmarkTools.Trial: memory estimate: 399.26 MiB allocs estimate: 4782960 -------------- minimum time: 233.679 ms (15.21% GC) median time: 253.098 ms (16.10% GC) mean time: 252.237 ms (16.27% GC) maximum time: 270.524 ms (16.39% GC) -------------- samples: 5 evals/sample: 1 ``` ### 修复 \#1：避免使用抽象类型 注意我们定义的 `ParticleAmbiguous` struct 中，属性没有指定类型。这应该会立即引起警觉，因为 Julia 中的动态类型非常低效。让我们用 `@code_warntype` 来确认一下这个问题： **输入 \[4\]：** ```julia Pi = particles[1] @code_warntype P2P_pythonic(particles, g_wnk, dgdr_wnk) ``` ``` Variables: #self# <optimized out> particles::Array{ParticleAmbiguous,1} g::getfield(Main, Symbol("##7#9")) dgdr::getfield(Main, Symbol("##8#10")) Pi@_4::Union{} Pj@_5::Union{} Pi@_6::ParticleAmbiguous Pj@_7::ParticleAmbiguous dX::Any r::Any gsgm::Any dgsgmdr::Any crss::Any j::Int64 Body::Nothing │╻ iterate 2 1 ─ %1 = (Base.arraylen)(particles)::Int64 │ %2 = (Base.sle_int)(1, %1)::Bool │ %3 = (Base.ifelse)(%2, 1, 0)::Int64 │ %4 = (Base.slt_int)(%3, 1)::Bool └── goto #13 if not %4 . . . │╻ getproperty 37 │ %35 = (Base.getfield)(%16, :X)::Any ││ │ %36 = (Base.getfield)(%33, :X)::Any │ │ %37 = (%35 - %36)::Any │ 38 │ %38 = (Main.norm)(%37)::Any . . . ``` 理解这种降级 AST 语法本身就是一门艺术，但你很快会发现 `@code_warntype` 是优化代码时最好的朋友。当我们向下滚动 AST 时，会看到代码在 `ParticleAmbiguous` 类型的属性中遇到了 `Any` 类型，这应该立刻引起警觉（`Any` 是抽象类型）。实际上，运行 `@code_warntype` 时，输出会自动将 `::Any` 断言以红色高亮显示。 我们可以通过[参数化定义 struct 属性](https://docs.julialang.org/en/v1/manual/performance-tips/index.html#Type-declarations-1)来解决这些抽象类型的问题： **输入 \[5\]：** ```julia """ 这是一个显式/参数化定义了属性类型的粒子 struct。 """ struct Particle{T} # 用户输入 X::Array{T, 1} # 位置 Gamma::Array{T, 1} # 矢量环量 sigma::T # 平滑半径 # 属性 U::Array{T, 1} # 粒子处的速度 J::Array{T, 2} # 粒子处的雅可比矩阵 J[i,j]=dUi/dxj end # 另一个初始化别名 Particle{T}(X, Gamma, sigma) where {T} = Particle(X, Gamma, sigma, zeros(T,3), zeros(T, 3, 3)) # 空初始化 Base.zero(::Type{<:Particle{T}}) where {T} = Particle(zeros(T, 3), zeros(T, 3), zero(T), zeros(T, 3), zeros(T, 3, 3)) ``` `P2P_pythonic` 函数无需进一步修改，因为 Julia 的多重分发和 JIT 会自动即时编译出一个专门针对我们新的 `Particle{T}` 类型的函数版本。不过，我们还是定义一个别名以便比较基准： **输入 \[6\]：** ```julia P2P_concretetypes(args...) = P2P_pythonic(args...) ``` **输入 \[7\]：** ```julia # 创建一个 6x6x6 的具体 Float64 粒子网格 particles = generate_particles(Particle{Float64}, n, lambda) # 运行基准测试 args = (particles, g_wnk, dgdr_wnk) compare(P2P_concretetypes, P2P_pythonic, args) @benchmark P2P_concretetypes(args...); ``` ``` P2P_concretetypes 比 P2P_pythonic 快 3.06 倍（76.488ms vs 233.679ms） BenchmarkTools.Trial: memory estimate: 189.93 MiB allocs estimate: 2275776 -------------- minimum time: 76.488 ms (13.63% GC) median time: 77.860 ms (13.79% GC) mean time: 84.567 ms (17.89% GC) maximum time: 145.918 ms (42.64% GC) -------------- samples: 12 evals/sample: 1 ``` 太棒了！通过在 `Particle` struct 中指定具体类型，我们获得了 3 倍的速度提升（还应该再次运行 `@code_warntype` 以确认所有抽象类型都已消除，但为简洁起见此处略去）。现在来看看与 C++ 的对比： **输入 \[8\]：** ```julia printcomparison(P2P_concretetypes, "C++", false) ``` ``` C++ 比 P2P_concretetypes 快 19.14 倍（3.996ms vs 76.488ms） ``` 使用具体类型大大加快了计算速度；然而，C++ 版本仍然比 Julia 快约 20 倍。让我们看看还能优化什么。 ### 修复 \#2：避免列表推导式 列表推导式的便利性可能会诱使你写出行数精简的代码；然而，随意使用可能导致非常低效的计算。以这个列表推导式求和为例： **输入 \[9\]：** ```julia sum_list(n) = sum([i for i in 1:n]) @btime sum_list(100); ``` ``` 80.975 ns (1 allocation: 896 bytes) ``` 下面是同一函数不使用列表推导式的展开版本，速度快了 60 倍： **输入 \[10\]：** ```julia function sum_unrolled(n) out = 0 for i in 1:n out += i end return out end @btime sum_unrolled(100); ``` ``` 1.374 ns (0 allocations: 0 bytes) ``` 在我们的 P2P 函数中，有一个 Kronecker delta 叉积原本用一行列表推导式计算： ```julia # ∂u∂xj(x) = ∑[ ∂gσ∂xj(x−xp) * K(x−xp)×Γp + gσ(x−xp) * ∂K∂xj(x−xp)×Γp ] for j in 1:3 Pi.J[:, j] += ( dX[j] / (Pj.sigma*r) * (dgsgmdr*crss) - gsgm * (3*dX[j]/r^2) * crss - gsgm * (const4/r^3) * cross([i==j for i in 1:3], Pj.Gamma) ) end ``` 替代方案是将其展开为几行： ```julia # ∂u∂xj(x) = ∑[ ∂gσ∂xj(x−xp) * K(x−xp)×Γp + gσ(x−xp) * ∂K∂xj(x−xp)×Γp ] for j in 1:3 Pi.J[:, j] += ( dX[j] / (Pj.sigma*r) * (dgsgmdr*crss) - gsgm * (3*dX[j]/r^2) * crss ) end # ∂u∂xj(x) = −∑gσ/(4πr^3) δij×Γp # 添加 Kronecker delta 项 aux = -const4 * gsgm / r^3 # j=1 Pi.J[2, 1] -= aux * Pj.Gamma[3] Pi.J[3, 1] += aux * Pj.Gamma[2] # j=2 Pi.J[1, 2] += aux * Pj.Gamma[3] Pi.J[3, 2] -= aux * Pj.Gamma[1] # j=3 Pi.J[1, 3] -= aux * Pj.Gamma[2] Pi.J[2, 3] += aux * Pj.Gamma[1] ``` 列表推导式操作的问题在于它需要分配内存来构建生成的数组。抵制使用列表推导式来节省几行代码的诱惑，直接将其展开即可。如下所示，通过展开这一行，我们获得了 1.5 倍的速度提升： **输入 \[11\]：** ```julia """ 展开列表推导式 """ function P2P_nocomprehension(particles, g, dgdr) for Pi in particles for Pj in particles dX = Pi.X - Pj.X r = norm(dX) if r != 0 # g_σ 和 ∂gσ∂r gsgm = g(r / Pj.sigma) dgsgmdr = dgdr(r / Pj.sigma) # K × Γp crss = cross(-const4 * (dX/r^3), Pj.Gamma) # U = ∑g_σ(x-xp) * K(x-xp) × Γp Pi.U[:] += gsgm * crss # ∂u∂xj(x) = ∑[ ∂gσ∂xj(x−xp) * K(x−xp)×Γp + gσ(x−xp) * ∂K∂xj(x−xp)×Γp ] for j in 1:3 Pi.J[:, j] += ( dX[j] / (Pj.sigma*r) * (dgsgmdr*crss) - gsgm * (3*dX[j]/r^2) * crss ) end # ∂u∂xj(x) += −∑gσ/(4πr^3) δij×Γp # 添加 Kronecker delta 项 aux = -const4 * gsgm / r^3 # j=1 Pi.J[2, 1] -= aux * Pj.Gamma[3] Pi.J[3, 1] += aux * Pj.Gamma[2] # j=2 Pi.J[1, 2] += aux * Pj.Gamma[3] Pi.J[3, 2] -= aux * Pj.Gamma[1] # j=3 Pi.J[1, 3] -= aux * Pj.Gamma[2] Pi.J[2, 3] += aux * Pj.Gamma[1] end end end end ``` **输入 \[12\]：** ```julia args = (particles, g_wnk, dgdr_wnk) compare(P2P_nocomprehension, P2P_concretetypes, args) @benchmark P2P_nocomprehension(args...); ``` ``` P2P_nocomprehension 比 P2P_concretetypes 快 1.52 倍（50.196ms vs 76.488ms） BenchmarkTools.Trial: memory estimate: 126.16 MiB allocs estimate: 1300536 -------------- minimum time: 50.196 ms (11.28% GC) median time: 51.929 ms (11.65% GC) mean time: 55.319 ms (15.41% GC) maximum time: 107.039 ms (50.23% GC) -------------- samples: 19 evals/sample: 1 ``` ### 修复 \#3：减少内存分配 接下来，我们注意到基准测试分配了大量内存（126MiB）和分配操作（130万次）。我怀疑这是 Julia 允许动态大小数组导致的问题。解决此问题的第一步是**避免创建任何数组类型的内部变量**。在下面的代码中，注意我将数组变量 `dX` 和 `crss` 替换为浮点变量 `dX1, dX2, dX3` 和 `crss1, crss2, crss3`，这就需要将内层 for 循环完全展开： **输入 \[13\]：** ```julia """ 减少内存分配 """ function P2P_noallocation(particles, g, dgdr) for Pi in particles for Pj in particles dX1 = Pi.X[1] - Pj.X[1] dX2 = Pi.X[2] - Pj.X[2] dX3 = Pi.X[3] - Pj.X[3] r = norm(Pi.X - Pj.X) if r != 0 # g_σ 和 ∂gσ∂r gsgm = g(r / Pj.sigma) dgsgmdr = dgdr(r / Pj.sigma) # K × Γp crss1, crss2, crss3 = -const4 / r^3 * cross(Pi.X - Pj.X, Pj.Gamma) # U = ∑g_σ(x-xp) * K(x-xp) × Γp Pi.U[1] += gsgm * crss1 Pi.U[2] += gsgm * crss2 Pi.U[3] += gsgm * crss3 # ∂u∂xj(x) = ∑[ ∂gσ∂xj(x−xp) * K(x−xp)×Γp + gσ(x−xp) * ∂K∂xj(x−xp)×Γp ] # ∂u∂xj(x) += ∑p[(Δxj∂gσ∂r/(σr) − 3Δxjgσ/r^2) K(Δx)×Γp aux = dgsgmdr/(Pj.sigma*r)* - 3*gsgm /r^2 # j=1 Pi.J[1, 1] += aux * crss1 * dX1 Pi.J[2, 1] += aux * crss2 * dX1 Pi.J[3, 1] += aux * crss3 * dX1 # j=2 Pi.J[1, 2] += aux * crss1 * dX2 Pi.J[2, 2] += aux * crss2 * dX2 Pi.J[3, 2] += aux * crss3 * dX2 # j=3 Pi.J[1, 3] += aux * crss1 * dX3 Pi.J[2, 3] += aux * crss2 * dX3 Pi.J[3, 3] += aux * crss3 * dX3 # 添加 Kronecker delta 项 aux = -const4 * gsgm / r^3 # j=1 Pi.J[2, 1] -= aux * Pj.Gamma[3] Pi.J[3, 1] += aux * Pj.Gamma[2] # j=2 Pi.J[1, 2] += aux * Pj.Gamma[3] Pi.J[3, 2] -= aux * Pj.Gamma[1] # j=3 Pi.J[1, 3] -= aux * Pj.Gamma[2] Pi.J[2, 3] += aux * Pj.Gamma[1] end end end end ``` **输入 \[14\]：** ```julia args = (particles, g_wnk, dgdr_wnk) compare(P2P_noallocation, P2P_nocomprehension, args) @benchmark P2P_noallocation(args...); ``` ``` P2P_noallocation 比 P2P_nocomprehension 快 3.68 倍（13.627ms vs 50.196ms） BenchmarkTools.Trial: memory estimate: 21.99 MiB allocs estimate: 325296 -------------- minimum time: 13.627 ms (7.41% GC) median time: 15.615 ms (8.40% GC) mean time: 16.582 ms (12.83% GC) maximum time: 59.011 ms (66.91% GC) -------------- samples: 61 evals/sample: 1 ``` 我们将内存分配从 126MiB 降低到了 22MiB，带来了 3.5 倍的速度提升。让我们看看还能做什么来进一步减少内存分配。 ### 修复 \#4：避免线性代数函数 接下来需要考虑的是，**使用 Julia 基础库中的线性代数运算（即 `dot(X,X)`、`cross(X,X)`、`norm(X,X)`）比显式地将操作展开为代码行代价更高**。我怀疑这也是一个内存分配问题，因为这些函数需要分配内部数组来存储计算过程中的中间结果，然后再输出。 以下是不使用任何基础库线性代数函数的代码（注意我不再使用 `norm()` 和 `cross()`）： **输入 \[15\]：** ```julia """ 不使用线性代数函数 """ function P2P_nolinalg(particles, g, dgdr) for Pi in particles for Pj in particles dX1 = Pi.X[1] - Pj.X[1] dX2 = Pi.X[2] - Pj.X[2] dX3 = Pi.X[3] - Pj.X[3] r = sqrt(dX1*dX1 + dX2*dX2 + dX3*dX3) if r != 0 # g_σ 和 ∂gσ∂r gsgm = g(r / Pj.sigma) dgsgmdr = dgdr(r / Pj.sigma) # K × Γp crss1 = -const4 / r^3 * ( dX2*Pj.Gamma[3] - dX3*Pj.Gamma[2] ) crss2 = -const4 / r^3 * ( dX3*Pj.Gamma[1] - dX1*Pj.Gamma[3] ) crss3 = -const4 / r^3 * ( dX1*Pj.Gamma[2] - dX2*Pj.Gamma[1] ) # U = ∑g_σ(x-xp) * K(x-xp) × Γp Pi.U[1] += gsgm * crss1 Pi.U[2] += gsgm * crss2 Pi.U[3] += gsgm * crss3 # ∂u∂xj(x) = ∑[ ∂gσ∂xj(x−xp) * K(x−xp)×Γp + gσ(x−xp) * ∂K∂xj(x−xp)×Γp ] # ∂u∂xj(x) += ∑p[(Δxj∂gσ∂r/(σr) − 3Δxjgσ/r^2) K(Δx)×Γp aux = dgsgmdr/(Pj.sigma*r)* - 3*gsgm /r^2 # j=1 Pi.J[1, 1] += aux * crss1 * dX1 Pi.J[2, 1] += aux * crss2 * dX1 Pi.J[3, 1] += aux * crss3 * dX1 # j=2 Pi.J[1, 2] += aux * crss1 * dX2 Pi.J[2, 2] += aux * crss2 * dX2 Pi.J[3, 2] += aux * crss3 * dX2 # j=3 Pi.J[1, 3] += aux * crss1 * dX3 Pi.J[2, 3] += aux * crss2 * dX3 Pi.J[3, 3] += aux * crss3 * dX3 # 添加 Kronecker delta 项 aux = -const4 * gsgm / r^3 # j=1 Pi.