Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Simplify dropdims(::Transpose) and insertdims(::PermutedDimsArray) etc. #55381

Open
wants to merge 1 commit into
base: master
Choose a base branch
from

Conversation

mcabbott
Copy link
Contributor

@mcabbott mcabbott commented Aug 5, 2024

Before, these wrap the wrapper in a ReshapedArray:

julia> dropdims([1,2,3]'; dims=1)
3-element reshape(adjoint(::Vector{Int64}), 3) with eltype Int64:
 1
 2
 3

julia> insertdims(PermutedDimsArray([1 2; 3 4], (2,1)); dims=2)
2×1×2 reshape(PermutedDimsArray(::Matrix{Int64}, (2, 1)), 2, 1, 2) with eltype Int64:
[:, :, 1] =
 1
 2

[:, :, 2] =
 3
 4

After:

julia> dropdims([1,2,3]'; dims=1)
3-element Vector{Int64}:
 1
 2
 3

julia> insertdims(PermutedDimsArray([1 2; 3 4], (2,1)); dims=2)
2×1×2 PermutedDimsArray(::Array{Int64, 3}, (2, 3, 1)) with eltype Int64:
[:, :, 1] =
 1
 2

[:, :, 2] =
 3
 4

This is not always faster, alone, but the hope is that the simpler return type will usually be faster downstream. (Especially with e.g. CuArrays, where wrapping twice causes them to miss specialised methods.)

julia> P13 = PermutedDimsArray(randn(2,1,3,1,4), (4,5,2,1,3));

julia> @btime dropdims($P13; dims=(1,3));
  189.628 ns (10 allocations: 192 bytes)  # before, ReshapedArray(PermutedDimsArray(Array{Float64, 5}))
  272.895 ns (10 allocations: 256 bytes)  # after, PermutedDimsArray(Array{Float64, 3})

julia> A13 = collect(P13);

julia> @btime dropdims($A13; dims=(1,3));
  172.333 ns (11 allocations: 240 bytes)

It appears that the compiler can predict the type when dims is constant. More stringent tests (or ways to do better) would be welcome:

julia> drop1(x) = dropdims(x; dims=(1,));

julia> @code_warntype drop1(P13)  # before
MethodInstance for drop1(::PermutedDimsArray{Float64, 5, (4, 5, 2, 1, 3), (4, 3, 5, 1, 2), Array{Float64, 5}})
  from drop1(x) @ Main REPL[12]:1
Arguments
  #self#::Core.Const(drop1)
  x::PermutedDimsArray{Float64, 5, (4, 5, 2, 1, 3), (4, 3, 5, 1, 2), Array{Float64, 5}}
Body::Base.ReshapedArray{Float64, 4, PermutedDimsArray{Float64, 5, (4, 5, 2, 1, 3), (4, 3, 5, 1, 2), Array{Float64, 5}}, NTuple{4, Base.MultiplicativeInverses.SignedMultiplicativeInverse{Int64}}}
1%1 = (:dims,)::Core.Const((:dims,))
...

julia> @code_warntype drop1(P13)  # after
MethodInstance for drop1(::PermutedDimsArray{Float64, 5, (4, 5, 2, 1, 3), (4, 3, 5, 1, 2), Array{Float64, 5}})
  from drop1(x) @ Main REPL[240]:1
Arguments
  #self#::Core.Const(drop1)
  x::PermutedDimsArray{Float64, 5, (4, 5, 2, 1, 3), (4, 3, 5, 1, 2), Array{Float64, 5}}
Body::PermutedDimsArray{Float64, 4, (4, 2, 1, 3), (3, 2, 4, 1), Array{Float64, 4}}
1%1 = (:dims,)::Core.Const((:dims,))
...

@mcabbott mcabbott added the arrays [a, r, r, a, y, s] label Aug 5, 2024
@mcabbott mcabbott requested a review from nsajko August 23, 2024 17:08
@mcabbott mcabbott requested a review from mbauman August 23, 2024 17:10
@mcabbott
Copy link
Contributor Author

Xref #45793 which added insertdims. And I requested review from a few of those who commented on that PR.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
arrays [a, r, r, a, y, s]
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants