engine/arena.js

// Dumb-but-sufficient run-scoped GPU allocator. Encoder/decoder passes
// allocate their activations through an arena and destroy everything at once
// after readback / run end. No reuse, no sub-allocation — perf tasks replace
// this later.

// Requested-byte accounting across every arena (encoder scratch/retained,
// decode state, compaction). Weights, staging rings and uniform-pool banks are
// NOT arenas and are not counted. This is not an OS/Metal residency counter:
// `live` is JavaScript-owned arena storage, while `retired` conservatively
// keeps post-submit destroys charged until that submit drains.
// One module-wide ledger: a session drives one device, and tests that need
// isolation reset the peaks explicitly.
const stats = {
  allocBytes: 0,
  freedBytes: 0,
  liveBytes: 0,
  highWaterBytes: 0,
  liveBuffers: 0,
  retiredBytes: 0,
  retiredBuffers: 0,
  submittedHighWaterBytes: 0,
};

export function getArenaStats() {
  return {
    ...stats,
    submittedBytes: stats.liveBytes + stats.retiredBytes,
    submittedBuffers: stats.liveBuffers + stats.retiredBuffers,
  };
}

// Start a fresh peak window (e.g. per file batch): each high-water mark
// restarts from its CURRENT footprint, not from zero.
export function resetArenaPeak() {
  stats.highWaterBytes = stats.liveBytes;
  stats.submittedHighWaterBytes = stats.liveBytes + stats.retiredBytes;
}

function trackAlloc(size) {
  stats.allocBytes += size;
  stats.liveBytes += size;
  stats.liveBuffers += 1;
  if (stats.liveBytes > stats.highWaterBytes) stats.highWaterBytes = stats.liveBytes;
  const submitted = stats.liveBytes + stats.retiredBytes;
  if (submitted > stats.submittedHighWaterBytes) {
    stats.submittedHighWaterBytes = submitted;
  }
}

function trackFree(size) {
  stats.freedBytes += size;
  stats.liveBytes -= size;
  stats.liveBuffers -= 1;
}

function trackRetire(size) {
  stats.liveBytes -= size;
  stats.liveBuffers -= 1;
  stats.retiredBytes += size;
  stats.retiredBuffers += 1;
}

function trackDrained(size) {
  stats.retiredBytes -= size;
  stats.retiredBuffers -= 1;
  stats.freedBytes += size;
}

export function createArena(device) {
  const buffers = []; // [buffer, size] — size recorded here so destroy() never
  //                     depends on GPUBuffer.size existing (mock devices).
  return {
    // Plain allocation (byteLength is rounded up to a multiple of 4).
    buf(byteLength, usage, label) {
      const size = Math.ceil(byteLength / 4) * 4;
      const buffer = device.createBuffer({ label, size, usage });
      buffers.push([buffer, size]);
      trackAlloc(size);
      return buffer;
    },
    // Uniform buffer written once at creation from a TypedArray view.
    uniform(view, label) {
      const size = Math.max(16, Math.ceil(view.byteLength / 16) * 16);
      const buffer = device.createBuffer({
        label,
        size,
        usage: GPUBufferUsage.UNIFORM,
        mappedAtCreation: true,
      });
      new Uint8Array(buffer.getMappedRange())
        .set(new Uint8Array(view.buffer, view.byteOffset, view.byteLength));
      buffer.unmap();
      buffers.push([buffer, size]);
      trackAlloc(size);
      return buffer;
    },
    destroy() {
      for (const [b, size] of buffers) {
        b.destroy();
        trackFree(size);
      }
      buffers.length = 0;
    },
    // Post-submit retirement: invalidate the handles immediately, but keep
    // their requested bytes in the conservative submitted footprint until
    // the queue work that references them settles. The rejection arm releases
    // accounting too — a lost queue cannot keep useful submitted work alive.
    destroyDeferred(drained) {
      if (!drained || typeof drained.then !== 'function') {
        throw new Error('arena.destroyDeferred needs a queue-drain promise');
      }
      if (buffers.length === 0) return Promise.resolve();
      const retired = buffers.splice(0);
      for (const [b, size] of retired) {
        b.destroy();
        trackRetire(size);
      }
      const release = () => {
        for (const [, size] of retired) trackDrained(size);
      };
      return Promise.resolve(drained).then(release, release);
    },
  };
}

// Two ownership classes behind one compatible destroy handle. Encoder
// scratch can be retired immediately after submit while cross-KV/lens (and an
// optional inspectable output) remain alive for decode/readback. Both cleanup
// methods are idempotent because createArena.destroy() empties its list.
export function createSplitArena(device) {
  const retained = createArena(device);
  const scratch = createArena(device);
  return {
    retained,
    scratch,
    destroyScratch() { scratch.destroy(); },
    retireScratch(drained) { return scratch.destroyDeferred(drained); },
    destroy() {
      scratch.destroy();
      retained.destroy();
    },
  };
}