webgpu: support AvgPoolGrad kernel (#7188)

Author: xhcao

tfjs-backend-webgpu/src/avg_pool2d_backprop_webgpu.ts

@@ -0,0 +1,87 @@
+/**
+ * @license
+ * Copyright 2022 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 {backend_util} from '@tensorflow/tfjs-core';
+import {getMainHeaderString as main, WebGPUProgram} from './webgpu_program';
+import {computeDispatch, flatDispatchLayout} from './webgpu_util';
+
+export class AvgPool2DBackpropProgram implements WebGPUProgram {
+  outputShape: number[];
+  shaderKey: string;
+  dispatchLayout: {x: number[]};
+  dispatch: [number, number, number];
+  variableNames = ['dy'];
+  uniforms =
+      `stride : vec2<i32>, pads : vec2<i32>, dilation : vec2<i32>, filterDims : vec2<i32>,
+       outHeight : i32, outWidth : i32, avgMultiplier : f32,`;
+  workgroupSize: [number, number, number] = [64, 1, 1];
+  size = true;
+
+  constructor(convInfo: backend_util.Conv2DInfo) {
+    this.outputShape = convInfo.inShape;
+
+    this.dispatchLayout = flatDispatchLayout(this.outputShape);
+
+    this.dispatch = computeDispatch(
+        this.dispatchLayout, this.outputShape, this.workgroupSize);
+
+    this.shaderKey = `avg_pool2d_backprop`;
+  }
+
+  getUserCode(): string {
+    const userCode = `
+      ${main('index')} {
+      if (index < uniforms.size) {
+        let coords = getCoordsFromIndex(index);
+        let batch = coords[0];
+        let d = coords[3];
+
+        let dyRCCorner = vec2<i32>(coords.yz) - uniforms.pads;
+        let dyRCorner = dyRCCorner.x;
+        let dyCCorner = dyRCCorner.y;
+
+        // Convolve dy(?, ?, d) with pos mask(:, :, d) to get dx(xR, xC, d).
+        // ? = to be determined. : = across all values in that axis.
+        var dotProd = 0.0;
+        for (var wR = 0; wR < uniforms.filterDims[0]; wR = wR + uniforms.dilation[0]) {
+          let dyR = f32(dyRCorner + wR) / f32(uniforms.stride[0]);
+
+          if (dyR < 0.0 || dyR >= f32(uniforms.outHeight) || fract(dyR) > 0.0) {
+            continue;
+          }
+          let idyR = i32(dyR);
+
+          for (var wC = 0; wC < uniforms.filterDims[1]; wC = wC + uniforms.dilation[1]) {
+            let dyC = f32(dyCCorner + wC) / f32(uniforms.stride[1]);
+
+            if (dyC < 0.0 || dyC >= f32(uniforms.outWidth) || fract(dyC) > 0.0) {
+              continue;
+            }
+            let idyC = i32(dyC);
+
+            let dyValue = getDy(batch, idyR, idyC, d);
+
+            dotProd = dotProd + dyValue * uniforms.avgMultiplier;
+          }
+        }
+        setOutputAtIndex(index, dotProd);
+      }
+    }
+    `;
+    return userCode;
+  }
+}

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

@@ -0,0 +1,64 @@
+/**
+ * @license
+ * Copyright 2022 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 {AvgPoolGrad, AvgPoolGradAttrs, AvgPoolGradInputs, backend_util, KernelConfig, KernelFunc, TensorInfo} from '@tensorflow/tfjs-core';
+
+import {AvgPool2DBackpropProgram} from '../avg_pool2d_backprop_webgpu';
+import {WebGPUBackend} from '../backend_webgpu';
+import {assertNotComplex} from '../webgpu_util';
+
+export function avgPoolGrad(args: {
+  inputs: AvgPoolGradInputs,
+  backend: WebGPUBackend,
+  attrs: AvgPoolGradAttrs
+}): TensorInfo {
+  const {inputs, backend, attrs} = args;
+  const {dy, input} = inputs;
+  const x = input;
+  assertNotComplex([dy, input], 'avgPoolGrad');
+  const {filterSize, strides, pad} = attrs;
+
+  const convInfo = backend_util.computePool2DInfo(
+      x.shape as [number, number, number, number], filterSize, strides,
+      1 /* dilations */, pad);
+  const avgPoolBackpropProgram = new AvgPool2DBackpropProgram(convInfo);
+  const avgMultiplier = 1 / (convInfo.filterHeight * convInfo.filterWidth);
+  const uniformData = [
+    {type: 'int32', data: [convInfo.strideHeight, convInfo.strideWidth]}, {
+      type: 'int32',
+      data: [
+        convInfo.effectiveFilterHeight - 1 - convInfo.padInfo.top,
+        convInfo.effectiveFilterWidth - 1 - convInfo.padInfo.left
+      ]
+    },
+    {type: 'int32', data: [convInfo.dilationHeight, convInfo.dilationWidth]}, {
+      type: 'int32',
+      data: [convInfo.effectiveFilterHeight, convInfo.effectiveFilterWidth]
+    },
+    {type: 'int32', data: [convInfo.outHeight]},
+    {type: 'int32', data: [convInfo.outWidth]},
+    {type: 'float32', data: [avgMultiplier]}
+  ];
+  return backend.runWebGPUProgram(
+      avgPoolBackpropProgram, [dy], x.dtype, uniformData);
+}
+
+export const avgPoolGradConfig: KernelConfig = {
+  kernelName: AvgPoolGrad,
+  backendName: 'webgpu',
+  kernelFunc: avgPoolGrad as unknown as KernelFunc
+};

