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

[js/webgpu] Fix scalar uniform (#18318)

Browse source
This commit is contained in:
Xu Xing 2023-11-11 02:12:22 +08:00 committed by GitHub
parent d955885791
commit dd1bb760eb
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
3 changed files with 14 additions and 13 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

3
js/web/lib/wasm/jsep/backend-webgpu.ts
View file

@ -345,6 +345,9 @@ export class WebGpuBackend {
let maxAlignmentOfField = 1;
programUniforms.forEach(v => {
const data = typeof v.data === 'number' ? [v.data] : v.data;
if (data.length === 0) {
return;
}
// https://www.w3.org/TR/WGSL/#alignof
let baseAlignment: number;
switch (data.length) {

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

@ -120,7 +120,6 @@ const createBinaryOpProgramInfo =
let vectorize = false;
// TODO: deal with zero-sized tensors (eg. dims=[1,0])
const cacheKeyAux = [isBroadcast];
if (isBroadcast) {
const calculatedShape = BroadcastUtil.calcShape(a.dims, b.dims, false);
@ -158,10 +157,7 @@ const createBinaryOpProgramInfo =
name,
shaderCache: {
hint: cacheKey + cacheKeyAux.map((x) => x.toString()).join('_'),
// If the input is scalar then use type instead of dims because useShapesUniforms is false.
inputDependencies: useShapesUniforms ?
['rank', 'rank'] :
[a.dims.length > 0 ? 'dims' : 'type', b.dims.length > 0 ? 'dims' : 'type'],
inputDependencies: useShapesUniforms ? ['rank', 'rank'] : ['dims', 'dims'],
},
getShaderSource: (shaderHelper) => createBinaryOpProgramShader(
shaderHelper, a.dims, b.dims, outputShape, vectorize, isBroadcast, funcCall, a.dataType, b.dataType,

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

@ -258,8 +258,8 @@ export const tensorTypeToWsglValueType = (type: DataType, components: 1|2|3|4 =
return typeof mappedType === 'string' ? mappedType : mappedType[1];
};
export const createTensorShapeVariables = (dims: readonly number[]):
ProgramUniform[] => [{type: 'uint32', data: dims}, {type: 'uint32', data: ShapeUtil.computeStrides(dims)}];
export const createTensorShapeVariables = (dims: readonly number[]): ProgramUniform[] =>
dims.length === 0 ? [] : [{type: 'uint32', data: dims}, {type: 'uint32', data: ShapeUtil.computeStrides(dims)}];
/**
* A helper function to get maximum vector size for specified data length
@ -732,11 +732,13 @@ class ShaderHelperImpl implements ShaderHelper {
private declareVariable(variable: IndicesHelper, bindingIndex: number): string {
this.indicesHelpers.push(variable);
if (variable.shape.startsWith('uniforms.')) {
this.uniforms.push({name: variable.shape.replace('uniforms.', ''), type: variable.type.indices});
}
if (variable.strides.startsWith('uniforms.')) {
this.uniforms.push({name: variable.strides.replace('uniforms.', ''), type: variable.type.indices});
if (variable.rank !== 0) {
if (variable.shape.startsWith('uniforms.')) {
this.uniforms.push({name: variable.shape.replace('uniforms.', ''), type: variable.type.indices});
}
if (variable.strides.startsWith('uniforms.')) {
this.uniforms.push({name: variable.strides.replace('uniforms.', ''), type: variable.type.indices});
}
}
const access = variable.usage === 'input' ? 'read' : 'read_write';
const storageType = variable.type.storage;
@ -805,4 +807,4 @@ export const getBroadcastDims = (inShape: readonly number[], outShape: readonly
};
// TODO: remove this limitation once >4D dims are supported by uniform.
export const enableShapesUniforms = (rank: number): boolean => rank <= 4 && rank > 0;
export const enableShapesUniforms = (rank: number): boolean => rank <= 4;