Reverse-mode over a CuArray aborts (`gc-transition` on `cuMemcpyDtoHAsync_v2`)

[rsenne]

On ParallelMCMC I have been trying to use true Second Order AD for hvps. I've reduced my issues via two MWREs and it seems like Reverse Mode AD on the GPU isn't working properly.

MWE 1

First i tried the native HVP.

using Enzyme, CUDA
@assert CUDA.functional()

logp(β) = -0.5f0 * sum(abs2, β)      

β = CUDA.CuArray(Float32.(1:8) ./ 8)
v = CUDA.CuArray(ones(Float32, 8))
Enzyme.hvp(logp, β, v)                 

This aborts.

MWE 2

So then, I was curious if maybe the reverse part of the HVP calculation was broken so i tried this:

using Enzyme, CUDA
@assert CUDA.functional()

logp(β) = -0.5f0 * sum(abs2, β) 

β = CUDA.CuArray(Float32.(1:8) ./ 8)
Enzyme.gradient(Enzyme.Reverse, logp, β)

This also aborts and interestingly in the same way as example 1.

Error

unsupported tag gc-transition for   %13 = call i32 @cuMemcpyDtoHAsync_v2(
  ptr %"dst::Ptr", i64 %"src::CuPtr", i64 4, ptr %.unbox)
  [ "jl_roots"(...), "gc-transition"() ]
UNREACHABLE executed at .../Enzyme/GradientUtils.cpp:309!

signal (6): Aborted
getInvertedBundles            at .../Enzyme/GradientUtils.cpp:309
recursivelyHandleSubfunction  at .../Enzyme/AdjointGenerator.h:5636
...

So seems related to my previous issue #3195 (gc-transition tag) and also #2139 (cuMemcpyDtoHAsync_v2).

Version Info

• Julia 1.12.6
• Enzyme.jl 0.13.153, EnzymeCore 0.8.20
• CUDA.jl 5.11.2 (CUDA runtime 13.0)
• GPUCompiler 1.16.1, LLVM.jl 9.8.2
• NVIDIA L40S