Discuss LLVM-level autodiff via Enzyme for cross-language composability #130
Description
Note: This issue was drafted with LLM assistance based on web research.
Context
The paper discusses backend flexibility and the potential for different computational backends. One avenue not currently discussed is LLVM-level automatic differentiation via Enzyme, which could enable cross-language composability.
What is Enzyme?
Enzyme is an LLVM plugin that performs automatic differentiation at the compiler level rather than the language level. By operating on LLVM intermediate representation, it can differentiate programs written in C, C++, Swift, Julia, Rust, Fortran, and TensorFlow using a single unified tool.
Key advantage: Enzyme differentiates code after LLVM optimisation occurs, producing substantially faster derivatives than tools that differentiate before optimisation (Moses & Churavy, 2020).
Current ecosystem status
Julia (Enzyme.jl)
- Enzyme.jl provides Julia bindings
- Turing.jl already supports Enzyme as an AD backend via DifferentiationInterface.jl
- Enzyme currently outperforms other AD packages for CPU performance, particularly for scalar code where Zygote struggles
- Most battle-tested backends in Julia are ForwardDiff or Enzyme (forward mode) and Zygote or Enzyme (reverse mode)
Rust
- Autodiff is being upstreamed into rustc - approved in lang triage meeting May 2024
- Now shipping in nightly for Linux
- GSoC 2025 project focused on making it stable
- Current limitations include: custom derivatives not yet exposed, limited enum support, checkpointing control not available, parallel code (rayon) not yet supported
- Rust 2024 H2 project goal to expose experimental LLVM features for autodiff and GPU offloading
Conference and community
- Enzyme Conference 2024 was held
- Enzyme is an LLVM Incubator project with plans for upstreaming
Potential benefits for composable ID modelling
- Cross-language components: Domain experts could write performance-critical components in Rust, C, or C++ whilst maintaining autodifferentiability within the framework
- Performance: LLVM-level optimisations applied before differentiation can yield substantial speedups
- Already integrated: Turing.jl (our current backend) already supports Enzyme via DifferentiationInterface.jl
- Broader contributor base: Researchers comfortable in Rust or C++ could contribute components without learning Julia
- Memory safety + performance: Rust components would bring memory safety guarantees with C-like performance
Relevance to the paper
This relates to:
- Discussion of backend flexibility and abstract backends (Section 3)
- Goal of lowering barriers for domain experts to contribute
- Cross-ecosystem accessibility (currently addressed via EpiAwareR)
- Backend-agnostic computational components (Section 3.2)
Questions to consider
- Is this worth mentioning in the discussion of alternative approaches or future work?
- Does it strengthen the argument about composability across computational paradigms?
- Is the Rust/cross-language angle too speculative given current limitations?
- Should we mention that Turing.jl already supports Enzyme, even if we don't use it in the case studies?