From 5e747071be882efd6b54d7a7421042e68dcd6aff Mon Sep 17 00:00:00 2001 From: Jiajia Qin Date: Sat, 2 Sep 2023 15:16:28 +0800 Subject: [PATCH] [js/webgpu] Fix bug in conv2dByMatMul path (#17369) ### Description 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. --- .../ops/3rd-party/matmul_packed_webgpu.ts | 11 ++-- js/web/lib/wasm/jsep/webgpu/ops/conv.ts | 54 +++++++++---------- js/web/lib/wasm/jsep/webgpu/ops/matmul.ts | 13 +++-- 3 files changed, 39 insertions(+), 39 deletions(-) diff --git a/js/web/lib/wasm/jsep/webgpu/ops/3rd-party/matmul_packed_webgpu.ts b/js/web/lib/wasm/jsep/webgpu/ops/3rd-party/matmul_packed_webgpu.ts index fee872f412..ab4f608451 100644 --- a/js/web/lib/wasm/jsep/webgpu/ops/3rd-party/matmul_packed_webgpu.ts +++ b/js/web/lib/wasm/jsep/webgpu/ops/3rd-party/matmul_packed_webgpu.ts @@ -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)); } diff --git a/js/web/lib/wasm/jsep/webgpu/ops/conv.ts b/js/web/lib/wasm/jsep/webgpu/ops/conv.ts index afac503290..3a83b1c2de 100644 --- a/js/web/lib/wasm/jsep/webgpu/ops/conv.ts +++ b/js/web/lib/wasm/jsep/webgpu/ops/conv.ts @@ -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; } diff --git a/js/web/lib/wasm/jsep/webgpu/ops/matmul.ts b/js/web/lib/wasm/jsep/webgpu/ops/matmul.ts index 2d5750c3e2..e4dae00db6 100644 --- a/js/web/lib/wasm/jsep/webgpu/ops/matmul.ts +++ b/js/web/lib/wasm/jsep/webgpu/ops/matmul.ts @@ -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) {