saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-23 19:32:23 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

[js/webgpu] Fix bug in conv2dByMatMul path (#17369)

Browse source 
### Description
<!-- Describe your changes. -->
For the conv2dByMatMul path, the simulated matmul output shape is the
reshape of the original conv2d. So we should pass this information to
`createMatmulProgramInfo` so that it can process it correctly.
This commit is contained in:
Jiajia Qin 2023-09-02 15:16:28 +08:00 committed by GitHub
parent e745575187
commit 5e747071be
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
3 changed files with 39 additions and 39 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

11
js/web/lib/wasm/jsep/webgpu/ops/3rd-party/matmul_packed_webgpu.ts vendored
View file

@ -416,24 +416,23 @@ const matMulReadWriteFnSource =
export const createMatmulProgramInfo =
(metadata: ProgramMetadata, inputs: readonly TensorView[], activationAttributes: InternalActivationAttributes,
outputShape: readonly number[]): ProgramInfo => {
outputShape: readonly number[], reshapedOutputShape?: readonly number[]): ProgramInfo => {
const aShape = inputs[0].dims;
const bShape = inputs[1].dims;
const outerDimsA = aShape.slice(0, -2);
const outerDimsB = bShape.slice(0, -2);
const outerDims = outputShape.slice(0, -2);
const outerDims = reshapedOutputShape ? reshapedOutputShape.slice(0, -2) : outputShape.slice(0, -2);
const batchDims = inputVariable('batchDims', inputs[0].dataType, outerDims);
const batchADims = inputVariable('batchADims', inputs[0].dataType, outerDimsA);
const batchBDims = inputVariable('batchBDims', inputs[0].dataType, outerDimsB);
const variables = [batchADims, batchBDims, batchDims];
const batchSize = ShapeUtil.size(outerDims);
const dimAOuter = outputShape[outputShape.length - 2];
const dimAOuter = aShape[aShape.length - 2];
const dimInner = aShape[aShape.length - 1];
const dimBOuter = outputShape[outputShape.length - 1];
const dimBOuter = bShape[bShape.length - 1];
const isVec4 = dimInner % 4 === 0 && dimBOuter % 4 === 0;
const component = isVec4 ? 4 : 1;
const {activationFunction, applyActivation} = getActicationSnippet(activationAttributes);
// TODO: fine tune size
@ -455,7 +454,7 @@ export const createMatmulProgramInfo =
variables.push(output);
const inputVariables = [A, B];
const hasBias = inputs.length > 2;
const declareFunctions = matMulReadWriteFnSource(component, hasBias, applyActivation, variables);
const declareFunctions = matMulReadWriteFnSource(components, hasBias, applyActivation, variables);
if (hasBias) {
inputVariables.push(inputVariable('bias', inputs[2].dataType, [dimBOuter / components], components));
}

54
js/web/lib/wasm/jsep/webgpu/ops/conv.ts
View file

@ -147,6 +147,10 @@ const conv2d = (context: ComputeContext, inputs: readonly TensorView[], attribut
const hasBias = inputs.length === 3;
// const hasPreluActivationWeights = false; /* TODO: add support for prelu activation weights */
const isChannelsLast = attributes.format === 'NHWC';
if (!isChannelsLast || attributes.group !== 1) {
context.compute(createGroupedConvProgramInfoLoader(inputs, adjustedAttributes));
return;
}
// const batchSize = context.inputs[0].dims[0];
const inputHeight = inputs[0].dims[isChannelsLast ? 1 : 2];
@ -169,36 +173,30 @@ const conv2d = (context: ComputeContext, inputs: readonly TensorView[], attribut
(weightHeight === 1 && weightWidth === 1 && attributes.dilations[0] === 1 && attributes.dilations[1] === 1 &&
attributes.strides[0] === 1 && attributes.strides[1] === 1 && attributes.pads[0] === 0 &&
attributes.pads[1] === 0)) {
if (isChannelsLast && attributes.group === 1) {
// conv2dByMatMul
const transposedWeight = (context.kernelCustomData.wT as TensorView | undefined) ??
context.compute(
{
...transposeProgramMetadata,
cacheHint: weightTransposeAttribute.cacheKey,
get: () => createTransposeProgramInfo(inputs[1], weightTransposeAttribute.perm)
},
{inputs: [1], outputs: [attributes.wIsConst ? -2 : -1]})[0];
if (attributes.wIsConst && !context.kernelCustomData.wT) {
context.kernelCustomData.wT = transposedWeight;
}
const matmulInputs = [];
matmulInputs.push(inputs[0].reshape([batch, inputHeight * inputWidth, inputChannels]));
matmulInputs.push(transposedWeight.reshape([1, inputChannels, outChannels]));
if (hasBias) {
matmulInputs.push(inputs[2]);
}
context.compute(
createMatmulProgramInfoLoader(matmulInputs, adjustedAttributes, outputShape), {inputs: matmulInputs});
} else {
context.compute(createGroupedConvProgramInfoLoader(inputs, adjustedAttributes));
// conv2dByMatMul
const transposedWeight = (context.kernelCustomData.wT as TensorView | undefined) ??
context.compute(
{
...transposeProgramMetadata,
cacheHint: weightTransposeAttribute.cacheKey,
get: () => createTransposeProgramInfo(inputs[1], weightTransposeAttribute.perm)
},
{inputs: [1], outputs: [attributes.wIsConst ? -2 : -1]})[0];
if (attributes.wIsConst && !context.kernelCustomData.wT) {
context.kernelCustomData.wT = transposedWeight;
}
return;
}
if (!isChannelsLast || attributes.group !== 1) {
context.compute(createGroupedConvProgramInfoLoader(inputs, adjustedAttributes));
const matmulInputs = [];
matmulInputs.push(inputs[0].reshape([batch, inputHeight * inputWidth, inputChannels]));
matmulInputs.push(transposedWeight.reshape([1, inputChannels, outChannels]));
if (hasBias) {
matmulInputs.push(inputs[2]);
}
const matmulOutputShape = [batch, outHeight * outWidth, outChannels];
context.compute(
createMatmulProgramInfoLoader(matmulInputs, adjustedAttributes, outputShape, matmulOutputShape),
{inputs: matmulInputs});
return;
}

13
js/web/lib/wasm/jsep/webgpu/ops/matmul.ts
View file

@ -18,11 +18,14 @@ const createMatmulProgramMetadata = (hasBias: boolean, cacheHint: string) => ({
});
export const createMatmulProgramInfoLoader =
(inputs: readonly TensorView[], activationAttributes: InternalActivationAttributes, outputShape: readonly number[]):
ProgramInfoLoader => {
const metadata = createMatmulProgramMetadata(inputs.length > 2, activationAttributes.activationCacheKey);
return {...metadata, get: () => createMatmulProgramInfo(metadata, inputs, activationAttributes, outputShape)};
};
(inputs: readonly TensorView[], activationAttributes: InternalActivationAttributes, outputShape: readonly number[],
reshapedOutputShape?: readonly number[]): ProgramInfoLoader => {
const metadata = createMatmulProgramMetadata(inputs.length > 2, activationAttributes.activationCacheKey);
return {
...metadata,
get: () => createMatmulProgramInfo(metadata, inputs, activationAttributes, outputShape, reshapedOutputShape)
};
};
const validateInputs = (inputs: readonly TensorView[]): void => {
if (!inputs || inputs.length !== 2) {