Instead of Rewriting Foreign Code for Machine Learning, Automatically Synthesize Fast Gradients

1 Introduction

• ML frameworks require differentiated code to use their DSLs, while scientific tools such as physics simulators and climate models often use other languages.The paper identifies rewriting as a central challenge in applying ML to scientific computing.Source: §1 · p. 1
• For norm, optimizing before AD yields an O(N) gradient, whereas optimizing after AD leaves an O(N^2) gradient.The figure attributes the difference to the adjoint remaining inside the loop when its input depends on a value computed there.Source: Figure 1 · p. 2, §1 · p. 2
• Operator overloading requires replacing standard utilities with differentiable operators, while source rewriting requires ahead-of-time source and a supported language subset.These constraints hinder use of libraries, including precompiled and header-only libraries.Source: §1 · p. 2
• Enzyme is an LLVM compiler plugin that synthesizes gradients for statically analyzable IR, including IR from C, C++, Fortran, Rust, and Swift.It performs AD on optimized IR, a low-level representation traditional wisdom considered difficult for efficient AD.Source: §1 · p. 2
• The contributions also include PyTorch and TensorFlow interfaces for foreign code, Enzyme.jl for Julia, and multisource AD and static-library support using LTO.The paper reports a benchmark study in which AD after optimization achieves state-of-the-art performance.Source: §1 · p. 2
• Unlike related systems described here, Enzyme operates on generic low-level IR after optimization, which the authors say provides performance and compatibility benefits.The related-work comparison includes differentiable DSLs and language-specific AD tools.Source: §1 · p. 3

• Enzyme runs type analysis, activity analysis, and gradient synthesis on optimized LLVM IR, producing forward and reverse passes linked by instruction adjoints.The reverse pass computes gradients by inverting forward-pass instructions.Source: §2 · p. 3
• Type analysis recovers underlying data types that LLVM IR can hide, using interprocedural byte-offset type trees propagated from TBAA metadata to a fixed point.This is needed because different underlying types require different gradient operations for copied memory.Source: §2 · p. 3
• Activity analysis uses LLVM alias and type analyses to skip instructions unable to propagate differentials, including read-only integer-returning calls.An instruction is active if it can propagate a differential to its return value or another memory location.Source: §2 · p. 4
• Shadow allocations mirror active values and data structures; Enzyme caches unrecomputable forward values for reverse adjoints and reduces caching through recomputation and reuse.When the number of cached values is statically bounded, Enzyme can allocate storage once; otherwise it may reallocate dynamically.Source: §2 · pp. 4–5
• For each original basic block, synthesis emits adjoints in reverse instruction order, branches to the predecessor’s reverse block, and redirects forward returns into reverse control flow.This constructs the gradient function from the analyses and initialized shadow values.Source: §2 · p. 4
• When legal, Enzyme merges a callee’s forward work into its adjoint; indirect calls obtain augmented-forward and gradient functions through shadow function pointers.For a virtual C++ class argument, the shadow must supply a modified virtual method table conforming to Enzyme’s calling convention.Source: §2 · p. 5
• Enzyme requires IR access and statically deducible types for active memory operations and phi nodes, and it does not implement adjoints for exception-handling instructions.The paper advises disabling exceptions, for example with -fno-exceptions for C++.Source: §2 · p. 6

3 Usage

• Enzyme is an LLVM compiler plugin designed to integrate into existing tools without requiring custom forks of LLVM, PyTorch, or TensorFlow.The plugin lets LLVM-based languages call __enzyme_autodiff; users can mark variables active through a variable or metadata.Source: §3 · p. 6
• For multi-source programs and static libraries, Enzyme requires the relevant IR at pass time, supplied through LTO or embedded bitcode.Single-source programs already have their IR available; static libraries can be compiled with -fembed-bitcode.Source: §3 · p. 6
• Users can specify custom forward and backward passes as function metadata, including for functions whose definitions are unavailable during AD.Enzyme provides rules for common functions in the C/C++ standard and math libraries.Source: §3 · p. 6
• Enzyme.jl runs the Enzyme pass on collected Julia IR and replaces some Julia math calls with LLVM intrinsics using Cassette.jl.Julia’s execution engine caches code, so Enzyme.jl collects function definitions reachable from the differentiated function.Source: §3 · p. 7
• Embedding Enzyme in Zygote combines high-level knowledge with low-level optimization, addressing Zygote’s poor performance on scalar programs and enabling derivatives of foreign functions.Enzyme.jl uses embedded bitcode to differentiate foreign functions.Source: §3 · p. 7
• PyTorch and TensorFlow users can call Enzyme through Python packages, avoiding the need to compile a custom operator themselves.Users can also call __enzyme_autodiff and compile a custom operator when using the lower-level workflow.Source: §3 · p. 7, Figure 7 · p. 7

4 Evaluation

• The seven tests comprise ADBench’s BA, LSTM, and GMM, plus Euler and RK4 integrators, a recursive FFT, and a discretized 2D Brusselator simulation.The integrators test indirect calls and complex C++ headers; FFT tests recursive functions, and Bruss represents a use for adjoint sensitivity analysis.Source: §4 · p. 7
• Experiments used a quiesced AWS c4.8xlarge instance and geometric means across inputs, excluding 21 of ADBench’s 92 inputs.The excluded inputs were cases where Adept exhausted system memory or a tool ran in under 0.01 seconds; integrator and FFT tests used 36 inputs, and Bruss used 10 trials.Source: §4 · p. 8
• 4.5x: Enzyme outperforms Ref by a factor of 4.5 in geometric mean runtime across all benchmarks and inputs.Ref performs AD before its first optimization round; Enzyme performs AD after that round, with a second optimization round in both pipelines.Source: §4 · p. 8
• Enzyme demonstrates state-of-the-art performance on all benchmarks; the reported advantages vary with optimization order and tape design.Optimization before AD appears to help BA, LSTM, Euler, and RK4; tape structure accounts for differences on GMM and Bruss, while reduced tape storage helps on FFT and Bruss.Source: §4 · p. 8, Figure 9 · p. 8

5 Conclusion

• Enzyme demonstrates the feasibility of efficient AD on low-level programs, opening the way to language-independent AD and AD after optimization.The conclusion frames this as a way to bring foreign code into machine-learning workflows without rewriting it or creating a new DSL.Source: §5 · p. 8
• Future directions include AD-specific LLVM optimizations, using GPU or parallel code generators, and extending support to GPU and CPU-parallel programs.The authors also identify forward-mode and mixed-mode AD, checkpointing beyond a simple heuristic, and porting physics engines and other codebases to ML frameworks.Source: §5 · p. 8

Broader Impact

• Enzyme reduces the workload researchers face when applying machine learning to new domains.The authors also note that this may negatively affect people whose job is manually rewriting existing code for ML frameworks.Source: Broader Impact · p. 9
• By making tools easier to share and interoperate, Enzyme may help bridge the machine-learning and scientific-computing communities.The paper gives projected-gradient descent on a climate simulator as a possible application.Source: Broader Impact · p. 9
• The authors cite improved climate-policy design through projected-gradient descent on a climate simulator as a possible benefit.They describe this as a potential application, not a demonstrated result.Source: Broader Impact · p. 9