oidn-web 0.1.0

package/src/WGPUComputePass.ts

@@ -0,0 +1,318 @@
+function isStorageParamsEqual(
+  params: WGPUComputePassOutput,
+  other: WGPUComputePassOutput
+) {
+  return params.channels === other.channels;
+}
+
+export const WORKGROUP_SIZE = 8;
+
+export interface Uniform {
+  label: string;
+  type: string;
+  data: Float32Array | Int32Array | Uint32Array;
+}
+export interface WGPUComputePassInput {
+  buffer: GPUBuffer;
+  channels: number;
+}
+
+export interface WGPUComputePassOutput {
+  channels: number;
+}
+
+export class WGPUComputePass<I extends string, O extends string> {
+  private _label;
+
+  private _device;
+  private _outputBuffers: Record<
+    string,
+    {
+      buffer: GPUBuffer;
+      params: WGPUComputePassOutput;
+    }
+  > = {};
+
+  private _pipeline!: GPUComputePipeline;
+  private _bindGroups: GPUBindGroup[] = [];
+  private _needsUpdatePipeline = true;
+
+  private _inputs: string[] = [];
+  private _outputs: string[] = [];
+  private _uniforms: Uniform[] = [];
+  private _uniformBuffers: Record<string, GPUBuffer> = {};
+
+  private _width = 10;
+  private _height = 10;
+
+  private _execWidth?: number;
+  private _execHeight?: number;
+
+  private _csCode = '';
+  private _csMain;
+  private _csDefine;
+
+  constructor(
+    label: string,
+    device: GPUDevice,
+    opts: {
+      inputs: I[];
+      outputs: O[];
+      csMain: string;
+      csDefine?: string;
+      uniforms: Uniform[];
+    }
+  ) {
+    this._label = label;
+    this._device = device;
+    this._csMain = opts.csMain;
+    this._csDefine = opts.csDefine;
+    this._inputs = opts.inputs;
+    this._outputs = opts.outputs;
+    this._uniforms = opts.uniforms;
+
+    opts.uniforms.forEach((uniform) => {
+      this._uniformBuffers[uniform.label] = device.createBuffer({
+        label: this._label,
+        size: uniform.data.byteLength,
+        usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST
+      });
+      this._device.queue.writeBuffer(
+        this._uniformBuffers[uniform.label],
+        0,
+        uniform.data
+      );
+    });
+  }
+
+  setSize(width: number, height: number) {
+    width = Math.ceil(width);
+    height = Math.ceil(height);
+    const sizeChanged = width !== this._width || height !== this._height;
+    this._width = width;
+    this._height = height;
+    if (sizeChanged) {
+      this._resizeOutputBuffers();
+      this._needsUpdatePipeline = true;
+    }
+  }
+
+  setExecuteSize(width: number, height: number) {
+    width = Math.ceil(width);
+    height = Math.ceil(height);
+    this._execWidth = width;
+    this._execHeight = height;
+  }
+
+  setOutputParams(outputParams: Record<O, WGPUComputePassOutput>) {
+    this._updateOutputBuffers(outputParams);
+    this._needsUpdatePipeline = true;
+  }
+
+  setUniform(label: string, data: Float32Array | Int32Array | Uint32Array) {
+    const buffer = this._uniformBuffers[label];
+    this._device.queue.writeBuffer(buffer, 0, data);
+  }
+
+  getOutputBuffer(name: O) {
+    return this._outputBuffers[name].buffer;
+  }
+
+  dispose() {
+    Object.keys(this._uniformBuffers).forEach((key) => {
+      (this._uniformBuffers as any)[key].destroy();
+    });
+    Object.keys(this._outputBuffers).forEach((key) => {
+      (this._outputBuffers as any)[key].texture.destroy();
+    });
+  }
+
+  private _createBuffer(params: WGPUComputePassOutput) {
+    // const byteLength = this._width * this._height * params.channels * 4;
+    // Buffer data needs to be aligned with 8, 16
+    const byteLength = this._width * this._height * 4 * 4;
+    return this._device.createBuffer({
+      label: this._label,
+      // webgpu needs buffer at least 80 bytes.
+      size: Math.max(byteLength, 80),
+      usage:
+        GPUBufferUsage.STORAGE |
+        GPUBufferUsage.COPY_DST |
+        GPUBufferUsage.COPY_SRC
+    });
+  }
+
+  private _resizeOutputBuffers() {
+    const outputBuffers = this._outputBuffers;
+    for (const key in outputBuffers) {
+      const { buffer, params } = outputBuffers[key];
+      buffer.destroy();
+      outputBuffers[key].buffer = this._createBuffer(params);
+    }
+  }
+
+  private _updateOutputBuffers(
+    outputParams: Record<string, WGPUComputePassOutput>
+  ) {
+    const outputBuffers = this._outputBuffers;
+    for (const key in outputParams) {
+      const params = outputParams[key];
+      if (
+        !isStorageParamsEqual(
+          params,
+          outputBuffers[key]?.params || ({} as WGPUComputePassOutput)
+        )
+      ) {
+        outputBuffers[key]?.buffer.destroy();
+        const buffer = this._createBuffer(params);
+        outputBuffers[key] = {
+          buffer,
+          params
+        };
+      }
+    }
+  }
+
+  private _updatePipeline(inputParams: Record<string, WGPUComputePassInput>) {
+    if (!this._needsUpdatePipeline) {
+      return;
+    }
+    this._needsUpdatePipeline = false;
+    const device = this._device;
+    const csCode = this._getFullCs(inputParams);
+    if (csCode === this._csCode) {
+      return;
+    }
+    this._csCode = csCode;
+    this._pipeline = device.createComputePipeline({
+      label: this._label,
+      layout: 'auto',
+      compute: {
+        module: device.createShaderModule({
+          label: this._label,
+          code: csCode
+        }),
+        entryPoint: 'main'
+      }
+    });
+    this._updateBindGroups();
+  }
+
+  private _getFullCs(inputParams: Record<string, WGPUComputePassInput>) {
+    const inputs = this._inputs;
+    const hasInputs = inputs.length > 0;
+    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>;`
+  )
+  .join('\n')}
+${this._uniforms
+  .map(
+    (uniform, idx) =>
+      `@group(${hasInputs ? 1 : 0}) @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>;`
+  )
+  .join('\n')}
+${this._csDefine ?? ''}
+@compute @workgroup_size(${WORKGROUP_SIZE}, ${WORKGROUP_SIZE}, 1)
+fn main(@builtin(global_invocation_id) globalId: vec3u) {
+${this._csMain}
+}
+`;
+
+    return cs;
+  }
+
+  private _updateBindGroups() {
+    const bindGroups: GPUBindGroup[] = [];
+    const device = this._device;
+
+    //TODO
+    const uniformBindGroupIndex = this._inputs.length > 0 ? 1 : 0;
+    if (this._uniforms.length > 0) {
+      bindGroups[uniformBindGroupIndex] = device.createBindGroup({
+        label: this._label,
+        layout: this._pipeline.getBindGroupLayout(uniformBindGroupIndex),
+        entries: this._uniforms.map(
+          (uniform, idx) =>
+            ({
+              binding: idx,
+              resource: {
+                buffer: this._uniformBuffers[uniform.label]
+              }
+            } as GPUBindGroupEntry)
+        )
+      });
+    }
+
+    this._bindGroups = bindGroups;
+  }
+
+  createPass(
+    commandEncoder: GPUCommandEncoder,
+    inputBuffers: Record<I, WGPUComputePassInput>
+  ) {
+    this._updatePipeline(inputBuffers);
+
+    const hasInputs = this._inputs.length > 0;
+    // TODO createBindGroup every time?
+    if (hasInputs) {
+      this._bindGroups[0] = this._device.createBindGroup({
+        label: this._label,
+        layout: this._pipeline.getBindGroupLayout(0),
+        entries: this._inputs.map((bufferName, idx) => ({
+          binding: idx,
+          // TODO
+          resource: {
+            buffer: inputBuffers[bufferName as I].buffer
+          }
+        }))
+      });
+    }
+
+    // Outputs
+    this._bindGroups[hasInputs ? 2 : 1] = this._device.createBindGroup({
+      label: this._label,
+      layout: this._pipeline.getBindGroupLayout(hasInputs ? 2 : 1),
+      entries: this._outputs.map((bufferName, idx) => ({
+        binding: idx,
+        resource: {
+          buffer: this._outputBuffers[bufferName].buffer
+        }
+      }))
+    });
+
+    // Begin the render pass
+    const computePass = commandEncoder.beginComputePass();
+
+    // Draw a full quad
+    computePass.setPipeline(this._pipeline);
+    // Bind groups
+    this._bindGroups.forEach((bindGroup, idx) => {
+      computePass.setBindGroup(idx, bindGroup);
+    });
+    computePass.dispatchWorkgroups(
+      Math.ceil((this._execWidth ?? this._width) / WORKGROUP_SIZE),
+      Math.ceil((this._execHeight ?? this._height) / WORKGROUP_SIZE),
+      1
+    );
+    // End the render pass
+    computePass.end();
+  }
+}

