@driftengine/core 3.63.0 → 4.1.0

package/src/render/inference/reuse.ts

@@ -0,0 +1,79 @@
+/**
+ * Which values of a graph may share one buffer, as one rule both sides plan by: the device's
+ * runner beside this file, and `@driftengine/texture`'s reference evaluator.
+ *
+ * **It lives here rather than beside that evaluator because of the direction of the arrow.**
+ * `texture` peers on core and core does not peer on `texture` — `deviceGraph.ts` says so and is
+ * built around it — so a rule both need has to sit on this side or be written twice. It was on the
+ * other side for a day, where it made core's runtime import a devDependency and
+ * `scripts/boundaries.test.mjs` refused it.
+ *
+ * **A value's buffer is given back after the last node that reads it, and not before** — so a
+ * node's output never shares a buffer with any of its own inputs, since the inputs are read while
+ * the output is written. What is `held` keeps a buffer of its own for the life of the plan, and so
+ * does every output; a value that is neither written by a node nor held is not planned at all,
+ * which is how the evaluator leaves weights and inputs in the caller's own arrays.
+ *
+ * **Two implementations of this would drift**, and the failure is a network that runs, validates
+ * and answers wrongly — a value overwritten before its last reader reads it.
+ *
+ * What it gives up: a free buffer is taken first-come rather than by the size that fits best, so a
+ * small value may hold a large buffer while a larger one waits. The transformers this runs repeat a
+ * block of identically shaped values, where first-come and best-fit choose the same buffers.
+ */
+
+export interface BufferReuse {
+  /** The buffer each planned value lives in. */
+  readonly slotOf: ReadonlyMap<string, number>;
+  /** Each buffer's size, in values: the largest value it ever holds. */
+  readonly sizes: readonly number[];
+}
+
+export function planReuse(
+  graph: {
+    readonly nodes: readonly { readonly inputs: readonly string[]; readonly output: string }[];
+    /** Values given a buffer each before any node runs, in this order. */
+    readonly held: readonly string[];
+    /** Values never given back, claimed where they are written. */
+    readonly outputs: readonly string[];
+  },
+  sizeOf: (name: string) => number,
+): BufferReuse {
+  const slotOf = new Map<string, number>();
+  const sizes: number[] = [];
+  const pinned = new Set<string>([...graph.held, ...graph.outputs]);
+  const claim = (name: string): number => {
+    sizes.push(sizeOf(name));
+    slotOf.set(name, sizes.length - 1);
+    return sizes.length - 1;
+  };
+  for (const name of graph.held) claim(name);
+
+  /* The last node reading each value; a value never read is last read where it is written. */
+  const lastRead = new Map<string, number>();
+  graph.nodes.forEach((node, at) => {
+    lastRead.set(node.output, at);
+    for (const input of node.inputs) lastRead.set(input, at);
+  });
+
+  const free: number[] = [];
+  graph.nodes.forEach((node, at) => {
+    if (pinned.has(node.output)) {
+      claim(node.output);
+    } else {
+      const reuse = free.shift();
+      if (reuse === undefined) claim(node.output);
+      else {
+        slotOf.set(node.output, reuse);
+        sizes[reuse] = Math.max(sizes[reuse] as number, sizeOf(node.output));
+      }
+    }
+    /* Only after the node has run: its inputs are read while its output is written. */
+    for (const value of new Set([...node.inputs, node.output])) {
+      const slot = slotOf.get(value);
+      if (slot === undefined || pinned.has(value) || lastRead.get(value) !== at) continue;
+      free.push(slot);
+    }
+  });
+  return { slotOf, sizes };
+}

package/src/render/inference/runner.ts

