`cudaconvert` for closures

[vchuravy]

A flatten broadcast is no longer amendable to cudaconvert since there is a closure involved (see SciML/RecursiveArrayTools.jl#40)

https://github.com/JuliaLang/julia/blob/26b6a5811c464bd66a371ddb341df140c677a2a3/base/broadcast.jl#L307-L313

julia> x.body
getfield(Base.Broadcast, Symbol("##1#2")){bc,makeargs} where makeargs

I don't know if it is feasible for us to define a version of cudaconvert that could inspect a closure and apply cudaconvert to its fields

[vchuravy]

MWE without ArrayPartition:

using CuArrays
import Base.Broadcast: flatten, broadcasted

# construct bc directly
bc = broadcasted(+, CuArray(zeros(3,3)), broadcasted(*, CuArray(ones(3, 3)), CuArray(ones(3))))
fbc = flatten(bc)

# this works
copy(bc)

# this doesn't
copy(fbc)

[maleadt]

Shorter repro

using CuArrays
import Base.Broadcast: flatten, broadcasted, Broadcasted

function main()
    # construct bc directly
    A = CuArray([1])
    bc = broadcasted(identity, broadcasted(identity, A))
    bc = convert(Broadcasted{Nothing}, bc)

    # this works
    @show copyto!(A, bc)

    # this doesn't
    fbc = flatten(bc)
    @show copyto!(A, fbc)
end

[maleadt]

https://discourse.julialang.org/t/is-there-a-way-to-modify-captured-variables-in-a-closure/31213/12

using CUDA

function main()
    a = CUDA.zeros(1)
    function kernel()
        a[1] = 1
        return
    end

    flds = Any[CUDA.cudaconvert(getfield(kernel, field)) for field in 1:nfields(kernel)]
    kernel′ = ccall(:jl_new_structt, Any, (Any, Any),
                    typeof(kernel).name.wrapper{typeof.(flds)...}, Tuple(flds))
    @cuda kernel′()

    Array(a)
end

julia> main()
1-element Vector{Float32}:
 1.0

🤯

Although I guess this would break calling custom objects where the typevars don't match up with the fields. @vchuravy, are you doing anything like that?