oidn-webgpu 0.0.0-stage → 0.6.0

package/src/process.ts

@@ -0,0 +1,368 @@
+import { WGPUComputePass } from './WGPUComputePass';
+import {
+  hdrForwardWGSL,
+  hdrInverseWGSL,
+  hdrTransferFuncCPU,
+  hdrTransferFuncInverseCPU,
+  type HDRTransfer
+} from './hdrTransfer';
+import { sanitizeWGSL } from './sanitize';
+export { avgLogLum, hdrTransferFuncCPU, hdrTransferFuncInverseCPU } from './hdrTransfer';
+export type { HDRTransfer } from './hdrTransfer';
+
+export class Tile {
+  constructor(
+    public x: number,
+    public y: number,
+    public width: number,
+    public height: number
+  ) {}
+}
+
+export class GPUDataProcess {
+  private _inputPassAux;
+  private _inputPassColor;
+  private _outputPass;
+  private _copyPass;
+
+  private _isInputTexture?: boolean;
+
+  constructor(
+    private _device: GPUDevice,
+    private _isHDR: boolean,
+    private _hdrTransfer: HDRTransfer = 'pu'
+  ) {
+    const commonUniforms = [
+      {
+        label: 'inputScale',
+        type: 'f32',
+        data: new Float32Array([1])
+      },
+      {
+        label: 'inputSize',
+        type: 'vec2i',
+        data: new Int32Array(2)
+      },
+      {
+        label: 'outputSize',
+        type: 'vec2i',
+        data: new Int32Array(2)
+      },
+      {
+        label: 'inputOffset',
+        type: 'vec2i',
+        data: new Int32Array(2)
+      }
+    ];
+    this._inputPassAux = new WGPUComputePass('inputPassAux', this._device, {
+      inputs: ['color', 'albedo', 'normal'],
+      outputs: ['color', 'albedo', 'normal'],
+      uniforms: commonUniforms,
+      csDefine: '',
+      csMain: ``
+    });
+    this._inputPassColor = new WGPUComputePass('inputPassColor', this._device, {
+      inputs: ['color'],
+      outputs: ['color'],
+      uniforms: commonUniforms,
+      csDefine: '',
+      csMain: ``
+    });
+    this._outputPass = new WGPUComputePass('outputPass', this._device, {
+      inputs: ['color', 'raw'],
+      outputs: ['color'],
+      uniforms: [
+        {
+          label: 'inputScale',
+          type: 'f32',
+          data: new Float32Array([1])
+        },
+        {
+          label: 'inputSize',
+          type: 'vec2i',
+          data: new Int32Array(2)
+        },
+        {
+          label: 'outputSize',
+          type: 'vec2i',
+          data: new Int32Array(2)
+        },
+        {
+          label: 'imageSize',
+          type: 'vec2i',
+          data: new Int32Array(2)
+        },
+        {
+          label: 'inputOffset',
+          type: 'vec2i',
+          data: new Int32Array(2)
+        },
+        {
+          label: 'outputOffset',
+          type: 'vec2i',
+          data: new Int32Array(2)
+        }
+      ],
+      csDefine: '',
+      csMain: ``
+    });
+    this._copyPass = new WGPUComputePass('copyPass', this._device, {
+      inputs: ['color'],
+      outputs: ['color'],
+      autoUpdateOutputBuffer: false,
+      uniforms: [
+        {
+          label: 'size',
+          type: 'vec2i',
+          data: new Int32Array(2)
+        }
+      ],
+      csMain: /*wgsl*/ `
+let outIdx = i32(globalId.x + globalId.y * u32(size.x));
+out_color[outIdx] = textureLoad(in_color, globalId.xy, 0);
+`
+    });
+    this._inputPassAux.setOutputParams({
+      color: { channels: 3 },
+      albedo: { channels: 3 },
+      normal: { channels: 3 }
+    });
+    this._inputPassColor.setOutputParams({
+      color: { channels: 3 }
+    });
+    this._outputPass.setOutputParams({
+      color: { channels: 4 }
+    });
+  }
+
+  private _updatePasses(isInputTexture: boolean, denoiseAlpha = false) {
+    if (
+      this._isInputTexture != undefined &&
+      this._isInputTexture === isInputTexture
+    ) {
+      return;
+    }
+
+    this._isInputTexture = isInputTexture;
+    const isHDR = this._isHDR;
+    const commonCSDefine = hdrForwardWGSL(this._hdrTransfer) + sanitizeWGSL;
+    function readInputCode(inputName: string) {
+      return isInputTexture
+        ? `textureLoad(in_${inputName}, vec2u(inputPosition), 0)`
+        : `in_${inputName}[inIdx]`;
+    }
+    const commonCSMain = /* wgsl */ `
+let x = i32(globalId.x);
+let y = i32(globalId.y);
+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;
+
+if (${denoiseAlpha}) {
+  // Denoise the inversed alpha. Or the anti aliased edge will be too dark after denoised
+  out_color[outIdx] = sanitize(vec3f(1.0 - col.a), 0.0, 1.0);
+}
+else if (${isHDR}) {
+  let scaled = sanitize(col.rgb * inputScale, 0.0, FLT_MAX);
+  out_color[outIdx] = vec3f(HDRForward(scaled.r), HDRForward(scaled.g), HDRForward(scaled.b));
+}
+else {
+  out_color[outIdx] = sanitize(col.rgb, 0.0, 1.0);
+}
+`;
+    this._inputPassAux.setCSCode({
+      csDefine: commonCSDefine,
+      csMain: /* wgsl */ `
+${commonCSMain}
+let alb = ${readInputCode('albedo')};
+let nor = ${readInputCode('normal')};
+// Normals arrive [0, 1]-encoded, so OIDN's [-1, 1] clamp becomes [0, 1].
+out_normal[outIdx] = sanitize(nor.rgb, 0.0, 1.0);
+out_albedo[outIdx] = sanitize(alb.rgb, 0.0, 1.0);
+  `
+    });
+
+    this._inputPassColor.setCSCode({
+      csDefine: commonCSDefine,
+      csMain: /* wgsl */ `
+${commonCSMain}
+`
+    });
+    this._outputPass.setCSCode({
+      csDefine: hdrInverseWGSL(this._hdrTransfer) + sanitizeWGSL,
+      csMain: /* wgsl */ `
+let x = i32(globalId.x);
+let y = i32(globalId.y);
+if (x >= outputSize.x || y >= outputSize.y) {
+  return;
+}
+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
+          ? 'textureLoad(in_raw, globalId.xy + vec2u(outputOffset), 0)'
+          : 'in_raw[outIdx]'
+      };
+
+// The network output may contain negative values or NaNs.
+let denoised = sanitize(col.rgb, 0.0, ${isHDR && !denoiseAlpha ? 'FLT_MAX' : '1.0'});
+if (${denoiseAlpha}) {
+  out_color[outIdx] = vec4f(raw.rgb, 1.0 - denoised.r);
+}
+else if (${isHDR}) {
+  out_color[outIdx] = vec4f(
+    vec3f(HDRInverse(denoised.r), HDRInverse(denoised.g), HDRInverse(denoised.b)) / inputScale,
+    // Pick the alpha
+    raw.a
+  );
+}
+else {
+  out_color[outIdx] = vec4f(denoised, raw.a);
+}
+`
+    });
+  }
+
+  setImageSize(w: number, h: number) {
+    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 Int32Array([w, h]));
+  }
+
+  setInputTile(tile: Tile) {
+    const size = new Int32Array([tile.width, tile.height]);
+    [this._inputPassAux, this._inputPassColor].forEach((inputPass) => {
+      inputPass.setUniform('inputOffset', new Int32Array([tile.x, tile.y]));
+      inputPass.setUniform('outputSize', size);
+      inputPass.setSize(size[0], size[1]);
+    });
+
+    this._outputPass.setUniform('inputSize', size);
+  }
+
+  setOutputTile(dstTile: Tile, srcTile: Tile) {
+    const outputPass = this._outputPass;
+    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 Int32Array([dx, dy]));
+    outputPass.setUniform(
+      'outputOffset',
+      new Int32Array([dstTile.x, dstTile.y])
+    );
+    outputPass.setExecuteSize(size[0], size[1]);
+  }
+
+  forward(
+    colorBuffer: GPUBuffer | GPUTexture,
+    // TODO optional albedo and normal.
+    albedoBuffer: GPUBuffer | GPUTexture | undefined,
+    normalBuffer: GPUBuffer | GPUTexture | undefined,
+
+    denoiseAlpha: boolean | undefined
+  ) {
+    const isInputTexture = colorBuffer instanceof GPUTexture;
+    this._updatePasses(isInputTexture, denoiseAlpha);
+
+    const inputPassAux = this._inputPassAux;
+    const inputPassColor = this._inputPassColor;
+    const commandEncoder = this._device.createCommandEncoder();
+
+    function createInput(bufferOrTex: GPUBuffer | GPUTexture) {
+      return bufferOrTex instanceof GPUTexture
+        ? {
+            texture: bufferOrTex,
+            channels: 4
+          }
+        : {
+            buffer: bufferOrTex,
+            channels: 4
+          };
+    }
+    if (albedoBuffer && normalBuffer) {
+      inputPassAux.createPass(commandEncoder, {
+        color: createInput(colorBuffer),
+        albedo: createInput(albedoBuffer),
+        normal: createInput(normalBuffer)
+      });
+    } else {
+      inputPassColor.createPass(commandEncoder, {
+        color: createInput(colorBuffer)
+      });
+    }
+    this._device.queue.submit([commandEncoder.finish()]);
+
+    return albedoBuffer && normalBuffer
+      ? {
+          color: inputPassAux.getOutput('color'),
+          albedo: inputPassAux.getOutput('albedo'),
+          normal: inputPassAux.getOutput('normal')
+        }
+      : {
+          color: inputPassColor.getOutput('color')
+        };
+  }
+
+  inverse(buffer: GPUBuffer, raw: GPUBuffer | GPUTexture) {
+    const device = this._device;
+
+    const commandEncoder = device.createCommandEncoder();
+    const outputGPUPass = this._outputPass;
+
+    outputGPUPass.createPass(commandEncoder, {
+      color: { buffer: buffer, channels: 4 },
+      raw:
+        raw instanceof GPUBuffer
+          ? { buffer: raw, channels: 4 }
+          : { texture: raw, channels: 4 }
+    });
+    this._device.queue.submit([commandEncoder.finish()]);
+
+    return outputGPUPass.getOutput('color');
+  }
+
+  copyInputDataToOutput(inputColorBuffer: GPUBuffer | GPUTexture) {
+    const encoder = this._device.createCommandEncoder();
+    const outputGPUPass = this._outputPass;
+    const colorBuffer = outputGPUPass.getOutput('color');
+    const copyPass = this._copyPass;
+
+    if (inputColorBuffer instanceof GPUTexture) {
+      copyPass.setOutputBuffers({
+        color: colorBuffer
+      });
+      copyPass.createPass(encoder, {
+        color: { texture: inputColorBuffer, channels: 4 }
+      });
+    } else {
+      encoder.copyBufferToBuffer(
+        inputColorBuffer,
+        0,
+        colorBuffer,
+        0,
+        colorBuffer.size
+      );
+    }
+
+    this._device.queue.submit([encoder.finish()]);
+  }
+
+  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/sanitize.ts

