Differentiating Compiled Code (Enzyme + Fortran)¶

Context¶

The Fortran Integration example wraps a compiled solver as a Tesseract, but stops at the forward pass. This example adds exact, machine-precision derivatives with no hand-written adjoint code: the solver is compiled to LLVM IR with LFortran, Enzyme generates forward- and reverse-mode derivatives directly from that IR, and the result is a shared library exposing all three differentiable endpoints. The pattern applies to any code that lowers to LLVM IR (C, C++, Rust).

Example Tesseract (examples/fortran_enzyme)¶

The Fortran kernel¶

The solver computes a single explicit Euler step of the 1D heat equation:

\[\frac{\partial T}{\partial t} = \alpha \frac{\partial^2 T}{\partial x^2}\]

discretized with central differences as \(T_\text{out}(i) = T_\text{in}(i) + r \, (T_\text{in}(i-1) - 2 T_\text{in}(i) + T_\text{in}(i+1))\), where \(r = \alpha \, \Delta t / \Delta x^2\):

subroutine heat_step(n, T_in, T_out, alpha, dx, dt)
  implicit none
  integer, intent(in) :: n
  double precision, intent(in) :: T_in(n), alpha, dx, dt
  double precision, intent(out) :: T_out(n)
  integer :: i
  double precision :: r

  r = alpha * dt / (dx * dx)

  ! Dirichlet boundary conditions
  T_out(1) = T_in(1)
  T_out(n) = T_in(n)

  ! Interior points: explicit finite difference stencil
  do i = 2, n-1
    T_out(i) = T_in(i) + r * (T_in(i-1) - 2.0d0*T_in(i) + T_in(i+1))
  end do
end subroutine

Two details keep the generated IR clean enough for Enzyme to differentiate reliably: the stencil uses explicit do loops (not whole-array intrinsics), and LFortran is run with --no-array-bounds-checking, since bounds checks emit runtime calls Enzyme cannot trace through.

Input and output schemas¶

Every differentiable field is marked with Differentiable[...]. The InputSchema also enforces positivity and the CFL stability condition \(r \leq 0.5\), skipping the check during abstract_eval when only shapes are known:

class InputSchema(BaseModel):
    """Input for a single explicit Euler step of the 1D heat equation."""

    T_in: Differentiable[Array[(None,), Float64]] = Field(
        description="Temperature profile at current time step [K]. Shape: (n,). "
        "Boundary values T_in[0] and T_in[-1] are held fixed (Dirichlet).",
    )
    alpha: Differentiable[Float64] = Field(
        default=0.01,
        description="Thermal diffusivity [m^2/s]. Must be > 0.",
    )
    dx: Differentiable[Float64] = Field(
        default=0.25,
        description="Grid spacing [m]. Must be > 0.",
    )
    dt: Differentiable[Float64] = Field(
        default=0.001,
        description="Time step size [s]. Must be > 0.",
    )

    @model_validator(mode="after")
    def check_stability(self) -> Self:
        """Verify positivity and CFL stability condition: r = alpha * dt / dx^2 <= 0.5."""
        if isinstance(self.alpha, ShapeDType):
            return self  # skip during abstract_eval
        for name in ("alpha", "dx", "dt"):
            if getattr(self, name) <= 0:
                raise ValueError(f"{name} must be > 0, got {getattr(self, name)}")
        r = self.alpha * self.dt / (self.dx**2)
        if r > 0.5:
            raise ValueError(
                f"CFL stability condition violated: r = {r:.4f} > 0.5. "
                f"Reduce dt or alpha, or increase dx."
            )
        return self

    @model_validator(mode="after")
    def check_min_points(self) -> Self:
        """Need at least 3 points for interior stencil."""
        if isinstance(self.T_in, ShapeDType):
            return self  # skip during abstract_eval
        if len(self.T_in) < 3:
            raise ValueError("T_in must have at least 3 points.")
        return self

class OutputSchema(BaseModel):
    """Output: temperature profile after one heat equation step."""

    T_out: Differentiable[Array[(None,), Float64]] = Field(
        description="Temperature profile after one time step [K]. Shape: (n,).",
    )

