Julia中的预分配

Question

我正在尝试通过预先分配数组来最小化Julia中的内存分配，如文档中所示。我的示例代码如下：

using BenchmarkTools

dim1 = 100
dim2 = 1000
A = rand(dim1,dim2)
B = rand(dim1,dim2)
C = rand(dim1,dim2)
D = rand(dim1,dim2)

M = Array{Float64}(undef,dim1,dim2)

function calc!(a, b, c, d, E)
     @. E = a * b * ((d-c)/d)
    nothing
end

function run_calc(A,B,C,D,M)
    for i in 1:dim2
        @views calc!(A[:,i], B[:,i], C[:,i], D[:,i], M[:,i])
    end
end

我的理解是，由于M在两个函数之外被预先分配，因此这应该本质上不会再进行分配。然而，当我对此进行基准测试时，仍然看到了大量的分配：

@btime run_calc(A,B,C,D,M)

1.209毫秒（14424次分配：397.27 KiB）

在这种情况下，我当然可以运行更简洁的代码

@btime @. M = A * B * ((D-C)/D)

这段代码只期望执行非常少的内存分配：

122.599 微秒 (6 次内存分配：144 字节)

然而，我的实际代码更复杂，无法像这样简化，因此我想知道第一个版本哪里出了问题。

- jul345

Answer 1

你没有做错任何事情。目前在Julia中创建视图会进行内存分配（正如Stefan所指出的，这方面已经比过去好了很多，但在这种情况下仍然会发生一些分配）。你看到的内存分配是由此引起的。

参见：

julia> @allocated view(M, 1:10, 1:10)
64

在这种情况下，编写一个适当的循环可能是最简单的解决方案（我假设您的代码中的循环会更加复杂，但我希望意图是清晰的），例如：

julia> function run_calc2(A,B,C,D,M)
           @inbounds for i in eachindex(A,B,C,D,M)
               M[i] = A[i] * B[i] * ((D[i] - C[i])/D[i])
           end
       end
run_calc2 (generic function with 1 method)

julia> @btime run_calc2($A,$B,$C,$D,$M)
  56.441 μs (0 allocations: 0 bytes)

julia> @btime run_calc($A,$B,$C,$D,$M)
  893.789 μs (14424 allocations: 397.27 KiB)

julia> @btime @. $M = $A * $B * (($D-$C)/$D);
  381.745 μs (0 allocations: 0 bytes)

编辑：所有时间均基于Julia版本1.6.0-DEV.1580。

编辑2：为了完整起见，提供一个将@views传递给内部函数的代码。它仍然会分配内存（但更好），并且比仅使用循环要慢：

julia> function calc2!(a, b, c, d, E, i)
            @inbounds @. @views E[:,i] = a[:,i] * b[:,i] * ((d[:,i]-c[:,i])/d[:,i])
           nothing
       end
calc2! (generic function with 1 method)

julia> function run_calc3(A,B,C,D,M)
           for i in 1:dim2
               calc2!(A,B,C,D,M,i)
           end
       end
run_calc3 (generic function with 1 method)

julia> @btime run_calc3($A,$B,$C,$D,$M);
  305.709 μs (1979 allocations: 46.56 KiB)

- Bogumił Kamiński

---

我特别喜欢你提出的run_calc2解决方案。有趣的是，我认为在数组的较大部分上执行更少的循环迭代会更好，但在这里显然不是这种情况。谢谢！ - jul345

重要的是，“eachindex”将选择遍历矩阵的顺序，并且该顺序是高效的（例如，对于密集数组，它将是CPU缓存友好的）。 - Bogumił Kamiński

Answer 2

在Julia 1.5之前，创建数组视图通常会为视图对象分配一些内存。在Julia 1.5之后，创建视图通常不会造成任何分配。您的帖子中没有包括您使用的Julia版本，因此我假设它比1.5版本旧。在您的代码中，您正在为潜在的大数组维度的每个索引创建一个视图，这肯定会累加。您可以重构此代码以通过内部计算将维度传递给视图。否则，您可以升级Julia并查看分配是否消失。

- StefanKarpinski

---

这段代码在 Julia 1.5 和 1.6 夜间版上都可以分配。 - Bogumił Kamiński

抱歉，我应该提到我正在使用的是Julia 1.4版本。我刚刚在1.5上尝试了相同的代码，可以确认它运行速度略快，但分配的内存量相同。 然而，我不知道视图确实会分配一些内存，所以感谢你指出这一点。像Bogumił建议的重构方式很好用。 - jul345