Add CUDA specific copy rules by rsenne · Pull Request #3363 · EnzymeAD/Enzyme.jl

rsenne · 2026-07-19T02:56:42Z

@wsmoses added these copy rules per your suggestion and my understanding. Here is what it now allows that prior would break:

Array(x) / collect(x) pulling a CuArray back to host mid-function. before this it'd just die with "no augmented/forward pass found for cuMemcpyDtoHAsync_v2". e.g., both fwd and rev sum(abs2, Array(x)) gives back 2x like you'd want.
CuArray(y) / copyto!(gpu, host)
copy(x::CuArray)
round-trips: Array(CuArray(y)) and back

however, this did not fix #3235 and need further fixes for sum(abs2, CuArray(x))

EnzymeNoDerivativeError:
No create nofree of empty function (ijl_eqtable_put) ijl_eqtable_put)
 at context:   call fastcc void @julia_task_local_storage_86635(i32 zeroext 0) ... (julia_task_local_storage_86635)

Stacktrace:
 [1] macro expansion
    @ GPUArraysCore/src/GPUArraysCore.jl:205   
 [2] #_mapreduce#48
    @ GPUArrays/src/host/mapreduce.jl:88

dot also fails

# reverse
EnzymeNoDerivativeError:
No augmented forward pass found for ejlstr$cuMemcpyDtoHAsync_v2$nvcuda
 at context:   %17 = call i32 @"ejlstr$cuMemcpyDtoHAsync_v2$nvcuda"(
     ptr %"dst::Ptr", i64 %"src::CuPtr", i64 4, ptr %.unbox)
     [ "jl_roots"(...), "gc-transition"() ]

# forward
EnzymeNoDerivativeError:
No forward mode derivative found for ejlstr$cuMemcpyDtoHAsync_v2$nvcuda
 at context:   %17 = call i32 @"ejlstr$cuMemcpyDtoHAsync_v2$nvcuda"( ... i64 4 ... )

I figured i wouldn't add anything beyond these copy rules unless you think its in scope for this PR otherwise i can continue to target this

wsmoses · 2026-07-19T03:30:35Z

the dot case implies that this is defined on the wrong method. For now don't worry about gpuarrays.

Define rules for https://github.com/JuliaGPU/CUDA.jl/blob/cac0ac97adcd43c23abf0b3c8bb3f99364d9b590/CUDACore/lib/cudadrv/memory.jl#L421

and

https://github.com/JuliaGPU/CUDA.jl/blob/cac0ac97adcd43c23abf0b3c8bb3f99364d9b590/CUDACore/lib/cudadrv/memory.jl#L432

and

https://github.com/JuliaGPU/CUDA.jl/blob/cac0ac97adcd43c23abf0b3c8bb3f99364d9b590/CUDACore/lib/cudadrv/memory.jl#L450

and similar unsafe_copyto down

rsenne · 2026-07-19T23:09:26Z

Define rules

So IIUC the current methods I defined rules for are a level too high. The current extension I wrote these in does not have CUDA as a dep --- should I add it as a weakdep, or move this PR to CUDA.jl and add it in the EnzymeCoreExt?

Also, assuming the pointer-level rules subsume it, should I drop the array-level rule I have here or keep it as a fallback?

wsmoses · 2026-07-19T23:16:58Z

Yeah basically.

Re where they live either is a reasonable choice -- mostly one of whose methods/internals seem more used. In the case of the current cuda kernel launch rules those require a lot of the cuda internals so make sense to live there.

These don't really have any internals so could live anywhere imo.

A relevant fyi, current cuda.jl main doesn't pkg resolve with enzyme atm due to the GPUCompiler 2 bump, though also ci there isn't running regardless and @kshyatt was looking into that .

Maybe makes sense to have here then. But I'm fine either way

codecov · 2026-07-21T15:11:16Z

Codecov Report

❌ Patch coverage is 0% with 48 lines in your changes missing coverage. Please review.
✅ Project coverage is 76.81%. Comparing base (6bc8421) to head (9890a93).
⚠️ Report is 10 commits behind head on main.

Files with missing lines	Patch %	Lines
ext/EnzymeGPUArraysCoreExt.jl	0.00%	48 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #3363      +/-   ##
==========================================
+ Coverage   75.04%   76.81%   +1.76%     
==========================================
  Files          66       66              
  Lines       22852    22326     -526     
==========================================
- Hits        17150    17149       -1     
+ Misses       5702     5177     -525

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

