mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-27 20:02:15 +00:00
[js/webgpu] Refactor attributes of pool (#18728)
This commit is contained in:
Xu Xing
GitHub
parent
dbb8680bdc
commit
0bc71b0c9b
1 changed files with 33 additions and 28 deletions
|
|
@ -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 = <AttributeType extends AveragePoolAttributes|MaxPoo
|
|||
const sw = attributes.strides[attributes.strides.length - 1];
|
||||
const pwStart = attributes.pads[attributes.pads.length / 2 - 1];
|
||||
const pwEnd = attributes.pads[attributes.pads.length - 1];
|
||||
const pwStartEnd = !!(pwStart + pwEnd);
|
||||
const pwStartEndNotZero = !!(pwStart + pwEnd);
|
||||
programUniforms.push(
|
||||
{type: 'uint32', data: kw},
|
||||
{type: 'uint32', data: sw},
|
||||
|
|
@ -74,13 +74,13 @@ const getUniformAndPadInfo = <AttributeType extends AveragePoolAttributes|MaxPoo
|
|||
{name: 'kw', type: 'u32'}, {name: 'sw', type: 'u32'}, {name: 'pwStart', type: 'u32'},
|
||||
{name: 'pwEnd', type: 'u32'});
|
||||
|
||||
let phStartEnd = false;
|
||||
let phStartEndNotZero = false;
|
||||
if (attributes.kernelShape.length === 2) {
|
||||
const kh = attributes.kernelShape[attributes.kernelShape.length - 2];
|
||||
const sh = attributes.strides[attributes.strides.length - 2];
|
||||
const phStart = attributes.pads[attributes.pads.length / 2 - 2];
|
||||
const phEnd = attributes.pads[attributes.pads.length - 2];
|
||||
phStartEnd = !!(phStart + phEnd);
|
||||
phStartEndNotZero = !!(phStart + phEnd);
|
||||
programUniforms.push(
|
||||
{type: 'uint32', data: kh}, {type: 'uint32', data: sh}, {type: 'uint32', data: phStart},
|
||||
{type: 'uint32', data: phEnd});
|
||||
|
|
@ -89,7 +89,7 @@ const getUniformAndPadInfo = <AttributeType extends AveragePoolAttributes|MaxPoo
|
|||
{name: 'kh', type: 'u32'}, {name: 'sh', type: 'u32'}, {name: 'phStart', type: 'u32'},
|
||||
{name: 'phEnd', type: 'u32'});
|
||||
}
|
||||
return [programUniforms, uniforms, true, pwStartEnd, phStartEnd];
|
||||
return [programUniforms, uniforms, true, pwStartEndNotZero, phStartEndNotZero];
|
||||
} else {
|
||||
if (isChannelsLast) {
|
||||
throw new Error('Pooling with kernelShape.length > 2 is not supported for NHWC format.');
|
||||
|
|
@ -110,8 +110,8 @@ const getUniformAndPadInfo = <AttributeType extends AveragePoolAttributes|MaxPoo
|
|||
|
||||
const generatePoolingCode = <AttributeType extends AveragePoolAttributes|MaxPoolAttributes>(
|
||||
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 = <AttributeType extends AveragePoolAttributes|MaxPool
|
|||
let codeH = '';
|
||||
let codeHEnd = '';
|
||||
const dimIdxW = rank - (isChannelsLast ? 2 : 1);
|
||||
if (pwStartEnd === true) {
|
||||
if (pwStartEndNotZero) {
|
||||
codeW = `
|
||||
for (var i: u32 = 0u; i < uniforms.kw; i++) {
|
||||
xIndices[${dimIdxW}] = indices[${dimIdxW}] * uniforms.sw - uniforms.pwStart + i;
|
||||
|
|
@ -144,7 +144,7 @@ const generatePoolingCode = <AttributeType extends AveragePoolAttributes|MaxPool
|
|||
|
||||
if (attributes.kernelShape.length === 2) {
|
||||
const dimIdxH = rank - (isChannelsLast ? 3 : 2);
|
||||
if (phStartEnd === true) {
|
||||
if (phStartEndNotZero) {
|
||||
codeH = `
|
||||
for (var j: u32 = 0u; j < uniforms.kh; j++) {
|
||||
xIndices[${dimIdxH}] = indices[${dimIdxH}] * uniforms.sh - uniforms.phStart + j;
|
||||
|
|
@ -258,6 +258,15 @@ export interface PoolCommonAttributes extends FormatAttributes {
|
|||
readonly pads: readonly number[];
|
||||
}
|
||||
|
||||
const createShaderKeyFromAttributes = (attributes: PoolCommonAttributes): string =>
|
||||
(`${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<string, unknown>): 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<string, unknown>):
|
|||
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<string, unknown>): 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<string, unknown>): 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<string, unknown>): MaxPoolAttributes => {
|
||||
|
|
|
|||
Loading…
Reference in a new issue