Remove Extra Copy Buffer in _integrate.loop

benchmarks/integrate_memory_usage.py

@@ -0,0 +1,123 @@
+"""does the t0 == t1 identity cond in diffrax's loop() cost a
+copy of the save buffer?
+
+Toy Workload: B trajectories of an n-dim linear ODE y' = A y, Tsit5 + PID.
+Reports XLA temp memory against the size of one batched save buffer
+(B * n_ts * n * 8 bytes for --saveat ts, B * max_steps * n * 8 bytes for --saveat steps)
+
+How t0 reaches the cond (--t0):
+    const    vmap over y0, t0 a Python constant: XLA can fold t0 == t1 and drop the cond
+    shared   vmap over y0, t0 a traced scalar shared by all trajectories: the inner cond
+             predicate is unbatched, so it stays a real cond
+    batched  vmap over (y0, t0): the inner cond predicate is batched, so vmap turns it
+             into a select over the whole save state
+    novmap   no vmap, the batch is solved as one (B, n) ODE with a traced scalar t0
+
+Compare temp memory between cases with the same --saveat and --consume flags, not
+against an absolute ratio: other loop state also lands in temp.
+
+    python repro_t0t1.py [--t0 {const,shared,batched,novmap}] [--saveat {ts,steps}]
+                         [--consume] [--B B] [--n N] [--n-ts N_TS]
+
+"""
+
+import argparse
+import json
+
+import jax
+import jax.numpy as jnp
+
+
+jax.config.update("jax_enable_x64", True)
+
+import diffrax
+
+
+parser = argparse.ArgumentParser(
+    description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
+)
+parser.add_argument(
+    "--consume", action="store_true", help="use sol.ys downstream of the solve"
+)
+parser.add_argument(
+    "--t0",
+    choices=["const", "shared", "batched", "novmap"],
+    default="const",
+    help="how t0 is passed in (see above)",
+)
+parser.add_argument(
+    "--saveat", choices=["ts", "steps"], default="ts", help="SaveAt(ts=...) or steps"
+)
+parser.add_argument("--B", type=int, default=1000, help="number of trajectories")
+parser.add_argument("--n", type=int, default=100, help="state dimension")
+parser.add_argument("--n-ts", type=int, default=1000, help="number of save points")
+cli = parser.parse_args()
+
+consume = cli.consume
+B, n, n_ts = cli.B, cli.n, cli.n_ts
+max_steps = 4096
+
+A = -jnp.eye(n) + 0.1 * jax.random.normal(jax.random.PRNGKey(0), (n, n)) / jnp.sqrt(n)
+ts = jnp.linspace(0.0, 10.0, n_ts)
+if cli.saveat == "ts":
+    saveat = diffrax.SaveAt(ts=ts)
+else:
+    saveat = diffrax.SaveAt(steps=True)
+
+
+def solve(y0, t0):
+    sol = diffrax.diffeqsolve(
+        diffrax.ODETerm(lambda t, y, args: y @ A.T),  # y is (n,) or, for novmap, (B, n)
+        diffrax.Tsit5(),
+        t0=t0,
+        t1=10.0,
+        dt0=0.01,
+        y0=y0,
+        saveat=saveat,
+        stepsize_controller=diffrax.PIDController(rtol=1e-8, atol=1e-8),
+        max_steps=max_steps,
+    )
+    ys = sol.ys
+    assert ys is not None
+    if consume:  # use the saved ys downstream to prevent XLA optimizing it away
+        first_half = jnp.arange(len(ys)) < len(ys) // 2
+        mask = first_half.reshape((-1,) + (1,) * (ys.ndim - 1))
+        ys = jnp.where(mask, ys, 2.0 * ys)
+    return ys
+
+
+y0 = jax.random.normal(jax.random.PRNGKey(1), (B, n))
+t0 = jnp.array(0.0)
+if cli.t0 == "const":
+    f, args = jax.jit(jax.vmap(lambda y: solve(y, 0.0))), (y0,)
+elif cli.t0 == "shared":
+    f, args = jax.jit(jax.vmap(solve, in_axes=(0, None))), (y0, t0)
+elif cli.t0 == "batched":
+    f, args = jax.jit(jax.vmap(solve)), (y0, jnp.zeros(B))
+else:
+    f, args = jax.jit(solve), (y0, t0)
+m = f.lower(*args).compile().memory_analysis()
+assert m is not None
+buf = B * (n_ts if cli.saveat == "ts" else max_steps) * n * 8
+ys = f(*args)  # also run it. Answer shouldnt change
+print(
+    json.dumps(
+        dict(
+            consume=consume,
+            t0=cli.t0,
+            saveat=cli.saveat,
+            B=B,
+            n=n,
+            n_ts=n_ts,
+            device=str(jax.devices()[0]),
+            save_buffer_GB=round(buf / 1e9, 3),
+            temp_GB=round(m.temp_size_in_bytes / 1e9, 3),
+            out_GB=round(m.output_size_in_bytes / 1e9, 3),
+            temp_in_save_buffers=round(m.temp_size_in_bytes / buf, 2),
+            # dirty check to make sure output doesnt degrade; steps pads unused rows
+            # with inf, so only sum the finite entries
+            checksum=float(jnp.sum(jnp.where(jnp.isfinite(ys), ys, 0.0))),
+        )
+    ),
+    flush=True,
+)