@@ -0,0 +1,204 @@
+/**
+ * A graph on a device: its buffers, its pipelines, and one compute pass that runs it.
+ *
+ * **A binder over `schedule.ts`**, which decides everything and is tested without a device; what is
+ * left here is creating what the schedule names and recording its dispatches. One pass holds the
+ * whole graph, because WebGPU orders a pass's dispatches and makes each one's writes visible to the
+ * next — a network is a chain of dependent dispatches and needs nothing more.
+ *
+ * **Creation is asynchronous, and it is where anything wrong is reported**: a value larger than the
+ * device can bind is refused by name before a buffer exists, and every pipeline is built with
+ * `createComputePipelineAsync`, which rejects on a kernel that does not compile. A runner that was
+ * created runs. Identical kernels share one pipeline — a transformer repeats one block's shapes
+ * once per layer, so its pipelines number its distinct shapes rather than its nodes.
+ *
+ * `encode` records the graph into a caller's encoder, so a network can run inside a frame or under
+ * a timestamp query; `run` is the whole round trip, and one is in flight at a time, because its
+ * readback buffers are the runner's and a second run would map them while the first holds them.
+ * **A run the device refuses is refused too**: WebGPU reports a validation failure at `submit`,
+ * and a refused command buffer runs nothing, so without the scope around it the readback would
+ * map and return whatever the buffers last held — a plausible answer to a question never run.
+ */
+import { shapeSize, type DeviceGraph } from './deviceGraph.ts';
+import type { InferenceDevice } from './device.ts';
+import { scheduleGraph } from './schedule.ts';
+
+/* The usage and stage flags by value: a host with no browser has no `GPUBufferUsage` global. */
+const MAP_READ = 0x0001;
+const COPY_SRC = 0x0004;
+const COPY_DST = 0x0008;
+const STORAGE = 0x0080;
+const COMPUTE = 0x4;
+const MODE_READ = 0x0001;
+
+export interface GraphRunner {
+  /** Records every dispatch of the graph into `encoder`, in one compute pass. */
+  encode(encoder: GPUCommandEncoder, timestampWrites?: GPUComputePassTimestampWrites): void;
+  /** Writes the inputs, runs the graph and reads its outputs back, each by name. */
+  run(inputs: ReadonlyMap<string, Float32Array>): Promise<Map<string, Float32Array>>;
+  dispose(): void;
+}
+
+export async function createGraphRunner(
+  inference: InferenceDevice,
+  graph: DeviceGraph,
+): Promise<GraphRunner> {
+  const { device } = inference;
+  const schedule = scheduleGraph(graph, { half: inference.half });
+  const namesOf = new Map<number, string[]>();
+  for (const [name, slot] of schedule.plan.slotOf) {
+    namesOf.set(slot, [...(namesOf.get(slot) ?? []), name]);
+  }
+  const limit = device.limits.maxStorageBufferBindingSize;
+  const oversized = schedule.bytes.flatMap((bytes, slot) =>
+    bytes > limit ? [`${(namesOf.get(slot) ?? []).join('/')} (${bytes} bytes)`] : [],
+  );
+  if (oversized.length > 0) {
+    throw new Error(
+      `${oversized.join(', ')} ${oversized.length === 1 ? 'is' : 'are'} larger than the ` +
+        `${limit} bytes this device can bind as one buffer`,
+    );
+  }
+
+  device.pushErrorScope('validation');
+  const buffers = schedule.bytes.map((size, slot) =>
+    device.createBuffer({
+      label: `inference ${(namesOf.get(slot) ?? []).join('/')}`,
+      size,
+      usage: STORAGE | COPY_DST | COPY_SRC,
+    }),
+  );
+  for (const upload of schedule.uploads) {
+    device.queue.writeBuffer(buffers[upload.slot] as GPUBuffer, 0, padded(upload.data));
+  }
+
+  const layouts = new Map<number, GPUBindGroupLayout>();
+  const layoutFor = (inputs: number): GPUBindGroupLayout => {
+    let layout = layouts.get(inputs);
+    if (layout === undefined) {
+      const entries: GPUBindGroupLayoutEntry[] = [];
+      for (let i = 0; i <= inputs; i += 1) {
+        entries.push({
+          binding: i,
+          visibility: COMPUTE,
+          buffer: { type: i < inputs ? 'read-only-storage' : 'storage' },
+        });
+      }
+      layout = device.createBindGroupLayout({ label: `inference ${inputs} in`, entries });
+      layouts.set(inputs, layout);
+    }
+    return layout;
+  };
+  const pipelines = new Map<string, Promise<GPUComputePipeline>>();
+  const steps = await Promise.all(
+    schedule.steps.map(async (step) => {
+      const layout = layoutFor(step.slots.length - 1);
+      let pipeline = pipelines.get(step.code);
+      if (pipeline === undefined) {
+        pipeline = device.createComputePipelineAsync({
+          label: `inference ${step.label}`,
+          layout: device.createPipelineLayout({ bindGroupLayouts: [layout] }),
+          compute: {
+            module: device.createShaderModule({ label: step.label, code: step.code }),
+            entryPoint: 'main',
+          },
+        });
+        pipelines.set(step.code, pipeline);
+      }
+      const group = device.createBindGroup({
+        label: `inference ${step.label}`,
+        layout,
+        entries: step.slots.map((slot, binding) => ({
+          binding,
+          resource: { buffer: buffers[slot] as GPUBuffer },
+        })),
+      });
+      return { pipeline: await pipeline, group, workgroups: step.workgroups };
+    }),
+  );
+  const staging = graph.outputs.map((name) => {
+    const bytes = Math.max(4, shapeSize(graph.shapes.get(name) ?? []) * 4);
+    return {
+      name,
+      source: buffers[schedule.plan.slotOf.get(name) as number] as GPUBuffer,
+      bytes,
+      buffer: device.createBuffer({
+        label: `inference readback ${name}`,
+        size: bytes,
+        usage: MAP_READ | COPY_DST,
+      }),
+    };
+  });
+  const invalid = await device.popErrorScope();
+  if (invalid !== null)
+    throw new Error(`the device refused a network's buffers: ${invalid.message}`);
+
+  const encode = (encoder: GPUCommandEncoder, timestampWrites?: GPUComputePassTimestampWrites) => {
+    const pass = encoder.beginComputePass(timestampWrites === undefined ? {} : { timestampWrites });
+    for (const step of steps) {
+      pass.setPipeline(step.pipeline);
+      pass.setBindGroup(0, step.group);
+      pass.dispatchWorkgroups(step.workgroups[0], step.workgroups[1], step.workgroups[2]);
+    }
+    pass.end();
+  };
+
+  let running = false;
+  const run = async (inputs: ReadonlyMap<string, Float32Array>) => {
+    if (running) throw new Error('a network runner runs one graph at a time; await the last run');
+    for (const input of graph.inputs) {
+      const values = inputs.get(input.name);
+      const expected = shapeSize(input.shape);
+      if (values === undefined || values.length !== expected) {
+        throw new RangeError(
+          `input "${input.name}" needs ${expected} values and was given ${values?.length ?? 'none'}`,
+        );
+      }
+    }
+    running = true;
+    try {
+      for (const input of graph.inputs) {
+        const slot = schedule.plan.slotOf.get(input.name) as number;
+        device.queue.writeBuffer(
+          buffers[slot] as GPUBuffer,
+          0,
+          inputs.get(input.name) as Float32Array,
+        );
+      }
+      device.pushErrorScope('validation');
+      const encoder = device.createCommandEncoder({ label: 'inference run' });
+      encode(encoder);
+      for (const out of staging)
+        encoder.copyBufferToBuffer(out.source, 0, out.buffer, 0, out.bytes);
+      device.queue.submit([encoder.finish()]);
+      const refused = await device.popErrorScope();
+      if (refused !== null)
+        throw new Error(`the device refused a network's run: ${refused.message}`);
+      await Promise.all(staging.map((out) => out.buffer.mapAsync(MODE_READ)));
+      const results = new Map<string, Float32Array>();
+      for (const out of staging) {
+        const values = new Float32Array(out.buffer.getMappedRange().slice(0));
+        const count = shapeSize(graph.shapes.get(out.name) ?? []);
+        results.set(out.name, count === values.length ? values : values.subarray(0, count));
+        out.buffer.unmap();
+      }
+      return results;
+    } finally {
+      running = false;
+    }
+  };
+
+  const dispose = () => {
+    for (const buffer of buffers) buffer.destroy();
+    for (const out of staging) out.buffer.destroy();
+  };
+  return { encode, run, dispose };
+}
+
+/* A queue write is sized in fours, so an odd count of halves is written with one of padding. */
+function padded(data: Float32Array | Uint16Array): Float32Array | Uint16Array {
+  if (data.byteLength % 4 === 0) return data;
+  const out = new Uint16Array(data.length + 1);
+  out.set(data);
+  return out;
+}

