Enable LayerNorm and SkipLayerNorm in OneDNN EP (#11128)

onnxruntime/core/providers/dnnl/dnnl_node_capability.cc

@@ -983,4 +983,64 @@ bool DnnlDequantizeLinearNodeCapability::Supported(const Node* node, const Graph
   return true;
 }
 
+bool DnnlLayerNormalizationNodeCapability::Supported(const Node* node, const GraphViewer& graph_viewer) const {
+  ORT_UNUSED_PARAMETER(graph_viewer);
+  if (!IsTypeSupported(node)) return false;
+  if (!IsTrainingSupported(node)) return false;
+  if (!IsDimensionSupported(node)) return false;
+  return true;
+}
+
+bool DnnlLayerNormalizationNodeCapability::IsAxisSupported(const Node* node) const {
+  // At the moment of implementation OneDNN does not support broadcasting 
+  // on LayerNorm so we can only accept the op when the normalization 
+  // is done on the last dim
+  const NodeAttributes& attr = node->GetAttributes();
+  auto axis = attr.find("axis");
+
+  if (axis != attr.end()) {
+    // Get the norm dim
+    auto norm_dim = axis->second().i();
+    // Get the input dims
+    auto input_dims = node->InputDefs()[0]->Shape()->dim_size();
+
+    // If the axis is not the last dim, it's not supported
+    if ((norm_dim != -1) && (norm_dim != (input_dims - 1))) {
+      return false;
+    }
+  }
+  // If no axis is provided or the conditions above were not met, 
+  // we normalize on the last dim
+  return true;
+}
+
+bool DnnlLayerNormalizationNodeCapability::IsTrainingSupported(const Node* node) const {
+  // Get the output defs
+  auto num_outputs = node->OutputDefs().size();
+  // Training support is in progress, for the moment we don't support multiple outputs
+  if (num_outputs > 1) {
+    return false;
+  } else {
+    return true;
+  }
+}
+bool DnnlLayerNormalizationNodeCapability::IsDimensionSupported(const Node* node) const {
+  // Get the input shape
+  auto input_shape = node->InputDefs()[0]->Shape();
+
+  if (input_shape != nullptr) {
+    auto input_dims = input_shape->dim_size();
+    // OneDNN is supports from 2D to 5D tensors
+    if ((input_dims <= 5) && (input_dims >= 2)) {
+      return true;
+
+    // If the tensor is 1D or >6D then we dont support it
+    } else {
+      return false;
+    }   
+  }
+  // If we dont have shape info, accept the data and catch the error on the implementation
+  return true;  
+}
+
 }  // namespace onnxruntime

onnxruntime/core/providers/dnnl/dnnl_node_capability.h

@@ -416,5 +416,30 @@ class DnnlDequantizeLinearNodeCapability : public DnnlDefaultOptionalMultiInputN
   bool Supported(const Node* node, const GraphViewer& graph_viewer) const override;
 };
 
+class DnnlLayerNormalizationNodeCapability : public DnnlDefaultNodeCapability {
+ public:
+ // Default constructor
+  DnnlLayerNormalizationNodeCapability() : DnnlDefaultNodeCapability({type_float16,
+                                                                      type_bfloat16,
+                                                                      type_float32,
+                                                                      type_double}) {}
+  // Constructor for other LayerNorm flavors
+  DnnlLayerNormalizationNodeCapability(std::vector<ORT_DataType> supported_dt)
+    : DnnlDefaultNodeCapability(supported_dt){}
+
+  bool Supported(const Node* node, const GraphViewer& graph_viewer) const override;
+
+ private:
+  bool IsAxisSupported(const Node* node) const;
+  bool IsTrainingSupported(const Node* node) const;
+  bool IsDimensionSupported(const Node* node) const;
+};
+
+class DnnlSkipLayerNormalizationNodeCapability : public DnnlLayerNormalizationNodeCapability {
+ public:
+  DnnlSkipLayerNormalizationNodeCapability() : DnnlLayerNormalizationNodeCapability({type_float32,
+                                                                                    type_float16}) {}
+};
+
                                  
 }  // namespace onnxruntime

onnxruntime/core/providers/dnnl/dnnl_op_manager.cc

