广播会分配大量内存，而 for 循环不会

Posted 2023-02-23

【中文标题】广播会分配大量内存，而 for 循环不会

【英文标题】：Broadcasting allocates a lot of memory whist for loop doesn't

【发布时间】：2021-12-30 20:20:47

【问题描述】：

如果我有以下两种变体来编写我的代码的一部分，我通常更喜欢较短的一种。这就是为什么我喜欢在 Julia 中进行点播的原因。但是我有这个例子表明它有问题。

function test1(A)
    n = size(A, 2)
    for k in 1:n
        for i in k+1:n
            A[i, k] /= A[k, k]
        end
    end
    A
end

function test2(A)
    n = size(A, 2)
    for k in 1:n
        A[k+1:n, k] ./= A[k,k]
    end
    A
end

A = rand(10000,10000)
B = copy(A)
@allocated test1(B) # returns 0
C = copy(A)
@allocated test2(C) # returns 401131136
C == B # returns true

此外，我试图消除语法糖并编写了一个更多的测试函数

function test3(A)
    n = size(A, 2)
    for k in 1:n
        broadcast!(/, A[k+1:n, k], A[k+1:n, k], A[k,k])
    end
    A
end

这个分配的内存是test2 的两倍，并且...返回错误的结果。

test3为什么会返回错误的结果，广播有什么问题，有没有指导说明如何使用？

【参考方案1】：

问题在于，在A[k+1:n, k] ./= A[k,k] 中，您可以认为它被重写为：

A[k+1:n, k] .= A[k+1:n, k] ./ A[k,k]

现在的重点是，在A[k+1:n, k] ./ A[k,k] 部分中，A[k+1:n, k] 部分在执行除法之前复制了数组。原因是安全。您可以通过显式创建视图而不是副本来覆盖它：

function test2_fix(A)
    n = size(A, 2)
    for k in 1:n
        A[k+1:n, k] .= view(A, k+1:n, k) ./ A[k,k]
    end
    A
end

broadcast! 调用更糟糕且不正确，因为A[k+1:n, k] 在将值传递给broadcast! 之前创建了两次副本（当您使用它两次时），因此您的源数组A 不会受到影响

【参考方案2】：

我将分解语法的最终含义。这是你的第一次广播尝试和它的等价物（好吧，Meta.@lower 中的步骤有点不同，但它对Arrays 做同样的事情）：

A[k+1:n, k] ./= A[k,k]
A[k+1:n, k] .= A[k+1:n, k] ./ A[k,k]
broadcast!(/, @view(A[k+1:n, k]), A[k+1:n, k], A[k,k])

注意.=broadcast! 是如何设置为将结果写入左侧数组A 的视图。视图访问数组的内存，从 v1.5 开始它是非分配的。您对broadcast! 的第二次尝试错过了这个细节，因此您最终写入了一个新数组而不是A。

问题在于分割括号语法，例如A[k+1:n, k]，不适合getindex。如果getindex 正在访问>1 个值，它会分配一个新数组并将值复制到其中，因此在这种情况下，您需要使用view 而不是getindex。以下是执行此操作的等效方法：

broadcast!(/, @view(A[k+1:n, k]), @view(A[k+1:n, k]), A[k,k])
A[k+1:n, k] .= @view(A[k+1:n, k]) ./ A[k,k]

@view(A[k+1:n, k]) ./= A[k,k]
@views A[k+1:n, k] ./= A[k,k]

@view 宏允许您使用切片括号语法来调用view 而不是getindex。 @views 宏是一种将以下表达式中的所有切片括号语法转换为使用 view 而不是 getindex 的方法。

【讨论】：

应该注意，您必须扩展 ./= 以避免所有额外的分配。但是，此答案中对@view 和@views 的讨论非常好。

您能否更详细地解释“扩展./=”，也许用一行代码？我不确定你的意思是什么以及它如何减少分配。

这条语句有点不精确：“如果getindex 正在访问>1 个值，它会分配一个新数组”。重要的是索引的类型，标量索引返回一个标量，数组索引分配一个数组。即使是单元素向量索引也会分配。