gpu: the mirror table and its buffer pool live once, in gpu_abi.h

cuda.h and webgpu.h each kept their own copy of the same shape: a
struct pairing a host buffer with a device handle and a residency
state, a map from the native handle to it, and a size-keyed free list
alloc()/release() recycle through. Only the device handle's type
differed (CUdeviceptr, wgpu::Buffer). gpu::mirror_table<Handle>
(gpu_abi.h) holds that shape once; each backend instantiates it for
its own handle type and keeps doing its own copying (before_kernel_,
sync_to_host), which is where the two drivers actually diverge.

mirror_table::insert moves its Handle parameter into the entry rather
than copying it, so a ref-counted handle (wgpu::Buffer) picks up only
one AddRef per allocation instead of two.

Verified on native Metal (117 cases, 11640 assertions with real
device access, plus check_qwen's greedy tokens against the numpy
oracle) and gpu_host's no-GPU reference build, under WebGPU via Deno,
and by a clean CUDA host-side trace diff against the pre-refactor
commit (no kernel call changed). A mutation that let take() hand out
an already-live pooled buffer twice was caught hard by the WebGPU
suite (32 of 117 cases failed).

Author: yhirose

docs/backends.md

@@ -60,6 +60,15 @@ buffer, so an output into a buffer whose live bytes are the host's has to bring
 them up or lose the rest, while an output into a fresh buffer, which is nearly
 every output, has nothing to bring.
 
+The table that holds those copies — host pointer, device handle, size,
+`residency` — and the size-keyed free list `alloc`/`release` recycle buffers
+through are `gpu::mirror_table<Handle>` (`gpu_abi.h`), one instantiation per
+mirrored backend (`mirror_table<CUdeviceptr>` in `cuda.h`,
+`mirror_table<wgpu::Buffer>` in `webgpu.h`). Only `Handle` — the backend's own
+device buffer type — varies; the map, the pool and the copying policy do not.
+A backend still does the actual copy (`before_kernel_`, `sync_to_host`) since
+that call differs by driver.
+
 ## The kernel ABI
 
 An op hands the backend a kernel id, an ordered list of views, a params
@@ -259,9 +268,6 @@ asks; a new backend edits no test.
 
 ## What is not shared yet
 
-- The mirror table (handle to host copy, device copy, size, `residency`) and
-  the buffer pool exist twice, in `cuda.h` and `webgpu.h`. The state machine
-  itself is shared.
 - WebGPU's kernels still predate the canonical ABI and go through `marshal_`
   (see above).
 

include/cuda.h

