Differentiate the CUDA memcpy family in forward mode - #3102
Open
wsmoses wants to merge 1 commit into
Open
Conversation
wsmoses
force-pushed
the
cuda-forward-mode
branch
2 times, most recently
from
August 8, 2026 04:51
8afba05 to
26fda3a
Compare
wsmoses
force-pushed
the
cuda-forward-mode
branch
from
August 9, 2026 01:22
26fda3a to
ee875a7
Compare
This was referenced Aug 9, 2026
wsmoses
force-pushed
the
cuda-forward-mode
branch
from
August 9, 2026 03:05
ee875a7 to
db3a88d
Compare
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
force-pushed
the
cuda-forward-mode
branch
from
August 9, 2026 04:29
db3a88d to
c34a67d
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Neither the driver API (
cuMemcpy{HtoD,DtoH,DtoD}[Async][_v2], and the unifiedcuMemcpy[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:
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 atleast one side generally lives in memory the host cannot address. So rather
than a parallel implementation,
visitMemTransferCommongrows an optionalemitter 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
CUdeviceptris a pointer despite being an integer, and to stop markingtrailing 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$libnaming afrontend may use — plus
Integration/ForwardMode/cudamemcpy.cpp, which checksthe tangent numerically against a mocked CUDA runtime and confirms both the
primal and the shadow are transferred each way. Passes
-O0through-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 unmodifiedmainwith this local LLVM 16 build. Format (clang-format 16) andcheck_emission_order.pyclean.Relationship to the other PRs
Split up at review request; this one is now just the transfer derivatives.
getOrInsertPerCallingConvand the pre-existing call sitesthat needed it.
because they free the device buffers they allocate.
_v2forward-mode scalar result. Independent of this;can land in either order.
Remaining gaps in EnzymeAD/Enzyme.jl#3442 after all four:
cublasGemmExis notrecognized by
extractBLAS; a forward-mode shadow placeholder for acmpxchgresult is never replaced, tripping the verifier inside CUDA.jl's memory pool;
and the shadow chain for a
CuArraybreaks atManaged{DeviceMemory}, which istype-inactive because
DeviceMemoryis an untyped byte buffer.🤖 Generated with Claude Code
https://claude.ai/code/session_01R6a8BiaAKpUTgP86mQ9ZKP