oidn-web 0.3.4 → 0.4.0

package/src/process.ts

@@ -1,4 +1,3 @@
-import { buffer, data } from '@tensorflow/tfjs';
 import { WGPUComputePass } from './WGPUComputePass';
 
 const a = 1.41283765e3;
@@ -157,18 +156,18 @@ export class GPUDataProcess {
       },
       {
         label: 'inputSize',
-        type: 'vec2<f32>',
-        data: new Float32Array(2)
+        type: 'vec2i',
+        data: new Int32Array(2)
       },
       {
         label: 'outputSize',
-        type: 'vec2<f32>',
-        data: new Float32Array(2)
+        type: 'vec2i',
+        data: new Int32Array(2)
       },
       {
         label: 'inputOffset',
-        type: 'vec2<f32>',
-        data: new Float32Array(2)
+        type: 'vec2i',
+        data: new Int32Array(2)
       }
     ];
     this._inputPassAux = new WGPUComputePass('inputPassAux', this._device, {
@@ -196,28 +195,28 @@ export class GPUDataProcess {
         },
         {
           label: 'inputSize',
-          type: 'vec2<f32>',
-          data: new Float32Array(2)
+          type: 'vec2i',
+          data: new Int32Array(2)
         },
         {
           label: 'outputSize',
-          type: 'vec2<f32>',
-          data: new Float32Array(2)
+          type: 'vec2i',
+          data: new Int32Array(2)
         },
         {
           label: 'imageSize',
-          type: 'vec2<f32>',
-          data: new Float32Array(2)
+          type: 'vec2i',
+          data: new Int32Array(2)
         },
         {
           label: 'inputOffset',
-          type: 'vec2<f32>',
-          data: new Float32Array(2)
+          type: 'vec2i',
+          data: new Int32Array(2)
         },
         {
           label: 'outputOffset',
-          type: 'vec2<f32>',
-          data: new Float32Array(2)
+          type: 'vec2i',
+          data: new Int32Array(2)
         }
       ],
       csDefine: '',
@@ -230,8 +229,8 @@ export class GPUDataProcess {
       uniforms: [
         {
           label: 'size',
-          type: 'vec2<f32>',
-          data: new Float32Array(2)
+          type: 'vec2i',
+          data: new Int32Array(2)
         }
       ],
       csMain: /*wgsl*/ `
@@ -279,12 +278,12 @@ fn PUForward(y: f32) -> f32 {
         : `in_${inputName}[inIdx]`;
     }
     const commonCSMain = /* wgsl */ `
-let x = f32(globalId.x);
-let y = f32(globalId.y);
-let inIdx = i32((y + inputOffset.y) * inputSize.x + (x + inputOffset.x));
+let x = i32(globalId.x);
+let y = i32(globalId.y);
+let inIdx = (y + inputOffset.y) * inputSize.x + (x + inputOffset.x);
 let col = ${readInputCode('color')};
 
-let outIdx = i32(y * outputSize.x + x);
+let outIdx = y * outputSize.x + x;
 
 if (${denoiseAlpha}) {
   // Denoise the inversed alpha. Or the anti aliased edge will be too dark after denoised
@@ -328,13 +327,13 @@ fn PUInverse(y: f32) -> f32 {
 }
 `,
       csMain: /* wgsl */ `
-let x = f32(globalId.x);
-let y = f32(globalId.y);
+let x = i32(globalId.x);
+let y = i32(globalId.y);
 if (x >= outputSize.x || y >= outputSize.y) {
   return;
 }
-let inIdx = i32((y + inputOffset.y) * inputSize.x + x + inputOffset.x);
-let outIdx = i32((y + outputOffset.y) * imageSize.x + x + outputOffset.x);
+let inIdx = (y + inputOffset.y) * inputSize.x + x + inputOffset.x;
+let outIdx = (y + outputOffset.y) * imageSize.x + x + outputOffset.x;
 let col = in_color[inIdx];
 let raw = ${
         isInputTexture
@@ -360,18 +359,18 @@ else {
   }
 
   setImageSize(w: number, h: number) {
-    this._inputPassAux.setUniform('inputSize', new Float32Array([w, h]));
-    this._inputPassColor.setUniform('inputSize', new Float32Array([w, h]));
-    this._outputPass.setUniform('imageSize', new Float32Array([w, h]));
+    this._inputPassAux.setUniform('inputSize', new Int32Array([w, h]));
+    this._inputPassColor.setUniform('inputSize', new Int32Array([w, h]));
+    this._outputPass.setUniform('imageSize', new Int32Array([w, h]));
     this._outputPass.setSize(w, h);
     this._copyPass.setSize(w, h);
-    this._copyPass.setUniform('size', new Float32Array([w, h]));
+    this._copyPass.setUniform('size', new Int32Array([w, h]));
   }
 
   setInputTile(tile: Tile) {
-    const size = new Float32Array([tile.width, tile.height]);
+    const size = new Int32Array([tile.width, tile.height]);
     [this._inputPassAux, this._inputPassColor].forEach((inputPass) => {
-      inputPass.setUniform('inputOffset', new Float32Array([tile.x, tile.y]));
+      inputPass.setUniform('inputOffset', new Int32Array([tile.x, tile.y]));
       inputPass.setUniform('outputSize', size);
       inputPass.setSize(size[0], size[1]);
     });
@@ -381,14 +380,14 @@ else {
 
   setOutputTile(dstTile: Tile, srcTile: Tile) {
     const outputPass = this._outputPass;
-    const size = new Float32Array([dstTile.width, dstTile.height]);
+    const size = new Int32Array([dstTile.width, dstTile.height]);
     const dx = dstTile.x - srcTile.x;
     const dy = dstTile.y - srcTile.y;
     outputPass.setUniform('outputSize', size);
-    outputPass.setUniform('inputOffset', new Float32Array([dx, dy]));
+    outputPass.setUniform('inputOffset', new Int32Array([dx, dy]));
     outputPass.setUniform(
       'outputOffset',
-      new Float32Array([dstTile.x, dstTile.y])
+      new Int32Array([dstTile.x, dstTile.y])
     );
     outputPass.setExecuteSize(size[0], size[1]);
   }
@@ -490,5 +489,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,138 @@
+export interface DynamicTileOptions {
+  /** Smallest output tile edge. Defaults to 256. */
+  minTileSize?: number;
+  /** Initial output tile edge. Defaults to 384. */
+  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 128. */
+  adjustmentStep?: number;
+}
+
+export type DynamicTileSetting = boolean | DynamicTileOptions;
+
+const defaultMinTileSize = 256;
+const defaultInitialTileSize = 384;
+const defaultTargetTileTimeMs = 16;
+const defaultAdjustmentStep = 128;
+export const OIDN_TILE_ALIGNMENT = 16;
+
+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 clamp(value: number, min: number, max: number) {
+  return Math.min(Math.max(value, min), max);
+}
+
+function median(values: number[]) {
+  const sorted = [...values].sort((a, b) => a - b);
+  const middle = Math.floor(sorted.length / 2);
+  return sorted.length % 2
+    ? sorted[middle]
+    : (sorted[middle - 1] + sorted[middle]) / 2;
+}
+
+/** 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;
+
+  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;
+
+    const observedTime = median(validTimes);
+    let nextTileSize = this._tileSize;
+    if (observedTime > this.targetTileTimeMs * 1.25) {
+      nextTileSize -= this._adjustmentStep;
+    } else if (observedTime < 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;
+    return true;
+  }
+}
+
+/** Waits for all work submitted before this call and tolerates device loss. */
+export async function waitForSubmittedGPUWork(
+  queue: Pick<GPUQueue, 'onSubmittedWorkDone'>
+) {
+  try {
+    await queue.onSubmittedWorkDone();
+  } catch {
+    // tileExecute has no error callback. Preserve its completion/cancellation
+    // behavior and let the owning GPUDevice report device loss separately.
+  }
+}