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

[js/webgpu] Optimize broadcast binary. (#18185)

Browse source 
### Description
Currently, the binary algorithms are divided into the vectorize one
(efficient) and non-vectorize one (less efficient). Below situations
will go to the vectorize one:
1) A or B's shape length is 1.
2) The shared dimensions length of A and B are divisible by 4.
3) A and B have same shape.

This PR adds another situation as below to go to the vectorize
algorithm.
4. A or B's last dimension is divisible by 4.

With this change, the aggerate time of Add in sam-b-encoder becomes
309.65 ms from 409.12 ms on Intel ADL.
This commit is contained in:
Jiajia Qin 2023-11-21 08:52:17 +08:00 committed by GitHub
parent cc542024ce
commit abdf8b7c3f
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
1 changed files with 24 additions and 6 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

30
js/web/lib/wasm/jsep/webgpu/ops/binary-op.ts
View file

@ -17,8 +17,9 @@ type BinaryFunctionCall = BuiltinFunctionName|BinaryCustomExpression|{
const createBinaryOpProgramShader =
(shaderHelper: ShaderHelper, dimsA: readonly number[], dimsB: readonly number[], dimsOutput: readonly number[],
vectorize: boolean, doBroadcast: boolean, funcCall: BinaryFunctionCall, typeA: number, typeB: number,
typeOutput: number, useShapesUniforms: boolean, additionalImplementation?: string) => {
vectorize: boolean, doBroadcast: boolean, sharedDimensionDivisibleBy4: boolean, funcCall: BinaryFunctionCall,
typeA: number, typeB: number, typeOutput: number, useShapesUniforms: boolean,
additionalImplementation?: string) => {
let expressionScalar: BinaryCustomExpression;
let expressionVector: BinaryCustomExpression;
if (typeof funcCall === 'string') {
@ -42,6 +43,8 @@ const createBinaryOpProgramShader =
if (doBroadcast) {
const isAOneElement = ShapeUtil.size(dimsA) === 1;
const isBOneElement = ShapeUtil.size(dimsB) === 1;
const aLastDimDivisibleBy4 = dimsA.length > 0 && dimsA[dimsA.length - 1] % 4 === 0;
const bLastDimDivisibleBy4 = dimsB.length > 0 && dimsB[dimsB.length - 1] % 4 === 0;
if (isAOneElement || isBOneElement) {
assignment = output.setByOffset(
'global_idx',
@ -55,7 +58,14 @@ const createBinaryOpProgramShader =
let offsetB = ${b.broadcastedIndicesToOffset('outputIndices', output)};
${
output.setByOffset(
'global_idx', expressionVector(a.getByOffset('offsetA / 4u'), b.getByOffset('offsetB / 4u')))}
'global_idx',
expressionVector(
sharedDimensionDivisibleBy4 || aLastDimDivisibleBy4 ?
a.getByOffset('offsetA / 4u') :
`${a.type.value}(${a.getByOffset('offsetA / 4u')}[offsetA % 4u])`,
sharedDimensionDivisibleBy4 || bLastDimDivisibleBy4 ?
b.getByOffset('offsetB / 4u') :
`${b.type.value}(${b.getByOffset('offsetB / 4u')}[offsetB % 4u])`))}
`;
}
} else {
@ -118,6 +128,7 @@ const createBinaryOpProgramInfo =
let outputSize = ShapeUtil.size(a.dims);
let vectorize = false;
let sharedDimensionDivisibleBy4 = false;
// TODO: deal with zero-sized tensors (eg. dims=[1,0])
const cacheKeyAux = [isBroadcast];
@ -130,8 +141,12 @@ const createBinaryOpProgramInfo =
outputSize = ShapeUtil.size(outputShape);
const isAOneElement = ShapeUtil.size(a.dims) === 1;
const isBOneElement = ShapeUtil.size(b.dims) === 1;
const aLastDimDivisibleBy4 = a.dims.length > 0 && a.dims[a.dims.length - 1] % 4 === 0;
const bLastDimDivisibleBy4 = b.dims.length > 0 && b.dims[b.dims.length - 1] % 4 === 0;
cacheKeyAux.push(isAOneElement);
cacheKeyAux.push(isBOneElement);
cacheKeyAux.push(aLastDimDivisibleBy4);
cacheKeyAux.push(bLastDimDivisibleBy4);
// check whether vectorize can be enabled
let sharedDimension = 1;
for (let i = 1; i < outputShape.length; i++) {
@ -143,7 +158,10 @@ const createBinaryOpProgramInfo =
break;
}
}
if (sharedDimension % 4 === 0 || isAOneElement || isBOneElement) {
if (sharedDimension % 4 === 0) {
sharedDimensionDivisibleBy4 = true;
vectorize = true;
} else if (isAOneElement || isBOneElement || aLastDimDivisibleBy4 || bLastDimDivisibleBy4) {
vectorize = true;
}
} else {
@ -160,8 +178,8 @@ const createBinaryOpProgramInfo =
inputDependencies: useShapesUniforms ? ['rank', 'rank'] : ['dims', 'dims'],
},
getShaderSource: (shaderHelper) => createBinaryOpProgramShader(
shaderHelper, a.dims, b.dims, outputShape, vectorize, isBroadcast, funcCall, a.dataType, b.dataType,
outputDataType, useShapesUniforms, additionalImplementation),
shaderHelper, a.dims, b.dims, outputShape, vectorize, isBroadcast, sharedDimensionDivisibleBy4, funcCall,
a.dataType, b.dataType, outputDataType, useShapesUniforms, additionalImplementation),
getRunData: () => ({
outputs: [{dims: outputShape, dataType: outputDataType}],
dispatchGroup: {x: Math.ceil(outputSize / 64 /* workgroup size */ / 4 /* component size */)},