optimize skiplayernorm (microsoft#20551)

SkipSimplifiedLayerNormalization used in phi3 comes down from 222usec to
14usec

Author: guschmue

js/web/lib/wasm/jsep/webgpu/ops/skip-layer-norm.ts

@@ -6,7 +6,7 @@ import {TensorView} from '../../tensor-view';
 import {ShapeUtil} from '../../util';
 import {ComputeContext, ProgramInfo, ProgramUniform} from '../types';
 
-import {castToF32, fillVector, getMaxComponents, inputVariable, outputVariable, ShaderHelper, sumVector, tensorTypeToWsglStorageType, UniformsArrayType} from './common';
+import {castToF32, getMaxComponents, inputVariable, outputVariable, ShaderHelper, sumVector, tensorTypeToWsglStorageType, UniformsArrayType} from './common';
 
 export interface SkipLayerNormAttributes {
   simplified: boolean;
@@ -58,7 +58,6 @@ const validateInputs = (inputs: readonly TensorView[]): void => {
       throw new Error('Beta must have the same hidden size as input');
     }
   }
-
   if (inputs.length > 4) {
     const bias: TensorView = inputs[4];
     if (bias.dims.length !== 1) {
@@ -86,6 +85,7 @@ const createSkipLayerNormProgramInfo =
           const hasMeanOutput = isTraining && outputCount > 1;
           const hasInvStdDevOutput = isTraining && outputCount > 2;
           const hasInputSkipBiasSumOutput = outputCount > 3;
+          const workgroupSize = 64;
 
           const components = getMaxComponents(hiddenSize);
 
@@ -124,35 +124,61 @@ const createSkipLayerNormProgramInfo =
               variables.push(outputVariable('input_skip_bias_sum', inputs[0].dataType, outputShape, components));
             }
             const dataType = tensorTypeToWsglStorageType(inputs[0].dataType);
+            const vecDataType = tensorTypeToWsglStorageType(DataType.float, components);
             return `
 
       ${shaderHelper.registerUniforms(uniformsArray).declareVariables(...variables)}
+      var<workgroup> sum_shared : array<${vecDataType}, ${workgroupSize}>;
+      var<workgroup> sum_squared_shared : array<${vecDataType}, ${workgroupSize}>;
+
+      ${shaderHelper.mainStart([
+              workgroupSize, 1, 1
+            ])}
+        let ix = local_id.x;
+        let iy = global_id.x / ${workgroupSize};
 
-      ${shaderHelper.mainStart()}
-        ${shaderHelper.guardAgainstOutOfBoundsWorkgroupSizes('uniforms.output_size / uniforms.hidden_size')}
         let hidden_size_vectorized: u32 = uniforms.hidden_size / uniforms.components;
-        let offset = global_idx * hidden_size_vectorized;
-        var sum = ${fillVector('f32', components)};
-        var squareSum = ${fillVector('f32', components)};
-        for (var i: u32 = 0; i < hidden_size_vectorized; i++) {
+        var stride = hidden_size_vectorized / ${workgroupSize};
+        let offset = ix * stride + iy * hidden_size_vectorized;
+        let offset1d = stride * ix;
+        if (ix == ${workgroupSize - 1}) {
+          stride = hidden_size_vectorized - stride * ix;
+        }
+        for (var i: u32 = 0; i < stride; i++) {
           let skip_value = skip[offset + i];
-          let bias_value = ${hasBiasInput ? 'bias[i]' : dataType + '(0.0)'};
+          let bias_value = ${hasBiasInput ? 'bias[offset1d + i]' : dataType + '(0.0)'};
           let input_value = x[offset + i];
           let value = input_value + skip_value + bias_value;
           ${hasInputSkipBiasSumOutput ? 'input_skip_bias_sum[offset + i] = value;' : ''}
           output[offset + i] = value;
           let f32_value = ${castToF32(dataType, components, 'value')};
-          sum += f32_value;
-          squareSum += f32_value * f32_value;
+          sum_shared[ix] += f32_value;
+          sum_squared_shared[ix] += f32_value * f32_value;
         }
+        workgroupBarrier();
+
+        var reduce_size : u32 = ${workgroupSize};
+        for (var curr_size = reduce_size >> 1;  curr_size > 0; curr_size = reduce_size >> 1) {
+          reduce_size = curr_size + (reduce_size & 1);
+          if (ix < curr_size) {
+            sum_shared[ix] += sum_shared[ix + reduce_size];
+            sum_squared_shared[ix] += sum_squared_shared[ix + reduce_size];
+          }
+          workgroupBarrier();
+        }
+
+        let sum = sum_shared[0];
+        let square_sum = sum_squared_shared[0];
         let mean = ${sumVector('sum', components)} / f32(uniforms.hidden_size);
-        let inv_std_dev = inverseSqrt(${sumVector('squareSum', components)} / f32(uniforms.hidden_size) ${
+        let inv_std_dev = inverseSqrt(${sumVector('square_sum', components)} / f32(uniforms.hidden_size) ${
                 simplified ? '' : '- mean * mean'} + uniforms.epsilon);
         ${hasMeanOutput ? 'mean_output[global_idx] = mean;' : ''}
         ${hasInvStdDevOutput ? 'inv_std_output[global_idx] = inv_std_dev;' : ''}
-        for (var i: u32 = 0; i < hidden_size_vectorized; i++) {
-          output[offset + i] = (output[offset + i] ${simplified ? '' : `- ${dataType}(mean)`}) * ${
-                dataType}(inv_std_dev) * gamma[i] ${hasBetaInput ? '+ beta[i]' : ''};
+
+        for (var i: u32 = 0; i < stride; i++) {
+          output[offset + i] = (output[offset + i] ${simplified ? '' : `- ${dataType}(mean)`}) *
+            ${dataType}(inv_std_dev) * gamma[offset1d + i]
+            ${hasBetaInput ? '+ beta[offset1d + i]' : ''};
         }
       }`;
           };
@@ -173,7 +199,13 @@ const createSkipLayerNormProgramInfo =
               inputDependencies: inputs.map((_input, _index) => 'type')
             },
             getShaderSource,
-            getRunData: () => ({outputs, dispatchGroup: {x: Math.ceil(outputSize / hiddenSize / 64)}, programUniforms}),
+            getRunData: () => ({
+              outputs,
+              dispatchGroup: {
+                x: Math.ceil(outputSize / hiddenSize),
+              },
+              programUniforms
+            }),
           };
         };