@@ -282,29 +282,19 @@ struct context {
   std::vector<timed_launch> timed;
   std::vector<CUevent> spare_events;
 
-  // Host/device mirror per allocation, keyed by the device pointer (== the
-  // `native` handle stored in storage). Views sharing a storage share the key,
-  // so one state serves every view. When to copy is gpu::residency's decision
-  // (gpu_abi.h); the copies are made here.
-  struct mirror {
-    float* host = nullptr;  // CPU-side buffer (storage.contents/ptr)
-    CUdeviceptr dev = 0;    // device buffer (storage.native)
-    size_t bytes = 0;
-    gpu::residency live;
-  };
-  std::unordered_map<CUdeviceptr, mirror> mirrors;
-
-  // Size-keyed free list (like Metal's MTLBuffer pool). Released buffers are
-  // recycled, not cuMemFree'd — repeated large alloc/free otherwise fragments
-  // the driver allocator (decode benches, training that churns activations).
-  // Buffers persist until the (leaked) context tears down. Keyed by exact byte
-  // size; the workloads that churn reuse identical shapes.
-  std::unordered_map<size_t, std::vector<std::pair<CUdeviceptr, float*>>> pool;
-
-  mirror* mirror_(void* native) {
-    auto it = mirrors.find(reinterpret_cast<CUdeviceptr>(native));
-    return it == mirrors.end() ? nullptr : &it->second;
-  }
+  // Host/device mirror per allocation and its size-keyed free list (shared
+  // shape with webgpu.h; `gpu::mirror_table`, gpu_abi.h). Keyed by the device
+  // pointer, reinterpreted as the `native` handle stored in storage. Views
+  // sharing a storage share the key, so one state serves every view. When to
+  // copy is gpu::residency's decision; the copies are made here. Released
+  // buffers are recycled, not cuMemFree'd — repeated large alloc/free
+  // otherwise fragments the driver allocator (decode benches, training that
+  // churns activations); they persist until the (leaked) context tears down.
+  using mirror = gpu::mirror_table<CUdeviceptr>::entry;
+  gpu::mirror_table<CUdeviceptr> mt;
+
+  mirror* mirror_(void* native) { return mt.find(native); }
+
   // A kernel is about to touch this buffer as `a`: bring the host copy up if
   // residency says so. Async on the stream like the meta uploads (a blocking
   // copy would wait out every kernel already queued and stall the pipeline
@@ -1141,35 +1131,26 @@ inline void* alloc(int64_t bytes, float** contents, bool host_fill = false) {
   size_t nb = bytes > 0 ? (size_t)bytes : 4;
   CUdeviceptr dev = 0;
   float* host = nullptr;
-  auto it = c.pool.find(nb);  // reuse a recycled buffer of this exact size
-  if (it != c.pool.end() && !it->second.empty()) {
-    dev = it->second.back().first;
-    host = it->second.back().second;
-    it->second.pop_back();
-  } else {
+  if (!c.mt.take(nb, dev, host)) {  // no recycled buffer of this exact size
     if (c.d.MemAlloc(&dev, nb) != 0) return nullptr;
     host = static_cast<float*>(std::malloc(nb));
     if (!host) {
       c.d.MemFree(dev);
       return nullptr;
     }
   }
-  c.mirrors[dev] = context::mirror{host, dev, nb, gpu::residency(host_fill)};
+  void* native = reinterpret_cast<void*>(dev);
+  c.mt.insert(native, dev, host, nb, host_fill);
   if (contents) *contents = host;
-  return reinterpret_cast<void*>(dev);
+  return native;
 }
 
 inline void release(void* buf, int64_t, float*) {
   auto& c = context::get();
   if (!c.ready || !buf) return;
-  CUdeviceptr dev = reinterpret_cast<CUdeviceptr>(buf);
-  auto it = c.mirrors.find(dev);
-  if (it == c.mirrors.end()) {
-    c.d.MemFree(dev);  // untracked (shouldn't happen); free outright
-    return;
+  if (!c.mt.release(buf)) {
+    c.d.MemFree(reinterpret_cast<CUdeviceptr>(buf));  // untracked (shouldn't happen); free outright
   }
-  c.pool[it->second.bytes].push_back({dev, it->second.host});  // recycle
-  c.mirrors.erase(it);
 }
 
 // Reconcile a buffer for a CPU access: when the device holds the live copy,

include/gpu_abi.h

@@ -19,6 +19,9 @@
 #include <cstddef>
 #include <cstdint>
 #include <string_view>
+#include <unordered_map>
+#include <utility>
+#include <vector>
 
 #include "profile.h"
 
@@ -160,6 +163,54 @@ struct residency {
   void uploaded() { where = both; }
 };
 
+// The host/device mirror and size-keyed buffer pool a mirrored backend (CUDA,
+// WebGPU) keeps per allocation. `Handle` is the backend's own device buffer
+// type (`CUdeviceptr`, `wgpu::Buffer`); the map, the pool and the entry shape
+// are otherwise identical between them, so only `Handle` varies. A backend
+// still does its own copying (`before_kernel_`, `sync_to_host`) and still
+// keys lookups by whatever `void*` it hands out as `native`; this only owns
+// where the state lives.
+template <class Handle>
+struct mirror_table {
+  struct entry {
+    float* host = nullptr;
+    Handle dev{};
+    size_t bytes = 0;
+    residency live;
+  };
+  std::unordered_map<void*, entry> mirrors;
+  std::unordered_map<size_t, std::vector<std::pair<Handle, float*>>> pool;
+
+  entry* find(void* native) {
+    auto it = mirrors.find(native);
+    return it == mirrors.end() ? nullptr : &it->second;
+  }
+
+  // A released buffer of this exact size, if the pool has one.
+  bool take(size_t bytes, Handle& dev, float*& host) {
+    auto it = pool.find(bytes);
+    if (it == pool.end() || it->second.empty()) return false;
+    dev = it->second.back().first;
+    host = it->second.back().second;
+    it->second.pop_back();
+    return true;
+  }
+
+  void insert(void* key, Handle dev, float* host, size_t bytes, bool host_fill) {
+    mirrors[key] = entry{host, std::move(dev), bytes, residency(host_fill)};
+  }
+
+  // Moves `native`'s entry into the free list, keyed by its size. False if
+  // `native` is not tracked (nothing to release).
+  bool release(void* native) {
+    auto it = mirrors.find(native);
+    if (it == mirrors.end()) return false;
+    pool[it->second.bytes].push_back({it->second.dev, it->second.host});
+    mirrors.erase(it);
+    return true;
+  }
+};
+
 // A launch's extent: how many groups, how many threads in each, and the bytes
 // of per-group scratch a kernel's reduction needs where the backend sizes it at
 // launch (CUDA's shared memory; Metal and WGSL size theirs in the kernel).

include/webgpu.h

@@ -196,21 +196,16 @@ struct context {
   void* meta_ring_(float** host_out);
   uint32_t meta_reserve_slot_();
 
-  // Host/device mirror per allocation, keyed by the opaque handle alloc()
-  // returns as `native`. Views sharing a storage share the key, so one dirty
-  // state serves every view. `where` tracks which copy is live.
-  struct mirror {
-    float* host = nullptr;   // CPU-side buffer (storage.ptr)
-    wgpu::Buffer dev;        // device buffer
-    size_t bytes = 0;
-    gpu::residency live;  // when to copy (gpu_abi.h); the copies are made here
-  };
-  std::unordered_map<void*, mirror> mirrors;
-
-  // Size-keyed free lists (like Metal's MTLBuffer pool and CUDA's). Repeated
-  // alloc/free of identical shapes is the common case, and per-dispatch
-  // allocation would compound the fixed dispatch floor.
-  std::unordered_map<size_t, std::vector<std::pair<wgpu::Buffer, float*>>> pool;
+  // Host/device mirror per allocation and its size-keyed free list (shared
+  // shape with cuda.h; `gpu::mirror_table`, gpu_abi.h), keyed by the opaque
+  // handle alloc() returns as `native`. Views sharing a storage share the
+  // key, so one dirty state serves every view.
+  using mirror = gpu::mirror_table<wgpu::Buffer>::entry;
+  gpu::mirror_table<wgpu::Buffer> mt;
+
+  // Mapped-for-readback staging buffers (sync_to_host's D2H), a distinct pool
+  // from the mirror table's device buffers: these are never bound as a
+  // kernel operand, only mapped.
   std::unordered_map<size_t, std::vector<wgpu::Buffer>> staging_pool;
 
   // Compute pipelines, keyed by WGSL entry point. Every one is built in the
@@ -347,10 +342,7 @@ struct context {
     return instance.WaitAny(f, UINT64_MAX) == wgpu::WaitStatus::Success;
   }
 
-  mirror* mirror_(void* native) {
-    auto it = mirrors.find(native);
-    return it == mirrors.end() ? nullptr : &it->second;
-  }
+  mirror* mirror_(void* native) { return mt.find(native); }
 
   // A kernel is about to touch this buffer as `a`: bring the host copy up if
   // residency says so.
@@ -512,12 +504,7 @@ inline void* alloc(int64_t bytes, float** contents, bool host_fill = false) {
 
   wgpu::Buffer dev;
   float* host = nullptr;
-  auto it = c.pool.find(nb);  // reuse a recycled buffer of this exact size
-  if (it != c.pool.end() && !it->second.empty()) {
-    dev = it->second.back().first;
-    host = it->second.back().second;
-    it->second.pop_back();
-  } else {
+  if (!c.mt.take(nb, dev, host)) {  // no recycled buffer of this exact size
     wgpu::BufferDescriptor d = {};
     d.size = nb;
     d.usage = wgpu::BufferUsage::Storage | wgpu::BufferUsage::CopyDst |
@@ -532,18 +519,15 @@ inline void* alloc(int64_t bytes, float** contents, bool host_fill = false) {
   // host pointer is both, and malloc will not hand out the same address twice
   // while it is live.
   void* token = host;
-  c.mirrors[token] = context::mirror{host, dev, nb, gpu::residency(host_fill)};
+  c.mt.insert(token, dev, host, nb, host_fill);
   if (contents) *contents = host;
   return token;
 }
 
 inline void release(void* buf, int64_t, float*) {
   auto& c = context::get();
   if (!c.ready || !buf) return;
-  auto it = c.mirrors.find(buf);
-  if (it == c.mirrors.end()) return;
-  c.pool[it->second.bytes].push_back({it->second.dev, it->second.host});
-  c.mirrors.erase(it);
+  c.mt.release(buf);
 }
 
 // Reconcile a buffer for a CPU access: flush pending kernels, then D2H if the