Broadcasting allocates a lot of memory whist for loop doesn't

Viewed 119

If I have following two variants to write a part of my code and I generally prefer a shorter one. This is the reason why I do like dot-broadcasting in Julia. However I have this example shows that there's something wrong with it.

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

Moreover I tried to eliminate syntactic sugar and wrote one more test function

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

And this one allocates twice as much memory as test2 and... returns wrong result.

Why does test3 return wrong result, and what's wrong with broadcasting, are there any guides explaining how to use it?

2 Answers

The issue is that in A[k+1:n, k] ./= A[k,k] you can think of as it being rewritten as:

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

Now the point is that in the A[k+1:n, k] ./ A[k,k] part the A[k+1:n, k] part makes a copy of the array before division is performed. The reason is safety. You can override it by explicitly making a view instead of a copy:

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

The broadcast! call is even worse and incorrect, because A[k+1:n, k] creates a copy twice (as you use it twice) before passing the values to broadcast! so your source array A will not be affected

I'll break down what the syntax ends up meaning. This was your first broadcasting attempt and its equivalents (well, the steps in Meta.@lower are a bit different but it does the same thing for 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])

Note how .= broadcast! is set up to write the results to a view of the left-hand side array A. A view accesses the memory of an array, and it is non-allocating as of v1.5. Your second attempt with broadcast! missed this detail, so you ended up writing to a new array instead of to A.

The issue is that slicing bracket syntax, such as A[k+1:n, k], does getindex. If getindex is accessing >1 value, it allocates a new array and copies values to it, so in such cases, you want to use view instead of getindex. Here are the equivalent ways of doing that:

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]

The @view macro allows you to use the slicing bracket syntax to call view instead of getindex. The @views macro is a way to turn all slicing bracket syntax in the following expression to use view instead of getindex.

Related