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] Optimize Expand (#22752)

Browse source 
Use components = 4 if possible.

llama3.2-1B becomes 20 tokens/s from 18 tokens/s on my iGPUs.
This commit is contained in:
Jiajia Qin 2024-11-13 04:37:19 +08:00 committed by GitHub
parent 05c8dc9d1c
commit 7e0dd9d433
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 63 additions and 5 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

18
js/web/lib/wasm/jsep/webgpu/ops/expand.ts
View file

@ -48,11 +48,18 @@ const createExpandProgramInfo = (inputs: readonly TensorView[]): ProgramInfo =>
const shape = Array.from(inputs[1].getBigInt64Array(), Number);
const outputShape: number[] = calculateOutputShape(inputShape, shape);
const dataType = inputs[0].dataType;
const components = dataType === DataType.bool ? 4 : 1;
const isBoolOrScalar = dataType === DataType.bool || ShapeUtil.size(inputShape) === 1;
const iComponents =
dataType === DataType.bool ? 4 : inputShape.length > 0 && inputShape[inputShape.length - 1] % 4 === 0 ? 4 : 1;
const components = isBoolOrScalar
? 4
: outputShape.length > 0 && outputShape[outputShape.length - 1] % 4 === 0
? 4
: 1;
const outputSize = Math.ceil(ShapeUtil.size(outputShape) / components);
const getShaderSource = (shaderHelper: ShaderHelper) => {
const input = inputVariable('input', dataType, inputShape.length, components);
const input = inputVariable('input', dataType, inputShape.length, iComponents);
const output = outputVariable('output', dataType, outputShape.length, components);
let assignment: string;
if (dataType === DataType.bool) {
@ -74,9 +81,10 @@ const createExpandProgramInfo = (inputs: readonly TensorView[]): ProgramInfo =>
}`;
} else {
assignment = `
let outputIndices = ${output.offsetToIndices('global_idx')};
let outputIndices = ${output.offsetToIndices(`global_idx * ${components}`)};
let inputOffset = ${input.broadcastedIndicesToOffset('outputIndices', output)};
${output.setByOffset('global_idx', input.getByOffset('inputOffset'))}
let data = ${output.type.value}(${input.getByOffset(`inputOffset / ${iComponents}`)});
${output.setByOffset('global_idx', 'data')}
}`;
}
return `
@ -92,7 +100,7 @@ const createExpandProgramInfo = (inputs: readonly TensorView[]): ProgramInfo =>
];
return {
name: 'Expand',
shaderCache: { hint: `${outputShape.length}`, inputDependencies: ['rank'] },
shaderCache: { hint: `${outputShape.length};${iComponents}${components}`, inputDependencies: ['rank'] },
getShaderSource,
getRunData: () => ({
outputs: [{ dims: outputShape, dataType: inputs[0].dataType }],

50
js/web/test/data/ops/expand.jsonc
View file

@ -134,6 +134,56 @@
"type": "float32"
}
]
},
{
"name": "Expand in components = 1, out components = 4",
"inputs": [
{
"data": [1, 2, 3, 4, 5, 6],
"dims": [3, 2, 1],
"type": "float32"
},
{
"data": [3, 1, 8],
"dims": [3],
"type": "int64"
}
],
"outputs": [
{
"data": [
1, 1, 1, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3, 4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5,
5, 5, 5, 5, 5, 6, 6, 6, 6, 6, 6, 6, 6
],
"dims": [3, 2, 8],
"type": "float32"
}
]
},
{
"name": "Expand in components = 4, out components = 4",
"inputs": [
{
"data": [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16],
"dims": [1, 1, 2, 8],
"type": "float32"
},
{
"data": [2, 1, 8],
"dims": [3],
"type": "int64"
}
],
"outputs": [
{
"data": [
1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15,
16
],
"dims": [1, 2, 2, 8],
"type": "float32"
}
]
}
]
},