Notebook

In [1]:

using BenchmarkTools
macro atime(expr) :(@btime $expr samples=1 evals=1) end

Out[1]:

@atime (macro with 1 method)

In [2]:

A = rand(0:9, 2, 3, 4, 5)
perm = (4, 3, 1, 2)
@code_warntype PermutedDimsArray(A, perm)

MethodInstance for PermutedDimsArray(::Array{Int64, 4}, ::NTuple{4, Int64})
  from PermutedDimsArray(data::AbstractArray{T, N}, perm) where {T, N} in Base.PermutedDimsArrays at permuteddimsarray.jl:42
Static Parameters
  T = Int64
  N = 4
Arguments
  #self#::Type{PermutedDimsArray}
  data::Array{Int64, 4}
  perm::NTuple{4, Int64}
Locals
  iperm::NTuple{4, Int64}
Body::PermutedDimsArray{Int64, 4, _A, _B, Array{Int64, 4}} where {_A, _B}
1 ─       Core.NewvarNode(:(iperm))
│   %2  = Base.PermutedDimsArrays.length(perm)::Core.Const(4)
│   %3  = (%2 == $(Expr(:static_parameter, 2)))::Core.Const(true)
│         Core.typeassert(%3, Core.Bool)
└──       goto #3
2 ─       Core.Const(:(Base.PermutedDimsArrays.string(perm, " is not a valid permutation of dimensions 1:", $(Expr(:static_parameter, 2)))))
│         Core.Const(:(Base.PermutedDimsArrays.ArgumentError(%6)))
└──       Core.Const(:(Base.PermutedDimsArrays.throw(%7)))
3 ┄       (iperm = Base.PermutedDimsArrays.invperm(perm))
│   %10 = $(Expr(:static_parameter, 1))::Core.Const(Int64)
│   %11 = $(Expr(:static_parameter, 2))::Core.Const(4)
│   %12 = Core._apply_iterate(Base.iterate, Core.tuple, perm)::NTuple{4, Int64}
│   %13 = Core._apply_iterate(Base.iterate, Core.tuple, iperm)::NTuple{4, Int64}
│   %14 = Base.PermutedDimsArrays.typeof(data)::Core.Const(Array{Int64, 4})
│   %15 = Core.apply_type(Base.PermutedDimsArrays.PermutedDimsArray, %10, %11, %12, %13, %14)::Type{PermutedDimsArray{Int64, 4, _A, _B, Array{Int64, 4}}} where {_A, _B}
│   %16 = (%15)(data)::PermutedDimsArray{Int64, 4, _A, _B, Array{Int64, 4}} where {_A, _B}
└──       return %16

In [3]:

@which PermutedDimsArray(zeros(1, 2, 3, 4), (4, 2, 1, 3))

Out[3]:

PermutedDimsArray(data::AbstractArray{T, N}, perm) where {T, N} in Base.PermutedDimsArrays at permuteddimsarray.jl:42

In [4]:

module O

struct PermutedDimsArray{T,N,AA<:AbstractArray{T,N},NT<:Dims{N}} <: AbstractArray{T,N}
    parent::AA
    perm::NT
    iperm::NT
end

function PermutedDimsArray(data::AbstractArray{T,N}, perm::Dims{N}) where {T,N}
    iperm = invperm(perm)
    PermutedDimsArray{T,N,typeof(data),typeof(perm)}(data, perm, iperm)
end
function PermutedDimsArray(data::AbstractArray{T,N}, perm) where {T,N}
    PermutedDimsArray(data, Dims{N}(perm))
end

Base.parent(A::PermutedDimsArray) = A.parent
Base.size(A::PermutedDimsArray{T,N}) where {T,N} = genperm(size(parent(A)), A.perm)
Base.axes(A::PermutedDimsArray{T,N}) where {T,N} = genperm(axes(parent(A)), A.perm)

Base.similar(A::PermutedDimsArray, T::Type, dims::Base.Dims) = similar(parent(A), T, dims)

