Julia: map and broadcast on symmetric matrices is inefficient

[MasonProtter]

Currently, map and broadcast on Symmetric arguments don’t take advantage of the fact that you only have to compute n*(n+1)/2 elements of the output matrix instead of n^2elements. For big matrices, this can leave a lot of performance on the table.

[StefanKarpinski]

That optimization kind of assumes function purity.

[MasonProtter]

That optimization kind of assumes function purity

That’s a good point that I didn’t consider. However, I feel like the kind of code that would break from this must be very, very rare so perhaps we could get away with just noting in the documentation that map(f, ::Symmetric) assumes some level of purity in f and if they don’t want that, they can do Symmetric(map(f, xs.data)). One could call that ‘surprising behaviour’ but honestly, I’d be more surprised that a function like map doesn't take advantage of the symmetry than the other way around.

I think I recall someone mentioning that since map and broadcast don’t actually guarantee a particular order of operations and its just an implementation detail that they do use a rigid ordering, so one could argue that these functions do actually already sort of assume some level of purity. I don’t know how valid that is though.

[mbauman]

That optimization kind of assumes function purity.

We already assume that for all the other structured (and sparce) matrices with respect to zero preservation. That said, in the SymTridiagonal implementation we do check (and enforce) symmetric purity:

https://github.com/JuliaLang/julia/blob/79a57931a563882bac00647dc4dba1765fe4a981/stdlib/LinearAlgebra/src/structuredbroadcast.jl#L142

The reasoning is that assuming zero purity is a _huge_ win — often transforming complexity from O(n^2) to O(n). In the case of SymTridiagonal, going avoiding computing the symmetric elements goes from O(3n) to O(2n) and isn't as big of a win.

[StefanKarpinski]

I'm not saying we shouldn't do it because of that, just pointing it out.

[MasonProtter]

So as a proof of concept, I hacked together a (somewhat generic) version of map for Symmetric matrices using for loops instead of Generator(f, A) and possibly abusing Core.Compiler.return_type.

using LinearAlgebra, BenchmarkTools
mymap(ars...) = map(args...)
function mymap(f, xs::Symmetric)
    if xs.uplo == :U
        data = xs.data
    else
        data = transpose(xs.data)
    end
    out = similar(data, Core.Compiler.return_type(f, Tuple{eltype(data)}))
    is, js = axes(out)
    for i in is
        @inbounds for j in js[i:end]
            out[i, j] = f(xs[i, j])
            out[j, i] = out[i, j]
        end
    end
    out
end

Rough benchmarks:

A = rand(1:10, 20, 20);
A = A + transpose(A)
SA = Symmetric(A)

# test with a slow mapping function
foo(x) = Float64(factorial(big(x))/factorial(big(x + 1))) 

@btime map($foo, $A)     # 216.150 μs (5607 allocations: 197.05 KiB)
@btime map($foo, $SA)    # 218.232 μs (5607 allocations: 197.05 KiB)
@btime mymap($foo, $SA)  # 115.385 μs (2944 allocations: 104.99 KiB)

# test with a fast mapping function
bar(x) = x + 1
@btime map($bar, $A)    # 273.894 ns (2 allocations: 3.27 KiB)
@btime map($bar, $SA)   # 1.171 μs (2 allocations: 3.27 KiB)
@btime mymap($bar, $SA) # 486.349 ns (1 allocation: 3.25 KiB)

So it seems that this does have some overhead for fast mapping functions but that overhead is less than the current overhead for map(f, ::Symmetric). For slow mapping functions, this gives a dramatic speed improvement even for relatively small matrices.

Now, as one can see, I did this without using the interface that Base uses for map(f, A::AbstractArray) which is based on Generator(f, A). Would a PR for map(f, SA::Symmetric) need to use Generator? If so, does anyone have tips on how I might go about doing something similar with Generator?

Any other comments, perhaps a way I can do this without angering Jeff by using Core.Compiler.return_type?

[MasonProtter]

Bumping this issue as I'm still interested in making a PR with it.

Now, as one can see, I did this without using the interface that Base uses for map(f, A::AbstractArray) which is based on Generator(f, A). Would a PR for map(f, SA::Symmetric) need to use Generator? If so, does anyone have tips on how I might go about doing something similar with Generator?