diff --git a/js/web/lib/wasm/jsep/backend-webgpu.ts b/js/web/lib/wasm/jsep/backend-webgpu.ts index eb40da0488..e4f1808057 100644 --- a/js/web/lib/wasm/jsep/backend-webgpu.ts +++ b/js/web/lib/wasm/jsep/backend-webgpu.ts @@ -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) { diff --git a/js/web/lib/wasm/jsep/webgpu/ops/binary-op.ts b/js/web/lib/wasm/jsep/webgpu/ops/binary-op.ts index 0992552eba..0841da11d9 100644 --- a/js/web/lib/wasm/jsep/webgpu/ops/binary-op.ts +++ b/js/web/lib/wasm/jsep/webgpu/ops/binary-op.ts @@ -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, diff --git a/js/web/lib/wasm/jsep/webgpu/ops/common.ts b/js/web/lib/wasm/jsep/webgpu/ops/common.ts index 7352517733..1917c85fe5 100644 --- a/js/web/lib/wasm/jsep/webgpu/ops/common.ts +++ b/js/web/lib/wasm/jsep/webgpu/ops/common.ts @@ -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;