Lattice Gauge Theory via LLVM-Level Automatic Differentiation

This paper presents a method for automatically generating Hybrid Monte Carlo forces in lattice gauge theory by applying reverse-mode automatic differentiation directly to optimized LLVM intermediate representation, enabling a single-source workflow that produces performance-portable, high-performance force implementations for both CPU and GPU backends without the need for manual derivation.

The Gist

Imagine you are trying to simulate the universe on a computer. Specifically, you are trying to simulate the "glue" that holds quarks together inside protons and neutrons. This field is called Lattice Gauge Theory.

To do this, scientists use a method called Hybrid Monte Carlo (HMC). Think of HMC as a giant, complex game of "Roll the Dice and Move."

1. You have a giant grid (the lattice) representing space.
2. You have a "score" (the Action) that tells you how good a specific arrangement of the grid is.
3. To move to a new arrangement, you need to know which direction to push. In physics, this push is called a Force.

The Old Problem: The "Double Work" Nightmare

In the past, calculating this "Force" was a massive headache.

• Step 1: A physicist would write a computer program to calculate the Score (the Action).
• Step 2: They would then have to sit down with a pen and paper, do complex calculus, and figure out exactly how the Score changes if they tweak the grid.
• Step 3: They would write a second, completely separate computer program to calculate that Force.

The Analogy: Imagine you are a chef. You write a recipe to bake a cake (the Action). But then, to figure out how to improve the cake, you have to hire a different mathematician to write a completely separate manual explaining how to tweak the ingredients. If you change the recipe later, you have to fire the mathematician, hire a new one, and rewrite their manual. If the manual has a typo, your cake burns, and you don't know why.

This "Double Work" was slow, prone to errors, and made it very hard to try new, complex recipes (actions) because the math was too difficult to rewrite every time.

The New Solution: The "Magic Mirror" (LLVM & Automatic Differentiation)

The authors of this paper, Yuki Nagai, Hiroshi Ohno, and Akio Tomiya, found a way to automate Step 2 and 3. They used a tool called Automatic Differentiation (AD), but they took it to a very deep level of the computer's brain.

The Analogy:
Instead of hiring a mathematician to rewrite the manual, they built a Magic Mirror.

1. You write your recipe (the Action code) just once.
2. You hold it up to the Magic Mirror (the LLVM-level AD system).
3. The mirror instantly reflects back a perfect, working manual on how to tweak the ingredients (the Force).

If you change the recipe, the mirror updates the manual instantly. No new math, no new code, no errors.

How Does the Mirror Work? (The "LLVM" Part)

Computers don't understand high-level code like Julia or C++ directly. They translate it into a lower-level language called LLVM IR (Intermediate Representation) before running it. This is like translating a novel into a universal code of "0s and 1s" that the computer's processor actually executes.

The authors realized that if you ask the mirror to look at the translated code (the LLVM level) rather than the original recipe, it can see exactly how the computer is moving data in memory.

• Old way: The mirror tries to guess the math based on the recipe.
• New way: The mirror watches the computer actually do the work, step-by-step, and then runs the movie in reverse to see how the changes happened.

This is called Reverse-Mode Automatic Differentiation. It's like watching a video of a building being constructed, then playing it backward to see exactly which bricks were laid first and how to take them apart.

Why is this a Big Deal?

1. One Source Code: You write the Action code once. The Force code is generated automatically from it. No more "Double Work."
2. No More Bugs: Since the Force is mathematically derived directly from the Action by the computer, it is impossible for the Force to be inconsistent with the Action. The "manual" always matches the "recipe."
3. Speed: The authors tested this on supercomputers (CPUs and GPUs). They found that the "Magic Mirror" generated a Force that was just as fast as the best human-written Force code.
4. Complexity Made Easy: Modern physics uses very complex "recipes" (like multi-layered smearing). Writing the math for these by hand is nearly impossible. With this method, you just write the code, and the mirror handles the complex calculus.

The Bottom Line

This paper introduces a way to let computers do the heavy lifting of calculus for them. By looking at the code at the lowest possible level (LLVM), they created a system where writing the rules of the universe automatically writes the rules for how to simulate it.

It turns a tedious, error-prone, double-work process into a single, streamlined workflow. It's like upgrading from a typewriter where you have to retype the whole book to fix a typo, to a word processor that can instantly calculate the "undo" history for you. This allows physicists to explore more complex theories of the universe faster and with greater confidence.

Technical Summary

1. Problem Statement

Lattice Gauge Theory (LGT) simulations, particularly those involving dynamical fermions, rely heavily on the Hybrid Monte Carlo (HMC) algorithm. A critical component of HMC is the Molecular Dynamics (MD) force, which is the derivative of the lattice action with respect to the gauge link variables ($U_\mu$).

Current workflows face significant challenges:

• Manual Derivation & Implementation: For complex, compositional actions (e.g., improved gauge actions, multi-level smearing like stout smearing, or rational approximations in RHMC), researchers must analytically derive the force equations and implement a separate, distinct code routine.
• Inconsistency Risks: Maintaining two separate code paths (one for action evaluation, one for force calculation) leads to frequent subtle bugs where the force does not perfectly match the action, compromising simulation stability.
• Portability Issues: As simulations migrate to heterogeneous High-Performance Computing (HPC) environments (CPUs and GPUs), maintaining separate, architecture-specific force implementations increases development costs and reduces software sustainability.
• Limitations of Existing AD: Standard Automatic Differentiation (AD) tools used in machine learning (e.g., PyTorch, TensorFlow) are optimized for functional, tensor-centric graphs. They often struggle with the imperative, in-place memory updates, explicit loops, and stencil operations characteristic of high-performance lattice codes. Furthermore, tape-based reverse-mode AD can incur prohibitive memory overhead for long HMC trajectories.

2. Methodology

The authors propose a novel framework that generates HMC forces by performing reverse-mode automatic differentiation (AD) directly at the LLVM Intermediate Representation (IR) level.

• Compiler-Level Viewpoint: Instead of treating the action as a symbolic functional, the action is viewed as a discrete-time dynamical system of program states ($U_0 \to U_1 \to \dots \to S$).
• SSA Form & In-Place Updates: The method leverages Static Single Assignment (SSA) form, which is the standard representation in LLVM. In source code, lattice simulations often use "in-place" updates (overwriting variables). In LLVM IR, these are represented as distinct SSA values ($\%U_0, \%U_1, \dots$). This allows the compiler to track the dataflow of every intermediate state without ambiguity, enabling a well-defined reverse pass.
• Enzyme Integration: The implementation uses Enzyme, an LLVM-level AD tool integrated into the Julia programming language.
  • The user writes the action evaluation routine in a conventional, imperative style (with loops and in-place updates).
  • Julia compiles this to optimized LLVM IR.
  • Enzyme analyzes the optimized instruction stream and automatically generates the adjoint program (the reverse pass) that propagates sensitivities from the final scalar action back to the initial link variables.
• Force Construction:
  1. Enzyme computes the matrix derivative $\partial S / \partial U$ (via Wirtinger calculus on complex numbers).
  2. The HMC force is constructed by right-multiplying by $U^\dagger$ and projecting the result onto the Lie algebra $\mathfrak{su}(N_c)$ to ensure the force is traceless and anti-Hermitian.
  3. This process is applied to both gauge actions and Wilson fermion actions (including smearing chains).

3. Key Contributions

• Single-Source Workflow: The framework enables a "write-once" workflow where the HMC force is generated automatically from the action code. No manual derivation or separate force kernel implementation is required.
• Performance Portability: The same Julia source code and differentiation pipeline target both CPU and GPU backends (via the JACC.jl library). This decouples the physics definition from the hardware-specific implementation.
• Handling Imperative Code: By operating at the LLVM IR level, the method naturally handles aggressive compiler optimizations, in-place memory reuse, and explicit control flow, which are standard in LGT but problematic for higher-level AD tools.
• Fermion Force Generation: The paper demonstrates the end-to-end generation of forces for Wilson fermions, a complex case involving linear solvers and nested operations, without differentiating the solver itself (using implicit differentiation logic derived from the Dirac operator structure).

4. Results and Validation

The authors validated the approach through rigorous testing:

• Link-Level Accuracy:
  • The automatically generated force ($F_{AD}$) for the Wilson fermion action was compared against a conventional, hand-written analytic force ($F_{ana}$) on an $SU(3)$ $16^4$ lattice.
  • Result: The relative Frobenius-norm error was uniform at the $10^{-9}$ level, consistent with double-precision roundoff error, confirming mathematical correctness.
• Dynamical Consistency (Energy Conservation):
  • HMC simulations were run using the generated forces. The energy violation $\Delta H$ was measured against the MD step size ($\Delta t$).
  • Result: The simulations exhibited the expected quadratic scaling ($\Delta H \propto \Delta t^2$) characteristic of a second-order symplectic integrator. This confirms that the generated forces are reversible and symplectic, ensuring stable long-term trajectories.
• Performance:
  • Benchmarks on a $24^4$ lattice (CPU single-thread and NVIDIA H100 GPU) showed that the LLVM-AD generated force is competitive with, and in some cases faster than, the hand-written implementation.
  • CPU: LLVM-AD took ~162s vs. 195s for hand-written.
  • GPU: LLVM-AD took ~3.4s vs. 4.6s for hand-written.
  • This demonstrates that IR-level differentiation introduces negligible overhead and fully inherits compiler optimizations.
• Memory Footprint:
  • On an H100 GPU, the memory footprint for the $24^4$ lattice was ~22 GB without special checkpointing techniques, indicating the approach is viable for current accelerator hardware.

5. Significance and Outlook

• Reduced Development Cost: By eliminating the need for analytic derivation and separate force coding, the method significantly lowers the barrier to implementing complex, compositional lattice actions (e.g., deep smearing chains).
• Bug Reduction: It removes a major source of errors in lattice simulations by ensuring the force is mathematically guaranteed to be the derivative of the specific action code being used.
• Future-Proofing: The approach facilitates rapid exploration of new action formulations and smearing strategies. It also bridges the gap between lattice gauge theory and modern differentiable programming ecosystems.
• Scalability: While memory reduction techniques (like checkpointing) are future work, the current results prove that compiler-generated forces are practical for large-scale simulations on modern HPC systems.

In summary, this paper presents a paradigm shift in lattice gauge theory software engineering, moving from manual, error-prone force derivation to a compiler-automated, single-source, and performance-portable workflow.