From 446a7f05d17833f0d0f37a9c3fccc6d0ba0308fb Mon Sep 17 00:00:00 2001 From: mtfishman Date: Fri, 15 Nov 2024 11:01:48 -0500 Subject: [PATCH 1/5] [NDTensors] Introduce RecursivePermutedDimsArrays submodule --- NDTensors/Project.toml | 2 +- NDTensors/src/imports.jl | 1 + .../src/RecursivePermutedDimsArrays.jl | 262 ++++++++++++++++++ .../test/Project.toml | 2 + .../test/runtests.jl | 23 ++ NDTensors/test/lib/runtests.jl | 1 + 6 files changed, 290 insertions(+), 1 deletion(-) create mode 100644 NDTensors/src/lib/RecursivePermutedDimsArrays/src/RecursivePermutedDimsArrays.jl create mode 100644 NDTensors/src/lib/RecursivePermutedDimsArrays/test/Project.toml create mode 100644 NDTensors/src/lib/RecursivePermutedDimsArrays/test/runtests.jl diff --git a/NDTensors/Project.toml b/NDTensors/Project.toml index 94626cdddb..80f68c17f0 100644 --- a/NDTensors/Project.toml +++ b/NDTensors/Project.toml @@ -1,7 +1,7 @@ name = "NDTensors" uuid = "23ae76d9-e61a-49c4-8f12-3f1a16adf9cf" authors = ["Matthew Fishman "] -version = "0.3.67" +version = "0.3.68" [deps] Accessors = "7d9f7c33-5ae7-4f3b-8dc6-eff91059b697" diff --git a/NDTensors/src/imports.jl b/NDTensors/src/imports.jl index 73d8ae3771..f528954602 100644 --- a/NDTensors/src/imports.jl +++ b/NDTensors/src/imports.jl @@ -40,6 +40,7 @@ for lib in [ :SparseArrayInterface, :SparseArrayDOKs, :DiagonalArrays, + :RecursivePermutedDimsArrays, :BlockSparseArrays, :NamedDimsArrays, :SmallVectors, diff --git a/NDTensors/src/lib/RecursivePermutedDimsArrays/src/RecursivePermutedDimsArrays.jl b/NDTensors/src/lib/RecursivePermutedDimsArrays/src/RecursivePermutedDimsArrays.jl new file mode 100644 index 0000000000..30f2a83f97 --- /dev/null +++ b/NDTensors/src/lib/RecursivePermutedDimsArrays/src/RecursivePermutedDimsArrays.jl @@ -0,0 +1,262 @@ +# Mostly copied from https://github.com/JuliaLang/julia/blob/master/base/permuteddimsarray.jl +# Like `PermutedDimsArrays` but recursive, similar to `Adjoint` and `Transpose`. +module RecursivePermutedDimsArrays + +import Base: permutedims, permutedims! +export RecursivePermutedDimsArray + +# Some day we will want storage-order-aware iteration, so put perm in the parameters +struct RecursivePermutedDimsArray{T,N,perm,iperm,AA<:AbstractArray} <: AbstractArray{T,N} + parent::AA + + function RecursivePermutedDimsArray{T,N,perm,iperm,AA}( + data::AA + ) where {T,N,perm,iperm,AA<:AbstractArray} + (isa(perm, NTuple{N,Int}) && isa(iperm, NTuple{N,Int})) || + error("perm and iperm must both be NTuple{$N,Int}") + isperm(perm) || + throw(ArgumentError(string(perm, " is not a valid permutation of dimensions 1:", N))) + all(d -> iperm[perm[d]] == d, 1:N) || + throw(ArgumentError(string(perm, " and ", iperm, " must be inverses"))) + return new(data) + end +end + +""" + RecursivePermutedDimsArray(A, perm) -> B + +Given an AbstractArray `A`, create a view `B` such that the +dimensions appear to be permuted. Similar to `permutedims`, except +that no copying occurs (`B` shares storage with `A`). + +See also [`permutedims`](@ref), [`invperm`](@ref). + +# Examples +```jldoctest +julia> A = rand(3,5,4); + +julia> B = RecursivePermutedDimsArray(A, (3,1,2)); + +julia> size(B) +(4, 3, 5) + +julia> B[3,1,2] == A[1,2,3] +true +``` +""" +Base.@constprop :aggressive function RecursivePermutedDimsArray( + data::AbstractArray{T,N}, perm +) where {T,N} + length(perm) == N || + throw(ArgumentError(string(perm, " is not a valid permutation of dimensions 1:", N))) + iperm = invperm(perm) + return RecursivePermutedDimsArray{ + recursivepermuteddimsarraytype(T, perm),N,(perm...,),(iperm...,),typeof(data) + }( + data + ) +end + +# Ideally we would use `Base.promote_op(recursivepermuteddimsarray, type, perm)` but +# that doesn't seem to preserve the `perm`/`iperm` type parameters. +function recursivepermuteddimsarraytype(type::Type{<:AbstractArray{<:AbstractArray}}, perm) + return RecursivePermutedDimsArray{ + recursivepermuteddimsarraytype(eltype(type), perm),ndims(type),perm,invperm(perm),type + } +end +function recursivepermuteddimsarraytype(type::Type{<:AbstractArray}, perm) + return PermutedDimsArray{eltype(type),ndims(type),perm,invperm(perm),type} +end +recursivepermuteddimsarraytype(type::Type, perm) = type + +function recursivepermuteddimsarray(A::AbstractArray{<:AbstractArray}, perm) + return RecursivePermutedDimsArray(A, perm) +end +recursivepermuteddimsarray(A::AbstractArray, perm) = PermutedDimsArray(A, perm) +# By default, assume scalar and don't permute. +recursivepermuteddimsarray(x, perm) = x + +Base.parent(A::RecursivePermutedDimsArray) = A.parent +function Base.size(A::RecursivePermutedDimsArray{T,N,perm}) where {T,N,perm} + return genperm(size(parent(A)), perm) +end +function Base.axes(A::RecursivePermutedDimsArray{T,N,perm}) where {T,N,perm} + return genperm(axes(parent(A)), perm) +end +Base.has_offset_axes(A::RecursivePermutedDimsArray) = Base.has_offset_axes(A.parent) +function Base.similar(A::RecursivePermutedDimsArray, T::Type, dims::Base.Dims) + return similar(parent(A), T, dims) +end +function Base.cconvert(::Type{Ptr{T}}, A::RecursivePermutedDimsArray{T}) where {T} + return Base.cconvert(Ptr{T}, parent(A)) +end + +# It's OK to return a pointer to the first element, and indeed quite +# useful for wrapping C routines that require a different storage +# order than used by Julia. But for an array with unconventional +# storage order, a linear offset is ambiguous---is it a memory offset +# or a linear index? +function Base.pointer(A::RecursivePermutedDimsArray, i::Integer) + throw( + ArgumentError( + "pointer(A, i) is deliberately unsupported for RecursivePermutedDimsArray" + ), + ) +end + +function Base.strides(A::RecursivePermutedDimsArray{T,N,perm}) where {T,N,perm} + s = strides(parent(A)) + return ntuple(d -> s[perm[d]], Val(N)) +end +function Base.elsize( + ::Type{<:RecursivePermutedDimsArray{<:Any,<:Any,<:Any,<:Any,P}} +) where {P} + return Base.elsize(P) +end + +@inline function Base.getindex( + A::RecursivePermutedDimsArray{T,N,perm,iperm}, I::Vararg{Int,N} +) where {T,N,perm,iperm} + @boundscheck checkbounds(A, I...) + @inbounds val = recursivepermuteddimsarray(getindex(A.parent, genperm(I, iperm)...), perm) + return val +end +@inline function Base.setindex!( + A::RecursivePermutedDimsArray{T,N,perm,iperm}, val, I::Vararg{Int,N} +) where {T,N,perm,iperm} + @boundscheck checkbounds(A, I...) + @inbounds setindex!(A.parent, recursivepermuteddimsarray(val, perm), genperm(I, iperm)...) + return val +end + +function Base.isassigned( + A::RecursivePermutedDimsArray{T,N,perm,iperm}, I::Vararg{Int,N} +) where {T,N,perm,iperm} + @boundscheck checkbounds(Bool, A, I...) || return false + @inbounds x = isassigned(A.parent, genperm(I, iperm)...) + return x +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...,)) + +function Base.copyto!( + dest::RecursivePermutedDimsArray{T,N}, src::AbstractArray{T,N} +) where {T,N} + checkbounds(dest, axes(src)...) + return _copy!(dest, src) +end +Base.copyto!(dest::RecursivePermutedDimsArray, src::AbstractArray) = _copy!(dest, src) + +function _copy!(P::RecursivePermutedDimsArray{T,N,perm}, src) where {T,N,perm} + # If dest/src are "close to dense," then it pays to be cache-friendly. + # Determine the first permuted dimension + d = 0 # d+1 will hold the first permuted dimension of src + while d < ndims(src) && perm[d + 1] == d + 1 + d += 1 + end + if d == ndims(src) + copyto!(parent(P), src) # it's not permuted + else + R1 = CartesianIndices(axes(src)[1:d]) + d1 = findfirst(isequal(d + 1), perm)::Int # first permuted dim of dest + R2 = CartesianIndices(axes(src)[(d + 2):(d1 - 1)]) + R3 = CartesianIndices(axes(src)[(d1 + 1):end]) + _permutedims!(P, src, R1, R2, R3, d + 1, d1) + end + return P +end + +@noinline function _permutedims!( + P::RecursivePermutedDimsArray, src, R1::CartesianIndices{0}, R2, R3, ds, dp +) + ip, is = axes(src, dp), axes(src, ds) + for jo in first(ip):8:last(ip), io in first(is):8:last(is) + for I3 in R3, I2 in R2 + for j in jo:min(jo + 7, last(ip)) + for i in io:min(io + 7, last(is)) + @inbounds P[i, I2, j, I3] = src[i, I2, j, I3] + end + end + end + end + return P +end + +@noinline function _permutedims!(P::RecursivePermutedDimsArray, src, R1, R2, R3, ds, dp) + ip, is = axes(src, dp), axes(src, ds) + for jo in first(ip):8:last(ip), io in first(is):8:last(is) + for I3 in R3, I2 in R2 + for j in jo:min(jo + 7, last(ip)) + for i in io:min(io + 7, last(is)) + for I1 in R1 + @inbounds P[I1, i, I2, j, I3] = src[I1, i, I2, j, I3] + end + end + end + end + end + return P +end + +const CommutativeOps = Union{ + typeof(+), + typeof(Base.add_sum), + typeof(min), + typeof(max), + typeof(Base._extrema_rf), + typeof(|), + typeof(&), +} + +function Base._mapreduce_dim( + f, + op::CommutativeOps, + init::Base._InitialValue, + A::RecursivePermutedDimsArray, + dims::Colon, +) + return Base._mapreduce_dim(f, op, init, parent(A), dims) +end +function Base._mapreduce_dim( + f::typeof(identity), + op::Union{typeof(Base.mul_prod),typeof(*)}, + init::Base._InitialValue, + A::RecursivePermutedDimsArray{<:Union{Real,Complex}}, + dims::Colon, +) + return Base._mapreduce_dim(f, op, init, parent(A), dims) +end + +function Base.mapreducedim!( + f, + op::CommutativeOps, + B::AbstractArray{T,N}, + A::RecursivePermutedDimsArray{S,N,perm,iperm}, +) where {T,S,N,perm,iperm} + C = RecursivePermutedDimsArray{T,N,iperm,perm,typeof(B)}(B) # make the inverse permutation for the output + Base.mapreducedim!(f, op, C, parent(A)) + return B +end +function Base.mapreducedim!( + f::typeof(identity), + op::Union{typeof(Base.mul_prod),typeof(*)}, + B::AbstractArray{T,N}, + A::RecursivePermutedDimsArray{<:Union{Real,Complex},N,perm,iperm}, +) where {T,N,perm,iperm} + C = RecursivePermutedDimsArray{T,N,iperm,perm,typeof(B)}(B) # make the inverse permutation for the output + Base.mapreducedim!(f, op, C, parent(A)) + return B +end + +function Base.showarg( + io::IO, A::RecursivePermutedDimsArray{T,N,perm}, toplevel +) where {T,N,perm} + print(io, "RecursivePermutedDimsArray(") + Base.showarg(io, parent(A), false) + print(io, ", ", perm, ')') + toplevel && print(io, " with eltype ", eltype(A)) + return nothing +end + +end diff --git a/NDTensors/src/lib/RecursivePermutedDimsArrays/test/Project.toml b/NDTensors/src/lib/RecursivePermutedDimsArrays/test/Project.toml new file mode 100644 index 0000000000..9b1d5ccd25 --- /dev/null +++ b/NDTensors/src/lib/RecursivePermutedDimsArrays/test/Project.toml @@ -0,0 +1,2 @@ +[deps] +NDTensors = "23ae76d9-e61a-49c4-8f12-3f1a16adf9cf" diff --git a/NDTensors/src/lib/RecursivePermutedDimsArrays/test/runtests.jl b/NDTensors/src/lib/RecursivePermutedDimsArrays/test/runtests.jl new file mode 100644 index 0000000000..f2fa4348b4 --- /dev/null +++ b/NDTensors/src/lib/RecursivePermutedDimsArrays/test/runtests.jl @@ -0,0 +1,23 @@ +@eval module $(gensym()) +using NDTensors.RecursivePermutedDimsArrays: RecursivePermutedDimsArray +using Test: @test, @testset +@testset "RecursivePermutedDimsArrays" for elt in ( + Float32, Float64, Complex{Float32}, Complex{Float64} +) + a = map(_ -> randn(elt, 2, 3, 4), CartesianIndices((2, 3, 4))) + perm = (3, 2, 1) + p = RecursivePermutedDimsArray(a, perm) + T = PermutedDimsArray{elt,3,perm,invperm(perm),eltype(a)} + @test typeof(p) === RecursivePermutedDimsArray{T,3,perm,invperm(perm),typeof(a)} + @test size(p) == (4, 3, 2) + @test eltype(p) === T + for I in eachindex(p) + @test size(p[I]) == (4, 3, 2) + @test p[I] == permutedims(a[CartesianIndex(reverse(Tuple(I)))], perm) + end + x = randn(elt, 4, 3, 2) + p[3, 2, 1] = x + @test p[3, 2, 1] == x + @test a[1, 2, 3] == permutedims(x, perm) +end +end diff --git a/NDTensors/test/lib/runtests.jl b/NDTensors/test/lib/runtests.jl index b1b47cf866..f89a282dc8 100644 --- a/NDTensors/test/lib/runtests.jl +++ b/NDTensors/test/lib/runtests.jl @@ -15,6 +15,7 @@ using Test: @testset "LabelledNumbers", "MetalExtensions", "NamedDimsArrays", + "RecursivePermutedDimsArrays", "SmallVectors", "SortedSets", "SparseArrayDOKs", From 2150d93e0f2877ce9c0289aed84ce1b927735457 Mon Sep 17 00:00:00 2001 From: mtfishman Date: Fri, 15 Nov 2024 11:30:32 -0500 Subject: [PATCH 2/5] Rename RecursivePermutedDimsArrays to NestedPermutedDimsArrays, make singly nested --- NDTensors/src/imports.jl | 2 +- .../src/NestedPermutedDimsArrays.jl} | 108 ++++++++---------- .../test/Project.toml | 0 .../test/runtests.jl | 8 +- NDTensors/test/lib/runtests.jl | 2 +- 5 files changed, 52 insertions(+), 68 deletions(-) rename NDTensors/src/lib/{RecursivePermutedDimsArrays/src/RecursivePermutedDimsArrays.jl => NestedPermutedDimsArrays/src/NestedPermutedDimsArrays.jl} (58%) rename NDTensors/src/lib/{RecursivePermutedDimsArrays => NestedPermutedDimsArrays}/test/Project.toml (100%) rename NDTensors/src/lib/{RecursivePermutedDimsArrays => NestedPermutedDimsArrays}/test/runtests.jl (69%) diff --git a/NDTensors/src/imports.jl b/NDTensors/src/imports.jl index f528954602..9869df6448 100644 --- a/NDTensors/src/imports.jl +++ b/NDTensors/src/imports.jl @@ -40,7 +40,7 @@ for lib in [ :SparseArrayInterface, :SparseArrayDOKs, :DiagonalArrays, - :RecursivePermutedDimsArrays, + :NestedPermutedDimsArrays, :BlockSparseArrays, :NamedDimsArrays, :SmallVectors, diff --git a/NDTensors/src/lib/RecursivePermutedDimsArrays/src/RecursivePermutedDimsArrays.jl b/NDTensors/src/lib/NestedPermutedDimsArrays/src/NestedPermutedDimsArrays.jl similarity index 58% rename from NDTensors/src/lib/RecursivePermutedDimsArrays/src/RecursivePermutedDimsArrays.jl rename to NDTensors/src/lib/NestedPermutedDimsArrays/src/NestedPermutedDimsArrays.jl index 30f2a83f97..1f5304ecb1 100644 --- a/NDTensors/src/lib/RecursivePermutedDimsArrays/src/RecursivePermutedDimsArrays.jl +++ b/NDTensors/src/lib/NestedPermutedDimsArrays/src/NestedPermutedDimsArrays.jl @@ -1,15 +1,16 @@ # Mostly copied from https://github.com/JuliaLang/julia/blob/master/base/permuteddimsarray.jl -# Like `PermutedDimsArrays` but recursive, similar to `Adjoint` and `Transpose`. -module RecursivePermutedDimsArrays +# Like `PermutedDimsArrays` but singly nested, similar to `Adjoint` and `Transpose` +# (though those are fully recursive). +module NestedPermutedDimsArrays import Base: permutedims, permutedims! -export RecursivePermutedDimsArray +export NestedPermutedDimsArray # Some day we will want storage-order-aware iteration, so put perm in the parameters -struct RecursivePermutedDimsArray{T,N,perm,iperm,AA<:AbstractArray} <: AbstractArray{T,N} +struct NestedPermutedDimsArray{T,N,perm,iperm,AA<:AbstractArray} <: AbstractArray{T,N} parent::AA - function RecursivePermutedDimsArray{T,N,perm,iperm,AA}( + function NestedPermutedDimsArray{T,N,perm,iperm,AA}( data::AA ) where {T,N,perm,iperm,AA<:AbstractArray} (isa(perm, NTuple{N,Int}) && isa(iperm, NTuple{N,Int})) || @@ -23,7 +24,7 @@ struct RecursivePermutedDimsArray{T,N,perm,iperm,AA<:AbstractArray} <: AbstractA end """ - RecursivePermutedDimsArray(A, perm) -> B + NestedPermutedDimsArray(A, perm) -> B Given an AbstractArray `A`, create a view `B` such that the dimensions appear to be permuted. Similar to `permutedims`, except @@ -35,7 +36,7 @@ See also [`permutedims`](@ref), [`invperm`](@ref). ```jldoctest julia> A = rand(3,5,4); -julia> B = RecursivePermutedDimsArray(A, (3,1,2)); +julia> B = NestedPermutedDimsArray(A, (3,1,2)); julia> size(B) (4, 3, 5) @@ -44,50 +45,44 @@ julia> B[3,1,2] == A[1,2,3] true ``` """ -Base.@constprop :aggressive function RecursivePermutedDimsArray( +Base.@constprop :aggressive function NestedPermutedDimsArray( data::AbstractArray{T,N}, perm ) where {T,N} length(perm) == N || throw(ArgumentError(string(perm, " is not a valid permutation of dimensions 1:", N))) iperm = invperm(perm) - return RecursivePermutedDimsArray{ - recursivepermuteddimsarraytype(T, perm),N,(perm...,),(iperm...,),typeof(data) + return NestedPermutedDimsArray{ + maybe_permuteddimsarraytype(T, perm),N,(perm...,),(iperm...,),typeof(data) }( data ) end -# Ideally we would use `Base.promote_op(recursivepermuteddimsarray, type, perm)` but -# that doesn't seem to preserve the `perm`/`iperm` type parameters. -function recursivepermuteddimsarraytype(type::Type{<:AbstractArray{<:AbstractArray}}, perm) - return RecursivePermutedDimsArray{ - recursivepermuteddimsarraytype(eltype(type), perm),ndims(type),perm,invperm(perm),type - } -end -function recursivepermuteddimsarraytype(type::Type{<:AbstractArray}, perm) +# Ideally would use `Base.promote_op(maybe_permuteddimsarraytype, type, perm)` +# but it doesn't handle `perm` properly. +function maybe_permuteddimsarraytype(type::Type{<:AbstractArray}, perm) return PermutedDimsArray{eltype(type),ndims(type),perm,invperm(perm),type} end -recursivepermuteddimsarraytype(type::Type, perm) = type +maybe_permuteddimsarraytype(type::Type, perm) = type -function recursivepermuteddimsarray(A::AbstractArray{<:AbstractArray}, perm) - return RecursivePermutedDimsArray(A, perm) +function maybe_permuteddimsarray(A::AbstractArray, perm) + return PermutedDimsArray(A, perm) end -recursivepermuteddimsarray(A::AbstractArray, perm) = PermutedDimsArray(A, perm) # By default, assume scalar and don't permute. -recursivepermuteddimsarray(x, perm) = x +maybe_permuteddimsarray(x, perm) = x -Base.parent(A::RecursivePermutedDimsArray) = A.parent -function Base.size(A::RecursivePermutedDimsArray{T,N,perm}) where {T,N,perm} +Base.parent(A::NestedPermutedDimsArray) = A.parent +function Base.size(A::NestedPermutedDimsArray{T,N,perm}) where {T,N,perm} return genperm(size(parent(A)), perm) end -function Base.axes(A::RecursivePermutedDimsArray{T,N,perm}) where {T,N,perm} +function Base.axes(A::NestedPermutedDimsArray{T,N,perm}) where {T,N,perm} return genperm(axes(parent(A)), perm) end -Base.has_offset_axes(A::RecursivePermutedDimsArray) = Base.has_offset_axes(A.parent) -function Base.similar(A::RecursivePermutedDimsArray, T::Type, dims::Base.Dims) +Base.has_offset_axes(A::NestedPermutedDimsArray) = Base.has_offset_axes(A.parent) +function Base.similar(A::NestedPermutedDimsArray, T::Type, dims::Base.Dims) return similar(parent(A), T, dims) end -function Base.cconvert(::Type{Ptr{T}}, A::RecursivePermutedDimsArray{T}) where {T} +function Base.cconvert(::Type{Ptr{T}}, A::NestedPermutedDimsArray{T}) where {T} return Base.cconvert(Ptr{T}, parent(A)) end @@ -96,41 +91,37 @@ end # order than used by Julia. But for an array with unconventional # storage order, a linear offset is ambiguous---is it a memory offset # or a linear index? -function Base.pointer(A::RecursivePermutedDimsArray, i::Integer) +function Base.pointer(A::NestedPermutedDimsArray, i::Integer) throw( - ArgumentError( - "pointer(A, i) is deliberately unsupported for RecursivePermutedDimsArray" - ), + ArgumentError("pointer(A, i) is deliberately unsupported for NestedPermutedDimsArray") ) end -function Base.strides(A::RecursivePermutedDimsArray{T,N,perm}) where {T,N,perm} +function Base.strides(A::NestedPermutedDimsArray{T,N,perm}) where {T,N,perm} s = strides(parent(A)) return ntuple(d -> s[perm[d]], Val(N)) end -function Base.elsize( - ::Type{<:RecursivePermutedDimsArray{<:Any,<:Any,<:Any,<:Any,P}} -) where {P} +function Base.elsize(::Type{<:NestedPermutedDimsArray{<:Any,<:Any,<:Any,<:Any,P}}) where {P} return Base.elsize(P) end @inline function Base.getindex( - A::RecursivePermutedDimsArray{T,N,perm,iperm}, I::Vararg{Int,N} + A::NestedPermutedDimsArray{T,N,perm,iperm}, I::Vararg{Int,N} ) where {T,N,perm,iperm} @boundscheck checkbounds(A, I...) - @inbounds val = recursivepermuteddimsarray(getindex(A.parent, genperm(I, iperm)...), perm) + @inbounds val = maybe_permuteddimsarray(getindex(A.parent, genperm(I, iperm)...), perm) return val end @inline function Base.setindex!( - A::RecursivePermutedDimsArray{T,N,perm,iperm}, val, I::Vararg{Int,N} + A::NestedPermutedDimsArray{T,N,perm,iperm}, val, I::Vararg{Int,N} ) where {T,N,perm,iperm} @boundscheck checkbounds(A, I...) - @inbounds setindex!(A.parent, recursivepermuteddimsarray(val, perm), genperm(I, iperm)...) + @inbounds setindex!(A.parent, maybe_permuteddimsarray(val, perm), genperm(I, iperm)...) return val end function Base.isassigned( - A::RecursivePermutedDimsArray{T,N,perm,iperm}, I::Vararg{Int,N} + A::NestedPermutedDimsArray{T,N,perm,iperm}, I::Vararg{Int,N} ) where {T,N,perm,iperm} @boundscheck checkbounds(Bool, A, I...) || return false @inbounds x = isassigned(A.parent, genperm(I, iperm)...) @@ -141,14 +132,14 @@ end @inline genperm(I, perm::AbstractVector{Int}) = genperm(I, (perm...,)) function Base.copyto!( - dest::RecursivePermutedDimsArray{T,N}, src::AbstractArray{T,N} + dest::NestedPermutedDimsArray{T,N}, src::AbstractArray{T,N} ) where {T,N} checkbounds(dest, axes(src)...) return _copy!(dest, src) end -Base.copyto!(dest::RecursivePermutedDimsArray, src::AbstractArray) = _copy!(dest, src) +Base.copyto!(dest::NestedPermutedDimsArray, src::AbstractArray) = _copy!(dest, src) -function _copy!(P::RecursivePermutedDimsArray{T,N,perm}, src) where {T,N,perm} +function _copy!(P::NestedPermutedDimsArray{T,N,perm}, src) where {T,N,perm} # If dest/src are "close to dense," then it pays to be cache-friendly. # Determine the first permuted dimension d = 0 # d+1 will hold the first permuted dimension of src @@ -168,7 +159,7 @@ function _copy!(P::RecursivePermutedDimsArray{T,N,perm}, src) where {T,N,perm} end @noinline function _permutedims!( - P::RecursivePermutedDimsArray, src, R1::CartesianIndices{0}, R2, R3, ds, dp + P::NestedPermutedDimsArray, src, R1::CartesianIndices{0}, R2, R3, ds, dp ) ip, is = axes(src, dp), axes(src, ds) for jo in first(ip):8:last(ip), io in first(is):8:last(is) @@ -183,7 +174,7 @@ end return P end -@noinline function _permutedims!(P::RecursivePermutedDimsArray, src, R1, R2, R3, ds, dp) +@noinline function _permutedims!(P::NestedPermutedDimsArray, src, R1, R2, R3, ds, dp) ip, is = axes(src, dp), axes(src, ds) for jo in first(ip):8:last(ip), io in first(is):8:last(is) for I3 in R3, I2 in R2 @@ -210,11 +201,7 @@ const CommutativeOps = Union{ } function Base._mapreduce_dim( - f, - op::CommutativeOps, - init::Base._InitialValue, - A::RecursivePermutedDimsArray, - dims::Colon, + f, op::CommutativeOps, init::Base._InitialValue, A::NestedPermutedDimsArray, dims::Colon ) return Base._mapreduce_dim(f, op, init, parent(A), dims) end @@ -222,19 +209,16 @@ function Base._mapreduce_dim( f::typeof(identity), op::Union{typeof(Base.mul_prod),typeof(*)}, init::Base._InitialValue, - A::RecursivePermutedDimsArray{<:Union{Real,Complex}}, + A::NestedPermutedDimsArray{<:Union{Real,Complex}}, dims::Colon, ) return Base._mapreduce_dim(f, op, init, parent(A), dims) end function Base.mapreducedim!( - f, - op::CommutativeOps, - B::AbstractArray{T,N}, - A::RecursivePermutedDimsArray{S,N,perm,iperm}, + f, op::CommutativeOps, B::AbstractArray{T,N}, A::NestedPermutedDimsArray{S,N,perm,iperm} ) where {T,S,N,perm,iperm} - C = RecursivePermutedDimsArray{T,N,iperm,perm,typeof(B)}(B) # make the inverse permutation for the output + C = NestedPermutedDimsArray{T,N,iperm,perm,typeof(B)}(B) # make the inverse permutation for the output Base.mapreducedim!(f, op, C, parent(A)) return B end @@ -242,17 +226,17 @@ function Base.mapreducedim!( f::typeof(identity), op::Union{typeof(Base.mul_prod),typeof(*)}, B::AbstractArray{T,N}, - A::RecursivePermutedDimsArray{<:Union{Real,Complex},N,perm,iperm}, + A::NestedPermutedDimsArray{<:Union{Real,Complex},N,perm,iperm}, ) where {T,N,perm,iperm} - C = RecursivePermutedDimsArray{T,N,iperm,perm,typeof(B)}(B) # make the inverse permutation for the output + C = NestedPermutedDimsArray{T,N,iperm,perm,typeof(B)}(B) # make the inverse permutation for the output Base.mapreducedim!(f, op, C, parent(A)) return B end function Base.showarg( - io::IO, A::RecursivePermutedDimsArray{T,N,perm}, toplevel + io::IO, A::NestedPermutedDimsArray{T,N,perm}, toplevel ) where {T,N,perm} - print(io, "RecursivePermutedDimsArray(") + print(io, "NestedPermutedDimsArray(") Base.showarg(io, parent(A), false) print(io, ", ", perm, ')') toplevel && print(io, " with eltype ", eltype(A)) diff --git a/NDTensors/src/lib/RecursivePermutedDimsArrays/test/Project.toml b/NDTensors/src/lib/NestedPermutedDimsArrays/test/Project.toml similarity index 100% rename from NDTensors/src/lib/RecursivePermutedDimsArrays/test/Project.toml rename to NDTensors/src/lib/NestedPermutedDimsArrays/test/Project.toml diff --git a/NDTensors/src/lib/RecursivePermutedDimsArrays/test/runtests.jl b/NDTensors/src/lib/NestedPermutedDimsArrays/test/runtests.jl similarity index 69% rename from NDTensors/src/lib/RecursivePermutedDimsArrays/test/runtests.jl rename to NDTensors/src/lib/NestedPermutedDimsArrays/test/runtests.jl index f2fa4348b4..d13b881540 100644 --- a/NDTensors/src/lib/RecursivePermutedDimsArrays/test/runtests.jl +++ b/NDTensors/src/lib/NestedPermutedDimsArrays/test/runtests.jl @@ -1,14 +1,14 @@ @eval module $(gensym()) -using NDTensors.RecursivePermutedDimsArrays: RecursivePermutedDimsArray +using NDTensors.NestedPermutedDimsArrays: NestedPermutedDimsArray using Test: @test, @testset -@testset "RecursivePermutedDimsArrays" for elt in ( +@testset "NestedPermutedDimsArrays" for elt in ( Float32, Float64, Complex{Float32}, Complex{Float64} ) a = map(_ -> randn(elt, 2, 3, 4), CartesianIndices((2, 3, 4))) perm = (3, 2, 1) - p = RecursivePermutedDimsArray(a, perm) + p = NestedPermutedDimsArray(a, perm) T = PermutedDimsArray{elt,3,perm,invperm(perm),eltype(a)} - @test typeof(p) === RecursivePermutedDimsArray{T,3,perm,invperm(perm),typeof(a)} + @test typeof(p) === NestedPermutedDimsArray{T,3,perm,invperm(perm),typeof(a)} @test size(p) == (4, 3, 2) @test eltype(p) === T for I in eachindex(p) diff --git a/NDTensors/test/lib/runtests.jl b/NDTensors/test/lib/runtests.jl index f89a282dc8..71bf0d46b7 100644 --- a/NDTensors/test/lib/runtests.jl +++ b/NDTensors/test/lib/runtests.jl @@ -15,7 +15,7 @@ using Test: @testset "LabelledNumbers", "MetalExtensions", "NamedDimsArrays", - "RecursivePermutedDimsArrays", + "NestedPermutedDimsArrays", "SmallVectors", "SortedSets", "SparseArrayDOKs", From acd274669b32c434f6a1923fd850f94fddc11ee7 Mon Sep 17 00:00:00 2001 From: mtfishman Date: Fri, 15 Nov 2024 11:35:27 -0500 Subject: [PATCH 3/5] Expect array of arrays for now --- .../src/NestedPermutedDimsArrays.jl | 23 ++++++------------- 1 file changed, 7 insertions(+), 16 deletions(-) diff --git a/NDTensors/src/lib/NestedPermutedDimsArrays/src/NestedPermutedDimsArrays.jl b/NDTensors/src/lib/NestedPermutedDimsArrays/src/NestedPermutedDimsArrays.jl index 1f5304ecb1..d9749215ff 100644 --- a/NDTensors/src/lib/NestedPermutedDimsArrays/src/NestedPermutedDimsArrays.jl +++ b/NDTensors/src/lib/NestedPermutedDimsArrays/src/NestedPermutedDimsArrays.jl @@ -52,25 +52,16 @@ Base.@constprop :aggressive function NestedPermutedDimsArray( throw(ArgumentError(string(perm, " is not a valid permutation of dimensions 1:", N))) iperm = invperm(perm) return NestedPermutedDimsArray{ - maybe_permuteddimsarraytype(T, perm),N,(perm...,),(iperm...,),typeof(data) + PermutedDimsArray{eltype(T),N,(perm...,),(iperm...,),T}, + N, + (perm...,), + (iperm...,), + typeof(data), }( data ) end -# Ideally would use `Base.promote_op(maybe_permuteddimsarraytype, type, perm)` -# but it doesn't handle `perm` properly. -function maybe_permuteddimsarraytype(type::Type{<:AbstractArray}, perm) - return PermutedDimsArray{eltype(type),ndims(type),perm,invperm(perm),type} -end -maybe_permuteddimsarraytype(type::Type, perm) = type - -function maybe_permuteddimsarray(A::AbstractArray, perm) - return PermutedDimsArray(A, perm) -end -# By default, assume scalar and don't permute. -maybe_permuteddimsarray(x, perm) = x - Base.parent(A::NestedPermutedDimsArray) = A.parent function Base.size(A::NestedPermutedDimsArray{T,N,perm}) where {T,N,perm} return genperm(size(parent(A)), perm) @@ -109,14 +100,14 @@ end A::NestedPermutedDimsArray{T,N,perm,iperm}, I::Vararg{Int,N} ) where {T,N,perm,iperm} @boundscheck checkbounds(A, I...) - @inbounds val = maybe_permuteddimsarray(getindex(A.parent, genperm(I, iperm)...), perm) + @inbounds val = PermutedDimsArray(getindex(A.parent, genperm(I, iperm)...), perm) return val end @inline function Base.setindex!( A::NestedPermutedDimsArray{T,N,perm,iperm}, val, I::Vararg{Int,N} ) where {T,N,perm,iperm} @boundscheck checkbounds(A, I...) - @inbounds setindex!(A.parent, maybe_permuteddimsarray(val, perm), genperm(I, iperm)...) + @inbounds setindex!(A.parent, PermutedDimsArray(val, perm), genperm(I, iperm)...) return val end From ac5c500c9b30c749d4b54b5acac75c684dca1c4d Mon Sep 17 00:00:00 2001 From: mtfishman Date: Fri, 15 Nov 2024 12:10:42 -0500 Subject: [PATCH 4/5] [NDTensors] Bump to v0.3.69 --- NDTensors/Project.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/NDTensors/Project.toml b/NDTensors/Project.toml index 80f68c17f0..c1568e650f 100644 --- a/NDTensors/Project.toml +++ b/NDTensors/Project.toml @@ -1,7 +1,7 @@ name = "NDTensors" uuid = "23ae76d9-e61a-49c4-8f12-3f1a16adf9cf" authors = ["Matthew Fishman "] -version = "0.3.68" +version = "0.3.69" [deps] Accessors = "7d9f7c33-5ae7-4f3b-8dc6-eff91059b697" From 6e6f3b14f0036468629a28879f37195c895cd079 Mon Sep 17 00:00:00 2001 From: Matt Fishman Date: Fri, 15 Nov 2024 12:27:24 -0500 Subject: [PATCH 5/5] Reorder loading libraries --- NDTensors/src/imports.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/NDTensors/src/imports.jl b/NDTensors/src/imports.jl index 9869df6448..77d0d6c1d6 100644 --- a/NDTensors/src/imports.jl +++ b/NDTensors/src/imports.jl @@ -37,10 +37,10 @@ for lib in [ :GradedAxes, :SymmetrySectors, :TensorAlgebra, + :NestedPermutedDimsArrays, :SparseArrayInterface, :SparseArrayDOKs, :DiagonalArrays, - :NestedPermutedDimsArrays, :BlockSparseArrays, :NamedDimsArrays, :SmallVectors,