@@ -0,0 +1,29 @@
+/**
+ * OIDN's input and output sanitizing: NaN becomes 0, then values are clamped
+ * to their valid range. HDR values are bounded by FLT_MAX, so +inf becomes
+ * FLT_MAX. Used before the network, by auto-exposure, and on the network
+ * output before the inverse transfer.
+ */
+export const FLT_MAX = 3.4028234663852886e38;
+
+export function sanitize(value: number, lo: number, hi: number) {
+  return Number.isNaN(value) ? 0 : Math.min(Math.max(value, lo), hi);
+}
+
+/**
+ * WGSL declarations of `FLT_MAX` and `fn sanitize(v: vec3f, lo: f32, hi: f32)`.
+ * WGSL compilers may assume floats are finite, so NaN and infinity are
+ * detected from the bits rather than with comparisons. FLT_MAX is spelled out
+ * so the native core does not depend on scriptc's number formatting.
+ */
+export const sanitizeWGSL = /* wgsl */ `
+const FLT_MAX = 3.4028234663852886e38;
+fn sanitize(v: vec3f, lo: f32, hi: f32) -> vec3f {
+  let bits = bitcast<vec3u>(v);
+  let magnitude = bits & vec3u(0x7fffffffu);
+  let negative = (bits & vec3u(0x80000000u)) != vec3u(0u);
+  let infinite = select(vec3f(hi), vec3f(lo), negative);
+  let clamped = select(clamp(v, vec3f(lo), vec3f(hi)), infinite, magnitude == vec3u(0x7f800000u));
+  return select(clamped, vec3f(0.0), magnitude > vec3u(0x7f800000u));
+}
+`;