package/src/render/inference/schedule.ts

@@ -0,0 +1,80 @@
+/**
+ * What a graph becomes before a device sees it: its dispatches, its buffers, and its weights.
+ *
+ * **Everything a runner decides is decided here, with no device**, so the decisions have tests and
+ * the runner is a binder: GELUs folded into their multiplies, a buffer for every value from the
+ * planner, one kernel per node generated at the shapes the node has, and the weights narrowed to
+ * half precision when that was asked for. A dispatch binds its inputs' buffers in order and then
+ * its output's, which is the binding order every kernel declares.
+ *
+ * **Half precision narrows the weights and nothing else.** A weight is stored as the nearest half
+ * and widened as it is read; the values between nodes stay single precision, because they are what
+ * a network's error accumulates through and a weight is read-only and read many times. What it
+ * gives up is half the memory of the activations as well — the step to take if a model's
+ * activations turn out to be what does not fit.
+ */
+import { shapeSize, type DeviceGraph } from './deviceGraph.ts';
+import { fuseActivations } from './fuse.ts';
+import { toHalfFloats } from '../halfFloat.ts';
+import { DEVICE_KERNELS, type ElementType } from './kernels.ts';
+import { planBuffers, type BufferPlan } from './plan.ts';
+
+export interface ScheduledStep {
+  /** The operator and the value it writes, which is what a device error names. */
+  readonly label: string;
+  readonly code: string;
+  readonly workgroups: readonly [number, number, number];
+  /** The buffers bound, in binding order: each input's, then the output's. */
+  readonly slots: readonly number[];
+}
+
+export interface Schedule {
+  readonly steps: readonly ScheduledStep[];
+  readonly plan: BufferPlan;
+  /** Each buffer's size in bytes, rounded up to the four a storage binding is sized in. */
+  readonly bytes: readonly number[];
+  /** The weights, each into its buffer once, as the precision asked for stores them. */
+  readonly uploads: readonly { readonly slot: number; readonly data: Float32Array | Uint16Array }[];
+}
+
+export function scheduleGraph(graph: DeviceGraph, options: { readonly half: boolean }): Schedule {
+  const missing = graph.nodes.filter((node) => !DEVICE_KERNELS.has(node.op));
+  if (missing.length > 0) {
+    const names = [...new Set(missing.map((node) => node.op))].join(', ');
+    throw new Error(`the device has no kernel for ${names}; the graph cannot run here`);
+  }
+  const fused = fuseActivations(graph);
+  const plan = planBuffers(fused);
+  const narrow = new Set(options.half ? fused.tensors.map((tensor) => tensor.name) : []);
+  const typeOf = (name: string): ElementType => (narrow.has(name) ? 'f16' : 'f32');
+  const shapeOf = (name: string): readonly number[] => fused.shapes.get(name) ?? [];
+
+  const bytes = plan.sizes.map((size) => Math.max(4, size * 4));
+  for (const name of narrow) {
+    const slot = plan.slotOf.get(name) as number;
+    bytes[slot] = Math.max(4, Math.ceil((shapeSize(shapeOf(name)) * 2) / 4) * 4);
+  }
+
+  const steps = fused.nodes.map((node): ScheduledStep => {
+    const generate = DEVICE_KERNELS.get(node.op);
+    if (generate === undefined) throw new Error(`the device has no kernel for ${node.op}`);
+    const kernel = generate({
+      inputShapes: node.inputs.map(shapeOf),
+      outputShape: shapeOf(node.output),
+      attributes: node.attributes,
+      inputTypes: node.inputs.map(typeOf),
+    });
+    return {
+      label: `${node.op} ${node.output}`,
+      code: kernel.code,
+      workgroups: kernel.workgroups,
+      slots: [...node.inputs, node.output].map((name) => plan.slotOf.get(name) as number),
+    };
+  });
+
+  const uploads = fused.tensors.map((tensor) => ({
+    slot: plan.slotOf.get(tensor.name) as number,
+    data: narrow.has(tensor.name) ? toHalfFloats(tensor.data) : tensor.data,
+  }));
+  return { steps, plan, bytes, uploads };
+}

