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

[js/webgpu] Optimize InstanceNorm in some shapes (#22637)

Browse source 
BUG #22031

Optimize below two situations:
1. Increase workgroupSize if only one workgroup is dispatched.
2. Avoid transpose if not necessary.

The overall time of demucs model becomes 106.36 ms from 154.60 ms on my
dGPUs with this PR and PR #22577
This commit is contained in:
Jiajia Qin 2024-10-30 08:10:14 +08:00 committed by GitHub
parent 4850bcd896
commit 04e696d8e0
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 16 additions and 7 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

23
js/web/lib/wasm/jsep/webgpu/ops/instance-norm.ts
View file

@ -36,7 +36,10 @@ const computeChannelScaleShift = (
const f32Type = components === 1 ? 'f32' : `vec${components}f`;
const wgType = components === 1 ? 'vec2f' : `mat2x${components}f`;
const unitsOfWork = n * c;
let workgroupSize = 64;
if (unitsOfWork === 1) {
workgroupSize = 256;
}
const inputShape = [n, c, h / components];
const outputShape = [n, c, 2];
const inputDependencies: ProgramInputTensorInfoDependency[] = ['rank', 'type', 'type'];
@ -49,7 +52,6 @@ const computeChannelScaleShift = (
const b = inputVariable('bias', bias.dataType, bias.dims);
const output = outputVariable('output', DataType.float, 3, 2);
const variables = [x, s, b, output];
const workgroupSize = 64;
return `
var<workgroup> workgroup_shared : array<${wgType}, ${workgroupSize}>;
const workgroup_size = ${workgroupSize}u;
@ -91,7 +93,7 @@ const computeChannelScaleShift = (
{
name: 'InstanceNormComputeChannelScaleShift',
// TODO: use epsilon as uniform. Currently epsilon as uniform fails test_instancenorm_epsilon.
shaderCache: { hint: `${components};${epsilon}`, inputDependencies },
shaderCache: { hint: `${components};${epsilon};${workgroupSize}`, inputDependencies },
getRunData: () => ({
outputs: [{ dims: outputShape, dataType: DataType.float }],
dispatchGroup: { x: unitsOfWork },
@ -187,14 +189,21 @@ const createInstanceNormNHWCProgramInfo = (
const inputDependencies: ProgramInputTensorInfoDependency[] = ['type', 'type'];
// 1. transpose x from NHWC to NCHW
let needTranspose = false;
const transposedXPerm = [0, xShape.length - 1];
for (let i = 0; i < xShape.length - 2; i++) {
needTranspose = needTranspose || xShape[i + 1] !== 1;
transposedXPerm.push(i + 1);
}
const transposedX = context.compute(createTransposeProgramInfo(context.inputs[0], transposedXPerm), {
inputs: [context.inputs[0]],
outputs: [-1],
})[0];
needTranspose = needTranspose && xShape[xShape.length - 1] !== 1;
const transposedX = needTranspose
? context.compute(createTransposeProgramInfo(context.inputs[0], transposedXPerm), {
inputs: [context.inputs[0]],
outputs: [-1],
})[0]
: context.inputs[0].reshape(Array.from({ length: xShape.length }, (_, i) => xShape[transposedXPerm[i]]));
// 2. compute channel scale and channel shift.
const channelScaleShift = computeChannelScaleShift(
context,