fix(core): free every pooled, dummy and scratch buffer on executor dispose

packages/runntime/src/core/gpu/bufferPool.ts

@@ -12,6 +12,9 @@ export class BufferPool {
   private free = new Map<string, GpuBufferRef[]>();
   /** Bucket keys acquired since the last trim(). */
   private used = new Set<string>();
+  /** Buffers acquired and not released yet, so dispose() frees the ones a
+   *  captured frame or a caller still holds. */
+  private inUse = new Set<GpuBufferRef>();
 
   constructor(private readonly createBuffer: (shape: ValueMeta) => GpuBufferRef) {}
 
@@ -22,11 +25,13 @@ export class BufferPool {
   acquire(shape: ValueMeta): GpuBufferRef {
     const k = this.key(shape);
     this.used.add(k);
-    const reused = this.free.get(k)?.pop();
-    return reused ?? this.createBuffer(shape);
+    const buffer = this.free.get(k)?.pop() ?? this.createBuffer(shape);
+    this.inUse.add(buffer);
+    return buffer;
   }
 
   release(buffer: GpuBufferRef, shape: ValueMeta): void {
+    this.inUse.delete(buffer);
     const k = this.key(shape);
     const bucket = this.free.get(k);
     if (bucket) bucket.push(buffer);
@@ -47,9 +52,12 @@ export class BufferPool {
     return destroyed;
   }
 
+  /** Destroys every free buffer and every buffer still in use. */
   dispose(): void {
     for (const bucket of this.free.values()) for (const buf of bucket) buf.destroy();
+    for (const buf of this.inUse) buf.destroy();
     this.free.clear();
+    this.inUse.clear();
     this.used.clear();
   }
 }

packages/runntime/src/core/gpu/eval.ts

@@ -105,7 +105,7 @@ export type { KernelTiming };
  *  without walking the graph again. Fixed-shape forwards only. */
 export interface CaptureExecutor {
   /** Runs `fn` once, recording every dispatch. The executor is replay-only
-   *  afterwards. Keep `targets` — their buffers stay valid across replays — and
+   *  afterwards, until dispose(), which also ends the replays. Keep `targets` — their buffers stay valid across replays — and
    *  overwrite inputs in place before each replay. */
   captureFrame<T>(fn: () => T): {
     targets: T;
@@ -151,6 +151,9 @@ export function gpuExecutor(root: TgpuRoot, opts: { perf?: GpuPerfSink } = {}):
   };
   /** Memo key to the handles a dispatch encoded. */
   const handleCache = new Map<string, KernelHandle[]>();
+  /** Dummies and scratches, destroyed with the executor. A replaced scratch
+   *  stays here: memoized handles may still bind it. */
+  const owned: { destroy(): void }[] = [];
 
   // Filler for bind slots a pipeline never reads. One per dtype, 4 elements,
   // since vec4 layouts need 16 bytes minimum.
@@ -161,17 +164,26 @@ export function gpuExecutor(root: TgpuRoot, opts: { perf?: GpuPerfSink } = {}):
   function dummy(dtype: FloatDtype): FloatBuffer;
   function dummy(dtype: FloatDtype): FloatBuffer {
     if (dtype === 'f16') {
-      dummyF16 ??= root.createBuffer(d.arrayOf(d.f16, 4)).$usage('storage');
+      if (dummyF16 === undefined) {
+        dummyF16 = root.createBuffer(d.arrayOf(d.f16, 4)).$usage('storage');
+        owned.push(dummyF16);
+      }
       return dummyF16;
     }
-    dummyF32 ??= root.createBuffer(d.arrayOf(d.f32, 4)).$usage('storage');
+    if (dummyF32 === undefined) {
+      dummyF32 = root.createBuffer(d.arrayOf(d.f32, 4)).$usage('storage');
+      owned.push(dummyF32);
+    }
     return dummyF32;
   }
 
   // Replay-only once captured: a fresh eval would acquire pooled buffers
   // aliasing the captured frame's intermediates.
   let capturing: KernelHandle[] | undefined;
   let frozen = false;
+  /** Bumped by dispose(), so a replay recorded before it throws instead of
+   *  binding destroyed buffers. */
+  let disposals = 0;
 
   const readbacks = new ReadbackQueue(root);
   const scratches = new Map<string, { buf: F32Buffer; elems: number }>();
@@ -185,6 +197,7 @@ export function gpuExecutor(root: TgpuRoot, opts: { perf?: GpuPerfSink } = {}):
         checkBindable(root, 4 * elems, `${key} scratch`);
         s = { buf: root.createBuffer(d.arrayOf(d.f32, elems)).$usage('storage'), elems };
         scratches.set(key, s);
+        owned.push(s.buf);
       }
       return s.buf;
     },
@@ -226,11 +239,24 @@ export function gpuExecutor(root: TgpuRoot, opts: { perf?: GpuPerfSink } = {}):
       capturing = undefined;
       frozen = true;
       if (recorded.length === 0) throw new Error('captureFrame: fn() submitted no GPU work');
+      const recordedAt = disposals;
+      const live = () => {
+        if (disposals !== recordedAt) throw new Error('captureFrame: replay after dispose()');
+      };
       return {
         targets,
-        replay: () => readbacks.submit((encoder) => encodePass(encoder, recorded)),
-        encodeReplay: (encoder) => encodePass(encoder, recorded),
-        profileReplay: () => profileHandles(root.device, recorded),
+        replay: () => {
+          live();
+          readbacks.submit((encoder) => encodePass(encoder, recorded));
+        },
+        encodeReplay: (encoder) => {
+          live();
+          encodePass(encoder, recorded);
+        },
+        profileReplay: () => {
+          live();
+          return profileHandles(root.device, recorded);
+        },
       };
     },
     acquireOutput: (shape) => {
@@ -242,7 +268,19 @@ export function gpuExecutor(root: TgpuRoot, opts: { perf?: GpuPerfSink } = {}):
       return pool.acquire(shape);
     },
     releaseOutput: (buffer, shape) => pool.release(buffer, shape),
