diff --git a/js/web/lib/wasm/jsep/webgpu/ops/pool.ts b/js/web/lib/wasm/jsep/webgpu/ops/pool.ts index 84d04efc37..9e9b361c1a 100644 --- a/js/web/lib/wasm/jsep/webgpu/ops/pool.ts +++ b/js/web/lib/wasm/jsep/webgpu/ops/pool.ts @@ -5,7 +5,7 @@ import {env} from 'onnxruntime-common'; import {TensorView} from '../../tensor-view'; import {PoolConvUtil, ShapeUtil} from '../../util'; -import {AttributeWithCacheKey, createAttributeWithCacheKey} from '../attribute-with-cache-key'; +import {AttributeWithCacheKey} from '../attribute-with-cache-key'; import {ComputeContext, ProgramInfo, ProgramInputTensorInfoDependency, ProgramUniform} from '../types'; import {createTensorShapeVariables, getElementAt, IndicesHelper, inputVariable, outputVariable, ShaderHelper, UniformsArrayType} from './common'; @@ -63,7 +63,7 @@ const getUniformAndPadInfo = 2 is not supported for NHWC format.'); @@ -110,8 +110,8 @@ const getUniformAndPadInfo = ( shaderHelper: ShaderHelper, x: IndicesHelper, rank: number, outputShapeRank: number, attributes: AttributeType, - op1: string, op2: string, start: number, uniforms: UniformsArrayType, hasPads: boolean, pwStartEnd: boolean, - phStartEnd: boolean): string => { + op1: string, op2: string, start: number, uniforms: UniformsArrayType, hasPads: boolean, pwStartEndNotZero: boolean, + phStartEndNotZero: boolean): string => { const isChannelsLast = attributes.format === 'NHWC'; const dataType = x.type.value; const output = outputVariable('output', x.type.tensor, outputShapeRank); @@ -121,7 +121,7 @@ const generatePoolingCode = + (`${attributes.format};${attributes.ceilMode};${attributes.autoPad};${attributes.kernelShape.length}`); + +const createAveragePoolShaderKeyFromAttributes = (attributes: AveragePoolAttributes): string => + (`${createShaderKeyFromAttributes(attributes)};${attributes.countIncludePad}`); + +const createMaxPoolShaderKeyFromAttributes = (attributes: MaxPoolAttributes): string => + (`${createShaderKeyFromAttributes(attributes)};${attributes.storageOrder};${attributes.dilations}`); + const parsePoolCommonAttributes = (attributes: Record): PoolCommonAttributes => ({ format: attributes.format as FormatAttributes['format'], autoPad: ['NOTSET', 'VALID', 'SAME_UPPER', 'SAME_LOWER'][attributes.auto_pad as number], @@ -285,17 +294,14 @@ const createAveragePoolProgramInfo = } else { op2 += `value /= ${dataType}(i32(uniforms.kernelSize) - pad);`; } - const [programUniforms, uniforms, hasPads, pwStartEnd, phStartEnd] = + const [programUniforms, uniforms, hasPads, pwStartEndNotZero, phStartEndNotZero] = getUniformAndPadInfo(outputShape, adjustedAttributes); - programUniforms.push(...createTensorShapeVariables(input.dims)); - programUniforms.push(...createTensorShapeVariables(outputShape)); + programUniforms.push(...createTensorShapeVariables(input.dims), ...createTensorShapeVariables(outputShape)); const inputDependencies: ProgramInputTensorInfoDependency[] = ['rank']; return { name, - shaderCache: { - hint: attributes.cacheKey + hasPads + pwStartEnd + phStartEnd + adjustedAttributes.countIncludePad, - inputDependencies - }, + shaderCache: + {hint: `${attributes.cacheKey};${hasPads};${pwStartEndNotZero};${phStartEndNotZero}`, inputDependencies}, getRunData: () => ({ outputs: [{dims: outputShape, dataType: input.dataType}], dispatchGroup: {x: Math.ceil(ShapeUtil.size(outputShape) / 64 /* workgroup size */)}, @@ -303,7 +309,7 @@ const createAveragePoolProgramInfo = }), getShaderSource: shaderHelper => generatePoolingCode( shaderHelper, x, input.dims.length, outputShape.length, adjustedAttributes, op1, op2, 0.0, uniforms, - hasPads, pwStartEnd, phStartEnd), + hasPads, pwStartEndNotZero, phStartEndNotZero), }; }; @@ -315,8 +321,8 @@ export const parseAveragePoolAttributes = (attributes: Record): if (attr.ceilMode !== 0) { throw new Error('using ceil() in shape computation is not yet supported for AveragePool'); } - - return createAttributeWithCacheKey({countIncludePad, ...attr}); + const averagePoolAttributes = {countIncludePad, ...attr, cacheKey: ''}; + return {...averagePoolAttributes, cacheKey: createAveragePoolShaderKeyFromAttributes(averagePoolAttributes)}; }; export const averagePool = (context: ComputeContext, attributes: AveragePoolAttributes): void => { @@ -332,8 +338,7 @@ const globalPoolAttributes = { strides: [], pads: [], storageOrder: 0, - dilations: [], - cacheKey: '' + dilations: [] }; export const parseGlobalAveragePoolAttributes = (attributes: Record): AveragePoolAttributes => { @@ -361,13 +366,13 @@ const createMaxPoolProgramInfo = const op2 = ''; const x = inputVariable('x', input.dataType, input.dims.length); const inputDependencies: ProgramInputTensorInfoDependency[] = ['rank']; - const [programUniforms, uniforms, hasPads, pwStartEnd, phStartEnd] = + const [programUniforms, uniforms, hasPads, pwStartEndNotZero, phStartEndNotZero] = getUniformAndPadInfo(outputShape, adjustedAttributes); - programUniforms.push(...createTensorShapeVariables(input.dims)); - programUniforms.push(...createTensorShapeVariables(outputShape)); + programUniforms.push(...createTensorShapeVariables(input.dims), ...createTensorShapeVariables(outputShape)); return { name, - shaderCache: {hint: attributes.cacheKey + hasPads, inputDependencies}, + shaderCache: + {hint: `${attributes.cacheKey};${hasPads};${pwStartEndNotZero};${phStartEndNotZero}`, inputDependencies}, getRunData: () => ({ outputs: [{dims: outputShape, dataType: input.dataType}], dispatchGroup: {x: Math.ceil(ShapeUtil.size(outputShape) / 64 /* workgroup size */)}, @@ -375,7 +380,7 @@ const createMaxPoolProgramInfo = }), getShaderSource: shaderHelper => generatePoolingCode( shaderHelper, x, input.dims.length, outputShape.length, adjustedAttributes, op1, op2, -1e5, uniforms, - hasPads, pwStartEnd, phStartEnd), + hasPads, pwStartEndNotZero, phStartEndNotZero), }; }; @@ -396,8 +401,8 @@ export const parseMaxPoolAttributes = (attributes: Record): Max if (attr.ceilMode !== 0) { throw new Error('using ceil() in shape computation is not yet supported for MaxPool'); } - - return createAttributeWithCacheKey({storageOrder, dilations, ...attr}); + const maxPoolAttributes = {storageOrder, dilations, ...attr, cacheKey: ''}; + return {...maxPoolAttributes, cacheKey: createMaxPoolShaderKeyFromAttributes(maxPoolAttributes)}; }; export const parseGlobalMaxPoolAttributes = (attributes: Record): MaxPoolAttributes => {