webgpu: support MaxPoolGrad kernel

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

@@ -0,0 +1,80 @@
+/**
+ * @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 {backend_util, KernelConfig, KernelFunc, MaxPoolGrad, MaxPoolGradAttrs, MaxPoolGradInputs, TensorInfo} from '@tensorflow/tfjs-core';
+
+import {WebGPUBackend} from '../backend_webgpu';
+import {MaxPool2DBackpropProgram} from '../max_pool2d_backprop_webgpu';
+import {Pool2DProgram} from '../pool2d_webgpu';
+import {assertNotComplex} from '../webgpu_util';
+
+export function maxPoolGrad(args: {
+  inputs: MaxPoolGradInputs,
+  backend: WebGPUBackend,
+  attrs: MaxPoolGradAttrs
+}): TensorInfo {
+  const {inputs, backend, attrs} = args;
+  const {dy, input, output} = inputs;
+  const x = input;
+  assertNotComplex([input, output], 'maxPoolGrad');
+  const {filterSize, strides, pad, dimRoundingMode} = attrs;
+
+  const convInfo = backend_util.computePool2DInfo(
+      x.shape as [number, number, number, number], filterSize, strides,
+      1 /* dilations */, pad, dimRoundingMode);
+
+  const maxPoolPositionsProgram = new Pool2DProgram(convInfo, 'max', true);
+  let uniformData = [
+    {type: 'int32', data: [convInfo.strideHeight, convInfo.strideWidth]},
+    {type: 'int32', data: [convInfo.padInfo.top, convInfo.padInfo.left]},
+    {type: 'int32', data: [convInfo.dilationHeight, convInfo.dilationWidth]},
+    {type: 'int32', data: [convInfo.inHeight, convInfo.inWidth]}, {
+      type: 'int32',
+      data: [convInfo.effectiveFilterHeight, convInfo.effectiveFilterWidth]
+    }
+  ];
+  const maxPoolPositions = backend.runWebGPUProgram(
+      maxPoolPositionsProgram, [x], 'int32', uniformData);
+
+  const maxPoolBackpropProgram = new MaxPool2DBackpropProgram(convInfo);
+  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]}
+  ];
+  const result = backend.runWebGPUProgram(
+      maxPoolBackpropProgram, [dy, maxPoolPositions], x.dtype, uniformData);
+  backend.disposeData(maxPoolPositions.dataId);
+
+  return result;
+}
+
+export const maxPoolGradConfig: KernelConfig = {
+  kernelName: MaxPoolGrad,
+  backendName: 'webgpu',
+  kernelFunc: maxPoolGrad as unknown as KernelFunc
+};

tfjs-backend-webgpu/src/max_pool2d_backprop_webgpu.ts

@@ -0,0 +1,93 @@
+/**
+ * @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 {backend_util} from '@tensorflow/tfjs-core';
+import {getMainHeaderString as main, WebGPUProgram} from './webgpu_program';
+import {computeDispatch, flatDispatchLayout} from './webgpu_util';
+
+export class MaxPool2DBackpropProgram implements WebGPUProgram {
+  outputShape: number[];
+  shaderKey: string;
+  dispatchLayout: {x: number[]};
+  dispatch: [number, number, number];
+  variableNames = ['dy', 'maxPos'];
+  uniforms =
+      `strides : vec2<i32>, pads : vec2<i32>, dilations : vec2<i32>, filterDims : vec2<i32>,
+       outHeight : i32, outWidth : i32`;
+  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 = 'maxPool2DBackprop';
+  }
+
+  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;
+        let lastIndex = uniforms.filterDims[0] * uniforms.filterDims[1] - 1;
+        for (var wR = 0; wR < uniforms.filterDims[0]; wR += uniforms.dilations[0]) {
+          let dyR = f32(dyRCorner + wR) / f32(uniforms.strides[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++) {
+            let dyC = f32(dyCCorner + wC) / f32(uniforms.strides[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);
+            let maxPosValue = lastIndex - i32(getMaxPos(batch, idyR, idyC, d));
+
+            // Get the current value, check it against the value from the
+            // position matrix.
+            let curPosValue = wR * uniforms.filterDims[1] + wC;
+            let mask = select(0.0, 1.0, maxPosValue == curPosValue);
+            dotProd += dyValue * mask;
+          }
+        }
+        setOutputAtIndex(index, dotProd);
+      }
+    }
+    `;
+    return userCode;
+  }
+}

tfjs-backend-webgpu/src/register_all_kernels.ts

@@ -97,6 +97,7 @@ import {lrnConfig} from './kernels/LRN';
 import {maxConfig} from './kernels/Max';
 import {maximumConfig} from './kernels/Maximum';
 import {maxPoolConfig} from './kernels/MaxPool';
+import {maxPoolGradConfig} from './kernels/MaxPoolGrad';
 import {maxPoolWithArgmaxConfig} from './kernels/MaxPoolWithArgmax';
 import {meanConfig} from './kernels/Mean';
 import {minConfig} from './kernels/Min';
@@ -243,6 +244,7 @@ const kernelConfigs: KernelConfig[] = [
   maxConfig,
   maximumConfig,
   maxPoolConfig,
+  maxPoolGradConfig,
   maxPoolWithArgmaxConfig,
   meanConfig,
   minConfig,

tfjs-backend-webgpu/src/setup_test.ts

@@ -247,7 +247,6 @@ const TEST_FILTERS: TestFilter[] = [
       'conv3dTranspose ',
       'maxPool3d ',
       'maxPool3dBackprop ',
-      'maxPoolBackprop ',
       'raggedGather ',
       'raggedRange ',
       'raggedTensorToTensor ',