-    dispose: () => pool.dispose(),
+    // Forgets every handle, dummy and scratch it destroys, so a later
+    // dispatch makes fresh ones instead of binding a destroyed buffer. A
+    // captured executor can eval and capture again; its old replays throw.
+    dispose: () => {
+      frozen = false;
+      disposals++;
+      pool.dispose();
+      for (const b of owned.splice(0)) b.destroy();
+      handleCache.clear();
+      scratches.clear();
+      dummyF32 = undefined;
+      dummyF16 = undefined;
+    },
     trim: () => {
       // Memoized handles bind the destroyed buffers.
       if (!frozen && pool.trim()) handleCache.clear();

packages/runntime/test/core/gpu/bufferPool.test.ts

@@ -74,9 +74,21 @@ describe('BufferPool', () => {
     expect(made.map((m) => m.destroyed)).toEqual([1, 1]);
   });
 
-  // dispose() walks the free list only, so a buffer still out (a captured
-  // target, a result nobody released) is never destroyed.
-  it.fails('dispose destroys the buffers still out', () => {
+  it('trim then dispose destroys each buffer once', () => {
+    const { pool, made } = stubPool();
+    const a = pool.acquire(f32(16));
+    const b = pool.acquire(f32(32));
+    pool.release(a, f32(16));
+    pool.release(b, f32(32));
+    pool.trim();
+    // 32 goes unused, so the next trim destroys it; 16 is taken and held.
+    pool.acquire(f32(16));
+    pool.trim();
+    pool.dispose();
+    expect(made.map((m) => m.destroyed)).toEqual([1, 1]);
+  });
+
+  it('dispose destroys the buffers still out', () => {
     const { pool, made } = stubPool();
     pool.acquire(f32(16));
     pool.dispose();

packages/runntime/test/core/gpu/eval.gpu.test.ts

@@ -72,12 +72,8 @@ interface GraphCase {
   leftAfterCapture?: string[];
 }
 
-/** A 16 B dummy the executor binds for a missing optional input. */
-const DUMMY = '16 B, usage 140';
 /** A uniform buffer a kernel makes itself instead of through cachedUniform. */
 const uniform = (bytes: number) => `${bytes} B, usage 76`;
-/** The [ROWS, COLS] f32 result a captured frame holds. */
-const TARGET = '512 B, usage 140';
 /** The 3-row gatherRows index list. */
 const INDEX_LIST = '12 B, usage 140';
 
@@ -97,105 +93,89 @@ const GRAPHS: GraphCase[] = [
   {
     name: 'an elementwise chain',
     build: ({ x }) => tanh(add(mul(x, x), 1)),
-    // dispose() keeps the captured target, which never went back to the pool.
-    leftAfterCapture: [TARGET],
   },
   {
     name: 'layerNorm without a bias',
     build: ({ x, w }) => layerNorm(add(x, 1), w),
-    // dispose() keeps the dummy bias, the layerNorm uniform and a captured target.
-    leftOnDispose: [DUMMY, uniform(16)],
-    leftAfterCapture: [DUMMY, TARGET, uniform(16)],
+    // dispose() keeps the layerNorm uniform.
+    leftOnDispose: [uniform(16)],
+    leftAfterCapture: [uniform(16)],
   },
   {
     name: 'softmax',
     build: ({ x }) => softmax(add(x, 1)),
-    // dispose() keeps the dummy segments buffer and a captured target.
-    leftOnDispose: [DUMMY],
-    leftAfterCapture: [DUMMY, TARGET],
   },
   {
     name: 'matmul',
     build: ({ x, m }) => tanh(matmul(x, m)),
-    leftAfterCapture: [TARGET],
   },
   {
     name: 'gatherRows',
     build: ({ x }) => gatherRows(add(x, 1), [3, 0, 3]),
     // gatherRows uploads its index list on every dispatch and never destroys it.
     leftPerRun: [INDEX_LIST],
     leftOnDispose: [INDEX_LIST],
-    // The captured target is the 3 gathered rows.
-    leftAfterCapture: [INDEX_LIST, '192 B, usage 140'],
+    leftAfterCapture: [INDEX_LIST],
   },
   {
     name: 'flash attention',
     needs: 'subgroups',
     // flash: qLen > 1, subgroups, headDim % 16 == 0
     build: ({ x, m }) => sdpa(x, m, m, FULL_ATTENTION),
-    // dispose() keeps the dummy sinks and segments buffer and a captured target.
-    leftOnDispose: [DUMMY],
-    leftAfterCapture: [DUMMY, TARGET],
   },
   {
     name: 'rows split attention',
     needs: 'subgroups',
     // rowsSubgroupSplit: qLen > 1, f32, subgroups, headDim % 16 != 0
     build: ({ x, m }) => sdpa(x, m, m, TWO_HEAD_ATTENTION),
-    // dispose() keeps the dummy, the 640 B split scratch, the split and merge
-    // uniforms and a captured target.
-    leftOnDispose: [DUMMY, '640 B, usage 140', uniform(16), uniform(8)],
-    leftAfterCapture: [DUMMY, '640 B, usage 140', uniform(16), TARGET, uniform(8)],
+    // dispose() keeps the split and merge uniforms.
+    leftOnDispose: [uniform(16), uniform(8)],
+    leftAfterCapture: [uniform(16), uniform(8)],
   },
   {
     name: 'decode attention',
     // decode: qLen == 1, kvLen <= ATTN_DECODE_CHUNK
     build: ({ w, m }) => sdpa(w, m, m, FULL_ATTENTION),
-    // dispose() keeps the dummy, the attnDecode uniform and a captured target,
-    // the one query row.
-    leftOnDispose: [DUMMY, uniform(12)],
-    leftAfterCapture: [DUMMY, '64 B, usage 140', uniform(12)],
+    // dispose() keeps the attnDecode uniform.
+    leftOnDispose: [uniform(12)],
+    leftAfterCapture: [uniform(12)],
   },
   {
     name: 'conv1d',
     build: ({ x, m }) => conv1d(x, m, { kernelSize: 1, stride: 1 }),
-    // dispose() keeps the conv1d uniform and a captured target.
+    // dispose() keeps the conv1d uniform.
     leftOnDispose: [uniform(8)],
-    leftAfterCapture: [TARGET, uniform(8)],
+    leftAfterCapture: [uniform(8)],
   },
   {
     name: 'swigluChunk',
     build: ({ x }) => swigluChunk(add(x, 1)),
-    // dispose() keeps the swigluChunk uniform and a captured target, half as
-    // wide as x.
+    // dispose() keeps the swigluChunk uniform.
     leftOnDispose: [uniform(8)],
-    leftAfterCapture: ['256 B, usage 140', uniform(8)],
+    leftAfterCapture: [uniform(8)],
   },
   {
     name: 'argmaxDot',
     build: ({ w, m }) => argmaxDot(m, w),
-    // dispose() keeps the 1024 B partials scratch, the argmaxDot uniform and a
-    // captured target, the one index.
-    leftOnDispose: ['1024 B, usage 140', uniform(4)],
-    leftAfterCapture: ['1024 B, usage 140', uniform(4), '4 B, usage 140'],
+    // dispose() keeps the argmaxDot uniform.
+    leftOnDispose: [uniform(4)],
+    leftAfterCapture: [uniform(4)],
   },
   {
     name: 'gatherRowsFrom',
     build: ({ x, idx }) => gatherRowsFrom(add(x, 1), idx),
-    // dispose() keeps the gatherRowsFrom uniform and a captured target, the 3
-    // gathered rows.
+    // dispose() keeps the gatherRowsFrom uniform.
     leftOnDispose: [uniform(8)],
-    leftAfterCapture: ['192 B, usage 140', uniform(8)],
+    leftAfterCapture: [uniform(8)],
   },
   {
     name: 'writeRows',
     // Writes into x in place, as a KV-cache append does, then reads it: a
     // released writeRows result would hand x's buffer to the pool.
     build: ({ x, w }) => tanh(writeRows(x, add(w, 1), 2)),
-    // dispose() keeps the writeRows uniform and a captured target, the view of
-    // rows 0 to 3.
+    // dispose() keeps the writeRows uniform.
     leftOnDispose: [uniform(8)],
-    leftAfterCapture: [uniform(8), '192 B, usage 140'],
+    leftAfterCapture: [uniform(8)],
   },
 ];
 
@@ -247,6 +227,36 @@ describe.for(GRAPHS)('gpuExecutor with $name', (graph) => {
     expect(buffersSince(before)).toEqual(graph.leftOnDispose ?? []);
   });
 
+  it('runs again after dispose', async () => {
+    const values = inputs();
+    const ex = gpuExecutor(root);
+    try {
+      await run(ex, graph.build(values));
+      ex.dispose();
+      await run(ex, graph.build(values));
+    } finally {
+      ex.dispose();
+    }
+  });
+
+  it('runs again after a captured frame is disposed', async () => {
+    const values = inputs();
+    const ex = gpuExecutor(root);
+    try {
+      const { replay } = ex.captureFrame(() => {
+        const out = graph.build(values);
+        evalValues([out], ex);
+        return out;
+      });
+      replay();
+      ex.dispose();
+      expect(replay).toThrow('replay after dispose()');
+      await run(ex, graph.build(values));
+    } finally {
+      ex.dispose();
+    }
+  });
+
   it('frees a captured frame on dispose', async (ctx) => {
     if (graph.needs) requireFeature(ctx, root, graph.needs);
     const values = inputs();

packages/runntime/test/core/gpu/replayCache.gpu.test.ts

@@ -23,10 +23,6 @@ function cache() {
   });
 }
 
-/** What dispose() leaves, a known leak: each bucket's executor.dispose()
- *  keeps its dummy segments buffer, for buckets 4 and 8. */
-const LEFT_ON_DISPOSE = ['16 B, usage 140', '16 B, usage 140'];
-
 const fill = (T: number, scratch: { x: Float32Array }) => scratch.x.set(randomF32(T * COLS, T));
 
 describe('createReplayCache', () => {
@@ -53,6 +49,6 @@ describe('createReplayCache', () => {
     const replays = cache();
     await replays.warmAll();
     replays.dispose();
-    expect(buffersSince(before)).toEqual(LEFT_ON_DISPOSE);
+    expect(buffersSince(before)).toEqual([]);
   });
 });

packages/runntime/test/zoo/leaks.gpu.test.ts

@@ -101,11 +101,10 @@ const CASES: LeakCase[] = [
         },
       };
     },
-    // The bucket executor's dispose() keeps its dummy and the kernels'
-    // uniforms (usage 76).
+    // dispose() keeps the kernels' uniforms.
     leftOnDispose: [
-      ...['8 B, usage 76', '8 B, usage 76', '16 B, usage 76', '16 B, usage 140'],
-      ...Array<string>(4).fill('16 B, usage 76'),
+      ...Array<string>(2).fill('8 B, usage 76'),
+      ...Array<string>(5).fill('16 B, usage 76'),
     ],
   },
   {
@@ -155,9 +154,7 @@ const CASES: LeakCase[] = [
         },
       };
     },
-    // executor.dispose() keeps its 8 B f16 dummy, and toArray keeps its
-    // 10-class result out of the pool.
-    leftOnDispose: ['8 B, usage 140', '40 B, usage 140'],
+    // toArray keeps its 10-class result out of the pool.
     leftPerCall: ['40 B, usage 140'],
   },
 ];