Forward-mode over a CuArray aborts on `-β .- β` (gc-transition bundle on `cuPointerGetAttribute`)

[rsenne]

@wsmoses asked me to open this on a PR thread.

I've been doing GPU Hessian-vector products as a forward-mode pushforward of an
analytic gradient, and I kept hitting a hard abort during compilation. After
bisecting it down I've got it to a one-liner that doesn't involve any of my own
code, so here it is.

Where I actually hit it

In real code this showed up in a log-density gradient that ends in a fused
broadcast over a matmul result, roughly:

gradlogp(β) = -β .- Xᵀ * (X * β) ./ Float32(N)   # aborts on GPU

and it went away when I split that last line into single broadcast ops:

function gradlogp(β)        # works on GPU
    Y = Xᵀ * (X * β)
    Y = Y ./ Float32(N)
    Y = Y .+ β
    return -Y
end

So from the outside it looked like "Enzyme can't fuse multiple ops into one GPU
broadcast kernel," but after working on reducing this to a MWRE, I think I overstated.

Reduced

Forward-mode over a CuVector aborts when the active (Duplicated) input shows
up both negated and direct in the same broadcast. The smallest version I found:

using Enzyme, CUDA, Random
@assert CUDA.functional()

β = CUDA.CuArray(randn(MersenneTwister(1), Float32, 8))
v = CUDA.CuArray(ones(Float32, 8))   # forward tangent

Enzyme.autodiff(Enzyme.Forward, b ->  b .- b, Enzyme.Duplicated(β, v))  # fine
Enzyme.autodiff(Enzyme.Forward, b -> -b .- b, Enzyme.Duplicated(β, v))  # aborts

What threw me at first is that it's pretty specific. It's not broadcasts are
broken and it's not you used the input twice, both of those work on their
own. It really seems to be the negation combined with the second use:

• -β alone -> fine
• β .- β, β .+ β -> fine
• -β .- β, -β .+ β -> abort

set_runtime_activity doesn't make a difference either way.

The abort itself is Enzyme's getInvertedBundles giving up on a CUDA ccall that
carries a "gc-transition"() operand bundle:

unsupported tag gc-transition for   %1 = call i32 @cuPointerGetAttribute(
  ptr nonnull %"new::RefValue", i32 9, i64 %"dst::CuPtr")
  [ "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:5166
visitCallInst                 at .../Enzyme/AdjointGenerator.h:6744
CreateForwardDiff             at .../Enzyme/EnzymeLogic.cpp:5093
...
enzyme!                       at .../Enzyme/src/compiler.jl:2697

I also see it through DifferentiationInterface.pushforward with
AutoEnzyme(; mode=Enzyme.Forward, function_annotation=Enzyme.Const), which is
how I actually hit it. The bare autodiff above is just the reduced version.

Versions:

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

Also looks like #2139 might be related somehow?

[wsmoses]

@vchuravy @gbaraldi what the heck is this gc-transition tag?

[wsmoses]

cc @oscardssmith @maleadt too

[rsenne]

As a note--the whole motivation behind #3148 and its related PR #3122 is also related to this gc-transition thing as I got similar abort messages there. LMK if you want me to put in any other MWE for reverse mode etc.

[vchuravy]

what the heck is this gc-transition tag?

JuliaLang/julia@85458a0