diff --git a/onnxruntime/core/providers/dnnl/dnnl_node_capability.cc b/onnxruntime/core/providers/dnnl/dnnl_node_capability.cc index 287c5f4148..5496e4672c 100644 --- a/onnxruntime/core/providers/dnnl/dnnl_node_capability.cc +++ b/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 diff --git a/onnxruntime/core/providers/dnnl/dnnl_node_capability.h b/onnxruntime/core/providers/dnnl/dnnl_node_capability.h index bb3916f12c..84c0ffc9da 100644 --- a/onnxruntime/core/providers/dnnl/dnnl_node_capability.h +++ b/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 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 diff --git a/onnxruntime/core/providers/dnnl/dnnl_op_manager.cc b/onnxruntime/core/providers/dnnl/dnnl_op_manager.cc index 0302febf96..ae98b04033 100644 --- a/onnxruntime/core/providers/dnnl/dnnl_op_manager.cc +++ b/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(new DnnlGemmNodeCapability()))); dnnl_ops_map_.emplace(std::make_pair("GlobalAveragePool", std::unique_ptr(new DnnlPoolNodeCapability()))); dnnl_ops_map_.emplace(std::make_pair("GlobalMaxPool", std::unique_ptr(new DnnlPoolNodeCapability()))); + dnnl_ops_map_.emplace(std::make_pair("LayerNormalization", std::unique_ptr(new DnnlLayerNormalizationNodeCapability()))); dnnl_ops_map_.emplace(std::make_pair("LeakyRelu", std::unique_ptr(new DnnlElementwiseCapability()))); dnnl_ops_map_.emplace(std::make_pair("Log", std::unique_ptr(new DnnlElementwiseCapability()))); dnnl_ops_map_.emplace(std::make_pair("LRN", std::unique_ptr(new DnnlDefaultNodeCapability()))); @@ -48,6 +49,7 @@ DnnlOpManager::DnnlOpManager() { dnnl_ops_map_.emplace(std::make_pair("Reshape", std::unique_ptr(new DnnlReshapeNodeCapability()))); dnnl_ops_map_.emplace(std::make_pair("Round", std::unique_ptr(new DnnlElementwiseCapability()))); dnnl_ops_map_.emplace(std::make_pair("Sigmoid", std::unique_ptr(new DnnlElementwiseCapability()))); + dnnl_ops_map_.emplace(std::make_pair("SkipLayerNormalization", std::unique_ptr(new DnnlSkipLayerNormalizationNodeCapability()))); dnnl_ops_map_.emplace(std::make_pair("Softmax", std::unique_ptr(new DnnlSoftmaxNodeCapability()))); dnnl_ops_map_.emplace(std::make_pair("Softplus", std::unique_ptr(new DnnlElementwiseCapability()))); dnnl_ops_map_.emplace(std::make_pair("Squeeze", std::unique_ptr(new DnnlSqueezeNodeCapability()))); diff --git a/onnxruntime/core/providers/dnnl/subgraph/dnnl_layernorm.cc b/onnxruntime/core/providers/dnnl/subgraph/dnnl_layernorm.cc new file mode 100644 index 0000000000..05e0c034b3 --- /dev/null +++ b/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 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 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 \ No newline at end of file diff --git a/onnxruntime/core/providers/dnnl/subgraph/dnnl_layernorm.h b/onnxruntime/core/providers/dnnl/subgraph/dnnl_layernorm.h new file mode 100644 index 0000000000..5fe27cb4f2 --- /dev/null +++ b/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> 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 \ No newline at end of file diff --git a/onnxruntime/core/providers/dnnl/subgraph/dnnl_subgraph_primitive.cc b/onnxruntime/core/providers/dnnl/subgraph/dnnl_subgraph_primitive.cc index 2ef3c08af4..d3e6636875 100644 --- a/onnxruntime/core/providers/dnnl/subgraph/dnnl_subgraph_primitive.cc +++ b/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") { diff --git a/onnxruntime/test/contrib_ops/layer_norm_op_test.cc b/onnxruntime/test/contrib_ops/layer_norm_op_test.cc index 2e71186b87..00ea1439ef 100644 --- a/onnxruntime/test/contrib_ops/layer_norm_op_test.cc +++ b/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("epsilon", 1e-05f); + + std::vector dims{1, 2, 3}; + test.AddInput("x", dims, {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f}); + test.AddInput("gamma", {3}, {1.0f, 1.0f, 1.0f}); + test.AddOutput("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("epsilon", 1e-05f); + + std::vector dims{2, 2, 2}; + test.AddInput("x", dims, {-10.264f, 8.6453f, 43.1561f, -0.641239f, -8.2164f, 0.11412f, 41.3156f, 3.0458f}); + test.AddInput("gamma", {2}, {-0.6953f, 5.1824f}); + test.AddOutput("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("epsilon", 1e-05f); + + std::vector dims{1, 3, 2}; + test.AddInput("x", dims, {1.2416f, 0.946123f, 13.1685f, 0.36423f, 21.145f, 0.03941f}); + test.AddInput("gamma", {2}, {-0.6953f, 5.1824f}); + test.AddInput("bias", {2}, {0.6435f, -0.3964f}); + test.AddOutput("output", dims, {-0.0516f, -5.5776f, -0.0518f, -5.5788f, -0.0518f, -5.5788f}); + test.Run(); +} + } // namespace test } // namespace onnxruntime