Separate arguments for map! / broadcasting kernel

[wsmoses]

This is part of ongoing work to make Enzyme + CUDA.jl work nicely from a user point of view.

cc @vchuravy

The implementation of map! (https://github.com/JuliaGPU/GPUArrays.jl/blob/ec9fe5b6f7522902e444c95a0c9248a4bc55d602/src/host/broadcast.jl#L120C46-L120C59) creates a broadcasted object which captures all of the arguments to map!. This is then passed to the kernel.

This makes things more difficult to successfully differentiate (as well as apply additional optimization to). Specifically if one of the arguments was non-differentiable (i.e. Const) and another argument was (i.e. Duplicated) then the resulting capturing variable would contain both differntiable and non-differentiable data (and any uses of it may potentially be activity-unstable). See https://enzyme.mit.edu/julia/stable/faq/#Activity-of-temporary-storage for a more thorough description as to how this presents difficulties.

Is it possible to have the implementation be something like:

function map_kernel(ctx, dest, nelem, args...)
   bc = Broadcast.instantiate(Broadcast.broadcasted(f, args...))
   ...

[maleadt]

Is it possible to have the implementation be something like:

function map_kernel(ctx, dest, nelem, args...)
   bc = Broadcast.instantiate(Broadcast.broadcasted(f, args...))
   ...

No, Broadcast is a host-side abstraction that's not supported on-device.