AI authored to showcase ironpad capabilities; human edited.
This notebook calls the real std::autodiff: the nightly feature where Enzyme differentiates your compiled LLVM IR. You write one attribute; the compiler does calculus.
This demo is a simulation: a cannonball with quadratic air drag, integrated step by step. That choice is the point. With drag there is no simple analytical formula for where the ball lands; the only way to know is to run the loop. Which raises the question: if there's no formula, where does the derivative come from?
From the compiler! 🎉
#[autodiff_reverse] on the simulator function gives us d(range)/d(angle) in a single reverse pass, straight through the timestepped loop. We'll verify it against finite differences, break it in a way that took me a while to diagnose, and then let Newton's method aim the cannon in 4 iterations.
If you want the theory first (dual numbers, tapes, why reverse mode powers all of ML), the companion notebook std::autodiff: Differentiation in the Compiler hand-rolls everything. This one uses the real thing.
In a vacuum, gravity is the only force: the horizontal and vertical motions separate cleanly, and two integrations hand you the parabola: R = v_0^2 \sin(2\theta)/g, optimum at 45°. This is a closed form solution: plug in an angle, get a range, differentiate with a pencil.
When we add air, drag on a fast projectile is roughly quadratic in speed and acts against the velocity vector:
\frac{dv_x}{dt} = -k\,|v|\,v_x \qquad\qquad \frac{dv_y}{dt} = -g - k\,|v|\,v_y \qquad\qquad |v| = \sqrt{v_x^2 + v_y^2}
That |v| term welds the two equations together: the horizontal deceleration now depends on the vertical speed, and vice versa, through a square root. Every standard differentiation trick dead-ends immediately:
\frac{dv_x}{v_x} = -k \sqrt{v_x^2 + v_y^2}\; dt
Clean separation fails here: v_y is trapped inside the square root.
The two-dimensional quadratic-drag problem has no elementary closed-form solution (purely vertical motion has one; the full trajectory does not). It is why armies published hand-computed range tables, and why computing more of them was ENIAC's first job.
And drag is not a small correction here: the ball leaves the muzzle at 100 m/s with k = 0.01, so the initial drag deceleration is k v^2 = 100 m/s², about ten times gravity. Ignoring it is not an option (if you want to hit your target).
So we do what the range-table computers did: integrate, step by step. The simulator below is the entire physics engine, and a cool way to show off the std::autodiff capabilities.
The attribute is doing a lot of heavy lifting here. #[autodiff_reverse(d_range, Active, Active)] asks the compiler to generate a second function, d_range, that returns (range, d(range)/d(angle)) in one reverse pass through the compiled loop. There's no real formula here: just the compiler's inbuilt capabilities.
The loop has a data-dependent trip count (the ball lands when it lands), and the landing branch interpolates the final partial step instead of stopping on a whole timestep. The interpolation gets its own section below. The integrator itself is explicit Euler on 10 ms steps, the plainest available; what that plainness costs, we will measure a few cells down.
The range function is not particularly special or beautiful, and the compiler doesn't care.
The rest of the code in this notebook is collapsed by default, since the meat is in the shared code above. You can expand the code to see the application, if you'd like.
With drag, range tops out near 0.570 rad (32.7°) at 168.3 m, not at the vacuum-blessed 45°.
Same muzzle speed, formula against simulation, head to head:
The integrator inside range is explicit Euler, the bluntest tool in numerical analysis: freeze the forces for one timestep, coast in a straight line, recompute. Each step commits an O(dt^2) local error; over the roughly 500 steps of a flight these accumulate into an O(dt) global error. First order: halve the step, halve the error.
The sweep below re-integrates one shot at step sizes from 160 ms down to 5 ms against a 0.1 ms reference. The error column reads 7.7 m, 3.8 m, 1.9 m, 0.95 m: each row half the one above, the first-order signature visible in raw digits.
The trade-off is the classic one: accuracy costs steps (30 at dt = 0.16, 989 at dt = 0.005). But there is a second bill that matters here: the reverse pass we are building toward walks every step backwards, and its stored state scales with the trip count. Your step size is also your derivative's memory budget.
range(angle) exists only as a loop, and Newton's method (where this is headed) will demand dR/d\theta. So where does a derivative come from when there is nothing to differentiate by hand?
The classical workaround is finite differences: nudge the angle, re-run the sim, divide. It works, but you pay a full simulation per probe and you inherit the step-size error valley: truncation error fighting roundoff, best case ~1e-11, and a step size h you must pick blind.
Enzyme just works: it takes the compiled simulator, loop and all, and emits a second function that runs it backwards, accumulating exact adjoints. No h to pick, no truncation term, no extra simulations per probe. The next cell puts the two side by side.
d_range does no symbolic algebra and nudges nothing. Reverse-mode AD generates a program with two halves. The forward pass runs the simulation exactly as written, recording what the second half will need (here, each step's velocities). The reverse pass starts at the landing point, seeds the sensitivity of the answer with respect to itself (exactly 1.0), and walks the loop in reverse iteration order, applying the transposed chain rule at every step.
For a mental model: the cannonball flies twice. The first flight is the physics, muzzle to crater. The second runs crater back to muzzle, and what flows along it is influence rather than position: at each timestep it carries the answer to "if this velocity were perturbed, how much would the landing point move?" When the reverse flight arrives at t = 0, all of that influence has landed on the launch angle, and that number is dR/d\theta.
The loop's trip count is decided at runtime, and the reverse pass simply replays however many steps the tape recorded, 494 for this shot. One reverse sweep also prices the gradient for any number of inputs, which is why reverse mode underlies every neural network.
The next cell hand-rolls the scheme in about 35 lines and checks it against the compiler-generated d_range. If the adjoint picture is right, the two columns should agree far past anything finite differences can reach.
The first version of this simulator produced gradients that were exactly right and completely useless. Enzyme and finite differences agreed to four decimals, and both said the range was decreasing at angles where it was plainly increasing. Newton's method promptly launched the cannon into orbit.
Enzyme was fine. The loop originally stopped at the first whole timestep below ground, so as the angle varied, the landing point jumped in 10 ms quanta. range(angle) was a microscopic sawtooth: globally rising, locally falling, and every derivative method faithfully reported the slope of the teeth.
The fix is one line, and it applies whenever you differentiate through a simulation with a stopping condition: smooth the event. Interpolate the final partial step to y = 0 (see the landing branch in the simulator cell up top) and the sawtooth collapses into the piecewise-smooth function the physics always described. The attribute differentiates whatever you actually wrote. If your events are discontinuous, your gradients will be too.
We want the angle that drops the ball on a target at 150 m, which is a root-finding problem: define \text{miss}(\theta) = \text{range}(\theta) - 150 and drive it to zero.
Newton's method linearizes at the current guess and jumps to where the tangent line crosses zero: \theta \leftarrow \theta - \text{miss}(\theta)/\text{miss}'(\theta). Each iteration costs one forward flight plus one reverse flight (a single d_range call), so the whole aiming procedure below runs in about the budget of eight plain simulations. Because the derivative is exact, Newton's method delivers its famous quadratic convergence: near the root the error is squared every iteration, so the count of correct digits roughly doubles per step. A finite-difference slope would also converge here, but at two simulations per probe for a noisier slope; exactness is what makes the tables below so clean.
Newton's rule is easier to watch than to read. Plot \text{range}(\theta) against launch angle, draw the target as a flat line, and at each guess draw the tangent to the curve: its slope is exactly dR/d\theta, the number Enzyme returns. Where the tangent crosses the target is the next guess.
Drag the slider to take the steps one at a time from \theta_0 = 0.15. The miss falls 35 m → 8.5 m → 0.9 m → 0.01 m, the count of correct digits roughly doubling each step: quadratic convergence, and by the fourth step the tangent lands on the root.
Each Newton attempt, plotted where it lands.
The table above stops when the miss drops under 0.1 mm, which undersells how good the convergence is. Below is the same iteration with the miss in scientific notation, plus the ratio |\text{miss}_n| / |\text{miss}_{n-1}|^2. Quadratic convergence predicts that ratio is roughly constant (it is the curvature-to-slope ratio at the root), and the exponent staircase runs 1e+1, 1e+0, 1e-2, 1e-7, 1e-13, correct digits doubling once the iteration gets close, with the ratio column parked near 0.005–0.014 the whole way down.
range(θ) climbs to its 168.3 m maximum at 0.570 rad and then falls, so every reachable target is crossed twice: once on the way up (a flat shot) and once on the way down (a lofted one). For the 150 m target, the roots sit near 17.1° and 50.5°.
Newton is a local algorithm: it converges to whichever root's basin you start in. Start at 0.15 rad and you get the flat shot. Start at 1.2 rad, where dR/d\theta is negative, and the identical update rule marches the other way and finds the mortar solution. Same crater, wildly different flights: apex 18 m vs 80 m, hang time 3.8 s vs 7.9 s. Artillery doctrine has names for the two (direct and indirect fire); the optimizer rediscovers the distinction from nothing but a sign.
On the crest, dR/d\theta passes through zero (at 0.570 rad), so a start there makes Newton divide by almost nothing and catapult the guess far away; that is where you would add a bracket or a line search. From sensible starts, a handful of iterations land either shot.
The two solutions are clearest as a picture. Plot the miss, |\text{range}(\theta) - 150|, against launch angle and it falls into a double well: zero at the flat shot, zero at the lofted shot, with an overshoot hump between them where the crest carries the ball past 150 m before drag reins it back in.
The shading answers the sequel question by brute force: from every starting angle, run Newton to convergence and color that column by the trough it fell into. The split lands right on the crest, where dR/d\theta = 0 and the update briefly divides by almost nothing and flings the guess out of the physical range (the gray seam). Start at 0.15 rad and you roll into the flat trough; start at 1.2 and the identical rule catches the lofted one.
d_range is a real function in the compiled artifact, and you can go read it. The excerpts below come from a one-off build of this simulator with -Zautodiff=Enable,PrintModAfter, which dumps the LLVM module right after Enzyme runs; the WASM text comes from wasm-tools print on the linked module. Mangled symbol names are shortened and the annotations are mine; the instructions are otherwise verbatim.
Start with the call side, because it is tiny. In the final WASM, d_range is a twenty-line wrapper: reserve 16 bytes for the return pair, seed the output adjoint with 1.0, call the function Enzyme synthesized, and read the second slot back.
(func $d_range (param f64) (result f64)
(local i32)
global.get $__stack_pointer
i32.const 16
i32.sub ;; room for the (range, d_range) pair
local.tee 1
global.set $__stack_pointer
local.get 1
local.get 0 ;; angle
f64.const 0x1p+0 (;=1;) ;; the output seed: d(range)/d(range) = 1
call $diffe_range ;; Enzyme's synthesized reverse pass
local.get 1
f64.load offset=8 ;; second slot: d(range)/d(angle)
local.set 0
;; ...restore the stack pointer...
local.get 0)
diffe_range itself is 743 lines of WASM, so the readable view is the LLVM IR it descends from (452 lines, three clean phases). Phase one replays the flight forward, recording two doubles per step: vx and vy. The tape is a growable buffer that doubles its capacity whenever the step count hits a power of two; the ctpop intrinsic is the power-of-two test, and the realloc is the same amortized-doubling trick a Vec uses.
41: ; the forward flight loop
%47 = add nuw nsw i64 %42, 1 ; step counter
%51 = call i64 @llvm.ctpop.i64(i64 %47) ; power-of-two test:
br i1 %53, label %54, label %62 ; time to grow the tape?
54:
%61 = call ptr @realloc(ptr %48, i64 %57) ; double it
62:
%65 = getelementptr inbounds double, ptr %64, i64 %42
store double %44, ptr %65 ; tape_vx[step] = vx
; ...the same grow-and-store dance for tape_vy...
store double %43, ptr %83 ; tape_vy[step] = vy
; then one ordinary physics step, constants intact:
%86 = fadd double %85, %84 ; vx*vx + vy*vy
%87 = tail call double @llvm.sqrt.f64(double %86)
%88 = fmul double %87, 1.000000e-02 ; speed * DRAG
%93 = fadd double %92, 9.810000e+00 ; ... + G
%99 = fcmp olt double %98, 0.000000e+00 ; landed?
br i1 %99, label %102, label %100
Phase two handles the landing interpolation and seeds the adjoint. Phase three is the sweep: walk the step index down, reload that step's vx and vy from the tapes, recompute the speed (cheaper to recompute one sqrt than to tape a third double), and accumulate every intermediate's sensitivity with fadd into shadow slots. The derivative rules sit right in the arithmetic; here is the chain rule for sqrt, guard and all.
136: ; the reverse sweep
%164 = load i64, ptr %3 ; step index, counting DOWN
%167 = load double, ptr %166 ; vy = tape_vy[step]
%171 = load double, ptr %170 ; vx = tape_vx[step]
%173 = fadd double %168, %172 ; vx*vx + vy*vy, recomputed
%174 = tail call double @llvm.sqrt.f64(double %173)
; ...adjoint accumulation: shadow += contribution, dozens of times...
%176 = fmul fast double %163, %175
%178 = fadd fast double %177, %176
store double %178, ptr %17
; ...and the calculus rule for sqrt, adjoint / (2 * sqrt(x)),
; with a select guarding the sqrt(0) case:
%206 = fmul fast double 2.000000e+00, %205
%207 = fdiv fast double %203, %206
%208 = select fast i1 %204, double 0.000000e+00, double %207
This is the hand-rolled tape from the rewind section, written by the compiler: same forward recording, same backwards walk, same chain rule. The differences are the ones you would want. Enzyme tapes only what the reverse pass actually needs (two doubles per step, so the 494-step shot from earlier carries about 8 KB of tape, freed at the end of the sweep), it recomputes what is cheaper to recompute, and it differentiates the optimized IR, so the derivative inherits every optimization the primal got.
We converged to the target in 4 iterations because Newton's method works best under a specific condition: an exact derivative. The cost was a single attribute on an ordinary Rust function. Enzyme differentiates the optimized IR, loop and all, which is why the gradient sails straight through a timestepped simulation that has no closed form solution.
The feature itself is still nightly and moving: tracking issue #124509, driven by Manuel Drehwald with LLNL and the University of Toronto behind it, with batching and std::offload (GPU) queued up behind autodiff. Compose those and this notebook's punchline gets an upgrade: differentiated simulations running on the GPU, with the compiler doing the calculus end to end (std::offload doesn't quite work in Wasm in any interesting way, so it is not shown here).
For the theory under the hood, the companion notebook builds dual numbers and a reverse-mode tape from scratch, then benchmarks them against exactly the machinery Enzyme just used.