JuliaGPU / JuliaGPU/GPUArrays.jl
Separate arguments for map! / broadcasting kernel
- Dominant language
- Julia
- Stars
- 450
- Forks
- 104
- Avg merge
- 1d 4h
- Merged PRs (30d)
- 10
Description
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...))
...
```
Contributor guide
No contributing guide indexed for this repository
Assessment
This issue has not been assessed yet.