Make StructArrayStyle track inputs dimension

[N5N3]

test added.
Close #185

[piever]

Could this maybe be written using the function StructArrayStyle(...) syntax and be split over a couple of lines, for the sake of readability?

[piever]

I confess I'm not fully familiar with the broadcast machinery, but I imagine this comment could be either expanded or removed, as it is it is a bit cryptic (it only makes sense if one knows the diff from this PR).

[piever]

Thanks for the PR, this is very helpful! I've left some minor comments about style.

[N5N3]

The style confilct is resolved with c7a4dc8
Now we have:

julia> using StaticArrays, StructArrays

julia> a = @SArray (randn(3,3))
3×3 SMatrix{3, 3, Float64, 9} with indices SOneTo(3)×SOneTo(3):
 -2.35452   -0.492535   0.401572
 -0.133826   1.66839    2.41737
  1.08633   -1.31099   -0.0771547

julia> b = StructArray([1;;1;;2+im])
1×3 StructArray(::Matrix{Int64}, ::Matrix{Int64}) with eltype Complex{Int64}:
 1+0im  1+0im  2+1im

julia> a .+ b
3×3 StructArray(::Matrix{Float64}, ::Matrix{Float64}) with eltype ComplexF64: # before 3 x 3 Matrix{ComplexF64}
 -1.35452+0.0im   0.507465+0.0im  2.40157+1.0im
 0.866174+0.0im    2.66839+0.0im  4.41737+1.0im
  2.08633+0.0im  -0.310988+0.0im  1.92285+1.0im

julia> a .+ cu(b) # they have style confict
3×3 StructArray(::Matrix{Float64}, ::Matrix{Float64}) with eltype ComplexF64:
     2.037+0.0im   3.35919+0.0im  3.90445+1.0im
   2.25239+0.0im  0.468944+0.0im   2.1828+1.0im
 -0.506096+0.0im   1.52528+0.0im  3.47241+1.0im

julia> collect(a) .+ cu(b) # This is still done on CPU!
3×3 StructArray(::CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}, ::CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}) with eltype ComplexF64:
     2.037+0.0im   3.35919+0.0im  3.90445+1.0im
   2.25239+0.0im  0.468944+0.0im   2.1828+1.0im
 -0.506096+0.0im   1.52528+0.0im  3.47241+1.0im

I only test it with StaticArrays.jl and CUDA.jl (only locally). do we have another wide-used package have it's own style defined?

This PR is more "aggressive" than the original version in #136.
As it hopes User's BroadcastStyle has a similar defined. (That should be true for our GPU ecosystem.)
StaticArray is a little different as it has a @generated based copy, so I just keep the original behaviour via 793d035.

[piever]

Getting this to work with CUDA would certainly be great... Pinging @timholy as he worked on the initial implementation, especially because in the intermediate comments, there were attempts to create such a fallback.

[N5N3]

At present, our GPU broadcast is dispatched based on dest's type:

@inline function Base.copyto!(dest::BroadcastGPUArray, bc::Broadcasted{Nothing})

Once we change it to something like:

@inline function Base.copyto!(dest::AbstractArray, bc::Broadcasted{<:AbstractGPUArrayStyle})

Then the unwrapped style above could activate the kernal automaticly.

Of cource there's still something left to handle manually. (e.g. backend(dest)) But I think this does help us reuse more code.

[piever]

Also fixes #189. Related: #150. Potentially related: #94. CC: @lcw (involved in both the issue and PR mentioned).

So, I do have some comments on the code itself, but I think first we should figure out exactly what we want here.

If I understand correctly, the key problem is that when computing the overall broadcast style, one should figure out the overall broadcast style B of the components (that is, both the components of the structarrays and the "non-structarrays" elements of the broadcast), then choose if the outcome has broadcast type B (if eltype is scalar) or StructArrayStyle{B}` (if eltype is struct).

At that point the issue becomes how to make sure that the correct mechanism for building the correct destination type is triggered. Expecting things other than StaticArray (which can in principle be immutable) to implement similar for their broadcasting style seems reasonable.

Things I don't understand:

• Does this fully address Add support of GPU Broadcast with StructArray LHS? #150, or does it need some changes on the GPUArrays.jl side?
• Does this machinery also work for type unstable broadcast?

If both work, I imagine we should test them (potentially using something like JLArray for the first one).

If this is not enough to get broadcast with GPU arrays to work (and we need to rethink the dispatch on copyto!), then maybe it's safer to figure that out (I imagine it'd require some changes at GPUArrays.jl) before merging here. (The first commit is of course fine to merge as a bugfix.)

[N5N3]

Does this machinery also work for type unstable broadcast?

Yes, as our general copy also use similar(bc′, typeof(val)) to form the return.

Does this fully address #150, or does it need some changes on the GPUArrays.jl side?

We need some change like JuliaGPU/GPUArrays.jl#295 to enable broadcast!. For broadcast, we need to add copyto!(dest, bc::Broadcasted{<:AbstractGPUArrayStyle}).
Also #150 needs to define backend for StructArray, so a glue package is needed if we dont want more dependency/@require.

I'll remove the 2nd/3rd commit, and add test for unstable broadcast.

[piever]

Yes, I think a good action plan would be to merge now something much simpler here (essentially the first commit). Then,

• mention over at use style dispatch in materialize JuliaGPU/GPUArrays.jl#295 that the change is indeed needed to get broadcast of StructArrays and GPU arrays working together
• maybe suggest moving backend to Adapt.jl? It is already a dependency and very lightweight

I'll remove the 2nd/3rd commit, and add test for unstable broadcast.

Thanks, I think that's best (there were also some comments on the first commit that were addressed somewhere within the later commits). Those 2 commits can be a draft pull request to master once the bugfix is merged (so we can make sure that things work with the changes in GPUArrays).

[piever]

Merged in #212, thanks again for the PR! If you want, feel free to add back the other changes against current master.