package/src/WGPUFullQuadPass.ts

@@ -0,0 +1,348 @@
+export const fullScreenQuadVertexShaderWGSL = /*wgsl */ `
+@vertex
+fn main(
+  @builtin(vertex_index) vertexIndex: u32
+) -> @builtin(position) vec4f {
+  const pos = array(
+    vec2(-1.0, -1.0), vec2(1.0, -1.0), vec2(-1.0, 1.0),
+    vec2(-1.0, 1.0), vec2(1.0, -1.0), vec2(1.0, 1.0),
+  );
+
+  return vec4<f32>(pos[vertexIndex], 0.0, 1.0);
+}
+`;
+
+function isTextureParamsEqual(
+  params: WGPUFullQuadPassOutput,
+  other: WGPUFullQuadPassOutput
+) {
+  return params.format === other.format;
+}
+
+export interface Uniform {
+  label: string;
+  type: string;
+  data: Float32Array | Int32Array | Uint32Array;
+}
+
+export interface WGPUFullQuadPassOutput {
+  format: GPUTextureFormat;
+}
+export class WGPUFullQuadPass<I extends string, O extends string> {
+  private _label;
+
+  private _device;
+  private _outputTextures: Record<
+    string,
+    {
+      texture: GPUTexture;
+      params: WGPUFullQuadPassOutput;
+    }
+  > = {};
+
+  private _pipeline!: GPURenderPipeline;
+  private _bindGroups: GPUBindGroup[] = [];
+  private _needsUpdatePipeline = true;
+  /**
+   * When render to canvas
+   */
+  private _renderToScreen?: {
+    screenTexture: GPUTexture;
+    presentationFormat: GPUTextureFormat;
+  };
+  private _inputs: string[] = [];
+  private _outputs: string[] = [];
+  private _uniforms: Uniform[] = [];
+  private _uniformBuffers: Record<string, GPUBuffer> = {};
+
+  private _width = 10;
+  private _height = 10;
+
+  private _fsCode = '';
+  private _fsMain;
+  private _fsDefine;
+
+  constructor(
+    label: string,
+    device: GPUDevice,
+    opts: {
+      inputs: I[];
+      outputs: O[];
+      fsMain: string;
+      fsDefine?: string;
+      uniforms: Uniform[];
+    }
+  ) {
+    this._label = label;
+    this._device = device;
+    this._fsMain = opts.fsMain;
+    this._fsDefine = opts.fsDefine;
+    this._inputs = opts.inputs;
+    this._outputs = opts.outputs;
+    this._uniforms = opts.uniforms;
+
+    opts.uniforms.forEach((uniform) => {
+      this._uniformBuffers[uniform.label] = device.createBuffer({
+        label: this._label,
+        size: uniform.data.byteLength,
+        usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST
+      });
+      this._device.queue.writeBuffer(
+        this._uniformBuffers[uniform.label],
+        0,
+        uniform.data
+      );
+    });
+  }
+
+  setSize(width: number, height: number) {
+    width = Math.ceil(width);
+    height = Math.ceil(height);
+    const sizeChanged = width !== this._width || height !== this._height;
+    this._width = width;
+    this._height = height;
+    if (sizeChanged) {
+      this._resizeOutputTextures();
+      this._needsUpdatePipeline = true;
+    }
+  }
+
+  setOutputParams(outputParams: Record<O, WGPUFullQuadPassOutput>) {
+    this._renderToScreen = undefined;
+    this._updateOutputTextures(outputParams);
+    this._needsUpdatePipeline = true;
+  }
+
+  setRenderToScreen(
+    screenTexture: GPUTexture,
+    presentationFormat: GPUTextureFormat
+  ) {
+    this._renderToScreen = {
+      screenTexture,
+      presentationFormat
+    };
+  }
+
+  setUniform(label: string, data: Float32Array | Int32Array | Uint32Array) {
+    const buffer = this._uniformBuffers[label];
+    this._device.queue.writeBuffer(buffer, 0, data);
+  }
+
+  getOutputTexture(name: O) {
+    return this._outputTextures[name].texture;
+  }
+
+  dispose() {
+    Object.keys(this._uniformBuffers).forEach((key) => {
+      (this._uniformBuffers as any)[key].destroy();
+    });
+    Object.keys(this._outputTextures).forEach((key) => {
+      (this._outputTextures as any)[key].texture.destroy();
+    });
+  }
+
+  private _createTexture(params: WGPUFullQuadPassOutput) {
+    return this._device.createTexture({
+      label: this._label,
+      size: {
+        width: this._width,
+        height: this._height,
+        depthOrArrayLayers: 1
+      },
+      format: params.format,
+      usage: GPUTextureUsage.RENDER_ATTACHMENT | GPUTextureUsage.TEXTURE_BINDING
+    });
+  }
+
+  private _resizeOutputTextures() {
+    const outputTextures = this._outputTextures;
+    for (const key in outputTextures) {
+      const { texture, params } = outputTextures[key];
+      texture.destroy();
+      outputTextures[key].texture = this._createTexture(params);
+    }
+  }
+
+  private _updateOutputTextures(
+    outputParams: Record<string, WGPUFullQuadPassOutput>
+  ) {
+    const outputTextures = this._outputTextures;
+    for (const key in outputParams) {
+      const params = outputParams[key];
+      if (
+        !isTextureParamsEqual(
+          params,
+          outputTextures[key]?.params || ({} as WGPUFullQuadPassOutput)
+        )
+      ) {
+        outputTextures[key]?.texture.destroy();
+        const texture = this._createTexture(params);
+        outputTextures[key] = {
+          texture,
+          params
+        };
+      }
+    }
+  }
+
+  private _updatePipeline() {
+    if (!this._needsUpdatePipeline) {
+      return;
+    }
+    this._needsUpdatePipeline = false;
+    const device = this._device;
+    const fsCode = this._getFullFs();
+    if (fsCode === this._fsCode) {
+      return;
+    }
+    const { screenTexture, presentationFormat } = this._renderToScreen || {};
+    this._fsCode = fsCode;
+    this._pipeline = device.createRenderPipeline({
+      label: this._label,
+      layout: 'auto',
+      vertex: {
+        module: device.createShaderModule({
+          label: this._label,
+          code: fullScreenQuadVertexShaderWGSL
+        }),
+        entryPoint: 'main'
+      },
+      fragment: {
+        module: device.createShaderModule({
+          label: this._label,
+          code: fsCode
+        }),
+        entryPoint: 'main',
+        targets: screenTexture
+          ? [
+              {
+                format: presentationFormat!
+              }
+            ]
+          : this._outputs.map((key) => ({
+              format: this._outputTextures[key].params.format
+            }))
+      },
+      primitive: {
+        topology: 'triangle-list'
+      }
+    });
+    this._updateBindGroups();
+  }
+
+  private _getFullFs() {
+    const inputs = this._inputs;
+    const hasInputs = inputs.length > 0;
+    const fs = `
+${inputs
+  .sort()
+  .map(
+    (textureName, idx) =>
+      `@group(0) @binding(${idx}) var ${textureName}: texture_2d<f32>;`
+  )
+  .join('\n')}
+${this._uniforms
+  .map(
+    (uniform, idx) =>
+      `@group(${hasInputs ? 1 : 0}) @binding(${idx}) var<uniform> ${
+        uniform.label
+      }: ${uniform.type};`
+  )
+  .join('\n')}
+
+struct FSOutput {
+${this._outputs
+  .map((name, idx) => `@location(${idx}) ${name}: vec4f,`)
+  .join('\n')}
+}
+${this._fsDefine ?? ''}
+@fragment
+fn main(
+  @builtin(position) coord: vec4f
+) -> FSOutput {
+  var uv = vec2i(floor(coord.xy));
+  var output: FSOutput;
+${this._fsMain}
+  return output;
+}
+`;
+
+    return fs;
+  }
+
+  private _updateBindGroups() {
+    const bindGroups: GPUBindGroup[] = [];
+    const device = this._device;
+
+    //TODO
+    const uniformBindGroupIndex = this._inputs.length > 0 ? 1 : 0;
+    if (this._uniforms.length > 0) {
+      bindGroups[uniformBindGroupIndex] = device.createBindGroup({
+        label: this._label,
+        layout: this._pipeline.getBindGroupLayout(uniformBindGroupIndex),
+        entries: this._uniforms.map(
+          (uniform, idx) =>
+            ({
+              binding: idx,
+              resource: {
+                buffer: this._uniformBuffers[uniform.label]
+              }
+            } satisfies GPUBindGroupEntry)
+        )
+      });
+    }
+
+    this._bindGroups = bindGroups;
+  }
+
+  createPass(
+    commandEncoder: GPUCommandEncoder,
+    inputTextures: Record<I, GPUTexture>
+  ) {
+    this._updatePipeline();
+
+    // TODO createBindGrou every time?
+    if (this._inputs.length > 0) {
+      this._bindGroups[0] = this._device.createBindGroup({
+        label: this._label,
+        layout: this._pipeline.getBindGroupLayout(0),
+        entries: this._inputs.map((textureName, idx) => ({
+          binding: idx,
+          // TODO
+          resource: inputTextures[textureName as I].createView()
+        }))
+      });
+    }
+    // Begin the render pass
+    const renderPass = commandEncoder.beginRenderPass({
+      colorAttachments: this._renderToScreen
+        ? [
+            {
+              view: this._renderToScreen.screenTexture.createView(),
+              clearValue: { r: 0, g: 0, b: 0, a: 0 },
+              storeOp: 'store' as GPUStoreOp,
+              loadOp: 'clear' as GPULoadOp
+            }
+          ]
+        : this._outputs.map(
+            (textureName) =>
+              ({
+                view: this._outputTextures[textureName].texture.createView(),
+                clearValue: { r: 0, g: 0, b: 0, a: 0 },
+                loadOp: 'clear' as GPULoadOp,
+                storeOp: 'store' as GPUStoreOp
+              } satisfies GPURenderPassColorAttachment)
+          )
+    });
+
+    // Draw a full quad
+    renderPass.setPipeline(this._pipeline);
+    // Bind groups
+    this._bindGroups.forEach((bindGroup, idx) => {
+      renderPass.setBindGroup(idx, bindGroup);
+    });
+    renderPass.draw(6, 1, 0, 0);
+    // End the render pass
+    renderPass.end();
+  }
+}

