From dc9dfcbcc890f7149cd83a55a10726809471bb6d Mon Sep 17 00:00:00 2001 From: cwlm Date: Mon, 26 Oct 2026 15:10:03 +0100 Subject: [PATCH] Add `scan_kind="unrolled"` for explicit Runge--Kutta solvers On GPU, every check of a while loop's predicate is a device-to-host copy that the host waits for, and the stage loop of `AbstractRungeKutta.step` makes one per stage: 7 a Tsit5 step, 14 a Dopri8 step. With `scan_kind="unrolled"`, the stages of an explicit tableau are traced one after another instead. Each stage is kept as its own array, and the stages are stacked once, at the end, for the error estimate and the dense output: written into the buffer, as the stage loop does, they made the step 1.6x slower on GPU, since XLA rebuilds the whole buffer at every write once the indices are static. An FSAL method's first stage is a `lax.cond`, so that no step evaluates the vector field once more. Opt-in, since the vector field is then traced once per stage, which costs compile time. Implicit tableaus are refused. Co-Authored-By: Claude Opus 5.5 --- diffrax/_solver/runge_kutta.py | 106 ++++++++++++++++----- test/test_solver.py | 165 +++++++++++++++++++++++++++++++++ 2 files changed, 250 insertions(+), 21 deletions(-) diff --git a/diffrax/_solver/runge_kutta.py b/diffrax/_solver/runge_kutta.py index 9473ab44..2096eb1d 100644 --- a/diffrax/_solver/runge_kutta.py +++ b/diffrax/_solver/runge_kutta.py @@ -353,11 +353,23 @@ class AbstractRungeKutta(AbstractAdaptiveSolver[_SolverState]): instance of [`diffrax.CalculateJacobian`][]. """ - scan_kind: None | Literal["lax", "checkpointed", "bounded"] = None + scan_kind: None | Literal["lax", "checkpointed", "bounded", "unrolled"] = None tableau: AbstractClassVar[ButcherTableau | MultiButcherTableau] calculate_jacobian: AbstractClassVar[CalculateJacobian] + def __check_init__(self): + if self.scan_kind == "unrolled": + if isinstance(self.tableau, MultiButcherTableau): + implicit = any(tab.implicit for tab in self.tableau.tableaus) + else: + implicit = self.tableau.implicit + if implicit: + raise ValueError( + "`scan_kind='unrolled'` is only supported by explicit Runge--Kutta " + "methods." + ) + if TYPE_CHECKING: # Pretend that we're implicit root_finder: ClassVar[optx.AbstractRootFinder] @@ -695,6 +707,7 @@ def prod(f): k0 = prod(f0) else: # Non-FSAL solvers just iterate over all stages. + eval_first_stage = True f0 = _unused k0 = _unused init_stage_index = 0 @@ -1141,27 +1154,78 @@ def buffers(val): ks, RESULTS.successful, ) - const_result = const_result_sentinel = object() - # Needs to be an `eqxi.while_loop` as: - # (a) we may have variable length: e.g. an FSAL explicit RK scheme will have one - # more stage on the first step. - # (b) to work around a limitation of JAX's autodiff being unable to express - # "triangular computations" (every stage depends on all previous stages) - # without spurious copies. - final_val = eqxi.while_loop( - cond_stage, - rk_stage, - init_val, - max_steps=num_stages, - buffers=buffers, - kind="checkpointed" if self.scan_kind is None else self.scan_kind, - checkpoints=num_stages, - base=num_stages, - ) - _, y1, f1_for_fsal, _, _, fs, ks, result = final_val - assert const_result is not const_result_sentinel - if const_result: + + def unrolled_stages(): + # Explicit tableaus only (see `__check_init__`). The stages are traced one + # after another rather than in a `while_loop`: on GPU, every check of a + # while loop's predicate is a device-to-host copy that the host waits for, + # once per stage. The price is tracing the vector field once per stage. + # Each stage is kept as its own array rather than written into a buffer: + # with static indices, XLA rebuilds the whole buffer at every write. + ks_list = [] + yi, fi = y0, _unused + for stage_index in range(num_stages): + increments = [] + + def _increment(tab, *k_js): + if stage_index > 0: + for a_ij, k_j in zip(tab.a_lower[stage_index - 1], k_js): + if a_ij != 0: + increments.append(y_map(lambda k: float(a_ij) * k, k_j)) + + def _time(tab): + c_i = tab.c1 if stage_index == 0 else tab.c[stage_index - 1] + # No floating point error + return t1 if c_i == 1 else t0 + float(c_i) * dt + + jtu.tree_map(_increment, tableaus, *ks_list) + yi = y_map(_sum, y0, *increments) if increments else y0 + ti = jtu.tree_map(_time, tableaus) + if eval_fs: + if fsal and stage_index == 0: + # Only the first step, or a step after a jump, evaluates the + # first stage of an FSAL method. A `lax.cond` rather than a + # select, so that no step evaluates the vector field once more. + fi = lax.cond( + eval_first_stage, + lambda: vf(ti, yi, implicit_val=sentinel), + lambda: f0, + ) + else: + fi = vf(ti, yi, implicit_val=sentinel) + ki = prod(fi) + else: + fi = _unused + ki = vf_prod(ti, yi, implicit_val=sentinel) + ks_list.append(ki) + ks = jtu.tree_map(lambda *k: jnp.stack(k), *ks_list) + return yi, fi if fsal else _unused, ks + + if self.scan_kind == "unrolled": + y1, f1_for_fsal, ks = unrolled_stages() result = RESULTS.successful + else: + const_result = const_result_sentinel = object() + # Needs to be an `eqxi.while_loop` as: + # (a) we may have variable length: e.g. an FSAL explicit RK scheme will + # have one more stage on the first step. + # (b) to work around a limitation of JAX's autodiff being unable to + # express "triangular computations" (every stage depends on all + # previous stages) without spurious copies. + final_val = eqxi.while_loop( + cond_stage, + rk_stage, + init_val, + max_steps=num_stages, + buffers=buffers, + kind="checkpointed" if self.scan_kind is None else self.scan_kind, + checkpoints=num_stages, + base=num_stages, + ) + _, y1, f1_for_fsal, _, _, fs, ks, result = final_val + assert const_result is not const_result_sentinel + if const_result: + result = RESULTS.successful # # Calculate outputs: the final `y1` from our step, any dense information, etc. diff --git a/test/test_solver.py b/test/test_solver.py index dfff3f6b..b32694bb 100644 --- a/test/test_solver.py +++ b/test/test_solver.py @@ -502,3 +502,168 @@ def test_adaptive_dt0_milstein(getkey): diffrax.diffeqsolve( terms, solver, 0, 1, None, 1, stepsize_controller=stepsize_controller ) + + +_explicit_rk = ( + diffrax.Heun, + diffrax.Midpoint, + diffrax.Ralston, + diffrax.Bosh3, + diffrax.Tsit5, + diffrax.Dopri5, + diffrax.Dopri8, +) + + +def _pytree_vf(t, y, args): + x, v = y + return v * jnp.ones_like(x), -jnp.sum(args * jnp.sin(x)) + jnp.cos(t) + + +_pytree_y0 = (jnp.array([0.3, -0.2]), jnp.array(0.5)) +_pytree_args = jnp.array([1.5, 0.7]) + + +def _unrolled(solver): + return eqx.tree_at( + lambda s: s.scan_kind, solver, "unrolled", is_leaf=lambda x: x is None + ) + + +@pytest.mark.parametrize("solver_ctr", _explicit_rk) +@pytest.mark.parametrize("adaptive", (False, True)) +def test_unrolled_stages(solver_ctr, adaptive): + # `scan_kind="unrolled"` computes the same steps as the stage loop, with the same + # dense output and gradients. The jump makes an FSAL method re-evaluate its first + # stage partway through the solve. + if adaptive: + controller = diffrax.ClipStepSizeController( + diffrax.PIDController(rtol=1e-6, atol=1e-8), jump_ts=[0.45] + ) + else: + controller = diffrax.ConstantStepSize() + saveat = diffrax.SaveAt(ts=jnp.linspace(0, 1, 7)) + + def solve(solver, y0, args): + sol = diffrax.diffeqsolve( + diffrax.ODETerm(_pytree_vf), + solver, + 0, + 1, + 0.05, + y0, + args, + saveat=saveat, + stepsize_controller=controller, + ) + return sol.ys, sol.stats["num_steps"] + + def loss(solver, y0, args): + return sum(jnp.sum(y) for y in jtu.tree_leaves(solve(solver, y0, args)[0])) + + solver = solver_ctr() + ys, steps = solve(solver, _pytree_y0, _pytree_args) + ys_unrolled, steps_unrolled = solve(_unrolled(solver), _pytree_y0, _pytree_args) + assert steps == steps_unrolled + assert tree_allclose(ys, ys_unrolled, rtol=1e-10, atol=1e-12) + grads = jax.grad(loss, argnums=(1, 2))(solver, _pytree_y0, _pytree_args) + grads_unrolled = jax.grad(loss, argnums=(1, 2))( + _unrolled(solver), _pytree_y0, _pytree_args + ) + assert tree_allclose(grads, grads_unrolled, rtol=1e-10, atol=1e-12) + + +@pytest.mark.parametrize( + "adjoint", + (diffrax.ForwardMode(), diffrax.DirectAdjoint(), diffrax.BacksolveAdjoint()), +) +def test_unrolled_stages_adjoints(adjoint): + # `DirectAdjoint` and `ForwardMode` replace a `scan_kind` of `None`, but not + # "unrolled"; `BacksolveAdjoint` solves its adjoint equation with the same solver. + def final(solver, y0): + sol = diffrax.diffeqsolve( + diffrax.ODETerm(_pytree_vf), + solver, + 0, + 1, + 0.05, + y0, + _pytree_args, + stepsize_controller=diffrax.PIDController(rtol=1e-6, atol=1e-8), + adjoint=adjoint, + ) + return sum(jnp.sum(y) for y in jtu.tree_leaves(sol.ys)) + + def derivative(solver): + if isinstance(adjoint, diffrax.ForwardMode): + tangent = jtu.tree_map(jnp.ones_like, _pytree_y0) + return jax.jvp(lambda y0: final(solver, y0), (_pytree_y0,), (tangent,))[1] + return jax.grad(final, argnums=1)(solver, _pytree_y0) + + solver = diffrax.Tsit5() + assert tree_allclose( + derivative(solver), derivative(_unrolled(solver)), rtol=1e-8, atol=1e-10 + ) + + +def test_unrolled_stages_vmap(): + def solve(solver, y0): + sol = diffrax.diffeqsolve( + diffrax.ODETerm(_pytree_vf), + solver, + 0, + 1, + 0.05, + y0, + _pytree_args, + stepsize_controller=diffrax.PIDController(rtol=1e-6, atol=1e-8), + ) + return sol.ys + + y0s = ( + jnp.array([[0.3, -0.2], [0.1, 0.4], [-0.5, 0.0]]), + jnp.array([0.5, 0.0, 1.0]), + ) + solver = diffrax.Tsit5() + ys = jax.vmap(lambda y0: solve(solver, y0))(y0s) + ys_unrolled = jax.vmap(lambda y0: solve(_unrolled(solver), y0))(y0s) + assert tree_allclose(ys, ys_unrolled, rtol=1e-10, atol=1e-12) + + +def test_unrolled_stages_sde(getkey): + bm = diffrax.VirtualBrownianTree(0, 1, 1e-3, (2,), key=getkey()) + drift = diffrax.ODETerm(lambda t, y, args: -args * y) + diffusion = diffrax.ControlTerm(lambda t, y, args: 0.3 * jnp.diag(y), bm) + terms = diffrax.MultiTerm(drift, diffusion) + + def solve(solver): + sol = diffrax.diffeqsolve( + terms, solver, 0, 1, 0.01, jnp.array([1.0, 2.0]), _pytree_args + ) + return sol.ys + + solver = diffrax.Heun() + assert tree_allclose( + solve(solver), solve(_unrolled(solver)), rtol=1e-10, atol=1e-12 + ) + + +@pytest.mark.parametrize("solver_ctr", (diffrax.Tsit5, diffrax.Heun)) +def test_unrolled_stages_have_no_loop(solver_ctr): + # The point of `scan_kind="unrolled"`: on GPU, every check of a while loop's + # predicate is a device-to-host copy, once per stage. + solver = solver_ctr(scan_kind="unrolled") + term = diffrax.ODETerm(_pytree_vf) + state = solver.init(term, 0.0, 0.1, _pytree_y0, _pytree_args) + + def step(y0): + return solver.step(term, 0.0, 0.1, y0, _pytree_args, state, made_jump=False) + + assert "while" not in str(jax.make_jaxpr(step)(_pytree_y0)) + + +def test_unrolled_stages_explicit_only(): + with pytest.raises(ValueError, match="unrolled"): + diffrax.Kvaerno5(scan_kind="unrolled") + with pytest.raises(ValueError, match="unrolled"): + diffrax.KenCarp4(scan_kind="unrolled")