[WIP] Provide an efficient way for toPixels

tfjs-backend-webgpu/src/backend_webgpu.ts

@@ -84,6 +84,7 @@ export class WebGPUBackend extends KernelBackend {
   private commandQueueOwnedIds = new WeakSet<DataId>();
   private layoutCache: {[key: number]: WebGPULayout};
   private pipelineCache: {[key: string]: GPUComputePipeline};
+  private renderPipeline: GPURenderPipeline;
   private bufferManager: BufferManager;
 
   private tensorDisposalQueue: DataId[] = [];
@@ -777,7 +778,7 @@ export class WebGPUBackend extends KernelBackend {
 
     const pipeline = this.getAndSavePipeline(key, () => {
       return webgpu_program.compileProgram(
-          this.device, program, pipelineLayout, inputsData, output);
+          this.device, program, pipelineLayout, inputsData, output.dtype);
     });
 
     const shouldTimeProgram = this.activeTimers != null;
@@ -884,6 +885,181 @@ export class WebGPUBackend extends KernelBackend {
     }
   }
 
+  private createBufferToTextureLayout(): WebGPULayout {
+    const bindGroupLayoutEntries: GPUBindGroupLayoutEntry[] = [];
+    // Output buffer binding layout.
+    bindGroupLayoutEntries.push({
+      binding: 0,
+      visibility: GPUShaderStage.COMPUTE,
+      storageTexture: {access: 'write-only', format: 'rgba8unorm'}
+    });
+    // Input buffer binding layout.
+    bindGroupLayoutEntries.push({
+      binding: 1,
+      visibility: GPUShaderStage.COMPUTE,
+      buffer: {type: 'read-only-storage' as const}
+    });
+    // Uniform buffer binding layout.
+    bindGroupLayoutEntries.push({
+      binding: 2,
+      visibility: GPUShaderStage.COMPUTE,
+      buffer: {type: 'uniform' as const}
+    });
+    const bindGroupLayout =
+        this.device.createBindGroupLayout({entries: bindGroupLayoutEntries});
+    const pipelineLayout =
+        this.device.createPipelineLayout({bindGroupLayouts: [bindGroupLayout]});
+    return {bindGroupLayout, pipelineLayout};
+  }
+
+  private runDrawTexture(ctx: GPUCanvasContext, srcTexture: GPUTexture) {
+    if (!this.renderPipeline) {
+      this.renderPipeline = this.device.createRenderPipeline({
+        vertex: {
+          module: this.device.createShaderModule({
+            code: `
+            [[stage(vertex)]]
+            fn main([[builtin(vertex_index)]] VertexIndex : u32) -> [[builtin(position)]] vec4<f32> {
+              var pos = array<vec2<f32>, 6>(
+                vec2<f32>(-1.0,  1.0),
+                vec2<f32>(-1.0, -1.0),
+                vec2<f32>( 1.0,  1.0),
+                vec2<f32>(-1.0, -1.0),
+                vec2<f32>( 1.0,  1.0),
+                vec2<f32>( 1.0, -1.0));
+              return vec4<f32>(pos[VertexIndex], 0.0, 1.0);
+            }`,
+          }),
+          entryPoint: 'main',
+        },
+        fragment: {
+          module: this.device.createShaderModule({
+            code: `
+            [[group(0), binding(0)]] var image0 : texture_2d<f32>;
+            [[stage(fragment)]]
+            fn main([[builtin(position)]] coord_in: vec4<f32>) -> [[location(0)]] vec4<f32> {
+              var coord_in_vec2 = vec2<i32>(i32(coord_in.x), i32(coord_in.y));
+              let value = textureLoad(image0, coord_in_vec2, 0);
+              return value;
+            }
+            `,
+          }),
+          entryPoint: 'main',
+          targets: [{format: 'bgra8unorm'}],
+        },
+      });
+    }
+
+    const bindGroup = this.device.createBindGroup({
+      layout: this.renderPipeline.getBindGroupLayout(0),
+      entries: [
+        {
+          binding: 0,
+          resource: srcTexture.createView(),
+        },
+      ],
+    });
+
+    const renderPassDescriptor: GPURenderPassDescriptor = {
+      colorAttachments: [
+        {
+          view: ctx.getCurrentTexture().createView(),
+
+          loadValue: {r: 0.0, g: 0.0, b: 0.0, a: 0.0},
+          storeOp: 'store',
+        },
+      ],
+    };
+
+    this.ensureCommandEncoderReady();
+    const passEncoder =
+        this.currentCommandEncoder.beginRenderPass(renderPassDescriptor);
+    passEncoder.setPipeline(this.renderPipeline);
+    passEncoder.setBindGroup(0, bindGroup);
+    passEncoder.draw(6);
+    passEncoder.endPass();
+  }
+
+  runToCanvasProgram(
+      program: webgpu_program.WebGPUProgram, input: TensorInfo,
+      ctx: GPUCanvasContext) {
+    this.uploadToGPU(input.dataId);
+    const srcBuffer = this.tensorMap.get(input.dataId).bufferInfo.buffer;
+    const width = input.shape[1];
+    const height = input.shape[0];
+
+    const interTexture = this.device.createTexture({
+      format: 'rgba8unorm',
+      size: [width, height, 1],
+      usage: GPUTextureUsage.COPY_SRC | GPUTextureUsage.STORAGE_BINDING |
+          GPUTextureUsage.TEXTURE_BINDING
+    });
+
+    const size = util.sizeFromShape(program.outputShape);
+    const depth = input.shape.length === 2 ? 1 : input.shape[2];
+    const strides = util.computeStrides(program.outputShape);
+    const uniformData = [size, depth, ...strides];
+    const uniformBuffer = this.bufferManager.acquireBuffer(
+        uniformData.length * 4,
+        GPUBufferUsage.COPY_DST | GPUBufferUsage.UNIFORM);
+    this.device.queue.writeBuffer(
+        uniformBuffer, 0, new Int32Array(uniformData));
+
+    const layout = this.createBufferToTextureLayout();
+    const pipeline = this.getAndSavePipeline(program.shaderKey, () => {
+      return webgpu_program.compileProgram(
+          this.device, program, layout.pipelineLayout, [], input.dtype, true);
+    });
+
+    const bindGroup = this.device.createBindGroup({
+      layout: layout.bindGroupLayout,
+      entries: [
+        {
+          binding: 0,
+          resource: interTexture.createView(),
+        },
+        {
+          binding: 1,
+          resource: {
+            buffer: srcBuffer,
+          }
+        },
+        {
+          binding: 2,
+          resource: {
+            buffer: uniformBuffer,
+          }
+        }
+      ],
+    });
+    this.ensureCommandEncoderReady();
+    const passEncoder = this.getComputePass();
+    const shouldTimeProgram = this.activeTimers != null;
+    if (shouldTimeProgram) {
+      if (this.supportTimeQuery) {
+        passEncoder.writeTimestamp(this.querySet, 0);
+      }
+    }
+    passEncoder.setPipeline(pipeline);
+    passEncoder.setBindGroup(0, bindGroup);
+    passEncoder.dispatch(
+        program.dispatch[0], program.dispatch[1], program.dispatch[2]);
+    if (shouldTimeProgram) {
+      if (this.supportTimeQuery) {
+        passEncoder.writeTimestamp(this.querySet, 1);
+      }
+    }
+    this.ensureComputePassEnded();
+    this.runDrawTexture(ctx, interTexture);
+    this.submitQueue();
+    if (shouldTimeProgram) {
+      this.activeTimers.push({
+        name: program.constructor.name,
+        query: this.getQueryTime(this.querySet)
+      });
+    }
+  }
+
   async getTimeFromQuerySet(querySet: GPUQuerySet) {
     const queryBuffer = this.acquireBuffer(
         16, GPUBufferUsage.COPY_SRC | GPUBufferUsage.QUERY_RESOLVE);

tfjs-backend-webgpu/src/kernels/FromPixels.ts

@@ -48,9 +48,9 @@ export function fromPixels(args: {
   const isImage = typeof (HTMLImageElement) !== 'undefined' &&
       pixels instanceof HTMLImageElement;
   const isCanvas = (typeof (HTMLCanvasElement) !== 'undefined' &&
-      pixels instanceof HTMLCanvasElement) ||
+                    pixels instanceof HTMLCanvasElement) ||
       (typeof (OffscreenCanvas) !== 'undefined' &&
-      pixels instanceof OffscreenCanvas);
+       pixels instanceof OffscreenCanvas);
   const isImageBitmap =
       typeof (ImageBitmap) !== 'undefined' && pixels instanceof ImageBitmap;
 

tfjs-backend-webgpu/src/kernels/FromPixelsExternalImage.ts

@@ -53,7 +53,7 @@ export function fromPixelsExternalImage(args: {
 
   const pipeline = backend.getAndSavePipeline(key, () => {
     return webgpu_program.compileProgram(
-        backend.device, program, layout.pipelineLayout, [], output, true);
+        backend.device, program, layout.pipelineLayout, [], output.dtype, true);
   });
 
   program.setPipeline(pipeline);
@@ -71,7 +71,7 @@ export function fromPixelsExternalImage(args: {
 
   info.bufferInfo.buffer = backend.acquireBuffer(info.bufferInfo.byteSize);
 
-  const uniformData = [size, numChannels, ...strides, ...program.dispatch];
+  const uniformData = [size, numChannels, ...strides];
   program.setUniform(backend.device, uniformData);
 
   let externalResource: GPUExternalTexture|GPUTextureView;

tfjs-backend-webgpu/src/kernels/FromPixels_utils/from_pixels_webgpu.ts

@@ -66,7 +66,8 @@ export class FromPixelsProgram implements WebGPUProgram {
         'textureLoad(src, vec2<i32>(coords.yx), 0)';
     const textureType = this.useImport ? 'texture_external' : 'texture_2d<f32>';
     return `
-      [[binding(1), group(0)]] var src: ${textureType};
+      [[group(0), binding(0)]] var<storage, write> result : Matrix0;
+      [[group(0), binding(1)]] var src: ${textureType};
 
       ${getMainHeaderAndGlobalIndexString()}
         let flatIndexBase = index * uniforms.numChannels;

tfjs-backend-webgpu/src/kernels/ToCanvas.ts

@@ -0,0 +1,60 @@
+/**
+ * @license
+ * Copyright 2021 Google LLC. All Rights Reserved.
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use backend file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ * =============================================================================
+ */
+
+import {KernelConfig, KernelFunc, TensorInfo} from '@tensorflow/tfjs-core';
+import {ToCanvas, ToCanvasInputs, ToCanvasOutput} from '@tensorflow/tfjs-core';
+
+import {WebGPUBackend} from '../backend_webgpu';
+
+import {ToCanvasProgram} from './to_canvas_webgpu';
+
+export const toCanvasConfig: KernelConfig = {
+  kernelName: ToCanvas,
+  backendName: 'webgpu',
+  kernelFunc: toCanvas as {} as KernelFunc,
+};
+
+export function toCanvas(args: {
+  inputs: ToCanvasInputs,
+  backend: WebGPUBackend,
+  attrs: ToCanvasOutput
+}): TensorInfo {
+  const {inputs, backend, attrs} = args;
+  const {$img} = inputs;
+  const {canvas} = attrs;
+  const [height, width] = $img.shape.slice(0, 2);
+
+  const outShape = [height, width, 4];
+  const program = new ToCanvasProgram(outShape, $img.dtype);
+  canvas.width = width;
+  canvas.height = height;
+  const gpuContext = canvas.getContext('webgpu');
+  //  'rgba8unorm' is not supported yet as the context format. Otherwise, we
+  //  can save the second render pass. Ideally, just one comput pass, we can
+  //  transfer the input tensor data to webgpu context canvas and then return
+  //  the canvas to user. https://bugs.chromium.org/p/dawn/issues/detail?id=1219
+  gpuContext.configure({
+    device: backend.device,
+    format: 'bgra8unorm',
+    compositingAlphaMode: 'opaque'
+  });
+
+  backend.runToCanvasProgram(program, $img, gpuContext);
+
+  const outTensor = backend.makeTensorInfo(program.outputShape, 'int32');
+  return outTensor;
+}

tfjs-backend-webgpu/src/kernels/to_canvas_webgpu.ts

@@ -0,0 +1,88 @@
+/**
+ * @license
+ * Copyright 2021 Google LLC. All Rights Reserved.
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ * =============================================================================
+ */
+
+import {DataType} from '@tensorflow/tfjs-core';
+
+import {getMainHeaderAndGlobalIndexString} from '../shader_preprocessor';
+import {computeDispatch, flatDispatchLayout} from '../webgpu_util';
+
+import {WebGPUProgram} from './webgpu_program';
+
+export class ToCanvasProgram implements WebGPUProgram {
+  variableNames = ['Image'];
+  outputShape: number[];
+  shaderKey: string;
+  dispatchLayout: {x: number[]};
+  dispatch: [number, number, number];
+  workGroupSize: [number, number, number] = [64, 1, 1];
+  workPerThread = 4;
+  type: DataType;
+
+  constructor(outShape: number[], type: DataType) {
+    this.outputShape = outShape;
+    this.dispatchLayout = flatDispatchLayout(this.outputShape);
+    this.dispatch = computeDispatch(
+        this.dispatchLayout, this.outputShape, this.workGroupSize,
+        [this.workPerThread, 1, 1]);
+    this.shaderKey = `toCanvas_${type}`;
+    this.type = type;
+  }
+
+  getUserCode(): string {
+    let calculateResult;
+    if (this.type === 'float32') {
+      calculateResult = `
+      if (uniforms.numChannels == 1) {
+        rgba[0] = value;
+        rgba[1] = value;
+        rgba[2] = value;
+      } else {
+        rgba[d] = value;
+      }`;
+    } else {
+      calculateResult = `
+      if (uniforms.numChannels == 1) {
+        rgba[0] = value / 255.0;
+        rgba[1] = value / 255.0;
+        rgba[2] = value / 255.0;
+      } else {
+        rgba[d] = value / 255.0;
+      }`;
+    }
+
+    const userCode = `
+      [[group(0), binding(0)]] var outImage : texture_storage_2d<rgba8unorm, write>;
+      [[group(0), binding(1)]] var<storage, read> inBuf : Matrix0;
+
+      ${getMainHeaderAndGlobalIndexString()}
+            let flatIndex = index * 4;
+            if (flatIndex < uniforms.size) {
+              var rgba = vec4<f32>(0.0, 0.0, 0.0, 1.0);
+
+              for (var d = 0; d < uniforms.numChannels; d = d + 1) {
+                let value = f32(inBuf.numbers[index * uniforms.numChannels + d]);
+                ${calculateResult}
+              }
+
+              let coords = getCoordsFromFlatIndex(flatIndex);
+              textureStore(outImage, vec2<i32>(coords.yx), rgba);
+            }
+          }
+        `;
+    return userCode;
+  }
+}