Add support of GPU Broadcast with StructArray LHS?

[N3N5]

Add something like

import GPUArrays: BroadcastGPUArray,backend,@cartesianidx,launch_heuristic,launch_configuration,gpu_call
const GPUStore = Tuple{Vararg{BroadcastGPUArray}}
const NamedGPUStore = NamedTuple{Name,<:GPUStore} where {Name}
const StructGPUArray = StructArray{T, N, <:Union{GPUStore,NamedGPUStore}} where {T,N}
## backend for StructArray
backend(A::StructGPUArray) = backend(fieldarrays(A))
backend(t::GPUStore) = backend(typeof(t))
backend(nt::NamedGPUStore) = backend(typeof(nt).parameters[2])
backend(::Type{T}) where {T<:GPUStore} = begin
    I = mapreduce(backend,===,tuple(T.parameters...))
    I || throw("device error")
    backend(T.parameters[1])
end

## copy from GPUArrays
@inline function Base.copyto!(dest::StructGPUArray, bc::Broadcast.Broadcasted{Nothing})
    axes(dest) == axes(bc) || Broadcast.throwdm(axes(dest), axes(bc))
    isempty(dest) && return dest
    bc′ = Broadcast.preprocess(dest, bc)

    # grid-stride kernel
    function broadcast_kernel(ctx, dest, bc′, nelem)
        for i in 1:nelem
            I = @cartesianidx(dest, i)
            @inbounds dest[I] = bc′[I]
        end
        return
    end
    heuristic = launch_heuristic(backend(dest), broadcast_kernel, dest, bc′, 1)
    config = launch_configuration(backend(dest), heuristic, length(dest), typemax(Int))
    gpu_call(broadcast_kernel, dest, bc′, config.elements_per_thread;
             threads=config.threads, blocks=config.blocks)

    return dest
end

Then we can write code like:

CUDA.allowscalar(false)
a_re = CuArray{Float32}(undef,4000,4000)
a_im = CuArray{Float32}(undef,4000,4000)
a = StructArray{ComplexF32}(tuple(a_re,a_im))
b =  CUDA.randn(size(a)...)
a .= cis.(b)
isapprox(replace_storage(Array,a),cis.(collect(b)))  # true

[lcw]

@N3N5 This is useful! Have you had any luck getting other broadcasting routines to work? For example, I would love it if the following worked

using StructArrays
using StaticArrays
using CUDA
using Adapt

StructArrays.staticschema(::Type{SVector{N, T}}) where {N, T} = NTuple{N, T}
StructArrays.createinstance(::Type{<:SVector}, args...) = SVector(args...)
StructArrays.component(s::SVector, i) = getindex(s, i)

CUDA.allowscalar(false)
N = 4
b = StructArray{SVector{3, Float64}}(adapt(CuArray, (rand(Float64, N),
                                                     rand(Float64, N),
                                                     rand(Float64, N))));

c = similar(b)
c .= b .+ b
d = b .+ b

[lcw]

Okay, with #177 the following now works

using StructArrays
using StaticArrays
using CUDA
using Adapt

import GPUArrays: BroadcastGPUArray,backend,@cartesianidx,launch_heuristic,launch_configuration,gpu_call
const GPUStore = Tuple{Vararg{BroadcastGPUArray}}
const NamedGPUStore = NamedTuple{Name,<:GPUStore} where {Name}
const StructGPUArray = StructArray{T, N, <:Union{GPUStore,NamedGPUStore}} where {T,N}
## backend for StructArray
backend(A::StructGPUArray) = backend(fieldarrays(A))
backend(t::GPUStore) = backend(typeof(t))
backend(nt::NamedGPUStore) = backend(typeof(nt).parameters[2])
function backend(::Type{T}) where {T<:GPUStore}
    bs = backend.(tuple(T.parameters...))
    I = all(map(isequal(first(bs)), bs))
    I || throw("device error")
    backend(T.parameters[1])
end

## copy from GPUArrays
@inline function Base.copyto!(dest::StructGPUArray, bc::Broadcast.Broadcasted{Nothing})
    axes(dest) == axes(bc) || Broadcast.throwdm(axes(dest), axes(bc))
    isempty(dest) && return dest
    bc′ = Broadcast.preprocess(dest, bc)

    # grid-stride kernel
    function broadcast_kernel(ctx, dest, bc′, nelem)
        for i in 1:nelem
            I = @cartesianidx(dest, i)
            @inbounds dest[I] = bc′[I]
        end
        return
    end
    heuristic = launch_heuristic(backend(dest), broadcast_kernel, dest, bc′, 1)
    config = launch_configuration(backend(dest), heuristic, length(dest), typemax(Int))
    gpu_call(broadcast_kernel, dest, bc′, config.elements_per_thread;
             threads=config.threads, blocks=config.blocks)

    return dest
end

function Base.similar(bc::Broadcast.Broadcasted{StructArrays.StructArrayStyle{S}}, ::Type{ElType}) where {S<:CUDA.CuArrayStyle,N,ElType}
    if isstructtype(ElType)
        return StructArrays.buildfromschema(typ -> similar(CuArray{typ}, axes(bc)), ElType)
    else
        return similar(CuArray{ElType}, axes(bc))
    end
end

StructArrays.staticschema(::Type{SVector{N, T}}) where {N, T} = NTuple{N, T}
StructArrays.createinstance(::Type{<:SVector}, args...) = SVector(args...)
StructArrays.component(s::SVector, i) = getindex(s, i)

CUDA.allowscalar(false)
N = 4
b = StructArray{SVector{3, Float64}}(adapt(CuArray, (rand(Float64, N),
                                                     rand(Float64, N),
                                                     rand(Float64, N))));

c = similar(b);
c .= b .+ b;
d = b .+ b;

nothing