在Julia中高效地将标量加到矩阵中

Question

我需要将一个标量添加到一个巨大矩阵的所有元素中。该矩阵可能会非常大。在这个例子中，我将使用2 GiB的大小，但在我的实际计算中它会更大。

A = rand(2^14, 2^14)

如果我执行

A += 1.0

Julia分配了额外的2 GiB内存。此操作大约需要1秒钟时间。我可以使用一个for循环：

for jj = 1:size(A, 2), ii = 1:size(A, 1)
  A[ii, jj] = A[ii, jj] + 1.0
end

这不会分配任何内存，但需要一分钟的时间。对于我来说，这两种方法都不可行，因为第一种违反了内存限制，而第二种明显效率低下。对于逐元素乘法，可以使用使用 BLAS 的 scal!。是否有任何方法可以像使用 scal! 一样高效地执行加法？

- Diphtong

Answer 1

@DSM的回答是很好的。然而，除此之外，还有一些事情我想补充解释一下。你的for循环慢的原因是因为A是一个非常量的全局变量，而且你的代码在直接改变这个全局变量。由于A不是常量，代码必须防止在循环执行期间A变成不同类型的不同值。代码必须在每次循环迭代中查找A的类型和位置，并动态分派表达式A[ii, jj] = A[ii, jj] + 1.0中的方法调用 - 这是一个对getindex、+和setindex!的调用，所有这些都依赖于静态未知的A类型。你可以通过将其封装在一个函数中来立即获得更好的性能：

julia> A = rand(2^10, 2^10);

julia> @time for jj = 1:size(A, 2), ii = 1:size(A, 1)
           A[ii, jj] += 1
       end
elapsed time: 0.288340785 seconds (84048040 bytes allocated, 15.59% gc time)

julia> function inc!(A)
           for jj = 1:size(A, 2), ii = 1:size(A, 1)
               A[ii, jj] += 1
           end
       end
inc! (generic function with 1 method)

julia> @time inc!(A)
elapsed time: 0.006076414 seconds (171336 bytes allocated)

julia> @time inc!(A)
elapsed time: 0.000888457 seconds (80 bytes allocated)

避免使用非常量全局变量是手册中第一个建议，详见 Performance Tips 部分。您可能还想浏览本章的其余内容。

我们可以进一步改善 inc! 函数的性能，使用 @inbounds 注释指示此代码不需要边界检查，并使用线性索引而非二维索引：

julia> function inc!(A)
           @inbounds for i = 1:length(A)
               A[i] += 1
           end
       end
inc! (generic function with 1 method)

julia> @time inc!(A)
elapsed time: 0.000637934 seconds (80 bytes allocated)

大部分加速来自于@inbounds注释，而不是线性索引，尽管它确实可以稍微提高一点速度。然而，应该只在确定索引不会越界且性能至关重要的情况下，谨慎使用@inbounds注释。如您所见，额外的性能提升虽然存在，但并不是压倒性的。大部分好处来自于不直接更改全局变量。

- StefanKarpinski

Answer 2

您可以进行原地广播操作：

julia> A = rand(2^14, 2^14); A[1:5, 1:5]
5x5 Array{Float64,2}:
 0.229662  0.680236    0.131202  0.111664   0.802698
 0.500575  0.580994    0.385844  0.983806   0.324382
 0.701694  0.577749    0.532591  0.0508955  0.94325 
 0.592929  0.00319653  0.759241  0.448704   0.706204
 0.867945  0.0413606   0.586151  0.82561    0.679233

julia> @time broadcast!(.+, A, A, 100);
elapsed time: 0.382669486 seconds (11490976 bytes allocated)

julia> A[1:5, 1:5]
5x5 Array{Float64,2}:
 100.23   100.68   100.131  100.112  100.803
 100.501  100.581  100.386  100.984  100.324
 100.702  100.578  100.533  100.051  100.943
 100.593  100.003  100.759  100.449  100.706
 100.868  100.041  100.586  100.826  100.679

这个程序只使用了大约2G的内存。

- DSM

---

值得注意的是，broadcast! 使用的大部分时间和内存都来自编译；第二次运行时，两者都会显著减少。 - tholy