@@ -25,6 +25,7 @@ DnnlOpManager::DnnlOpManager() {
   dnnl_ops_map_.emplace(std::make_pair("Gemm", std::unique_ptr<DnnlNodeCapability>(new DnnlGemmNodeCapability())));
   dnnl_ops_map_.emplace(std::make_pair("GlobalAveragePool", std::unique_ptr<DnnlNodeCapability>(new DnnlPoolNodeCapability())));
   dnnl_ops_map_.emplace(std::make_pair("GlobalMaxPool", std::unique_ptr<DnnlNodeCapability>(new DnnlPoolNodeCapability())));
+  dnnl_ops_map_.emplace(std::make_pair("LayerNormalization", std::unique_ptr<DnnlNodeCapability>(new DnnlLayerNormalizationNodeCapability())));
   dnnl_ops_map_.emplace(std::make_pair("LeakyRelu", std::unique_ptr<DnnlNodeCapability>(new DnnlElementwiseCapability())));
   dnnl_ops_map_.emplace(std::make_pair("Log", std::unique_ptr<DnnlNodeCapability>(new DnnlElementwiseCapability())));
   dnnl_ops_map_.emplace(std::make_pair("LRN", std::unique_ptr<DnnlNodeCapability>(new DnnlDefaultNodeCapability())));
@@ -48,6 +49,7 @@ DnnlOpManager::DnnlOpManager() {
   dnnl_ops_map_.emplace(std::make_pair("Reshape", std::unique_ptr<DnnlNodeCapability>(new DnnlReshapeNodeCapability())));
   dnnl_ops_map_.emplace(std::make_pair("Round", std::unique_ptr<DnnlNodeCapability>(new DnnlElementwiseCapability())));
   dnnl_ops_map_.emplace(std::make_pair("Sigmoid", std::unique_ptr<DnnlNodeCapability>(new DnnlElementwiseCapability())));
+  dnnl_ops_map_.emplace(std::make_pair("SkipLayerNormalization", std::unique_ptr<DnnlNodeCapability>(new DnnlSkipLayerNormalizationNodeCapability())));
   dnnl_ops_map_.emplace(std::make_pair("Softmax", std::unique_ptr<DnnlNodeCapability>(new DnnlSoftmaxNodeCapability())));
   dnnl_ops_map_.emplace(std::make_pair("Softplus", std::unique_ptr<DnnlNodeCapability>(new DnnlElementwiseCapability())));
   dnnl_ops_map_.emplace(std::make_pair("Squeeze", std::unique_ptr<DnnlNodeCapability>(new DnnlSqueezeNodeCapability())));

onnxruntime/core/providers/dnnl/subgraph/dnnl_layernorm.cc

@@ -0,0 +1,314 @@
+// Copyright(C) 2022 Intel Corporation
+// Licensed under the MIT License
+
+#include "dnnl_layernorm.h"
+#include "dnnl_subgraph.h"
+#include "dnnl_subgraph_primitive.h"
+
+namespace onnxruntime {
+namespace ort_dnnl {
+
+DnnlLayerNorm::DnnlLayerNorm() {}
+
+/*
+Layer Normalization and Skip Layer Normalization implementation:
+Layer Norm:
+  Inputs:
+    0) X - Input Tensor
+    1) G - Gamma Tensor (Scaling factor in the LN formula)
+    2) B - Bias Tensor (Shift value in the LN formula. Optional)
+  Outputs:
+    0) Y - Output Tensor
+    1) M - Mean Tensor (Optional)
+    2) I - Inverse std Tensor (Optional)
+
+               +-----------+
+(X) ---------->+           +----------> (Y)
+               |           |                  
+(G) ---------->+ LayerNorm +----------> (M)
+               |           |     
+(B) ---------->+           +----------> (I) 
+               +-----------+                     
+                                                                 
+    
+
+Skip Layer Norm:
+  Inputs:
+    0) X - Input Tensor
+    1) S - Skip Tensor
+    2) G - Gamma Tensor (Scaling factor in the LN formula)
+    4) E - Beta Tensor (Shift value in the LN formula. Optional)
+    4) B - Bias Tensor (Bias when adding X + S. Optional)
+  Outputs:
+    0) Y - Output Tensor
+    1) M - Mean Tensor (Optional)
+    2) I - Inverse std Tensor (Optional)
+
+               +-----------+
+(X) ---------->+           |                   +-----------+
+               |           |   (X + S + B)     |           |         
+(S) ---------->+ BuildSLN  +------------------>+           +----------> (Y)
+               |           |                   |           |
+(B) ---------->+           |      (G) -------->+ LayerNorm +----------> (M)
+               +-----------+                   |           |
+                                  (E) -------->+           +----------> (I)         
+                                               |           |
+                                               +-----------+
+                                 
+Attributes (epsilon)
+*/
+void DnnlLayerNorm::CreatePrimitive(DnnlSubgraphPrimitive& sp, DnnlNode& node) {
+
+  // Get engine
+  auto dnnl_engine = sp.GetEngine();
+
+  // Make sure every input's dimension follows the spec
+  ValidateDims(sp, node);
+
+  // Optional input flag
+  bool shift_exists = false;
+
+  // Input positions
+  int shift_pos, scale_pos;
+
+  // Get src mem
+  auto src_mem = sp.GetMemory(node.Input(IN_INPUT));
+  auto src_md = src_mem.get_desc();
+
+  // This contains the layer norm op and its parameters
+  ln_components op_comps;
+  if (node.OpType() == "SkipLayerNormalization") {
+    
+    // Check if shift is available
+    shift_exists = node.Input(IN_BETA).Exists();
+
+    // Fix positions for arguments
+    shift_pos = IN_BETA;
+    scale_pos = IN_SLN_GAMMA;
+    
+    // Build SLN and get modified mem
+    src_mem = BuildSLN(sp, node, dnnl_engine);
+
+  } else if (node.OpType() == "LayerNormalization") {
+
+    // Check if shift is available
+    shift_exists = node.Input(IN_LN_BIAS).Exists();
+
+    // Fix positions for arguments
+    shift_pos = IN_LN_BIAS;
+    scale_pos = IN_LN_GAMMA;
+
+    // Move the src to GPU if needed
+    src_mem = sp.GetMemoryAndReshape(node.Input(IN_INPUT), src_mem.get_desc(), dnnl_engine);
+
+  } else {
+    ORT_THROW("Unknown LayerNormalization flavor");
+  }
+
+  // X = LayerNornm(X)
+  // Check if we are training and need the extra outputs for backprop
+  dnnl::prop_kind prop_kind;
+#if defined(ENABLE_TRAINING)
+  prop_kind = dnnl::prop_kind::forward_training;
+#else
+  prop_kind = dnnl::prop_kind::forward_inference;
+#endif  // ENABLE_TRAINING
+
+  // If beta is available use shift, else only scale
+  dnnl::normalization_flags op_flags = dnnl::normalization_flags::use_scale;
+  if (shift_exists) {
+    op_flags |= dnnl::normalization_flags::use_shift;
+  }
+
+  // Get epsilon to avoid zero division
+  float epsilon = GetEpsilon(node);
+  // Operation desciptor
+  auto lnorm_desc = dnnl::layer_normalization_forward::desc(prop_kind, src_md, epsilon, op_flags);
+  // Primitive desciptor
+  auto lnorm_pd = dnnl::layer_normalization_forward::primitive_desc(lnorm_desc, dnnl_engine);
+  // Primitive
+  auto lnorm_prim = dnnl::layer_normalization_forward(lnorm_pd);
+
+  // Get gamma
+  auto gamma_mem = sp.GetMemory(node.Input(scale_pos));
+  gamma_mem = sp.GetMemoryAndReshape(node.Input(scale_pos), gamma_mem.get_desc(), dnnl_engine);
+
+  // Define primitive arguments
+  std::unordered_map<int, dnnl::memory> lnorm_args = {{DNNL_ARG_SRC, src_mem},
+                                                      {DNNL_ARG_SCALE, gamma_mem},
+                                                      {DNNL_ARG_DST, src_mem}};
+
+  // Get Beta and add shift if available
+  if (shift_exists) {
+    auto beta_mem = sp.GetMemory(node.Input(shift_pos));
+    beta_mem = sp.GetMemoryAndReshape(node.Input(shift_pos), beta_mem.get_desc(), dnnl_engine);
+    lnorm_args.insert({DNNL_ARG_SHIFT, beta_mem});
+  }
+
+// Check outputs used for training
+#if defined(ENABLE_TRAINING)
+  // If Mean exists
+  if (node.Output(OUT_MEAN).Exists()) {
+    auto mean_mem = dnnl::memory(lnorm_pd.mean_desc(), dnnl_engine);
+    lnorm_args.insert({DNNL_ARG_MEAN, mean_mem});
+    sp.SetMemory(node.Output(OUT_MEAN), mean_mem);
+  }
+  // If Variance exists
+  if (node.Output(OUT_INV_STD_VAR).Exists()) {
+    auto variance_mem = dnnl::memory(lnorm_pd.variance_desc(), dnnl_engine);
+    lnorm_args.insert({DNNL_ARG_VARIANCE, variance_mem});
+    sp.SetMemory(node.Output(OUT_INV_STD_VAR), variance_mem);
+  }
+#endif  // ENABLE_TRAINING
+
+  sp.AddPrimitive(lnorm_prim, lnorm_args);
+
+  sp.SetMemory(node.Output(OUT_OUTPUT), src_mem, true);
+}
+
+dnnl::memory DnnlLayerNorm::BuildSLN(DnnlSubgraphPrimitive& sp, DnnlNode& node, dnnl::engine dnnl_engine) {
+  // X += SKIP
+  // Get input and skip info
+  auto input_md = sp.GetMemory(node.Input(IN_INPUT)).get_desc();
+  auto skip_md = sp.GetMemory(node.Input(IN_SKIP)).get_desc();
+  auto skip_dims = skip_md.dims();
+
+  // Create primitive input map
+  std::unordered_map<int, dnnl::memory> skip_bias_args;
+
+  // Create md for the add op, according to the spec the output type should be
+  // the same as the input, so we can support inplace ops
+  auto add_skip_dst_md = dnnl::memory::desc(skip_dims, node.Output(OUT_OUTPUT).Type(), dnnl::memory::format_tag::any);
+  // Create desc for the op and primitive
+  auto add_skip_d = dnnl::binary::desc(dnnl::algorithm::binary_add, input_md, skip_md, add_skip_dst_md);
+
+  // Create primitive descriptor container
+  dnnl::binary::primitive_desc add_skip_pd;
+  // Add post op bias
+  if (node.Input(IN_SLN_BIAS).Exists()) {
+    // X += BIAS
+    // Get bias md
+    auto bias_md = sp.GetMemory(node.Input(IN_SLN_BIAS)).get_desc();
+    auto bias_dims = bias_md.dims();
+    // To follow the spec means our bias will always have less dimensions that our input}
+    // so we add the extra dimensions, reshape it and let OneDNN broadcast the value
+    while (bias_dims.size() < skip_dims.size()) {
+      bias_dims.insert(bias_dims.begin(), 1);
+    }
+    bias_md = bias_md.reshape(bias_dims);
+
+    dnnl::post_ops bias_add;
+    dnnl::primitive_attr binary_attr;
+    bias_add.append_binary(dnnl::algorithm::binary_add, bias_md);
+    binary_attr.set_post_ops(bias_add);
+    // Add post op to scale result
+    add_skip_pd = dnnl::binary::primitive_desc(add_skip_d, binary_attr, dnnl_engine);
+
+    // Get bias mem
+    auto bias_mem = sp.GetMemoryAndReshape(node.Input(IN_SLN_BIAS), bias_md, dnnl_engine);
+    // Add bias arg
+    skip_bias_args.insert({DNNL_ARG_ATTR_MULTIPLE_POST_OP(0) | DNNL_ARG_SRC_1, bias_mem});
+
+  } else {
+    add_skip_pd = dnnl::binary::primitive_desc(add_skip_d, dnnl_engine);
+  }
+
+  // Move the memory to the target device
+  auto src_mem = sp.GetMemoryAndReshape(node.Input(IN_INPUT), add_skip_pd.src0_desc(), dnnl_engine);
+  auto skip_mem = sp.GetMemoryAndReshape(node.Input(IN_SKIP), add_skip_pd.src1_desc(), dnnl_engine);
+
+  // Add args
+  skip_bias_args.insert({DNNL_ARG_SRC_0, src_mem});
+  skip_bias_args.insert({DNNL_ARG_SRC_1, skip_mem});
+  skip_bias_args.insert({DNNL_ARG_DST, src_mem});
+
+  // Create and add primitive
+  auto add_skip_prim = dnnl::binary(add_skip_pd);
+  sp.AddPrimitive(add_skip_prim, skip_bias_args);
+
+  // Return src
+  return src_mem;
+}
+
+void DnnlLayerNorm::ValidateDims(DnnlSubgraphPrimitive& sp, DnnlNode& node) {
+  // Get input and evaluate
+  auto input_dims = sp.GetMemory(node.Input(IN_INPUT)).get_desc().dims();
+  auto input_dims_size = input_dims.size();
+
+  // Check the inputs are supported by OneDNN, this is mandatory since sometimes
+  // we can't check the input size in the node capability
+  if ((input_dims_size > 5) || (input_dims_size < 2)) {
+    ORT_THROW("Input tensor dimensionality is not supported by OneDNN, got ", input_dims_size);
+  }
+
+  // To make this function compliant with all possible layernorm flavors,
+  // define gamma and shift input position, depending on the operation
+  int gamma_pos, shift_pos;
+  if (node.OpType() == "SkipLayerNormalization") {
+    // For SkipLayerNorm the spec defines the input as a 3D tensor
+    if (input_dims_size != 3) {
+      // We support 2D arrays but the expected is 3D
+      ORT_THROW("Input tensor is expected to have 3 dimensions, got ", input_dims_size);
+    }
+
+    // Get skip and evaluate
+    auto skip_dims = sp.GetMemory(node.Input(IN_SKIP)).get_desc().dims();
+    if (input_dims != skip_dims) {
+      ORT_THROW("Input and skip dimmentions do not match");
+    }
+
+    // Check if bias was provided and evaluate
+    if (node.Input(IN_SLN_BIAS).Exists()) {
+      auto bias_dims = sp.GetMemory(node.Input(IN_SLN_BIAS)).get_desc().dims();
+      if (bias_dims.size() != 1) {
+        ORT_THROW("Bias is expected to have 1 dimension, got ", bias_dims.size());
+      }
+      if (bias_dims[0] != input_dims[2]) {
+        ORT_THROW("Last dimension of bias and input does not match");
+      }
+    }
+
+    // Define the input position when using SLN
+    gamma_pos = IN_SLN_GAMMA;
+    shift_pos = IN_BETA;
+
+  // If the op is LayerNorm
+  } else{
+    // Define the input position when using LN
+    gamma_pos = IN_LN_GAMMA;
+    shift_pos = IN_LN_BIAS;
+  }
+
+  // Get gamma and evaluate
+  auto gamma_dims = sp.GetMemory(node.Input(gamma_pos)).get_desc().dims();
+  if (gamma_dims.size() != 1) {
+    ORT_THROW("Gamma is expected to have 1 dimension, got ", gamma_dims.size());
+  }
+  if (gamma_dims[0] != input_dims[input_dims_size - 1]) {
+    ORT_THROW("Last dimension of gamma and input does not match");
+  }
+
+  // Check if shift was provided and evaluate
+  if (node.Input(shift_pos).Exists()) {
+    auto beta_dims = sp.GetMemory(node.Input(shift_pos)).get_desc().dims();
+    if (beta_dims.size() != 1) {
+      ORT_THROW("Beta is expected to have 1 dimension, got ", beta_dims.size());
+    }
+    if (beta_dims[0] != input_dims[input_dims_size - 1]) {
+      ORT_THROW("Last dimension of beta and input does not match");
+    }
+  }
+}
+
+float DnnlLayerNorm::GetEpsilon(DnnlNode& node) {
+  auto attr = node.Attributes().find("epsilon");
+  float epsilon = 1e-05f;  // Default value according to ONNX spec
+  if (attr != node.Attributes().end() &&
+      attr->second().type() == ::ONNX_NAMESPACE::AttributeProto_AttributeType::AttributeProto_AttributeType_FLOAT) {
+    epsilon = attr->second().f();
+  }
+  return epsilon;
+}
+
+}  // namespace ort_dnnl
+}  // namespace onnxruntime

onnxruntime/core/providers/dnnl/subgraph/dnnl_layernorm.h

@@ -0,0 +1,46 @@
+// Copyright(C) 2022 Intel Corporation
+// Licensed under the MIT License
+
+#pragma once
+#include "dnnl_subgraph.h"
+#include "dnnl_subgraph_primitive.h"
+
+namespace onnxruntime {
+namespace ort_dnnl {
+
+class DnnlLayerNorm {
+ public:
+
+  typedef std::pair<dnnl::layer_normalization_forward, std::unordered_map<int, dnnl::memory>> ln_components;
+
+  enum InputTensorsSLN : int {
+    IN_INPUT = 0,
+    IN_SKIP = 1,
+    IN_SLN_GAMMA = 2,
+    IN_BETA = 3,        // Optional
+    IN_SLN_BIAS = 4     // Optional
+  };
+
+  enum InputTensorsLN : int {
+    // IN_INPUT = 0,
+    IN_LN_GAMMA = 1,
+    IN_LN_BIAS = 2      // Optional
+  };
+
+  enum OutputTensors : int {
+    OUT_OUTPUT = 0,
+    OUT_MEAN = 1,        // Optional
+    OUT_INV_STD_VAR = 2  // Optional
+  };
+  DnnlLayerNorm();
+
+  void CreatePrimitive(DnnlSubgraphPrimitive& sp, DnnlNode& node);
+
+ private:
+  dnnl::memory BuildSLN(DnnlSubgraphPrimitive& sp, DnnlNode& node, dnnl::engine dnnl_engine);
+  void ValidateDims(DnnlSubgraphPrimitive& sp, DnnlNode& node);
+  float GetEpsilon(DnnlNode& node);
+};
+
+}  // namespace ort_dnnl
+}  // namespace onnxruntime

onnxruntime/core/providers/dnnl/subgraph/dnnl_subgraph_primitive.cc

@@ -12,6 +12,7 @@
 #include "dnnl_elementwise.h"
 #include "dnnl_gelu.h"
 #include "dnnl_gemm.h"
+#include "dnnl_layernorm.h"
 #include "dnnl_lrn.h"
 #include "dnnl_matmul.h"
 #include "dnnl_matmul_integer.h"
@@ -151,6 +152,8 @@ void DnnlSubgraphPrimitive::AddKernels() {
       DnnlGelu().CreatePrimitive(*this, node);
     } else if (node.OpType() == "Gelu" || node.OpType() == "BiasGelu") {
       DnnlGelu().CreatePrimitive(*this, node);
+    } else if (node.OpType() == "LayerNormalization" || node.OpType() == "SkipLayerNormalization") {
+      DnnlLayerNorm().CreatePrimitive(*this, node);
     } else if (node.OpType() == "Gemm") {
       DnnlGemm().CreatePrimitive(*this, node);
     } else if (node.OpType() == "LRN") {

onnxruntime/test/contrib_ops/layer_norm_op_test.cc

@@ -66,5 +66,39 @@ TEST(LayerNormTest, BERTLayerNorm_NoBias) {
   tester.Run();
 }
 
+TEST(LayerNormTest, LayerNorm) {
+  OpTester test("LayerNormalization");
+  test.AddAttribute<float>("epsilon", 1e-05f);
+
+  std::vector<int64_t> dims{1, 2, 3};
+  test.AddInput<float>("x", dims, {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f});
+  test.AddInput<float>("gamma", {3}, {1.0f, 1.0f, 1.0f});
+  test.AddOutput<float>("output", dims, {-1.2247f, 0.0f, 1.2247f, -1.2247f, 0.0f, 1.2247f});
+  test.Run();
+}
+
+TEST(LayerNormTest, LayerNorm_Scale) {
+  OpTester test("LayerNormalization");
+  test.AddAttribute<float>("epsilon", 1e-05f);
+
+  std::vector<int64_t> dims{2, 2, 2};
+  test.AddInput<float>("x", dims, {-10.264f, 8.6453f, 43.1561f, -0.641239f, -8.2164f, 0.11412f, 41.3156f, 3.0458f});
+  test.AddInput<float>("gamma", {2}, {-0.6953f, 5.1824f});
+  test.AddOutput<float>("output", dims, {0.6953f, 5.1824f, -0.6953f, -5.1824f, 0.6953f, 5.1824f, -0.6953f, -5.1824f});
+  test.Run();
+}
+
+TEST(LayerNormTest, LayerNorm_Scale_Bias) {
+  OpTester test("LayerNormalization");
+  test.AddAttribute<float>("epsilon", 1e-05f);
+
+  std::vector<int64_t> dims{1, 3, 2};
+  test.AddInput<float>("x", dims, {1.2416f, 0.946123f, 13.1685f, 0.36423f, 21.145f, 0.03941f});
+  test.AddInput<float>("gamma", {2}, {-0.6953f, 5.1824f});
+  test.AddInput<float>("bias", {2}, {0.6435f, -0.3964f});
+  test.AddOutput<float>("output", dims, {-0.0516f, -5.5776f, -0.0518f, -5.5788f, -0.0518f, -5.5788f});
+  test.Run();
+}
+
 }  // namespace test
 }  // namespace onnxruntime