The compilation pipeline¶

The Enzyme magic happens at image build time. A C wrapper (wrapper.c) bridges the Fortran ABI and Enzyme: it copies scalar arguments to the stack (Fortran passes everything by pointer) and declares __enzyme_autodiff / __enzyme_fwddiff calls annotated with enzyme_dup (differentiate) or enzyme_const (hold fixed). These sentinel calls are replaced by generated derivative code when the Enzyme pass runs:

/* ── Reverse mode (VJP) ────────────────────────────────────────────── */

void heat_step_vjp(int n,
                   const double* T_in,  double* dT_in,
                   const double* T_out, double* dT_out,
                   double alpha,  double* dalpha,
                   double dx,     double* ddx,
                   double dt,     double* ddt)
{
    int n_ = n;
    double alpha_ = alpha, dx_ = dx, dt_ = dt;

    __enzyme_autodiff((void*)heat_step,
        enzyme_const, &n_,
        enzyme_dup, (double*)T_in,  dT_in,
        enzyme_dup, (double*)T_out, dT_out,
        enzyme_dup, &alpha_,        dalpha,
        enzyme_dup, &dx_,           ddx,
        enzyme_dup, &dt_,           ddt);
}

The build.sh script chains the toolchain together, emitting libheat_ad.so with three entry points—heat_step_forward (primal), heat_step_jvp (forward-mode), and heat_step_vjp (reverse-mode):

echo "=== Step 1: LFortran -> LLVM IR ==="
lfortran --show-llvm --no-array-bounds-checking \
    "${SCRIPT_DIR}/heat_step.f90" > /tmp/heat_step.ll

echo "=== Step 2: Optimize IR ==="
opt -O1 -S /tmp/heat_step.ll -o /tmp/heat_step_opt.ll

echo "=== Step 3: Compile C wrapper -> LLVM IR ==="
clang -emit-llvm -S -O1 "${SCRIPT_DIR}/wrapper.c" -o /tmp/wrapper.ll

echo "=== Step 4: Link IR modules ==="
llvm-link /tmp/wrapper.ll /tmp/heat_step_opt.ll -S -o /tmp/combined.ll

echo "=== Step 5: Enzyme AD pass ==="
opt --load-pass-plugin="${ENZYME_LIB}" -passes=enzyme \
    -S /tmp/combined.ll -o /tmp/ad.ll

echo "=== Step 6: Compile to shared library ==="
clang -shared -O2 /tmp/ad.ll -o "${OUTPUT}" -lm

echo "=== Built ${OUTPUT} ==="

Wiring the library into the Tesseract API¶

tesseract_api.py loads the shared library via ctypes and declares the signatures of the three entry points. The apply function calls the primal kernel:

def apply(inputs: InputSchema) -> OutputSchema:
    """Compute one explicit Euler step of the 1D heat equation."""
    T_in = np.ascontiguousarray(inputs.T_in, dtype=np.float64)
    n = len(T_in)
    T_out = np.zeros(n, dtype=np.float64)

    _lib.heat_step_forward(
        n, _as_ptr(T_in), _as_ptr(T_out), inputs.alpha, inputs.dx, inputs.dt
    )

    return OutputSchema(T_out=T_out)

The differentiable endpoints call the Enzyme-generated wrappers. Reverse mode threads cotangents into shadow arrays that Enzyme accumulates gradients into; the vector_jacobian_product endpoint then returns only the requested inputs. The jacobian_vector_product endpoint is the forward-mode mirror image, calling heat_step_jvp with input tangents instead.

def vector_jacobian_product(
    inputs: InputSchema,
    vjp_inputs: set[str],
    vjp_outputs: set[str],
    cotangent_vector: dict[str, Any],
):
    """Reverse-mode AD via Enzyme: compute v^T @ J."""
    cotangent_T_out = cotangent_vector.get("T_out", np.zeros_like(inputs.T_in))
    dT_in, dalpha, ddx, ddt = _run_vjp(inputs, cotangent_T_out)

    result = {}
    if "T_in" in vjp_inputs:
        result["T_in"] = dT_in
    if "alpha" in vjp_inputs:
        result["alpha"] = dalpha
    if "dx" in vjp_inputs:
        result["dx"] = ddx
    if "dt" in vjp_inputs:
        result["dt"] = ddt
    return result

