使用 Julia 更新运算符进行不必要的分配答案

【问题标题】：Unnecessary allocations using Julia update operators使用 Julia 更新运算符进行不必要的分配
【发布时间】：2014-11-26 16:57:32
【问题描述】：

考虑以下函数：

function mytest(x, b)
    y = zeros(x[:,:,1])
    for i in 1:length(b)
        y += b[i] * x[:,:,i]
    end
    return y
end

当我运行它时，我得到以下信息：

x = rand(30,30,100000)
b = rand(100000)
@time mytest(x,b)

elapsed time: 0.571765222 seconds (727837732 bytes allocated, 66.49% gc time)

为什么要分配这么多内存并花费这么多时间进行垃圾回收？代码应该是类型稳定的，我希望+= 运算符不会执行重新分配。但是，似乎每次添加两个矩阵时都会重新分配。

我应该认为这是 Julia 中的一个错误吗？更重要的是，我怎样才能以不重新分配的方式编写这段代码？

编辑：修正了错字。

【问题讨论】：

右侧b[i] * x[:,:,i] 首先分配一个临时数组，其中包含乘积的结果，然后将其就地添加到左侧。然后需要对临时数组进行垃圾收集。至少......这就是它与numpy 一起使用的方式。
@moarningsun，感谢您的提示。我之前在性能关键代码方面的经验主要是使用 C++ 和 Eigen，它们在获取切片时不会分配临时空间。我已将代码更改为具有三个嵌套的 for 循环，并且分配问题已得到修复。（在此之前我尝试使用 ArrayViews 包，但这似乎并没有解决问题。）我仍然很好奇是否有更好的方法。
@JimGarrison：查看讨论here。 += 只是语法糖（至少现在是这样），所以我认为它总是重新分配（但如果我错了，请有人纠正我）。另外，我很想看看三个循环如何解决您的问题。
看起来这可能在 Julia 0.4 中得到修复：github.com/JuliaLang/julia/pull/9150

标签： memory-management garbage-collection profiling julia

【解决方案1】：

@cd98 请求了我的三嵌套循环解决方案，它解决了分配问题，但我认为它的性能不如等效的矢量化版本。这里是：

function mytest(x, b)
    d1, d2, d3 = size(x)
    y = zeros(eltype(x), d1, d2)
    for i in 1:d3
        for j in 1:d2
            for k in 1:d1
                y[k,j] += b[i] * x[k,j,i]
            end
        end
    end
    return y
end

x = rand(30,30,100000)
b = rand(100000)
@time mytest(x,b)
@time mytest(x,b)

还有输出：

elapsed time: 0.218220119 seconds (767172 bytes allocated)
elapsed time: 0.197181799 seconds (7400 bytes allocated)

【讨论】：

【解决方案2】：

不能解决（原始）分配问题，但我只需使用 @inbounds 将循环包装在后一种解决方案中，即可获得大约 1.8 倍的加速：

function mytest_inbounds(x, b)
    d1, d2, d3 = size(x)
    y = zeros(eltype(x), d1, d2)
    @inbounds begin
        for i in 1:d3
            for j in 1:d2
                for k in 1:d1
                    y[k,j] += b[i] * x[k,j,i]
                end
            end
        end
    end
    return y
end

x = rand(30, 30, 100000)
b = rand(100000)
@time mytest(x, b)
@time mytest(x, b)
@time mytest_inbounds(x, b)
@time mytest_inbounds(x, b)

输出：

elapsed time: 0.39144919 seconds (767212 bytes allocated)
elapsed time: 0.353495867 seconds (7400 bytes allocated)
elapsed time: 0.202614643 seconds (396972 bytes allocated)
elapsed time: 0.193425902 seconds (7400 bytes allocated)

另外，这里有很多关于相关问题的很好的讨论：

https://groups.google.com/forum/#!msg/julia-users/aYS_AvKqPCI/DyTiq4lKIAoJ

【讨论】：

【解决方案3】：

还有使用Base.Cartesian的可能：using Base.Cartesian之后可以写

@nloops 3 i x begin
  (@nref 2 y i) += b[i_3] * (@nref 3 x i)
end

它扩展到与 Jim 的答案基本相同的循环。

【讨论】：

但就个人而言，我认为这很难看，我更喜欢 Jim 的解决方案。