Broadcasting on `ArrayPartition`s is not applied recursively

[efaulhaber]

julia> x = CUDA.zeros(10_000_000);

julia> y = ArrayPartition(x);

julia> y .= 0
ERROR: Scalar indexing is disallowed.
Invocation of setindex! resulted in scalar indexing of a GPU array.
This is typically caused by calling an iterating implementation of a method.
Such implementations *do not* execute on the GPU, but very slowly on the CPU,
and therefore should be avoided.

This should recursively apply a broadcast to the elements of the partition, which should then correctly dispatch to the GPU code. This doesn't work, so some CPU code is called instead.

Edit: It looks like a simple

function Base.fill!(m::ArrayPartition, x)
    for i in eachindex(m.x)
        fill!(m.x[i], x)
    end

    return m
end

might do the trick.

[mateuszbaran]

I've worked a bit with recursive broadcasting and it's a very tough problem. The code is very complex and Julia's optimizer gives up with inlining fairly quickly when dealing with recursion. And then performance for small ArrayPartitions get's bad. The performance hit would likely be acceptable when you keep inside such large arrays but in other cases it would not.

My advice would be to start special-casing broadcasting ArrayPartition that holds GPU arrays inside and don't touch the generic case. But if fill! is enough for you then I wouldn't bother with full broadcasting.

[efaulhaber]

Looks like broadcasting ArrayPartitions to ArrayPartitions already works with the generic copyto! code. Only .= 0 didn't work because it is not going through RecursiveArrayTools.jl's code but is directly translated to fill!.