[webgpu] Add support for scatterND (#5643)

FEATURE

Author: haoyunfeix

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

@@ -0,0 +1,76 @@
+/**
+ * @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 {backend_util, KernelConfig, KernelFunc, ScatterNd, ScatterNdAttrs, ScatterNdInputs, TensorInfo} from '@tensorflow/tfjs-core';
+
+import {WebGPUBackend} from '../backend_webgpu';
+
+import {reshape} from './Reshape';
+import {ScatterProgram} from './scatter_webgpu';
+
+export function scatterNd(args: {
+  inputs: ScatterNdInputs,
+  backend: WebGPUBackend,
+  attrs: ScatterNdAttrs
+}): TensorInfo {
+  const {inputs, backend, attrs} = args;
+  const {indices, updates} = inputs;
+  const {shape} = attrs;
+
+  const {sliceRank, numUpdates, sliceSize, strides, outputSize} =
+      backend_util.calculateShapes(updates, indices, shape);
+
+  const flattenShape = [outputSize / sliceSize, sliceSize];
+
+  if (outputSize === 0) {
+    return backend.makeTensorInfo(shape, indices.dtype);
+  }
+
+  const flattenIndices = reshape(
+      {inputs: {x: indices}, backend, attrs: {shape: [numUpdates, sliceRank]}});
+  const flattenX = reshape(
+      {inputs: {x: updates}, backend, attrs: {shape: [numUpdates, sliceSize]}});
+
+  const defaultValue = backend.makeTensorInfo(
+      [], 'float32', new Float32Array([0]));  // scalar(0)
+  const uniformData = [
+    {type: 'int32', data: [numUpdates]},
+    {type: 'int32', data: [sliceRank]},
+    {type: 'int32', data: strides},
+  ];
+  const program = new ScatterProgram(
+      numUpdates, sliceRank, flattenIndices.shape.length, flattenX.shape.length,
+      strides, flattenShape);
+  const res = backend.runWebGPUProgram(
+      program, [flattenX, flattenIndices, defaultValue], flattenX.dtype,
+      uniformData);
+
+  const reshaped = reshape({inputs: {x: res}, backend, attrs: {shape}});
+
+  backend.disposeData(flattenIndices.dataId);
+  backend.disposeData(flattenX.dataId);
+  backend.disposeData(res.dataId);
+  backend.disposeData(defaultValue.dataId);
+
+  return reshaped;
+}
+
+export const scatterNdConfig: KernelConfig = {
+  kernelName: ScatterNd,
+  backendName: 'webgpu',
+  kernelFunc: scatterNd as {} as KernelFunc
+};

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

@@ -42,7 +42,9 @@ export class ScatterProgram implements WebGPUProgram {
     this.dispatchLayout = flatDispatchLayout(this.outputShape);
     this.dispatch = computeDispatch(
         this.dispatchLayout, this.outputShape, this.workGroupSize);
-    this.shaderKey = `scatter_${indicesRank}_${updatesRank}`;
+    const sliceDimGreaterThanOne = sliceDim > 1;
+    this.shaderKey =
+        `scatter_${indicesRank}_${updatesRank}_${sliceDimGreaterThanOne}`;
     this.size = util.sizeFromShape(this.outputShape);
     const stridesType = getCoordsDataType(strides.length);
     this.uniforms =
@@ -64,7 +66,7 @@ export class ScatterProgram implements WebGPUProgram {
     this.updatesSnippet = `getUpdates(${updatesString})`;
 
     this.strideString =
-        sliceDim > 1 ? 'uniforms.strides[j]' : 'uniforms.strides';
+        sliceDimGreaterThanOne ? 'uniforms.strides[j]' : 'uniforms.strides';
   }
 
   getUserCode(): string {

tfjs-backend-webgpu/src/register_all_kernels.ts

@@ -90,6 +90,7 @@ import {resizeBilinearConfig} from './kernels/ResizeBilinear';
 import {resizeNearestNeighborConfig} from './kernels/ResizeNearestNeighbor';
 import {rotateWithOffsetConfig} from './kernels/RotateWithOffset';
 import {rsqrtConfig} from './kernels/Rsqrt';
+import {scatterNdConfig} from './kernels/ScatterNd';
 import {selectConfig} from './kernels/Select';
 import {sigmoidConfig} from './kernels/Sigmoid';
 import {sinConfig} from './kernels/Sin';
@@ -189,6 +190,7 @@ const kernelConfigs: KernelConfig[] = [
   resizeNearestNeighborConfig,
   rotateWithOffsetConfig,
   rsqrtConfig,
+  scatterNdConfig,
   selectConfig,
   sigmoidConfig,
   sinConfig,

tfjs-backend-webgpu/src/setup_test.ts

@@ -584,6 +584,12 @@ const TEST_FILTERS: TestFilter[] = [
       'gradient'  // gradient function not found.
     ]
   },
+  {
+    include: 'scatterND',
+    excludes: [
+      'gradient'  // gradient function not found.
+    ]
+  },
   {
     startsWith: 'logicalAnd ',
   },