[webgpu] Add crop_and_resize kernel (#2636)

FEATURE

tfjs-backend-webgpu/src/backend_webgpu.ts

@@ -19,7 +19,7 @@
 
 import './flags_webgpu';
 
-import {backend_util, DataStorage, DataType, engine, env, findBackend, KernelBackend, Rank, RecursiveArray, ShapeMap, slice_util, sumOutType, Tensor, Tensor2D, Tensor3D, Tensor4D, TimingInfo, util} from '@tensorflow/tfjs-core';
+import {backend_util, DataStorage, DataType, engine, env, findBackend, KernelBackend, Rank, RecursiveArray, ShapeMap, slice_util, sumOutType, Tensor, Tensor1D, Tensor2D, Tensor3D, Tensor4D, TimingInfo, util} from '@tensorflow/tfjs-core';
 import {Glslang} from '@webgpu/glslang/dist/web-devel/glslang.onefile';
 
 import {BufferManager} from './buffer_manager';
@@ -30,6 +30,7 @@ import {ClipProgram} from './kernels/clip_webgpu';
 import {ConcatProgram} from './kernels/concat_webgpu';
 import {Conv2DMMProgram} from './kernels/conv2d_mm_webgpu';
 import {Conv2DNaiveProgram} from './kernels/conv2d_naive_webgpu';
+import {CropAndResizeProgram} from './kernels/crop_and_resize_webgpu';
 import {DepthwiseConv2DProgram} from './kernels/depthwise_conv2d_webgpu';
 import {FillProgram} from './kernels/fill_webgpu';
 import {Im2ColProgram} from './kernels/im2col_webgpu';
@@ -974,6 +975,19 @@ export class WebGPUBackend extends KernelBackend {
     return this.compileAndRun(program, [condition, a, b], output);
   }
 
+  cropAndResize(
+      image: Tensor4D, boxes: Tensor2D, boxIndex: Tensor1D,
+      cropSize: [number, number], method: 'bilinear'|'nearest',
+      extrapolationValue: number): Tensor4D {
+    const program = new CropAndResizeProgram(
+        image.shape, boxes.shape, cropSize, method, extrapolationValue);
+    const dataId =
+        this.write(null /*values*/, program.outputShape, image.dtype);
+    const output = engine().makeTensorFromDataId(
+        dataId, program.outputShape, image.dtype, this);
+    return this.compileAndRun(program, [image, boxes, boxIndex], output);
+  }
+
   fill<R extends Rank>(
       shape: ShapeMap[R], value: number|string, dtype?: DataType): Tensor<R> {
     dtype = dtype || util.inferDtype(value);

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

@@ -0,0 +1,133 @@
+/**
+ * @license
+ * Copyright 2019 Google Inc. 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 {computeDispatch} from '../webgpu_util';
+import {WebGPUProgram} from './webgpu_program';
+
+export class CropAndResizeProgram implements WebGPUProgram {
+  outputShape: number[];
+  shaderKey: string;
+  userCode: string;
+  dispatchLayout: {x: number[], y: number[], z: number[]};
+  dispatch: [number, number, number];
+  variableNames = ['Image', 'Boxes', 'BoxInd'];
+  workGroupSize: [number, number, number] = [4, 4, 4];
+
+  constructor(
+      imageShape: [number, number, number, number], boxShape: [number, number],
+      cropSize: [number, number], method: 'bilinear'|'nearest',
+      extrapolationValue: number) {
+    const [batch, imageHeight, imageWidth, depth] = imageShape;
+    const [numBoxes, ] = boxShape;
+    const [cropHeight, cropWidth] = cropSize;
+    this.outputShape = [numBoxes, cropHeight, cropWidth, depth];
+    const methodId = method === 'bilinear' ? 1 : 0;
+
+    this.dispatchLayout = {x: [1, 2], y: [0], z: [3]};
+    this.dispatch = computeDispatch(
+        this.dispatchLayout, this.outputShape, this.workGroupSize);
+
+    const [inputHeightFloat, inputWidthFloat] =
+        [`${imageHeight - 1}.0`, `${imageWidth - 1}.0`];
+
+    const [heightRatio, heightScale, inY] = cropHeight > 1 ?
+        [
+          `${(imageHeight - 1) / (cropHeight - 1)}`,
+          '(y2-y1) * height_ratio',
+          `y1*${inputHeightFloat} + float(y)*(height_scale)`,
+        ] :
+        [
+          '0.0',
+          '0.0',
+          `0.5 * (y1+y2) * ${inputHeightFloat}`,
+        ];
+    const [widthRatio, widthScale, inX] = cropWidth > 1 ?
+        [
+          `${(imageWidth - 1) / (cropWidth - 1)}`,
+          '(x2-x1) * width_ratio',
+          `x1*${inputWidthFloat} + float(x)*(width_scale)`,
+        ] :
+        [
+          '0.0',
+          '0.0',
+          `0.5 * (x1+x2) * ${inputWidthFloat}`,
+        ];
+
+    // Reference implementation
+    // tslint:disable-next-line:max-line-length
+    // https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/kernels/crop_and_resize_op_gpu.cu.cc
+    this.userCode = `
+      const float height_ratio = float(${heightRatio});
+      const float width_ratio = float(${widthRatio});
+      void writeResult(ivec4 coords,float value) {
+        if (coordsInBounds(coords, outShape)) {
+          setOutput(coords[0], coords[1], coords[2], coords[3], value);
+        }
+      }
+      void main() {
+        ivec4 coords = getOutputCoords();
+        int b = coords[0];
+        int y = coords[1];
+        int x = coords[2];
+        int d = coords[3];
+        // get box vals
+        float y1 = getBoxes(b,0);
+        float x1 = getBoxes(b,1);
+        float y2 = getBoxes(b,2);
+        float x2 = getBoxes(b,3);
+        // get image in batch index
+        int bInd = int(round(getBoxInd(b)));
+        if(bInd < 0 || bInd >= ${batch}) {
+          return;
+        }
+        float height_scale = ${heightScale};
+        float width_scale = ${widthScale};
+        float in_y = ${inY};
+        if( in_y < 0.0 || in_y > ${inputHeightFloat} ) {
+          writeResult(coords,float(${extrapolationValue}));
+          return;
+        }
+        float in_x = ${inX};
+        if( in_x < 0.0 || in_x > ${inputWidthFloat} ) {
+          writeResult(coords,float(${extrapolationValue}));
+          return;
+        }
+        vec2 sourceFracIndexCR = vec2(in_x,in_y);
+        if(${methodId} == 1) {
+          // Compute the four integer indices.
+          ivec2 sourceFloorCR = ivec2(sourceFracIndexCR);
+          ivec2 sourceCeilCR = ivec2(ceil(sourceFracIndexCR));
+          float topLeft = getImage(bInd, sourceFloorCR.y, sourceFloorCR.x, d);
+          float bottomLeft = getImage(bInd, sourceCeilCR.y, sourceFloorCR.x, d);
+          float topRight = getImage(bInd, sourceFloorCR.y, sourceCeilCR.x, d);
+          float bottomRight = getImage(bInd, sourceCeilCR.y, sourceCeilCR.x, d);
+          vec2 fracCR = sourceFracIndexCR - vec2(sourceFloorCR);
+          float top = topLeft + (topRight - topLeft) * fracCR.x;
+          float bottom = bottomLeft + (bottomRight - bottomLeft) * fracCR.x;
+          float newValue = top + (bottom - top) * fracCR.y;
+          writeResult(coords,newValue);
+        } else {
+          // Compute the coordinators of nearest neighbor point.
+          ivec2 sourceNearestCR = ivec2(floor(
+            sourceFracIndexCR + vec2(0.5,0.5)));
+          float newValue = getImage(bInd, sourceNearestCR.y, sourceNearestCR.x, d);
+          writeResult(coords,newValue);
+        }
+      }
+    `;
+  }
+}

tfjs-backend-webgpu/src/setup_test.ts

@@ -323,6 +323,14 @@ const TEST_FILTERS: TestFilter[] = [
       'gradient',                  // zerosLike not yet implemented.
       'absoluteDifference',        // absoluteDifference not yet implemented
     ]
+  },
+  {
+    include: 'cropAndResize',
+    excludes: [
+      '2x2to3x3-NoCrop',  // The operation failed for an operation-specific
+                          // reason
+      'MultipleBoxes-DifferentBoxes',  // TimeOut
+    ]
   }
 ];