oidn-web 0.2.2 → 0.3.0

package/src/WGPUComputePass.ts

@@ -13,7 +13,8 @@ export interface Uniform {
   data: Float32Array | Int32Array | Uint32Array;
 }
 export interface WGPUComputePassInput {
-  buffer: GPUBuffer;
+  buffer?: GPUBuffer;
+  texture?: GPUTexture;
   channels: number;
 }
 
@@ -22,6 +23,8 @@ export interface WGPUComputePassOutput {
 }
 
 export class WGPUComputePass<I extends string, O extends string> {
+  autoUpdateOutputBuffer = true;
+
   private _label;
 
   private _device;
@@ -36,6 +39,7 @@ export class WGPUComputePass<I extends string, O extends string> {
   private _pipeline!: GPUComputePipeline;
   private _bindGroups: GPUBindGroup[] = [];
   private _needsUpdatePipeline = true;
+  private _needsResizeBuffer = true;
 
   private _inputs: string[] = [];
   private _outputs: string[] = [];
@@ -52,6 +56,12 @@ export class WGPUComputePass<I extends string, O extends string> {
   private _csMain;
   private _csDefine;
 
+  private _groupOffsets = {
+    inputs: 0,
+    uniforms: 1,
+    outputs: 2
+  };
+
   constructor(
     label: string,
     device: GPUDevice,
@@ -61,6 +71,7 @@ export class WGPUComputePass<I extends string, O extends string> {
       csMain: string;
       csDefine?: string;
       uniforms: Uniform[];
+      autoUpdateOutputBuffer?: boolean;
     }
   ) {
     this._label = label;
@@ -70,6 +81,7 @@ export class WGPUComputePass<I extends string, O extends string> {
     this._inputs = opts.inputs;
     this._outputs = opts.outputs;
     this._uniforms = opts.uniforms;
+    this.autoUpdateOutputBuffer = opts.autoUpdateOutputBuffer ?? true;
 
     opts.uniforms.forEach((uniform) => {
       this._uniformBuffers[uniform.label] = device.createBuffer({
@@ -85,6 +97,12 @@ export class WGPUComputePass<I extends string, O extends string> {
     });
   }
 
+  setCSCode({ csDefine, csMain }: { csDefine: string; csMain: string }) {
+    this._csDefine = csDefine;
+    this._csMain = csMain;
+    this._needsUpdatePipeline = true;
+  }
+
   setSize(width: number, height: number) {
     width = Math.ceil(width);
     height = Math.ceil(height);
@@ -92,7 +110,7 @@ export class WGPUComputePass<I extends string, O extends string> {
     this._width = width;
     this._height = height;
     if (sizeChanged) {
-      this._resizeOutputBuffers();
+      this._needsResizeBuffer = true;
       this._needsUpdatePipeline = true;
     }
   }
@@ -105,16 +123,32 @@ export class WGPUComputePass<I extends string, O extends string> {
   }
 
   setOutputParams(outputParams: Record<O, WGPUComputePassOutput>) {
-    this._updateOutputBuffers(outputParams);
+    if (this.autoUpdateOutputBuffer) {
+      this._updateOutputBuffers(outputParams);
+    }
     this._needsUpdatePipeline = true;
   }
 
+  setOutputBuffers(outputBuffers: Record<O, GPUBuffer>) {
+    this._outputBuffers = Object.keys(outputBuffers).reduce((obj, key) => {
+      obj[key] = {
+        buffer: outputBuffers[key as O],
+        params: { channels: 4 }
+      };
+      return obj;
+    }, {} as WGPUComputePass<I, O>['_outputBuffers']);
+  }
+
   setUniform(label: string, data: Float32Array | Int32Array | Uint32Array) {
     const buffer = this._uniformBuffers[label];
     this._device.queue.writeBuffer(buffer, 0, data);
   }
 
-  getOutputBuffer(name: O) {
+  getOutput(name: O) {
+    if (this._needsResizeBuffer && this.autoUpdateOutputBuffer) {
+      this._resizeOutputBuffers();
+      this._needsResizeBuffer = false;
+    }
     return this._outputBuffers[name].buffer;
   }
 
@@ -123,7 +157,7 @@ export class WGPUComputePass<I extends string, O extends string> {
       (this._uniformBuffers as any)[key].destroy();
     });
     Object.keys(this._outputBuffers).forEach((key) => {
-      (this._outputBuffers as any)[key].texture.destroy();
+      (this._outputBuffers as any)[key].buffer.destroy();
     });
   }
 
@@ -173,13 +207,16 @@ export class WGPUComputePass<I extends string, O extends string> {
     }
   }
 
-  private _updatePipeline(inputParams: Record<string, WGPUComputePassInput>) {
+  private _updatePipeline(
+    inputParams: Record<string, WGPUComputePassInput>,
+    inputTypes: Record<string, 'texture' | 'buffer'>
+  ) {
     if (!this._needsUpdatePipeline) {
       return;
     }
     this._needsUpdatePipeline = false;
     const device = this._device;
-    const csCode = this._getFullCs(inputParams);
+    const csCode = this._getFullCs(inputParams, inputTypes);
     if (csCode === this._csCode) {
       return;
     }
@@ -198,35 +235,49 @@ export class WGPUComputePass<I extends string, O extends string> {
     this._updateBindGroups();
   }
 
-  private _getFullCs(inputParams: Record<string, WGPUComputePassInput>) {
+  private _getFullCs(
+    inputParams: Record<string, WGPUComputePassInput>,
+    inputTypes: Record<string, 'texture' | 'buffer'>
+  ) {
     const inputs = this._inputs;
-    const hasInputs = inputs.length > 0;
+    const uniforms = this._uniforms;
+    let offset = 0;
+    const groupOffsets = (this._groupOffsets = {
+      inputs: 0,
+      uniforms: 0,
+      outputs: 0
+    });
+    if (inputs.length > 0) {
+      offset++;
+    }
+    if (uniforms.length > 0) {
+      groupOffsets.uniforms = offset;
+      offset++;
+    }
+    groupOffsets.outputs = offset;
     const cs = `
 ${inputs
   .sort()
-  .map(
-    (bufferName, idx) =>
-      // TODO more channels option.
-      `@group(0) @binding(${idx}) var<storage, read> in_${bufferName}: array<vec${inputParams[bufferName].channels}f>;`
-  )
+  .map((inputName, idx) => {
+    const bindingPrefix = `@group(${groupOffsets.inputs}) @binding(${idx}) `;
+    const varName = `in_${inputName}`;
+    // TODO more channels option.
+    return inputTypes[inputName] === 'texture'
+      ? `${bindingPrefix} var ${varName}: texture_2d<f32>;`
+      : `${bindingPrefix} var<storage, read> ${varName}: array<vec${inputParams[inputName].channels}f>;`;
+  })
   .join('\n')}
 ${this._uniforms
   .map(
     (uniform, idx) =>
-      `@group(${hasInputs ? 1 : 0}) @binding(${idx}) var<uniform> ${
-        uniform.label
-      }: ${uniform.type};`
+      `@group(${groupOffsets.uniforms}) @binding(${idx}) var<uniform> ${uniform.label}: ${uniform.type};`
   )
   .join('\n')}
 
 ${this._outputs
   .map(
     (name, idx) =>
-      `@group(${
-        hasInputs ? 2 : 1
-      }) @binding(${idx}) var<storage, read_write> out_${name}: array<vec${
-        this._outputBuffers[name].params.channels
-      }f>;`
+      `@group(${groupOffsets.outputs}) @binding(${idx}) var<storage, read_write> out_${name}: array<vec${this._outputBuffers[name].params.channels}f>;`
   )
   .join('\n')}
 ${this._csDefine ?? ''}
@@ -242,13 +293,12 @@ ${this._csMain}
   private _updateBindGroups() {
     const bindGroups: GPUBindGroup[] = [];
     const device = this._device;
+    const groupOffsets = this._groupOffsets;
 
-    //TODO
-    const uniformBindGroupIndex = this._inputs.length > 0 ? 1 : 0;
     if (this._uniforms.length > 0) {
-      bindGroups[uniformBindGroupIndex] = device.createBindGroup({
+      bindGroups[groupOffsets.uniforms] = device.createBindGroup({
         label: this._label,
-        layout: this._pipeline.getBindGroupLayout(uniformBindGroupIndex),
+        layout: this._pipeline.getBindGroupLayout(groupOffsets.uniforms),
         entries: this._uniforms.map(
           (uniform, idx) =>
             ({
@@ -266,30 +316,42 @@ ${this._csMain}
 
   createPass(
     commandEncoder: GPUCommandEncoder,
-    inputBuffers: Record<I, WGPUComputePassInput>
+    inputs: Record<I, WGPUComputePassInput>
   ) {
-    this._updatePipeline(inputBuffers);
+    if (this._needsResizeBuffer && this.autoUpdateOutputBuffer) {
+      this._resizeOutputBuffers();
+      this._needsResizeBuffer = false;
+    }
+
+    const inputTypes = this._inputs.reduce((obj, inputName) => {
+      obj[inputName] = inputs[inputName as I].buffer ? 'buffer' : 'texture';
+      return obj;
+    }, {} as Record<string, 'texture' | 'buffer'>);
 
-    const hasInputs = this._inputs.length > 0;
+    this._updatePipeline(inputs, inputTypes);
+
+    const groupOffsets = this._groupOffsets;
     // TODO createBindGroup every time?
-    if (hasInputs) {
-      this._bindGroups[0] = this._device.createBindGroup({
+    if (this._inputs.length > 0) {
+      this._bindGroups[groupOffsets.inputs] = this._device.createBindGroup({
         label: this._label,
-        layout: this._pipeline.getBindGroupLayout(0),
-        entries: this._inputs.map((bufferName, idx) => ({
+        layout: this._pipeline.getBindGroupLayout(groupOffsets.inputs),
+        entries: this._inputs.map((inputName, idx) => ({
           binding: idx,
           // TODO
-          resource: {
-            buffer: inputBuffers[bufferName as I].buffer
-          }
+          resource: inputs[inputName as I].buffer
+            ? {
+                buffer: inputs[inputName as I].buffer!
+              }
+            : inputs[inputName as I].texture!.createView()
         }))
       });
     }
 
     // Outputs
-    this._bindGroups[hasInputs ? 2 : 1] = this._device.createBindGroup({
+    this._bindGroups[groupOffsets.outputs] = this._device.createBindGroup({
       label: this._label,
-      layout: this._pipeline.getBindGroupLayout(hasInputs ? 2 : 1),
+      layout: this._pipeline.getBindGroupLayout(groupOffsets.outputs),
       entries: this._outputs.map((bufferName, idx) => ({
         binding: idx,
         resource: {

package/src/backend.ts

@@ -45,9 +45,15 @@ export async function initWebGPUBackendWithDevice(
   device: GPUDevice,
   adapter: GPUAdapterInfo
 ) {
-  const backend = new WebGPUBackend(device, adapter);
+  // TODO multiple device and adapter in one backend
+  let backend = ENGINE.findBackend('webgpu-oidn');
+  if (backend != null) {
+    return backend as WebGPUBackend;
+  }
+
+  backend = new WebGPUBackend(device, adapter);
   ENGINE.registerBackend('webgpu-oidn', () => backend);
   await ENGINE.setBackend('webgpu-oidn');
 
-  return backend;
+  return backend as WebGPUBackend;
 }

package/src/helper.ts

@@ -33,3 +33,11 @@ export function profileAndLogKernelCode(execute: () => void, disabled = true) {
       console.log(code);
     });
 }
+
+export function memory() {
+  return tfjs.memory();
+}
+
+export function tidy(f: () => void) {
+  return tfjs.tidy(f);
+}