Speeding up mapslices

General Usage Performance
, ,
1

Hello!

I was trying to perform a matrix inversion in a 3-D array along the third dimension (more info here) and reached two possible solutions.

First solution – Manual inversion

I am not worrying too much about strange types (the inverse of Int are Float64 and so on, but it could be easily solved).

# Create example array
S = (rand(4,4,1000) + 1im .* rand(4,4,1000)) ./ sqrt(2);

# Manual inversion
function f(A)
       B = Array{eltype(A),3}(undef,size(A))
       for i in 1:size(A,3)
           B[:,:,i] = inv(A[:,:,i])
       end
       return B
end

julia> @btime f($S);
    1.650 ms (5002 allocations: 5.05 MiB)

Second solution – mapslices 

julia> @btime mapslices(inv,$S; dims=(1,2));
    2.192 ms (11504 allocations: 4.91 MiB)

Problem

When I use @code_warntype I see there is an Any in the second solution, which may slow things down. Apart from it, there are twice as many allocations.

Is there a way to fix this in Julia internals? Is it on purpose?

Thanks!