数组循环优化

最近我遇到了一些Julia数组循环的速度问题。一开始采用的是dot fusion的写法,发现可能有内存分配的速度影响,于是查询文档后决定改用loop的写法,结果却更慢了…

以下是我截取的函数例子:

x = range(0.0, 1.0, length=10000)
y = 0.5:0.5:1.5
z = 0.5:0.5:1.5
var = zeros(10000,3,3,3)

function foo(hx::StepRangeLen,hy::StepRangeLen,hz::StepRangeLen,
   var::Array{Float64,4})::Array{Float64,3}

   siz = size(var)::NTuple{4,Int64}

   px = zeros(Float64,siz[1:3]...)
   qy = zeros(Float64,siz[1:3]...)
   rz = zeros(Float64,siz[1:3]...)

   n = size(hx,1)
   # Right now do nothing for the ghost cells maybe needed later!
   # Take central differences on interior points
   if n > 2
      @time for k=1:siz[3], j=1:siz[2], i=1:n-2
         px[i+1,j,k] = (var[i+2,j,k,1] - var[i,j,k,1])/(hx[i+2] - hx[i])
      end
      @time @. px[2:n-1,:,:] = (var[3:n,:,:,1] - var[1:n-2,:,:,1])/(hx[3:n] - hx[1:n-2])
   end

   px
end

px = foo(x,y,z,var);

速度上显示,dot fusion写法比 explicit loop快了5倍,内存分配也更少:

px = foo(x,y,z,var);
  0.006217 seconds (175.37 k allocations: 2.676 MiB)
  0.001222 seconds (18 allocations: 1.374 MiB)

我尝试了用@code_warntype 查看生成的type stability,有红色部分,但是我不太能看懂。

我的实际代码中有大量类似的数组操作。请问我哪里做得不对?

应该是缓存的原因,用了 @. 的代码里 hx[3:n] - hx[1:n-2] 会被缓存,而手写的循环会重复计算 hx[3:n] - hx[1:n-2] 中元素的值。使用

@macroexpand @. px[2:n-1,:,:] = (var[3:n,:,:,1] - var[1:n-2,:,:,1])/(hx[3:n] - hx[1:n-2])

即可查看展开后的表达式。

最终我还是在英文论坛上询问了一样的问题,得到了一些有用的建议。有兴趣请参看:
loop-over-array-optimization

在使用@views以后性能下降可能是view现在依然在heap上进行内存分配导致的

https://github.com/JuliaLang/julia/issues/14955

在这个issue解决之前,在能够假设你多次使用的view在内存上连续的前提下(没有这个前提那么应该区别不会太大),使用UnsafeArrays.jl 或者直接将所有的循环展开写(不要使用slice)将会提高很多性能。

@views @. px[2:n-1,:,:] = (var[3:n,:,:,1] - var[1:n-2,:,:,1])/(hx[3:n] - hx[1:n-2])

我之前的确没有注意到这一点,毕竟接触时间和经验还不是很够。所以比如以上这个例子,@view目前会对等号右边的所有变量在heap上分配内存进行操作?我还是不是很明白。

在我之前的尝试中,有过另外一个类似的性能问题。对一个explicit loop,如果里面所有的操作都是常规数组的运算,效率很高,@time显示内存分配为0。但是当我把explicit loop里面等号右边的数组替换为其他数组的view时,这些explicit loop会显著地变慢,并且有大量的内存分配。不知道是否跟这个performance issue里的问题是否是相关的?

你这里的代码里没有这个问题,你的问题是类型不稳定。不要单纯的把东西塞进函数里,注意好类型稳定性能解决90%以上的问题。

是的,time现实的gc allocation就是heap allocation。view本身是有allocation的,但是你的这个例子并不会频繁的在for loop里使用view,所以并不是主要问题。

1 个赞