From 5e978f351cced9388e519b3cf36c26b24865254f Mon Sep 17 00:00:00 2001 From: sumitsays Date: Fri, 9 Sep 2022 15:52:33 -0700 Subject: [PATCH] [DML EP] LayerNorm Kernel (#12809) * LayerNorm * Initialze tensors with inputDimensionCount and add 2 versions Co-authored-by: Sumit Agarwal --- .../DmlOperatorLayerNormalization.cpp | 70 +++++++++++++++++++ .../src/Operators/OperatorRegistration.cpp | 7 ++ .../dml/OperatorAuthorHelper/OperatorHelper.h | 2 + .../OperatorAuthorHelper/OperatorVersions.h | 8 ++- 4 files changed, 86 insertions(+), 1 deletion(-) create mode 100644 onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlOperatorLayerNormalization.cpp diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlOperatorLayerNormalization.cpp b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlOperatorLayerNormalization.cpp new file mode 100644 index 0000000000..5883fc62c3 --- /dev/null +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlOperatorLayerNormalization.cpp @@ -0,0 +1,70 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "precomp.h" + +namespace Dml +{ + +class DmlOperatorLayerNormalization : public DmlOperator +{ +public: + DmlOperatorLayerNormalization(const MLOperatorKernelCreationContext& kernelCreationContext) + : DmlOperator(kernelCreationContext) + { + std::vector> kernelInputIndices = {0, 1, 2}; + + // Initialize Input, Scale and Bias tensors with same dimension count as Input tensor + // because DML MVN1 has a validation which requires all 3 needs to have same dimension count + // due to historical artifact. + DmlOperator::Initialize( + kernelCreationContext, + kernelInputIndices, + std::nullopt, + std::nullopt, + std::nullopt, + kernelCreationContext.GetTensorShapeDescription().GetInputTensorDimensionCount(0)); + + const float epsilon = kernelCreationContext.GetOptionalAttribute(AttrName::Epsilon, DefaultEpsilon); + + int32_t onnxAxis = kernelCreationContext.GetOptionalAttribute(AttrName::Axis, -1); + uint32_t inputDimCount = kernelCreationContext.GetTensorShapeDescription().GetInputTensorDimensionCount(0); + onnxAxis = OperatorHelper::HandleNegativeAxis(onnxAxis, inputDimCount); + std::vector onnxAxes(inputDimCount - onnxAxis); + std::iota(onnxAxes.begin(), onnxAxes.end(), onnxAxis); + + std::vector inputDescs = GetDmlInputDescs(); + std::vector outputDescs = GetDmlOutputDescs(); + + DML_MEAN_VARIANCE_NORMALIZATION1_OPERATOR_DESC operatorDesc = {}; + operatorDesc.InputTensor = &inputDescs[0]; + operatorDesc.ScaleTensor = &inputDescs[1]; + operatorDesc.BiasTensor = inputDescs[2].Desc != nullptr ? &inputDescs[2] : nullptr; + operatorDesc.OutputTensor = outputDescs.data(); + operatorDesc.Axes = onnxAxes.data(); + operatorDesc.AxisCount = gsl::narrow_cast(onnxAxes.size()); + operatorDesc.NormalizeVariance = true; + operatorDesc.Epsilon = epsilon; + operatorDesc.FusedActivation = nullptr; + + DML_OPERATOR_DESC opDesc = { DML_OPERATOR_MEAN_VARIANCE_NORMALIZATION1, &operatorDesc }; + SetDmlOperatorDesc(opDesc, kernelCreationContext); + } +}; + +void CALLBACK QueryLayerNormalization(IMLOperatorSupportQueryContextPrivate* context, /*out*/ bool* isSupported) +{ + *isSupported = true; + + // Mean and InvStdDev are not supported outputs. + if (context->GetOutputCount() > 1) + { + *isSupported = false; + return; + } +} + +DML_OP_DEFINE_CREATION_FUNCTION(LayerNormalization, DmlOperatorLayerNormalization); +DML_OP_DEFINE_CREATION_FUNCTION(LayerNormalization17, DmlOperatorLayerNormalization); + +} // namespace Dml diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/OperatorRegistration.cpp b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/OperatorRegistration.cpp index 9f52cf5279..b39d886177 100644 --- a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/OperatorRegistration.cpp +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/OperatorRegistration.cpp @@ -96,6 +96,8 @@ DML_OP_EXTERN_CREATION_FUNCTION(RoiAlign10); DML_OP_EXTERN_CREATION_FUNCTION(InstanceNormalization); DML_OP_EXTERN_CREATION_FUNCTION(BatchNormalization); DML_OP_EXTERN_CREATION_FUNCTION(BatchNormalization15); +DML_OP_EXTERN_CREATION_FUNCTION(LayerNormalization); +DML_OP_EXTERN_CREATION_FUNCTION(LayerNormalization17); DML_OP_EXTERN_CREATION_FUNCTION(LRN); DML_OP_EXTERN_CREATION_FUNCTION(MeanVarianceNormalization); DML_OP_EXTERN_CREATION_FUNCTION(LpNormalization); @@ -266,9 +268,11 @@ DML_OP_EXTERN_QUERY_FUNCTION(EinSum); DML_OP_EXTERN_QUERY_FUNCTION(RecurrentNeuralNetwork); DML_OP_EXTERN_QUERY_FUNCTION(BatchNormalization); DML_OP_EXTERN_QUERY_FUNCTION(Pad); +DML_OP_EXTERN_QUERY_FUNCTION(LayerNormalization); constexpr static std::array typeNameListDefault = {"T"}; constexpr static std::array typeNameListTwo = { "T1", "T2" }; +constexpr static std::array typeNameListLayerNorm = { "T", "U" }; constexpr static std::array typeNameListThree = { "T1", "T2", "T3" }; constexpr static std::array typeNameListFour = { "T1", "T2", "T3", "T4" }; constexpr static std::array typeNameListTopK = { "T", "I" }; @@ -324,6 +328,7 @@ constexpr static std::array supportedTypeListIntege constexpr static std::array supportedTypeListInteger8 = {SupportedTensorDataTypes::Int8|SupportedTensorDataTypes::UInt8 }; constexpr static std::array supportedTypeListRoiAlign = {SupportedTensorDataTypes::Float16to32, SupportedTensorDataTypes::Int32|SupportedTensorDataTypes::Int64 }; constexpr static std::array supportedTypeListArgMinMax = {SupportedTensorDataTypes::Float16to32|SupportedTensorDataTypes::Ints8to64}; +constexpr static std::array supportedTypeListLayerNormalization = {SupportedTensorDataTypes::Float16to32, SupportedTensorDataTypes::Float32}; constexpr static std::array supportedTypeListQLinearMatMul = { SupportedTensorDataTypes::Int8|SupportedTensorDataTypes::UInt8, @@ -403,6 +408,8 @@ constexpr static OperatorRegistrationInformation operatorRegistrationInformation {REG_INFO( 9, BatchNormalization, typeNameListDefault, supportedTypeListFloat16to32, DmlGraphSupport::Supported)}, // v9 just removes 'spatial' attribute. {REG_INFO( 14, BatchNormalization, typeNameListDefault, supportedTypeListFloat16to32, DmlGraphSupport::Supported, requiredConstantCpuInputs(), std::nullopt, QueryBatchNormalization)}, // v14 adds training_mode attribute {REG_INFO( 15, BatchNormalization, typeNameListDefault, supportedTypeListFloat16to32, DmlGraphSupport::Supported, requiredConstantCpuInputs(), std::nullopt, QueryBatchNormalization)}, // v15 adds differing types for scale and bias vs input. + {REG_INFO( 7, LayerNormalization, typeNameListLayerNorm, supportedTypeListLayerNormalization, DmlGraphSupport::Supported, requiredConstantCpuInputs(), std::nullopt, QueryLayerNormalization)}, + {REG_INFO_VER( 17, LayerNormalization, typeNameListLayerNorm, supportedTypeListLayerNormalization, DmlGraphSupport::Supported, requiredConstantCpuInputs(), std::nullopt, QueryLayerNormalization)}, {REG_INFO( 7, LRN, typeNameListDefault, supportedTypeListFloat16to32, DmlGraphSupport::Supported)}, {REG_INFO( 13, LRN, typeNameListDefault, supportedTypeListFloat16to32, DmlGraphSupport::Supported)}, {REG_INFO( 7, MeanVarianceNormalization, typeNameListDefault, supportedTypeListFloat16to32, DmlGraphSupport::Supported)}, diff --git a/onnxruntime/core/providers/dml/OperatorAuthorHelper/OperatorHelper.h b/onnxruntime/core/providers/dml/OperatorAuthorHelper/OperatorHelper.h index 3dadd4c01b..f14ef30e72 100644 --- a/onnxruntime/core/providers/dml/OperatorAuthorHelper/OperatorHelper.h +++ b/onnxruntime/core/providers/dml/OperatorAuthorHelper/OperatorHelper.h @@ -1368,6 +1368,8 @@ using ShapeInferenceHelper_BatchNormalization15 = BatchNormalizationHelper; using ShapeInferenceHelper_LRN = GetOutputShapeAsInputShapeHelper; using ShapeInferenceHelper_MeanVarianceNormalization = GetOutputShapeAsInputShapeHelper; +using ShapeInferenceHelper_LayerNormalization = GetOutputShapeAsInputShapeHelper; +using ShapeInferenceHelper_LayerNormalization17 = GetOutputShapeAsInputShapeHelper; using ShapeInferenceHelper_LpNormalization = GetOutputShapeAsInputShapeHelper; using ShapeInferenceHelper_RNN = RecurrentHelper; using ShapeInferenceHelper_GRU = RecurrentHelper; diff --git a/onnxruntime/core/providers/dml/OperatorAuthorHelper/OperatorVersions.h b/onnxruntime/core/providers/dml/OperatorAuthorHelper/OperatorVersions.h index 1e03a49068..0904312542 100644 --- a/onnxruntime/core/providers/dml/OperatorAuthorHelper/OperatorVersions.h +++ b/onnxruntime/core/providers/dml/OperatorAuthorHelper/OperatorVersions.h @@ -123,6 +123,7 @@ namespace OperatorHelper static const int sc_sinceVer_Tan = 7; static const int sc_sinceVer_Upsample = 7; static const int sc_sinceVer_Xor = 7; + static const int sc_sinceVer_LayerNormalization = 1; // Special operators static const int sc_sinceVer_MemcpyToHost = 1; @@ -363,7 +364,12 @@ namespace OperatorHelper static const int sc_sinceVer_CastLike = 15; static const int sc_sinceVer_BatchNormalization = 15; static const int sc_sinceVer_Pow = 15; - } // namespace OnnxOperatorSet14 + } // namespace OnnxOperatorSet15 + + namespace OnnxOperatorSet17 + { + static const int sc_sinceVer_LayerNormalization = 17; + } // namespace OnnxOperatorSet17 namespace MsftOperatorSet1 {