JuliaGPU / JuliaGPU/GPUArrays.jl

Separate arguments for map! / broadcasting kernel

Open
#535 1 comment 0 reactions 0 assignees View on GitHub
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.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.