Improve SimplifiedLayerNorm by using same techniques as SkipSimplifiedLayerNorm (microsoft#25850)

### Description
Use similar shaders as SkipSimplifiedLayerNorm in SimplifiedLayerNorm,
to fix the performance issues with SimplifiedLayerNorm.

### Motivation and Context
Prior to this change, generation in Bitnet was bottlenecked on
SimplifiedLayerNorm
<img width="332" height="378" alt="image"
src="https://github.com/user-attachments/assets/3bc16ac1-ef7d-46bf-b403-92fc9192a2df"
/>
with this change performance has now improved to match
SkipSimplifiedLayerNorm
<img width="699" height="179" alt="image"
src="https://github.com/user-attachments/assets/30009d85-d5d9-4585-987a-b39ecf52e0b5"
/>

Author: sushraja-msft

onnxruntime/core/providers/webgpu/nn/layer_norm.cc

@@ -26,7 +26,7 @@ static TensorShape GetOverrideShape(const TensorShape& shape, int components) {
 }
 
 Status LayerNormProgram::GenerateShaderCode(ShaderHelper& shader) const {
-  const auto& x = shader.AddInput("x", ShaderUsage::UseUniform | ShaderUsage::UseValueTypeAlias);
+  const auto& x = shader.AddInput("x", ShaderUsage::UseUniform | ShaderUsage::UseValueTypeAlias | ShaderUsage::UseElementTypeAlias);
   shader.AddInput("scale", ShaderUsage::UseUniform);
   if (has_bias_) {
     shader.AddInput("bias", ShaderUsage::UseUniform);
@@ -39,35 +39,113 @@ Status LayerNormProgram::GenerateShaderCode(ShaderHelper& shader) const {
     shader.AddOutput("inv_std_dev_output", ShaderUsage::None);
   }
 
-  int components = x.NumComponents();
-  std::string bias = (has_bias_) ? " + bias[j]" : "";
-  std::string simpl1 = (simplified_) ? "" : " - mean * mean";
-  std::string simpl2 = (simplified_) ? "" : " - mean";
-
-  shader.AdditionalImplementation() << "alias element_t = " << (is_fp16_ ? "f16;\n" : "f32;\n")
-                                    << "alias f32_val_t = " << (components == 4 ? "vec4<f32>" : (components == 2 ? "vec2<f32>" : "f32")) << ";\n";
-
-  shader.MainFunctionBody() << shader.GuardAgainstOutOfBoundsWorkgroupSizes("uniforms.norm_count")
-                            << "let offset = global_idx * uniforms.norm_size_vectorized;\n"
-                            << "var mean_vector = f32_val_t(0);\n"
-                            << "var mean_square_vector = f32_val_t(0);\n"
-                            << "for (var h: u32 = 0u; h < uniforms.norm_size_vectorized; h++) {\n"
-                            << "   let value = f32_val_t(x[h + offset]);\n"
-                            << "   mean_vector += value;\n"
-                            << "   mean_square_vector += value * value;\n"
-                            << "}\n"
-                            << "let mean = " << SumVector("mean_vector", components) << " / f32(uniforms.norm_size);\n"
-                            << "let inv_std_dev = inverseSqrt(" << SumVector("mean_square_vector", components) << " / f32(uniforms.norm_size)" << simpl1 << " + uniforms.epsilon);\n"
-                            << "for (var j: u32 = 0; j < uniforms.norm_size_vectorized; j++) {\n"
-                            << "   let f32input = f32_val_t(x[j + offset]);\n"
-                            << "   let f32scale = f32_val_t(scale[j]);\n"
-                            << "   y[j + offset] =  x_value_t((f32input" << simpl2 << ") * inv_std_dev * f32scale)" << bias << ";\n"
-                            << "}\n";
-  if (has_mean_output_) {
-    shader.MainFunctionBody() << "mean_output[global_idx] = mean;\n";
-  }
-  if (has_inv_std_dev_output_) {
-    shader.MainFunctionBody() << "inv_std_dev_output[global_idx] = inv_std_dev;\n";
+  std::string simpl1 = (simplified_) ? "" : "- mean * mean ";
+  std::string simpl2 = (simplified_) ? "" : "- x_element_t(mean) ";
+
+  if (split_norm_dim_) {
+    shader.AdditionalImplementation()
+        << "var<workgroup> sum_shared : array<f32, workgroup_size_x>;\n"
+        << "var<workgroup> sum_squared_shared : array<f32, workgroup_size_x>;\n";
+
+    shader.MainFunctionBody()
+        << "  var sum_vec4 = vec4<f32>(0);\n"
+        << "  var sum_squared_vec4 = vec4<f32>(0);\n"
+        << "  var cur_input = x_value_t(0);\n"
+        << "  for (var i: u32 = 0; i < uniforms.norm_size / (workgroup_size_x * 4); i++) {\n"
+        << "    let input_offset = i * workgroup_size_x + local_idx;\n"
+        << "    let input_value = x[input_offset];\n"
+        << "    if (i == workgroup_idx) {\n"
+        << "      cur_input = input_value;\n"
+        << "    }\n"
+        << "    let f32_value = vec4<f32>(input_value);\n"
+        << "    sum_vec4 += f32_value;\n"
+        << "    sum_squared_vec4 += f32_value * f32_value;\n"
+        << "  }\n"
+        << "  var sum = " << SumVector("sum_vec4", 4) << ";\n"
+        << "  var sum_squared = " << SumVector("sum_squared_vec4", 4) << ";\n"
+        << "  sum_shared[local_idx] = sum;\n"
+        << "  sum_squared_shared[local_idx] = sum_squared;\n"
+        << "  workgroupBarrier();\n"
+        << "  var reduce_size : u32 = workgroup_size_x;\n"
+        << "  for (var curr_size = reduce_size >> 1;  curr_size > 0; curr_size = reduce_size >> 1) {\n"
+        << "    reduce_size = curr_size + (reduce_size & 1);\n"
+        << "    if (local_idx < curr_size) {\n"
+        << "      sum_shared[local_idx] += sum_shared[local_idx + reduce_size];\n"
+        << "      sum_squared_shared[local_idx] += sum_squared_shared[local_idx + reduce_size];\n"
+        << "    }\n"
+        << "    workgroupBarrier();\n"
+        << "  }\n"
+        << "  let mean = sum_shared[0] / f32(uniforms.norm_size);\n"
+        << "  let inv_std_dev = inverseSqrt(sum_squared_shared[0] / f32(uniforms.norm_size) " << simpl1 << "+ uniforms.epsilon);\n"
+        << "  let offset = workgroup_idx * workgroup_size_x + local_idx;\n"
+        << "  y[offset] = ((cur_input " << simpl2 << ") * x_element_t(inv_std_dev) * scale[offset]" << (has_bias_ ? " + bias[offset] " : "") << ");\n";
+
+    if (has_mean_output_) {
+      shader.MainFunctionBody() << "  if (local_idx == 0 && workgroup_idx == 0) {\n"
+                                << "    mean_output[global_idx / uniforms.norm_size] = mean;\n"
+                                << "  }\n";
+    }
+    if (has_inv_std_dev_output_) {
+      shader.MainFunctionBody() << "  if (local_idx == 0 && workgroup_idx == 0) {\n"
+                                << "    inv_std_dev_output[global_idx / uniforms.norm_size] = inv_std_dev;\n"
+                                << "  }\n";
+    }
+  } else {
+    int components = x.NumComponents();
+    std::string bias = (has_bias_) ? " + bias[offset1d + i] " : "";
+
+    shader.AdditionalImplementation()
+        << "alias f32_val_t = " << (components == 4 ? "vec4<f32>" : (components == 2 ? "vec2<f32>" : "f32")) << ";\n"
+        << "var<workgroup> sum_shared : array<f32_val_t, workgroup_size_x>;\n"
+        << "var<workgroup> sum_squared_shared : array<f32_val_t, workgroup_size_x>;\n";
+
+    shader.MainFunctionBody()
+        << "let ix = local_idx;\n"
+        << "let iy = global_idx / workgroup_size_x;\n"
+        << "let norm_size_vectorized: u32 = uniforms.norm_size / uniforms.components;\n"
+        << "var stride = norm_size_vectorized / workgroup_size_x;\n"
+        << "let offset = ix * stride + iy * norm_size_vectorized;\n"
+        << "let offset1d = stride * ix;\n"
+        << "sum_shared[ix] = f32_val_t(0);\n"
+        << "sum_squared_shared[ix] = f32_val_t(0);\n"
+        << "if (ix == workgroup_size_x - 1) {\n"
+        << " stride = norm_size_vectorized - stride * ix;\n"
+        << "}\n"
+        << "for (var i: u32 = 0; i < stride; i++) {\n"
+        << " let input_value = x[offset + i];\n"
+        << " y[offset + i] = input_value;\n"
+        << " let f32_value = f32_val_t(input_value);\n"
+        << " sum_shared[ix] += f32_value;\n"
+        << " sum_squared_shared[ix] += f32_value * f32_value;\n"
+        << "}\n"
+        << "workgroupBarrier();\n"
+        << "var reduce_size : u32 = workgroup_size_x;\n"
+        << "for (var curr_size = reduce_size >> 1;  curr_size > 0; curr_size = reduce_size >> 1) {\n"
+        << " reduce_size = curr_size + (reduce_size & 1);\n"
+        << " if (ix < curr_size) {\n"
+        << "  sum_shared[ix] += sum_shared[ix + reduce_size];\n"
+        << "  sum_squared_shared[ix] += sum_squared_shared[ix + reduce_size];\n"
+        << " }\n"
+        << " workgroupBarrier();\n"
+        << "}\n"
+        << "let sum = sum_shared[0];\n"
+        << "let square_sum = sum_squared_shared[0];\n"
+        << "let mean = " << SumVector("sum", components) << " / f32(uniforms.norm_size);\n"
+        << "let inv_std_dev = inverseSqrt(" << SumVector("square_sum", components) << " / f32(uniforms.norm_size) " << simpl1 << "+ uniforms.epsilon);\n"
+        << "for (var i: u32 = 0; i < stride; i++) {\n"
+        << " y[offset + i] = (y[offset + i] " << simpl2 << ") * x_element_t(inv_std_dev) * scale[offset1d + i]" << bias << ";\n"
+        << "};\n";
+
+    if (has_mean_output_) {
+      shader.MainFunctionBody() << "if (ix == 0) {\n"
+                                << "  mean_output[iy] = mean;\n"
+                                << "}\n";
+    }
+    if (has_inv_std_dev_output_) {
+      shader.MainFunctionBody() << "if (ix == 0) {\n"
+                                << "  inv_std_dev_output[iy] = inv_std_dev;\n"
+                                << "}\n";
+    }
   }
 
   return Status::OK();
@@ -81,8 +159,6 @@ Status LayerNorm<simplified>::ComputeInternal(onnxruntime::webgpu::ComputeContex
 
   const auto x_shape = x->Shape();
 
-  const bool is_fp16 = x->GetElementType() == ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16;
-
   const size_t axis = NormalizeAxis(axis_, x_shape.NumDimensions());
   const uint32_t norm_count = onnxruntime::narrow<uint32_t>(x_shape.SizeToDimension(axis));
   const int64_t norm_size = x_shape.SizeFromDimension(axis);
@@ -116,14 +192,19 @@ Status LayerNorm<simplified>::ComputeInternal(onnxruntime::webgpu::ComputeContex
     return Status::OK();
   }
 
-  LayerNormProgram program{bias != nullptr, is_fp16, simplified, mean != nullptr, inv_std_dev != nullptr};
+  // Check if we should use split norm dimension optimization
+  const bool split_norm_dim = norm_size % 512 == 0 && norm_count == 1;
+
+  LayerNormProgram program{bias != nullptr, simplified, mean != nullptr, inv_std_dev != nullptr, split_norm_dim};
 
-  program.CacheHint(components, simplified)
+  program.CacheHint(components, simplified, split_norm_dim)
       .AddInputs({{x, ProgramTensorMetadataDependency::Type, GetOverrideShape(x->Shape(), components), components}})
       .AddInputs(
           {{scale, ProgramTensorMetadataDependency::Type, GetOverrideShape(scale->Shape(), components), components}})
       .AddOutputs({{y, ProgramTensorMetadataDependency::None, GetOverrideShape(y->Shape(), components), components}})
-      .SetDispatchGroupSize((norm_count + WORKGROUP_SIZE - 1) / WORKGROUP_SIZE)
+      .AddUniformVariables({
+          {static_cast<uint32_t>(components)},
+      })
       .AddUniformVariables({
           {static_cast<uint32_t>(norm_count)},
       })
@@ -137,6 +218,15 @@ Status LayerNorm<simplified>::ComputeInternal(onnxruntime::webgpu::ComputeContex
           {static_cast<float>(epsilon_)},
       });
 
+  if (split_norm_dim) {
+    const uint32_t workgroup_size_x = 128;
+    const uint32_t dispatch_size_x = onnxruntime::narrow<uint32_t>(norm_size / (workgroup_size_x * components));
+    program.SetDispatchGroupSize(dispatch_size_x, 1, 1)
+        .SetWorkgroupSize(workgroup_size_x);
+  } else {
+    program.SetDispatchGroupSize(norm_count);
+  }
+
   if (bias != nullptr) {
     program.AddInput(
         {bias, ProgramTensorMetadataDependency::Type, GetOverrideShape(bias->Shape(), components), components});

onnxruntime/core/providers/webgpu/nn/layer_norm.h

@@ -11,28 +11,29 @@ namespace webgpu {
 
 class LayerNormProgram final : public Program<LayerNormProgram> {
  public:
-  LayerNormProgram(bool has_bias, bool is_fp16, bool simplified, bool has_mean_output,
-                   bool has_inv_std_dev_output)
+  LayerNormProgram(bool has_bias, bool simplified, bool has_mean_output,
+                   bool has_inv_std_dev_output, bool split_norm_dim = false)
       : Program{"LayerNorm"},
         has_bias_{has_bias},
-        is_fp16_{is_fp16},
         simplified_{simplified},
         has_mean_output_{has_mean_output},
-        has_inv_std_dev_output_{has_inv_std_dev_output} {}
+        has_inv_std_dev_output_{has_inv_std_dev_output},
+        split_norm_dim_{split_norm_dim} {}
 
   Status GenerateShaderCode(ShaderHelper& sh) const override;
 
-  WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES({"norm_count", ProgramUniformVariableDataType::Uint32},
+  WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES({"components", ProgramUniformVariableDataType::Uint32},
+                                          {"norm_count", ProgramUniformVariableDataType::Uint32},
                                           {"norm_size", ProgramUniformVariableDataType::Uint32},
                                           {"norm_size_vectorized", ProgramUniformVariableDataType::Uint32},
                                           {"epsilon", ProgramUniformVariableDataType::Float32});
 
  private:
   bool has_bias_;
-  bool is_fp16_;
   bool simplified_;
   bool has_mean_output_;
   bool has_inv_std_dev_output_;
+  bool split_norm_dim_;
 };
 
 template <bool simplified>