From 7cf624ee6cacd8627e8bad6fa636be77c9fdd9f4 Mon Sep 17 00:00:00 2001 From: Rafael Schouten Date: Thu, 28 Nov 2024 00:27:43 +0100 Subject: [PATCH 1/2] unrol map for large stacks (#873) --- src/stack/indexing.jl | 12 +++++++++--- src/stack/stack.jl | 21 +++++++++++++++------ src/utils.jl | 21 +++++++++++++++++++++ 3 files changed, 45 insertions(+), 9 deletions(-) diff --git a/src/stack/indexing.jl b/src/stack/indexing.jl index e3d41eb0a..f5fc0fea1 100644 --- a/src/stack/indexing.jl +++ b/src/stack/indexing.jl @@ -128,11 +128,12 @@ for f in (:getindex, :view, :dotview) I = length(layerdims) > 0 ? layerdims : map(_ -> :, size(A)) Base.$f(A, I...) end - newlayers = map(f, values(s)) + newlayers = unrolled_map(f, values(s)) # Decide to rewrap as an AbstractDimStack, or return a scalar - if any(map(v -> v isa AbstractDimArray, newlayers)) + if _any_dimarray(newlayers) + # TODO rethink this for many-layered stacks # Some scalars, re-wrap them as zero dimensional arrays - non_scalar_layers = map(values(s), newlayers) do l, nl + non_scalar_layers = unrolled_map(values(s), newlayers) do l, nl nl isa AbstractDimArray ? nl : rebuild(l, fill(nl), ()) end rebuildsliced(Base.$f, s, NamedTuple{K}(non_scalar_layers), (dims2indices(dims(s), D))) @@ -144,6 +145,11 @@ for f in (:getindex, :view, :dotview) end end +@generated function _any_dimarray(v::Union{NamedTuple,Tuple}) + any(T -> T <: AbstractDimArray, v.types) +end + + #### setindex #### @propagate_inbounds Base.setindex!(s::AbstractDimStack, xs, I...; kw...) = diff --git a/src/stack/stack.jl b/src/stack/stack.jl index 68a645faf..c1b4ea730 100644 --- a/src/stack/stack.jl +++ b/src/stack/stack.jl @@ -53,7 +53,7 @@ end @assume_effects :foldable DD.layers(s::AbstractDimStack, k::Symbol) = s[k] @assume_effects :foldable data_eltype(nt::NamedTuple{K}) where K = - NamedTuple{K,Tuple{map(eltype, Tuple(nt))...}} + NamedTuple{K,Tuple{unrolled_map(eltype, Tuple(nt))...}} stacktype(s, data, dims, layerdims::NamedTuple{K}) where K = basetypeof(s){K,data_eltype(data),length(dims)} @@ -82,10 +82,15 @@ function rebuild(s::AbstractDimStack; return T(data, dims, refdims, layerdims, metadata, layermetadata) end -function rebuildsliced(f::Function, s::AbstractDimStack, layers, I) - layerdims = map(basedims, layers) +function rebuildsliced(f::Function, s::AbstractDimStack, layers::NamedTuple, I) + layerdims = unrolled_map(basedims, layers) dims, refdims = slicedims(f, s, I) - return rebuild(s; data=map(parent, layers), dims, refdims, layerdims) + return rebuild(s; data=unrolled_map(parent, layers), dims, refdims, layerdims) +end +function rebuildsliced(f::Function, s::AbstractDimStack{K}, layers::Tuple, I) where K + layerdims = NamedTuple{K}(unrolled_map(basedims, layers)) + dims, refdims = slicedims(f, s, I) + return rebuild(s; data=unrolled_map(parent, layers), dims, refdims, layerdims) end """ @@ -169,7 +174,10 @@ Base.copy(s::AbstractDimStack) = modify(copy, s) # NamedTuple-like @assume_effects :foldable Base.getproperty(s::AbstractDimStack, x::Symbol) = s[x] Base.haskey(s::AbstractDimStack{K}, k) where K = k in K -Base.values(s::AbstractDimStack) = values(layers(s)) +Base.values(s::AbstractDimStack) = _values_gen(s) +@generated function _values_gen(s::AbstractDimStack{K}) where K + Expr(:tuple, map(k -> :(s[$(QuoteNode(k))]), K)...) +end Base.checkbounds(s::AbstractDimStack, I...) = checkbounds(CartesianIndices(s), I...) Base.checkbounds(T::Type, s::AbstractDimStack, I...) = checkbounds(T, CartesianIndices(s), I...) @@ -216,7 +224,8 @@ Base.map(f, s::AbstractDimStack) = error("Use maplayers(f, stack)) instad of map Base.map(f, ::Union{AbstractDimStack,NamedTuple}, xs::Union{AbstractDimStack,NamedTuple}...) = error("Use maplayers(f, stack, args...)) instad of map(f, stack, args...)") -maplayers(f, s::AbstractDimStack) = _maybestack(s, map(f, values(s))) +maplayers(f, s::AbstractDimStack) = + _maybestack(s, unrolled_map(f, values(s))) function maplayers( f, x1::Union{AbstractDimStack,NamedTuple}, xs::Union{AbstractDimStack,NamedTuple}... ) diff --git a/src/utils.jl b/src/utils.jl index fe3eb4172..31001f7c5 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -215,3 +215,24 @@ function _as_extended_nts(nt::NamedTuple{K1}, st::AbstractDimStack{K2}, As...) w return (extended_layers, _as_extended_nts(nt, As...)...) end _as_extended_nts(::NamedTuple) = () + + +# Tuple map that is always unrolled +# mostly for stack indexing performance +_unrolled_map_inner(f, v::Type{T}) where T = + Expr(:tuple, (:(f(v[$i])) for i in eachindex(T.types))...) +_unrolled_map_inner(f, v1::Type{T}, v2::Type) where T = + Expr(:tuple, (:(f(v1[$i], v2[$i])) for i in eachindex(T.types))...) + +@generated function unrolled_map(f, v::NamedTuple{K}) where K + exp = _unrolled_map_inner(f, v) + :(NamedTuple{K}($exp)) +end +@generated function unrolled_map(f, v1::NamedTuple{K}, v2::NamedTuple{K}) where K + exp = _unrolled_map_inner(f, v1, v2) + :(NamedTuple{K}($exp)) +end +@generated unrolled_map(f, v::Tuple) = + _unrolled_map_inner(f, v) +@generated unrolled_map(f, v1::Tuple, v2::Tuple) = + _unrolled_map_inner(f, v1, v2) From 51e94b15a97162feb3ee15cdf8788893faefff24 Mon Sep 17 00:00:00 2001 From: Rafael Schouten Date: Thu, 28 Nov 2024 00:30:05 +0100 Subject: [PATCH 2/2] using isdefined (#872) --- ext/DimensionalDataDiskArraysExt.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ext/DimensionalDataDiskArraysExt.jl b/ext/DimensionalDataDiskArraysExt.jl index 0ac69651f..5d9c5b860 100644 --- a/ext/DimensionalDataDiskArraysExt.jl +++ b/ext/DimensionalDataDiskArraysExt.jl @@ -12,7 +12,7 @@ import DiskArrays # cache was only introduced in DiskArrays v0.4, so # we lock out the method definition if the method does # not exist. -@static if :cache in names(DiskArrays) +@static if isdefined(DiskArrays, :cache) DiskArrays.cache(x::Union{AbstractDimStack,AbstractDimArray}; kw...) = modify(A -> DiskArrays.cache(A; kw...), x) end