mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-17 18:40:28 +00:00
[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:
parent
05c65a54b3
commit
5e978f351c
4 changed files with 86 additions and 1 deletions
|
|
@ -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
|
||||
|
|
@ -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)},
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
|
|
|
|||
|
|
@ -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
|
||||
{
|
||||
|
|
|
|||
Loading…
Reference in a new issue