Build configuration¶

The tesseract_config.yaml installs the LLVM 19 toolchain, LFortran (via micromamba from conda-forge), and the prebuilt Enzyme plugin as custom_build_steps, then runs build.sh to produce the differentiated library. All tools are installed from prebuilt binaries—the toolchain itself is never compiled from source:

name: "enzyme-ad"
version: "1.0.0"
description: |
  Differentiable 1D heat equation solver using Enzyme automatic differentiation.

  Demonstrates how to obtain exact (machine-precision) derivatives of a Fortran
  simulation without writing manual adjoint code.  The pipeline is:

    Fortran -> LFortran -> LLVM IR -> Enzyme AD pass -> shared library

  Enzyme generates both forward-mode (JVP) and reverse-mode (VJP) derivatives
  directly from the compiled Fortran code at the LLVM IR level.

  Industry relevance: gradient-based optimization of thermal parameters,
  inverse problems, sensitivity analysis, and differentiable physics.

build_config:
  base_image: "debian:bookworm-slim"
  target_platform: "linux/amd64"

  extra_packages:
    - wget
    - gnupg
    - ca-certificates
    - bzip2

  package_data:
    - ["enzyme/heat_step.f90", "enzyme/heat_step.f90"]
    - ["enzyme/wrapper.c", "enzyme/wrapper.c"]
    - ["enzyme/build.sh", "enzyme/build.sh"]

  custom_build_steps:
    # Install LLVM 19 toolchain
    - |
      RUN wget -qO- https://apt.llvm.org/llvm-snapshot.gpg.key | gpg --dearmor -o /etc/apt/keyrings/llvm.gpg && \
          echo "deb [signed-by=/etc/apt/keyrings/llvm.gpg] http://apt.llvm.org/bookworm/ llvm-toolchain-bookworm-19 main" > /etc/apt/sources.list.d/llvm.list && \
          apt-get update && apt-get install -y --no-install-recommends llvm-19 clang-19 && \
          rm -rf /var/lib/apt/lists/* && \
          for tool in opt llvm-as llvm-link llvm-dis llc clang clang++; do \
            ln -sf /usr/bin/${tool}-19 /usr/local/bin/${tool} 2>/dev/null || true; \
          done

    # Install LFortran via micromamba (prebuilt from conda-forge)
    - |
      RUN wget -q https://github.com/mamba-org/micromamba-releases/releases/latest/download/micromamba-linux-64 \
            -O /usr/local/bin/micromamba && chmod +x /usr/local/bin/micromamba && \
          MAMBA_ROOT_PREFIX=/opt/conda micromamba create -y -n base -c conda-forge lfortran=0.61.0 && \
          ln -sf $(find /opt/conda -name lfortran -type f | head -1) /usr/local/bin/lfortran && \
          echo /opt/conda/lib > /etc/ld.so.conf.d/conda.conf && ldconfig

    # Download prebuilt Enzyme LLVM plugin
    - |
      RUN wget -q https://github.com/EnzymeAD/Enzyme/releases/download/nightly/LLVMEnzyme-19.so \
            -O /usr/local/lib/LLVMEnzyme-19.so

    # Build the differentiated shared library
    - |
      RUN chmod +x /tesseract/enzyme/build.sh && \
          /tesseract/enzyme/build.sh /tesseract/enzyme/libheat_ad.so

Adapting this pattern¶

To differentiate your own compiled code with Enzyme:

Write a C wrapper declaring __enzyme_autodiff / __enzyme_fwddiff calls, annotating each argument enzyme_dup or enzyme_const, against a clean kernel (explicit loops, no runtime checks Enzyme can’t trace through).
Assemble the pipeline as custom_build_steps: lower to LLVM IR, link the wrapper, run the Enzyme pass, and compile to a shared library.
Wire it up: load the library with ctypes, mark differentiable fields with Differentiable[...], and call the JVP/VJP wrappers from the gradient endpoints.