diffrax/_integrate.py

@@ -820,28 +820,34 @@ def unsave(subsaveat: SubSaveAt, save_state: SaveState) -> SaveState:
             unsave, saveat.subs, save_state, is_leaf=_is_subsaveat
         )
 
-    def _save_if_t0_equals_t1(subsaveat: SubSaveAt, save_state: SaveState) -> SaveState:
-        if subsaveat.ts is not None:
-            save_state = _save(
-                t0, yfinal, args, subsaveat.fn, save_state, repeat=len(subsaveat.ts)
-            )
-        return save_state
-
     # if t0 == t1 then we don't enter the integration loop. In this case we have to
     # manually update the saved ts and ys if we want to save at "intermediate"
     # times specified by saveat.subs.ts
-    save_state = jax.lax.cond(
-        eqxi.unvmap_any(t0 == t1),
-        lambda __save_state: jax.lax.cond(
-            t0 == t1,
-            lambda _save_state: jtu.tree_map(
-                _save_if_t0_equals_t1, saveat.subs, _save_state, is_leaf=_is_subsaveat
-            ),
-            lambda _save_state: _save_state,
-            __save_state,
-        ),
-        lambda __save_state: __save_state,
-        save_state,
+    #
+    # For subsaveat states, only 't0' row has been written. Hence, the 'ts' rows start
+    # at a static offset. Instead of lax.cond, we can statically calculate what slice
+    # of the state to update. This coaxes XLA optimizer to not emit an extra copy buffer
+    # of the full difeq state needlessly
+    #
+    def _save_if_t0_equals_t1(subsaveat: SubSaveAt, save_state: SaveState) -> SaveState:
+        if subsaveat.ts is None:
+            return save_state
+        pred = t0 == t1
+        start = 1 if subsaveat.t0 else 0
+        stop = start + len(subsaveat.ts)
+
+        def _update(buffer, value):
+            return buffer.at[start:stop].set(jnp.where(pred, value, buffer[start:stop]))
+
+        ts = _update(save_state.ts, t0)
+        ys = jtu.tree_map(_update, save_state.ys, subsaveat.fn(t0, yfinal, args))
+        save_index = jnp.where(pred, stop, save_state.save_index)
+        return eqx.tree_at(
+            lambda s: [s.ts, s.ys, s.save_index], save_state, [ts, ys, save_index]
+        )
+
+    save_state = jtu.tree_map(
+        _save_if_t0_equals_t1, saveat.subs, save_state, is_leaf=_is_subsaveat
     )
 
     def _save_t1(subsaveat, save_state):

test/test_saveat_solution.py

@@ -428,6 +428,48 @@ def _solve(tf):
     assert tree_allclose(yc, true_yc)
 
 
+@pytest.mark.parametrize("batch_t0", [False, True])
+def test_t0_eq_t1_no_save_buffer_copy(batch_t0):
+    """Handling t0 == t1 must not cost an extra copy of the save buffer.
+    Compares XLA's planned temporary memory against a constant t0, for which the
+    t0 == t1 check is constant-folded away. `throw=False` keeps the `error_if` on the
+    solution, which can cost a copy of its own on GPU, out of the comparison.
+    """
+    num_batch, num_ts, size = 16, 200, 16
+    ts = jnp.linspace(0.0, 1.0, num_ts)
+
+    def _solve(y0, t0):
+        sol = diffrax.diffeqsolve(
+            diffrax.ODETerm(lambda t, y, args: -y),
+            diffrax.Tsit5(),
+            t0=t0,
+            t1=1.0,
+            dt0=0.1,
+            y0=y0,
+            saveat=diffrax.SaveAt(ts=ts),
+            # On CPU, the old cond-based handling only cost a copy with an adaptive
+            # controller; with a constant step size XLA happened to fold it.
+            stepsize_controller=diffrax.PIDController(rtol=1e-6, atol=1e-6),
+            throw=False,
+        )
+        return sol.ys
+
+    def _temp_bytes(fn, *args):
+        analysis = jax.jit(fn).lower(*args).compile().memory_analysis()
+        if analysis is None:
+            pytest.skip("memory analysis is unavailable on this backend")
+        return analysis.temp_size_in_bytes
+
+    y0 = jnp.ones((num_batch, size))
+    constant = _temp_bytes(jax.vmap(lambda y: _solve(y, 0.0)), y0)
+    if batch_t0:
+        traced = _temp_bytes(jax.vmap(_solve), y0, jnp.zeros(num_batch))
+    else:
+        traced = _temp_bytes(jax.vmap(_solve, in_axes=(0, None)), y0, jnp.array(0.0))
+    save_buffer = num_batch * num_ts * size * y0.dtype.itemsize
+    assert traced - constant < save_buffer / 2
+
+
 def test_trivial_dense():
     term = diffrax.ODETerm(lambda t, y, args: -0.5 * y)
     y0 = jnp.array([2.1])