oidn-webgpu 0.0.0-stage → 0.6.0

package/src/native/denoise.ts

@@ -0,0 +1,268 @@
+import {
+  BufferUsage,
+  type OIDNGPUBindGroupEntry,
+  type OIDNGPUBuffer,
+  type OIDNGPUDevice
+} from '../gpuTypes.js';
+import { hdrForwardWGSL, hdrInverseWGSL, type HDRTransfer } from '../hdrTransfer.js';
+import { sanitizeWGSL } from '../sanitize.js';
+import type { ValidatedUNetModel } from '../modelSpec.js';
+import type { NativeUNetCore, NativeUNetPrecision } from '../nativeUNetCore.js';
+import { OIDN_TILE_ALIGNMENT, planTileGrid } from '../tileScheduler.js';
+import type { RecordedBuffer, RecordingDevice } from './recordingDevice.js';
+
+export interface NativeDenoiseOptions {
+  width: number;
+  height: number;
+  hdr: boolean;
+  transfer: HDRTransfer;
+  /** Maximum output tile edge; 0 denoises the whole image as one tile. */
+  maxTileSize: number;
+  inputScale: number;
+}
+
+const WORKGROUP = 8;
+
+function alignUp(value: number, alignment: number) {
+  return Math.ceil(value / alignment) * alignment;
+}
+
+/**
+ * Input processing of one tile for buffer images: clamp-to-edge reads of the
+ * tile's input rectangle, HDR transfer, and conversion directly into the
+ * network's packed tensor. Auxiliary RGB channels cross vec4 boundaries;
+ * unused lanes must be zero just as in createInputPackShader.
+ */
+function inputShader(sourceCount: number, hdr: boolean, transfer: HDRTransfer, precision: NativeUNetPrecision) {
+  const names = ['color', 'albedo', 'normal'].slice(0, sourceCount);
+  const bindings = names.map((name, index) => `
+@group(0) @binding(${index}) var<storage, read> in_${name}: array<vec4f>;`).join('');
+  // Like OIDN, callers pass signed [-1, 1] normals; the network reads n * 0.5 + 0.5.
+  const copies = names.slice(1).map((name) => name === 'normal'
+    ? 'let normal = sanitize(in_normal[inIdx].rgb, -1.0, 1.0) * 0.5 + 0.5;'
+    : `let ${name} = sanitize(in_${name}[inIdx].rgb, 0.0, 1.0);`).join('\n  ');
+  const channels = names.flatMap((name) => ['r', 'g', 'b'].map((channel) => `${name}.${channel}`));
+  const blocks = Math.ceil(channels.length / 4);
+  const type = precision === 'fp16' ? 'vec4<f16>' : 'vec4<f32>';
+  const stores = Array.from({ length: blocks }, (_, block) =>
+    `packed[outIdx * ${blocks} + ${block}] = ${type}(vec4f(${[0, 1, 2, 3].map((lane) => channels[block * 4 + lane] ?? '0.0').join(', ')}));`).join('\n  ');
+  return /* wgsl */ `${precision === 'fp16' ? 'enable f16;\n' : ''}${hdrForwardWGSL(transfer)}${sanitizeWGSL}
+struct Params {
+  imageSize: vec2i,
+  tileSize: vec2i,
+  offset: vec2i,
+
+}${bindings}
+@group(0) @binding(${sourceCount}) var<storage, read_write> packed: array<${type}>;
+@group(0) @binding(${sourceCount + 1}) var<uniform> params: Params;
+${hdr ? `struct Frame { inputScale: f32, }
+@group(0) @binding(${sourceCount + 2}) var<uniform> frame: Frame;` : ''}
+
+@compute @workgroup_size(${WORKGROUP}, ${WORKGROUP}, 1)
+fn main(@builtin(global_invocation_id) globalId: vec3u) {
+  let position = vec2i(globalId.xy);
+  if (position.x >= params.tileSize.x || position.y >= params.tileSize.y) {
+    return;
+  }
+  let inputPosition = clamp(position + params.offset, vec2i(0), params.imageSize - vec2i(1));
+  let inIdx = inputPosition.y * params.imageSize.x + inputPosition.x;
+  let outIdx = position.y * params.tileSize.x + position.x;
+  let col = in_color[inIdx];
+  ${hdr
+    ? `let scaled = sanitize(col.rgb * frame.inputScale, 0.0, FLT_MAX);
+  let color = vec3f(HDRForward(scaled.r), HDRForward(scaled.g), HDRForward(scaled.b));`
+    : 'let color = sanitize(col.rgb, 0.0, 1.0);'}
+  ${copies}
+  ${stores}
+}
+`;
+}
+
+/**
+ * Output processing of one tile: the inverse HDR transfer of the network
+ * output inside the tile's output rectangle, written with the input color's
+ * alpha. Matches the browser runtime's GPU output pass.
+ */
+function outputShader(hdr: boolean, transfer: HDRTransfer) {
+  return /* wgsl */ `${hdrInverseWGSL(transfer)}${sanitizeWGSL}
+struct Params {
+  imageWidth: i32,
+  tileWidth: i32,
+  outputOffset: vec2i,
+  outputSize: vec2i,
+  inputOffset: vec2i,
+
+}
+@group(0) @binding(0) var<storage, read> in_color: array<vec4f>;
+@group(0) @binding(1) var<storage, read> in_raw: array<vec4f>;
+@group(0) @binding(2) var<storage, read_write> out_color: array<vec4f>;
+@group(0) @binding(3) var<uniform> params: Params;
+${hdr ? `struct Frame { inputScale: f32, }
+@group(0) @binding(4) var<uniform> frame: Frame;` : ''}
+
+@compute @workgroup_size(${WORKGROUP}, ${WORKGROUP}, 1)
+fn main(@builtin(global_invocation_id) globalId: vec3u) {
+  let x = i32(globalId.x);
+  let y = i32(globalId.y);
+  if (x >= params.outputSize.x || y >= params.outputSize.y) {
+    return;
+  }
+  let inIdx = (y + params.inputOffset.y) * params.tileWidth + x + params.inputOffset.x;
+  let outIdx = (y + params.outputOffset.y) * params.imageWidth + x + params.outputOffset.x;
+  let col = in_color[inIdx];
+  let raw = in_raw[outIdx];
+  // The network output may contain negative values or NaNs.
+  let denoised = sanitize(col.rgb, 0.0, ${hdr ? 'FLT_MAX' : '1.0'});
+  ${hdr
+    ? 'out_color[outIdx] = vec4f(vec3f(HDRInverse(denoised.r), HDRInverse(denoised.g), HDRInverse(denoised.b)) / frame.inputScale, raw.a);'
+    : 'out_color[outIdx] = vec4f(denoised, raw.a);'}
+}
+`;
+}
+
+function uniformBuffer(device: OIDNGPUDevice, label: string, words: Int32Array) {
+  const buffer = device.createBuffer({
+    label,
+    size: alignUp(words.byteLength, 16),
+    usage: BufferUsage.UNIFORM | BufferUsage.COPY_DST,
+    mappedAtCreation: true
+  });
+  new Uint8Array(buffer.getMappedRange()).set(new Uint8Array(words.buffer));
+  buffer.unmap();
+  return buffer;
+}
+
+function storageBuffer(device: OIDNGPUDevice, label: string, size: number) {
+  return device.createBuffer({
+    label,
+    size,
+    usage: BufferUsage.STORAGE | BufferUsage.COPY_SRC | BufferUsage.COPY_DST
+  });
+}
+
+/**
+ * Records the commands of one denoising of a `width` x `height` image: for
+ * every planned tile the input pass, the network and the output pass. The
+ * tiles follow `planTileGrid` with the model's receptive field as overlap,
+ * like `UNet.tileExecute`.
+ */
+export function recordDenoiseExecution(
+  device: RecordingDevice,
+  executor: NativeUNetCore,
+  model: ValidatedUNetModel,
+  options: NativeDenoiseOptions
+) {
+  const { width, height } = options;
+  if (width <= 0 || height <= 0) throw new Error('Image dimensions must be positive');
+  const sourceCount = model.inputChannels / 3;
+  const imageBytes = width * height * 16;
+  const inputs: RecordedBuffer[] = [];
+  for (const name of ['color', 'albedo', 'normal'].slice(0, sourceCount)) {
+    inputs.push(storageBuffer(device, `oidn/external/${name}`, imageBytes) as RecordedBuffer);
+  }
+  const output = storageBuffer(device, 'oidn/external/output', imageBytes) as RecordedBuffer;
+
+  const overlap = alignUp(model.spec.receptiveField / 2, OIDN_TILE_ALIGNMENT);
+  const plan = planTileGrid(
+    width,
+    height,
+    options.maxTileSize > 0 ? options.maxTileSize : alignUp(Math.max(width, height), OIDN_TILE_ALIGNMENT),
+    overlap
+  );
+
+  const inputPipeline = device.createComputePipeline({
+    label: 'oidn/native/input',
+    layout: 'auto',
+    compute: {
+      module: device.createShaderModule({
+        label: 'oidn/native/input',
+        code: inputShader(sourceCount, options.hdr, options.transfer, executor.precision)
+      }),
+      entryPoint: 'main'
+    }
+  });
+  const outputPipeline = device.createComputePipeline({
+    label: 'oidn/native/output',
+    layout: 'auto',
+    compute: {
+      module: device.createShaderModule({
+        label: 'oidn/native/output',
+        code: outputShader(options.hdr, options.transfer)
+      }),
+      entryPoint: 'main'
+    }
+  });
+  const scaleWord = new Int32Array(new Float32Array([options.inputScale]).buffer)[0];
+
+  const frame = uniformBuffer(device, 'oidn/native/frame', new Int32Array([scaleWord, 0, 0, 0])) as RecordedBuffer;
+
+  // Network inputs are per tile shape; tiles run in order on one queue.
+  const tileInputs = new Map<string, OIDNGPUBuffer>();
+  for (const tile of plan.tiles) {
+    const key = `${tile.input.width}x${tile.input.height}`;
+    if (tileInputs.has(key)) continue;
+    tileInputs.set(key, executor.packedInputBuffer(tile.input.width, tile.input.height));
+  }
+  device.operations = [];
+
+  plan.tiles.forEach((tile, index) => {
+    const packed = tileInputs.get(`${tile.input.width}x${tile.input.height}`)!;
+    const inputParams = uniformBuffer(device, `oidn/native/input/${index}`, new Int32Array([
+      width, height, tile.input.width, tile.input.height, tile.input.x, tile.input.y, 0, 0
+    ]));
+    const entries: OIDNGPUBindGroupEntry[] = [];
+    for (let source = 0; source < sourceCount; source++) {
+      entries.push({ binding: source, resource: { buffer: inputs[source] } });
+    }
+    entries.push({ binding: sourceCount, resource: { buffer: packed } });
+    entries.push({ binding: sourceCount + 1, resource: { buffer: inputParams } });
+    if (options.hdr) entries.push({ binding: sourceCount + 2, resource: { buffer: frame } });
+    const inputBindings = device.createBindGroup({
+      label: `oidn/native/input/${index}`,
+      layout: inputPipeline.getBindGroupLayout(0),
+      entries
+    });
+    const inputEncoder = device.createCommandEncoder({ label: `oidn/native/input/${index}` });
+    const inputPass = inputEncoder.beginComputePass();
+    inputPass.setPipeline(inputPipeline);
+    inputPass.setBindGroup(0, inputBindings);
+    inputPass.dispatchWorkgroups(
+      Math.ceil(tile.input.width / WORKGROUP),
+      Math.ceil(tile.input.height / WORKGROUP)
+    );
+    inputPass.end();
+    device.queue.submit([inputEncoder.finish()]);
+
+    const networkOutput = executor.executePacked(tile.input.width, tile.input.height);
+
+    const outputParams = uniformBuffer(device, `oidn/native/output/${index}`, new Int32Array([
+      width, tile.input.width,
+      tile.output.x, tile.output.y,
+      tile.output.width, tile.output.height,
+      tile.output.x - tile.input.x, tile.output.y - tile.input.y,
+      0, 0, 0, 0
+    ]));
+    const outputBindings = device.createBindGroup({
+      label: `oidn/native/output/${index}`,
+      layout: outputPipeline.getBindGroupLayout(0),
+      entries: [
+        { binding: 0, resource: { buffer: networkOutput } },
+        { binding: 1, resource: { buffer: inputs[0] } },
+        { binding: 2, resource: { buffer: output } },
+        { binding: 3, resource: { buffer: outputParams } },
+        ...(options.hdr ? [{ binding: 4, resource: { buffer: frame } }] : [])
+      ]
+    });
+    const outputEncoder = device.createCommandEncoder({ label: `oidn/native/output/${index}` });
+    const outputPass = outputEncoder.beginComputePass();
+    outputPass.setPipeline(outputPipeline);
+    outputPass.setBindGroup(0, outputBindings);
+    outputPass.dispatchWorkgroups(
+      Math.ceil(tile.output.width / WORKGROUP),
+      Math.ceil(tile.output.height / WORKGROUP)
+    );
+    outputPass.end();
+    device.queue.submit([outputEncoder.finish()]);
+  });
+  return { inputs, output, tiles: plan.tiles.length, frame };
+}

package/src/native/program.ts

@@ -0,0 +1,326 @@
+import { BufferUsage } from '../gpuTypes.js';
+import { detectUNetModelSpec, validateUNetModel } from '../modelSpec.js';
+import {
+  NativeUNetCore,
+  type NativeUNetKernelSetting,
+  type NativeUNetPrecisionSetting
+} from '../nativeUNetCore.js';
+import { parseTZA } from '../tza.js';
+import { RecordingDevice, type RecordedBuffer } from './recordingDevice.js';
+import { recordDenoiseExecution } from './denoise.js';
+
+/**
+ * Records the native U-Net for one input size as a replayable program.
+ *
+ * Config (little-endian u32 words):
+ *   0 version (1), 1 width, 2 height,
+ *   3 precision (0 auto, 1 fp32, 2 fp16),
+ *   4 kernel (0 auto, 1 implicit-gemm, 2 direct, 3 subgroup-matrix),
+ *   5 features (1 shader-f16, 2 subgroups, 4 chromium-experimental-subgroup-matrix,
+ *     8 wgpu-cooperative-matrix),
+ *   6 maxComputeWorkgroupStorageSize, 7 maxComputeInvocationsPerWorkgroup,
+ *   8 maxComputeWorkgroupSizeX, 9 maxComputeWorkgroupSizeY,
+ *   10 maxComputeWorkgroupsPerDimension, 11 subgroupMinSize, 12 subgroupMaxSize,
+ *   13 matrix configuration count, then per configuration:
+ *      component type (0 f32, 1 f16), result type, M, N, K,
+ *   then subgroup-matrix options: double buffering (0 default, 1 off, 2 on),
+ *   accumulation (0 f32, 1 f16), sub-pixel decoders (0 on, 1 off).
+ *
+ * Program (little-endian u32 words, strings and byte blobs padded to 4 bytes):
+ *   magic 0x4e44494f ("OIDN"), version 2,
+ *   input count, input buffer ids..., output buffer id,
+ *   output external (1 when the host may supply the output buffer),
+ *   output width, output height, output channels,
+ *   buffers: count, then size, usage, data length (0 = zero-filled), data, label
+ *   shaders: count, then code
+ *   pipelines: count, then shader id, entry point
+ *   bind groups: count, then pipeline id, group, entry count,
+ *     entries (binding, buffer id, offset, size; size 0 = rest of buffer)
+ *   operations: count, then kind
+ *     0 submit: command count, commands as kind
+ *       0 dispatch: pipeline id, bind group id, x, y, z
+ *       1 copy: source, source offset, destination, destination offset, size
+ *     1 write: buffer id, offset, data length, data
+ *
+ * Input buffers are RGBA f32 images in the transfer space of the network
+ * (as produced by the browser runtime's GPU preprocessing); the output buffer
+ * holds RGBA f32 pixels in the same space.
+ */
+export function recordNetwork(tza: Uint8Array, config: Uint8Array): Uint8Array {
+  const { words, device, model, executor } = createNativeExecutor(tza, config);
+  const width = words[1];
+  const height = words[2];
+  const inputs: RecordedBuffer[] = [];
+  for (let index = 0; index < model.inputChannels / 3; index++) {
+    inputs.push(device.createBuffer({
+      label: `oidn/external/input${index}`,
+      size: width * height * 16,
+      usage: BufferUsage.STORAGE | BufferUsage.COPY_DST
+    }) as RecordedBuffer);
+  }
+  // The first execution creates the per-shape resources; the second records
+  // exactly what every later execution submits.
+  executor.execute(inputs, width, height);
+  device.operations = [];
+  const output = executor.execute(inputs, width, height) as RecordedBuffer;
+  return serializeProgram(device, inputs, output, false, width, height, model.outputChannels, executor, 0, undefined);
+}
+
+/**
+ * Records a whole denoising of one image size: tiles, input processing,
+ * the network and output processing (see `denoise.ts`). Config: the network
+ * config with width and height of the image, followed after the
+ * subgroup-matrix options by:
+ *   hdr (0/1), HDR transfer (0 PU, 1 log), maximum tile edge (0 = whole image),
+ *   input scale (f32 bits; 0 = 1).
+ * Inputs are color, then albedo and normal for auxiliary models, as RGBA f32
+ * images; the external output is an RGBA f32 image (color alpha preserved).
+ */
+export function recordDenoise(tza: Uint8Array, config: Uint8Array): Uint8Array {
+  const { words, device, model, executor, next } = createNativeExecutor(tza, config);
+  const width = words[1];
+  const height = words[2];
+  const scaleBits = new Uint32Array([words[next + 3]]);
+  const inputScale = words[next + 3] === 0 ? 1 : new Float32Array(scaleBits.buffer)[0];
+  const recorded = recordDenoiseExecution(device, executor, model, {
+    width,
+    height,
+    hdr: words[next] === 1,
+    transfer: words[next + 1] === 1 ? 'log' : 'pu',
+    maxTileSize: words[next + 2],
+    inputScale
+  });
+  return serializeProgram(device, recorded.inputs, recorded.output, true, width, height, 4, executor, recorded.tiles, recorded.frame);
+}
+
+function createNativeExecutor(tza: Uint8Array, config: Uint8Array) {
+  const words = new Uint32Array(config.buffer, config.byteOffset, config.byteLength / 4);
+  if (words[0] !== 1) throw new Error(`Unsupported native config version ${words[0]}`);
+  const precisions: NativeUNetPrecisionSetting[] = ['auto', 'fp32', 'fp16'];
+  const kernels: NativeUNetKernelSetting[] = ['auto', 'implicit-gemm', 'direct', 'subgroup-matrix'];
+  const features: string[] = [];
+  if (words[5] & 1) features.push('shader-f16');
+  if (words[5] & 2) features.push('subgroups');
+  if (words[5] & 4) features.push('chromium-experimental-subgroup-matrix');
+  if (words[5] & 8) features.push('wgpu-cooperative-matrix');
+  const configs: { componentType: string; resultComponentType: string; M: number; N: number; K: number }[] = [];
+  for (let index = 0; index < words[13]; index++) {
+    const base = 14 + index * 5;
+    configs.push({
+      componentType: words[base] === 1 ? 'f16' : 'f32',
+      resultComponentType: words[base + 1] === 1 ? 'f16' : 'f32',
+      M: words[base + 2],
+      N: words[base + 3],
+      K: words[base + 4]
+    });
+  }
+  const options = 14 + words[13] * 5;
+  const device = new RecordingDevice({
+    features,
+    limits: {
+      maxComputeWorkgroupStorageSize: words[6],
+      maxComputeInvocationsPerWorkgroup: words[7],
+      maxComputeWorkgroupSizeX: words[8],
+      maxComputeWorkgroupSizeY: words[9],
+      maxComputeWorkgroupsPerDimension: words[10]
+    },
+    adapterInfo: {
+      subgroupMinSize: words[11],
+      subgroupMaxSize: words[12],
+      subgroupMatrixConfigs: configs
+    }
+  });
+
+  const tensors = parseTZA(
+    tza.buffer.slice(tza.byteOffset, tza.byteOffset + tza.byteLength) as ArrayBuffer
+  );
+  const model = validateUNetModel(tensors, detectUNetModelSpec(tensors));
+  const doubleBuffer = words[options] === 0 ? undefined : words[options] === 2;
+  const executor = new NativeUNetCore(device, model, {
+    precision: precisions[words[3]],
+    kernel: kernels[words[4]],
+    gemm: { unrollRows: true },
+    subgroupMatrix: {
+      ...(doubleBuffer === undefined ? {} : { doubleBuffer }),
+      accumulation: words[options + 1] === 1 ? 'f16' : 'f32',
+      subPixel: words[options + 2] !== 1
+    }
+  });
+  return { words, device, model, executor, next: options + 3 };
+}
+
+class ProgramWriter {
+  private _bytes = new Uint8Array(1 << 16);
+  private _view = new DataView(this._bytes.buffer);
+  length = 0;
+
+  private _reserve(count: number) {
+    if (this.length + count <= this._bytes.length) return;
+    let capacity = this._bytes.length * 2;
+    while (capacity < this.length + count) capacity *= 2;
+    const bytes = new Uint8Array(capacity);
+    bytes.set(this._bytes.subarray(0, this.length));
+    this._bytes = bytes;
+    this._view = new DataView(bytes.buffer);
+  }
+
+  u32(value: number) {
+    this._reserve(4);
+    this._view.setUint32(this.length, value, true);
+    this.length += 4;
+  }
+
+  bytes(data: Uint8Array) {
+    this.u32(data.length);
+    const padded = (data.length + 3) & ~3;
+    this._reserve(padded);
+    this._bytes.set(data, this.length);
+    this._bytes.fill(0, this.length + data.length, this.length + padded);
+    this.length += padded;
+  }
+
+  string(value: string) {
+    this.bytes(new TextEncoder().encode(value));
+  }
+
+  finish() {
+    return this._bytes.slice(0, this.length);
+  }
+}
+
+// Device-independent model validation for native Rust hosts. Packet: MODEL
+// magic, version, image input count; errors use the existing ERRO packet.
+export function inspectModelPacket(tza: Uint8Array): Uint8Array {
+  const writer = new ProgramWriter();
+  try {
+    const tensors = parseTZA(tza.buffer.slice(tza.byteOffset, tza.byteOffset + tza.byteLength) as ArrayBuffer);
+    const model = validateUNetModel(tensors, detectUNetModelSpec(tensors));
+    writer.u32(0x4c444f4d);
+    writer.u32(1);
+    writer.u32(model.inputChannels / 3);
+  } catch (error) {
+    writer.u32(0x4f525245);
+    writer.string(error instanceof Error ? error.message : String(error));
+  }
+  return writer.finish();
+}
+
+function serializeProgram(
+  device: RecordingDevice,
+  inputs: readonly RecordedBuffer[],
+  output: RecordedBuffer,
+  outputExternal: boolean,
+  width: number,
+  height: number,
+  outputChannels: number,
+  executor: NativeUNetCore,
+  tiles: number,
+  frame: RecordedBuffer | undefined
+) {
+  const writer = new ProgramWriter();
+  writer.u32(0x4e44494f);
+  writer.u32(2);
+  writer.u32(inputs.length);
+  for (const input of inputs) writer.u32(input.id);
+  writer.u32(output.id);
+  writer.u32(outputExternal ? 1 : 0);
+  writer.u32(width);
+  writer.u32(height);
+  writer.u32(outputChannels);
+
+  writer.u32(device.buffers.length);
+  for (const buffer of device.buffers) {
+    if (!Number.isSafeInteger(buffer.size) || buffer.size <= 0 || buffer.size > 0xffffffff) {
+      throw new Error(`Native program buffer ${buffer.label} exceeds the 32-bit size limit`);
+    }
+    writer.u32(buffer.size);
+    writer.u32(buffer.usage);
+    writer.bytes(buffer.data ?? new Uint8Array(0));
+    writer.string(buffer.label);
+  }
+  writer.u32(device.shaders.length);
+  for (const shader of device.shaders) writer.string(shader.code);
+  writer.u32(device.pipelines.length);
+  for (const pipeline of device.pipelines) {
+    writer.u32(pipeline.shader);
+    writer.string(pipeline.entryPoint);
+  }
+  writer.u32(device.bindGroups.length);
+  for (const bindGroup of device.bindGroups) {
+    writer.u32(bindGroup.pipeline);
+    writer.u32(bindGroup.group);
+    writer.u32(bindGroup.entries.length);
+    for (const entry of bindGroup.entries) {
+      writer.u32(entry.binding);
+      writer.u32(entry.buffer);
+      writer.u32(entry.offset);
+      writer.u32(entry.size);
+    }
+  }
+  writer.u32(device.operations.length);
+  for (const operation of device.operations) {
+    writer.u32(operation.kind);
+    if (operation.kind === 0) {
+      writer.u32(operation.commands.length);
+      for (const recorded of operation.commands) {
+        writer.u32(recorded.kind);
+        if (recorded.kind === 0) {
+          writer.u32(recorded.pipeline);
+          writer.u32(recorded.bindGroup);
+          writer.u32(recorded.x);
+          writer.u32(recorded.y);
+          writer.u32(recorded.z);
+        } else {
+          writer.u32(recorded.source);
+          writer.u32(recorded.sourceOffset);
+          writer.u32(recorded.destination);
+          writer.u32(recorded.destinationOffset);
+          writer.u32(recorded.size);
+        }
+      }
+    } else {
+      writer.u32(operation.buffer);
+      writer.u32(operation.offset);
+      writer.bytes(operation.data ?? new Uint8Array(0));
+    }
+  }
+  // Optional extension after v2 operations. Old replayers can ignore it.
+  writer.u32(0x4154454d); // META
+  writer.u32(1);
+  writer.u32(executor.precision === 'fp16' ? 2 : 1);
+  writer.u32(tiles);
+  let kernels = 0;
+  for (const kernel of executor.usedKernels) {
+    kernels |= kernel === 'direct' ? 1 : kernel === 'implicit-gemm' ? 2
+      : kernel === 'subgroup-matrix' ? 4 : kernel === 'spatial' ? 8 : 16;
+  }
+  writer.u32(kernels);
+  writer.u32(frame ? frame.id : 0xffffffff);
+  const weights = executor.weightBuffers as RecordedBuffer[];
+  writer.u32(weights.length);
+  for (const buffer of weights) writer.u32(buffer.id);
+  return writer.finish();
+}
+
+/** Error packet: magic 0x4f525245 ("ERRO") followed by the message. */
+export function recordNetworkPacket(tza: Uint8Array, config: Uint8Array): Uint8Array {
+  try {
+    return recordNetwork(tza, config);
+  } catch (error) {
+    const writer = new ProgramWriter();
+    writer.u32(0x4f525245);
+    writer.string(error instanceof Error ? error.message : String(error));
+    return writer.finish();
+  }
+}
+
+export function recordDenoisePacket(tza: Uint8Array, config: Uint8Array): Uint8Array {
+  try {
+    return recordDenoise(tza, config);
+  } catch (error) {
+    const writer = new ProgramWriter();
+    writer.u32(0x4f525245);
+    writer.string(error instanceof Error ? error.message : String(error));
+    return writer.finish();
+  }
+}