diff --git a/README.md b/README.md index 1f2d721..69756fb 100644 --- a/README.md +++ b/README.md @@ -37,6 +37,105 @@ probably slower than an ordinary array, since a new object must be heap allocated every time the StructOfArrays is indexed. In practice, StructsOfArrays works best with `isbits` immutables such as `Complex{T}`. +## Advanced Usage + +When embedding a StructOfArrays in a larger, data structure, it can be useful +to automatically compute the SoA type corresponding to a regular arrays type. +This can be accomplished using the `similar` function: + +``` +struct Vectors + x::similar(StructOfArrays, Vector{Complex{Float64}}) + y::similar(StructOfArrays, Vector{Complex{Float64}}) +end +# equivalent in layout to +struct Vectors + x_real::Vector{Float64} + x_imag::Vector{Float64} + y_real::Vector{Float64} + y_imag::Vector{Float64} +end +``` + +Note that this feature also works with parameterized types. However, if +later a composite type is substituted for type type parameter, it will not +be unpacked into separate arrays: + +``` +struct Vector3D{T} + x::T + y::T + z::T +end +struct Container{T} + soa::similar(StructOfArrays, Vector{Vector3D{T}}) +end +# Container{Float64} is equivalent to +struct ContainerFloat64 + soa_x::Vector{Float64} + soa_y::Vector{Float64} + soa_z::Vector{Float64} +end +# Container{Complex{Float64}} is equivalent to +struct ContainerFloat64 + soa_x::Vector{Complex{Float64}} + soa_y::Vector{Complex{Float64}} + soa_z::Vector{Complex{Float64}} +end +# Note that this is different from similar(StructOfArrays, Vector{Vector3D{Complex{Float64}}}), which would expand to +struct ContainerFloat64 + soa_x_real::Vector{Float64} + soa_x_imag::Vector{Float64} + soa_y_real::Vector{Float64} + soa_y_imag::Vector{Float64} + soa_y_real::Vector{Float64} + soa_y_imag::Vector{Float64} +end +``` + +This behavior was chosen to accomodate julia's handling of parameterize element, +types. If future versions of julia expand these capabilities, the default behavior +may need to be revisited. + +Lastly, note that it is possible to choose control the recursion explicitly, +by providing a type (as the second type parameters) whose subtypes should be +considered leaves for the purpose of recursion. E.g.: +``` +struct Bundle + x::Complex{Float64} + y::Complex{Int64} + z::Rational{Int64} +end +# Consider +A = StructOfArrays{Bundle, Any}(2,2) +# Since all types are <: Any, no recusion will occur, and the SoaA is equivalent to +struct SoAAny + x::Matrix{Complex{Float64}} + y::Matrix{Complex{Int64}} + z::Matrix{Rational{Int64}} +end +# Next, let's say we want to have separate SoA for the complex values. Consider +B = StructOfArrays{Bundle, Complex}(2,2) +# which will be equivalent to +struct SoAComplex + x_real::Matrix{Float64} + x_imag::Matrix{Float64} + y_real::Matrix{Int64} + y_imag::Matrix{Int64} + z::Matrix{Rational{Int64}} +end +# Lastly it is of course possible to specify a union: +C = StructOfArrays{Bundle, Union{Complex{Float64},Rational}}(2,2) +struct SoAUnion + x_real::Matrix{Float64} + x_imag::Matrix{Float64} + y::Matrox{Complex{Int64}} + z_num::Matrix{Int64} + z_den::Matrix{Int64} +end +``` + + ## Benchmark ```julia diff --git a/src/StructsOfArrays.jl b/src/StructsOfArrays.jl index eb4226a..e7457b3 100644 --- a/src/StructsOfArrays.jl +++ b/src/StructsOfArrays.jl @@ -1,40 +1,97 @@ +__precompile__() module StructsOfArrays using Compat export StructOfArrays -immutable StructOfArrays{T,N,U<:Tuple} <: AbstractArray{T,N} +immutable StructOfArrays{T,Leaves,N,U<:Tuple} <: AbstractArray{T,N} arrays::U + StructOfArrays{T,Leaves,N,U}(arrays::U) where {T,Leaves,N,U} = new{T,Leaves,N,U}(arrays) end -function gather_eltypes(T, visited = Set{Type}()) - (!isleaftype(T) || T.mutable) && throw(ArgumentError("can only create an StructOfArrays of leaf type immutables")) +function gather_eltypes(T, Leaves=Union{}, visited = Set{Type}(); allow_typevar=false) + @assert !isa(T, UnionAll) + T.mutable && throw(ArgumentError("can not create StructOfArray for a mutable type")) isempty(T.types) && throw(ArgumentError("cannot create an StructOfArrays of an empty or bitstype")) - types = Type[] + types = Any[] + typevars = TypeVar[] push!(visited, T) for S in T.types - sizeof(S) == 0 && continue (S in visited) && throw(ArgumentError("Recursive types are not allowed for SoA conversion")) - if isempty(S.types) + if isa(S, TypeVar) + !allow_typevar && throw(ArgumentError("TypeVars are not allowed when constructing an SoA array")) + push!(typevars, S) + elseif isa(S, UnionAll) + throw(ArgumentError("Interior UnionAlls are not allowed")) + elseif isleaftype(S) && S.size == 0 + continue + end + if isa(S, DataType) && S.name.wrapper == Vararg + n = S.parameters[2] + elT = S.parameters[1] + isa(n, Integer) || throw(ArgumentError("Tuple are only supported with definite lengths")) + isa(elT, TypeVar) && push!(typevars, elT) + if isa(elT, TypeVar) || isempty(eT.types) || elT <: Leaves + append!(types, [elT for i = 1:n]) + else + ntypes, ntypevars = gather_eltypes(S, Leaves, copy(visited); allow_typevar=allow_typevar) + for i = 1:n + append!(types, ntypes) + end + append!(typevars, ntypevars) + end + elseif isa(S, TypeVar) || isempty(S.types) || S <: Leaves push!(types, S) else - append!(types, gather_eltypes(S, copy(visited))) + ntypes, ntypevars = gather_eltypes(S, Leaves, copy(visited); allow_typevar=allow_typevar) + append!(types, ntypes) + append!(typevars, ntypevars) end end - types + types, typevars end -@generated function StructOfArrays{T}(::Type{T}, dims::Integer...) - N = length(dims) - types = gather_eltypes(T) - arrtuple = Tuple{[Array{S,N} for S in types]...} - :(StructOfArrays{T,$N,$arrtuple}(($([:(Array{$(S)}(dims)) for S in types]...),))) +@generated function StructOfArrays{T,Leaves,N,arrtuple}(dims::Tuple{Vararg{Integer}}) where {T,Leaves,N,arrtuple} + Expr(:new, StructOfArrays{T,Leaves,N,arrtuple}, + Expr(:tuple, + (:($arr(dims)) for arr in arrtuple.parameters)...)) +end + +@generated function StructOfArrays{T,Leaves}(dims::Tuple{Vararg{Integer}}) where {T,Leaves} + types = gather_eltypes(T, Leaves)[1] + arrtuple = Expr(:curly, Tuple, [:(Array{$S,length(dims)}) for S in types]...) + :(StructOfArrays{T,Leaves,length(dims),$arrtuple}(dims...)) end -StructOfArrays(T::Type, dims::Tuple{Vararg{Integer}}) = StructOfArrays(T, dims...) + +@generated function StructOfArrays{T}(dims::Tuple{Vararg{Integer}}) where {T} + types, typevars = gather_eltypes(T) + @assert isempty(typevars) + arrtuple = Expr(:curly, Tuple, [:(Array{$S,length(dims)}) for S in types]...) + :(StructOfArrays{T,Union{},length(dims),$arrtuple}(($([:(Array{$(S)}(dims)) for S in types]...),))) +end +StructOfArrays(T::Type, dims::Integer...) = StructOfArrays{T}(dims) +StructOfArrays(T::Type, dims::Tuple{Vararg{Integer}}) = StructOfArrays{T}(dims) +StructOfArrays{T}(dims::Integer...) where {T} = StructOfArrays{T}(dims) +StructOfArrays{T,Leaves}(dims::Integer...) where {T,Leaves} = StructOfArrays{T,Leaves}(dims) +StructOfArrays{T,Leaves,N,U}(dims::Integer...) where {T,Leaves,N,U} = StructOfArrays{T,Leaves,N,U}(dims) +StructOfArrays{T,<:Any,N}() where {T,N} = StructOfArrays{T}((0 for i=1:N)...) +StructOfArrays{T,Leaves,N,U}() where {T,Leaves,N,U} = StructOfArrays{T,Leaves,N,U}((0 for i=1:N)...) @compat Base.IndexStyle{T<:StructOfArrays}(::Type{T}) = IndexLinear() +function Base.similar{T,N}(::Type{<:StructOfArrays}, ::Type{<:AbstractArray{T,N}}) + uT = Base.unwrap_unionall(T) + types, typevars = gather_eltypes(uT, allow_typevar = true) + Leaves = Union{typevars...} + arrtuple = Tuple{[Array{S,N} for S in types]...} + Base.rewrap_unionall(StructOfArrays{uT, Leaves, N, arrtuple}, T) +end + +function Base.similar{T,N}(::Type{StructOfArrays}, A::AbstractArray{T,N}) + StructOfArrays{T}(size(A)) +end + @generated function Base.similar{T}(A::StructOfArrays, ::Type{T}, dims::Dims) if isbits(T) && length(T.types) > 1 :(StructOfArrays(T, dims)) @@ -53,44 +110,58 @@ Base.convert{T,N}(::Type{StructOfArrays}, A::AbstractArray{T,N}) = Base.size(A::StructOfArrays) = size(A.arrays[1]) Base.size(A::StructOfArrays, d) = size(A.arrays[1], d) -function generate_getindex(T, arraynum) - members = Expr[] - for S in T.types +function generate_elementwise(leaff, recursef, combinef, state, T, Leaves, arraynum=1) + for (el,S) in enumerate(T.types) sizeof(S) == 0 && push!(members, :($(S()))) - if isempty(S.types) - push!(members, :(A.arrays[$arraynum][i...])) + if isempty(S.types) || S <: Leaves + leaff(arraynum, el, state) arraynum += 1 else - member, arraynum = generate_getindex(S, arraynum) - push!(members, member) + nstate, arraynum = generate_elementwise(leaff, recursef, combinef, recursef(S, arraynum, el, state), S, Leaves, arraynum) + combinef(S, arraynum, el, state, nstate) end end - Expr(:new, T, members...), arraynum + state, arraynum end -@generated function Base.getindex{T}(A::StructOfArrays{T}, i::Integer...) - strct, _ = generate_getindex(T, 1) - Expr(:block, Expr(:meta, :inline), strct) +@generated function Base.getindex{T,Leaves}(A::StructOfArrays{T,Leaves}, i::Integer...) + leaf(arraynum, el, members) = push!(members, :(A.arrays[$arraynum][i...])) + recursef(S, arraynum, el, state) = Expr[] + combinef(S, arraynum, el, members, eltmems) = push!(members, Expr(:new, S, eltmems...)) + mems, _ = generate_elementwise(leaf, recursef, combinef, Expr[], T, Leaves) + Expr(:block, Expr(:meta, :inline, :propagate_inbounds), Expr(:new, T, mems...)) end -function generate_setindex(T, x, arraynum) - s = gensym() - exprs = Expr[:($s = $x)] - for (el,S) in enumerate(T.types) - sizeof(S) == 0 && push!(members, :($(S()))) - if isempty(S.types) - push!(exprs, :(A.arrays[$arraynum][i...] = getfield($s, $el))) - arraynum += 1 - else - nexprs, arraynum = generate_setindex(S, :(getfield($s, $el)), arraynum) - append!(exprs, nexprs) - end +function setindex_recusion(exprs) + function recursef(S, arraynum, el, state) + s = gensym() + push!(exprs, :($s = getfield($state, $el))) + s + end +end + +no_combinef(args...) = nothing + +@generated function Base.setindex!{T, Leaves}(A::StructOfArrays{T, Leaves}, x, i::Integer...) + exprs = Expr[] + function leaf(arraynum, el, state) + push!(exprs, :(A.arrays[$arraynum][i...] = getfield($state, $el))) + end + generate_elementwise(leaf, setindex_recusion(exprs), no_combinef, :v, T, Leaves) + exprs = Expr(:block, exprs...) + quote + $(Expr(:meta, :inline, :propagate_inbounds)) + v = convert(T, x) + $exprs + x end - exprs, arraynum end -@generated function Base.setindex!{T}(A::StructOfArrays{T}, x, i::Integer...) - exprs = Expr(:block, generate_setindex(T, :x, 1)[1]...) +@generated function Base.push!{T, Leaves}(A::StructOfArrays{T, Leaves}, x) + exprs = Expr[] + leaf(arraynum, el, state) = push!(exprs, :(push!(A.arrays[$arraynum],getfield($state,$el)))) + generate_elementwise(leaf, setindex_recusion(exprs), no_combinef, :v, T, Leaves) + exprs = Expr(:block, exprs...) quote $(Expr(:meta, :inline)) v = convert(T, x) @@ -98,4 +169,5 @@ end x end end + end # module diff --git a/test/runtests.jl b/test/runtests.jl index 98b30bc..9591081 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -42,3 +42,108 @@ end small = StructOfArrays(TwoField, 2, 2) small[1,1] = TwoField(OneField(1), OneField(2)) @test small[1,1] == TwoField(OneField(1), OneField(2)) + +# SoA with explicit leaves +immutable Complex3D + x::Complex{Float64} + y::Complex{Float64} + z::Complex{Float64} +end + +ComplexA = Complex3D(Complex{Float64}(1.0,2.0), + Complex{Float64}(3.0,4.0), + Complex{Float64}(5.0,6.0)) +full_rec = StructOfArrays{Complex3D}(2, 2) +@test length(full_rec.arrays) == 6 +full_rec[1,1] = ComplexA +@test full_rec[1,1] == ComplexA + +partial_rec = StructOfArrays{Complex3D,Complex}(2, 2) +@test length(partial_rec.arrays) == 3 +partial_rec[2,2] = ComplexA +@test partial_rec[2,2] == ComplexA + +struct Bundle + x::Complex{Float64} + y::Complex{Int64} + z::Rational{Int64} +end + +BundleA = Bundle(Complex{Float64}(1.0,2.0), + Complex{Int64}(1.0,2.0), + 1//2) +full_rec = StructOfArrays{Bundle}(2, 2) +@test length(full_rec.arrays) == 6 +full_rec[1,1] = BundleA +@test full_rec[1,1] == BundleA + +partial_rec = StructOfArrays{Bundle,Any}(2, 2) +@test length(partial_rec.arrays) == 3 +partial_rec[2,2] = BundleA +@test partial_rec[2,2] == BundleA + +# SoA with parameterized leaves +immutable ParamComplex3D{T} + x::Complex{T} + y::Complex{T} + z::Complex{T} +end + +ParamComplexA = ParamComplex3D{Float64}(ComplexA.x, ComplexA.y, ComplexA.z) +full_rec = StructOfArrays{ParamComplex3D{Float64}}(2, 2) +@test length(full_rec.arrays) == 6 +full_rec[1,2] = ParamComplexA +@test full_rec[1,2] == ParamComplexA + +partial_rec = StructOfArrays{ParamComplex3D{Float64},Complex}(2, 2) +@test length(partial_rec.arrays) == 3 +full_rec[2,1] = ParamComplexA +@test full_rec[2,1] == ParamComplexA + +const SoAParamComplex3D{T} = similar(StructOfArrays, Matrix{ParamComplex3D{T}}) +soa = SoAParamComplex3D{Float64}(2, 2) +@test length(soa.arrays) == 6 +@test eltype(soa.arrays[1]) == Float64 +full_rec[2,2] = ParamComplexA +@test full_rec[2,2] == ParamComplexA + +immutable Param3D{T} + x::T + y::T + z::T +end +OtherParamComplexA = Param3D{Complex{Float64}}(ComplexA.x, ComplexA.y, ComplexA.z) +const OtherSoAParamComplex3D{T} = similar(StructOfArrays, Matrix{Param3D{T}}) +soa = OtherSoAParamComplex3D{Complex{Float64}}(2, 2) +@test length(soa.arrays) == 3 +@test eltype(soa.arrays[1]) == Complex{Float64} +soa[1,1] = OtherParamComplexA +@test soa[1,1] == OtherParamComplexA + +NesetedA = Param3D{Param3D{Complex{Float64}}}(OtherParamComplexA, OtherParamComplexA, OtherParamComplexA) +nestedsoa = OtherSoAParamComplex3D{Param3D{Complex{Float64}}}(2, 2) +@test length(soa.arrays) == 3 +@test eltype(soa.arrays[1]) == Complex{Float64} +nestedsoa[1,1] = NesetedA +@test nestedsoa[1,1] == NesetedA + +# NTuples +V = Vector{NTuple{4, Int64}}() +ntuplesoa = similar(StructOfArrays, V) +@test length(ntuplesoa.arrays) == 4 +push!(ntuplesoa, (1,2,3,4)) +@test ntuplesoa[1] == (1,2,3,4) +ntuplesoa[1] = (5,6,7,8) +@test ntuplesoa[1] == (5,6,7,8) + +# Parameterized NTuples +immutable SVec{T} + data::NTuple{3, T} +end +SVecSoA{T} = similar(StructOfArrays, Vector{SVec{T}}) +svecsoa = SVecSoA{Float64}() +@test length(svecsoa.arrays) == 3 +push!(svecsoa, SVec((1.,2.,3.))) +@test svecsoa[1] == SVec((1.,2.,3.)) +svecsoa[1] = SVec((4.,5.,6.)) +@test svecsoa[1] == SVec((4.,5.,6.))