tfjs-backend-webgpu/src/register_all_kernels.ts

@@ -32,6 +32,7 @@ import {atanConfig} from './kernels/Atan';
 import {atan2Config} from './kernels/Atan2';
 import {atanhConfig} from './kernels/Atanh';
 import {avgPoolConfig} from './kernels/AvgPool';
+import {avgPoolGradConfig} from './kernels/AvgPoolGrad';
 import {batchMatMulConfig} from './kernels/BatchMatMul';
 import {batchToSpaceNDConfig} from './kernels/BatchToSpaceND';
 import {bincountConfig} from './kernels/Bincount';
@@ -170,6 +171,7 @@ const kernelConfigs: KernelConfig[] = [
   atan2Config,
   atanhConfig,
   avgPoolConfig,
+  avgPoolGradConfig,
   batchMatMulConfig,
   batchToSpaceNDConfig,
   bincountConfig,

tfjs-backend-webgpu/src/setup_test.ts

@@ -45,12 +45,6 @@ const TEST_FILTERS: TestFilter[] = [
       'gradient',  // Not yet implemented.
     ]
   },
-  {
-    startsWith: 'avgPool ',
-    excludes: [
-      'gradient',  // Not yet implemented.
-    ]
-  },
   {
     startsWith: 'batchToSpaceND ',
     excludes: [
@@ -186,6 +180,12 @@ const TEST_FILTERS: TestFilter[] = [
       'poolBackprop',  // maxPoolBackprop not yet implemented.
     ]
   },
+  {
+    startsWith: 'poolBackprop ',
+    excludes: [
+      'max',  // maxPoolBackprop not yet implemented.
+    ]
+  },
   {
     startsWith: 'prod ',
     excludes: [
@@ -264,7 +264,6 @@ const TEST_FILTERS: TestFilter[] = [
       'maxPoolBackprop ',
       'maxPoolWithArgmax ',
       'multinomial ',
-      'poolBackprop ',
       'raggedGather ',
       'raggedRange ',
       'raggedTensorToTensor ',

tfjs-backend-webgpu/src/webgpu_util.ts

@@ -14,7 +14,7 @@
  * limitations under the License.
  * =============================================================================
  */
-import {DataType} from '@tensorflow/tfjs-core';
+import {DataType, TensorInfo, util} from '@tensorflow/tfjs-core';
 
 const arrayProduct = (arr: number[]) => {
   let product = 1;
@@ -161,6 +161,21 @@ export function isWebGPUSupported(): boolean {
       !!navigator.gpu;
 }
 
+export function assertNotComplex(
+    tensor: TensorInfo|TensorInfo[], opName: string): void {
+  if (!Array.isArray(tensor)) {
+    tensor = [tensor];
+  }
+  tensor.forEach(t => {
+    if (t != null) {
+      util.assert(
+          t.dtype !== 'complex64',
+          () => `${opName} does not support complex64 tensors ` +
+              'in the WebGPU backend.');
+    }
+  });
+}
+
 export enum MatMulProgramType {
   MatMulReduceProgram,
   MatMulSplitKProgram,