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

[JS/Web] Added Unifroms support to unary ops. (#18223)

Browse source 
### Description
Added uniforms support to unary ops.


### Motivation and Context
Improve performance
This commit is contained in:
satyajandhyala 2023-11-03 09:30:54 -07:00 committed by GitHub
parent 90f205e79c
commit e207060ac9
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
1 changed files with 7 additions and 4 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

11
js/web/lib/wasm/jsep/webgpu/ops/unary-op.ts
View file

@ -29,12 +29,12 @@ const createElementwiseProgramShader =
const output = outputVariable('outputData', outputDataType, [vecSize], 4);
return `
${shaderHelper.declareVariables(input, output)}
${shaderHelper.registerUniform('vec_size', 'u32').declareVariables(input, output)}
${additionalImplementation ?? ''}
${shaderHelper.mainStart()}
${shaderHelper.guardAgainstOutOfBoundsWorkgroupSizes(vecSize)}
${shaderHelper.guardAgainstOutOfBoundsWorkgroupSizes('uniforms.vec_size')}
let a = ${input.getByOffset('global_idx')};
${output.setByOffset('global_idx', expression)}
@ -45,13 +45,16 @@ const createElementwiseProgramInfo =
(input: TensorView, name: string, funcCall: ElementwiseFunctionCall, additionalImplementation?: string,
cacheKey?: string, outputDataType: number = input.dataType): ProgramInfo => ({
name,
shaderCache: {hint: cacheKey},
shaderCache: {hint: cacheKey, inputDependencies: ['type']},
getShaderSource: shaderHelper => createElementwiseProgramShader(
shaderHelper, ShapeUtil.size(input.dims), input.dataType, outputDataType, funcCall, additionalImplementation),
getRunData: (inputTensors) => ({
outputs: [{dims: input.dims, dataType: outputDataType}],
dispatchGroup:
{x: Math.ceil(ShapeUtil.size(inputTensors[0].dims) / 64 /* workgroup size */ / 4 /* vec size */)}
{x: Math.ceil(ShapeUtil.size(inputTensors[0].dims) / 64 /* workgroup size */ / 4 /* vec size */)},
programUniforms: [
{type: 'uint32', data: Math.ceil(ShapeUtil.size(input.dims) / 4)},
],
})
});