rsenne · 2026-07-22T00:32:16Z

@wsmoses i think this is ready

wsmoses · 2026-07-22T00:39:51Z

+    return EnzymeRules.width(config) == 1 ? x.dval : x.dval[batch]
+end
+
+function EnzymeRules.augmented_primal(


we shouldn't need this?

I think this is needed---I tried removing it and all the copy-rule gradients come back zero. Without it the reverse unsafe_copyto! rules accumulate into an unconnected pointer because unlike say pointer(::Array), pointer(::CuArray) isn't an intrinsic Enzyme shadows (at least that is how i understand it)

Er that is a bug then, can you upload an isolated as mwe possible case with enzyme.api.printall!(true) on

kshyatt · 2026-07-22T07:08:38Z

+using LinearAlgebra: dot
 using Test

+@testset "CUDA memory copies" begin


Would be good to check this works with all the CUDA memory types (Device, Host, Unified)

Added device / unified / host coverage. Device passes; unified and host are @test_broken for now. The blocker is the pointer(::StridedCuArray) rule: for unified/host memory pointer(::CuArray). This is related to the other issue i am trying to figure out so leaving it as a flag for now

vchuravy · 2026-07-22T12:30:56Z

Same comment as #3375 (comment) the CUDA API ones should live in EnzymeAD/Enzyme?

rsenne · 2026-07-22T21:00:38Z

@wsmoses here is the MWE.

using CUDA, Enzyme

f(x) = sum(abs2, Array(CuArray(x)))

x  = Float32[1, 2, 3]
dx = zeros(Float32, 3)

get(ENV, "PRINTALL", "0") == "1" && Enzyme.API.printall!(true)
Enzyme.autodiff(Reverse, f, Active, Duplicated(x, dx))
Enzyme.API.printall!(false)

println("dx       = ", dx)
println("expected = ", 2 .* x)
println("correct? = ", dx == 2 .* x)

For me, with x = Float32[1,2,3], keeping the pointer(::StridedCuArray) rule gives dx = [2,4,6] and removing it gives dx = [3,6,9].

wsmoses · 2026-07-22T21:01:19Z

+function _zero!(ptr::CuArrayPtr{T}, off::Integer, n::Integer) where {T}
+    n == 0 && return nothing
+    buffer = _stage(ptr, off, n)
+    fill!(buffer, zero(eltype(buffer)))


this can call the same zero as CuPtr{T}, right?

i do not think so as memset is defined only for CuPtr{T}

i would assume a convert or something exists given https://github.com/JuliaGPU/CUDA.jl/blob/8fc49feda7cf563fa9fb4900e329dfecc99ed0fd/CUDACore/src/array.jl#L848

hmm i tried the following:

using CUDA mem = CUDA.alloc(CUDA.ArrayMemory{Float32}, (4,)) convert(CuPtr{Float32}, mem) # ERROR: ArgumentError: Illegal conversion of a ArrayMemory{Float32, 1} to a CuPtr{Float32} Base.unsafe_convert(CuPtr{Float32}, mem) # ERROR: ArgumentError: Illegal conversion of a ArrayMemory{Float32, 1} to a CuPtr{Float32} CUDA.memset(reinterpret(CuPtr{UInt8}, pointer(mem)), UInt8(0), 16) # ERROR: CUDA error: invalid argument (code 1, ERROR_INVALID_VALUE)

If i understand the code snippet linked, pointer(A) is already a CuPtr, because A is a DenseCuArray. So convert(CuPtr{U}, pointer(A)) is a CuPtr -> CuPtr retype

this just feels like a missing memset impl in cuda.jl [and we should open an issue for]?

### Array memory Array memory is a special kind of memory that is optimized for 2D and 3D access patterns. The memory is opaquely managed by the CUDA runtime, and is typically only used on combination with texture intrinsics. ``@docs CUDA.ArrayMemory CUDA.alloc(::Type{CUDA.ArrayMemory{T}}, ::Dims) where T `` ### Pointers To work with these buffers, you need to `convert` them to a `Ptr`, `CuPtr`, or in the case of `ArrayMemory` an `CuArrayPtr`. You can then use common Julia methods on these pointers, such as `unsafe_copyto!`. CUDA.jl also provides some specialized functionality that does not match standard Julia functionality: ``@docs CUDA.unsafe_copy2d! CUDA.unsafe_copy3d! CUDA.memset ``

In any case let's make our lives easier then and just leave the method unimplemented for CuArrayPtr for now

wsmoses · 2026-07-22T21:02:08Z

+end
+
+@inline _commit!(::Union{Ptr, CuPtr}, offset, buffer, n) = nothing
+function _commit!(ptr::CuArrayPtr, offset, buffer, n)


what is this whole commit thing and why is it needed?

wsmoses · 2026-07-22T21:02:38Z

+    n == 0 && return nothing
+    typeof(src) === typeof(dest) && src == dest && soff == doff && return nothing
+
+    dsrc = _stage(src, soff, n)


get rid of the stage/commit stuff

wsmoses · 2026-07-22T21:03:44Z

+    axpy!(one(eltype(dst)), src, dst)
+    return nothing
+end
+_accumulate!(dst::Array, src::Array) = (dst .+= src; nothing)


we should benchmark if we should still use axpy here, I would hope it to be better than broadcasting?

eltype n broadcast(us) axpy(us) speedup(bcast/axpy) Float32 1 7.896 13.734 0.57 Float32 16 8.115 13.123 0.62 Float32 256 8.668 13.615 0.64 Float32 4096 8.096 13.318 0.61 Float32 65536 8.483 13.476 0.63 Float32 1048576 8.266 13.597 0.61 Float32 16777216 96.727 95.242 1.02 Float64 1 9.326 13.458 0.69 Float64 16 9.092 13.423 0.68 Float64 256 9.014 14.109 0.64 Float64 4096 8.579 13.794 0.62 Float64 65536 8.256 14.423 0.57 Float64 1048576 9.157 13.786 0.66 Float64 16777216 256.104 257.62 0.99

apparently not ?

s a v a g e.

but sure, add a comment that we could do axpy [called like X], but it is faster [copying this table] to use a bcast

wsmoses · 2026-07-22T21:03:57Z

+    dsrc = _stage(src, soff, n)
+    ddest = _stage(dest, doff, n)
+    if dsrc isa CuArray && ddest isa CuArray
+        _accumulate!(dsrc, ddest)


just do this inside overrides of _accumulate!

wsmoses · 2026-07-22T21:04:05Z

+
+function _accumulate_and_zero!(src, soff, dest, doff, n::Integer)
+    n == 0 && return nothing
+    typeof(src) === typeof(dest) && src == dest && soff == doff && return nothing


why would this ever be true?

Add copy rules

9890a93

wsmoses reviewed Jul 19, 2026

View reviewed changes

Comment thread ext/EnzymeGPUArraysCoreExt.jl Outdated

cossio mentioned this pull request Jul 19, 2026

Add easy_rules for exponentially scaled Bessel functions (besselix etc.) #3365

Merged

implement cuda rules

0d50873

rsenne changed the title ~~Add copy rules~~ Add CUDA specific copy rules Jul 22, 2026

rsenne added 2 commits July 21, 2026 20:31

Add more tests plus format

8a95a85

Merge branch 'main' into fix/copy-rule

fbcbc0f

rsenne requested a review from wsmoses July 22, 2026 00:32

wsmoses reviewed Jul 22, 2026

View reviewed changes

Comment thread Project.toml Outdated

wsmoses reviewed Jul 22, 2026

View reviewed changes

Comment thread ext/EnzymeCUDAExt.jl Outdated

wsmoses reviewed Jul 22, 2026

View reviewed changes

Comment thread ext/EnzymeCUDAExt.jl Outdated

wsmoses reviewed Jul 22, 2026

View reviewed changes

Comment thread ext/EnzymeCUDAExt.jl Outdated

wsmoses reviewed Jul 22, 2026

View reviewed changes

pr review

ccefc10

rsenne requested a review from wsmoses July 22, 2026 03:04

kshyatt reviewed Jul 22, 2026

View reviewed changes

Comment thread ext/EnzymeCUDAExt.jl Outdated

kshyatt reviewed Jul 22, 2026

View reviewed changes

Pr comments

2591948

wsmoses reviewed Jul 22, 2026

View reviewed changes

Uh oh!

Conversation

rsenne commented Jul 19, 2026

Uh oh!

Uh oh!

wsmoses commented Jul 19, 2026

Uh oh!

rsenne commented Jul 19, 2026

Uh oh!

wsmoses commented Jul 19, 2026

Uh oh!

codecov Bot commented Jul 21, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

rsenne commented Jul 22, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

vchuravy commented Jul 22, 2026

Uh oh!

rsenne commented Jul 22, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

codecov Bot commented Jul 21, 2026 •

edited

Loading