Base.unsafe_convert(::Type{Ptr{T}}, A::PermutedDimsArray{T}) where {T} = Base.unsafe_convert(Ptr{T}, parent(A))

Base.pointer(A::PermutedDimsArray, i::Integer) = throw(ArgumentError("pointer(A, i) is deliberately unsupported for PermutedDimsArray"))

function Base.strides(A::PermutedDimsArray{T,N}) where {T,N}
    s = strides(parent(A))
    ntuple(d->s[A.perm[d]], Val(N))
end
Base.elsize(::Type{<:PermutedDimsArray{<:Any, <:Any, P}}) where {P} = Base.elsize(P)

@inline function Base.getindex(A::PermutedDimsArray{T,N}, I::Vararg{Int,N}) where {T,N}
    @boundscheck checkbounds(A, I...)
    @inbounds val = getindex(A.parent, genperm(I, A.iperm)...)
    val
end
@inline function Base.setindex!(A::PermutedDimsArray{T,N}, val, I::Vararg{Int,N}) where {T,N}
    @boundscheck checkbounds(A, I...)
    @inbounds setindex!(A.parent, val, genperm(I, A.iperm)...)
    val
end

@inline genperm(I::NTuple{N,Any}, perm::Dims{N}) where {N} = ntuple(d -> I[perm[d]], Val(N))
@inline genperm(I, perm::AbstractVector{Int}) = genperm(I, (perm...,))

end

Out[4]:

Main.O

In [5]:

A = rand(0:9, 2, 3, 4, 5)
perm = (4, 3, 1, 2)
@code_warntype O.PermutedDimsArray(A, perm)

MethodInstance for Main.O.PermutedDimsArray(::Array{Int64, 4}, ::NTuple{4, Int64})
  from Main.O.PermutedDimsArray(data::AbstractArray{T, N}, perm::Tuple{Vararg{Int64, N}}) where {T, N} in Main.O at In[4]:9
Static Parameters
  T = Int64
  N = 4
Arguments
  #self#::Type{Main.O.PermutedDimsArray}
  data::Array{Int64, 4}
  perm::NTuple{4, Int64}
Locals
  iperm::NTuple{4, Int64}
Body::Main.O.PermutedDimsArray{Int64, 4, Array{Int64, 4}, NTuple{4, Int64}}
1 ─      (iperm = Main.O.invperm(perm))
│   %2 = $(Expr(:static_parameter, 1))::Core.Const(Int64)
│   %3 = $(Expr(:static_parameter, 2))::Core.Const(4)
│   %4 = Main.O.typeof(data)::Core.Const(Array{Int64, 4})
│   %5 = Main.O.typeof(perm)::Core.Const(NTuple{4, Int64})
│   %6 = Core.apply_type(Main.O.PermutedDimsArray, %2, %3, %4, %5)::Core.Const(Main.O.PermutedDimsArray{Int64, 4, Array{Int64, 4}, NTuple{4, Int64}})
│   %7 = (%6)(data, perm, iperm)::Main.O.PermutedDimsArray{Int64, 4, Array{Int64, 4}, NTuple{4, Int64}}
└──      return %7

In [6]:

A = rand(2, 3, 4, 5)
perm = (4, 3, 1, 2)
B = @atime PermutedDimsArray($A, $perm)
C = @atime O.PermutedDimsArray($A, $perm)
B == C

  5.900 μs (4 allocations: 176 bytes)
  200.000 ns (0 allocations: 0 bytes)

Out[6]:

true

In [7]:

A = rand(2, 3, 4, 5)
perm = [4, 3, 1, 2]
B = @atime PermutedDimsArray($A, $perm)
C = @atime O.PermutedDimsArray($A, $perm)
B == C

  8.500 μs (5 allocations: 272 bytes)
  300.000 ns (0 allocations: 0 bytes)

Out[7]:

true

In [ ]: