[Web/JS] Add ConvTranspose support (#16433)

### Description
Add ConvTranspose support for WebGPU


### Motivation and Context
<!-- - Why is this change required? What problem does it solve?
- If it fixes an open issue, please link to the issue here. -->
This commit is contained in:
satyajandhyala 2023-07-08 11:10:50 -07:00 committed by GitHub
parent 5c6613875c
commit 00e8f2a2a9
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
8 changed files with 846 additions and 9 deletions

View file

@ -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+) | |

View file

@ -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<string, OperatorImplementation> = new
['Clip', [unaryOps.clip]],
['Concat', [concat, parseConcatAttributes]],
['Conv', [conv, parseConvAttributes]],
['ConvTranspose', [convTranspose, parseConvTransposeAttributes]],
['Cos', [unaryOps.cos]],
['Cosh', [unaryOps.cosh]],
['Div', [binaryOps.div]],

View file

@ -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<storage, read> Dy: array<${
isVec4 && innerElementSize === 4 ? 'vec4<f32>' : 'f32'}>;`,
`@group(0) @binding(1) var<storage, read> W: array<${isVec4 ? 'vec4<f32>' : 'f32'}>;`
];
let declareFunctions = `
fn setOutputAtIndex(flatIndex : u32, value : ${isVec4 ? 'vec4<f32>' : 'f32'}) {
result[flatIndex] = ${isVec4 ? 'vec4<f32>' : 'f32'}(value);
}`;
if (hasBias) {
declareInputs.push(`@group(0) @binding(2) var<storage, read> bias: array<${isVec4 ? 'vec4<f32>' : 'f32'}>;`);
declareFunctions += `
fn getBiasByOutputCoords(coords : vec4<u32>) -> ${isVec4 ? 'vec4<f32>' : '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>(i32(r), i32(c)) - vec2<i32>(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<vec4<f32>, ${workPerThread}>;
for (var i = 0; i < ${workPerThread}; i++) {
dotProd[i] = vec4<f32>(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<f32>(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<f32>(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<f32>(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<f32>(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>(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<storage, read_write> result: array<${isVec4 ? 'vec4<f32>' : 'f32'}>;
const outShape : vec4<u32> = vec4<u32>(${outputShape.join(',')});
const outBackprop : vec4<u32> = vec4<u32>(${inputs[0].dims.join(',')});
const strides : vec2<u32> = vec2<u32>(${attributes.strides[0]}, ${attributes.strides[1]});
const filterDims : vec2<u32> = vec2<u32>(${attributes.kernelShape[isChannelsLast ? 1 : 2]}, ${
attributes.kernelShape[isChannelsLast ? 2 : 3]});
const dilations : vec2<u32> = vec2<u32>(${attributes.dilations[0]}, ${attributes.dilations[1]});
const effectiveFilterDims : vec2<u32> = filterDims + vec2<u32>(
${
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<i32> = vec2<i32>(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),
};
};

View file

@ -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 =
<T extends ConvTransposeAttributes>(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<string, unknown>): 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);
}
};

View file

@ -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",

View file

@ -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<KernelRegistry> RegisterKernels() {
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kJsExecutionProvider, kOnnxDomain, 13, Transpose)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kJsExecutionProvider, kMSInternalNHWCDomain, 11, float, Conv)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kJsExecutionProvider, kMSInternalNHWCDomain, 11, float, ConvTranspose)>,
BuildKernelCreateInfo<ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kJsExecutionProvider, kMSInternalNHWCDomain, 11, 11, float, MaxPool)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kJsExecutionProvider, kMSInternalNHWCDomain, 12, float, MaxPool)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kJsExecutionProvider, kMSInternalNHWCDomain, 11, float, AveragePool)>,
@ -391,6 +395,8 @@ std::unique_ptr<KernelRegistry> RegisterKernels() {
BuildKernelCreateInfo<ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kJsExecutionProvider, kOnnxDomain, 1, 10, float, Conv)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kJsExecutionProvider, kOnnxDomain, 11, float, Conv)>,
BuildKernelCreateInfo<ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kJsExecutionProvider, kOnnxDomain, 1, 10, float, ConvTranspose)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kJsExecutionProvider, kOnnxDomain, 11, float, ConvTranspose)>,
BuildKernelCreateInfo<ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kJsExecutionProvider, kOnnxDomain, 7, 8, float, Gemm)>,
BuildKernelCreateInfo<ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kJsExecutionProvider, kOnnxDomain, 9, 10, float, Gemm)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kJsExecutionProvider, kOnnxDomain, 11, float, Gemm)>,

