[DML EP] LayerNorm Kernel (#12809)

* LayerNorm

* Initialze tensors with inputDimensionCount and add 2 versions

Co-authored-by: Sumit Agarwal <sumitagarwal@microsoft.com>
This commit is contained in:
sumitsays 2022-09-09 15:52:33 -07:00 committed by GitHub
parent 05c65a54b3
commit 5e978f351c
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
4 changed files with 86 additions and 1 deletions

View file

@ -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<std::optional<uint32_t>> 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<float>(AttrName::Epsilon, DefaultEpsilon);
int32_t onnxAxis = kernelCreationContext.GetOptionalAttribute<int32_t>(AttrName::Axis, -1);
uint32_t inputDimCount = kernelCreationContext.GetTensorShapeDescription().GetInputTensorDimensionCount(0);
onnxAxis = OperatorHelper::HandleNegativeAxis(onnxAxis, inputDimCount);
std::vector<uint32_t> onnxAxes(inputDimCount - onnxAxis);
std::iota(onnxAxes.begin(), onnxAxes.end(), onnxAxis);
std::vector<DML_TENSOR_DESC> inputDescs = GetDmlInputDescs();
std::vector<DML_TENSOR_DESC> 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<uint32_t>(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

View file

@ -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<const char*, 1> typeNameListDefault = {"T"};
constexpr static std::array<const char*, 2> typeNameListTwo = { "T1", "T2" };
constexpr static std::array<const char*, 2> typeNameListLayerNorm = { "T", "U" };
constexpr static std::array<const char*, 3> typeNameListThree = { "T1", "T2", "T3" };
constexpr static std::array<const char*, 4> typeNameListFour = { "T1", "T2", "T3", "T4" };
constexpr static std::array<const char*, 2> typeNameListTopK = { "T", "I" };
@ -324,6 +328,7 @@ constexpr static std::array<SupportedTensorDataTypes, 3> supportedTypeListIntege
constexpr static std::array<SupportedTensorDataTypes, 1> supportedTypeListInteger8 = {SupportedTensorDataTypes::Int8|SupportedTensorDataTypes::UInt8 };
constexpr static std::array<SupportedTensorDataTypes, 2> supportedTypeListRoiAlign = {SupportedTensorDataTypes::Float16to32, SupportedTensorDataTypes::Int32|SupportedTensorDataTypes::Int64 };
constexpr static std::array<SupportedTensorDataTypes, 1> supportedTypeListArgMinMax = {SupportedTensorDataTypes::Float16to32|SupportedTensorDataTypes::Ints8to64};
constexpr static std::array<SupportedTensorDataTypes, 2> supportedTypeListLayerNormalization = {SupportedTensorDataTypes::Float16to32, SupportedTensorDataTypes::Float32};
constexpr static std::array<SupportedTensorDataTypes, 3> 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)},

View file

@ -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;

View file

@ -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
{