From d1f657e472fc34c58006bd0668ac6618a0934087 Mon Sep 17 00:00:00 2001 From: Xinghua Cao Date: Mon, 19 Dec 2022 15:55:43 +0800 Subject: [PATCH 1/3] webgpu: support MaxPoolGrad kernel --- .../src/kernels/MaxPoolGrad.ts | 80 ++++++++++++++++ .../src/max_pool2d_backprop_webgpu.ts | 93 +++++++++++++++++++ .../src/register_all_kernels.ts | 2 + tfjs-backend-webgpu/src/setup_test.ts | 1 - 4 files changed, 175 insertions(+), 1 deletion(-) create mode 100644 tfjs-backend-webgpu/src/kernels/MaxPoolGrad.ts create mode 100644 tfjs-backend-webgpu/src/max_pool2d_backprop_webgpu.ts diff --git a/tfjs-backend-webgpu/src/kernels/MaxPoolGrad.ts b/tfjs-backend-webgpu/src/kernels/MaxPoolGrad.ts new file mode 100644 index 00000000000..d837e911d56 --- /dev/null +++ b/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 avgPoolBackpropProgram = 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( + avgPoolBackpropProgram, [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 +}; diff --git a/tfjs-backend-webgpu/src/max_pool2d_backprop_webgpu.ts b/tfjs-backend-webgpu/src/max_pool2d_backprop_webgpu.ts new file mode 100644 index 00000000000..0fc59180b70 --- /dev/null +++ b/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 = + `stride : vec2, pad : vec2, dilation : vec2, filterDims : vec2, + 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 = `max_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(coords.yz) - uniforms.pad; + 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 = 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 + 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); + 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 = dotProd + dyValue * mask; + } + } + setOutputAtIndex(index, dotProd); + } + } + `; + return userCode; + } +} diff --git a/tfjs-backend-webgpu/src/register_all_kernels.ts b/tfjs-backend-webgpu/src/register_all_kernels.ts index 7071f0c997f..014fd6a2180 100644 --- a/tfjs-backend-webgpu/src/register_all_kernels.ts +++ b/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, diff --git a/tfjs-backend-webgpu/src/setup_test.ts b/tfjs-backend-webgpu/src/setup_test.ts index 55b6a7386a8..4af0b7f1796 100644 --- a/tfjs-backend-webgpu/src/setup_test.ts +++ b/tfjs-backend-webgpu/src/setup_test.ts @@ -247,7 +247,6 @@ const TEST_FILTERS: TestFilter[] = [ 'conv3dTranspose ', 'maxPool3d ', 'maxPool3dBackprop ', - 'maxPoolBackprop ', 'raggedGather ', 'raggedRange ', 'raggedTensorToTensor ', From 65de9118328c0353bbf7156c1213008453ff8a31 Mon Sep 17 00:00:00 2001 From: Xinghua Cao Date: Thu, 12 Jan 2023 10:34:21 +0800 Subject: [PATCH 2/3] Address Jiajia's comments --- tfjs-backend-webgpu/src/kernels/MaxPoolGrad.ts | 4 ++-- .../src/max_pool2d_backprop_webgpu.ts | 14 +++++++------- 2 files changed, 9 insertions(+), 9 deletions(-) diff --git a/tfjs-backend-webgpu/src/kernels/MaxPoolGrad.ts b/tfjs-backend-webgpu/src/kernels/MaxPoolGrad.ts index d837e911d56..2f0fc25cc2a 100644 --- a/tfjs-backend-webgpu/src/kernels/MaxPoolGrad.ts +++ b/tfjs-backend-webgpu/src/kernels/MaxPoolGrad.ts @@ -50,7 +50,7 @@ export function maxPoolGrad(args: { const maxPoolPositions = backend.runWebGPUProgram( maxPoolPositionsProgram, [x], 'int32', uniformData); - const avgPoolBackpropProgram = new MaxPool2DBackpropProgram(convInfo); + const maxPoolBackpropProgram = new MaxPool2DBackpropProgram(convInfo); uniformData = [ {type: 'int32', data: [convInfo.strideHeight, convInfo.strideWidth]}, { type: 'int32', @@ -67,7 +67,7 @@ export function maxPoolGrad(args: { {type: 'int32', data: [convInfo.outWidth]} ]; const result = backend.runWebGPUProgram( - avgPoolBackpropProgram, [dy, maxPoolPositions], x.dtype, uniformData); + maxPoolBackpropProgram, [dy, maxPoolPositions], x.dtype, uniformData); backend.disposeData(maxPoolPositions.dataId); return result; diff --git a/tfjs-backend-webgpu/src/max_pool2d_backprop_webgpu.ts b/tfjs-backend-webgpu/src/max_pool2d_backprop_webgpu.ts index 0fc59180b70..e204c81cc60 100644 --- a/tfjs-backend-webgpu/src/max_pool2d_backprop_webgpu.ts +++ b/tfjs-backend-webgpu/src/max_pool2d_backprop_webgpu.ts @@ -26,7 +26,7 @@ export class MaxPool2DBackpropProgram implements WebGPUProgram { dispatch: [number, number, number]; variableNames = ['dy', 'maxPos']; uniforms = - `stride : vec2, pad : vec2, dilation : vec2, filterDims : vec2, + `strides : vec2, pads : vec2, dilations : vec2, filterDims : vec2, outHeight : i32, outWidth : i32`; workgroupSize: [number, number, number] = [64, 1, 1]; size = true; @@ -50,7 +50,7 @@ export class MaxPool2DBackpropProgram implements WebGPUProgram { let batch = coords[0]; let d = coords[3]; - let dyRCCorner = vec2(coords.yz) - uniforms.pad; + let dyRCCorner = vec2(coords.yz) - uniforms.pads; let dyRCorner = dyRCCorner.x; let dyCCorner = dyRCCorner.y; @@ -58,16 +58,16 @@ export class MaxPool2DBackpropProgram implements WebGPUProgram { // ? = 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 = wR + uniforms.dilation[0]) { - let dyR = f32(dyRCorner + wR) / f32(uniforms.stride[0]); + 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 = wC + 1) { - let dyC = f32(dyCCorner + wC) / f32(uniforms.stride[1]); + 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; @@ -81,7 +81,7 @@ export class MaxPool2DBackpropProgram implements WebGPUProgram { // position matrix. let curPosValue = wR * uniforms.filterDims[1] + wC; let mask = select(0.0, 1.0, maxPosValue == curPosValue); - dotProd = dotProd + dyValue * mask; + dotProd += dyValue * mask; } } setOutputAtIndex(index, dotProd); From 790008882e8937d71c03e3318f88417abd29ce00 Mon Sep 17 00:00:00 2001 From: Xinghua Cao Date: Tue, 17 Jan 2023 10:57:19 +0800 Subject: [PATCH 3/3] Address Yang's comment --- tfjs-backend-webgpu/src/max_pool2d_backprop_webgpu.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tfjs-backend-webgpu/src/max_pool2d_backprop_webgpu.ts b/tfjs-backend-webgpu/src/max_pool2d_backprop_webgpu.ts index e204c81cc60..79ba76f6d44 100644 --- a/tfjs-backend-webgpu/src/max_pool2d_backprop_webgpu.ts +++ b/tfjs-backend-webgpu/src/max_pool2d_backprop_webgpu.ts @@ -39,7 +39,7 @@ export class MaxPool2DBackpropProgram implements WebGPUProgram { this.dispatch = computeDispatch( this.dispatchLayout, this.outputShape, this.workgroupSize); - this.shaderKey = `max_pool2d_backprop`; + this.shaderKey = 'maxPool2DBackprop'; } getUserCode(): string {