package/src/backend.ts

@@ -0,0 +1,53 @@
+import { WebGPUBackend } from '@tensorflow/tfjs-backend-webgpu/dist/base';
+import { ENGINE } from '@tensorflow/tfjs-core/dist/engine';
+
+import './kernels';
+
+export async function initWebGPUBackend() {
+  try {
+    const gpuDescriptor: GPURequestAdapterOptions = {
+      powerPreference: 'high-performance'
+    };
+
+    const adapter = (await navigator.gpu.requestAdapter(gpuDescriptor))!;
+    const deviceDescriptor: GPUDeviceDescriptor = {};
+
+    const requiredFeatures = [];
+    if (adapter.features.has('timestamp-query')) {
+      requiredFeatures.push('timestamp-query');
+    }
+    if (adapter.features.has('bgra8unorm-storage')) {
+      requiredFeatures.push(['bgra8unorm-storage']);
+    }
+    deviceDescriptor.requiredFeatures =
+      requiredFeatures as Iterable<GPUFeatureName>;
+
+    const adapterLimits = adapter.limits;
+    deviceDescriptor.requiredLimits = {
+      maxComputeWorkgroupStorageSize:
+        adapterLimits.maxComputeWorkgroupStorageSize,
+      maxComputeWorkgroupsPerDimension:
+        adapterLimits.maxComputeWorkgroupsPerDimension,
+      maxStorageBufferBindingSize: adapterLimits.maxStorageBufferBindingSize,
+      maxBufferSize: adapterLimits.maxBufferSize,
+      maxComputeWorkgroupSizeX: adapterLimits.maxComputeWorkgroupSizeX,
+      maxComputeInvocationsPerWorkgroup:
+        adapterLimits.maxComputeInvocationsPerWorkgroup
+    };
+    const device = await adapter.requestDevice(deviceDescriptor);
+    const adapterInfo = await adapter.requestAdapterInfo();
+
+    return initWebGPUBackendWithDevice(device, adapterInfo);
+  } catch (e) {}
+}
+
+export async function initWebGPUBackendWithDevice(
+  device: GPUDevice,
+  adapter: GPUAdapterInfo
+) {
+  const backend = new WebGPUBackend(device, adapter);
+  ENGINE.registerBackend('webgpu-oidn', () => backend);
+  await ENGINE.setBackend('webgpu-oidn');
+
+  return backend;
+}