[webgpu] Implement draw API (#7749)

FEATURE

* [webgpu] Implement draw API

* Support bgra8-unorm

* Mark texture externael and rename to_pixels to draw

* Nit

---------

Co-authored-by: Linchenn <40653845+Linchenn@users.noreply.github.com>

tfjs-backend-webgl/src/setup_test.ts

@@ -39,7 +39,7 @@ const customInclude = (testName: string) => {
  'throws when index is out of bound',
  // otsu tests for threshold op is failing on windows
  'method otsu',
- 'Draw on 2d context',
+ 'draw on canvas context',
  // https://github.com/tensorflow/tfjs/issues/7618
  'numbers exceed float32 precision',
  ];

tfjs-backend-webgpu/src/backend_webgpu.ts

@@ -909,7 +909,8 @@ export class WebGPUBackend extends KernelBackend {
  // program size, program defined uniforms.
  let programUniform: ProgramUniform = [];
  let bufferShapes: number[][] = [];
- if (!program.isFromPixels) {
+ const uniformsType = 'int32';
+ if (program.pixelsOpType == null) {
  programUniform.push(
  {type: 'float32', data: [NaN]}, {type: 'float32', data: [Infinity]});
  bufferShapes = inputs.concat(output).map(d => d.shape);
@@ -919,14 +920,16 @@ export class WebGPUBackend extends KernelBackend {
  const strides = util.computeStrides(d);
  programUniform.push({type: uniformsType, data: strides});
  });
- if (program.size) {
- const size = util.sizeFromShape(program.outputShape);
- programUniform.push({
- type: uniformsType,
- data:
- [program.outputComponent ? size / program.outputComponent : size]
- });
- }
+ } else {
+ const strides = util.computeStrides(output.shape);
+ programUniform.push({type: uniformsType, data: strides});
+ }
+ if (program.size) {
+ const size = util.sizeFromShape(program.outputShape);
+ programUniform.push({
+ type: uniformsType,
+ data: [program.outputComponent ? size / program.outputComponent : size]
+ });
  }
 
  if (programDefinedUniform) {
@@ -986,7 +989,8 @@ export class WebGPUBackend extends KernelBackend {
 
  if (shouldTimeProgram ||
  env().get('WEBGPU_DEFERRED_SUBMIT_BATCH_SIZE') as
- number <= this.dispatchCountInPass) {
+ number <= this.dispatchCountInPass ||
+ program.pixelsOpType === webgpu_program.PixelsOpType.DRAW) {
  this.endComputePassEncoder();
  if (shouldTimeProgram) {
  this.activeTimers.push(

tfjs-backend-webgpu/src/base.ts

@@ -33,9 +33,15 @@ if (isWebGPUSupported()) {
  const adapter = await navigator.gpu.requestAdapter(gpuDescriptor);
  const deviceDescriptor: GPUDeviceDescriptor = {};
 
+ const requiredFeatures = [];
  if (adapter.features.has('timestamp-query')) {
- deviceDescriptor.requiredFeatures = ['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 = {

tfjs-backend-webgpu/src/draw_webgpu.ts

@@ -0,0 +1,79 @@
+/**
+ * @license
+ * Copyright 2023 Google LLC.
+ * 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 {getMainHeaderString as main, PixelsOpType, WebGPUProgram} from './webgpu_program';
+import {computeDispatch, flatDispatchLayout} from './webgpu_util';
+
+export class DrawProgram implements WebGPUProgram {
+ variableNames = ['Image'];
+ uniforms = 'alpha: f32,';
+ outputShape: number[];
+ shaderKey: string;
+ dispatchLayout: {x: number[]};
+ dispatch: [number, number, number];
+ workgroupSize: [number, number, number] = [64, 1, 1];
+ type: DataType;
+ textureFormat: GPUTextureFormat;
+ pixelsOpType = PixelsOpType.DRAW;
+ size = true;
+
+ constructor(
+ outShape: number[], type: DataType, textureFormat: GPUTextureFormat) {
+ this.outputShape = outShape;
+ this.dispatchLayout = flatDispatchLayout(this.outputShape);
+ this.dispatch = computeDispatch(
+ this.dispatchLayout, this.outputShape, this.workgroupSize);
+ this.type = type;
+ this.textureFormat = textureFormat;
+ this.shaderKey = `draw_${type}_${textureFormat}`;
+ }
+
+ getUserCode(): string {
+ let calculateResult;
+ const value = this.type === 'float32' ? 'value' : 'value / 255.0';
+ calculateResult = `
+ if (uniforms.numChannels == 1) {
+ rgba[0] = ${value};
+ rgba[1] = ${value};
+ rgba[2] = ${value};
+ } else {
+ rgba[d] = ${value};
+ }`;
+
+ const userCode = `
+ @group(0) @binding(0) var outImage : texture_storage_2d<${
+ this.textureFormat}, write>;
+ ${main('index')} {
+ if (index < uniforms.size) {
+ var rgba = vec4<f32>(0.0, 0.0, 0.0, uniforms.alpha);
+ for (var d = 0; d < uniforms.numChannels; d = d + 1) {
+ let value = f32(inBuf[index * uniforms.numChannels + d]);
+ ${calculateResult}
+ }
+ rgba.x = rgba.x * rgba.w;
+ rgba.y = rgba.y * rgba.w;
+ rgba.z = rgba.z * rgba.w;
+ let coords = getCoordsFromIndex(index);
+ textureStore(outImage, vec2<i32>(coords.yx), rgba);
+ }
+ }
+ `;
+ return userCode;
+ }
+}

tfjs-backend-webgpu/src/from_pixels_webgpu.ts

@@ -15,13 +15,13 @@
  * =============================================================================
  */
 
-import {getMainHeaderString as main, WebGPUProgram} from './webgpu_program';
+import {getMainHeaderString as main, PixelsOpType, WebGPUProgram} from './webgpu_program';
 import {computeDispatch, flatDispatchLayout} from './webgpu_util';
 
 export class FromPixelsProgram implements WebGPUProgram {
  dispatch: [number, number, number];
  dispatchLayout: {x: number[]};
- isFromPixels = true;
+ pixelsOpType = PixelsOpType.FROM_PIXELS;
  outputShape: number[] = [0];
  shaderKey: string;
  importVideo: boolean;

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

@@ -0,0 +1,85 @@
+/**
+ * @license
+ * Copyright 2023 Google LLC.
+ * 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 {Draw, DrawAttrs, DrawInputs,} from '@tensorflow/tfjs-core';
+
+import {WebGPUBackend} from '../backend_webgpu';
+import {DrawProgram} from '../draw_webgpu';
+
+export function draw(
+ args: {inputs: DrawInputs, backend: WebGPUBackend, attrs: DrawAttrs}):
+ TensorInfo {
+ const {inputs, backend, attrs} = args;
+ const {image} = inputs;
+ const {canvas, options} = attrs;
+ const [height, width] = image.shape.slice(0, 2);
+ const {imageOptions} = options || {};
+ const alpha = imageOptions ?.alpha || 1;
+
+ // 'rgba8unorm' should work on macOS according to
+ // https://bugs.chromium.org/p/chromium/issues/detail?id=1298618. But
+ // failed on macOS/M2. So use 'bgra8unorm' first when available.
+ const format = backend.device.features.has('bgra8unorm-storage') ?
+ 'bgra8unorm' :
+ 'rgba8unorm';
+ const outShape = [height, width];
+ const program = new DrawProgram(outShape, image.dtype, format);
+ canvas.width = width;
+ canvas.height = height;
+ const backendName = 'webgpu';
+ let gpuContext = canvas.getContext(backendName);
+ let canvasWebGPU;
+ if (!gpuContext) {
+ canvasWebGPU = new OffscreenCanvas(width, height);
+ gpuContext = canvasWebGPU.getContext(backendName);
+ }
+ const numChannels = image.shape.length === 3 ? image.shape[2] : 1;
+ gpuContext.configure({
+ device: backend.device,
+ format,
+ usage: GPUTextureUsage.STORAGE_BINDING,
+ alphaMode: 'premultiplied'
+ });
+
+ const outputDtype = 'int32';
+ const output = backend.makeTensorInfo(outShape, outputDtype);
+ const info = backend.tensorMap.get(output.dataId);
+ info.resource = gpuContext.getCurrentTexture();
+ info.external = true;
+
+ const uniformData =
+ [{type: 'uint32', data: [numChannels]}, {type: 'float32', data: [alpha]}];
+ backend.runWebGPUProgram(program, [image], outputDtype, uniformData, output);
+
+ if (canvasWebGPU) {
+ const canvas2dContext = canvas.getContext('2d');
+ if (!canvas2dContext) {
+ throw new Error(
+ `Please make sure this canvas has only been used for 2d or webgpu context!`);
+ }
+ canvas2dContext.drawImage(canvasWebGPU, 0, 0);
+ }
+ backend.disposeData(output.dataId);
+ return image;
+}
+
+export const drawConfig: KernelConfig = {
+ kernelName: Draw,
+ backendName: 'webgpu',
+ kernelFunc: draw as unknown as KernelFunc
+};

tfjs-backend-webgpu/src/register_all_kernels.ts

@@ -65,6 +65,7 @@ import {diagConfig} from './kernels/Diag';
 import {dilation2DConfig} from './kernels/Dilation2D';
 import {dilation2DBackpropFilterConfig} from './kernels/Dilation2DBackpropFilter';
 import {dilation2DBackpropInputConfig} from './kernels/Dilation2DBackpropInput';
+import {drawConfig} from './kernels/Draw';
 import {einsumConfig} from './kernels/Einsum';
 import {eluConfig} from './kernels/Elu';
 import {eluGradConfig} from './kernels/EluGrad';
@@ -229,6 +230,7 @@ const kernelConfigs: KernelConfig[] = [
  dilation2DConfig,
  dilation2DBackpropFilterConfig,
  dilation2DBackpropInputConfig,
+ drawConfig,
  einsumConfig,
  eluConfig,
  eluGradConfig,

tfjs-backend-webgpu/src/setup_test.ts

@@ -126,12 +126,6 @@ const TEST_FILTERS: TestFilter[] = [
  'canvas and image match', // Failing on Linux
  ],
  },
- {
- startsWith: 'Draw',
- excludes: [
- 'on 2d context',
- ]
- },
  {
  startsWith: 'sign ',
  excludes: [

tfjs-backend-webgpu/src/webgpu_program.ts

@@ -19,6 +19,11 @@ import {backend_util, DataType, DataTypeMap, env, Rank, TensorInfo, util} from '
 
 import {symbolicallyComputeStrides} from './shader_util';
 
+export enum PixelsOpType {
+ FROM_PIXELS,
+ DRAW
+}
+
 export interface WebGPUProgram {
  // Whether to use atomic built-in functions.
  atomic?: boolean;
@@ -27,10 +32,10 @@ export interface WebGPUProgram {
  // dispatchLayout enumerates how tensor dimensions are distributed among
  // dispatch x,y,z dimensions.
  dispatchLayout: {x: number[], y?: number[], z?: number[]};
- isFromPixels?: boolean;
  // By default, the output data component is 1.
  outputComponent?: number;
  outputShape: number[];
+ pixelsOpType?: PixelsOpType;
  // The unique key to distinguish different shader source code.
  shaderKey: string;
  // Whether to use output size for bounds checking.
@@ -219,16 +224,23 @@ function makeShader(
  }
  `);
 
- if (program.isFromPixels) {
+ if (program.pixelsOpType != null) {
+ const inoutSnippet = program.pixelsOpType === PixelsOpType.FROM_PIXELS ?
+ `@group(0) @binding(0) var<storage, read_write> result: array<${
+ dataTypeToGPUType(outputData.dtype, program.outputComponent)}>;` :
+ `@group(0) @binding(1) var<storage, read> inBuf : array<${
+ dataTypeToGPUType(inputInfo[0].dtype, program.outputComponent)}>;`;
+ const outShapeStridesType =
+ outputData.shape.length === 3 ? 'vec2<i32>' : 'i32';
  prefixSnippets.push(`
  struct Uniform {
+ outShapeStrides : ${outShapeStridesType},
  size : i32,
  numChannels : i32,
- outShapeStrides : vec2<i32>,
+ alpha : f32,
  };
 
- @group(0) @binding(0) var<storage, read_write> result: array<${
- dataTypeToGPUType(outputData.dtype, program.outputComponent)}>;
+ ${inoutSnippet}
  @group(0) @binding(2) var<uniform> uniforms: Uniform;
  `);
  const useGlobalIndex = isFlatDispatchLayout(program);
@@ -339,7 +351,7 @@ export function makeShaderKey<R extends Rank>(
  program: WebGPUProgram, inputsData: InputInfo[],
  output: TensorInfo): string {
  let key = program.shaderKey;
- if (program.isFromPixels) {
+ if (program.pixelsOpType != null) {
  return key;
  }