diff --git a/onnxruntime/contrib_ops/cpu/bert/embed_layer_norm.cc b/onnxruntime/contrib_ops/cpu/bert/embed_layer_norm.cc new file mode 100644 index 0000000000..b3b77aa0e1 --- /dev/null +++ b/onnxruntime/contrib_ops/cpu/bert/embed_layer_norm.cc @@ -0,0 +1,114 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "embed_layer_norm.h" +#include "embed_layer_norm_helper.h" +#include "core/util/math_cpuonly.h" + +namespace onnxruntime { +namespace contrib { +// These ops are internal-only, so register outside of onnx +#define REGISTER_KERNEL_TYPED(T) \ + ONNX_OPERATOR_TYPED_KERNEL_EX( \ + EmbedLayerNormalization, \ + kMSDomain, \ + 1, \ + T, \ + kCpuExecutionProvider, \ + KernelDefBuilder() \ + .TypeConstraint("T", DataTypeImpl::GetTensorType()), \ + EmbedLayerNorm); + +REGISTER_KERNEL_TYPED(float) + +template +EmbedLayerNorm::EmbedLayerNorm(const OpKernelInfo& info) : OpKernel(info) {} + +template +Status EmbedLayerNorm::Compute(OpKernelContext* context) const { + ORT_RETURN_IF_ERROR(embed_layer_norm::CheckInputs(context)); + + const Tensor* input_ids = context->Input(0); + const Tensor* segment_ids = context->Input(1); + const Tensor* mask = context->Input(2); + const Tensor* word_embedding = context->Input(3); + const Tensor* position_embedding = context->Input(4); + const Tensor* segment_embedding = context->Input(5); + const Tensor* gamma = context->Input(6); + const Tensor* beta = context->Input(7); + + const auto input_dims = input_ids->Shape().GetDims(); + int64_t hidden_size = word_embedding->Shape()[1]; + + std::vector out_dims; + out_dims.reserve(3); + out_dims.push_back(input_dims[0]); + out_dims.push_back(input_dims[1]); + out_dims.push_back(hidden_size); + TensorShape output_shape(out_dims); + Tensor* output = context->Output(0, output_shape); + + std::vector mask_index_dims; + mask_index_dims.push_back(input_dims[0]); + TensorShape mask_index_shape(mask_index_dims); + Tensor* mask_index = context->Output(1, mask_index_shape); + + int batch_size = static_cast(input_dims[0]); + int sequence_length = static_cast(input_dims[1]); + + int word_embedding_length = static_cast(word_embedding->Shape()[0]); + int position_embedding_length = static_cast(position_embedding->Shape()[0]); + int segment_embedding_length = static_cast(segment_embedding->Shape()[0]); + + ConstEigenArrayMap word_embedding_arr(word_embedding->template Data(), hidden_size, word_embedding_length); + ConstEigenArrayMap position_embedding_arr(position_embedding->template Data(), hidden_size, position_embedding_length); + ConstEigenArrayMap segment_embedding_arr(segment_embedding->template Data(), hidden_size, segment_embedding_length); + ConstEigenVectorMap gamma_vector(gamma->template Data(), hidden_size); + ConstEigenVectorMap beta_vector(beta->template Data(), hidden_size); + EigenArrayMap output_arr(output->template MutableData(), hidden_size, batch_size * sequence_length); + + // Calculate output + { + size_t index = 0; + for (int b = 0; b < batch_size; b++) { + for (int s = 0; s < sequence_length; s++) { + int word_col_index = input_ids->template Data()[index]; + if (word_col_index < 0 || word_col_index >= word_embedding_length) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "word_col_index out of range"); + } + int position_col_index = s; + if (position_col_index >= position_embedding_length) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "position_col_index out of range"); + } + int segment_col_index = segment_ids->template Data()[index]; + if (segment_col_index < 0 || segment_col_index >= segment_embedding_length) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "segment_col_index out of range"); + } + + output_arr.col(index) = word_embedding_arr.col(word_col_index) + + position_embedding_arr.col(position_col_index) + + segment_embedding_arr.col(segment_col_index); + output_arr.col(index) -= output_arr.col(index).mean(); + output_arr.col(index) /= static_cast(sqrt(output_arr.col(index).pow(2).mean() + 1.0e-13)); + output_arr.col(index) *= gamma_vector.array(); + output_arr.col(index) += beta_vector.array(); + index++; + } + } + } + + // Calculate mask + { + const int* mask_data = mask->template Data(); + for (int b = 0; b < batch_size; b++) { + mask_index->template MutableData()[b] = static_cast(std::count_if(mask_data + (b * sequence_length), + mask_data + (b * sequence_length) + sequence_length, + [](int v) { return v == 1; })); + } + } + + return Status::OK(); +} + +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cpu/bert/embed_layer_norm.h b/onnxruntime/contrib_ops/cpu/bert/embed_layer_norm.h new file mode 100644 index 0000000000..dcd7683c87 --- /dev/null +++ b/onnxruntime/contrib_ops/cpu/bert/embed_layer_norm.h @@ -0,0 +1,18 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include "core/common/common.h" +#include "core/framework/op_kernel.h" + +namespace onnxruntime { +namespace contrib { +template +class EmbedLayerNorm : public OpKernel { + public: + explicit EmbedLayerNorm(const OpKernelInfo& info); + Status Compute(OpKernelContext* context) const override; +}; +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cpu/bert/embed_layer_norm_helper.cc b/onnxruntime/contrib_ops/cpu/bert/embed_layer_norm_helper.cc new file mode 100644 index 0000000000..59d4f0b0a1 --- /dev/null +++ b/onnxruntime/contrib_ops/cpu/bert/embed_layer_norm_helper.cc @@ -0,0 +1,83 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "embed_layer_norm_helper.h" +#include "core/framework/tensorprotoutils.h" +#include "onnx/defs/tensor_proto_util.h" + +namespace onnxruntime { +namespace contrib { +namespace embed_layer_norm { + +Status CheckInputs(const OpKernelContext* context) { + const Tensor* input_ids = context->Input(0); + const Tensor* segment_ids = context->Input(1); + const Tensor* mask = context->Input(2); + const Tensor* word_embedding = context->Input(3); + const Tensor* position_embedding = context->Input(4); + const Tensor* segment_embedding = context->Input(5); + const Tensor* gamma = context->Input(6); + const Tensor* beta = context->Input(7); + + if (input_ids->Shape() != segment_ids->Shape() || input_ids->Shape() != mask->Shape()) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "Input 0, 1 and 2 shall have same shape"); + } + + const auto input_dims = input_ids->Shape().GetDims(); + if (input_dims.size() != 2) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "input_ids is expected to have 2 dimensions, got ", input_dims.size()); + } + + const auto word_embedding_dims = word_embedding->Shape().GetDims(); + if (word_embedding_dims.size() != 2) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "word_embedding is expected to have 2 dimensions, got ", word_embedding_dims.size()); + } + + const auto position_embedding_dims = position_embedding->Shape().GetDims(); + if (position_embedding_dims.size() != 2) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "position_embedding is expected to have 2 dimensions, got ", position_embedding_dims.size()); + } + + const auto segment_embedding_dims = segment_embedding->Shape().GetDims(); + if (segment_embedding_dims.size() != 2) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "segment_embedding is expected to have 2 dimensions, got ", segment_embedding_dims.size()); + } + + if (word_embedding_dims[1] != position_embedding_dims[1]) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "word_embedding and position_embedding shall have same dimension 1"); + } + + const auto beta_dims = beta->Shape().GetDims(); + if (beta_dims.size() != 1) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "beta is expected to have 1 dimensions, got ", beta_dims.size()); + } + + if (beta_dims[0] != word_embedding_dims[1]) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "beta is expected to have size of ", word_embedding_dims[1], ", got ", beta_dims[0]); + } + + const auto gamma_dims = gamma->Shape().GetDims(); + if (gamma_dims.size() != 1) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "gamma is expected to have 1 dimensions, got ", gamma_dims.size()); + } + + if (gamma_dims[0] != word_embedding_dims[1]) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "gamma is expected to have size of ", word_embedding_dims[1], ", got ", gamma_dims[0]); + } + + return Status::OK(); +} + +} // namespace embed_layer_norm +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cpu/bert/embed_layer_norm_helper.h b/onnxruntime/contrib_ops/cpu/bert/embed_layer_norm_helper.h new file mode 100644 index 0000000000..09eb6b4ddf --- /dev/null +++ b/onnxruntime/contrib_ops/cpu/bert/embed_layer_norm_helper.h @@ -0,0 +1,17 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include "core/common/common.h" +#include "core/framework/op_kernel.h" + +namespace onnxruntime { +namespace contrib { +namespace embed_layer_norm { + +Status CheckInputs(const OpKernelContext* context); + +} // namespace embed_layer_norm +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cpu_contrib_kernels.cc b/onnxruntime/contrib_ops/cpu_contrib_kernels.cc index fc51291cac..d7542a207b 100644 --- a/onnxruntime/contrib_ops/cpu_contrib_kernels.cc +++ b/onnxruntime/contrib_ops/cpu_contrib_kernels.cc @@ -9,6 +9,7 @@ namespace onnxruntime { namespace contrib { class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, float, SampleOp); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, float, EmbedLayerNormalization); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, float, ExpandDims); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, float, FusedConv); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, float, FusedGemm); @@ -87,6 +88,7 @@ Status RegisterCpuContribKernels(KernelRegistry& kernel_registry) { BuildKernelCreateInfo, // add more kernels here + BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, diff --git a/onnxruntime/contrib_ops/cuda/bert/embed_layer_norm.cc b/onnxruntime/contrib_ops/cuda/bert/embed_layer_norm.cc index bf88fa9ca0..e401677fd4 100644 --- a/onnxruntime/contrib_ops/cuda/bert/embed_layer_norm.cc +++ b/onnxruntime/contrib_ops/cuda/bert/embed_layer_norm.cc @@ -5,6 +5,7 @@ #include "core/providers/cuda/cudnn_common.h" #include "core/framework/tensorprotoutils.h" #include "onnx/defs/tensor_proto_util.h" +#include "contrib_ops/cpu/bert/embed_layer_norm_helper.h" #include "embed_layer_norm.h" #include "embed_layer_norm_impl.h" @@ -34,6 +35,8 @@ EmbedLayerNorm::EmbedLayerNorm(const OpKernelInfo& op_kernel_info) : CudaKern template Status EmbedLayerNorm::ComputeInternal(OpKernelContext* context) const { + ORT_RETURN_IF_ERROR(embed_layer_norm::CheckInputs(context)); + const Tensor* input_ids = context->Input(0); const Tensor* segment_ids = context->Input(1); const Tensor* mask = context->Input(2); @@ -43,40 +46,8 @@ Status EmbedLayerNorm::ComputeInternal(OpKernelContext* context) const { const Tensor* gamma = context->Input(6); const Tensor* beta = context->Input(7); - if (input_ids->Shape() != segment_ids->Shape() || input_ids->Shape() != mask->Shape()) { - return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, - "Input 0, 1 and 2 shall have same shape"); - } - const auto input_dims = input_ids->Shape().GetDims(); - if (input_dims.size() != 2) { - return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, - "input_ids is expected to have 2 dimensions, got ", input_dims.size()); - } - - const auto word_embedding_dims = word_embedding->Shape().GetDims(); - if (word_embedding_dims.size() != 2) { - return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, - "word_embedding is expected to have 2 dimensions, got ", word_embedding_dims.size()); - } - - const auto position_embedding_dims = position_embedding->Shape().GetDims(); - if (position_embedding_dims.size() != 2) { - return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, - "position_embedding is expected to have 2 dimensions, got ", position_embedding_dims.size()); - } - - const auto segment_embedding_dims = segment_embedding->Shape().GetDims(); - if (segment_embedding_dims.size() != 2) { - return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, - "segment_embedding is expected to have 2 dimensions, got ", segment_embedding_dims.size()); - } - - if (word_embedding_dims[1] != position_embedding_dims[1]) { - return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, - "word_embedding and position_embedding shall have same dimension 1"); - } - int64_t hidden_size = word_embedding_dims[1]; + int64_t hidden_size = word_embedding->Shape()[1]; std::vector out_dims; out_dims.reserve(3); diff --git a/onnxruntime/test/contrib_ops/embedlayernorm_op_test.cc b/onnxruntime/test/contrib_ops/embedlayernorm_op_test.cc index 05a4650431..f73fbf7958 100644 --- a/onnxruntime/test/contrib_ops/embedlayernorm_op_test.cc +++ b/onnxruntime/test/contrib_ops/embedlayernorm_op_test.cc @@ -24,7 +24,11 @@ static void RunTest( int hidden_size, bool use_float16 = false) { int min_cuda_architecture = use_float16 ? 530 : 0; - if (HasCudaEnvironment(min_cuda_architecture)) { + + bool enable_cuda = HasCudaEnvironment(min_cuda_architecture); + bool enable_cpu = !use_float16; + + if (enable_cpu || enable_cuda) { // Input and output shapes // Input 0 - input_ids : (batch_size, sequence_length) // Input 1 - segment_ids : (batch_size, sequence_length) @@ -76,9 +80,7 @@ static void RunTest( } tester.AddOutput("mask_index", mask_index_dims, mask_index_data); - std::vector> execution_providers; - execution_providers.push_back(DefaultCudaExecutionProvider()); - tester.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); + tester.Run(); } }