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

[js/webgpu] Support range uniforms (#19055)

Browse source
This commit is contained in:
Xu Xing 2024-01-10 01:33:57 +08:00 committed by GitHub
parent c1367ae553
commit eb92681bfb
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
1 changed files with 27 additions and 14 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

41
js/web/lib/wasm/jsep/webgpu/ops/range.ts
View file

@ -3,10 +3,10 @@
import {env} from 'onnxruntime-common';
import {DataType} from '../../../wasm-common';
import {ComputeContext, ProgramInfo} from '../types';
import {DataType, tensorDataTypeEnumToString} from '../../../wasm-common';
import {ComputeContext, ProgramInfo, ProgramUniform} from '../types';
import {outputVariable, ShaderHelper} from './common';
import {createTensorShapeVariables, outputVariable, ShaderHelper, UniformDataElementType, UniformsArrayType} from './common';
const validateInputsContent = (start: number, limit: number, delta: number): void => {
const sameStartLimit = start === limit;
@ -22,23 +22,36 @@ const createRangeProgramInfo = (start: number, limit: number, delta: number, dat
const numElements = Math.abs(Math.ceil((limit - start) / delta));
const outputShape: number[] = [numElements];
const outputSize = numElements;
const tensorDataType = tensorDataTypeEnumToString(dataType) as ProgramUniform['type'];
const programUniforms: ProgramUniform[] = [
{type: 'uint32', data: outputSize}, {type: tensorDataType, data: start}, {type: tensorDataType, data: delta},
...createTensorShapeVariables(outputShape)
];
const output = outputVariable('output', dataType, outputShape);
const wgslType = output.type.storage;
const getShaderSource = (shaderHelper: ShaderHelper) => `
${shaderHelper.declareVariables(output)}
const getShaderSource = (shaderHelper: ShaderHelper) => {
const output = outputVariable('output', dataType, outputShape.length);
const wgslType = output.type.value;
const uniforms: UniformsArrayType = [
{name: 'outputSize', type: 'u32'}, {name: 'start', type: wgslType as UniformDataElementType},
{name: 'delta', type: wgslType as UniformDataElementType}
];
return `
${shaderHelper.registerUniforms(uniforms).declareVariables(output)}
${shaderHelper.mainStart()}
${shaderHelper.guardAgainstOutOfBoundsWorkgroupSizes(outputSize)}
output[global_idx] = ${wgslType}(${start}) + ${wgslType}(global_idx) * ${wgslType}(${delta});
${shaderHelper.guardAgainstOutOfBoundsWorkgroupSizes('uniforms.outputSize')}
output[global_idx] = uniforms.start + ${wgslType}(global_idx) * uniforms.delta;
}`;
};
return {
name: 'Range',
shaderCache: {hint: [start, limit, delta].map(x => x.toString()).join('_')},
shaderCache: {hint: `${dataType}`},
getShaderSource,
getRunData: () => (
{outputs: [{dims: outputShape, dataType}],
dispatchGroup: {x: Math.ceil(outputSize / 64 /* workgroup size */)}})
getRunData: () => ({
outputs: [{dims: outputShape, dataType}],
dispatchGroup: {x: Math.ceil(outputSize / 64 /* workgroup size */)},
programUniforms
})
};
};