oidn-web 0.3.5 → 0.5.0

package/src/process.ts

@@ -1,5 +1,11 @@
-import { buffer, data } from '@tensorflow/tfjs';
 import { WGPUComputePass } from './WGPUComputePass';
+import {
+  hdrTransferFuncCPU,
+  hdrTransferFuncInverseCPU,
+  type HDRTransfer
+} from './hdrTransfer';
+export { hdrTransferFuncCPU, hdrTransferFuncInverseCPU } from './hdrTransfer';
+export type { HDRTransfer } from './hdrTransfer';
 
 const a = 1.41283765e3;
 const b = 1.64593172;
@@ -38,6 +44,8 @@ const yMax = 65504;
 const xMax = PUForward(yMax);
 const normScale = 1 / xMax;
 const rcpNormScale = xMax;
+const logXMax = Math.log(yMax + 1);
+const logNormScale = 1 / logXMax;
 
 export class Tile {
   constructor(
@@ -76,53 +84,6 @@ export function avgLogLum({
   return key / Math.pow(2, averageLuminance);
 }
 
-export function hdrTransferFuncCPU({
-  data,
-  channels,
-  inputScale
-}: {
-  data: Float32Array;
-  channels: number;
-  inputScale: number;
-}) {
-  const newData = new Float32Array(data.length);
-  newData.set(data);
-  // https://github.com/RenderKit/oidn/blob/713ec7838ba650f99e0a896549c0dca5eeb3652d/core/color.h#L71
-  for (let i = 0; i < newData.length; i += channels) {
-    // First three are color
-    for (let c = 0; c < 3; c++) {
-      let y = newData[i + c] * inputScale;
-      // https://github.com/RenderKit/oidn/blob/713ec7838ba650f99e0a896549c0dca5eeb3652d/devices/cpu/color.ispc#L135
-      newData[i + c] = PUForward(y) * normScale;
-    }
-  }
-
-  return newData;
-}
-
-export function hdrTransferFuncInverseCPU({
-  data,
-  channels,
-  inputScale
-}: {
-  data: Float32Array;
-  channels: number;
-  inputScale: number;
-}) {
-  const newData = new Float32Array(data.length);
-  newData.set(data);
-
-  const outputScale = 1 / inputScale;
-  for (let i = 0; i < newData.length; i += channels) {
-    for (let c = 0; c < 3; c++) {
-      let x = newData[i + c] * rcpNormScale;
-      newData[i + c] = PUInverse(x) * outputScale;
-    }
-  }
-
-  return newData;
-}
-
 const constsCode = `
 const a = ${a};
 const b = ${b};
@@ -148,7 +109,11 @@ export class GPUDataProcess {
 
   private _isInputTexture?: boolean;
 
-  constructor(private _device: GPUDevice, private _isHDR: boolean) {
+  constructor(
+    private _device: GPUDevice,
+    private _isHDR: boolean,
+    private _hdrTransfer: HDRTransfer = 'pu'
+  ) {
     const commonUniforms = [
       {
         label: 'inputScale',
@@ -262,26 +227,40 @@ out_color[outIdx] = textureLoad(in_color, globalId.xy, 0);
 
     this._isInputTexture = isInputTexture;
     const isHDR = this._isHDR;
+    const hdrForwardCode = this._hdrTransfer === 'log'
+      ? `fn HDRForward(y: f32) -> f32 { return log(y + 1.0) * logNormScale; }`
+      : `fn HDRForward(y: f32) -> f32 {
+  if (y <= y0) { return a * y * normScale; }
+  else if (y <= y1) { return (b * pow(y, c) + d) * normScale; }
+  else { return (e * log(y + f) + g) * normScale; }
+}`;
+    const hdrInverseCode = this._hdrTransfer === 'log'
+      ? `fn HDRInverse(x: f32) -> f32 { return exp(x * logXMax) - 1.0; }`
+      : `fn HDRInverse(x: f32) -> f32 {
+  let y = x * rcpNormScale;
+  if (y <= x0) { return y / a; }
+  else if (y <= x1) { return pow((y - d) / b, 1 / c); }
+  else { return exp((y - g) / e) - f; }
+}`;
     const commonCSDefine = /* wgsl */ `
 ${constsCode}
-fn PUForward(y: f32) -> f32 {
-  if (y <= y0) {
-    return a * y;
-  } else if (y <= y1) {
-    return b * pow(y, c) + d;
-  } else {
-    return e * log(y + f) + g;
-  }
-}`;
+const logXMax = ${logXMax};
+const logNormScale = ${logNormScale};
+${hdrForwardCode}`;
     function readInputCode(inputName: string) {
       return isInputTexture
-        ? `textureLoad(in_${inputName}, globalId.xy + vec2u(inputOffset), 0)`
+        ? `textureLoad(in_${inputName}, vec2u(inputPosition), 0)`
         : `in_${inputName}[inIdx]`;
     }
     const commonCSMain = /* wgsl */ `
 let x = i32(globalId.x);
 let y = i32(globalId.y);
-let inIdx = (y + inputOffset.y) * inputSize.x + (x + inputOffset.x);
+let inputPosition = clamp(
+  vec2i(x, y) + inputOffset,
+  vec2i(0),
+  inputSize - vec2i(1)
+);
+let inIdx = inputPosition.y * inputSize.x + inputPosition.x;
 let col = ${readInputCode('color')};
 
 let outIdx = y * outputSize.x + x;
@@ -291,7 +270,7 @@ if (${denoiseAlpha}) {
   out_color[outIdx] = vec3f(1.0 - col.a);
 }
 else if (${isHDR}) {
-  out_color[outIdx] = vec3f(PUForward(col.r * inputScale), PUForward(col.g * inputScale), PUForward(col.b * inputScale)) * normScale;
+  out_color[outIdx] = vec3f(HDRForward(col.r * inputScale), HDRForward(col.g * inputScale), HDRForward(col.b * inputScale));
 }
 else {
   out_color[outIdx] = col.rgb;
@@ -317,15 +296,8 @@ ${commonCSMain}
     this._outputPass.setCSCode({
       csDefine: /* wgsl */ `
 ${constsCode}
-fn PUInverse(y: f32) -> f32 {
-  if (y <= x0) {
-    return y / a;
-  } else if (y <= x1) {
-    return pow((y - d) / b, 1 / c);
-  } else {
-    return exp((y - g) / e) - f;
-  }
-}
+const logXMax = ${logXMax};
+${hdrInverseCode}
 `,
       csMain: /* wgsl */ `
 let x = i32(globalId.x);
@@ -347,7 +319,7 @@ if (${denoiseAlpha}) {
 }
 else if (${isHDR}) {
   out_color[outIdx] = vec4f(
-    vec3f(PUInverse(col.r * rcpNormScale), PUInverse(col.g * rcpNormScale), PUInverse(col.b * rcpNormScale)) / inputScale,
+    vec3f(HDRInverse(col.r), HDRInverse(col.g), HDRInverse(col.b)) / inputScale,
     // Pick the alpha
     raw.a
   );
@@ -490,5 +462,8 @@ else {
   dispose() {
     this._outputPass.dispose();
     this._inputPassAux.dispose();
+    this._inputPassColor.dispose();
+    // copyPass borrows outputPass's target and must not destroy it again.
+    this._copyPass.dispose(false);
   }
 }

package/src/resourceTracker.ts

@@ -0,0 +1,94 @@
+export type OIDNResourceKind =
+  | 'gpu-buffer'
+  | 'gpu-query-set'
+  | 'ml-context'
+  | 'ml-graph'
+  | 'ml-tensor';
+
+export interface OIDNResourceStats {
+  created: number;
+  destroyed: number;
+  live: number;
+  peakLive: number;
+}
+
+export interface OIDNResourceSnapshot {
+  live: number;
+  created: number;
+  destroyed: number;
+  peakLive: number;
+  pending: number;
+  byKind: Partial<Record<OIDNResourceKind, OIDNResourceStats>>;
+}
+
+interface MutableResourceStats extends OIDNResourceStats {
+  resources: Set<object>;
+}
+
+/**
+ * Counts explicitly owned browser resources. Keeping the ownership registry in
+ * production makes cache eviction and failure cleanup testable without relying
+ * on JS heap measurements, which do not include WebGPU allocations.
+ */
+export class OIDNResourceTracker {
+  private _stats = new Map<OIDNResourceKind, MutableResourceStats>();
+
+  track<T extends object>(kind: OIDNResourceKind, resource: T): T {
+    let stats = this._stats.get(kind);
+    if (!stats) {
+      stats = {
+        created: 0,
+        destroyed: 0,
+        live: 0,
+        peakLive: 0,
+        resources: new Set()
+      };
+      this._stats.set(kind, stats);
+    }
+    if (stats.resources.has(resource)) return resource;
+    stats.resources.add(resource);
+    stats.created++;
+    stats.live++;
+    stats.peakLive = Math.max(stats.peakLive, stats.live);
+    return resource;
+  }
+
+  release(
+    kind: OIDNResourceKind,
+    resource: object | undefined,
+    destroy: () => void
+  ) {
+    if (!resource) return false;
+    const stats = this._stats.get(kind);
+    if (!stats?.resources.delete(resource)) return false;
+    try {
+      destroy();
+    } finally {
+      stats.destroyed++;
+      stats.live--;
+    }
+    return true;
+  }
+
+  snapshot(pending = 0): OIDNResourceSnapshot {
+    let live = 0;
+    let created = 0;
+    let destroyed = 0;
+    let peakLive = 0;
+    const byKind: OIDNResourceSnapshot['byKind'] = {};
+    for (const [kind, stats] of this._stats) {
+      const entry = {
+        created: stats.created,
+        destroyed: stats.destroyed,
+        live: stats.live,
+        peakLive: stats.peakLive
+      };
+      byKind[kind] = entry;
+      live += entry.live;
+      created += entry.created;
+      destroyed += entry.destroyed;
+      peakLive += entry.peakLive;
+    }
+    return { live, created, destroyed, peakLive, pending, byKind };
+  }
+}

package/src/tileScheduler.ts

@@ -0,0 +1,330 @@
+export interface DynamicTileOptions {
+  /** Smallest output tile edge. Defaults to 256. */
+  minTileSize?: number;
+  /** Initial maximum output tile edge. Defaults to 432. */
+  initialTileSize?: number;
+  /** Desired upper bound for one tile's GPU latency. Defaults to 16 ms. */
+  targetTileTimeMs?: number;
+  /** Amount added or removed after each completed execution. Defaults to 16. */
+  adjustmentStep?: number;
+}
+
+export type DynamicTileSetting = boolean | DynamicTileOptions;
+
+const defaultMinTileSize = 256;
+const defaultInitialTileSize = 432;
+const defaultTargetTileTimeMs = 16;
+const defaultAdjustmentStep = 16;
+export const OIDN_TILE_ALIGNMENT = 16;
+
+export interface TileRect {
+  x: number;
+  y: number;
+  width: number;
+  height: number;
+}
+
+export interface PlannedTile {
+  column: number;
+  row: number;
+  input: TileRect;
+  output: TileRect;
+}
+
+export interface TilePlan {
+  columns: number;
+  rows: number;
+  overlap: number;
+  maxOutputWidth: number;
+  maxOutputHeight: number;
+  inputPixelCount: number;
+  inputShapeCount: number;
+  tiles: PlannedTile[];
+}
+
+function alignUp(value: number, alignment: number) {
+  return Math.ceil(value / alignment) * alignment;
+}
+
+function alignDown(value: number, alignment: number) {
+  return Math.floor(value / alignment) * alignment;
+}
+
+function partitionRange(length: number, parts: number, index: number) {
+  const start = Math.round(index * length / parts);
+  const end = Math.round((index + 1) * length / parts);
+  return { start, end };
+}
+
+function alignRangeWithinImage(
+  start: number,
+  end: number,
+  imageSize: number,
+  minimumSize = 0
+) {
+  const currentSize = end - start;
+  const alignedSize = Math.max(
+    minimumSize,
+    alignUp(currentSize, OIDN_TILE_ALIGNMENT)
+  );
+  let alignedStart = start - Math.floor((alignedSize - currentSize) / 2);
+  // At an image edge, keep the padded portion on the right/bottom. Input
+  // readers replicate the nearest image edge into that portion, while output
+  // tiles remain strictly inside the original image.
+  alignedStart = clamp(alignedStart, 0, Math.max(0, imageSize - alignedSize));
+  return { start: alignedStart, end: alignedStart + alignedSize };
+}
+
+/**
+ * Splits an image into balanced rectangular output regions. Only sides that
+ * touch another tile receive a halo; image edges rely on the model's normal
+ * boundary behavior instead of doing redundant off-edge work.
+ */
+export function planTileGrid(
+  width: number,
+  height: number,
+  maxTileSize: number,
+  overlap: number
+): TilePlan {
+  if (!Number.isInteger(width) || width <= 0 ||
+      !Number.isInteger(height) || height <= 0) {
+    throw new Error('Tile grid dimensions must be positive integers');
+  }
+  if (!Number.isFinite(maxTileSize) || maxTileSize <= 0) {
+    throw new Error('Maximum tile size must be positive');
+  }
+  if (!Number.isFinite(overlap) || overlap < 0) {
+    throw new Error('Tile overlap must be non-negative');
+  }
+
+  const alignedMaximum = Math.max(
+    OIDN_TILE_ALIGNMENT,
+    alignDown(maxTileSize, OIDN_TILE_ALIGNMENT)
+  );
+  const alignedOverlap = alignUp(overlap, OIDN_TILE_ALIGNMENT);
+  const columns = Math.max(1, Math.ceil(width / alignedMaximum));
+  const rows = Math.max(1, Math.ceil(height / alignedMaximum));
+  const tiles: PlannedTile[] = [];
+  let maxOutputWidth = 0;
+  let maxOutputHeight = 0;
+
+  for (let row = 0; row < rows; row++) {
+    const outputY = partitionRange(height, rows, row);
+    for (let column = 0; column < columns; column++) {
+      const outputX = partitionRange(width, columns, column);
+      const desiredInputX = {
+        start: column === 0
+          ? outputX.start
+          : Math.max(0, outputX.start - alignedOverlap),
+        end: column === columns - 1
+          ? outputX.end
+          : Math.min(width, outputX.end + alignedOverlap)
+      };
+      const desiredInputY = {
+        start: row === 0
+          ? outputY.start
+          : Math.max(0, outputY.start - alignedOverlap),
+        end: row === rows - 1
+          ? outputY.end
+          : Math.min(height, outputY.end + alignedOverlap)
+      };
+      const inputX = alignRangeWithinImage(
+        desiredInputX.start,
+        desiredInputX.end,
+        width
+      );
+      const inputY = alignRangeWithinImage(
+        desiredInputY.start,
+        desiredInputY.end,
+        height
+      );
+      const output: TileRect = {
+        x: outputX.start,
+        y: outputY.start,
+        width: outputX.end - outputX.start,
+        height: outputY.end - outputY.start
+      };
+      const input: TileRect = {
+        x: inputX.start,
+        y: inputY.start,
+        width: inputX.end - inputX.start,
+        height: inputY.end - inputY.start
+      };
+      maxOutputWidth = Math.max(maxOutputWidth, output.width);
+      maxOutputHeight = Math.max(maxOutputHeight, output.height);
+      tiles.push({ column, row, input, output });
+    }
+  }
+
+  const inputShapes = () => new Set(
+    tiles.map(({ input }) => `${input.width}x${input.height}`)
+  );
+  if (inputShapes().size > 2) {
+    const maximumInputWidth = Math.max(
+      ...tiles.map(({ input }) => input.width)
+    );
+    const maximumInputHeight = Math.max(
+      ...tiles.map(({ input }) => input.height)
+    );
+    const uniformWidthPixelCount = tiles.reduce(
+      (sum, { input }) => sum + maximumInputWidth * input.height,
+      0
+    );
+    const uniformHeightPixelCount = tiles.reduce(
+      (sum, { input }) => sum + input.width * maximumInputHeight,
+      0
+    );
+    if (uniformWidthPixelCount <= uniformHeightPixelCount) {
+      for (const tile of tiles) {
+        const range = alignRangeWithinImage(
+          tile.input.x,
+          tile.input.x + tile.input.width,
+          width,
+          maximumInputWidth
+        );
+        tile.input.x = range.start;
+        tile.input.width = range.end - range.start;
+      }
+    } else {
+      for (const tile of tiles) {
+        const range = alignRangeWithinImage(
+          tile.input.y,
+          tile.input.y + tile.input.height,
+          height,
+          maximumInputHeight
+        );
+        tile.input.y = range.start;
+        tile.input.height = range.end - range.start;
+      }
+    }
+  }
+
+  const finalInputShapes = inputShapes();
+  const inputPixelCount = tiles.reduce(
+    (sum, { input }) => sum + input.width * input.height,
+    0
+  );
+
+  return {
+    columns,
+    rows,
+    overlap: alignedOverlap,
+    maxOutputWidth,
+    maxOutputHeight,
+    inputPixelCount,
+    inputShapeCount: finalInputShapes.size,
+    tiles
+  };
+}
+
+function clamp(value: number, min: number, max: number) {
+  return Math.min(Math.max(value, min), max);
+}
+
+function percentile(values: number[], fraction: number) {
+  const sorted = [...values].sort((a, b) => a - b);
+  return sorted[Math.ceil(sorted.length * fraction) - 1];
+}
+
+/** Fits one output dimension without exceeding the configured tile cap. */
+export function fitTileDimension(imageSize: number, maxTileSize: number) {
+  return imageSize <= maxTileSize
+    ? Math.min(alignUp(imageSize, OIDN_TILE_ALIGNMENT), maxTileSize)
+    : maxTileSize;
+}
+
+/**
+ * Chooses an output tile size from completed GPU timings. A size change is
+ * applied to the next execution so one in-flight tiled pass keeps a stable
+ * overlap, allocation plan, and pipeline set.
+ */
+export class DynamicTileController {
+  readonly enabled: boolean;
+  readonly maxTileSize: number;
+  readonly minTileSize: number;
+  readonly targetTileTimeMs: number;
+
+  private _tileSize: number;
+  private _adjustmentStep: number;
+  private _smoothedTileTimeMs?: number;
+  private _completeExecutionsSinceChange = 0;
+
+  constructor(maxTileSize: number, setting: DynamicTileSetting = true) {
+    const options = typeof setting === 'object' ? setting : {};
+    this.enabled = setting !== false;
+    this.maxTileSize = Math.max(
+      OIDN_TILE_ALIGNMENT,
+      alignDown(maxTileSize, OIDN_TILE_ALIGNMENT)
+    );
+    this.minTileSize = clamp(
+      alignUp(options.minTileSize ?? defaultMinTileSize, OIDN_TILE_ALIGNMENT),
+      OIDN_TILE_ALIGNMENT,
+      this.maxTileSize
+    );
+    this.targetTileTimeMs = Math.max(
+      1,
+      options.targetTileTimeMs ?? defaultTargetTileTimeMs
+    );
+    this._adjustmentStep = Math.max(
+      OIDN_TILE_ALIGNMENT,
+      alignUp(
+        options.adjustmentStep ?? defaultAdjustmentStep,
+        OIDN_TILE_ALIGNMENT
+      )
+    );
+    this._tileSize = this.enabled
+      ? clamp(
+          alignUp(
+            options.initialTileSize ?? defaultInitialTileSize,
+            OIDN_TILE_ALIGNMENT
+          ),
+          this.minTileSize,
+          this.maxTileSize
+        )
+      : this.maxTileSize;
+  }
+
+  get tileSize() {
+    return this._tileSize;
+  }
+
+  /** Returns true when the next execution should use a different tile size. */
+  observe(tileTimesMs: number[]) {
+    if (!this.enabled || tileTimesMs.length === 0) return false;
+
+    const validTimes = tileTimesMs.filter(
+      (duration) => Number.isFinite(duration) && duration >= 0
+    );
+    if (validTimes.length === 0) return false;
+
+    // The first tile may still pay one-time GPU residency costs even after
+    // shape preparation. It should not make a stable layout look too slow.
+    const representativeTimes = validTimes.length >= 3
+      ? validTimes.slice(1)
+      : validTimes;
+    const observedTime = percentile(representativeTimes, 0.75);
+    this._smoothedTileTimeMs = this._smoothedTileTimeMs === undefined
+      ? observedTime
+      : this._smoothedTileTimeMs * 0.65 + observedTime * 0.35;
+    this._completeExecutionsSinceChange++;
+    if (this._completeExecutionsSinceChange < 2) return false;
+
+    let nextTileSize = this._tileSize;
+    if (this._smoothedTileTimeMs > this.targetTileTimeMs * 1.25) {
+      nextTileSize -= this._adjustmentStep;
+    } else if (this._smoothedTileTimeMs < this.targetTileTimeMs * 0.65) {
+      nextTileSize += this._adjustmentStep;
+    }
+
+    nextTileSize = clamp(
+      alignUp(nextTileSize, OIDN_TILE_ALIGNMENT),
+      this.minTileSize,
+      this.maxTileSize
+    );
+    if (nextTileSize === this._tileSize) return false;
+    this._tileSize = nextTileSize;
+    this._smoothedTileTimeMs = undefined;
+    this._completeExecutionsSinceChange = 0;
+    return true;
+  }
+}