[webgpu] Fix test_layer_normalization_2d_axis0 (#24223)

The optional 'Mean' and 'InvStdDev' outputs of the LayerNormalization
were not implemented.

---------

Co-authored-by: Yulong Wang <7679871+fs-eire@users.noreply.github.com>

Author: jchen10

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

@@ -24,7 +24,13 @@ Status LayerNormProgram::GenerateShaderCode(ShaderHelper& shader) const {
   if (has_bias_) {
     shader.AddInput("bias", ShaderUsage::UseUniform);
   }
-  shader.AddOutput("output", ShaderUsage::UseUniform);
+  shader.AddOutput("y", ShaderUsage::UseUniform);
+  if (has_mean_output_) {
+    shader.AddOutput("mean_output", ShaderUsage::None);
+  }
+  if (has_inv_std_dev_output_) {
+    shader.AddOutput("inv_std_dev_output", ShaderUsage::None);
+  }
 
   int components = x.NumComponents();
   std::string bias = (has_bias_) ? " + bias[j]" : "";
@@ -48,8 +54,14 @@ Status LayerNormProgram::GenerateShaderCode(ShaderHelper& shader) const {
                             << "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"
-                            << "   output[j + offset] =  x_value_t((f32input" << simpl2 << ") * inv_std_dev * f32scale)" << bias << ";\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";
+  }
 
   return Status::OK();
 }
@@ -62,8 +74,6 @@ Status LayerNorm<simplified>::ComputeInternal(onnxruntime::webgpu::ComputeContex
 
   const auto x_shape = x->Shape();
 
-  auto* output = context.Output(0, x_shape);
-
   if (x_shape.Size() == 0) {
     return Status::OK();
   }
@@ -85,13 +95,27 @@ Status LayerNorm<simplified>::ComputeInternal(onnxruntime::webgpu::ComputeContex
                            scale_size, " and bias size of ", bias_size);
   }
 
-  LayerNormProgram program{bias != nullptr, is_fp16, simplified};
+  TensorShapeVector mean_dim;
+  for (size_t i = 0; i < x_shape.NumDimensions(); ++i) {
+    if (i < axis) {
+      mean_dim.push_back(x_shape[i]);
+    } else {
+      mean_dim.push_back(1);
+    }
+  }
+  TensorShape mean_shape(mean_dim);
+
+  auto* y = context.Output(0, x_shape);
+  auto* mean = context.Output(1, mean_shape);
+  auto* inv_std_dev = context.Output(2, mean_shape);
+
+  LayerNormProgram program{bias != nullptr, is_fp16, simplified, mean != nullptr, inv_std_dev != nullptr};
 
   program
       .CacheHint(simplified)
       .AddInputs({{x, ProgramTensorMetadataDependency::Type, components}})
       .AddInputs({{scale, ProgramTensorMetadataDependency::Type, components}})
-      .AddOutputs({{output, ProgramTensorMetadataDependency::None, components}})
+      .AddOutputs({{y, ProgramTensorMetadataDependency::None, components}})
       .SetDispatchGroupSize((norm_count + WORKGROUP_SIZE - 1) / WORKGROUP_SIZE)
       .AddUniformVariables({
           {static_cast<uint32_t>(norm_count)},
@@ -109,25 +133,26 @@ Status LayerNorm<simplified>::ComputeInternal(onnxruntime::webgpu::ComputeContex
   if (bias != nullptr) {
     program.AddInput({bias, ProgramTensorMetadataDependency::Type, components});
   }
+
+  if (mean != nullptr) {
+    program.AddOutputs({{mean, ProgramTensorMetadataDependency::None}});
+  }
+  if (inv_std_dev != nullptr) {
+    program.AddOutputs({{inv_std_dev, ProgramTensorMetadataDependency::None}});
+  }
+
   return context.RunProgram(program);
 }
 
-ONNX_OPERATOR_KERNEL_EX(
-    LayerNormalization,
-    kOnnxDomain,
-    17,
-    kWebGpuExecutionProvider,
-    (*KernelDefBuilder::Create())
-        .TypeConstraint("T", WebGpuSupportedFloatTypes()),
-    LayerNorm<false>);
-
-ONNX_OPERATOR_KERNEL_EX(
-    SimplifiedLayerNormalization,
-    kOnnxDomain,
-    1,
-    kWebGpuExecutionProvider,
-    (*KernelDefBuilder::Create()).TypeConstraint("T", WebGpuSupportedFloatTypes()),
-    LayerNorm<true>);
+ONNX_OPERATOR_KERNEL_EX(LayerNormalization, kOnnxDomain, 17, kWebGpuExecutionProvider,
+                        (*KernelDefBuilder::Create()).TypeConstraint("T", WebGpuSupportedFloatTypes()),
+                        LayerNorm<false>);
+
+ONNX_OPERATOR_KERNEL_EX(SimplifiedLayerNormalization, kOnnxDomain, 1, kWebGpuExecutionProvider,
+                        (*KernelDefBuilder::Create())
+                            .TypeConstraint("T", WebGpuSupportedFloatTypes())
+                            .TypeConstraint("U", WebGpuSupportedFloatTypes()),
+                        LayerNorm<true>);
 
 }  // namespace webgpu
 }  // namespace onnxruntime

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

@@ -11,25 +11,28 @@ namespace webgpu {
 
 class LayerNormProgram final : public Program<LayerNormProgram> {
  public:
-  LayerNormProgram(bool has_bias,
-                   bool is_fp16,
-                   bool simplified) : Program{"LayerNorm"},
-                                      has_bias_{has_bias},
-                                      is_fp16_{is_fp16},
-                                      simplified_{simplified} {}
+  LayerNormProgram(bool has_bias, bool is_fp16, bool simplified, bool has_mean_output,
+                   bool has_inv_std_dev_output)
+      : 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} {}
 
   Status GenerateShaderCode(ShaderHelper& sh) const override;
 
-  WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES(
-      {"norm_count", ProgramUniformVariableDataType::Uint32},
-      {"norm_size", ProgramUniformVariableDataType::Uint32},
-      {"norm_size_vectorized", ProgramUniformVariableDataType::Uint32},
-      {"epsilon", ProgramUniformVariableDataType::Float32});
+  WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES({"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_;
 };
 
 template <bool simplified>