oidn-web 0.4.0 → 0.5.0

package/src/backend.ts

@@ -32,18 +32,5 @@ export async function initWebGPUBackend() {
     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);
-}
-
-export async function initWebGPUBackendWithDevice(
-  device: GPUDevice,
-  adapterInfo: GPUAdapterInfo
-) {
-  return { device, adapterInfo };
+  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

@@ -112,8 +112,7 @@ export function optimizeModelGraph(
       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
-      // and can use the compatibility engine until a scalar-tail kernel exists.
+      // 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;

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;
+}

package/src/main.ts

@@ -1,16 +1,25 @@
 import { parseTZA } from './tza';
 import UNet from './UNet';
 import type { UNetEngineSetting, UNetExecutionStats } from './UNet';
-import { initWebGPUBackend, initWebGPUBackendWithDevice } from './backend';
+import { initWebGPUBackend } from './backend';
 import type { DynamicTileSetting } from './tileScheduler';
 import type { UNetModelSpec } from './modelSpec';
 import type {
+  NativeUNetGemmOptions,
   NativeUNetKernelSetting,
   NativeUNetPrecisionSetting
 } from './nativeUNet';
+import type { HDRTransfer } from './process';
 
 export { parseTZA, UNet };
-export type { DynamicTileOptions, DynamicTileSetting } from './tileScheduler';
+export { planTileGrid } from './tileScheduler';
+export type {
+  DynamicTileOptions,
+  DynamicTileSetting,
+  PlannedTile,
+  TilePlan,
+  TileRect
+} from './tileScheduler';
 export {
   detectUNetModelSpec,
   OIDN_UNET_LARGE_SPEC,
@@ -36,6 +45,8 @@ export {
   resolveNativeUNetPrecision
 } from './nativeUNet';
 export type {
+  NativeUNetGemmWorkgroup,
+  NativeUNetGemmOptions,
   NativeUNetExecutionProfile,
   NativeUNetKernel,
   NativeUNetKernelSetting,
@@ -48,21 +59,30 @@ export type {
 export interface UNetOptions {
   aux?: boolean;
   hdr?: boolean;
+  /** HDR transfer function expected by the trained model. Defaults to PU. */
+  hdrTransfer?: HDRTransfer;
   /** Hard upper bound for an output tile edge. Defaults to 512. */
   maxTileSize?: number;
   /** Adaptive GPU-time-based tile sizing. Enabled by default. */
   dynamicTile?: DynamicTileSetting;
-  /** `auto` uses stable WGSL; `webnn` opts into the experimental WebNN backend. */
+  /** `auto` uses native WGSL; `webnn` opts into the WebNN backend. */
   engine?: UNetEngineSetting;
   /** `auto` selects FP16 when shader-f16 was enabled on the GPUDevice. */
   precision?: NativeUNetPrecisionSetting;
-  /** `auto` selects kernels from precision, operation shape, and GPU limits. */
+  /** `auto` uses implicit GEMM for FP16/FP32 convolutions, except the direct output layer. */
   kernel?: NativeUNetKernelSetting;
+  /**
+   * Experimental implicit-GEMM tuning for native WGSL execution. These knobs
+   * exist for benchmarking and may change or be removed in any release.
+   * @experimental
+   */
+  gemm?: NativeUNetGemmOptions;
   /** Versioned topology descriptor for future/custom OIDN TZA models. */
   modelSpec?: UNetModelSpec;
 }
 
 export type { UNetEngineSetting, UNetExecutionStats } from './UNet';
+export type { HDRTransfer } from './process';
 export { WebNNUNetExecutor } from './webnnUNet';
 export type { WebNNRuntimeSupport, WebNNUNetOptions } from './webnnUNet';
 export type {
@@ -73,24 +93,19 @@ export type {
 
 export async function initUNetFromBuffer(
   tzaBuffer: ArrayBuffer,
-  backendParams?: { device: GPUDevice; adapterInfo: GPUAdapterInfo },
+  backendParams?: { device: GPUDevice },
   opts?: UNetOptions
 ) {
-  const backend = await (backendParams
-    ? initWebGPUBackendWithDevice(
-        backendParams.device,
-        backendParams.adapterInfo
-      )
-    : initWebGPUBackend());
+  const device = backendParams?.device ?? await initWebGPUBackend();
   const tensors = parseTZA(tzaBuffer);
-  const unet = new UNet(tensors, backend, opts);
+  const unet = new UNet(tensors, device, opts);
   await unet.prepare();
   return unet;
 }
 
 export async function initUNetFromURL(
   modelPath: string,
-  backendParams?: { device: GPUDevice; adapterInfo: GPUAdapterInfo },
+  backendParams?: { device: GPUDevice },
   opts?: UNetOptions
 ) {
   return fetch(modelPath)