Differentiate the CUDA memcpy family in forward mode
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
Enzyme is a plugin that performs automatic differentiation (AD) of statically analyzable LLVM and MLIR.
Enzyme can be used by calling __enzyme_autodiff on a function to be differentiated as shown below. Running the Enzyme transformation pass then replaces the call to __enzyme_autodiff with the gradient of its first argument.
double foo(double); double grad_foo(double x) { return __enzyme_autodiff(foo, x); }
Enzyme is highly-efficient and its ability to perform AD on optimized code allows Enzyme to meet or exceed the performance of state-of-the-art AD tools.
Detailed information on installing and using Enzyme can be found on our website: https://enzyme.mit.edu.
A short example of how to install Enzyme is below:
cd /path/to/Enzyme/enzyme mkdir build && cd build cmake -G Ninja .. -DLLVM_DIR=/path/to/llvm/lib/cmake/llvm -DLLVM_EXTERNAL_LIT=/path/to/lit/lit.py ninja
Or, install Enzyme using a package manager:
brew install enzyme
spack install enzyme
nix-shell -p enzyme
To get involved or if you have questions, please join our mailing list.
If using this code in an academic setting, please cite the following three papers (first for Enzyme as a whole, second for GPU+optimizations, and third for AD of all other parallel programs (OpenMP, MPI, Julia Tasks, etc.)):
@inproceedings{NEURIPS2020_9332c513,
author = {Moses, William and Churavy, Valentin},
booktitle = {Advances in Neural Information Processing Systems},
editor = {H. Larochelle and M. Ranzato and R. Hadsell and M. F. Balcan and H. Lin},
pages = {12472--12485},
publisher = {Curran Associates, Inc.},
title = {Instead of Rewriting Foreign Code for Machine Learning, Automatically Synthesize Fast Gradients},
url = {https://proceedings.neurips.cc/paper/2020/file/9332c513ef44b682e9347822c2e457ac-Paper.pdf},
volume = {33},
year = {2020}
}
@inproceedings{10.1145/3458817.3476165,
author = {Moses, William S. and Churavy, Valentin and Paehler, Ludger and H\"{u}ckelheim, Jan and Narayanan, Sri Hari Krishna and Schanen, Michel and Doerfert, Johannes},
title = {Reverse-Mode Automatic Differentiation and Optimization of GPU Kernels via Enzyme},
year = {2021},
isbn = {9781450384421},
publisher = {Association for Computing Machinery},
address = {New York, NY, USA},
url = {https://doi.org/10.1145/3458817.3476165},
doi = {10.1145/3458817.3476165},
booktitle = {Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis},
articleno = {61},
numpages = {16},
keywords = {CUDA, LLVM, ROCm, HPC, AD, GPU, automatic differentiation},
location = {St. Louis, Missouri},
series = {SC '21}
}
@inproceedings{10.5555/3571885.3571964,
author = {Moses, William S. and Narayanan, Sri Hari Krishna and Paehler, Ludger and Churavy, Valentin and Schanen, Michel and H\"{u}ckelheim, Jan and Doerfert, Johannes and Hovland, Paul},
title = {Scalable Automatic Differentiation of Multiple Parallel Paradigms through Compiler Augmentation},
year = {2022},
isbn = {9784665454445},
publisher = {IEEE Press},
booktitle = {Proceedings of the International Conference on High Performance Computing, Networking, Storage and Analysis},
articleno = {60},
numpages = {18},
keywords = {automatic differentiation, tasks, OpenMP, compiler, Julia, parallel, Enzyme, C++, RAJA, hybrid parallelization, MPI, distributed, LLVM},
location = {Dallas, Texas},
series = {SC '22}
}
Julia bindings, Rust bindings, and Fortran bindings are available for Enzyme.