diff --git a/js/web/lib/wasm/jsep/webgpu/ops/layer-norm.ts b/js/web/lib/wasm/jsep/webgpu/ops/layer-norm.ts index 36eefa2179..48627bfaec 100644 --- a/js/web/lib/wasm/jsep/webgpu/ops/layer-norm.ts +++ b/js/web/lib/wasm/jsep/webgpu/ops/layer-norm.ts @@ -15,7 +15,7 @@ export interface LayerNormAttributes extends AttributeWithCacheKey { } const validateInputs = (inputs: readonly TensorView[]): void => { - if (!inputs || inputs.length <= 2) { + if (!inputs || inputs.length < 2) { throw new Error('layerNorm requires at least 2 inputs.'); } @@ -41,7 +41,7 @@ const createLayerNormProgramInfo = const biasSize = bias ? ShapeUtil.size(bias.dims) : 0; if (scaleSize !== normSize || (bias && biasSize !== normSize)) { throw new Error(`Size of X.shape()[axis:] == ${normSize}. - Size of scale and bias (if provided) must match this. + Size of scale and bias (if provided) must match this. Got scale size of ${scaleSize} and bias size of ${biasSize}`); } @@ -58,17 +58,24 @@ const createLayerNormProgramInfo = const hasMeanDataOutput = outputCount > 1; const hasInvStdOutput = outputCount > 2; + let bindingIndex = 0; const getShaderSource = (shaderHelper: ShaderHelper) => ` const normSize: u32 = ${normSize}; const normSizeTyped: ${dataType} = ${normSize}; const epsilon: f32 = ${attributes.epsilon}; - @group(0) @binding(0) var x : array<${dataType}>; - @group(0) @binding(1) var scale : array<${dataType}>; - ${bias ? `@group(0) @binding(2) var bias : array<${dataType}>;` : ''} - @group(0) @binding(3) var output : array<${dataType}>; - ${hasMeanDataOutput ? `@group(0) @binding(4) var meanDataOutput : array<${dataType}>` : ''}; - ${hasInvStdOutput ? `@group(0) @binding(5) var invStdOutput : array<${dataType}>` : ''}; + @group(0) @binding(${bindingIndex++}) var x : array<${dataType}>; + @group(0) @binding(${bindingIndex++}) var scale : array<${dataType}>; + ${bias ? `@group(0) @binding(${bindingIndex++}) var bias : array<${dataType}>;` : ''} + @group(0) @binding(${bindingIndex++}) var output : array<${dataType}>; + ${ + hasMeanDataOutput ? + `@group(0) @binding(${bindingIndex++}) var meanDataOutput : array<${dataType}>` : + ''}; + ${ + hasInvStdOutput ? + `@group(0) @binding(${bindingIndex++}) var invStdOutput : array<${dataType}>` : + ''}; ${shaderHelper.mainStart()} let offset = global_idx * normSize; @@ -118,7 +125,8 @@ export const layerNorm = (context: ComputeContext, attributes: LayerNormAttribut const metadata = { name: 'LayerNormalization', - inputTypes: [GpuDataType.default, GpuDataType.default, GpuDataType.default], + inputTypes: context.inputs.length === 2 ? [GpuDataType.default, GpuDataType.default] : + [GpuDataType.default, GpuDataType.default, GpuDataType.default], cacheHint: attributes.cacheKey + context.outputCount.toString(10) + context.inputs.length.toString(10), }; diff --git a/js/web/test/data/ops/layer-norm.jsonc b/js/web/test/data/ops/layer-norm.jsonc new file mode 100644 index 0000000000..d27d7e7083 --- /dev/null +++ b/js/web/test/data/ops/layer-norm.jsonc @@ -0,0 +1,33 @@ +[ + { + "name": "Simple test without bias", + "operator": "LayerNormalization", + "cases": [ + { + "name": "Simple test without bias", + "inputs": [ + { + "data": [1, 2, 3, 4, 5, 6, 7, 8], + "dims": [1, 2, 4], + "type": "float32" + }, + { + "data": [1, 2, 3, 4], + "dims": [4], + "type": "float32" + } + ], + "outputs": [ + { + "data": [ + -1.3416354656219482, -0.8944236636161804, 1.3416354656219482, 5.366541862487793, -1.3416354656219482, + -0.8944236636161804, 1.3416354656219482, 5.366541862487793 + ], + "dims": [1, 2, 4], + "type": "float32" + } + ] + } + ] + } +] diff --git a/js/web/test/suite-test-list.jsonc b/js/web/test/suite-test-list.jsonc index a28cd7a615..72300ebb3e 100644 --- a/js/web/test/suite-test-list.jsonc +++ b/js/web/test/suite-test-list.jsonc @@ -1345,6 +1345,7 @@ //"neg.jsonc", //"not.jsonc", //"or.jsonc", + "layer-norm.jsonc", "leaky-relu.jsonc", "reduce-min.jsonc", "relu.jsonc",