Enzyme.jl -- Reverse mode differentiation on LLVM IR for Julia
2021-07-28 , Green

Enzyme (https://enzyme.mit.edu) is a reverse mode auto-differentiation tool that performs automatic differentiation over LLVM intermediate representation and synthesis high-performance reverse-mode functions. We will discuss how Enzyme.jl integrates with the Julia compiler and special considerations required for differentiating a dynamic programming language such as Julia.


Automatic differentiation (AD) is key to training neural networks, bayesian inference, and scientific computing. This talk presents Enzyme.jl, a Julia frontend for the Enzyme high performance LLVM automatic differentiation (AD) toolkit. By operating at a low level, Enzyme is able to run optimizations prior to differentiation and is therefore highly efficient on scalar code and can support mutation out of the box. We explain how Enzyme.jl integrates with the Julia compiler, supports synthesis for Julia GPU kernels, and propagates Julia knowledge of types to the lower-level tool. We will discuss ongoing work to extend Enzyme.jl to be able to differentiate through Julia language features like dynamic calls and garbage collection. We will conclude by describing the potential of combining high level and low level systems to get the benefit of both algebraic and instruction level optimizations, and using Enzyme.jl in other AD systems such as Zygote.jl or Diffractor.jl to perform differentiation of foreign function calls, enabling cross-language AD.

PhD Candidate @ MIT