From c81c30892d64d37160e7b3ec529e7c93f8265e17 Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Tue, 4 Apr 2023 09:07:49 +0200 Subject: [PATCH 1/9] Use weakdeps on Julia v1.9 Update Project.toml --- Project.toml | 12 +++++++++-- ext/StructArraysGPUArraysCoreExt.jl | 21 +++++++++++++++++++ .../StructArraysStaticArraysCoreExt.jl | 11 +++++++++- src/StructArrays.jl | 13 +++--------- 4 files changed, 44 insertions(+), 13 deletions(-) create mode 100644 ext/StructArraysGPUArraysCoreExt.jl rename src/staticarrays_support.jl => ext/StructArraysStaticArraysCoreExt.jl (89%) diff --git a/Project.toml b/Project.toml index 29d085ac..8716f10b 100644 --- a/Project.toml +++ b/Project.toml @@ -6,9 +6,15 @@ version = "0.6.16" Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" ConstructionBase = "187b0558-2788-49d3-abe0-74a17ed4e7c9" DataAPI = "9a962f9c-6df0-11e9-0e5d-c546b8b5ee8a" +Tables = "bd369af6-aec1-5ad0-b16a-f7cc5008161c" + +[weakdeps] GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527" StaticArraysCore = "1e83bf80-4336-4d27-bf5d-d5a4f845583c" -Tables = "bd369af6-aec1-5ad0-b16a-f7cc5008161c" + +[extensions] +StructArraysGPUArraysCoreExt = "GPUArraysCore" +StructArraysStaticArraysCoreExt = "StaticArraysCore" [compat] Adapt = "1, 2, 3" @@ -22,14 +28,16 @@ julia = "1.6" [extras] Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4" +GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527" JLArrays = "27aeb0d3-9eb9-45fb-866b-73c2ecf80fcb" OffsetArrays = "6fe1bfb0-de20-5000-8ca7-80f57d26f881" PooledArrays = "2dfb63ee-cc39-5dd5-95bd-886bf059d720" SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" +StaticArraysCore = "1e83bf80-4336-4d27-bf5d-d5a4f845583c" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" TypedTables = "9d95f2ec-7b3d-5a63-8d20-e2491e220bb9" WeakRefStrings = "ea10d353-3f73-51f8-a26c-33c1cb351aa5" [targets] -test = ["Test", "JLArrays", "StaticArrays", "OffsetArrays", "PooledArrays", "TypedTables", "WeakRefStrings", "Documenter", "SparseArrays"] +test = ["Test", "JLArrays", "StaticArrays", "OffsetArrays", "PooledArrays", "TypedTables", "WeakRefStrings", "Documenter", "SparseArrays", "GPUArraysCore", "StaticArraysCore"] diff --git a/ext/StructArraysGPUArraysCoreExt.jl b/ext/StructArraysGPUArraysCoreExt.jl new file mode 100644 index 00000000..b05d3082 --- /dev/null +++ b/ext/StructArraysGPUArraysCoreExt.jl @@ -0,0 +1,21 @@ +module StructArraysGPUArraysCoreExt + +using StructArrays +using StructArrays: map_params, array_types + +using Base: tail + +import GPUArraysCore + +# for GPU broadcast +import GPUArraysCore +function GPUArraysCore.backend(::Type{T}) where {T<:StructArray} + backends = map_params(GPUArraysCore.backend, array_types(T)) + backend, others = backends[1], tail(backends) + isconsistent = mapfoldl(isequal(backend), &, others; init=true) + isconsistent || throw(ArgumentError("all component arrays must have the same GPU backend")) + return backend +end +StructArrays.always_struct_broadcast(::GPUArraysCore.AbstractGPUArrayStyle) = true + +end # module diff --git a/src/staticarrays_support.jl b/ext/StructArraysStaticArraysCoreExt.jl similarity index 89% rename from src/staticarrays_support.jl rename to ext/StructArraysStaticArraysCoreExt.jl index 1af186e8..28e73762 100644 --- a/src/staticarrays_support.jl +++ b/ext/StructArraysStaticArraysCoreExt.jl @@ -1,3 +1,10 @@ +module StructArraysStaticArraysCoreExt + +using StructArrays +using StructArrays: StructArrayStyle, createinstance, replace_structarray, isnonemptystructtype + +using Base.Broadcast: Broadcasted + using StaticArraysCore: StaticArray, FieldArray, tuple_prod """ @@ -40,7 +47,7 @@ Broadcast._axes(bc::Broadcasted{<:StructStaticArrayStyle}, ::Nothing) = axes(rep # StaticArrayStyle has no similar defined. # Overload `Base.copy` instead. -@inline function try_struct_copy(bc::Broadcasted{StaticArrayStyle{M}}) where {M} +@inline function StructArrays.try_struct_copy(bc::Broadcasted{StaticArrayStyle{M}}) where {M} sa = copy(bc) ET = eltype(sa) isnonemptystructtype(ET) || return sa @@ -66,3 +73,5 @@ end return map(Base.Fix2(getfield, i), x) end end + +end # module diff --git a/src/StructArrays.jl b/src/StructArrays.jl index 55d77542..65c6b736 100644 --- a/src/StructArrays.jl +++ b/src/StructArrays.jl @@ -14,7 +14,6 @@ include("collect.jl") include("sort.jl") include("lazy.jl") include("tables.jl") -include("staticarrays_support.jl") # Implement refarray and refvalue to deal with pooled arrays and weakrefstrings effectively import DataAPI: refarray, refvalue @@ -30,15 +29,9 @@ end import Adapt Adapt.adapt_structure(to, s::StructArray) = replace_storage(x->Adapt.adapt(to, x), s) -# for GPU broadcast -import GPUArraysCore -function GPUArraysCore.backend(::Type{T}) where {T<:StructArray} - backends = map_params(GPUArraysCore.backend, array_types(T)) - backend, others = backends[1], tail(backends) - isconsistent = mapfoldl(isequal(backend), &, others; init=true) - isconsistent || throw(ArgumentError("all component arrays must have the same GPU backend")) - return backend +@static if !isdefined(Base, :get_extension) + include("../ext/StructArraysGPUArraysCoreExt.jl") + include("../ext/StructArraysStaticArraysCoreExt.jl") end -always_struct_broadcast(::GPUArraysCore.AbstractGPUArrayStyle) = true end # module From 7171baba33c923d5a62dfd58bfcf5fad38fe2295 Mon Sep 17 00:00:00 2001 From: N5N3 <2642243996@qq.com> Date: Sat, 14 Oct 2023 23:28:05 +0800 Subject: [PATCH 2/9] move StructStaticArray broadcast to ext --- Project.toml | 8 +-- ext/StructArraysStaticArraysCoreExt.jl | 77 ------------------------ ext/StructArraysStaticArraysExt.jl | 82 ++++++++++++++++++++++++++ src/StructArrays.jl | 2 +- src/structarray.jl | 2 +- test/runtests.jl | 2 + 6 files changed, 89 insertions(+), 84 deletions(-) delete mode 100644 ext/StructArraysStaticArraysCoreExt.jl create mode 100644 ext/StructArraysStaticArraysExt.jl diff --git a/Project.toml b/Project.toml index 8716f10b..2310f14a 100644 --- a/Project.toml +++ b/Project.toml @@ -10,11 +10,11 @@ Tables = "bd369af6-aec1-5ad0-b16a-f7cc5008161c" [weakdeps] GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527" -StaticArraysCore = "1e83bf80-4336-4d27-bf5d-d5a4f845583c" +StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" [extensions] StructArraysGPUArraysCoreExt = "GPUArraysCore" -StructArraysStaticArraysCoreExt = "StaticArraysCore" +StructArraysStaticArraysExt = "StaticArrays" [compat] Adapt = "1, 2, 3" @@ -22,7 +22,6 @@ ConstructionBase = "1" DataAPI = "1" GPUArraysCore = "0.1.2" StaticArrays = "1.5.6" -StaticArraysCore = "1.3" Tables = "1" julia = "1.6" @@ -34,10 +33,9 @@ OffsetArrays = "6fe1bfb0-de20-5000-8ca7-80f57d26f881" PooledArrays = "2dfb63ee-cc39-5dd5-95bd-886bf059d720" SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" -StaticArraysCore = "1e83bf80-4336-4d27-bf5d-d5a4f845583c" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" TypedTables = "9d95f2ec-7b3d-5a63-8d20-e2491e220bb9" WeakRefStrings = "ea10d353-3f73-51f8-a26c-33c1cb351aa5" [targets] -test = ["Test", "JLArrays", "StaticArrays", "OffsetArrays", "PooledArrays", "TypedTables", "WeakRefStrings", "Documenter", "SparseArrays", "GPUArraysCore", "StaticArraysCore"] +test = ["Test", "JLArrays", "StaticArrays", "OffsetArrays", "PooledArrays", "TypedTables", "WeakRefStrings", "Documenter", "SparseArrays", "GPUArraysCore"] diff --git a/ext/StructArraysStaticArraysCoreExt.jl b/ext/StructArraysStaticArraysCoreExt.jl deleted file mode 100644 index 28e73762..00000000 --- a/ext/StructArraysStaticArraysCoreExt.jl +++ /dev/null @@ -1,77 +0,0 @@ -module StructArraysStaticArraysCoreExt - -using StructArrays -using StructArrays: StructArrayStyle, createinstance, replace_structarray, isnonemptystructtype - -using Base.Broadcast: Broadcasted - -using StaticArraysCore: StaticArray, FieldArray, tuple_prod - -""" - StructArrays.staticschema(::Type{<:StaticArray{S, T}}) where {S, T} - -The `staticschema` of a `StaticArray` element type is the `staticschema` of the underlying `Tuple`. -```julia -julia> StructArrays.staticschema(SVector{2, Float64}) -Tuple{Float64, Float64} -``` -The one exception to this rule is `<:StaticArrays.FieldArray`, since `FieldArray` is based on a -struct. In this case, `staticschema(<:FieldArray)` returns the `staticschema` for the struct -which subtypes `FieldArray`. -""" -@generated function StructArrays.staticschema(::Type{<:StaticArray{S, T}}) where {S, T} - return quote - Base.@_inline_meta - return NTuple{$(tuple_prod(S)), T} - end -end -StructArrays.createinstance(::Type{T}, args...) where {T<:StaticArray} = T(args) -StructArrays.component(s::StaticArray, i) = getindex(s, i) - -# invoke general fallbacks for a `FieldArray` type. -@inline function StructArrays.staticschema(T::Type{<:FieldArray}) - invoke(StructArrays.staticschema, Tuple{Type{<:Any}}, T) -end -StructArrays.component(s::FieldArray, i) = invoke(StructArrays.component, Tuple{Any, Any}, s, i) -StructArrays.createinstance(T::Type{<:FieldArray}, args...) = invoke(createinstance, Tuple{Type{<:Any}, Vararg}, T, args...) - -# Broadcast overload -using StaticArraysCore: StaticArrayStyle, similar_type -StructStaticArrayStyle{N} = StructArrayStyle{StaticArrayStyle{N}, N} -function Broadcast.instantiate(bc::Broadcasted{StructStaticArrayStyle{M}}) where {M} - bc′ = Broadcast.instantiate(replace_structarray(bc)) - return convert(Broadcasted{StructStaticArrayStyle{M}}, bc′) -end -# This looks costly, but the compiler should be able to optimize them away -Broadcast._axes(bc::Broadcasted{<:StructStaticArrayStyle}, ::Nothing) = axes(replace_structarray(bc)) - -# StaticArrayStyle has no similar defined. -# Overload `Base.copy` instead. -@inline function StructArrays.try_struct_copy(bc::Broadcasted{StaticArrayStyle{M}}) where {M} - sa = copy(bc) - ET = eltype(sa) - isnonemptystructtype(ET) || return sa - elements = Tuple(sa) - @static if VERSION >= v"1.7" - arrs = ntuple(Val(fieldcount(ET))) do i - similar_type(sa, fieldtype(ET, i))(_getfields(elements, i)) - end - else - _fieldtype(::Type{T}) where {T} = i -> fieldtype(T, i) - __fieldtype = _fieldtype(ET) - arrs = ntuple(Val(fieldcount(ET))) do i - similar_type(sa, __fieldtype(i))(_getfields(elements, i)) - end - end - return StructArray{ET}(arrs) -end - -@inline function _getfields(x::Tuple, i::Int) - if @generated - return Expr(:tuple, (:(getfield(x[$j], i)) for j in 1:fieldcount(x))...) - else - return map(Base.Fix2(getfield, i), x) - end -end - -end # module diff --git a/ext/StructArraysStaticArraysExt.jl b/ext/StructArraysStaticArraysExt.jl new file mode 100644 index 00000000..63401539 --- /dev/null +++ b/ext/StructArraysStaticArraysExt.jl @@ -0,0 +1,82 @@ +module StructArraysStaticArraysExt + +using StructArrays +using StaticArrays: StaticArray, FieldArray, tuple_prod + +""" + StructArrays.staticschema(::Type{<:StaticArray{S, T}}) where {S, T} + +The `staticschema` of a `StaticArray` element type is the `staticschema` of the underlying `Tuple`. +```julia +julia> StructArrays.staticschema(SVector{2, Float64}) +Tuple{Float64, Float64} +``` +The one exception to this rule is `<:StaticArrays.FieldArray`, since `FieldArray` is based on a +struct. In this case, `staticschema(<:FieldArray)` returns the `staticschema` for the struct +which subtypes `FieldArray`. +""" +@generated function StructArrays.staticschema(::Type{<:StaticArray{S, T}}) where {S, T} + return quote + Base.@_inline_meta + return NTuple{$(tuple_prod(S)), T} + end +end +StructArrays.createinstance(::Type{T}, args...) where {T<:StaticArray} = T(args) +StructArrays.component(s::StaticArray, i) = getindex(s, i) + +# invoke general fallbacks for a `FieldArray` type. +@inline function StructArrays.staticschema(T::Type{<:FieldArray}) + invoke(StructArrays.staticschema, Tuple{Type{<:Any}}, T) +end +StructArrays.component(s::FieldArray, i) = invoke(StructArrays.component, Tuple{Any, Any}, s, i) +StructArrays.createinstance(T::Type{<:FieldArray}, args...) = invoke(StructArrays.createinstance, Tuple{Type{<:Any}, Vararg}, T, args...) + +# Broadcast overload +using StaticArrays: StaticArrayStyle, similar_type, Size, SOneTo +using StaticArrays: broadcast_flatten, broadcast_sizes, first_statictype, __broadcast +using StructArrays: isnonemptystructtype +using Base.Broadcast: Broadcasted + +# StaticArrayStyle has no similar defined. +# Overload `try_struct_copy` instead. +@inline function StructArrays.try_struct_copy(bc::Broadcasted{StaticArrayStyle{M}}) where {M} + flat = broadcast_flatten(bc); as = flat.args; f = flat.f + argsizes = broadcast_sizes(as...) + ax = axes(bc) + ax isa Tuple{Vararg{SOneTo}} || error("Dimension is not static. Please file a bug at `StaticArrays.jl`.") + return _broadcast(f, Size(map(length, ax)), argsizes, as...) +end + +@inline function _broadcast(f, sz::Size{newsize}, s::Tuple{Vararg{Size}}, a...) where {newsize} + first_staticarray = first_statictype(a...) + elements, ET = if prod(newsize) == 0 + # Use inference to get eltype in empty case (see also comments in _map) + eltys = Tuple{map(eltype, a)...} + (), Core.Compiler.return_type(f, eltys) + else + temp = __broadcast(f, sz, s, a...) + temp, eltype(temp) + end + if isnonemptystructtype(ET) + @static if VERSION >= v"1.7" + arrs = ntuple(Val(fieldcount(ET))) do i + @inbounds similar_type(first_staticarray, fieldtype(ET, i), sz)(_getfields(elements, i)) + end + else + similarET(::Type{SA}, ::Type{T}) where {SA, T} = i -> @inbounds similar_type(SA, fieldtype(T, i), sz)(_getfields(elements, i)) + arrs = ntuple(similarET(first_staticarray, ET), Val(fieldcount(ET))) + end + return StructArray{ET}(arrs) + end + @inbounds return similar_type(first_staticarray, ET, sz)(elements) +end + +@inline function _getfields(x::Tuple, i::Int) + if @generated + return Expr(:tuple, (:(getfield(x[$j], i)) for j in 1:fieldcount(x))...) + else + return map(Base.Fix2(getfield, i), x) + end +end + +end diff --git a/src/StructArrays.jl b/src/StructArrays.jl index 65c6b736..8cfbdc48 100644 --- a/src/StructArrays.jl +++ b/src/StructArrays.jl @@ -31,7 +31,7 @@ Adapt.adapt_structure(to, s::StructArray) = replace_storage(x->Adapt.adapt(to, x @static if !isdefined(Base, :get_extension) include("../ext/StructArraysGPUArraysCoreExt.jl") - include("../ext/StructArraysStaticArraysCoreExt.jl") + include("../ext/StructArraysStaticArraysExt.jl") end end # module diff --git a/src/structarray.jl b/src/structarray.jl index ee361c39..c8231bc3 100644 --- a/src/structarray.jl +++ b/src/structarray.jl @@ -551,7 +551,7 @@ See also [`always_struct_broadcast`](@ref). """ try_struct_copy(bc::Broadcasted) = copy(bc) -function Base.copy(bc::Broadcasted{StructArrayStyle{S, N}}) where {S, N} +@inline function Base.copy(bc::Broadcasted{StructArrayStyle{S, N}}) where {S, N} if always_struct_broadcast(S()) return invoke(copy, Tuple{Broadcasted}, bc) else diff --git a/test/runtests.jl b/test/runtests.jl index 85a3637d..61d9be34 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -1297,8 +1297,10 @@ Base.BroadcastStyle(::Broadcast.ArrayStyle{MyArray2}, S::Broadcast.DefaultArrayS @testset "allocation test" begin a = StructArray{ComplexF64}(undef, 1) + sa = StructArray{ComplexF64}((SizedVector{1}(a.re), SizedVector{1}(a.re))) allocated(a) = @allocated a .+ 1 @test allocated(a) == 2allocated(a.re) + @test allocated(sa) == 2allocated(sa.re) allocated2(a) = @allocated a .= complex.(a.im, a.re) @test allocated2(a) == 0 end From b0d509b52b50d929add471f5d4caae94eb29a567 Mon Sep 17 00:00:00 2001 From: N5N3 <2642243996@qq.com> Date: Sat, 14 Oct 2023 23:28:32 +0800 Subject: [PATCH 3/9] fix doctest --- docs/src/advanced.md | 97 +++++++++++++++----------------------------- docs/src/index.md | 70 ++++++++------------------------ src/utils.jl | 4 +- 3 files changed, 52 insertions(+), 119 deletions(-) diff --git a/docs/src/advanced.md b/docs/src/advanced.md index 8f70caac..b529fb99 100644 --- a/docs/src/advanced.md +++ b/docs/src/advanced.md @@ -6,89 +6,56 @@ StructArrays support structures with custom data layout. The user is required to Here is an example of a type `MyType` that has as custom fields either its field `data` or fields of its field `rest` (which is a named tuple): -```jldoctest advanced1 -julia> using StructArrays +```@repl advanced1 +using StructArrays -julia> struct MyType{T, NT<:NamedTuple} - data::T - rest::NT - end +struct MyType{T, NT<:NamedTuple} + data::T + rest::NT +end -julia> MyType(x; kwargs...) = MyType(x, values(kwargs)) -MyType +MyType(x; kwargs...) = MyType(x, values(kwargs)) ``` Let's create a small array of these objects: -```jldoctest advanced1 -julia> s = [MyType(i/5, a=6-i, b=2) for i in 1:5] -5-element Vector{MyType{Float64, NamedTuple{(:a, :b), Tuple{Int64, Int64}}}}: - MyType{Float64, NamedTuple{(:a, :b), Tuple{Int64, Int64}}}(0.2, (a = 5, b = 2)) - MyType{Float64, NamedTuple{(:a, :b), Tuple{Int64, Int64}}}(0.4, (a = 4, b = 2)) - MyType{Float64, NamedTuple{(:a, :b), Tuple{Int64, Int64}}}(0.6, (a = 3, b = 2)) - MyType{Float64, NamedTuple{(:a, :b), Tuple{Int64, Int64}}}(0.8, (a = 2, b = 2)) - MyType{Float64, NamedTuple{(:a, :b), Tuple{Int64, Int64}}}(1.0, (a = 1, b = 2)) +```@repl advanced1 +s = [MyType(i/5, a=6-i, b=2) for i in 1:5] ``` The default `StructArray` does not unpack the `NamedTuple`: -```jldoctest advanced1 -julia> sa = StructArray(s); - -julia> sa.rest -5-element Vector{NamedTuple{(:a, :b), Tuple{Int64, Int64}}}: - (a = 5, b = 2) - (a = 4, b = 2) - (a = 3, b = 2) - (a = 2, b = 2) - (a = 1, b = 2) - -julia> sa.a -ERROR: type NamedTuple has no field a -Stacktrace: - [1] component -[...] +```@repl advanced1 +sa = StructArray(s); +sa.rest +sa.a ``` Suppose we wish to give the keywords their own fields. We can define custom `staticschema`, `component`, and `createinstance` methods for `MyType`: -```jldoctest advanced1 -julia> function StructArrays.staticschema(::Type{MyType{T, NamedTuple{names, types}}}) where {T, names, types} - # Define the desired names and eltypes of the "fields" - return NamedTuple{(:data, names...), Base.tuple_type_cons(T, types)} - end; - -julia> function StructArrays.component(m::MyType, key::Symbol) - # Define a component-extractor - return key === :data ? getfield(m, 1) : getfield(getfield(m, 2), key) - end; - -julia> function StructArrays.createinstance(::Type{MyType{T, NT}}, x, args...) where {T, NT} - # Generate an instance of MyType from components - return MyType(x, NT(args)) - end; +```@repl advanced1 +function StructArrays.staticschema(::Type{MyType{T, NamedTuple{names, types}}}) where {T, names, types} + # Define the desired names and eltypes of the "fields" + return NamedTuple{(:data, names...), Base.tuple_type_cons(T, types)} +end; + +function StructArrays.component(m::MyType, key::Symbol) + # Define a component-extractor + return key === :data ? getfield(m, 1) : getfield(getfield(m, 2), key) +end; + +function StructArrays.createinstance(::Type{MyType{T, NT}}, x, args...) where {T, NT} + # Generate an instance of MyType from components + return MyType(x, NT(args)) +end; ``` and now: -```jldoctest advanced1 -julia> sa = StructArray(s); - -julia> sa.a -5-element Vector{Int64}: - 5 - 4 - 3 - 2 - 1 - -julia> sa.b -5-element Vector{Int64}: - 2 - 2 - 2 - 2 - 2 +```@repl advanced1 +sa = StructArray(s); +sa.a +sa.b ``` The above strategy has been tested and implemented in [GeometryBasics.jl](https://github.com/JuliaGeometry/GeometryBasics.jl). diff --git a/docs/src/index.md b/docs/src/index.md index df7caa65..73dfed5c 100644 --- a/docs/src/index.md +++ b/docs/src/index.md @@ -9,73 +9,39 @@ The package was largely inspired by the `Columns` type in [IndexedTables](https: ## Collection and initialization One can create a `StructArray` by providing the struct type and a tuple or NamedTuple of field arrays: -```jldoctest intro -julia> using StructArrays - -julia> struct Foo{T} - a::T - b::T - end - -julia> adata = [1 2; 3 4]; bdata = [10 20; 30 40]; - -julia> x = StructArray{Foo}((adata, bdata)) -2×2 StructArray(::Matrix{Int64}, ::Matrix{Int64}) with eltype Foo: - Foo{Int64}(1, 10) Foo{Int64}(2, 20) - Foo{Int64}(3, 30) Foo{Int64}(4, 40) +```@repl intro +using StructArrays +struct Foo{T} + a::T + b::T +end +adata = [1 2; 3 4]; bdata = [10 20; 30 40]; +x = StructArray{Foo}((adata, bdata)) ``` You can also initialze a StructArray by passing in a NamedTuple, in which case the name (rather than the order) specifies how the input arrays are assigned to fields: -```jldoctest intro -julia> x = StructArray{Foo}((b = adata, a = bdata)) # initialize a with bdata and vice versa -2×2 StructArray(::Matrix{Int64}, ::Matrix{Int64}) with eltype Foo: - Foo{Int64}(10, 1) Foo{Int64}(20, 2) - Foo{Int64}(30, 3) Foo{Int64}(40, 4) +```@repl intro +x = StructArray{Foo}((b = adata, a = bdata)) # initialize a with bdata and vice versa ``` If a struct is not specified, a StructArray with Tuple or NamedTuple elements will be created: -```jldoctest intro -julia> x = StructArray((adata, bdata)) -2×2 StructArray(::Matrix{Int64}, ::Matrix{Int64}) with eltype Tuple{Int64, Int64}: - (1, 10) (2, 20) - (3, 30) (4, 40) - -julia> x = StructArray((a = adata, b = bdata)) -2×2 StructArray(::Matrix{Int64}, ::Matrix{Int64}) with eltype NamedTuple{(:a, :b), Tuple{Int64, Int64}}: - (a = 1, b = 10) (a = 2, b = 20) - (a = 3, b = 30) (a = 4, b = 40) +```@repl intro +x = StructArray((adata, bdata)) +x = StructArray((a = adata, b = bdata)) ``` It's also possible to create a `StructArray` by choosing a particular dimension to interpret as the components of a struct: -```jldoctest intro -julia> x = StructArray{Complex{Int}}(adata; dims=1) # along dimension 1, the first item `re` and the second is `im` -2-element StructArray(view(::Matrix{Int64}, 1, :), view(::Matrix{Int64}, 2, :)) with eltype Complex{Int64}: - 1 + 3im - 2 + 4im - -julia> x = StructArray{Complex{Int}}(adata; dims=2) # along dimension 2, the first item `re` and the second is `im` -2-element StructArray(view(::Matrix{Int64}, :, 1), view(::Matrix{Int64}, :, 2)) with eltype Complex{Int64}: - 1 + 2im - 3 + 4im +```@repl intro +x = StructArray{Complex{Int}}(adata; dims=1) # along dimension 1, the first item `re` and the second is `im` +x = StructArray{Complex{Int}}(adata; dims=2) # along dimension 2, the first item `re` and the second is `im` ``` One can also create a `StructArray` from an iterable of structs without creating an intermediate `Array`: -```jldoctest intro -julia> StructArray(log(j+2.0*im) for j in 1:10) -10-element StructArray(::Vector{Float64}, ::Vector{Float64}) with eltype ComplexF64: - 0.8047189562170501 + 1.1071487177940904im - 1.0397207708399179 + 0.7853981633974483im - 1.2824746787307684 + 0.5880026035475675im - 1.4978661367769954 + 0.4636476090008061im - 1.683647914993237 + 0.3805063771123649im - 1.8444397270569681 + 0.3217505543966422im - 1.985145956776061 + 0.27829965900511133im - 2.1097538525880535 + 0.24497866312686414im - 2.2213256282451583 + 0.21866894587394195im - 2.3221954495706862 + 0.19739555984988078im +```@repl intro +StructArray(log(j+2.0*im) for j in 1:10) ``` Another option is to create an uninitialized `StructArray` and then fill it with data. Just like in normal arrays, this is done with the `undef` syntax: diff --git a/src/utils.jl b/src/utils.jl index c4874d72..741e630d 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -144,9 +144,9 @@ julia> s_pooled = StructArrays.replace_storage(s) do v isbitstype(eltype(v)) ? v : convert(PooledArray, v) end $(if VERSION < v"1.6-" - "3-element StructArray(::UnitRange{Int64}, ::PooledArray{String,UInt32,1,Array{UInt32,1}}) with eltype NamedTuple{(:a, :b),Tuple{Int64,String}}:" + "3-element StructArray(::UnitRange{Int64}, ::PooledArray{String,UInt32,1,Array{UInt32,1}}) with eltype $(NamedTuple{(:a, :b),Tuple{Int64,String}}):" else - "3-element StructArray(::UnitRange{Int64}, ::PooledVector{String, UInt32, Vector{UInt32}}) with eltype NamedTuple{(:a, :b), Tuple{Int64, String}}:" + "3-element StructArray(::UnitRange{Int64}, ::PooledVector{String, UInt32, Vector{UInt32}}) with eltype $(NamedTuple{(:a, :b), Tuple{Int64, String}}):" end) (a = 1, b = "string") (a = 2, b = "string") From 58a54f9469f5385705b052c4539aafc41582c29d Mon Sep 17 00:00:00 2001 From: N5N3 <2642243996@qq.com> Date: Thu, 26 Oct 2023 22:10:47 +0800 Subject: [PATCH 4/9] move `Adapt` to ext And use curried adapter to avoid possible instability --- Project.toml | 8 +++++--- ext/StructArraysAdaptExt.jl | 16 ++++++++++++++++ src/StructArrays.jl | 5 +---- 3 files changed, 22 insertions(+), 7 deletions(-) create mode 100644 ext/StructArraysAdaptExt.jl diff --git a/Project.toml b/Project.toml index 2310f14a..fa804120 100644 --- a/Project.toml +++ b/Project.toml @@ -3,21 +3,22 @@ uuid = "09ab397b-f2b6-538f-b94a-2f83cf4a842a" version = "0.6.16" [deps] -Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" ConstructionBase = "187b0558-2788-49d3-abe0-74a17ed4e7c9" DataAPI = "9a962f9c-6df0-11e9-0e5d-c546b8b5ee8a" Tables = "bd369af6-aec1-5ad0-b16a-f7cc5008161c" [weakdeps] +Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527" StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" [extensions] +StructArraysAdaptExt = "Adapt" StructArraysGPUArraysCoreExt = "GPUArraysCore" StructArraysStaticArraysExt = "StaticArrays" [compat] -Adapt = "1, 2, 3" +Adapt = "2, 3" ConstructionBase = "1" DataAPI = "1" GPUArraysCore = "0.1.2" @@ -26,6 +27,7 @@ Tables = "1" julia = "1.6" [extras] +Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4" GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527" JLArrays = "27aeb0d3-9eb9-45fb-866b-73c2ecf80fcb" @@ -38,4 +40,4 @@ TypedTables = "9d95f2ec-7b3d-5a63-8d20-e2491e220bb9" WeakRefStrings = "ea10d353-3f73-51f8-a26c-33c1cb351aa5" [targets] -test = ["Test", "JLArrays", "StaticArrays", "OffsetArrays", "PooledArrays", "TypedTables", "WeakRefStrings", "Documenter", "SparseArrays", "GPUArraysCore"] +test = ["Test", "JLArrays", "StaticArrays", "OffsetArrays", "PooledArrays", "TypedTables", "WeakRefStrings", "Documenter", "SparseArrays", "GPUArraysCore", "Adapt"] diff --git a/ext/StructArraysAdaptExt.jl b/ext/StructArraysAdaptExt.jl new file mode 100644 index 00000000..2607105d --- /dev/null +++ b/ext/StructArraysAdaptExt.jl @@ -0,0 +1,16 @@ +module StructArraysAdaptExt +# Use Adapt allows for automatic conversion of CPU to GPU StructArrays +using Adapt, StructArrays +@static if !applicable(Adapt.adapt, Int) + # Adapt.jl has curried support, implement it ourself + adpat(to) = Base.Fix1(Adapt.adapt, to) + if VERSION < v"1.9.0-DEV.857" + @eval function adapt(to::Type{T}) where {T} + (@isdefined T) || return Base.Fix1(Adapt.adapt, to) + AT = Base.Fix1{typeof(Adapt.adapt),Type{T}} + return $(Expr(:new, :AT, :(Adapt.adapt), :to)) + end + end +end +Adapt.adapt_structure(to, s::StructArray) = replace_storage(adapt(to), s) +end diff --git a/src/StructArrays.jl b/src/StructArrays.jl index 8cfbdc48..4130d059 100644 --- a/src/StructArrays.jl +++ b/src/StructArrays.jl @@ -25,11 +25,8 @@ function refvalue(s::StructArray{T}, v::Tup) where {T} createinstance(T, map(refvalue, components(s), v)...) end -# Use Adapt allows for automatic conversion of CPU to GPU StructArrays -import Adapt -Adapt.adapt_structure(to, s::StructArray) = replace_storage(x->Adapt.adapt(to, x), s) - @static if !isdefined(Base, :get_extension) + include("../ext/StructArraysAdaptExt.jl") include("../ext/StructArraysGPUArraysCoreExt.jl") include("../ext/StructArraysStaticArraysExt.jl") end From bc629daaf2ce444d2b1537c095df7821c9c45a3e Mon Sep 17 00:00:00 2001 From: N5N3 <2642243996@qq.com> Date: Thu, 2 Nov 2023 09:14:40 +0800 Subject: [PATCH 5/9] Apply suggestions from code review --- Project.toml | 3 +++ ext/StructArraysAdaptExt.jl | 2 +- 2 files changed, 4 insertions(+), 1 deletion(-) diff --git a/Project.toml b/Project.toml index fa804120..2673d99b 100644 --- a/Project.toml +++ b/Project.toml @@ -3,8 +3,11 @@ uuid = "09ab397b-f2b6-538f-b94a-2f83cf4a842a" version = "0.6.16" [deps] +Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" ConstructionBase = "187b0558-2788-49d3-abe0-74a17ed4e7c9" DataAPI = "9a962f9c-6df0-11e9-0e5d-c546b8b5ee8a" +GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527" +StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" Tables = "bd369af6-aec1-5ad0-b16a-f7cc5008161c" [weakdeps] diff --git a/ext/StructArraysAdaptExt.jl b/ext/StructArraysAdaptExt.jl index 2607105d..7809abe9 100644 --- a/ext/StructArraysAdaptExt.jl +++ b/ext/StructArraysAdaptExt.jl @@ -2,7 +2,7 @@ module StructArraysAdaptExt # Use Adapt allows for automatic conversion of CPU to GPU StructArrays using Adapt, StructArrays @static if !applicable(Adapt.adapt, Int) - # Adapt.jl has curried support, implement it ourself + # Adapt.jl has no curried support, implement it ourself adpat(to) = Base.Fix1(Adapt.adapt, to) if VERSION < v"1.9.0-DEV.857" @eval function adapt(to::Type{T}) where {T} From 05421525a78b6eae6adb3efdf4182f6ed87370a6 Mon Sep 17 00:00:00 2001 From: N5N3 <2642243996@qq.com> Date: Thu, 9 Nov 2023 00:34:51 +0800 Subject: [PATCH 6/9] Adopt code style suggestion. --- ext/StructArraysStaticArraysExt.jl | 34 +++++++++++++++--------------- 1 file changed, 17 insertions(+), 17 deletions(-) diff --git a/ext/StructArraysStaticArraysExt.jl b/ext/StructArraysStaticArraysExt.jl index 63401539..fab643a0 100644 --- a/ext/StructArraysStaticArraysExt.jl +++ b/ext/StructArraysStaticArraysExt.jl @@ -47,10 +47,23 @@ using Base.Broadcast: Broadcasted return _broadcast(f, Size(map(length, ax)), argsizes, as...) end +# A functor generates the ith component of StructStaticBroadcast. +struct Similar_ith{SA, E<:Tuple} + elements::E + Similar_ith{SA}(elements::Tuple) where {SA} = new{SA, typeof(elements)}(elements) +end +function (s::Similar_ith{SA})(i::Int) where {SA} + ith_elements = ntuple(Val(length(s.elements))) do j + getfield(s.elements[j], i) + end + ith_SA = similar_type(SA, fieldtype(eltype(SA), i)) + return @inbounds ith_SA(ith_elements) +end + @inline function _broadcast(f, sz::Size{newsize}, s::Tuple{Vararg{Size}}, a...) where {newsize} first_staticarray = first_statictype(a...) elements, ET = if prod(newsize) == 0 - # Use inference to get eltype in empty case (see also comments in _map) + # Use inference to get eltype in empty case (following StaticBroadcast defined in StaticArrays.jl) eltys = Tuple{map(eltype, a)...} (), Core.Compiler.return_type(f, eltys) else @@ -58,24 +71,11 @@ end temp, eltype(temp) end if isnonemptystructtype(ET) - @static if VERSION >= v"1.7" - arrs = ntuple(Val(fieldcount(ET))) do i - @inbounds similar_type(first_staticarray, fieldtype(ET, i), sz)(_getfields(elements, i)) - end - else - similarET(::Type{SA}, ::Type{T}) where {SA, T} = i -> @inbounds similar_type(SA, fieldtype(T, i), sz)(_getfields(elements, i)) - arrs = ntuple(similarET(first_staticarray, ET), Val(fieldcount(ET))) - end + SA = similar_type(first_staticarray, ET, sz) + arrs = ntuple(Similar_ith{SA}(elements), Val(fieldcount(ET))) return StructArray{ET}(arrs) - end - @inbounds return similar_type(first_staticarray, ET, sz)(elements) -end - -@inline function _getfields(x::Tuple, i::Int) - if @generated - return Expr(:tuple, (:(getfield(x[$j], i)) for j in 1:fieldcount(x))...) else - return map(Base.Fix2(getfield, i), x) + @inbounds return similar_type(first_staticarray, ET, sz)(elements) end end From 02257b92a42a6ddaff5b3a81c32424119f93949b Mon Sep 17 00:00:00 2001 From: N5N3 <2642243996@qq.com> Date: Thu, 9 Nov 2023 00:40:51 +0800 Subject: [PATCH 7/9] restrict Adapt compat --- Project.toml | 2 +- ext/StructArraysAdaptExt.jl | 11 ----------- 2 files changed, 1 insertion(+), 12 deletions(-) diff --git a/Project.toml b/Project.toml index 2673d99b..53946677 100644 --- a/Project.toml +++ b/Project.toml @@ -21,7 +21,7 @@ StructArraysGPUArraysCoreExt = "GPUArraysCore" StructArraysStaticArraysExt = "StaticArrays" [compat] -Adapt = "2, 3" +Adapt = "3.4" ConstructionBase = "1" DataAPI = "1" GPUArraysCore = "0.1.2" diff --git a/ext/StructArraysAdaptExt.jl b/ext/StructArraysAdaptExt.jl index 7809abe9..44e99f89 100644 --- a/ext/StructArraysAdaptExt.jl +++ b/ext/StructArraysAdaptExt.jl @@ -1,16 +1,5 @@ module StructArraysAdaptExt # Use Adapt allows for automatic conversion of CPU to GPU StructArrays using Adapt, StructArrays -@static if !applicable(Adapt.adapt, Int) - # Adapt.jl has no curried support, implement it ourself - adpat(to) = Base.Fix1(Adapt.adapt, to) - if VERSION < v"1.9.0-DEV.857" - @eval function adapt(to::Type{T}) where {T} - (@isdefined T) || return Base.Fix1(Adapt.adapt, to) - AT = Base.Fix1{typeof(Adapt.adapt),Type{T}} - return $(Expr(:new, :AT, :(Adapt.adapt), :to)) - end - end -end Adapt.adapt_structure(to, s::StructArray) = replace_storage(adapt(to), s) end From b4de96bfcb644f33af9dc1b19cb46c2f1561b35b Mon Sep 17 00:00:00 2001 From: N5N3 <2642243996@qq.com> Date: Thu, 9 Nov 2023 00:44:43 +0800 Subject: [PATCH 8/9] Add empty bc test. --- test/runtests.jl | 2 ++ 1 file changed, 2 insertions(+) diff --git a/test/runtests.jl b/test/runtests.jl index 61d9be34..611dfc65 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -1312,6 +1312,8 @@ Base.BroadcastStyle(::Broadcast.ArrayStyle{MyArray2}, S::Broadcast.DefaultArrayS b = @SMatrix [0. for i in 1:10, j in 1:10] s = StructArray{ComplexF64}((a , b)) @test (@inferred bclog(s)) isa typeof(s) + s0 = StructArray{ComplexF64}((similar(a, Size(0,0)), similar(a, Size(0,0)))) + @test (@inferred bclog(s0)) isa typeof(s0) test_allocated(bclog, s) @test abs.(s) .+ ((1,) .+ (1,2,3,4,5,6,7,8,9,10)) isa SMatrix bc = Base.broadcasted(+, s, s, ntuple(identity, 10)); From 60a8c8cbaf395fe4960967f500f467723ff1d5e5 Mon Sep 17 00:00:00 2001 From: N5N3 <2642243996@qq.com> Date: Thu, 9 Nov 2023 01:05:31 +0800 Subject: [PATCH 9/9] define `__broadcast` ourselves --- ext/StructArraysStaticArraysExt.jl | 29 +++++++++++++++++++++++++++-- 1 file changed, 27 insertions(+), 2 deletions(-) diff --git a/ext/StructArraysStaticArraysExt.jl b/ext/StructArraysStaticArraysExt.jl index fab643a0..0c84f5b7 100644 --- a/ext/StructArraysStaticArraysExt.jl +++ b/ext/StructArraysStaticArraysExt.jl @@ -33,9 +33,9 @@ StructArrays.createinstance(T::Type{<:FieldArray}, args...) = invoke(StructArray # Broadcast overload using StaticArrays: StaticArrayStyle, similar_type, Size, SOneTo -using StaticArrays: broadcast_flatten, broadcast_sizes, first_statictype, __broadcast +using StaticArrays: broadcast_flatten, broadcast_sizes, first_statictype using StructArrays: isnonemptystructtype -using Base.Broadcast: Broadcasted +using Base.Broadcast: Broadcasted, _broadcast_getindex # StaticArrayStyle has no similar defined. # Overload `try_struct_copy` instead. @@ -79,4 +79,29 @@ end end end +# The `__broadcast` kernal is copied from `StaticArrays.jl`. +# see https://github.com/JuliaArrays/StaticArrays.jl/blob/master/src/broadcast.jl +@generated function __broadcast(f, ::Size{newsize}, s::Tuple{Vararg{Size}}, a...) where newsize + sizes = [sz.parameters[1] for sz ∈ s.parameters] + + indices = CartesianIndices(newsize) + exprs = similar(indices, Expr) + for (j, current_ind) ∈ enumerate(indices) + exprs_vals = (broadcast_getindex(sz, i, current_ind) for (i, sz) in enumerate(sizes)) + exprs[j] = :(f($(exprs_vals...))) + end + + return quote + Base.@_inline_meta + return tuple($(exprs...)) + end +end + +broadcast_getindex(::Tuple{}, i::Int, I::CartesianIndex) = return :(_broadcast_getindex(a[$i], $I)) +function broadcast_getindex(oldsize::Tuple, i::Int, newindex::CartesianIndex) + li = LinearIndices(oldsize) + ind = _broadcast_getindex(li, newindex) + return :(a[$i][$ind]) +end + end