package/src/render/inference/shapeKernels.ts

@@ -0,0 +1,112 @@
+/**
+ * The shape operators on the device: they move values and compute nothing, or write a zero.
+ *
+ * **Each is a copy, one invocation per output element**, with the index arithmetic baked from the
+ * shapes. `reshape` is a copy too rather than an alias of its input's buffer: aliasing would make
+ * two values share a buffer and the planner's rule — a buffer is reused once nothing reads what it
+ * holds — would have to follow aliases, which is a second rule for one saving the size of a copy.
+ */
+import { type KernelGenerator, list, num, perElement, read, size } from './kernelKit.ts';
+
+/* Strides of a row-major shape. */
+function strides(shape: readonly number[]): number[] {
+  const out = shape.map(() => 1);
+  for (let d = shape.length - 2; d >= 0; d -= 1)
+    out[d] = (out[d + 1] as number) * (shape[d + 1] as number);
+  return out;
+}
+
+const permute: KernelGenerator = (request) => {
+  const input = request.inputShapes[0] as readonly number[];
+  const order = list(request.attributes, 'order');
+  const from = strides(input);
+  const outStrides = strides(request.outputShape);
+  /* Output axis d walks input axis order[d]: its index times that axis's input stride. */
+  const terms = order.map(
+    (axis, d) =>
+      `((i / ${outStrides[d]}u) % ${request.outputShape[d]}u) * ${from[axis as number]}u`,
+  );
+  return perElement(
+    request,
+    size(request.outputShape),
+    `  output[i] = ${read(request, 0, terms.join(' + ') || '0u')};`,
+  );
+};
+
+const reshape: KernelGenerator = (request) =>
+  perElement(request, size(request.outputShape), `  output[i] = ${read(request, 0, 'i')};`);
+
+const concat: KernelGenerator = (request) => {
+  const axis = num(request.attributes, 'axis');
+  const out = request.outputShape;
+  const inner = size(out.slice(axis + 1));
+  const along = out[axis] as number;
+  /* Walk the inputs along the axis: each owns a run of it, starting where the one before ended. */
+  let start = 0;
+  const branches = request.inputShapes.map((shape, n) => {
+    const length = shape[axis] as number;
+    const branch = `if (a < ${start + length}u) { output[i] = ${read(request, n, `(outer * ${length}u + a - ${start}u) * ${inner}u + rest`)}; return; }`;
+    start += length;
+    return branch;
+  });
+  return perElement(
+    request,
+    size(out),
+    `  let rest = i % ${inner}u;
+  let a = (i / ${inner}u) % ${along}u;
+  let outer = i / ${inner * along}u;
+  ${branches.join('\n  ')}`,
+  );
+};
+
+const slice: KernelGenerator = (request) => {
+  const axis = num(request.attributes, 'axis');
+  const begin = num(request.attributes, 'start');
+  const input = request.inputShapes[0] as readonly number[];
+  const inner = size(input.slice(axis + 1));
+  const length = input[axis] as number;
+  const taken = request.outputShape[axis] as number;
+  return perElement(
+    request,
+    size(request.outputShape),
+    `  let rest = i % ${inner}u;
+  let a = (i / ${inner}u) % ${taken}u;
+  let outer = i / ${inner * taken}u;
+  output[i] = ${read(request, 0, `(outer * ${length}u + a + ${begin}u) * ${inner}u + rest`)};`,
+  );
+};
+
+/* An output element past the input's end on any axis is a zero; the rest read their own index. */
+const pad: KernelGenerator = (request) => {
+  const input = request.inputShapes[0] as readonly number[];
+  const from = strides(input);
+  const to = strides(request.outputShape);
+  const at = (axis: number): string => `((i / ${to[axis]}u) % ${request.outputShape[axis]}u)`;
+  const inside = input.map((d, axis) => `${at(axis)} < ${d}u`).join(' && ');
+  const index = input.map((_, axis) => `${at(axis)} * ${from[axis]}u`).join(' + ');
+  return perElement(
+    request,
+    size(request.outputShape),
+    `  if (${inside}) { output[i] = ${read(request, 0, index)}; } else { output[i] = 0.0; }`,
+  );
+};
+
+/* Row `indices[r]` of the table, clamped to the last row, since a device cannot refuse an index. */
+const gather: KernelGenerator = (request) => {
+  const [rows, width] = request.inputShapes[0] as readonly number[] as [number, number];
+  return perElement(
+    request,
+    size(request.outputShape),
+    `  let row = min(u32(${read(request, 1, `i / ${width}u`)}), ${rows - 1}u);
+  output[i] = ${read(request, 0, `row * ${width}u + i % ${width}u`)};`,
+  );
+};
+
+export const SHAPE_KERNELS: readonly (readonly [string, KernelGenerator])[] = [
+  ['permute', permute],
+  ['reshape', reshape],
+  ['concat', concat],
+  ['slice', slice],
+  ['pad', pad],
+  ['gather', gather],
+];