PermutedDimsArray, and layout for dense 1st index (#12)

* add PermutedDimsArray, and FirstMajor layout

* change to UnitStride{D} for any dimension D

* use Compat.jl

* tidy & fix

* smarter ntuple, rm Compat.jl

* more tests

Co-authored-by: Michael Abbott <me@escbook>
Co-authored-by: Sheehan Olver <[email protected]>

Author: 3 people

src/ArrayLayouts.jl

@@ -48,7 +48,7 @@ else
 end     
 
 export materialize, materialize!, MulAdd, muladd!, Ldiv, Rdiv, Lmul, Rmul, lmul, rmul, ldiv, rdiv, mul, MemoryLayout, AbstractStridedLayout,
-        DenseColumnMajor, ColumnMajor, ZerosLayout, FillLayout, AbstractColumnMajor, RowMajor, AbstractRowMajor,
+        DenseColumnMajor, ColumnMajor, ZerosLayout, FillLayout, AbstractColumnMajor, RowMajor, AbstractRowMajor, UnitStride,
         DiagonalLayout, ScalarLayout, SymTridiagonalLayout, HermitianLayout, SymmetricLayout, TriangularLayout, 
         UnknownLayout, AbstractBandedLayout, ApplyBroadcastStyle, ConjLayout, AbstractFillLayout,
         colsupport, rowsupport, layout_getindex, QLayout, LayoutArray, LayoutMatrix, LayoutVector
@@ -213,4 +213,4 @@ end
     return A
 end
 
-end
+end

src/memorylayout.jl

@@ -18,6 +18,7 @@ abstract type AbstractRowMajor <: AbstractDecreasingStrides end
 struct DenseRowMajor <: AbstractRowMajor end
 struct RowMajor <: AbstractRowMajor end
 struct DecreasingStrides <: AbstractIncreasingStrides end
+struct UnitStride{D} <: AbstractStridedLayout end
 struct StridedLayout <: AbstractStridedLayout end
 struct ScalarLayout <: MemoryLayout end
 
@@ -42,6 +43,31 @@ dispatch to BLAS and LAPACK routines if the memory layout is BLAS compatible and
 the element type is a `Float32`, `Float64`, `ComplexF32`, or `ComplexF64`.
 In this case, one must implement the strided array interface, which requires
 overrides of `strides(A::MyMatrix)` and `unknown_convert(::Type{Ptr{T}}, A::MyMatrix)`.
+
+The complete list of more specialised types is as follows:
+```
+julia> using ArrayLayouts, AbstractTrees
+
+julia> AbstractTrees.children(x::Type) = subtypes(x)
+
+julia> print_tree(AbstractStridedLayout)
+AbstractStridedLayout
+├─ AbstractDecreasingStrides
+│  └─ AbstractRowMajor
+│     ├─ DenseRowMajor
+│     └─ RowMajor
+├─ AbstractIncreasingStrides
+│  ├─ AbstractColumnMajor
+│  │  ├─ ColumnMajor
+│  │  └─ DenseColumnMajor
+│  ├─ DecreasingStrides
+│  └─ IncreasingStrides
+├─ StridedLayout
+└─ UnitStride
+
+julia> Base.show_supertypes(AbstractStridedLayout)
+AbstractStridedLayout <: MemoryLayout <: Any
+```
 """
 AbstractStridedLayout
 
@@ -157,7 +183,7 @@ MemoryLayout(::Type{<:ReshapedArray{T,N,A,DIMS}}) where {T,N,A,DIMS} = reshapedl
 @inline reshapedlayout(::DenseColumnMajor, _) = DenseColumnMajor()
 
 
-@inline MemoryLayout(A::Type{<:SubArray{T,N,P,I}}) where {T,N,P,I} = 
+@inline MemoryLayout(A::Type{<:SubArray{T,N,P,I}}) where {T,N,P,I} =
     sublayout(MemoryLayout(P), I)
 sublayout(_1, _2) = UnknownLayout()
 sublayout(_1, _2, _3)= UnknownLayout()
@@ -257,6 +283,59 @@ transposelayout(::ConjLayout{ML}) where ML = ConjLayout{typeof(transposelayout(M
 adjointlayout(::Type{T}, M::MemoryLayout) where T = transposelayout(conjlayout(T, M))
 
 
+# Layouts of PermutedDimsArrays
+"""
+    UnitStride{D}()
+
+is returned by `MemoryLayout(A)` for arrays of `ndims(A) >= 3` which have `stride(A,D) == 1`.
+
+`UnitStride{1}` is weaker than `ColumnMajor` in that it does not demand that the other
+strides are increasing, hence it is not a subtype of `AbstractIncreasingStrides`.
+To ensure that `stride(A,1) == 1`, you may dispatch on `Union{UnitStride{1}, AbstractColumnMajor}`
+to allow for both options. (With complex numbers, you may also need their `ConjLayout` versions.)
+
+Likewise, both `UnitStride{ndims(A)}` and `AbstractRowMajor` have `stride(A, ndims(A)) == 1`.
+"""
+UnitStride
+
+MemoryLayout(::Type{PermutedDimsArray{T,N,P,Q,S}}) where {T,N,P,Q,S} = permutelayout(MemoryLayout(S), Val(P))
+
+permutelayout(::Any, perm) = UnknownLayout()
+permutelayout(::StridedLayout, perm) = StridedLayout()
+permutelayout(::ConjLayout{ML}, perm) where ML = ConjLayout{typeof(permutelayout(ML(), perm))}()
+
+function permutelayout(layout::AbstractColumnMajor, ::Val{perm}) where {perm}
+    issorted(perm) && return layout
+    issorted(reverse(perm)) && return reverse(layout)
+    D = sum(ntuple(dim -> perm[dim] == 1 ? dim : 0, length(perm)))
+    return UnitStride{D}()
+end
+function permutelayout(layout::AbstractRowMajor, ::Val{perm}) where {perm}
+    issorted(perm) && return layout
+    issorted(reverse(perm)) && return reverse(layout)
+    N = length(perm) # == ndims(A)
+    D = sum(ntuple(dim -> perm[dim] == N ? dim : 0, N))
+    return UnitStride{D}()
+end
+function permutelayout(layout::UnitStride{D0}, ::Val{perm}) where {D0, perm}
+    D = sum(ntuple(dim -> perm[dim] == D0 ? dim : 0, length(perm)))
+    return UnitStride{D}()
+end
+function permutelayout(layout::Union{IncreasingStrides,DecreasingStrides}, ::Val{perm}) where {perm}
+    issorted(perm) && return layout
+    issorted(reverse(perm)) && return reverse(layout)
+    return StridedLayout()
+end
+
+Base.reverse(::DenseRowMajor) = DenseColumnMajor()
+Base.reverse(::RowMajor) = ColumnMajor()
+Base.reverse(::DenseColumnMajor) = DenseRowMajor()
+Base.reverse(::ColumnMajor) = RowMajor()
+Base.reverse(::IncreasingStrides) = DecreasingStrides()
+Base.reverse(::DecreasingStrides) = IncreasingStrides()
+Base.reverse(::AbstractStridedLayout) = StridedLayout()
+
+
 # MemoryLayout of Symmetric/Hermitian
 """
     SymmetricLayout{layout}()

test/test_layouts.jl

@@ -1,9 +1,9 @@
 using ArrayLayouts, LinearAlgebra, FillArrays, Test
 import ArrayLayouts: MemoryLayout, DenseRowMajor, DenseColumnMajor, StridedLayout,
-                        ConjLayout, RowMajor, ColumnMajor, UnknownLayout,
+                        ConjLayout, RowMajor, ColumnMajor, UnitStride,
                         SymmetricLayout, HermitianLayout, UpperTriangularLayout,
                         UnitUpperTriangularLayout, LowerTriangularLayout,
-                        UnitLowerTriangularLayout, ScalarLayout,
+                        UnitLowerTriangularLayout, ScalarLayout, UnknownLayout,
                         hermitiandata, symmetricdata, FillLayout, ZerosLayout,
                         DiagonalLayout, colsupport, rowsupport
 
@@ -178,4 +178,58 @@ struct FooNumber <: Number end
         @test colsupport(LowerTriangular(A),3) ≡ 3:5
         @test rowsupport(LowerTriangular(A),3) ≡ Base.OneTo(3)
     end
+
+    @testset "PermutedDimsArray" begin
+        A = [1.0 2; 3 4]
+        @test MemoryLayout(PermutedDimsArray(A, (1,2))) == DenseColumnMajor()
+        @test MemoryLayout(PermutedDimsArray(A, (2,1))) == DenseRowMajor()
+        @test MemoryLayout(transpose(PermutedDimsArray(A, (2,1)))) == DenseColumnMajor()
+        @test MemoryLayout(adjoint(PermutedDimsArray(A, (2,1)))) == DenseColumnMajor()
+        B = [1.0+im 2; 3 4]
+        @test MemoryLayout(PermutedDimsArray(B, (2,1))) == DenseRowMajor()
+        @test MemoryLayout(transpose(PermutedDimsArray(B, (2,1)))) == DenseColumnMajor()
+        @test MemoryLayout(adjoint(PermutedDimsArray(B, (2,1)))) == ConjLayout{DenseColumnMajor}()
+
+        C = view(ones(10,20,30), 2:9, 3:18, 4:27);
+        @test MemoryLayout(C) == ColumnMajor()
+        @test MemoryLayout(PermutedDimsArray(C, (1,2,3))) == ColumnMajor()
+        @test MemoryLayout(PermutedDimsArray(C, (1,3,2))) == UnitStride{1}()
+
+        @test MemoryLayout(PermutedDimsArray(C, (3,1,2))) == UnitStride{2}()
+        @test MemoryLayout(PermutedDimsArray(C, (2,1,3))) == UnitStride{2}()
+
+        @test MemoryLayout(PermutedDimsArray(C, (3,2,1))) == RowMajor()
+        @test MemoryLayout(PermutedDimsArray(C, (2,3,1))) == UnitStride{3}()
+
+        revC = PermutedDimsArray(C, (3,2,1));
+        @test MemoryLayout(PermutedDimsArray(revC, (3,2,1))) == ColumnMajor()
+        @test MemoryLayout(PermutedDimsArray(revC, (3,1,2))) == UnitStride{1}()
+
+        D = ones(10,20,30,40);
+        @test MemoryLayout(D) == DenseColumnMajor()
+        @test MemoryLayout(PermutedDimsArray(D, (1,2,3,4))) == DenseColumnMajor()
+        @test MemoryLayout(PermutedDimsArray(D, (1,4,3,2))) == UnitStride{1}()
+
+        @test MemoryLayout(PermutedDimsArray(D, (4,1,3,2))) == UnitStride{2}()
+        @test MemoryLayout(PermutedDimsArray(D, (2,1,4,3))) == UnitStride{2}()
+
+        @test MemoryLayout(PermutedDimsArray(D, (4,3,2,1))) == DenseRowMajor()
+        @test MemoryLayout(PermutedDimsArray(D, (4,2,1,3))) == UnitStride{3}()
+
+        twoD = PermutedDimsArray(D, (3,1,2,4));
+        MemoryLayout(PermutedDimsArray(twoD, (2,1,4,3))) == UnitStride{1}()
+
+        revD = PermutedDimsArray(D, (4,3,2,1));
+        MemoryLayout(PermutedDimsArray(revD, (4,3,2,1))) == DenseColumnMajor()
+        MemoryLayout(PermutedDimsArray(revD, (4,2,3,1))) == UnitStride{1}()
+
+
+        issorted((1,2,3,4))
+        # Fails on Julia 1.4, in tests. Could use BenchmarkTools.@ballocated instead.
+        @test_skip 0 == @allocated issorted((1,2,3,4))
+        reverse((1,2,3,4))
+        @test_skip 0 == @allocated reverse((1,2,3,4))
+        MemoryLayout(revD)
+        @test 0 == @allocated MemoryLayout(revD)
+    end
 end