Julia中使用向量值函数的性能开销及优化方法

问题现象 #

在Julia中使用向量值函数时，代码执行速度出现明显下降，以下是最小可复现示例：

using BenchmarkTools

function f_scalar(x::Float64)::Float64
    return -x
end

function f_vector(x::Float64)::Array{Float64,1}
    return [-x,0.0]
end

function trajectory_scalar(x0::Float64,t::Float64,dt::Float64)
   x = x0
   nsteps = convert(Int,t/dt)
   xvec = zeros(nsteps)
   for k in 1:nsteps
      x = x + f_scalar(x)*dt
      xvec[k] = x
   end
   return xvec
end


function trajectory_vector(x0::Float64,t::Float64,dt::Float64)
   x = x0
   nsteps = convert(Int,t/dt)
   xvec = zeros(nsteps)
   for k in 1:nsteps
      x = x + f_vector(x)[1]*dt
      xvec[k] = x
   end
   return xvec
end

@btime trajectory_scalar(2.0,10.0,0.01)      #1.140 μs (1 allocation: 7.94 KiB)

@btime trajectory_vector(2.0,10.0,0.01)      #12.800 μs (1001 allocations: 86.06 KiB)

基准测试结果显示，向量值函数版本的执行速度比标量版本慢一个数量级，且内存分配量大幅增加。

原因分析 #

核心原因是每次调用f_vector函数时都会分配新的动态数组（Array），循环执行1000次就会产生1000次额外的内存分配，内存管理的开销直接导致了性能下降。实际场景中需要模拟向量值随机微分方程，必须保留向量值函数的同时消除该性能瓶颈。

优化方案：使用StaticArrays #

改用StaticArrays.jl提供的静态向量类型SVector后，向量版本的性能与标量版本持平。修改后的向量值函数如下：

function f_vector(x::Float64)::SVector{2,Float64}
    return SVector(-x,0.0)
end

再次运行基准测试，结果如下：

@btime trajectory_vector(2.0,10.0,0.01)      # 1.200 μs (1 allocation: 7.94 KiB)

内容的提问来源于stack exchange，提问作者Andreas