View file

@ -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<T>()), \
ConvTranspose<T, true>); \
ONNX_OPERATOR_TYPED_KERNEL_EX( \
ConvTranspose, \
kOnnxDomain, \
11, \
T, \
kJsExecutionProvider, \
(*KernelDefBuilder::Create()).TypeConstraint("T", DataTypeImpl::GetTensorType<T>()), \
ConvTranspose<T, false>); \
ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_EX( \
ConvTranspose, \
kOnnxDomain, \
1, 10, \
T, \
kJsExecutionProvider, \
(*KernelDefBuilder::Create()).TypeConstraint("T", DataTypeImpl::GetTensorType<T>()), \
ConvTranspose<T, false>);
REGISTER_KERNEL_TYPED(float)
} // namespace js
} // namespace onnxruntime

View file

@ -0,0 +1,117 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#pragma once
#include <algorithm>
#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 <typename T, bool is_channels_last>
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<int64_t>("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<int32_t>(conv_transpose_attrs_.auto_pad),
static_cast<int32_t>(conv_transpose_attrs_.dilations.size() > 0 ? conv_transpose_attrs_.dilations[0] : 0),
static_cast<int32_t>(conv_transpose_attrs_.group),
static_cast<int32_t>(conv_transpose_attrs_.kernel_shape_specified && kernel_shape.size() > 0) ? kernel_shape[0] : 0,
static_cast<int32_t>(conv_transpose_attrs_.pads.size()),
static_cast<int32_t>(conv_transpose_attrs_.pads.size() > 1) ? conv_transpose_attrs_.pads[1] : 0,
static_cast<int32_t>(conv_transpose_attrs_.strides.size() > 0) ? conv_transpose_attrs_.strides[0] : 0,
static_cast<int32_t>(channels_last),
reinterpret_cast<int32_t>(&w_is_const_),
gsl::narrow_cast<int32_t>(conv_transpose_attrs_.output_shape.size()),
reinterpret_cast<int32_t>(conv_transpose_attrs_.output_padding.size() > 0 ? conv_transpose_attrs_.output_padding.data() : nullptr) >> 2,
gsl::narrow_cast<int32_t>(conv_transpose_attrs_.output_shape.size()),
reinterpret_cast<int32_t>(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<int32_t> local_pads(pads_vec_size, 0);
std::vector<int32_t> local_strides(strides_vec_size, 0);
std::vector<int32_t> local_dilations(dialations_vec_size, 0);
std::vector<int32_t> local_kernel_shape;
std::vector<int32_t> local_output_shape(conv_transpose_attrs_.output_shape.begin(), conv_transpose_attrs_.output_shape.end());
std::vector<int32_t> 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<int32_t>(conv_transpose_attrs_.auto_pad),
reinterpret_cast<int32_t>(local_dilations.data()) >> 2,
static_cast<int32_t>(conv_transpose_attrs_.group),
reinterpret_cast<int32_t>(local_kernel_shape.data()) >> 2,
reinterpret_cast<int32_t>(local_pads.data()) >> 2,
reinterpret_cast<int32_t>(local_strides.data()) >> 2,
static_cast<int32_t>(channels_last),
reinterpret_cast<int32_t>(&w_is_const_),
gsl::narrow_cast<int32_t>(local_output_padding.size()),
reinterpret_cast<int32_t>(local_output_padding.size() > 0 ? local_output_padding.data() : nullptr) >> 2,
gsl::narrow_cast<int32_t>(local_output_shape.size()),
reinterpret_cast<int32_t>(local_output_shape.size() > 0 ? local_output_shape.data() : nullptr) >> 2);
}
}
protected:
ConvTransposeAttributes conv_transpose_attrs_;
bool w_is_const_;
};
} // namespace js
} // namespace onnxruntime