EnzymeAD / EnzymeAD/Enzyme.jl

Reverse-mode failures differentiating through `reshape` of a CuArray (RuntimeActivity / NoDerivative / retain-freed-data)

Open
#3,370 1 comment 0 reactions 0 assignees View on GitHub
Dominant language
Julia
Stars
587
Forks
109
Avg merge
1d 5h
Merged PRs (30d)
44

Description

# `EnzymeRuntimeActivityError` when differentiating through `reshape` of a `CuArray`

Reshaping a 3-D `CuArray` (as e.g. Flux's `Dense` does to a batched input) inside an
Enzyme-differentiated function throws `EnzymeRuntimeActivityError` when the array is
`Const` (inactive) but the reshape result feeds into an active computation. The error
originates in the GPUArrays `reshape` → `DataRef` → `copy` → `derive` refcount path, not
in any user code. The identical function with a plain CPU `Array` differentiates fine.

Originally hit through Flux (https://github.com/FluxML/Flux.jl/issues/2448); reduced
below to a Flux-free / NNlib-free reproducer.

## MWE

```julia
using CUDA, Enzyme

f(a, x) = a * sum(reshape(x, size(x, 1), :))

x = CUDA.rand(Float32, 64, 100, 512) # inactive constant input

# Fails: EnzymeRuntimeActivityError at reshape(::CuArray{Float32,3}, ::Tuple{Int64,Colon})
Enzyme.autodiff(Reverse, f, Active, Active(2.0f0), Const(x))

# The same code with a plain CPU Array works fine:
xcpu = rand(Float32, 64, 100, 512)
Enzyme.autodiff(Reverse, f, Active, Active(2.0f0), Const(xcpu)) # OK
```

## Stacktrace

```
EnzymeRuntimeActivityError: Detected potential need for runtime activity.

Constant memory is stored (or returned) to a differentiable variable and correctness cannot be guaranteed with static activity analysis.
This might be due to the use of a constant variable as temporary storage for active memory (https://enzyme.mit.edu/julia/stable/faq/#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 performance, slower to setup), 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.

Failure within method:

reshape(::CuArray{Float32, 3, CUDACore.DeviceMemory}, ::Tuple{Int64, Colon})
@ Base reshapedarray.jl:133

Hint: catch this exception as `err` and call `code_typed(err)` to inspect the surrounding code.

Mismatched activity for: store atomic ptr addrspace(10) %"parent::CuArray.data.rc14", ptr addrspace(11) %25 release, align 8, !dbg !1022, !tbaa !994, !alias.scope !988, !noalias !989 const val: %"parent::CuArray.data.rc14" = load atomic ptr addrspace(10), ptr addrspace(11) %18 unordered, align 8, !dbg !1020, !tbaa !994, !alias.scope !988, !noalias !996, !nonnull !0, !dereferenceable !1018, !align !998, !enzyme_type !1019, !enzymejl_byref_MUT_REF !0, !enzymejl_source_type_GPUArrays.RefCounted\7BCUDACore.Managed\7BCUDACore.DeviceMemory\7D\7D !0
Julia value causing error: Unknown object of type GPUArrays.RefCounted{CUDACore.Managed{CUDACore.DeviceMemory}}
LLVM view of erring value: %"parent::CuArray.data.rc14" = load atomic ptr addrspace(10), ptr addrspace(11) %18 unordered, align 8, !dbg !1020, !tbaa !994, !alias.scope !988, !noalias !996, !nonnull !0, !dereferenceable !1018, !align !998, !enzyme_type !1019, !enzymejl_byref_MUT_REF !0, !enzymejl_source_type_GPUArrays.RefCounted\7BCUDACore.Managed\7BCUDACore.DeviceMemory\7D\7D !0

Stacktrace:
[1] DataRef
@ ~/.julia/packages/GPUArrays/VNJ6B/src/host/abstractarray.jl:58
[2] copy
@ ~/.julia/packages/GPUArrays/VNJ6B/src/host/abstractarray.jl:85
[3] derive
@ ~/.julia/packages/CUDACore/NlVPI/src/array.jl:858
[4] reshape
@ ~/.julia/packages/GPUArrays/VNJ6B/src/host/base.jl:195
[5] reshape
@ ./reshapedarray.jl:133
```

## Additional findings

The reshape of a `CuArray` seems to be broken for Enzyme in several distinct ways
depending on activity annotations. The reduced function used below is

```julia
g(W, x) = sum(W * reshape(x, size(x, 1), :))
W = CUDA.rand(Float32, 64, 64); x = CUDA.rand(Float32, 64, 100, 512)
```

### 1. `Const` input → `EnzymeRuntimeActivityError` (the primary MWE above)

- The trigger is specifically a **`Const` (inactive) array reshaped into an active
computation**. It reproduces with no matmul and no `Duplicated` array at all — an
`Active` scalar times `sum` of the reshaped `Const` array is enough (see the MWE).
- The same code on a plain CPU `Array` differentiates fine, so this is specific to the
GPU `reshape` → `DataRef`/refcount path.

### 2. `set_runtime_activity(Reverse)` → `EnzymeNoDerivativeError`

The documented workaround (b) gets past the activity error but then hits a *second* bug —
no derivative rule for the host→device copy:

```julia
Enzyme.autodiff(set_runtime_activity(Reverse), g, Active, Duplicated(W, dW), Const(x))
```

```
EnzymeNoDerivativeError:
No augmented forward pass found for ejlstr$cuMemcpyHtoDAsync_v2$.../libcuda.so
at context: %13 = call i32 @"ejlstr$cuMemcpyHtoDAsync_v2$.../libcuda.so"(i64 %"dst::CuPtr", ptr %"src::Ptr", i64 4, ptr %.unbox) #499 [ "jl_roots"(ptr addrspace(10) %value_phi), "gc-transition"() ], !dbg !504

Stacktrace:
[1] macro expansion
@ ~/.julia/packages/CUDACore/NlVPI/lib/cudadrv/libcuda.jl:530
[2] #cuMemcpyHtoDAsync_v2##0
@ ~/.julia/packages/GPUToolbox/9Elch/src/ccalls.jl:34
[3] check
@ ~/.julia/packages/CUDACore/NlVPI/lib/cudadrv/libcuda.jl:35
[4] cuMemcpyHtoDAsync_v2
@ ~/.julia/packages/GPUToolbox/9Elch/src/ccalls.jl:33
[5] #unsafe_copyto!#537
@ ~/.julia/packages/CUDACore/NlVPI/lib/cudadrv/memory.jl:426
[6] unsafe_copyto!
@ ~/.julia/packages/CUDACore/NlVPI/lib/cudadrv/memory.jl:421

Failure within method:

unsafe_copyto!(::CuPtr{Float32}, ::Ptr{Float32}, ::Int64; async::Bool)
@ CUDACore ~/.julia/packages/CUDACore/NlVPI/lib/cudadrv/memory.jl:421
```

### 3. `Duplicated` input → `ArgumentError: Attempt to retain freed data.`

Making the reshaped array `Duplicated` (instead of `Const`) fails with yet another, third
error:

```julia
Enzyme.autodiff(Reverse, g, Active, Duplicated(W, dW), Duplicated(x, dx))
```

```
ArgumentError: Attempt to retain freed data.
```

## Version info

```
Julia 1.12.6, LLVM 18.1.7
Enzyme 0.13.190, EnzymeCore 0.8.21
CUDA 6.2.1, CUDACore 6.2.1, GPUArrays 11.5.8, GPUCompiler 1.23.0
CUDA runtime 13.3.0, driver 595.71.5
NVIDIA GeForce RTX 5090 (sm_120)
```

Contributor guide

Open the contributing guide

Research direction

Run the reduced Enzyme.autodiff examples first, comparing the CuArray and CPU Array cases and the Const, runtime-activity, and Duplicated annotations. Then inspect the reported reshape path through GPUArrays' host/abstractarray.jl and host/base.jl, plus CUDACore's array.jl and memory.jl frames. Done means the reproduced CuArray cases no longer raise the reported activity, derivative, or retain-freed-data errors.

Written by the indexing model from the issue text.

Assessment

Tech stack
julia
Domain
compilers
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
42/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.