fix: release a plan's buffers after their last reader

run() held every kernel's output to the end-of-step barrier, so a step's
footprint was the sum of everything the graph computes rather than what
is live between two kernels, and the pool then kept that high-water mark
for the process. The plan now records each position's last reader and
run() drops the buffer there, which lets the pool hand the bytes to a
later output in the same run. The reuse is safe on the asynchronous
backend because kernels, memsets and commits all go down one in-order
stream.

Author: jiffies64

limn/jit.py

@@ -84,6 +84,7 @@ class Plan:
     buffers: tuple[int, ...]  # the BUFFER positions, whose bytes come from this step's graph
     sinks: tuple[int, ...]  # realized() resolved to positions, so alias chains are walked once
     positions: int  # how many nodes the graph had: the size of a run's buffer table
+    releases: tuple[tuple[int, ...], ...]  # per call: the positions whose bytes no later call reads
 
 
 @dataclass(frozen=True, slots=True)
@@ -164,19 +165,29 @@ def run(self, plan: Plan, sources: Mapping[int, Buffer]) -> list[Buffer]:
         repointed at the fresh value so sinks read what was just written rather than a copy the
         commit never touched. execute() and capture.replay both come through here, so the
         transaction rule has one owner.
+
+        A position is dropped right after its last reader's call: the run's peak footprint is
+        what is live between two kernels, not the sum of everything the graph computes, which
+        for a training step is most of a backward pass held to the barrier for no reason. The
+        reuse is safe on an asynchronous backend because kernels, memsets and commits all go
+        down one in-order stream: whatever still reads the old bytes was enqueued first. A
+        deferred value stays alive in `deferred` regardless, and a source buffer's bytes belong
+        to its tensor, so dropping the reference here frees at most prepare()'s migrated copy.
         """
         bufs: list[Buffer] = [None] * plan.positions
         for p, buf in sources.items():
             bufs[p] = self.prepare(buf)
         deferred: list[tuple[int, Buffer]] = []
-        for call in plan.calls:
+        for call, drop in zip(plan.calls, plan.releases, strict=True):
             outs = [self.out_alloc(nb, call.zero_fill) for nb in call.out_nbytes]
             call.fn([bufs[p] for p in call.inputs], outs)
             for p, out in zip(call.outputs, outs, strict=True):
                 if p is not None:  # a multi-output kernel writes this one, but nothing here reads it
                     bufs[p] = out
             if call.assign_target is not None:
                 deferred.append((call.assign_target, outs[0]))
+            for p in drop:
+                bufs[p] = None
         for p, value in deferred:
             self.commit(sources[p], value)
             bufs[p] = value
@@ -234,9 +245,19 @@ def plan_of(self, order: list[Node], position: dict[Node, int], sinks: list[Node
                     assign_target=position[kernel.target] if root.op is Op.ASSIGN else None,
                 )
             )
+        # each position's last reader, so run() can drop it there; a sink is read after the
+        # last call, and an assign value is deferred until the commit, but `deferred` holds
+        # its own reference, so only the sinks need shielding from release
+        sunk = tuple(position[realized(sink)] for sink in sinks)
+        last_read = {p: k for k, call in enumerate(calls) for p in call.inputs}
+        released: list[list[int]] = [[] for _ in calls]
+        for p, k in last_read.items():
+            if p not in sunk:
+                released[k].append(p)
         return Plan(
             tuple(calls),
             tuple(p for p, node in enumerate(order) if node.op is Op.BUFFER),
-            tuple(position[realized(sink)] for sink in sinks),
+            sunk,
             len(order),
+            tuple(map(tuple, released)),
         )

tests/test_backend_cuda.py

@@ -166,3 +166,20 @@ def test_the_pool_recycles_dropped_buffers():
     dev.trim()
     assert not dev.pool
     assert dev._alloc(1024).ptr != 0
+
+
+def test_a_run_releases_each_link_of_a_chain_as_its_reader_runs():
+    """A chain of n kernels must run in the memory of a couple of links, not all of them.
+
+    Each gather cuts a kernel, so this is n calls each reading only the link before it.
+    run() drops a buffer after its last reader and the pool hands the bytes to the next
+    output, so two link-sized allocations circulate for the whole chain; without the
+    release, every link stays live to the barrier and the pool ends holding all n."""
+    dev = CudaDevice()
+    rows = Tensor(np.arange(1 << 18, dtype=np.int32))
+    y = Tensor(np.ones((1 << 18, 4), dtype=np.float32))
+    for _ in range(8):
+        y = (y + 1.0)[rows]
+    out = dev.copyout(dev.execute([y.node])[0]).view(np.float32)
+    np.testing.assert_array_equal(out, np.full(1 << 20, 9.0, dtype=np.float32))
+    assert len(dev.pool.get(4 << 20, [])) <= 4

tests/test_jit.py

@@ -257,6 +257,22 @@ def test_an_output_nothing_reads_still_gets_its_buffer():
     assert runs == ["pair"]
 
 
+@needs_cc
+def test_a_plan_releases_positions_after_their_last_reader():
+    set_device("c")
+    rows = Tensor(np.arange(8, dtype=np.int32))
+    y = Tensor(randf(8, 2))
+    for _ in range(3):
+        y = (y + 1.0)[rows]  # each gather cuts a kernel, so the chain is one call per link
+    y.numpy()
+    dev = device.active()
+    assert isinstance(dev, CompiledDevice)
+    plan = next(iter(dev.plans.values()))
+    outs = [call.outputs[0] for call in plan.calls]
+    assert outs[0] in plan.releases[1] and outs[1] in plan.releases[2]
+    assert not any(s in drop for drop in plan.releases for s in plan.sinks)
+
+
 def test_on_an_interpreting_device_the_function_just_runs():
     data = batches(4)
     counter = {"ran": 0}