PermutedDimsArray, and layout for dense 1st index

Project.toml

@@ -1,13 +1,15 @@
 name = "ArrayLayouts"
 uuid = "4c555306-a7a7-4459-81d9-ec55ddd5c99a"
 authors = ["Sheehan Olver <[email protected]>"]
-version = "0.2.1"
+version = "0.2.2"
 
 [deps]
+Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"
 FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 
 [compat]
+Compat = "3.1"
 FillArrays = "0.8"
 julia = "1"
 

src/ArrayLayouts.jl

@@ -39,6 +39,8 @@ import LinearAlgebra.BLAS: BlasFloat, BlasReal, BlasComplex
 
 import FillArrays: AbstractFill, getindex_value
 
+import Compat: Compat # 3.1, for Base.filter(f, ::Tuple)
+
 if VERSION < v"1.2-"
     import Base: has_offset_axes
     require_one_based_indexing(A...) = !has_offset_axes(A...) || throw(ArgumentError("offset arrays are not supported but got an array with index other than 1"))
@@ -47,7 +49,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
@@ -187,4 +189,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,65 @@ 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 for which `stride(A,D) == 1`.
+
+`UnitStride{1}` is weaker than `ColumnMajor` in that it does not demand that the other
+strides are increasing, hence is is not a subtype of `AbstractIncreasingStrides`.
+To ensure that `stride(A,D) == 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)
+    tup = ntuple(length(perm)) do D
+        perm[D] == 1 && return UnitStride{D}()
+    end
+    only(filter(x -> x isa MemoryLayout, tup))
+end
+function permutelayout(layout::AbstractRowMajor, ::Val{perm}) where {perm}
+    issorted(perm) && return layout
+    issorted(reverse(perm)) && return reverse(layout)
+    N = length(perm) # == ndims(A)
+    tup = ntuple(N) do D
+        perm[D] == N && return UnitStride{D}()
+    end
+    only(filter(x -> x isa MemoryLayout, tup))
+end
+function permutelayout(layout::UnitStride{D0}, ::Val{perm}) where {D0, perm}
+    tup = ntuple(length(perm)) do D
+        perm[D] == D0 && return UnitStride{D}()
+    end
+    only(filter(x -> x isa MemoryLayout, tup))
+end
+function permutelayout(layout::T, ::Val{perm}) where {T <: Union{IncreasingStrides,DecreasingStrides}, 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
 
@@ -16,9 +16,9 @@ struct FooNumber <: Number end
         @test MemoryLayout(FooBar()) == MemoryLayout(FooBar) == UnknownLayout()
 
         A = randn(6)
-        @test MemoryLayout(typeof(A)) == MemoryLayout(typeof(Base.ReshapedArray(A,(2,3),()))) == 
+        @test MemoryLayout(typeof(A)) == MemoryLayout(typeof(Base.ReshapedArray(A,(2,3),()))) ==
             MemoryLayout(typeof(reinterpret(Float32,A))) == MemoryLayout(A) == DenseColumnMajor()
-        
+
         @test MemoryLayout(typeof(view(A,1:3))) == DenseColumnMajor()
         @test MemoryLayout(typeof(view(A,Base.OneTo(3)))) == DenseColumnMajor()
         @test MemoryLayout(typeof(view(A,:))) == DenseColumnMajor()
@@ -178,4 +178,46 @@ 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}()
+
+        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}()
+
+        issorted((1,2,3,4))
+        @test_skip 0 == @allocated issorted((1,2,3,4)) # fails on Julia 1.4, in tests
+        reverse((1,2,3,4))
+        @test_skip 0 == @allocated reverse((1,2,3,4))
+        revD = PermutedDimsArray(D, (4,3,2,1));
+        MemoryLayout(revD)
+        @test 0 == @allocated MemoryLayout(revD)
+    end
 end