Explicit Runge-Kutta solvers: an opt-in scan_kind="unrolled", without the stage loop

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.

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")