Julia 中的预分配

Pre-allocation in Julia

我正在尝试通过预分配数组来最小化 Julia 中的内存分配,如 in the documentation 所示。我的示例代码如下所示:

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 ms (14424 allocations: 397.27 KiB)

这样的话我当然可以运行更简洁

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

按预期执行很少的分配:

122.599 μs (6 allocations: 144 bytes)

但是我的实际代码更复杂,不能像这样减少,因此我想知道第一个版本哪里出了问题。

在 Julia 1.5 之前,创建数组视图通常会为视图对象分配一些内存。在 Julia 1.5 之后,创建视图通常不会引起任何分配。您的 post 不包括您使用的 Julia 版本,因此我假设它早于 1.5。在您的代码中,您正在为一个可能很大的数组维度的每个索引创建一个视图,这肯定会加起来。您可以重构此代码以将维度传递给内部计算。否则,您可以升级 Julia 并查看分配是否消失。

你没有做错任何事。目前在 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 的所有计时

EDIT2:为了完整起见,将 @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)