diff --git a/js/web/docs/webgpu-operators.md b/js/web/docs/webgpu-operators.md index 85848753f7..87b13cba8f 100644 --- a/js/web/docs/webgpu-operators.md +++ b/js/web/docs/webgpu-operators.md @@ -23,6 +23,7 @@ Do not modify directly.* | Clip | ai.onnx(6-10,11,12,13+) | | | Concat | ai.onnx(1-3,4-10,11-12,13+) | | | Conv | ai.onnx(1-10,11+); com.ms.internal.nhwc(11+) | need perf optimization; conv3d not supported; need implementing activation | +| ConvTranspose | ai.onnx(1-10,11+); com.ms.internal.nhwc(11+) | | | Cos | ai.onnx(7+) | | | Cosh | ai.onnx(9+) | | | Div | ai.onnx(7-12,13,14+) | | diff --git a/js/web/lib/wasm/jsep/webgpu/op-resolve-rules.ts b/js/web/lib/wasm/jsep/webgpu/op-resolve-rules.ts index 34fc8c0ae9..7a6d2927eb 100644 --- a/js/web/lib/wasm/jsep/webgpu/op-resolve-rules.ts +++ b/js/web/lib/wasm/jsep/webgpu/op-resolve-rules.ts @@ -4,6 +4,7 @@ import * as binaryOps from './ops/binary-op'; import {concat, parseConcatAttributes} from './ops/concat'; import {conv, parseConvAttributes} from './ops/conv'; +import {convTranspose, parseConvTransposeAttributes} from './ops/conv-transpose'; import {gemm, parseGemmAttributes} from './ops/gemm'; import {matMul} from './ops/matmul'; import * as pool from './ops/pool'; @@ -33,6 +34,7 @@ export const WEBGPU_OP_RESOLVE_RULES: Map = new ['Clip', [unaryOps.clip]], ['Concat', [concat, parseConcatAttributes]], ['Conv', [conv, parseConvAttributes]], + ['ConvTranspose', [convTranspose, parseConvTransposeAttributes]], ['Cos', [unaryOps.cos]], ['Cosh', [unaryOps.cosh]], ['Div', [binaryOps.div]], diff --git a/js/web/lib/wasm/jsep/webgpu/ops/3rd-party/conv_backprop_webgpu.ts b/js/web/lib/wasm/jsep/webgpu/ops/3rd-party/conv_backprop_webgpu.ts new file mode 100644 index 0000000000..03b4311a04 --- /dev/null +++ b/js/web/lib/wasm/jsep/webgpu/ops/3rd-party/conv_backprop_webgpu.ts @@ -0,0 +1,385 @@ +/** + * @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. + * ============================================================================= + */ + +// sampled from [@tensorflow/tfjs] tfjs-backend-webgpu/src/conv_backprop_webgpu.ts + +import {LOG_DEBUG} from '../../../log'; +import {TensorView} from '../../../tensor'; +import {ShapeUtil} from '../../../util'; +import {GpuDataType, ProgramInfo, ProgramMetadata} from '../../types'; +import {createIndicesHelper, ShaderHelper} from '../common'; +import {ConvTransposeAttributes} from '../conv-transpose'; + +const createConvTranspose2DOpProgramShaderSource = + (shaderHelper: ShaderHelper, inputs: readonly TensorView[], attributes: ConvTransposeAttributes, + outputShape: readonly number[], hasBias: boolean, elementsPerThread: readonly number[]): string => { + const isChannelsLast = attributes.format === 'NHWC'; + const rowDim = isChannelsLast ? 1 : 2; + const colDim = isChannelsLast ? 2 : 3; + const channelDim = isChannelsLast ? 3 : 1; + const outputSize = ShapeUtil.size(outputShape); + const outChannels = outputShape[isChannelsLast ? 3 : 1]; + const inChannels = inputs[0].dims[isChannelsLast ? 3 : 1]; + const isVec4 = inChannels % 4 === 0 && outChannels % 4 === 0; + const workPerThread = isVec4 ? 2 : 1; + + const innerElementSize = isVec4 ? (isChannelsLast && inChannels % 4 !== 0 ? 3 : 4) : elementsPerThread[0]; + + const declareInputs = [ + `@group(0) @binding(0) var Dy: array<${ + isVec4 && innerElementSize === 4 ? 'vec4' : 'f32'}>;`, + `@group(0) @binding(1) var W: array<${isVec4 ? 'vec4' : 'f32'}>;` + ]; + let declareFunctions = ` + fn setOutputAtIndex(flatIndex : u32, value : ${isVec4 ? 'vec4' : 'f32'}) { + result[flatIndex] = ${isVec4 ? 'vec4' : 'f32'}(value); + }`; + if (hasBias) { + declareInputs.push(`@group(0) @binding(2) var bias: array<${isVec4 ? 'vec4' : 'f32'}>;`); + declareFunctions += ` + fn getBiasByOutputCoords(coords : vec4) -> ${isVec4 ? 'vec4' : 'f32'} { + return bias[coords.${isChannelsLast ? 'w' : 'y'}${isVec4 ? '/ 4' : ''}]; + }`; + } + const wIndicesHelper = createIndicesHelper('W', inputs[1].dims); + const dyIndicesHelper = createIndicesHelper('Dy', inputs[0].dims); + const outputIndicesHelper = createIndicesHelper('result', outputShape); + const codeSnippet4 = `{ + let batch: u32 = global_id.z / outShape[1]; + let r = global_id.z % outShape[1]; + let c = global_id.y * ${workPerThread}; + let d1: u32 = global_id.x * 4; + + let dyCorner = vec2(i32(r), i32(c)) - vec2(pads); + + // Convolve dy(?, ?, d2) with w(:, :, d1, d2) to compute dx(xR, xC, d1). + // ? = to be determined. : = across all values in that axis. + var dotProd: array, ${workPerThread}>; + for (var i = 0; i < ${workPerThread}; i++) { + dotProd[i] = vec4(0.0); + } + for (var wR: u32 = 0; wR < filterDims[0]; wR = wR + 1) { + var dyR = f32(dyCorner.x + wR) / f32(strides.x); + let wRPerm: u32= filterDims[0] - 1 - wR; + if (dyR < 0.0 || dyR >= f32(outBackprop[1]) || + fract(dyR) > 0.0) { + continue; + } + let idyR: u32 = u32(dyR); + + for (var wC: u32 = 0; wC < filterDims[1]; wC = wC + 1) { + let dyC = f32(dyCorner.y + wC) / f32(strides.y); + let dyC2 = f32(dyCorner.y + 1 + wC) / f32(strides.y); + let wCPerm: u32 = filterDims[1] - 1 - wC; + var bDyCVal = true; + var bDyCVal2 = true; + if (dyC < 0.0 || dyC >= f32(outBackprop[2]) || + fract(dyC) > 0.0) { + bDyCVal = false; + } + if (dyC2 < 0.0 || dyC2 >= f32(outBackprop[2]) || + fract(dyC2) > 0.0) { + bDyCVal2 = false; + } + + let idyC: u32 = u32(dyC); + let idyC2: u32 = u32(dyC2); + if (bDyCVal && bDyCVal2) { + let d2Length = outBackprop[3]; + for (var d2 :u32 = 0; d2 < d2Length; d2 = d2 + 4) { + ${ + wIndicesHelper.indicesVariableDeclaration( + 'wIndices0', + [ + 'd2', 'd1', 'wRPerm', 'wCPerm' + ])}; + ${ + wIndicesHelper.indicesVariableDeclaration( + 'wIndices1', + [ + 'd2', 'd1+1', 'wRPerm', 'wCPerm' + ])}; + ${ + wIndicesHelper.indicesVariableDeclaration( + 'wIndices2', + [ + 'd2', 'd1+2', 'wRPerm', 'wCPerm' + ])}; + ${ + wIndicesHelper.indicesVariableDeclaration( + 'wIndices3', + [ + 'd2', 'd1+3', 'wRPerm', 'wCPerm' + ])}; + let wValue0 = W[${wIndicesHelper.i2oExpression('wIndices0')}]; + let wValue1 = W[${wIndicesHelper.i2oExpression('wIndices1')}]; + let wValue2 = W[${wIndicesHelper.i2oExpression('wIndices2')}]; + let wValue3 = W[${wIndicesHelper.i2oExpression('wIndices3')}]; + ${ + dyIndicesHelper.indicesVariableDeclaration( + 'dyIndices', + isChannelsLast ? ['batch', 'idyR', 'idyC', 'd2'] : + [ + 'batch', 'd2', 'idyR', 'idyC' + ])}; + var xValue = Dy[${dyIndicesHelper.i2oExpression('dyIndices')}]; + let tmpval = vec4(xValue * wValue0, + xValue * wValue1, + xValue * wValue2, + xValue * wValue3); + dotProd[0] = dotProd[0] + tmpval; + + ${ + dyIndicesHelper.indicesVariableDeclaration( + 'dyIndices2', + isChannelsLast ? ['batch', 'idyR', 'idyC2', 'd2'] : + [ + 'batch', 'd2', 'idyR', 'idyC2' + ])}; + xValue = Dy[${dyIndicesHelper.i2oExpression('dyIndices')}]; + + dotProd[1] = dotProd[1] + vec4(xValue * wValue0, + xValue * wValue1, + xValue * wValue2, + xValue * wValue3); + } + } else if (bDyCVal) { + let d2Length = outBackprop[3]; + for (var d2: u32 = 0; d2 < d2Length; d2 = d2 + 4) { + ${ + wIndicesHelper.indicesVariableDeclaration( + 'wIndices0', + [ + 'd2', 'd1', 'wRPerm', 'wCPerm' + ])}; + ${ + wIndicesHelper.indicesVariableDeclaration( + 'wIndices1', + [ + 'd2', 'd1+1', 'wRPerm', 'wCPerm' + ])}; + ${ + wIndicesHelper.indicesVariableDeclaration( + 'wIndices2', + [ + 'd2', 'd1+2', 'wRPerm', 'wCPerm' + ])}; + ${ + wIndicesHelper.indicesVariableDeclaration( + 'wIndices3', + [ + 'd2', 'd1+3', 'wRPerm', 'wCPerm' + ])}; + let wValue0 = W[${wIndicesHelper.i2oExpression('wIndices0')}]; + let wValue1 = W[${wIndicesHelper.i2oExpression('wIndices1')}]; + let wValue2 = W[${wIndicesHelper.i2oExpression('wIndices2')}]; + let wValue3 = W[${wIndicesHelper.i2oExpression('wIndices3')}]; + ${ + dyIndicesHelper.indicesVariableDeclaration( + 'dyIndices', + isChannelsLast ? ['batch', 'idyR', 'idyC', 'd2'] : + [ + 'batch', 'd2', 'idyR', 'idyC' + ])}; + var xValue = Dy[${dyIndicesHelper.i2oExpression('dyIndices')}]; + let tmpval = vec4(xValue * wValue0, + xValue * wValue1, + xValue * wValue2, + xValue * wValue3); + dotProd[0] = dotProd[0] + tmpval; + } + } else if (bDyCVal2) { + let d2Length = outBackprop[3]; + for (var d2: u32 = 0; d2 < d2Length; d2 = d2 + 4) { + ${ + wIndicesHelper.indicesVariableDeclaration( + 'wIndices0', + [ + 'd2', 'd1', 'wRPerm', 'wCPerm' + ])}; + ${ + wIndicesHelper.indicesVariableDeclaration( + 'wIndices1', + [ + 'd2', 'd1+1', 'wRPerm', 'wCPerm' + ])}; + ${ + wIndicesHelper.indicesVariableDeclaration( + 'wIndices2', + [ + 'd2', 'd1+2', 'wRPerm', 'wCPerm' + ])}; + ${ + wIndicesHelper.indicesVariableDeclaration( + 'wIndices3', + [ + 'd2', 'd1+3', 'wRPerm', 'wCPerm' + ])}; + let wValue0 = W[${wIndicesHelper.i2oExpression('wIndices0')}]; + let wValue1 = W[${wIndicesHelper.i2oExpression('wIndices1')}]; + let wValue2 = W[${wIndicesHelper.i2oExpression('wIndices2')}]; + let wValue3 = W[${wIndicesHelper.i2oExpression('wIndices3')}]; + ${ + dyIndicesHelper.indicesVariableDeclaration( + 'dyIndices', + isChannelsLast ? ['batch', 'idyR', 'idyC', 'd2'] : + [ + 'batch', 'd2', 'idyR', 'idyC' + ])}; + var xValue = Dy[${dyIndicesHelper.i2oExpression('dyIndices')}]; + let tmpval = vec4(xValue * wValue0, + xValue * wValue1, + xValue * wValue2, + xValue * wValue3); + dotProd[1] = dotProd[1] + tmpval; + } + } + } + } + + for (var i: u32 = 0; i < ${workPerThread}; i = i + 1) { + ${ + outputIndicesHelper.indicesVariableDeclaration('outputIndices', [ + 'batch', 'r', 'c+i', 'd1' + ])}; + result[${outputIndicesHelper.i2oExpression('outputIndices')}] = dotProd[i]; + } + }`; + const codeSnippet = ` + ${outputIndicesHelper.o2iCall('global_idx', 'outputIndices')} + let batch = outputIndices[0]; + let d1 = outputIndices[${channelDim}]; + let dyCorner = vec2(i32(outputIndices[${rowDim}]), i32(outputIndices[${colDim}])) - pads; + let dyRCorner = dyCorner.x; + let dyCCorner = dyCorner.y; + // Convolve dy(?, ?, d2) with w(:, :, d1, d2) to compute dx(xR, xC, d1). + // ? = to be determined. : = across all values in that axis. + var dotProd = 0.0; + for (var wR: u32 = 0; wR < effectiveFilterDims.x; wR = wR + 1) { + if (wR % dilations.x != 0) { + continue; + } + let dyR = (f32(dyRCorner) + f32(wR)) / f32(strides[0]); + let wRPerm = filterDims.x - 1 - wR / dilations.x; + if (dyR < 0.0 || dyR >= f32(outBackprop[1]) || fract(dyR) > 0.0 || + wRPerm < 0) { + continue; + } + let idyR: u32 = u32(dyR); + + for (var wC: u32 = 0; wC < effectiveFilterDims.y; wC = wC + 1) { + if (wC % dilations.y != 0) { + continue; + } + let dyC = (f32(dyCCorner) + f32(wC)) / f32(strides.y); + let wCPerm = filterDims.y - 1 - wC / dilations.y; + if (dyC < 0.0 || dyC >= f32(outBackprop[2]) || + fract(dyC) > 0.0 || wCPerm < 0) { + continue; + } + let idyC: u32 = u32(dyC); + + for (var d2: u32 = 0; d2 < outBackprop[3]; d2 = d2 + 1) { + ${ + dyIndicesHelper.indicesVariableDeclaration( + 'dyIndices', + isChannelsLast ? ['batch', 'idyR', 'idyC', 'd2'] : + [ + 'batch', 'd2', 'idyR', 'idyC' + ])}; + let xValue = Dy[${dyIndicesHelper.i2oExpression('dyIndices')}]; + ${ + wIndicesHelper.indicesVariableDeclaration('wIndices', [ + 'd2', 'd1', 'wRPerm', 'wCPerm' + ])}; + + let wValue = W[${wIndicesHelper.i2oExpression('wIndices')}]; + dotProd = dotProd + xValue * wValue; + } + } + } + result[global_idx] = dotProd; + `; + + return ` +${wIndicesHelper.i2oImpl} + ${dyIndicesHelper.i2oImpl} + ${outputIndicesHelper.o2iImpl} + ${declareFunctions} + ${declareInputs.join('\n')} + @group(0) @binding(${declareInputs.length}) var result: array<${isVec4 ? 'vec4' : 'f32'}>; + const outShape : vec4 = vec4(${outputShape.join(',')}); + const outBackprop : vec4 = vec4(${inputs[0].dims.join(',')}); + const strides : vec2 = vec2(${attributes.strides[0]}, ${attributes.strides[1]}); + const filterDims : vec2 = vec2(${attributes.kernelShape[isChannelsLast ? 1 : 2]}, ${ + attributes.kernelShape[isChannelsLast ? 2 : 3]}); + const dilations : vec2 = vec2(${attributes.dilations[0]}, ${attributes.dilations[1]}); + const effectiveFilterDims : vec2 = filterDims + vec2( + ${ + attributes.dilations[0] <= 1 ? + 0 : + (attributes.kernelShape[isChannelsLast ? 1 : 2] - 1) * (attributes.dilations[0] - 1)}, + ${ + attributes.dilations[1] <= 1 ? + 0 : + (attributes.kernelShape[isChannelsLast ? 2 : 3] - 1) * (attributes.dilations[1] - 1)}); + const pads : vec2 = vec2(i32(effectiveFilterDims[0]) - 1 - (${attributes.pads[0] + attributes.pads[2]})/2, + i32(effectiveFilterDims[1]) - 1 - (${attributes.pads[1] + attributes.pads[3]})/2); + ${shaderHelper.mainStart()} + ${outputIndicesHelper.indicesVariableDeclaration('outputIndices')} + ${shaderHelper.guardAgainstOutOfBoundsWorkgroupSizes(outputSize)}; + ${isVec4 ? codeSnippet4 : codeSnippet}}`; + }; + +export const createConvTranspose2DProgramInfo = + (inputs: readonly TensorView[], metadata: ProgramMetadata, attributes: ConvTransposeAttributes, + squeezeOutputShapeFunction?: (shape: readonly number[]) => number[]): ProgramInfo => { + const hasBias = inputs.length > 2; + const isChannelsLast = attributes.format === 'NHWC'; + const outputShape = attributes.outputShape; + const batchSize = outputShape[0]; + const outWidth = outputShape[isChannelsLast ? 1 : 2]; + const outHeight = outputShape[isChannelsLast ? 2 : 3]; + const outChannels = outputShape[isChannelsLast ? 3 : 1]; + const inChannels = inputs[0].dims[isChannelsLast ? 3 : 1]; + const isVec4 = inChannels % 4 === 0 && outChannels % 4 === 0; + + const dispatchX = isChannelsLast ? outChannels : outWidth * outHeight; + const dispatchY = isChannelsLast ? outWidth * outHeight : outChannels; + const workGroupSize: [number, number, number] = + isVec4 ? [8, 8, 1] : [dispatchX <= 4 ? 4 : 16, dispatchX > 4 && dispatchY <= 4 ? 4 : 16, 1]; + const elementsPerThread = + isVec4 ? [4, 4, 1] : [dispatchX <= 4 ? 1 : 2, dispatchX > 4 && dispatchY <= 4 ? 1 : 2, 1]; + const dispatch = [ + Math.ceil(dispatchX / workGroupSize[0] / elementsPerThread[0]), + Math.ceil(dispatchY / workGroupSize[1] / elementsPerThread[1]), + Math.ceil(batchSize / workGroupSize[2] / elementsPerThread[1]) + ]; + LOG_DEBUG('verbose', () => `[conv2d_backprop_webgpu] dispatch = ${dispatch}`); + + return { + ...metadata, + outputs: [{ + dims: squeezeOutputShapeFunction ? squeezeOutputShapeFunction(outputShape) : outputShape, + dataType: inputs[0].dataType, + gpuDataType: GpuDataType.default + }], + dispatchGroup: () => ({x: dispatch[0], y: dispatch[1], z: dispatch[2]}), + getShaderSource: (shaderHelper: ShaderHelper) => createConvTranspose2DOpProgramShaderSource( + shaderHelper, inputs, attributes, outputShape, hasBias, elementsPerThread), + }; + }; diff --git a/js/web/lib/wasm/jsep/webgpu/ops/conv-transpose.ts b/js/web/lib/wasm/jsep/webgpu/ops/conv-transpose.ts new file mode 100644 index 0000000000..acdfd7e40f --- /dev/null +++ b/js/web/lib/wasm/jsep/webgpu/ops/conv-transpose.ts @@ -0,0 +1,287 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +import {DataType} from '../../../wasm-common'; +import {TensorView} from '../../tensor'; +import {createAttributeWithCacheKey} from '../attribute-with-cache-key'; +import {ComputeContext, GpuDataType, ProgramInfoLoader, ProgramMetadata} from '../types'; + +import {createConvTranspose2DProgramInfo} from './3rd-party/conv_backprop_webgpu'; +import {ConvAttributes} from './conv'; +import {parseInternalActivationAttributes} from './fuse-utils'; + +const computeTotalPad = + (inDim: number, stride: number, adj: number, kernel: number, dilation: number, outSize: number) => + (inDim - 1) * stride + adj + (kernel - 1) * dilation + 1 - outSize; + +const distributePadding = (totalPad: number, autoPad: string, pads: number[], head: number, tail: number) => { + const smallPad = Math.floor(totalPad / 2); + if (autoPad === 'SAME_UPPER') { + pads[head] = smallPad; + pads[tail] = totalPad - smallPad; + } else if (autoPad === 'SAME_LOWER') { + pads[head] = totalPad - smallPad; + pads[tail] = smallPad; + } +}; + +const calculateOutputShapeAndPads = + (inputShape: readonly number[], kernelShape: readonly number[], dilations: readonly number[], autoPad: string, + group: number, pads: number[], strides: readonly number[], isChannelLast: boolean, outputPadding: number[], + outputShape: number[]) => { + const spatialRank = inputShape.length - 2; + const updateOutputShape = outputShape.length === 0; + if (outputPadding.length === 0) { + for (let i = 0; i < spatialRank; ++i) { + outputPadding.push(0); + } + } + const batchSize = inputShape[0]; + const outChannels = kernelShape[isChannelLast ? 3 : 1] * group; + for (let i = 0, j = inputShape.length - spatialRank - (isChannelLast ? 1 : 0); i < spatialRank; ++i, ++j) { + const inSize = inputShape[j]; + const outSize = updateOutputShape ? inSize * strides[i] : outputShape[i]; + const totalPad = computeTotalPad(inSize, strides[i], pads[i], kernelShape[j], dilations[i], outSize); + distributePadding(totalPad, autoPad, pads, i, i + spatialRank); + if (updateOutputShape) { + outputShape.push( + strides[i] * (inSize - 1) + outputPadding[i] + (kernelShape[j] - 1) * dilations[i] + 1 - pads[i] - + pads[i + spatialRank]); + } + } + outputShape.splice(0, 0, batchSize); + outputShape.splice(isChannelLast ? 3 : 1, 0, outChannels); + }; + +export interface ConvTransposeAttributes extends ConvAttributes { + readonly outputPadding: readonly number[]; + readonly outputShape: readonly number[]; +} + + +const getAdjustedConvTransposeAttributes = + (attributes: T, inputs: readonly TensorView[]): T => { + const kernelShape = attributes.kernelShape.slice(); + // if kernelShape is not specified in the attributes of this op, infer it from the weight tensor dims + if (attributes.kernelShape.length === 0 || attributes.kernelShape.reduce((a, b) => a * b, 0) === 0) { + kernelShape.length = 0; + for (let i = 2; i < inputs[1].dims.length; ++i) { + kernelShape.push(inputs[1].dims[i]); + } + } + const isChannelsLast = attributes.format === 'NHWC'; + kernelShape.splice(0, 0, inputs[1].dims[0]); + kernelShape.splice(isChannelsLast ? 3 : 1, 0, inputs[1].dims[1]); + + const pads = attributes.pads.slice(); + const outputShape = attributes.outputShape.slice(); + const outputPadding = attributes.outputPadding.slice(); + const inputShape = inputs[0].dims; + let dilations = attributes.dilations.slice(); + if (dilations.reduce((a, b) => a + b, 0) === 0) { + const spatialRank = inputs[0].dims.length - 2; + dilations = new Array(spatialRank).fill(1); + } + let strides = attributes.strides.slice(); + if (strides.reduce((a, b) => a + b, 0) === 0) { + const spatialRank = inputs[0].dims.length - 2; + strides = new Array(spatialRank).fill(1); + } + // If outputShape is not specified in the attributes of this op, infer it from the parameters + // Similarly, automatically infer pads if not specified + calculateOutputShapeAndPads( + inputShape, kernelShape, dilations, attributes.autoPad, attributes.group, pads, strides, isChannelsLast, + outputPadding, outputShape); + + // always return a new object so does not modify the original attributes + const newAttributes: T = Object.assign({}, attributes); + Object.assign( + newAttributes, + {kernelShape, pads, outputPadding, outputShape, dilations, strides, cacheKey: attributes.cacheKey}); + return newAttributes; + }; + +export const parseConvTransposeAttributes = (attributes: Record): ConvTransposeAttributes => { + const activationAttributes = parseInternalActivationAttributes(attributes); + // TODO : Make this generic enough to compute default attributes for multi-dimensional conv + const format = attributes.format as 'NHWC' | 'NCHW'; + const autoPad = + ['NOTSET', 'VALID', 'SAME_UPPER', + 'SAME_LOWER'][typeof attributes.autoPad == 'undefined' ? 0 : attributes.autoPad as number]; + const dilations = attributes.dilations as [number, number]; + const group = attributes.group as number; + const kernelShape = attributes.kernelShape as [number, number]; + const pads = attributes.pads as [number, number, number, number]; + const strides = attributes.strides as [number, number]; + const wIsConst = (attributes.wIsConst as () => boolean)(); + const outputPadding = attributes.outputPadding as [number, number, number, number]; + const outputShape = attributes.outputShape as [number, number]; + return createAttributeWithCacheKey({ + autoPad, + format, + dilations, + group, + kernelShape, + outputPadding, + outputShape, + pads, + strides, + wIsConst, + ...activationAttributes + }); +}; + +const validateInputs = (inputs: readonly TensorView[], attributes: ConvTransposeAttributes): void => { + // Refer to the below link for all input checks + // https://github.com/onnx/onnx/blob/main/docs/Operators.md#ConvTranspose + if (!inputs || (inputs.length !== 2 && inputs.length !== 3)) { + throw new Error('Conv requires 2 or 3 inputs'); + } + + // TODO : Need to add support for multi-dimensional conv + if (inputs[0].dims.length !== 4 && inputs[0].dims.length !== 3) { + throw new Error('currently only support 2-dimensional conv'); + } + + if (inputs[0].dims.length !== inputs[1].dims.length) { + throw new Error('filter does not have same dimension as input'); + } + + // FILTER_IN_CHANNEL should be equal to DATA_CHANNEL + const dataChannel = inputs[0].dims[attributes.format === 'NHWC' ? inputs[0].dims.length - 1 : 1]; + const filterInChannel = inputs[1].dims[0]; + if (dataChannel !== filterInChannel) { + throw new Error('FILTER_IN_CHANNEL should be equal to DATA_CHANNEL'); + } + + const featureMaps = inputs[1].dims[1] * attributes.group; + + // if bias is provided it should be 1D and the number of elements should be equal to the number of feature maps + if (inputs.length === 3 && (inputs[2].dims.length !== 1 || inputs[2].dims[0] !== featureMaps)) { + throw new Error('invalid bias'); + } + + const spatialRank = inputs[0].dims.length - 2; + const dilationsSet = attributes.dilations.reduce((a, b) => a + b, 0) > 0; + // wrong dilations dimension + if (dilationsSet && attributes.dilations.length !== spatialRank) { + throw new Error(`dilations should be ${spatialRank}D`); + } + + const stridesSet = attributes.strides.reduce((a, b) => a + b, 0) > 0; + // Wrong strides dimension + if (stridesSet && attributes.strides.length !== spatialRank) { + throw new Error(`strides should be ${spatialRank}D`); + } + + // Wrong pads dimension + const padsSet = attributes.pads.reduce((a, b) => a + b, 0) > 0; + if (padsSet && attributes.pads.length !== spatialRank * 2) { + throw new Error(`pads should be ${spatialRank * 2}D`); + } + + // Wrong output padding dimension + if (attributes.outputPadding.length !== spatialRank && attributes.outputPadding.length !== 0) { + throw new Error(`output_padding should be ${spatialRank}D`); + } + + // if kernelShape is specified, it's data length must be 2 less than dims length of the weights tensor + // (the first 2 dims are batch_size and channels) + const kernelShapeSet = attributes.kernelShape.reduce((a, b) => a + b, 0) > 0; + if (kernelShapeSet && attributes.kernelShape.length !== 0 && + attributes.kernelShape.length !== inputs[1].dims.length - 2) { + throw new Error('invalid kernel shape'); + } + + // as with kernelShape, must have same number of spatial dims as input + if (attributes.outputShape.length !== 0 && attributes.outputShape.length !== inputs[0].dims.length - 2) { + throw new Error('invalid output shape'); + } + + // TODO : Need to add support for float64 + if (inputs[0].dataType !== DataType.float || inputs[1].dataType !== DataType.float) { + throw new Error('ConvTranspose input(X,W) should be float tensor'); + } + + if (inputs.length === 3 && inputs[2].dataType !== DataType.float) { + throw new Error('ConvTranspose input(bias) should be float tensor'); + } +}; + +const createConvTranspose2DProgramMetadata = (hasBias: boolean, cacheHint: string): ProgramMetadata => ({ + name: 'ConvTranspose2D', + inputTypes: hasBias ? [GpuDataType.default, GpuDataType.default, GpuDataType.default] : + [GpuDataType.default, GpuDataType.default], + cacheHint +}); + +const createConvTranspose2DProgramInfoLoader = + (inputs: readonly TensorView[], attributes: ConvTransposeAttributes, + squeezeOutputShapeFunction?: (shape: readonly number[]) => number[]): ProgramInfoLoader => { + const hasBias = inputs.length === 3; + const metadata = createConvTranspose2DProgramMetadata(hasBias, attributes.cacheKey); + return { + ...metadata, + get: () => createConvTranspose2DProgramInfo(inputs, metadata, attributes, squeezeOutputShapeFunction) + }; + }; + +const convTranspose2d = + (context: ComputeContext, inputs: readonly TensorView[], attributes: ConvTransposeAttributes): void => { + const adjustedAttributes = getAdjustedConvTransposeAttributes(attributes, inputs); + + context.compute(createConvTranspose2DProgramInfoLoader(inputs, adjustedAttributes)); + }; +const convTranspose1d = (context: ComputeContext, attributes: ConvTransposeAttributes): void => { + // extend the input to 2D by adding H dimension + const isChannelLast = attributes.format === 'NHWC'; + + const inputs = [ + context.inputs[0].reshape( + isChannelLast ? + // [N, W, C] -> [N, H=1, W, C] + [context.inputs[0].dims[0], 1, context.inputs[0].dims[1], context.inputs[0].dims[2]] : + // [N, C, W] -> [N, C, H=1, W] + [context.inputs[0].dims[0], context.inputs[0].dims[1], 1, context.inputs[0].dims[2]]), + //[FILTER_OUT_CHANNEL, FILTER_IN_CHANNEL, kW] -> [FILTER_OUT_CHANNEL, FILTER_IN_CHANNEL, kH=1, kW] + context.inputs[1].reshape([context.inputs[1].dims[0], context.inputs[1].dims[1], 1, context.inputs[1].dims[2]]) + ]; + if (inputs.length === 3) { + inputs.push(context.inputs[2]); + } + let kernelShape = attributes.kernelShape; + if (kernelShape.length === 0 || kernelShape[0] === 0) { + kernelShape = [context.inputs[1].dims[2]]; + } + let dilations = attributes.dilations; + if (dilations.length === 0 || dilations[0] === 0) { + dilations = [1]; + } + let strides = attributes.strides; + if (strides.length === 0 || strides[0] === 0) { + strides = [1]; + } + let pads = attributes.pads; + if (pads.length === 0) { + pads = [0, 0]; + } + pads = [0, pads[0], 0, pads[1]]; + strides = [1].concat(strides); + dilations = [1].concat(dilations); + kernelShape = [1].concat(kernelShape); + const adjustedAttributes = + getAdjustedConvTransposeAttributes({...attributes, pads, strides, dilations, kernelShape}, inputs); + context.compute(createConvTranspose2DProgramInfoLoader( + inputs, adjustedAttributes, + outputShape => isChannelLast ? [outputShape[0], outputShape[2], outputShape[3]] : + [outputShape[0], outputShape[1], outputShape[3]])); +}; + +export const convTranspose = (context: ComputeContext, attributes: ConvTransposeAttributes): void => { + validateInputs(context.inputs, attributes); + if (context.inputs[0].dims.length === 3) { + convTranspose1d(context, attributes); + } else { + convTranspose2d(context, context.inputs, attributes); + } +}; diff --git a/js/web/test/suite-test-list.jsonc b/js/web/test/suite-test-list.jsonc index b818d4b5cb..d45e461994 100644 --- a/js/web/test/suite-test-list.jsonc +++ b/js/web/test/suite-test-list.jsonc @@ -450,16 +450,16 @@ "test_conv_with_strides_padding", // // "test_convinteger_with_padding", // // "test_convinteger_without_padding", - // // "test_convtranspose_1d", + "test_convtranspose_1d", // // "test_convtranspose_3d", - // // "test_convtranspose_autopad_same", - // // "test_convtranspose_dilations", - // // "test_convtranspose_kernel_shape", - // // "test_convtranspose_output_shape", - // // "test_convtranspose_pad", - // // "test_convtranspose_pads", - // // "test_convtranspose_with_kernel", - // // "test_convtranspose", + "test_convtranspose_autopad_same", + "test_convtranspose_dilations", + "test_convtranspose_kernel_shape", + "opset{9,17}/test_convtranspose_output_shape", + "test_convtranspose_pad", + "test_convtranspose_pads", + "test_convtranspose_with_kernel", + "test_convtranspose", "test_cos_example", "test_cos", "test_cosh_example", diff --git a/onnxruntime/core/providers/js/js_execution_provider.cc b/onnxruntime/core/providers/js/js_execution_provider.cc index 0b8431edf9..112501e193 100644 --- a/onnxruntime/core/providers/js/js_execution_provider.cc +++ b/onnxruntime/core/providers/js/js_execution_provider.cc @@ -205,6 +205,7 @@ class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kJsExecutionProvider, kOnnxDomai class ONNX_OPERATOR_KERNEL_CLASS_NAME(kJsExecutionProvider, kOnnxDomain, 13, Transpose); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kJsExecutionProvider, kMSInternalNHWCDomain, 11, float, Conv); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kJsExecutionProvider, kMSInternalNHWCDomain, 11, float, ConvTranspose); class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kJsExecutionProvider, kMSInternalNHWCDomain, 11, 11, float, MaxPool); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kJsExecutionProvider, kMSInternalNHWCDomain, 12, float, MaxPool); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kJsExecutionProvider, kMSInternalNHWCDomain, 11, float, AveragePool); @@ -213,6 +214,8 @@ class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kJsExecutionProvider, kMSInternalNHW class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kJsExecutionProvider, kOnnxDomain, 1, 10, float, Conv); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kJsExecutionProvider, kOnnxDomain, 11, float, Conv); +class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kJsExecutionProvider, kOnnxDomain, 1, 10, float, ConvTranspose); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kJsExecutionProvider, kOnnxDomain, 11, float, ConvTranspose); class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kJsExecutionProvider, kOnnxDomain, 7, 8, float, Gemm); class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kJsExecutionProvider, kOnnxDomain, 9, 10, float, Gemm); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kJsExecutionProvider, kOnnxDomain, 11, float, Gemm); @@ -383,6 +386,7 @@ std::unique_ptr RegisterKernels() { BuildKernelCreateInfo, BuildKernelCreateInfo, + BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, @@ -391,6 +395,8 @@ std::unique_ptr RegisterKernels() { BuildKernelCreateInfo, BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, diff --git a/onnxruntime/core/providers/js/operators/conv_transpose.cc b/onnxruntime/core/providers/js/operators/conv_transpose.cc new file mode 100644 index 0000000000..1a2fc99ead --- /dev/null +++ b/onnxruntime/core/providers/js/operators/conv_transpose.cc @@ -0,0 +1,39 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/providers/js/js_kernel.h" +#include "core/providers/cpu/nn/conv_attributes.h" + +#include "conv_transpose.h" +namespace onnxruntime { +namespace js { +#define REGISTER_KERNEL_TYPED(T) \ + ONNX_OPERATOR_TYPED_KERNEL_EX( \ + ConvTranspose, \ + kMSInternalNHWCDomain, \ + 11, \ + T, \ + kJsExecutionProvider, \ + (*KernelDefBuilder::Create()).TypeConstraint("T", DataTypeImpl::GetTensorType()), \ + ConvTranspose); \ + ONNX_OPERATOR_TYPED_KERNEL_EX( \ + ConvTranspose, \ + kOnnxDomain, \ + 11, \ + T, \ + kJsExecutionProvider, \ + (*KernelDefBuilder::Create()).TypeConstraint("T", DataTypeImpl::GetTensorType()), \ + ConvTranspose); \ + ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_EX( \ + ConvTranspose, \ + kOnnxDomain, \ + 1, 10, \ + T, \ + kJsExecutionProvider, \ + (*KernelDefBuilder::Create()).TypeConstraint("T", DataTypeImpl::GetTensorType()), \ + ConvTranspose); + +REGISTER_KERNEL_TYPED(float) + +} // namespace js +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/js/operators/conv_transpose.h b/onnxruntime/core/providers/js/operators/conv_transpose.h new file mode 100644 index 0000000000..3e412f2f62 --- /dev/null +++ b/onnxruntime/core/providers/js/operators/conv_transpose.h @@ -0,0 +1,117 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include +#include "core/common/gsl.h" +#include "core/providers/cpu/nn/conv_transpose_attributes.h" +#include "core/providers/js/js_kernel.h" +namespace onnxruntime { +namespace js { +template +class ConvTranspose : public JsKernel { + public: + ConvTranspose(const OpKernelInfo& info) : JsKernel(info), conv_transpose_attrs_(info), w_is_const_(false) { + TensorShapeVector kernel_shape; + if (conv_transpose_attrs_.kernel_shape_specified) { + ORT_ENFORCE(info.GetAttrs("kernel_shape", kernel_shape).IsOK()); + } + + int64_t channels_last = is_channels_last ? 1 : info.GetAttrOrDefault("channels_last", 0); + + // currently only support Conv 1D/2D. TODO: support Conv3D and other + if (conv_transpose_attrs_.dilations.size() == 1 || + (conv_transpose_attrs_.kernel_shape_specified && kernel_shape.size() == 1) || + conv_transpose_attrs_.strides.size() == 1) { + JSEP_INIT_KERNEL_ATTRIBUTE(ConvTranspose, ({ + "format" : $8 ? "NHWC" : "NCHW", + "autoPad" : $1, + "dilations" : [$2], + "group" : $3, + "kernel_shape" : [$4], + "pads" : [ $5, $6 ], + "strides" : [$7], + "wIsConst" : () JS_ARROW(!!HEAP8[$9]), + "outputPadding" : $10 ? Array.from(HEAP32.subarray($11, $11 + $10)) : [], + "outputShape" : $12 ? Array.from(HEAP32.subarray($12, $13 + $12)) : [] + }), + static_cast(conv_transpose_attrs_.auto_pad), + static_cast(conv_transpose_attrs_.dilations.size() > 0 ? conv_transpose_attrs_.dilations[0] : 0), + static_cast(conv_transpose_attrs_.group), + static_cast(conv_transpose_attrs_.kernel_shape_specified && kernel_shape.size() > 0) ? kernel_shape[0] : 0, + static_cast(conv_transpose_attrs_.pads.size()), + static_cast(conv_transpose_attrs_.pads.size() > 1) ? conv_transpose_attrs_.pads[1] : 0, + static_cast(conv_transpose_attrs_.strides.size() > 0) ? conv_transpose_attrs_.strides[0] : 0, + static_cast(channels_last), + reinterpret_cast(&w_is_const_), + gsl::narrow_cast(conv_transpose_attrs_.output_shape.size()), + reinterpret_cast(conv_transpose_attrs_.output_padding.size() > 0 ? conv_transpose_attrs_.output_padding.data() : nullptr) >> 2, + gsl::narrow_cast(conv_transpose_attrs_.output_shape.size()), + reinterpret_cast(conv_transpose_attrs_.output_shape.size() > 0 ? conv_transpose_attrs_.output_shape.data() : nullptr) >> 2); + } else { + constexpr size_t pads_vec_size = 4; + constexpr size_t strides_vec_size = 2; + constexpr size_t dialations_vec_size = 2; + constexpr size_t kernel_shape_vec_size = 2; + // First set default values for pads, strides and dialations + std::vector local_pads(pads_vec_size, 0); + std::vector local_strides(strides_vec_size, 0); + std::vector local_dilations(dialations_vec_size, 0); + std::vector local_kernel_shape; + std::vector local_output_shape(conv_transpose_attrs_.output_shape.begin(), conv_transpose_attrs_.output_shape.end()); + std::vector local_output_padding(conv_transpose_attrs_.output_padding.begin(), conv_transpose_attrs_.output_padding.end()); + if (conv_transpose_attrs_.kernel_shape_specified) { + for (size_t i = 0; i < kernel_shape.size() && i < kernel_shape_vec_size; ++i) { + local_kernel_shape.push_back(kernel_shape[i]); + } + } else { + for (size_t i = 0; i < kernel_shape_vec_size; ++i) { + local_kernel_shape.push_back(0); + } + } + for (size_t i = 0; i < conv_transpose_attrs_.pads.size() && i < pads_vec_size; ++i) { + local_pads[i] = conv_transpose_attrs_.pads[i]; + } + for (size_t i = 0; i < conv_transpose_attrs_.dilations.size() && i < dialations_vec_size; ++i) { + local_dilations[i] = conv_transpose_attrs_.dilations[i]; + } + for (size_t i = 0; i < conv_transpose_attrs_.strides.size() && i < strides_vec_size; ++i) { + local_strides[i] = conv_transpose_attrs_.strides[i]; + } + LOGS_DEFAULT(VERBOSE) << "output_shape = " << conv_transpose_attrs_.output_shape << std::endl; + LOGS_DEFAULT(VERBOSE) << "output_padding = " << conv_transpose_attrs_.output_padding << std::endl; + JSEP_INIT_KERNEL_ATTRIBUTE(ConvTranspose, ({ + "format" : $7 ? "NHWC" : "NCHW", + "autoPad" : $1, + "dilations" : Array.from(HEAP32.subarray($2, $2 + /* dialations_vec_size */ 2)), + "group" : $3, + "kernelShape" : Array.from(HEAP32.subarray($4, $4 + /* kernel_shape_vec_size */ 2)), + "pads" : Array.from(HEAP32.subarray($5, $5 + /* pads_vec_size */ 4)), + "strides" : Array.from(HEAP32.subarray($6, $6 + /* strides_vec_size */ 2)), + "wIsConst" : () JS_ARROW(!!HEAP8[$8]), + "outputPadding" : ($9 > 0) ? Array.from(HEAP32.subarray($10, $10 + $9)) : [], + "outputShape" : ($11 > 0) ? Array.from(HEAP32.subarray($12, $12 + $11)) : [] + }), + static_cast(conv_transpose_attrs_.auto_pad), + reinterpret_cast(local_dilations.data()) >> 2, + static_cast(conv_transpose_attrs_.group), + reinterpret_cast(local_kernel_shape.data()) >> 2, + reinterpret_cast(local_pads.data()) >> 2, + reinterpret_cast(local_strides.data()) >> 2, + static_cast(channels_last), + reinterpret_cast(&w_is_const_), + gsl::narrow_cast(local_output_padding.size()), + reinterpret_cast(local_output_padding.size() > 0 ? local_output_padding.data() : nullptr) >> 2, + gsl::narrow_cast(local_output_shape.size()), + reinterpret_cast(local_output_shape.size() > 0 ? local_output_shape.data() : nullptr) >> 2); + } + } + + protected: + ConvTransposeAttributes conv_transpose_attrs_; + bool w_is_const_; +}; + +} // namespace js +} // namespace onnxruntime