oidn-web 0.3.5 → 0.5.0

package/src/WGPUComputePass.ts

@@ -152,13 +152,15 @@ export class WGPUComputePass<I extends string, O extends string> {
     return this._outputBuffers[name].buffer;
   }
 
-  dispose() {
+  dispose(destroyOutputBuffers = true) {
     Object.keys(this._uniformBuffers).forEach((key) => {
       (this._uniformBuffers as any)[key].destroy();
     });
-    Object.keys(this._outputBuffers).forEach((key) => {
-      (this._outputBuffers as any)[key].buffer.destroy();
-    });
+    if (destroyOutputBuffers) {
+      Object.keys(this._outputBuffers).forEach((key) => {
+        (this._outputBuffers as any)[key].buffer.destroy();
+      });
+    }
   }
 
   private _createBuffer(params: WGPUComputePassOutput) {

package/src/backend.ts

@@ -1,62 +1,36 @@
-import { WebGPUBackend } from '@tensorflow/tfjs-backend-webgpu/dist/base';
-import { ENGINE } from '@tensorflow/tfjs-core/dist/engine';
-
-import './kernels';
-
 export async function initWebGPUBackend() {
-  try {
-    const gpuDescriptor: GPURequestAdapterOptions = {
-      powerPreference: 'high-performance'
-    };
-
-    const adapter = (await navigator.gpu.requestAdapter(gpuDescriptor))!;
-    const deviceDescriptor: GPUDeviceDescriptor = {};
-
-    const requiredFeatures = [];
-    if (adapter.features.has('timestamp-query')) {
-      requiredFeatures.push('timestamp-query');
-    }
-    if (adapter.features.has('bgra8unorm-storage')) {
-      requiredFeatures.push(['bgra8unorm-storage']);
-    }
-    deviceDescriptor.requiredFeatures =
-      requiredFeatures as Iterable<GPUFeatureName>;
-
-    const adapterLimits = adapter.limits;
-    deviceDescriptor.requiredLimits = {
-      maxComputeWorkgroupStorageSize:
-        adapterLimits.maxComputeWorkgroupStorageSize,
-      maxComputeWorkgroupsPerDimension:
-        adapterLimits.maxComputeWorkgroupsPerDimension,
-      maxStorageBufferBindingSize: adapterLimits.maxStorageBufferBindingSize,
-      maxBufferSize: adapterLimits.maxBufferSize,
-      maxComputeWorkgroupSizeX: adapterLimits.maxComputeWorkgroupSizeX,
-      maxComputeInvocationsPerWorkgroup:
-        adapterLimits.maxComputeInvocationsPerWorkgroup
-    };
-    const device = await adapter.requestDevice(deviceDescriptor);
-    const adapterInfo =
-      // requestAdapterInfo is deprecated
-      // @ts-ignore
-      adapter!.info ?? (await adapter!.requestAdapterInfo?.());
-
-    return initWebGPUBackendWithDevice(device, adapterInfo);
-  } catch (e) {}
-}
-
-export async function initWebGPUBackendWithDevice(
-  device: GPUDevice,
-  adapter: GPUAdapterInfo
-) {
-  // TODO multiple device and adapter in one backend
-  let backend = ENGINE.findBackend('webgpu-oidn');
-  if (backend != null) {
-    return backend as WebGPUBackend;
+  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');
   }
-
-  backend = new WebGPUBackend(device, adapter);
-  ENGINE.registerBackend('webgpu-oidn', () => backend);
-  await ENGINE.setBackend('webgpu-oidn');
-
-  return backend as WebGPUBackend;
+  if (adapter.features.has('bgra8unorm-storage')) {
+    requiredFeatures.push('bgra8unorm-storage');
+  }
+  if (adapter.features.has('shader-f16')) {
+    requiredFeatures.push('shader-f16');
+  }
+  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/graphOptimizer.ts

@@ -0,0 +1,300 @@
+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;
+}
+
+function onlyConsumer<T extends ModelNodeSpec['op']>(
+  consumers: ReadonlyMap<string, ModelNodeSpec[]>,
+  value: string,
+  op: T
+): Extract<ModelNodeSpec, { op: T }> | undefined {
+  const list = consumers.get(value);
+  if (list?.length !== 1 || list[0].op !== op) return undefined;
+  return list[0] as Extract<ModelNodeSpec, { op: T }>;
+}
+
+/**
+ * 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.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') !== node) 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') === concat
+        ) {
+          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 };
+}

package/src/hdrTransfer.ts

@@ -0,0 +1,88 @@
+/** HDR transfer functions used by the OIDN input and output processing passes. */
+export type HDRTransfer = 'pu' | 'log';
+
+const a = 1.41283765e3;
+const b = 1.64593172;
+const c = 4.31384981e-1;
+const d = -2.94139609e-3;
+const e = 1.92653254e-1;
+const f = 6.26026094e-3;
+const g = 9.98620152e-1;
+const y0 = 1.5794576e-6;
+const y1 = 3.22087631e-2;
+const x0 = 2.23151711e-3;
+const x1 = 3.70974749e-1;
+const yMax = 65504;
+const puXMax = puForward(yMax);
+const puNormScale = 1 / puXMax;
+const puRcpNormScale = puXMax;
+const logXMax = Math.log(yMax + 1);
+const logNormScale = 1 / logXMax;
+
+function puForward(y: number) {
+  if (y <= y0) return a * y;
+  if (y <= y1) return b * Math.pow(y, c) + d;
+  return e * Math.log(y + f) + g;
+}
+
+function puInverse(x: number) {
+  if (x <= x0) return x / a;
+  if (x <= x1) return Math.pow((x - d) / b, 1 / c);
+  return Math.exp((x - g) / e) - f;
+}
+
+function forward(y: number, transfer: HDRTransfer) {
+  return transfer === 'log'
+    ? Math.log(y + 1) * logNormScale
+    : puForward(y) * puNormScale;
+}
+
+function inverse(x: number, transfer: HDRTransfer) {
+  return transfer === 'log'
+    ? Math.exp(x * logXMax) - 1
+    : puInverse(x * puRcpNormScale);
+}
+
+export function hdrTransferFuncCPU({
+  data,
+  channels,
+  inputScale,
+  transfer = 'pu'
+}: {
+  data: Float32Array;
+  channels: number;
+  inputScale: number;
+  transfer?: HDRTransfer;
+}) {
+  const newData = new Float32Array(data);
+  for (let i = 0; i < newData.length; i += channels) {
+    for (let channel = 0; channel < 3; channel++) {
+      newData[i + channel] = forward(
+        newData[i + channel] * inputScale,
+        transfer
+      );
+    }
+  }
+  return newData;
+}
+
+export function hdrTransferFuncInverseCPU({
+  data,
+  channels,
+  inputScale,
+  transfer = 'pu'
+}: {
+  data: Float32Array;
+  channels: number;
+  inputScale: number;
+  transfer?: HDRTransfer;
+}) {
+  const newData = new Float32Array(data);
+  const outputScale = 1 / inputScale;
+  for (let i = 0; i < newData.length; i += channels) {
+    for (let channel = 0; channel < 3; channel++) {
+      newData[i + channel] = inverse(newData[i + channel], transfer) * outputScale;
+    }
+  }
+  return newData;
+}