`sum(abs2.(x))` does not work on GPU
- Dominant language
- Julia
- Stars
- 586
- Forks
- 108
- Avg merge
- 1d 5h
- Merged PRs (30d)
- 44
Description
Enzyme doesn't seem to work with a broadcast and reduction on the GPU. Here is MWE
```julia
using Enzyme
using CUDA
f(x) = sum(abs2.(x))
x = CUDA.ones(64)
dx = zero(x)
autodiff(Reverse, f, Active, Duplicated(x, dx))
No create nofree of empty function (jl_gc_safe_enter) jl_gc_safe_enter)
at context: call fastcc void @julia__launch_configuration_979_6800([2 x i64]* noalias nocapture nofree noundef nonnull writeonly sret([2 x i64]) align 8 dereferenceable(16) %6, { i64, {} addrspace(10)* } addrspace(11)* nocapture nofree noundef nonnull readonly align 8 dereferenceable(32) %32) #723, !dbg !1546 (julia__launch_configuration_979_6800)
Stacktrace:
[1] launch_configuration
@ ~/.julia/packages/CUDA/2kjXI/lib/cudadrv/occupancy.jl:56
[2] #launch_heuristic#1200
@ ~/.julia/packages/CUDA/2kjXI/src/gpuarrays.jl:22
[3] launch_heuristic
@ ~/.julia/packages/CUDA/2kjXI/src/gpuarrays.jl:15
[4] _copyto!
@ ~/.julia/packages/GPUArrays/qt4ax/src/host/broadcast.jl:78
[5] copyto!
@ ~/.julia/packages/GPUArrays/qt4ax/src/host/broadcast.jl:44
[6] copy
@ ~/.julia/packages/GPUArrays/qt4ax/src/host/broadcast.jl:29
[7] materialize
@ ./broadcast.jl:903
[8] f
@ ./REPL[4]:1
Stacktrace:
[1] launch_configuration
@ ~/.julia/packages/CUDA/2kjXI/lib/cudadrv/occupancy.jl:56 [inlined]
[2] #launch_heuristic#1200
@ ~/.julia/packages/CUDA/2kjXI/src/gpuarrays.jl:22 [inlined]
[3] launch_heuristic
@ ~/.julia/packages/CUDA/2kjXI/src/gpuarrays.jl:15 [inlined]
[4] _copyto!
@ ~/.julia/packages/GPUArrays/qt4ax/src/host/broadcast.jl:78 [inlined]
[5] copyto!
@ ~/.julia/packages/GPUArrays/qt4ax/src/host/broadcast.jl:44 [inlined]
[6] copy
@ ~/.julia/packages/GPUArrays/qt4ax/src/host/broadcast.jl:29 [inlined]
[7] materialize
@ ./broadcast.jl:903 [inlined]
[8] f
@ ./REPL[4]:1 [inlined]
[9] diffejulia_f_5251wrap
@ ./REPL[4]:0
[10] macro expansion
@ ~/.julia/packages/Enzyme/RTS5U/src/compiler.jl:8398 [inlined]
[11] enzyme_call
@ ~/.julia/packages/Enzyme/RTS5U/src/compiler.jl:7950 [inlined]
[12] CombinedAdjointThunk
@ ~/.julia/packages/Enzyme/RTS5U/src/compiler.jl:7723 [inlined]
[13] autodiff
@ ~/.julia/packages/Enzyme/RTS5U/src/Enzyme.jl:491 [inlined]
[14] autodiff(mode::ReverseMode{false, false, FFIABI, false, false}, f::typeof(f), ::Type{Active}, args::Duplicated{CuArray{Float32, 1, CUDA.DeviceMemory}})
@ Enzyme ~/.julia/packages/Enzyme/RTS5U/src/Enzyme.jl:512
[15] top-level scope
@ REPL[13]:1
```
If I instead do
```julia
g(x) = sum(abs2, x)
autodiff(Reverse, f, Active, Duplicated(x, dx))
ERROR: Constant memory is stored (or returned) to a differentiable variable.
As a result, Enzyme cannot provably ensure correctness and throws this error.
This might be due to the use of a constant variable as temporary storage for active memory (https://enzyme.mit.edu/julia/stable/faq/#Runtime-Activity).
If Enzyme should be able to prove this use non-differentable, open an issue!
To work around this issue, either:
a) rewrite this variable to not be conditionally active (fastest, but requires a code change), or
b) set the Enzyme mode to turn on runtime activity (e.g. autodiff(set_runtime_activity(Reverse), ...) ). This will maintain correctness, but may slightly reduce performance.
Mismatched activity for: %getfield26 = phi {} addrspace(10)* [ %newstruct, %L45 ], [ %getfield13, %L41 ], !dbg !1602 const val: %getfield13 = load atomic {} addrspace(10)*, {} addrspace(10)** %18 unordered, align 8, !dbg !1583, !tbaa !1587, !alias.scope !1552, !noalias !1555, !nonnull !0, !enzyme_type !1588
Type tree: {[-1]:Pointer, [-1,0]:Pointer, [-1,0,0]:Pointer, [-1,0,8]:Integer, [-1,0,9]:Integer, [-1,0,10]:Integer, [-1,0,11]:Integer, [-1,0,12]:Integer, [-1,0,13]:Integer, [-1,0,14]:Integer, [-1,0,15]:Integer, [-1,0,16]:Integer, [-1,0,17]:Integer, [-1,0,18]:Integer, [-1,0,19]:Integer, [-1,0,20]:Integer, [-1,0,21]:Integer, [-1,0,22]:Integer, [-1,0,23]:Integer, [-1,0,24]:Integer, [-1,0,25]:Integer, [-1,0,26]:Integer, [-1,0,27]:Integer, [-1,0,28]:Integer, [-1,0,29]:Integer, [-1,0,30]:Integer, [-1,0,31]:Integer, [-1,0,32]:Integer, [-1,0,33]:Integer, [-1,0,34]:Integer, [-1,0,35]:Integer, [-1,0,36]:Integer, [-1,0,37]:Integer, [-1,0,38]:Integer, [-1,0,39]:Integer, [-1,8]:Integer, [-1,9]:Integer, [-1,10]:Integer, [-1,11]:Integer, [-1,12]:Integer, [-1,13]:Integer, [-1,14]:Integer, [-1,15]:Integer, [-1,16]:Integer, [-1,17]:Integer, [-1,18]:Integer, [-1,19]:Integer, [-1,20]:Integer, [-1,21]:Integer, [-1,22]:Integer, [-1,23]:Integer}
llvalue= %3 = call {}*** @julia.get_pgcstack() #791
Stacktrace:
[1] getproperty
@ ./task.jl:189
[2] get_task_tls
@ ./task.jl:274
[3] task_local_storage
@ ./task.jl:269
[4] macro expansion
@ ~/.julia/packages/GPUArraysCore/GMsgk/src/GPUArraysCore.jl:208
[5] #_mapreduce#42
@ ~/.julia/packages/GPUArrays/qt4ax/src/host/mapreduce.jl:71
Stacktrace:
[1] get_task_tls
@ ./task.jl:271 [inlined]
[2] task_local_storage
@ ./task.jl:269 [inlined]
[3] macro expansion
@ ~/.julia/packages/GPUArraysCore/GMsgk/src/GPUArraysCore.jl:208 [inlined]
[4] #_mapreduce#42
@ ~/.julia/packages/GPUArrays/qt4ax/src/host/mapreduce.jl:71
[5] _mapreduce
@ ~/.julia/packages/GPUArrays/qt4ax/src/host/mapreduce.jl:33 [inlined]
[6] mapreduce
@ ~/.julia/packages/GPUArrays/qt4ax/src/host/mapreduce.jl:28 [inlined]
[7] _sum
@ ./reducedim.jl:1017 [inlined]
[8] sum
@ ./reducedim.jl:1013 [inlined]
[9] f
@ ./REPL[15]:1 [inlined]
[10] diffejulia_f_30099wrap
@ ./REPL[15]:0
[11] macro expansion
@ ~/.julia/packages/Enzyme/RTS5U/src/compiler.jl:8398 [inlined]
[12] enzyme_call
@ ~/.julia/packages/Enzyme/RTS5U/src/compiler.jl:7950 [inlined]
[13] CombinedAdjointThunk
@ ~/.julia/packages/Enzyme/RTS5U/src/compiler.jl:7723 [inlined]
[14] autodiff
@ ~/.julia/packages/Enzyme/RTS5U/src/Enzyme.jl:491 [inlined]
[15] autodiff(mode::ReverseMode{false, false, FFIABI, false, false}, f::typeof(f), ::Type{Active}, args::Duplicated{CuArray{Float32, 1, CUDA.DeviceMemory}})
@ Enzyme ~/.julia/packages/Enzyme/RTS5U/src/Enzyme.jl:512
[16] top-level scope
@ REPL[16]:1
```
and then trying with runtime activity gives
```julia
ERROR:
No augmented forward pass found for ijl_eqtable_pop
at context: %82 = call nonnull "enzyme_type"="{[-1]:Pointer}" {} addrspace(10)* @ijl_eqtable_pop({} addrspace(10)* noundef nonnull %getfield94, {} addrspace(10)* noundef addrspacecast ({}* inttoptr (i64 133898749444488 to {}*) to {} addrspace(10)*), {} addrspace(10)* noundef addrspacecast ({}* inttoptr (i64 133901345014144 to {}*) to {} addrspace(10)*), i64 noundef %bitcast_coercion) #791 [ "jl_roots"({} addrspace(10)* null) ], !dbg !1776
Stacktrace:
[1] pop!
@ ./iddict.jl:114
[2] delete!
@ ./iddict.jl:131
[3] macro expansion
@ ~/.julia/packages/GPUArraysCore/GMsgk/src/GPUArraysCore.jl:210
[4] #_mapreduce#42
@ ~/.julia/packages/GPUArrays/qt4ax/src/host/mapreduce.jl:71
Stacktrace:
[1] pop!
@ ./iddict.jl:114 [inlined]
[2] delete!
@ ./iddict.jl:131 [inlined]
[3] macro expansion
@ ~/.julia/packages/GPUArraysCore/GMsgk/src/GPUArraysCore.jl:210 [inlined]
[4] #_mapreduce#42
@ ~/.julia/packages/GPUArrays/qt4ax/src/host/mapreduce.jl:71
[5] _mapreduce
@ ~/.julia/packages/GPUArrays/qt4ax/src/host/mapreduce.jl:33 [inlined]
[6] mapreduce
@ ~/.julia/packages/GPUArrays/qt4ax/src/host/mapreduce.jl:28 [inlined]
[7] _sum
@ ./reducedim.jl:1017 [inlined]
[8] sum
@ ./reducedim.jl:1013 [inlined]
[9] g
@ ./REPL[20]:1 [inlined]
[10] diffejulia_g_58392wrap
@ ./REPL[20]:0
[11] macro expansion
@ ~/.julia/packages/Enzyme/RTS5U/src/compiler.jl:8398 [inlined]
[12] enzyme_call
@ ~/.julia/packages/Enzyme/RTS5U/src/compiler.jl:7950 [inlined]
[13] CombinedAdjointThunk
@ ~/.julia/packages/Enzyme/RTS5U/src/compiler.jl:7723 [inlined]
[14] autodiff
@ ~/.julia/packages/Enzyme/RTS5U/src/Enzyme.jl:491 [inlined]
[15] autodiff(mode::ReverseMode{false, true, FFIABI, false, false}, f::typeof(g), ::Type{Active}, args::Duplicated{CuArray{Float32, 1, CUDA.DeviceMemory}})
@ Enzyme ~/.julia/packages/Enzyme/RTS5U/src/Enzyme.jl:512
[16] top-level scope
@ REPL[22]:1
caused by:
No augmented forward pass found for cuMemcpyDtoHAsync_v2
at context: %12 = call i32 @cuMemcpyDtoHAsync_v2(i64 %1, i64 %2, i64 %7, i64 %unbox) #788, !dbg !1566
Stacktrace:
[1] macro expansion
@ ~/.julia/packages/CUDA/2kjXI/lib/utils/call.jl:214
[2] macro expansion
@ ~/.julia/packages/CUDA/2kjXI/lib/cudadrv/libcuda.jl:368
[3] #49
@ ~/.julia/packages/CUDA/2kjXI/lib/utils/call.jl:35
[4] check
@ ~/.julia/packages/CUDA/2kjXI/lib/cudadrv/libcuda.jl:35
[5] cuMemcpyDtoHAsync_v2
@ ~/.julia/packages/CUDA/2kjXI/lib/utils/call.jl:34
[6] #unsafe_copyto!#937
@ ~/.julia/packages/CUDA/2kjXI/lib/cudadrv/memory.jl:414
Stacktrace:
[1] macro expansion
@ ~/.julia/packages/CUDA/2kjXI/lib/utils/call.jl:214 [inlined]
[2] macro expansion
@ ~/.julia/packages/CUDA/2kjXI/lib/cudadrv/libcuda.jl:368 [inlined]
[3] #49
@ ~/.julia/packages/CUDA/2kjXI/lib/utils/call.jl:35 [inlined]
[4] check
@ ~/.julia/packages/CUDA/2kjXI/lib/cudadrv/libcuda.jl:35 [inlined]
[5] cuMemcpyDtoHAsync_v2
@ ~/.julia/packages/CUDA/2kjXI/lib/utils/call.jl:34 [inlined]
[6] #unsafe_copyto!#937
@ ~/.julia/packages/CUDA/2kjXI/lib/cudadrv/memory.jl:414
[7] unsafe_copyto!
@ ~/.julia/packages/CUDA/2kjXI/lib/cudadrv/memory.jl:411 [inlined]
[8] #1127
@ ~/.julia/packages/CUDA/2kjXI/src/array.jl:558
[9] #context!#990
@ ~/.julia/packages/CUDA/2kjXI/lib/cudadrv/state.jl:168 [inlined]
[10] context!
@ ~/.julia/packages/CUDA/2kjXI/lib/cudadrv/state.jl:163 [inlined]
[11] unsafe_copyto!
@ ~/.julia/packages/CUDA/2kjXI/src/array.jl:550
[12] copyto!
@ ~/.julia/packages/CUDA/2kjXI/src/array.jl:503 [inlined]
[13] getindex
@ ~/.julia/packages/GPUArrays/qt4ax/src/host/indexing.jl:52 [inlined]
[14] scalar_getindex
@ ~/.julia/packages/GPUArrays/qt4ax/src/host/indexing.jl:36 [inlined]
[15] _getindex
@ ~/.julia/packages/GPUArrays/qt4ax/src/host/indexing.jl:19 [inlined]
[16] getindex
@ ~/.julia/packages/GPUArrays/qt4ax/src/host/indexing.jl:17 [inlined]
[17] macro expansion
@ ~/.julia/packages/GPUArraysCore/GMsgk/src/GPUArraysCore.jl:210 [inlined]
[18] #_mapreduce#42
@ ~/.julia/packages/GPUArrays/qt4ax/src/host/mapreduce.jl:71
[19] _mapreduce
@ ~/.julia/packages/GPUArrays/qt4ax/src/host/mapreduce.jl:33 [inlined]
[20] mapreduce
@ ~/.julia/packages/GPUArrays/qt4ax/src/host/mapreduce.jl:28 [inlined]
[21] _sum
@ ./reducedim.jl:1017 [inlined]
[22] sum
@ ./reducedim.jl:1013 [inlined]
[23] g
@ ./REPL[20]:1 [inlined]
[24] diffejulia_g_58392wrap
@ ./REPL[20]:0
[25] macro expansion
@ ~/.julia/packages/Enzyme/RTS5U/src/compiler.jl:8398 [inlined]
[26] enzyme_call
@ ~/.julia/packages/Enzyme/RTS5U/src/compiler.jl:7950 [inlined]
[27] CombinedAdjointThunk
@ ~/.julia/packages/Enzyme/RTS5U/src/compiler.jl:7723 [inlined]
[28] autodiff
@ ~/.julia/packages/Enzyme/RTS5U/src/Enzyme.jl:491 [inlined]
[29] autodiff(mode::ReverseMode{false, true, FFIABI, false, false}, f::typeof(g), ::Type{Active}, args::Duplicated{CuArray{Float32, 1, CUDA.DeviceMemory}})
@ Enzyme ~/.julia/packages/Enzyme/RTS5U/src/Enzyme.jl:512
[30] top-level scope
@ REPL[22]:1
```
Contributor guide
Assessment
This issue has not been assessed yet.