oidn-webgpu 0.0.0-stage → 0.6.0

package/src/backend.ts

@@ -0,0 +1,42 @@
+export async function initWebGPUBackend() {
+  if (!navigator.gpu) throw new Error('WebGPU is not available');
+  const gpuDescriptor: GPURequestAdapterOptions = {
+    powerPreference: 'high-performance'
+  };
+
+  const adapter = await navigator.gpu.requestAdapter(gpuDescriptor);
+  if (!adapter) throw new Error('No WebGPU adapter is available');
+  const deviceDescriptor: GPUDeviceDescriptor = {};
+
+  const requiredFeatures: GPUFeatureName[] = [];
+  if (adapter.features.has('timestamp-query')) {
+    requiredFeatures.push('timestamp-query');
+  }
+  if (adapter.features.has('bgra8unorm-storage')) {
+    requiredFeatures.push('bgra8unorm-storage');
+  }
+  if (adapter.features.has('shader-f16')) {
+    requiredFeatures.push('shader-f16');
+  }
+  // Experimental: only exposed by Chrome with unsafe WebGPU features enabled.
+  for (const feature of ['subgroups', 'chromium-experimental-subgroup-matrix']) {
+    if (adapter.features.has(feature)) {
+      requiredFeatures.push(feature as GPUFeatureName);
+    }
+  }
+  deviceDescriptor.requiredFeatures = requiredFeatures;
+
+  const adapterLimits = adapter.limits;
+  deviceDescriptor.requiredLimits = {
+    maxComputeWorkgroupStorageSize:
+      adapterLimits.maxComputeWorkgroupStorageSize,
+    maxComputeWorkgroupsPerDimension:
+      adapterLimits.maxComputeWorkgroupsPerDimension,
+    maxStorageBufferBindingSize: adapterLimits.maxStorageBufferBindingSize,
+    maxBufferSize: adapterLimits.maxBufferSize,
+    maxComputeWorkgroupSizeX: adapterLimits.maxComputeWorkgroupSizeX,
+    maxComputeInvocationsPerWorkgroup:
+      adapterLimits.maxComputeInvocationsPerWorkgroup
+  };
+  return adapter.requestDevice(deviceDescriptor);
+}

package/src/finalRgbShader.ts

@@ -0,0 +1,186 @@
+export type FinalRgbPrecision = 'fp16' | 'fp32';
+export type FinalRgbActivation = 'relu' | 'identity';
+
+const WORKGROUP_SIZE = 8;
+const PATCH_SIZE = WORKGROUP_SIZE + 2;
+const KERNEL_ELEMENTS = 3 * 3;
+
+/** Shared storage required by the 8x8 final-RGB workgroup. */
+export function sharedMemoryBytes(
+  precision: FinalRgbPrecision,
+  inputBlocks: number,
+  cacheWeights = false
+) {
+  const bytesPerScalar = precision === 'fp16' ? 2 : 4;
+  const inputVec4Count = PATCH_SIZE * PATCH_SIZE * inputBlocks;
+  const weightVec4Count = cacheWeights
+    ? KERNEL_ELEMENTS * inputBlocks * 4
+    : 0;
+  const vec4Count = inputVec4Count + weightVec4Count;
+  return vec4Count * 4 * bytesPerScalar;
+}
+
+function storageVecType(precision: FinalRgbPrecision) {
+  return precision === 'fp16' ? 'vec4<f16>' : 'vec4<f32>';
+}
+
+function shaderPreamble(precision: FinalRgbPrecision) {
+  return precision === 'fp16' ? 'enable f16;\n' : '';
+}
+
+function accumulationCode(
+  precision: FinalRgbPrecision,
+  inputExpression: string,
+  weightExpression: string
+) {
+  const weight = (lane: number) =>
+    `${weightExpression}[weightBase + ${lane}u]`;
+  if (precision === 'fp16') {
+    return /* wgsl */ `
+        let inputValue = vec4<f16>(${inputExpression});
+        var partial = vec4<f16>(0.0h);
+        partial = fma(${weight(0)}, vec4<f16>(inputValue.x), partial);
+        partial = fma(${weight(1)}, vec4<f16>(inputValue.y), partial);
+        partial = fma(${weight(2)}, vec4<f16>(inputValue.z), partial);
+        partial = fma(${weight(3)}, vec4<f16>(inputValue.w), partial);
+        acc += vec4<f32>(partial);
+`;
+  }
+  return /* wgsl */ `
+        let inputValue = vec4<f32>(${inputExpression});
+        acc = fma(vec4<f32>(${weight(0)}), vec4<f32>(inputValue.x), acc);
+        acc = fma(vec4<f32>(${weight(1)}), vec4<f32>(inputValue.y), acc);
+        acc = fma(vec4<f32>(${weight(2)}), vec4<f32>(inputValue.z), acc);
+        acc = fma(vec4<f32>(${weight(3)}), vec4<f32>(inputValue.w), acc);
+`;
+}
+
+/**
+ * Builds the final three-channel same-padded convolution shader.
+ *
+ * The bind-group layout and Params block intentionally match createConvShader:
+ * input, weights, bias, output, then uniform params. Weights retain the
+ * output-major packed ABI, with four vec4 values per (kernel position, input
+ * block), one for each input lane.
+ */
+export function createFinalRgbShader(
+  precision: FinalRgbPrecision,
+  activation: FinalRgbActivation,
+  inputBlocks: number,
+  cacheWeights = false
+) {
+  if (!Number.isInteger(inputBlocks) || inputBlocks < 1) {
+    throw new Error(`Final RGB shader requires positive input blocks, got ${inputBlocks}`);
+  }
+  const inputType = storageVecType(precision);
+  const outputType = 'vec4<f32>';
+  const inputTileValues = PATCH_SIZE * PATCH_SIZE * inputBlocks;
+  const weightTileValues = KERNEL_ELEMENTS * inputBlocks * 4;
+  const stored = activation === 'relu'
+    ? 'max(acc, vec4<f32>(0.0))'
+    : 'acc';
+  const accumulation = accumulationCode(
+    precision,
+    'inputTile[patchBase + inputBlock]',
+    cacheWeights ? 'weightTile' : 'weights'
+  );
+
+  return /* wgsl */ `${shaderPreamble(precision)}
+struct Params {
+  inputWidth: u32,
+  inputHeight: u32,
+  outputWidth: u32,
+  outputHeight: u32,
+  inputBlocks: u32,
+  outputBlocks: u32,
+}
+
+@group(0) @binding(0) var<storage, read> inputData: array<${inputType}>;
+@group(0) @binding(1) var<storage, read> weights: array<${inputType}>;
+@group(0) @binding(2) var<storage, read> bias: array<vec4<f32>>;
+@group(0) @binding(3) var<storage, read_write> outputData: array<${outputType}>;
+@group(0) @binding(4) var<uniform> params: Params;
+
+var<workgroup> inputTile: array<${inputType}, ${inputTileValues}>;
+${cacheWeights
+    ? `var<workgroup> weightTile: array<${inputType}, ${weightTileValues}>;`
+    : ''}
+
+@compute @workgroup_size(${WORKGROUP_SIZE}, ${WORKGROUP_SIZE}, 1)
+fn main(
+  @builtin(local_invocation_id) localId: vec3<u32>,
+  @builtin(global_invocation_id) gid: vec3<u32>,
+  @builtin(workgroup_id) workgroupId: vec3<u32>
+) {
+  let localLinear = localId.y * ${WORKGROUP_SIZE}u + localId.x;
+  for (
+    var loadIndex = localLinear;
+    loadIndex < ${inputTileValues}u;
+    loadIndex += ${WORKGROUP_SIZE * WORKGROUP_SIZE}u
+  ) {
+    let tilePixel = loadIndex / ${inputBlocks}u;
+    let inputBlock = loadIndex % ${inputBlocks}u;
+    let tileX = tilePixel % ${PATCH_SIZE}u;
+    let tileY = tilePixel / ${PATCH_SIZE}u;
+    let inputX = i32(workgroupId.x * ${WORKGROUP_SIZE}u + tileX) - 1;
+    let inputY = i32(workgroupId.y * ${WORKGROUP_SIZE}u + tileY) - 1;
+    var value = ${inputType}(0.0);
+    if (
+      inputX >= 0 && inputX < i32(params.inputWidth) &&
+      inputY >= 0 && inputY < i32(params.inputHeight)
+    ) {
+      let inputIndex =
+        (u32(inputY) * params.inputWidth + u32(inputX)) *
+        ${inputBlocks}u + inputBlock;
+      value = inputData[inputIndex];
+    }
+    inputTile[loadIndex] = value;
+  }
+
+${cacheWeights ? `  for (
+    var loadIndex = localLinear;
+    loadIndex < ${weightTileValues}u;
+    loadIndex += ${WORKGROUP_SIZE * WORKGROUP_SIZE}u
+  ) {
+    weightTile[loadIndex] = weights[loadIndex];
+  }
+
+` : ''}  // Out-of-range invocations must reach this barrier before returning.
+  workgroupBarrier();
+
+  let outputInBounds =
+    gid.x < params.outputWidth &&
+    gid.y < params.outputHeight &&
+    gid.z < params.outputBlocks;
+  if (!outputInBounds) {
+    return;
+  }
+
+  var acc = bias[gid.z];
+  for (var ky = 0u; ky < 3u; ky++) {
+    let inputY = i32(gid.y) + i32(ky) - 1;
+    if (inputY < 0 || inputY >= i32(params.inputHeight)) {
+      continue;
+    }
+    for (var kx = 0u; kx < 3u; kx++) {
+      let inputX = i32(gid.x) + i32(kx) - 1;
+      if (inputX < 0 || inputX >= i32(params.inputWidth)) {
+        continue;
+      }
+      let patchBase =
+        ((localId.y + ky) * ${PATCH_SIZE}u + localId.x + kx) *
+        ${inputBlocks}u;
+      for (var inputBlock = 0u; inputBlock < ${inputBlocks}u; inputBlock++) {
+        let weightBase =
+          ((ky * 3u + kx) * ${inputBlocks}u + inputBlock) * 4u;
+        ${accumulation}
+      }
+    }
+  }
+
+  let outputIndex =
+    (gid.y * params.outputWidth + gid.x) * ${1}u + gid.z;
+  outputData[outputIndex] = ${stored};
+}
+`;
+}

package/src/gpuTypes.ts

@@ -0,0 +1,179 @@
+/**
+ * The WebGPU subset used by the native WGSL runtime, as structural types and
+ * spec-defined flag values. A browser GPUDevice implements these operations,
+ * but its branded TypeScript types need the conversion in nativeUNet.ts;
+ * they also let an embedder supply its own device (for example one that
+ * records the work for replay in a native wgpu host) without lib.dom or
+ * WebGPU globals.
+ */
+
+/** GPUBufferUsage flag values from the WebGPU specification. */
+export const BufferUsage = {
+  MAP_READ: 0x0001,
+  MAP_WRITE: 0x0002,
+  COPY_SRC: 0x0004,
+  COPY_DST: 0x0008,
+  INDEX: 0x0010,
+  VERTEX: 0x0020,
+  UNIFORM: 0x0040,
+  STORAGE: 0x0080,
+  INDIRECT: 0x0100,
+  QUERY_RESOLVE: 0x0200
+} as const;
+
+/** GPUMapMode flag values from the WebGPU specification. */
+export const MapMode = {
+  READ: 0x0001,
+  WRITE: 0x0002
+} as const;
+
+// The synchronous core types exclude promises so scriptc's library mode can
+// compile the core; browser-only code uses the asynchronous extensions below.
+
+export interface OIDNGPUBuffer {
+  readonly size: number;
+  readonly mapState: string;
+  getMappedRange(offset?: number, size?: number): ArrayBuffer;
+  unmap(): void;
+  destroy(): void;
+}
+
+export interface OIDNGPUMappableBuffer extends OIDNGPUBuffer {
+  mapAsync(mode: number, offset?: number, size?: number): Promise<void>;
+}
+
+export interface OIDNGPUQuerySet {
+  destroy(): void;
+}
+
+export interface OIDNGPUShaderModule {
+  readonly label: string;
+}
+
+export interface OIDNGPUBindGroupLayout {
+  readonly label: string;
+}
+
+export interface OIDNGPUBindGroup {
+  readonly label: string;
+}
+
+export interface OIDNGPUComputePipeline {
+  getBindGroupLayout(index: number): OIDNGPUBindGroupLayout;
+}
+
+export interface OIDNGPUCommandBuffer {
+  readonly label: string;
+}
+
+export interface OIDNGPUBindGroupEntry {
+  binding: number;
+  resource: { buffer: OIDNGPUBuffer; offset?: number; size?: number };
+}
+
+export interface OIDNGPUComputePassTimestampWrites {
+  querySet: OIDNGPUQuerySet;
+  beginningOfPassWriteIndex?: number;
+  endOfPassWriteIndex?: number;
+}
+
+export interface OIDNGPUComputePassEncoder {
+  setPipeline(pipeline: OIDNGPUComputePipeline): void;
+  setBindGroup(index: number, bindGroup: OIDNGPUBindGroup): void;
+  dispatchWorkgroups(x: number, y?: number, z?: number): void;
+  end(): void;
+}
+
+export interface OIDNGPUCommandEncoder {
+  beginComputePass(descriptor?: {
+    label?: string;
+    timestampWrites?: OIDNGPUComputePassTimestampWrites;
+  }): OIDNGPUComputePassEncoder;
+  copyBufferToBuffer(
+    source: OIDNGPUBuffer,
+    sourceOffset: number,
+    destination: OIDNGPUBuffer,
+    destinationOffset: number,
+    size: number
+  ): void;
+  resolveQuerySet(
+    querySet: OIDNGPUQuerySet,
+    firstQuery: number,
+    queryCount: number,
+    destination: OIDNGPUBuffer,
+    destinationOffset: number
+  ): void;
+  finish(): OIDNGPUCommandBuffer;
+}
+
+export interface OIDNGPUQueue {
+  submit(commandBuffers: OIDNGPUCommandBuffer[]): void;
+  writeBuffer(
+    buffer: OIDNGPUBuffer,
+    bufferOffset: number,
+    data: Uint8Array
+  ): void;
+}
+
+export interface OIDNGPUAsyncQueue extends OIDNGPUQueue {
+  onSubmittedWorkDone(): Promise<void>;
+}
+
+export interface OIDNGPULimits {
+  readonly maxComputeWorkgroupStorageSize: number;
+  readonly maxComputeInvocationsPerWorkgroup: number;
+  readonly maxComputeWorkgroupSizeX: number;
+  readonly maxComputeWorkgroupSizeY: number;
+  readonly maxComputeWorkgroupsPerDimension: number;
+}
+
+export interface OIDNGPUFeatures {
+  has(name: string): boolean;
+}
+
+/** Subset of GPUAdapterInfo used to select subgroup-matrix configurations. */
+export interface OIDNGPUAdapterInfo {
+  readonly subgroupMinSize?: number;
+  readonly subgroupMaxSize?: number;
+  readonly subgroupMatrixConfigs?: readonly {
+    componentType: string;
+    resultComponentType: string;
+    M: number;
+    N: number;
+    K: number;
+  }[];
+}
+
+export interface OIDNGPUDevice {
+  readonly features: OIDNGPUFeatures;
+  readonly limits: OIDNGPULimits;
+  readonly queue: OIDNGPUQueue;
+  createBuffer(descriptor: {
+    label?: string;
+    size: number;
+    usage: number;
+    mappedAtCreation?: boolean;
+  }): OIDNGPUBuffer;
+  createShaderModule(descriptor: { label?: string; code: string }): OIDNGPUShaderModule;
+  createComputePipeline(descriptor: {
+    label?: string;
+    layout: 'auto';
+    compute: { module: OIDNGPUShaderModule; entryPoint: string };
+  }): OIDNGPUComputePipeline;
+  createBindGroup(descriptor: {
+    label?: string;
+    layout: OIDNGPUBindGroupLayout;
+    entries: OIDNGPUBindGroupEntry[];
+  }): OIDNGPUBindGroup;
+  createCommandEncoder(descriptor?: { label?: string }): OIDNGPUCommandEncoder;
+  createQuerySet(descriptor: { label?: string; type: 'timestamp'; count: number }): OIDNGPUQuerySet;
+}
+
+export interface OIDNGPUAsyncDevice extends OIDNGPUDevice {
+  readonly queue: OIDNGPUAsyncQueue;
+  createComputePipelineAsync(descriptor: {
+    label?: string;
+    layout: 'auto';
+    compute: { module: OIDNGPUShaderModule; entryPoint: string };
+  }): Promise<OIDNGPUComputePipeline>;
+}

package/src/graphOptimizer.ts

@@ -0,0 +1,301 @@
+import type {
+  Conv2DNodeSpec,
+  MaxPool2DNodeSpec,
+  ModelNodeSpec,
+  UNetModelSpec,
+  Upsample2DNodeSpec,
+  UNetModelGraph
+} from './modelSpec';
+
+export interface FusedConvPoolNodeSpec {
+  op: 'fusedConvReluMaxPool2d';
+  id: string;
+  input: string;
+  conv: Conv2DNodeSpec;
+  pool: MaxPool2DNodeSpec;
+}
+
+export interface FusedUpsampleConcatConvNodeSpec {
+  op: 'fusedUpsampleConcatConv2d';
+  id: string;
+  /** Inputs stay in concat order because that order selects weight channels. */
+  inputs: readonly {
+    value: string;
+    upsample?: Upsample2DNodeSpec;
+  }[];
+  conv: Conv2DNodeSpec;
+}
+
+export type ExecutableModelNode =
+  | ModelNodeSpec
+  | FusedConvPoolNodeSpec
+  | FusedUpsampleConcatConvNodeSpec;
+
+export interface GraphOptimizationOptions {
+  fuseConvPool?: boolean;
+  fuseUpsampleConcatConv?: boolean;
+}
+
+export interface OptimizedModelGraph {
+  spec: UNetModelSpec;
+  nodes: readonly ExecutableModelNode[];
+  fusions: {
+    convPool: number;
+    upsampleConcatConv: number;
+  };
+}
+
+function nodeInputs(node: ModelNodeSpec): readonly string[] {
+  return node.op === 'concat' ? node.inputs : [node.input];
+}
+
+function buildConsumers(spec: UNetModelSpec) {
+  const consumers = new Map<string, ModelNodeSpec[]>();
+  for (const node of spec.nodes) {
+    for (const input of nodeInputs(node)) {
+      const list = consumers.get(input) ?? [];
+      list.push(node);
+      consumers.set(input, list);
+    }
+  }
+  return consumers;
+}
+
+// Callers narrow on `op`; node ids are unique, so they compare consumers by id.
+function onlyConsumer(
+  consumers: ReadonlyMap<string, ModelNodeSpec[]>,
+  value: string,
+  op: ModelNodeSpec['op']
+): ModelNodeSpec | undefined {
+  const list = consumers.get(value);
+  if (list === undefined || list.length !== 1 || list[0].op !== op) return undefined;
+  return list[0];
+}
+
+/**
+ * Applies topology-only fusions. It never depends on a particular OIDN model
+ * name, so new descriptors automatically benefit from known graph patterns.
+ */
+export function optimizeModelGraph(
+  validated: UNetModelGraph,
+  options: GraphOptimizationOptions = {}
+): OptimizedModelGraph {
+  const spec = validated.spec;
+  const consumers = buildConsumers(spec);
+  const nodesById = new Map(spec.nodes.map((node) => [node.id, node]));
+  const eliminated = new Set<string>();
+  const fusedAt = new Map<string, ExecutableModelNode>();
+  let convPool = 0;
+  let upsampleConcatConv = 0;
+
+  if (options.fuseConvPool !== false) {
+    for (const node of spec.nodes) {
+      if (node.op !== 'conv2d' || node.activation !== 'relu') continue;
+      const pool = onlyConsumer(consumers, node.id, 'maxPool2d');
+      if (!pool || pool.op !== 'maxPool2d' || pool.size !== 2 || pool.stride !== 2) continue;
+
+      eliminated.add(node.id);
+      fusedAt.set(pool.id, {
+        op: 'fusedConvReluMaxPool2d',
+        id: pool.id,
+        input: node.input,
+        conv: node,
+        pool
+      });
+      convPool++;
+    }
+  }
+
+  if (options.fuseUpsampleConcatConv !== false) {
+    for (const node of spec.nodes) {
+      if (node.op !== 'conv2d') continue;
+      const concat = nodesById.get(node.input);
+      if (concat?.op !== 'concat' || concat.inputs.length !== 2) continue;
+      if (onlyConsumer(consumers, concat.id, 'conv2d')?.id !== node.id) continue;
+      // The native blocked layout can remove concat only when the source
+      // boundary is also a vec4 boundary. Other graphs keep the generic ops.
+      const firstInputChannels = validated.channelsByValue.get(concat.inputs[0]);
+      if (firstInputChannels === undefined || firstInputChannels % 4 !== 0) {
+        continue;
+      }
+
+      const inputs = concat.inputs.map((value) => {
+        const candidate = nodesById.get(value);
+        if (
+          candidate?.op === 'upsample2d' &&
+          candidate.scale === 2 &&
+          candidate.mode === 'nearest' &&
+          onlyConsumer(consumers, candidate.id, 'concat')?.id === concat.id
+        ) {
+          return { value: candidate.input, upsample: candidate };
+        }
+        return { value };
+      });
+      const upsampleCount = inputs.filter((input) => input.upsample).length;
+      if (upsampleCount !== 1) continue;
+
+      eliminated.add(concat.id);
+      for (const value of concat.inputs) {
+        const candidate = nodesById.get(value);
+        if (candidate?.op === 'upsample2d') eliminated.add(candidate.id);
+      }
+      fusedAt.set(node.id, {
+        op: 'fusedUpsampleConcatConv2d',
+        id: node.id,
+        inputs,
+        conv: node
+      });
+      upsampleConcatConv++;
+    }
+  }
+
+  const nodes: ExecutableModelNode[] = [];
+  for (const node of spec.nodes) {
+    const fused = fusedAt.get(node.id);
+    if (fused) {
+      nodes.push(fused);
+    } else if (!eliminated.has(node.id)) {
+      nodes.push(node);
+    }
+  }
+
+  return {
+    spec,
+    nodes,
+    fusions: { convPool, upsampleConcatConv }
+  };
+}
+
+export interface ModelValueShape {
+  width: number;
+  height: number;
+  channels: number;
+}
+
+export interface PlannedModelNode {
+  node: ExecutableModelNode;
+  outputShape: ModelValueShape;
+  /** Last planned node that reads the output; output itself uses nodes.length. */
+  lastUse: number;
+}
+
+export interface ModelExecutionPlan extends OptimizedModelGraph {
+  inputShape: ModelValueShape;
+  valueShapes: ReadonlyMap<string, ModelValueShape>;
+  plannedNodes: readonly PlannedModelNode[];
+}
+
+function executableInputs(node: ExecutableModelNode): readonly string[] {
+  if (node.op === 'concat') return node.inputs;
+  if (node.op === 'fusedUpsampleConcatConv2d') {
+    return node.inputs.map((input) => input.value);
+  }
+  return [node.input];
+}
+
+function sameSpatialShape(
+  left: ModelValueShape,
+  right: ModelValueShape
+): boolean {
+  return left.width === right.width && left.height === right.height;
+}
+
+/** Resolve all runtime shapes and value lifetimes before allocating GPU data. */
+export function planModelExecution(
+  validated: UNetModelGraph,
+  width: number,
+  height: number,
+  options?: GraphOptimizationOptions
+): ModelExecutionPlan {
+  if (!Number.isInteger(width) || width <= 0 || !Number.isInteger(height) || height <= 0) {
+    throw new Error(`Invalid model input size ${width}x${height}`);
+  }
+
+  const graph = optimizeModelGraph(validated, options);
+  const inputShape = { width, height, channels: validated.inputChannels };
+  const valueShapes = new Map<string, ModelValueShape>([
+    [validated.spec.input, inputShape]
+  ]);
+  const outputShapes: ModelValueShape[] = [];
+
+  const shapeOf = (value: string, nodeId: string) => {
+    const shape = valueShapes.get(value);
+    if (!shape) throw new Error(`Planned node ${nodeId} reads missing value ${value}`);
+    return shape;
+  };
+
+  for (const node of graph.nodes) {
+    let outputShape: ModelValueShape;
+    if (node.op === 'conv2d') {
+      const input = shapeOf(node.input, node.id);
+      outputShape = {
+        width: input.width,
+        height: input.height,
+        channels: validated.convChannels.get(node.id)!.outputChannels
+      };
+    } else if (node.op === 'maxPool2d') {
+      const input = shapeOf(node.input, node.id);
+      outputShape = {
+        width: Math.ceil(input.width / 2),
+        height: Math.ceil(input.height / 2),
+        channels: input.channels
+      };
+    } else if (node.op === 'upsample2d') {
+      const input = shapeOf(node.input, node.id);
+      outputShape = {
+        width: input.width * 2,
+        height: input.height * 2,
+        channels: input.channels
+      };
+    } else if (node.op === 'concat') {
+      const inputs = node.inputs.map((value) => shapeOf(value, node.id));
+      if (inputs.some((shape) => !sameSpatialShape(shape, inputs[0]))) {
+        throw new Error(`Concat ${node.id} has mismatched spatial shapes`);
+      }
+      outputShape = {
+        width: inputs[0].width,
+        height: inputs[0].height,
+        channels: inputs.reduce((sum, shape) => sum + shape.channels, 0)
+      };
+    } else if (node.op === 'fusedConvReluMaxPool2d') {
+      const input = shapeOf(node.input, node.id);
+      outputShape = {
+        width: Math.ceil(input.width / 2),
+        height: Math.ceil(input.height / 2),
+        channels: validated.convChannels.get(node.conv.id)!.outputChannels
+      };
+    } else {
+      const inputs = node.inputs.map((input) => {
+        const source = shapeOf(input.value, node.id);
+        return input.upsample
+          ? { ...source, width: source.width * 2, height: source.height * 2 }
+          : source;
+      });
+      if (inputs.some((shape) => !sameSpatialShape(shape, inputs[0]))) {
+        throw new Error(`Fused decoder ${node.id} has mismatched spatial shapes`);
+      }
+      outputShape = {
+        width: inputs[0].width,
+        height: inputs[0].height,
+        channels: validated.convChannels.get(node.conv.id)!.outputChannels
+      };
+    }
+
+    valueShapes.set(node.id, outputShape);
+    outputShapes.push(outputShape);
+  }
+
+  const lastUses = new Map<string, number>();
+  graph.nodes.forEach((node, index) => {
+    for (const input of executableInputs(node)) lastUses.set(input, index);
+  });
+  lastUses.set(validated.spec.output, graph.nodes.length);
+
+  const plannedNodes = graph.nodes.map((node, index) => ({
+    node,
+    outputShape: outputShapes[index],
+    lastUse: lastUses.get(node.id) ?? index
+  }));
+
+  return { ...graph, inputShape, valueShapes, plannedNodes };
+}