Skip to content

Differentiate the CUDA memcpy family in forward mode - #3102

Open
wsmoses wants to merge 1 commit into
mainfrom
cuda-forward-mode
Open

Differentiate the CUDA memcpy family in forward mode#3102
wsmoses wants to merge 1 commit into
mainfrom
cuda-forward-mode

Conversation

@wsmoses

@wsmoses wsmoses commented Aug 8, 2026

Copy link
Copy Markdown
Member

Neither the driver API (cuMemcpy{HtoD,DtoH,DtoD}[Async][_v2], and the unified
cuMemcpy[Async]) nor the runtime API (cudaMemcpy[Async]) had a derivative,
so staging data through device memory failed outright. This is the error
reported in EnzymeAD/Enzyme.jl#3442:

No forward mode derivative found for cuMemcpyDtoHAsync_v2

A CUDA transfer behaves exactly like a memcpy, except that the shadow copy has
to go back through the CUDA API rather than becoming an llvm.memcpy, since at
least one side generally lives in memory the host cannot address. So rather
than a parallel implementation, visitMemTransferCommon grows an optional
emitter hook: the CUDA path reuses its type segmentation, activity handling and
zeroing of inactive float sources, and only the emission of the copy differs.

Type analysis learns to propagate the pointee type between the two sides, that
a CUdeviceptr is a pointer despite being an integer, and to stop marking
trailing arguments as integers before the stream pointer a CUDA transfer
carries where a memcpy carries only volatility.

Reverse mode still reports a missing derivative rather than answering
incorrectly; accumulating into device memory needs a copy that runs on the
device, which cannot be emitted from the host.

Testing

Three lit tests — driver API, runtime API, and the ejlstr$fn$lib naming a
frontend may use — plus Integration/ForwardMode/cudamemcpy.cpp, which checks
the tangent numerically against a mocked CUDA runtime and confirms both the
primal and the shadow are transferred each way. Passes -O0 through -O3.

Full unit suite on LLVM 16: 1193 pass, 11 xfail, 1 fail — that failure is
ReverseModeVector/partial_int_window.ll, which fails identically on unmodified
main with this local LLVM 16 build. Format (clang-format 16) and
check_emission_order.py clean.

Relationship to the other PRs

Split up at review request; this one is now just the transfer derivatives.

Remaining gaps in EnzymeAD/Enzyme.jl#3442 after all four: cublasGemmEx is not
recognized by extractBLAS; a forward-mode shadow placeholder for a cmpxchg
result is never replaced, tripping the verifier inside CUDA.jl's memory pool;
and the shadow chain for a CuArray breaks at Managed{DeviceMemory}, which is
type-inactive because DeviceMemory is an untyped byte buffer.

🤖 Generated with Claude Code

https://claude.ai/code/session_01R6a8BiaAKpUTgP86mQ9ZKP

@wsmoses
wsmoses force-pushed the cuda-forward-mode branch 2 times, most recently from 8afba05 to 26fda3a Compare August 8, 2026 04:51
@wsmoses
wsmoses force-pushed the cuda-forward-mode branch from 26fda3a to ee875a7 Compare August 9, 2026 01:22
@wsmoses
wsmoses changed the base branch from main to per-calling-conv August 9, 2026 01:22
Base automatically changed from per-calling-conv to main August 9, 2026 02:27
@wsmoses wsmoses changed the title Forward-mode CUDA: memcpy family derivatives and cuBLAS scalar results Differentiate the CUDA memcpy family in forward mode Aug 9, 2026
@wsmoses
wsmoses changed the base branch from main to cuda-dealloc August 9, 2026 03:05
@wsmoses
wsmoses force-pushed the cuda-forward-mode branch from ee875a7 to db3a88d Compare August 9, 2026 03:05
Base automatically changed from cuda-dealloc to main August 9, 2026 04:21
Neither the driver API (cuMemcpy{HtoD,DtoH,DtoD}[Async][_v2], and the
unified cuMemcpy[Async]) nor the runtime API (cudaMemcpy[Async]) had a
derivative, so staging data through device memory failed outright with
"No forward mode derivative found". A CUDA transfer behaves exactly like
a memcpy, except that the shadow copy has to go back through the CUDA API
rather than becoming an llvm.memcpy, since at least one side generally
lives in memory the host cannot address.

visitMemTransferCommon grows an optional emitter hook so the CUDA path
reuses its type segmentation, activity handling and zeroing of inactive
float sources, and only the emission of the copy itself differs. Type
analysis learns to propagate the pointee type between the two sides, that
a CUdeviceptr is a pointer despite being an integer, and to stop marking
trailing arguments as integers before the stream pointer that a CUDA
transfer carries where a memcpy carries only volatility.

Two things this uncovered:

  - cuMemFree{,_v2,Async} and cudaFree{,Async,Host} were not recognized as
    deallocations, so freeing a device allocation inside a differentiated
    function reported a missing derivative immediately after the transfer
    was handled. Recognizing them also required guarding the nocapture
    attribute in getOrInsertCheckedFree, which is invalid on the integer a
    CUdeviceptr is passed as.

  - A frontend may reach a library through a renamed declaration -- Julia
    names its lazily bound ccalls "ejlstr$<function>$<library>" and loads
    those libraries RTLD_LOCAL -- so a plainly named helper declared next
    to such a call is not reachable via dlsym and fails when the module is
    JIT linked. getOrInsertPerCallingConv declares a helper under the
    convention of the call it accompanies, which also fixes the
    pre-existing memset emitted for a shadow allocation. That memset now
    matches the ABI of the allocation it pairs with as well: cuMemsetD8
    takes an unsigned int length where cuMemsetD8_v2 takes a size_t.

Reverse mode still reports a missing derivative rather than answering
incorrectly; accumulating into device memory needs a copy that runs on the
device, which cannot be emitted from the host.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01R6a8BiaAKpUTgP86mQ9ZKP
@wsmoses
wsmoses force-pushed the cuda-forward-mode branch from db3a88d to c34a67d Compare August 9, 2026 04:29
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant