Reverse mode: `ijl_eqtable_put` via `task_local_storage` blocks CuArray reductions
- Dominant language
- Julia
- Stars
- 586
- Forks
- 108
- Avg merge
- 1d 5h
- Merged PRs (30d)
- 44
Description
Reverse-mode differentiation of a `CuArray` reduction fails on the scalar-indexing
guard in `GPUArraysCore`, which reads/writes `task_local_storage`.
```julia
using CUDA, Enzyme
f(x) = sum(abs2, x)
x = CuArray(randn(Float32, 8))
Enzyme.gradient(set_runtime_activity(Reverse), f, x)
```
```
EnzymeNoDerivativeError:
No create nofree of empty function (ijl_eqtable_put) ijl_eqtable_put)
at context: call fastcc void @julia_task_local_storage_16473(i32 zeroext 0) #929, !dbg !1032 (julia_task_local_storage_16473)
Stacktrace:
[1] macro expansion
@ GPUArraysCore/src/GPUArraysCore.jl:205
[2] #_mapreduce#48
@ GPUArrays/src/host/mapreduce.jl:88
Failure within method:
var"#_mapreduce#48"(::Colon, ::Nothing, ::typeof(GPUArrays._mapreduce), ::typeof(abs2),
::typeof(Base.add_sum), ::CuArray{Float32, 1, CUDA.DeviceMemory})
@ GPUArrays GPUArrays/src/host/mapreduce.jl:40
```
Contributor guide
Assessment
This issue has not been assessed yet.