From 8246ff015a7b44bb60e89d88181d13527f72eb1e Mon Sep 17 00:00:00 2001 From: Patrice Vignola Date: Tue, 13 Dec 2022 13:23:53 -0800 Subject: [PATCH] [DML EP] Add EmbedLayerNorm (#13868) ### Description Add EmbedLayerNorm to the DML EP --- docs/OperatorKernels.md | 1 + .../core/optimizer/graph_transformer_utils.cc | 2 +- .../DmlOperatorEmbedLayerNormalization.cpp | 495 ++++++++++++++++++ .../src/Operators/OperatorRegistration.cpp | 2 + .../OperatorAuthorHelper/OperatorHelper.cpp | 33 +- .../dml/OperatorAuthorHelper/OperatorHelper.h | 14 + .../OperatorAuthorHelper/OperatorVersions.h | 1 + .../contrib_ops/embed_layer_norm_op_test.cc | 7 +- 8 files changed, 551 insertions(+), 4 deletions(-) create mode 100644 onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlOperatorEmbedLayerNormalization.cpp diff --git a/docs/OperatorKernels.md b/docs/OperatorKernels.md index 9c345ee5f8..bc2a761600 100644 --- a/docs/OperatorKernels.md +++ b/docs/OperatorKernels.md @@ -1124,6 +1124,7 @@ Do not modify directly.* |BiasGelu|*in* A:**T**
*in* B:**T**
*out* C:**T**|1+|**T** = tensor(float), tensor(float16)| |ConvTransposeWithDynamicPads|*in* X:**T**
*in* W:**T**
*in* Pads:**tensor(int64)**
*in* B:**T**
*out* Y:**T**|1+|**T** = tensor(float), tensor(float16)| |DequantizeLinear|*in* x:**T1**
*in* x_scale:**T2**
*in* x_zero_point:**T1**
*out* y:**T2**|1+|**T1** = tensor(float)
**T2** = tensor(uint8)| +|EmbedLayerNormalization|*in* input_ids:**T1**
*in* segment_ids:**T1**
*in* word_embedding:**T**
*in* position_embedding:**T**
*in* segment_embedding:**T**
*in* gamma:**T**
*in* beta:**T**
*in* mask:**T1**
*in* position_ids:**T1**
*out* output:**T**
*out* mask_index:**T1**
*out* embedding_sum:**T**|1+|**T** = tensor(float), tensor(float16)| |FusedMatMul|*in* A:**T**
*in* B:**T**
*out* Y:**T**|1+|**T** = tensor(float), tensor(float16)| |Gelu|*in* X:**T**
*out* Y:**T**|1+|**T** = tensor(float), tensor(float16)| |QLinearAdd|*in* A:**T**
*in* A_scale:**tensor(float)**
*in* A_zero_point:**T**
*in* B:**T**
*in* B_scale:**tensor(float)**
*in* B_zero_point:**T**
*in* C_scale:**tensor(float)**
*in* C_zero_point:**T**
*out* C:**T**|1+|**T** = tensor(int8), tensor(uint8)| diff --git a/onnxruntime/core/optimizer/graph_transformer_utils.cc b/onnxruntime/core/optimizer/graph_transformer_utils.cc index 58d4c08b1a..2231eac0bd 100644 --- a/onnxruntime/core/optimizer/graph_transformer_utils.cc +++ b/onnxruntime/core/optimizer/graph_transformer_utils.cc @@ -266,7 +266,7 @@ InlinedVector> GenerateTransformers( transformers.emplace_back(std::make_unique(cpu_cuda_dml_rocm_eps)); transformers.emplace_back(std::make_unique(cpu_cuda_rocm_eps)); transformers.emplace_back(std::make_unique(cpu_cuda_dml_rocm_eps)); - transformers.emplace_back(std::make_unique(cpu_cuda_rocm_eps)); + transformers.emplace_back(std::make_unique(cpu_cuda_dml_rocm_eps)); transformers.emplace_back(std::make_unique(cpu_cuda_rocm_eps)); transformers.emplace_back(std::make_unique(cpu_cuda_rocm_eps)); diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlOperatorEmbedLayerNormalization.cpp b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlOperatorEmbedLayerNormalization.cpp new file mode 100644 index 0000000000..3efddb9392 --- /dev/null +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlOperatorEmbedLayerNormalization.cpp @@ -0,0 +1,495 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "precomp.h" + +/* +WordEmbeddings InputIds SegmentEmbeddings SegmentIds PositionEmbeddings PositionIds Gamma Beta Mask + │ │ │ │ │ │ │ │ │ + │ │ │ │ │ │ │ │ │ + │ │ │ │ │ │ │ │ │ + │ │ │ │ │ │ │ │ │ + │ │ │ │ │ │ │ │ │ OnesConstant + │ │ │ │ │ │ │ │ │ │ + └───────┬───────┘ └────────┬─────────┘ └─────────┬─────────┘ │ │ │ │ + │ │ │ │ │ │ │ + ▼ ▼ ▼ │ │ └───────┬────────┘ + Gather Gather Gather │ │ │ + │ │ │ │ │ ▼ + │ │ │ │ │ Equals + │ │ │ │ │ │ + │ │ │ │ │ │ + │ │ │ │ │ │ + │ │ │ │ │ │ + └─────────────┬─────────────────┘ │ │ │ │ + │ │ │ │ │ + ▼ │ │ │ │ + Add │ │ │ │ + │ │ │ │ ▼ + │ │ │ │ Reduce + └────────────────────────────┬─────────────────────────┘ │ │ │ + │ │ │ │ + ▼ │ │ │ + Add ┌──────────────────────────────────────────┘ │ │ + │ │ │ │ + ┌────────────────────────────┤ │ ┌──────────────────────────────────────────────┘ │ + │ │ │ │ │ + │ ▼ ▼ ▼ │ + │ MeanVarianceNormalization │ + │ │ ┌───────────────────────────────────────────────────────┘ + │ │ │ + │ │ │ + ▼ ▼ ▼ + EmbeddingSum Output MaskIndex + + This kernel creates a DML_GRAPH, as mentioned above. + For reference, refer to this Doc: + https://github.com/microsoft/onnxruntime/blob/main/docs/ContribOperators.md#com.microsoft.EmbedLayerNormalization + */ + +namespace Dml +{ + +class DmlOperatorEmbedLayerNormalization : public DmlOperator +{ +public: + DmlOperatorEmbedLayerNormalization(const MLOperatorKernelCreationContext& kernelCreationContext) + : DmlOperator(kernelCreationContext) + { + std::vector> kernelInputIndices = {0, 1, 2, 3, 4, 5, 6, 7, 8}; + std::vector> kernelOutputIndices = {0, 1, 2}; + + DmlOperator::Initialize(kernelCreationContext, kernelInputIndices, kernelOutputIndices); + const float epsilon = kernelCreationContext.GetOptionalAttribute(AttrName::Epsilon, DefaultEpsilon); + + assert(m_inputTensorDescs.size() == 9); + assert(m_outputTensorDescs.size() == 3); + + auto inputIdsDesc = m_inputTensorDescs[0].GetDmlDesc(); + auto segmentIdsDesc = m_inputTensorDescs[1].GetDmlDesc(); + auto wordEmbeddingDesc = m_inputTensorDescs[2].GetDmlDesc(); + auto positionEmbeddingDesc = m_inputTensorDescs[3].GetDmlDesc(); + auto segmentEmbeddingDesc = m_inputTensorDescs[4].GetDmlDesc(); + auto gammaDesc = m_inputTensorDescs[5].GetDmlDesc(); + auto betaDesc = m_inputTensorDescs[6].GetDmlDesc(); + auto maskDesc = m_inputTensorDescs[7].GetDmlDesc(); + auto positionIdsDesc = m_inputTensorDescs[8].GetDmlDesc(); + auto outputDesc = m_outputTensorDescs[0].GetDmlDesc(); + auto maskIndexDesc = m_outputTensorDescs[1].GetDmlDesc(); + auto embeddingSumDesc = m_outputTensorDescs[2].GetDmlDesc(); + + const DML_TENSOR_DATA_TYPE indicesDataType = m_inputTensorDescs[0].GetDmlDataType(); + const DML_TENSOR_DATA_TYPE valuesDataType = m_inputTensorDescs[2].GetDmlDataType(); + const uint32_t batchSize = m_inputTensorDescs[0].GetSizes()[2]; + const uint32_t sequenceLength = m_inputTensorDescs[0].GetSizes()[3]; + + // When position ids are not given, the indices are simply the sequence ids in ascending order + TensorDesc positionSequenceIdsTensorDesc(indicesDataType, std::vector({1u, 1u, 1u, sequenceLength})); + DML_TENSOR_DESC positionSequenceIdsDmlTensorDesc = positionSequenceIdsTensorDesc.GetDmlDesc(); + + DML_FILL_VALUE_SEQUENCE_OPERATOR_DESC positionSequenceIdsDesc = {}; + positionSequenceIdsDesc.ValueStart.Int32 = 0; + positionSequenceIdsDesc.ValueDelta.Int32 = 1; + positionSequenceIdsDesc.ValueDataType = indicesDataType; + positionSequenceIdsDesc.OutputTensor = &positionSequenceIdsDmlTensorDesc; + DML_OPERATOR_DESC positionSequenceIdsOpDesc = { DML_OPERATOR_FILL_VALUE_SEQUENCE, &positionSequenceIdsDesc }; + + // Gather the word embeddings + TensorDesc gatheredTensorDesc(valuesDataType, m_outputTensorDescs[0].GetSizes()); + DML_TENSOR_DESC gatheredDmlTensorDesc = gatheredTensorDesc.GetDmlDesc(); + + DML_GATHER_OPERATOR_DESC wordEmbeddingGatherDesc = {}; + wordEmbeddingGatherDesc.InputTensor = &wordEmbeddingDesc; + wordEmbeddingGatherDesc.IndicesTensor = &inputIdsDesc; + wordEmbeddingGatherDesc.OutputTensor = &gatheredDmlTensorDesc; + wordEmbeddingGatherDesc.Axis = 2; + wordEmbeddingGatherDesc.IndexDimensions = 2; + DML_OPERATOR_DESC wordEmbeddingGatherOpDesc = { DML_OPERATOR_GATHER, &wordEmbeddingGatherDesc }; + + // Gather the position embeddings + std::optional> positionIdsStrides; + if (positionIdsDesc.Desc && m_inputTensorDescs[8].GetSizes()[2] == 1 || !positionIdsDesc.Desc) + { + positionIdsStrides = std::vector({0, 0, 0, 1}); + } + + TensorDesc positionIdsTensorDesc(indicesDataType, m_inputTensorDescs[0].GetSizes(), std::move(positionIdsStrides)); + DML_TENSOR_DESC positionIdsDmlTensorDesc = positionIdsTensorDesc.GetDmlDesc(); + + DML_GATHER_OPERATOR_DESC positionEmbeddingGatherDesc = {}; + positionEmbeddingGatherDesc.InputTensor = &positionEmbeddingDesc; + positionEmbeddingGatherDesc.IndicesTensor = &positionIdsDmlTensorDesc; + positionEmbeddingGatherDesc.OutputTensor = &gatheredDmlTensorDesc; + positionEmbeddingGatherDesc.Axis = 2; + positionEmbeddingGatherDesc.IndexDimensions = 2; + DML_OPERATOR_DESC positionEmbeddingGatherOpDesc = { DML_OPERATOR_GATHER, &positionEmbeddingGatherDesc }; + + // Gather the segment embeddings + DML_GATHER_OPERATOR_DESC segmentEmbeddingGatherDesc = {}; + segmentEmbeddingGatherDesc.InputTensor = &segmentEmbeddingDesc; + segmentEmbeddingGatherDesc.IndicesTensor = &segmentIdsDesc; + segmentEmbeddingGatherDesc.OutputTensor = &gatheredDmlTensorDesc; + segmentEmbeddingGatherDesc.Axis = 2; + segmentEmbeddingGatherDesc.IndexDimensions = 2; + DML_OPERATOR_DESC segmentEmbeddingGatherOpDesc = { DML_OPERATOR_GATHER, &segmentEmbeddingGatherDesc }; + + // Add the embeddings together + DML_ELEMENT_WISE_ADD_OPERATOR_DESC embeddingsAddDesc = {}; + embeddingsAddDesc.ATensor = &gatheredDmlTensorDesc; + embeddingsAddDesc.BTensor = &gatheredDmlTensorDesc; + embeddingsAddDesc.OutputTensor = &gatheredDmlTensorDesc; + DML_OPERATOR_DESC embeddingsAddOpDesc = { DML_OPERATOR_ELEMENT_WISE_ADD, &embeddingsAddDesc }; + + // Execute MVN + std::vector mvnReductionAxes({m_inputTensorDescs[0].GetDimensionCount() - 1}); + + DML_MEAN_VARIANCE_NORMALIZATION1_OPERATOR_DESC mvnDesc = {}; + mvnDesc.InputTensor = &gatheredDmlTensorDesc; + mvnDesc.ScaleTensor = &gammaDesc; + mvnDesc.BiasTensor = &betaDesc; + mvnDesc.OutputTensor = &outputDesc; + mvnDesc.Axes = mvnReductionAxes.data(); + mvnDesc.AxisCount = gsl::narrow_cast(mvnReductionAxes.size()); + mvnDesc.NormalizeVariance = true; + mvnDesc.Epsilon = epsilon; + mvnDesc.FusedActivation = nullptr; + DML_OPERATOR_DESC mvnOpDesc = { DML_OPERATOR_MEAN_VARIANCE_NORMALIZATION1, &mvnDesc }; + + TensorDesc scalarTensorDesc(indicesDataType, std::vector(m_inputTensorDescs[0].GetDimensionCount(), 1)); + DML_TENSOR_DESC scalarDmlTensorDesc = scalarTensorDesc.GetDmlDesc(); + + // Create a tensor full of zeros + DML_FILL_VALUE_CONSTANT_OPERATOR_DESC zerosDesc = {}; + zerosDesc.Value.Int32 = 0; + zerosDesc.ValueDataType = indicesDataType; + zerosDesc.OutputTensor = &scalarDmlTensorDesc; + DML_OPERATOR_DESC zerosOpDesc = { DML_OPERATOR_FILL_VALUE_CONSTANT, &zerosDesc }; + + // Create a tensor full of ones + DML_FILL_VALUE_CONSTANT_OPERATOR_DESC onesDesc = {}; + onesDesc.Value.Int32 = 1; + onesDesc.ValueDataType = indicesDataType; + onesDesc.OutputTensor = &scalarDmlTensorDesc; + DML_OPERATOR_DESC onesOpDesc = { DML_OPERATOR_FILL_VALUE_CONSTANT, &onesDesc }; + + TensorDesc broadcastedOnesTensorDesc(indicesDataType, m_inputTensorDescs[0].GetSizes(), std::vector(m_inputTensorDescs[0].GetDimensionCount())); + DML_TENSOR_DESC broadcastedOnesDmlTensorDesc = broadcastedOnesTensorDesc.GetDmlDesc(); + + // Create the equal operator to keep all values in the mask that are 1 + TensorDesc equalOutputTensorDesc(DML_TENSOR_DATA_TYPE_UINT32, m_inputTensorDescs[0].GetSizes()); + DML_TENSOR_DESC equalOutputDmlTensorDesc = equalOutputTensorDesc.GetDmlDesc(); + + DML_ELEMENT_WISE_LOGICAL_EQUALS_OPERATOR_DESC equalDesc = {}; + equalDesc.ATensor = &maskDesc; + equalDesc.BTensor = &broadcastedOnesDmlTensorDesc; + equalDesc.OutputTensor = &equalOutputDmlTensorDesc; + DML_OPERATOR_DESC equalOpDesc = { DML_OPERATOR_ELEMENT_WISE_LOGICAL_EQUALS, &equalDesc }; + + // Reinterpret the uint32 tensor to an int32 tensor + TensorDesc sparseMaskTensorDesc(indicesDataType, m_inputTensorDescs[0].GetSizes()); + DML_TENSOR_DESC sparseMaskDmlTensorDesc = sparseMaskTensorDesc.GetDmlDesc(); + + // Create the reduce operator to sum the values of the mask for each batch + TensorDesc reducedMaskTensorDesc(indicesDataType, std::vector({1, 1, batchSize, 1})); + DML_TENSOR_DESC reducedMaskDmlTensorDesc = reducedMaskTensorDesc.GetDmlDesc(); + + uint32_t reduceAxes[] = {3}; + DML_REDUCE_OPERATOR_DESC reduceDesc = {}; + reduceDesc.Axes = reduceAxes; + reduceDesc.AxisCount = 1; + reduceDesc.Function = DML_REDUCE_FUNCTION_SUM; + reduceDesc.InputTensor = &sparseMaskDmlTensorDesc; + reduceDesc.OutputTensor = &reducedMaskDmlTensorDesc; + DML_OPERATOR_DESC reduceOpDesc = { DML_OPERATOR_REDUCE, &reduceDesc }; + + // Construct the graph + std::vector opDescs; + opDescs.reserve(11); + + std::vector inputEdges; + inputEdges.reserve(9); + + std::vector intermediateEdges; + intermediateEdges.reserve(8); + + std::vector outputEdges; + outputEdges.reserve(3); + + uint32_t currentNodeIndex = 0; + + // Insert the zeros operation into the graph + const uint32_t zerosNodeIndex = currentNodeIndex; + if (!maskDesc.Desc) + { + opDescs.push_back(&zerosOpDesc); + currentNodeIndex++; + } + + // Insert the sequence operation into the graph + const uint32_t sequenceIdsNodeIndex = currentNodeIndex; + if (!positionIdsDesc.Desc) + { + opDescs.push_back(&positionSequenceIdsOpDesc); + currentNodeIndex++; + } + + // Insert the word embeddings Gather operation into the graph + const uint32_t gatherWordsNodeIndex = currentNodeIndex; + opDescs.push_back(&wordEmbeddingGatherOpDesc); + currentNodeIndex++; + + // Insert the position embeddings Gather operation into the graph + const uint32_t gatherPositionsNodeIndex = currentNodeIndex; + opDescs.push_back(&positionEmbeddingGatherOpDesc); + currentNodeIndex++; + + // Insert the segment embeddings Gather opetation into the graph + const uint32_t gatherSegmentsNodeIndex = currentNodeIndex; + if (segmentEmbeddingDesc.Desc) + { + opDescs.push_back(&segmentEmbeddingGatherOpDesc); + currentNodeIndex++; + } + + // Insert the word+position embeddings operation into the graph + const uint32_t wordsPositionsAddNodeIndex = currentNodeIndex; + opDescs.push_back(&embeddingsAddOpDesc); + currentNodeIndex++; + + // Insert the word+position+segment embeddings operation into the graph + const uint32_t wordsPositionsSegmentsAddNodeIndex = currentNodeIndex; + if (segmentEmbeddingDesc.Desc) + { + opDescs.push_back(&embeddingsAddOpDesc); + currentNodeIndex++; + } + + // Insert the MVN operation into the graph + const uint32_t mvnNodeIndex = currentNodeIndex; + opDescs.push_back(&mvnOpDesc); + currentNodeIndex++; + + // Insert the Ones operation into the graph + const uint32_t onesNodeIndex = currentNodeIndex; + if (maskDesc.Desc) + { + opDescs.push_back(&onesOpDesc); + currentNodeIndex++; + } + + // Insert the Equal operation into the graph + const uint32_t equalNodeIndex = currentNodeIndex; + if (maskDesc.Desc) + { + opDescs.push_back(&equalOpDesc); + currentNodeIndex++; + } + + // Insert the Reduce operation into the graph + const uint32_t reduceNodeIndex = currentNodeIndex; + if (maskDesc.Desc) + { + opDescs.push_back(&reduceOpDesc); + currentNodeIndex++; + } + + // Insert the edges feeding into the words' gather operation + DML_INPUT_GRAPH_EDGE_DESC wordEmbeddingsInputEdge = {}; + wordEmbeddingsInputEdge.GraphInputIndex = 2; + wordEmbeddingsInputEdge.ToNodeIndex = gatherWordsNodeIndex; + wordEmbeddingsInputEdge.ToNodeInputIndex = 0; + inputEdges.push_back(std::move(wordEmbeddingsInputEdge)); + + DML_INPUT_GRAPH_EDGE_DESC inputIdsInputEdge = {}; + inputIdsInputEdge.GraphInputIndex = 0; + inputIdsInputEdge.ToNodeIndex = gatherWordsNodeIndex; + inputIdsInputEdge.ToNodeInputIndex = 1; + inputEdges.push_back(std::move(inputIdsInputEdge)); + + // Insert the edges feeding into the positions' gather operation + DML_INPUT_GRAPH_EDGE_DESC positionEmbeddingsInputEdge = {}; + positionEmbeddingsInputEdge.GraphInputIndex = 3; + positionEmbeddingsInputEdge.ToNodeIndex = gatherPositionsNodeIndex; + positionEmbeddingsInputEdge.ToNodeInputIndex = 0; + inputEdges.push_back(std::move(positionEmbeddingsInputEdge)); + + if (positionIdsDesc.Desc) + { + DML_INPUT_GRAPH_EDGE_DESC positionIdsInputEdge = {}; + positionIdsInputEdge.GraphInputIndex = 8; + positionIdsInputEdge.ToNodeIndex = gatherPositionsNodeIndex; + positionIdsInputEdge.ToNodeInputIndex = 1; + inputEdges.push_back(std::move(positionIdsInputEdge)); + } + else + { + DML_INTERMEDIATE_GRAPH_EDGE_DESC sequenceIntermediateEdge = {}; + sequenceIntermediateEdge.FromNodeIndex = sequenceIdsNodeIndex; + sequenceIntermediateEdge.FromNodeOutputIndex = 0; + sequenceIntermediateEdge.ToNodeIndex = gatherPositionsNodeIndex; + sequenceIntermediateEdge.ToNodeInputIndex = 1; + intermediateEdges.push_back(std::move(sequenceIntermediateEdge)); + } + + // Insert the edges feeding into the word+position operation + DML_INTERMEDIATE_GRAPH_EDGE_DESC gatheredWordsIntermediateEdge = {}; + gatheredWordsIntermediateEdge.FromNodeIndex = gatherWordsNodeIndex; + gatheredWordsIntermediateEdge.FromNodeOutputIndex = 0; + gatheredWordsIntermediateEdge.ToNodeIndex = wordsPositionsAddNodeIndex; + gatheredWordsIntermediateEdge.ToNodeInputIndex = 0; + intermediateEdges.push_back(std::move(gatheredWordsIntermediateEdge)); + + DML_INTERMEDIATE_GRAPH_EDGE_DESC gatheredPositionsIntermediateEdge = {}; + gatheredPositionsIntermediateEdge.FromNodeIndex = gatherPositionsNodeIndex; + gatheredPositionsIntermediateEdge.FromNodeOutputIndex = 0; + gatheredPositionsIntermediateEdge.ToNodeIndex = wordsPositionsAddNodeIndex; + gatheredPositionsIntermediateEdge.ToNodeInputIndex = 1; + intermediateEdges.push_back(std::move(gatheredPositionsIntermediateEdge)); + + if (segmentEmbeddingDesc.Desc) + { + // Insert the edges feeding into the segments' gather operation + DML_INPUT_GRAPH_EDGE_DESC segmentEmbeddingsInputEdge = {}; + segmentEmbeddingsInputEdge.GraphInputIndex = 4; + segmentEmbeddingsInputEdge.ToNodeIndex = gatherSegmentsNodeIndex; + segmentEmbeddingsInputEdge.ToNodeInputIndex = 0; + inputEdges.push_back(std::move(segmentEmbeddingsInputEdge)); + + DML_INPUT_GRAPH_EDGE_DESC segmentIdsInputEdge = {}; + segmentIdsInputEdge.GraphInputIndex = 1; + segmentIdsInputEdge.ToNodeIndex = gatherSegmentsNodeIndex; + segmentIdsInputEdge.ToNodeInputIndex = 1; + inputEdges.push_back(std::move(segmentIdsInputEdge)); + + // Insert the edges feeding into the word+position+segment operation + DML_INTERMEDIATE_GRAPH_EDGE_DESC wordsPositionsAddIntermediateEdge = {}; + wordsPositionsAddIntermediateEdge.FromNodeIndex = wordsPositionsAddNodeIndex; + wordsPositionsAddIntermediateEdge.FromNodeOutputIndex = 0; + wordsPositionsAddIntermediateEdge.ToNodeIndex = wordsPositionsSegmentsAddNodeIndex; + wordsPositionsAddIntermediateEdge.ToNodeInputIndex = 0; + intermediateEdges.push_back(std::move(wordsPositionsAddIntermediateEdge)); + + DML_INTERMEDIATE_GRAPH_EDGE_DESC gatheredSegmentsIntermediateEdge = {}; + gatheredSegmentsIntermediateEdge.FromNodeIndex = gatherSegmentsNodeIndex; + gatheredSegmentsIntermediateEdge.FromNodeOutputIndex = 0; + gatheredSegmentsIntermediateEdge.ToNodeIndex = wordsPositionsSegmentsAddNodeIndex; + gatheredSegmentsIntermediateEdge.ToNodeInputIndex = 1; + intermediateEdges.push_back(std::move(gatheredSegmentsIntermediateEdge)); + + // Insert the edges feeding into the MVN operation + DML_INTERMEDIATE_GRAPH_EDGE_DESC wordsPositionsSegmentsAddIntermediateEdge = {}; + wordsPositionsSegmentsAddIntermediateEdge.FromNodeIndex = wordsPositionsSegmentsAddNodeIndex; + wordsPositionsSegmentsAddIntermediateEdge.FromNodeOutputIndex = 0; + wordsPositionsSegmentsAddIntermediateEdge.ToNodeIndex = mvnNodeIndex; + wordsPositionsSegmentsAddIntermediateEdge.ToNodeInputIndex = 0; + intermediateEdges.push_back(std::move(wordsPositionsSegmentsAddIntermediateEdge)); + + if (embeddingSumDesc.Desc) + { + // Insert the edge feeding into the EmbeddingSum output + DML_OUTPUT_GRAPH_EDGE_DESC embeddingSumOutputEdge = {}; + embeddingSumOutputEdge.GraphOutputIndex = 2; + embeddingSumOutputEdge.FromNodeIndex = wordsPositionsSegmentsAddNodeIndex; + embeddingSumOutputEdge.FromNodeOutputIndex = 0; + outputEdges.push_back(std::move(embeddingSumOutputEdge)); + } + } + else + { + // Insert the edges feeding into the MVN operation + DML_INTERMEDIATE_GRAPH_EDGE_DESC wordsPositionsAddIntermediateEdge = {}; + wordsPositionsAddIntermediateEdge.FromNodeIndex = wordsPositionsAddNodeIndex; + wordsPositionsAddIntermediateEdge.FromNodeOutputIndex = 0; + wordsPositionsAddIntermediateEdge.ToNodeIndex = mvnNodeIndex; + wordsPositionsAddIntermediateEdge.ToNodeInputIndex = 0; + intermediateEdges.push_back(std::move(wordsPositionsAddIntermediateEdge)); + + if (embeddingSumDesc.Desc) + { + // Insert the edge feeding into the EmbeddingSum output + DML_OUTPUT_GRAPH_EDGE_DESC embeddingSumOutputEdge = {}; + embeddingSumOutputEdge.GraphOutputIndex = 2; + embeddingSumOutputEdge.FromNodeIndex = wordsPositionsAddNodeIndex; + embeddingSumOutputEdge.FromNodeOutputIndex = 0; + outputEdges.push_back(std::move(embeddingSumOutputEdge)); + } + } + + // Insert the remaining edges feeding into the MVN operation + DML_INPUT_GRAPH_EDGE_DESC gammaInputEdge = {}; + gammaInputEdge.GraphInputIndex = 5; + gammaInputEdge.ToNodeIndex = mvnNodeIndex; + gammaInputEdge.ToNodeInputIndex = 1; + inputEdges.push_back(std::move(gammaInputEdge)); + + DML_INPUT_GRAPH_EDGE_DESC betaInputEdge = {}; + betaInputEdge.GraphInputIndex = 6; + betaInputEdge.ToNodeIndex = mvnNodeIndex; + betaInputEdge.ToNodeInputIndex = 2; + inputEdges.push_back(std::move(betaInputEdge)); + + if (maskDesc.Desc) + { + // Insert the edges feeding into the Equal operation + DML_INPUT_GRAPH_EDGE_DESC maskInputEdge = {}; + maskInputEdge.GraphInputIndex = 7; + maskInputEdge.ToNodeIndex = equalNodeIndex; + maskInputEdge.ToNodeInputIndex = 0; + inputEdges.push_back(std::move(maskInputEdge)); + + DML_INTERMEDIATE_GRAPH_EDGE_DESC onesIntermediateEdge = {}; + onesIntermediateEdge.FromNodeIndex = onesNodeIndex; + onesIntermediateEdge.FromNodeOutputIndex = 0; + onesIntermediateEdge.ToNodeIndex = equalNodeIndex; + onesIntermediateEdge.ToNodeInputIndex = 1; + intermediateEdges.push_back(std::move(onesIntermediateEdge)); + + // Insert the edges feeding into the Reduce operation + DML_INTERMEDIATE_GRAPH_EDGE_DESC equalIntermediateEdge = {}; + equalIntermediateEdge.FromNodeIndex = equalNodeIndex; + equalIntermediateEdge.FromNodeOutputIndex = 0; + equalIntermediateEdge.ToNodeIndex = reduceNodeIndex; + equalIntermediateEdge.ToNodeInputIndex = 0; + intermediateEdges.push_back(std::move(equalIntermediateEdge)); + + // Insert the edge feeding into the MaskIndex output + DML_OUTPUT_GRAPH_EDGE_DESC maskIndexOutputEdge = {}; + maskIndexOutputEdge.GraphOutputIndex = 1; + maskIndexOutputEdge.FromNodeIndex = reduceNodeIndex; + maskIndexOutputEdge.FromNodeOutputIndex = 0; + outputEdges.push_back(std::move(maskIndexOutputEdge)); + } + else + { + // Insert the edge feeding into the MaskIndex output + DML_OUTPUT_GRAPH_EDGE_DESC maskIndexOutputEdge = {}; + maskIndexOutputEdge.GraphOutputIndex = 1; + maskIndexOutputEdge.FromNodeIndex = zerosNodeIndex; + maskIndexOutputEdge.FromNodeOutputIndex = 0; + outputEdges.push_back(std::move(maskIndexOutputEdge)); + } + + // Insert the edge feeding into the values output + DML_OUTPUT_GRAPH_EDGE_DESC outputEdge = {}; + outputEdge.GraphOutputIndex = 0; + outputEdge.FromNodeIndex = mvnNodeIndex; + outputEdge.FromNodeOutputIndex = 0; + outputEdges.push_back(std::move(outputEdge)); + + MLOperatorGraphDesc operatorGraphDesc = {}; + operatorGraphDesc.inputEdgeCount = gsl::narrow_cast(inputEdges.size()); + operatorGraphDesc.inputEdges = inputEdges.data(); + operatorGraphDesc.intermediateEdgeCount = gsl::narrow_cast(intermediateEdges.size()); + operatorGraphDesc.intermediateEdges = intermediateEdges.data(); + operatorGraphDesc.outputEdgeCount = gsl::narrow_cast(outputEdges.size()); + operatorGraphDesc.outputEdges = outputEdges.data(); + operatorGraphDesc.nodeCount = gsl::narrow_cast(opDescs.size()); + operatorGraphDesc.nodesAsOpDesc = opDescs.data(); + + SetDmlOperatorGraphDesc(std::move(operatorGraphDesc), kernelCreationContext); + } +}; + +DML_OP_DEFINE_CREATION_FUNCTION(EmbedLayerNormalization, DmlOperatorEmbedLayerNormalization); + +} // 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 78fb9f7339..1a3cbc9665 100644 --- a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/OperatorRegistration.cpp +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/OperatorRegistration.cpp @@ -101,6 +101,7 @@ DML_OP_EXTERN_CREATION_FUNCTION(BatchNormalization15); DML_OP_EXTERN_CREATION_FUNCTION(LayerNormalization); DML_OP_EXTERN_CREATION_FUNCTION(LayerNormalization17); DML_OP_EXTERN_CREATION_FUNCTION(SkipLayerNormalization); +DML_OP_EXTERN_CREATION_FUNCTION(EmbedLayerNormalization); DML_OP_EXTERN_CREATION_FUNCTION(LRN); DML_OP_EXTERN_CREATION_FUNCTION(MeanVarianceNormalization); DML_OP_EXTERN_CREATION_FUNCTION(LpNormalization); @@ -750,6 +751,7 @@ constexpr static OperatorRegistrationInformation operatorRegistrationInformation {REG_INFO( 11, DynamicQuantizeLinear, typeNameListTwo, supportedTypeListDynamicQuantizeLinear, DmlGraphSupport::Supported)}, {REG_INFO( 7, LayerNormalization, typeNameListLayerNormContrib, supportedTypeListLayerNormalizationContrib, DmlGraphSupport::Supported, requiredConstantCpuInputs(), std::nullopt, QueryLayerNormalization)}, {REG_INFO_MS( 1, SkipLayerNormalization, typeNameListDefault, supportedTypeListFloat16to32, DmlGraphSupport::Supported)}, + {REG_INFO_MS( 1, EmbedLayerNormalization, typeNameListDefault, supportedTypeListFloat16to32, DmlGraphSupport::Supported)}, }; template diff --git a/onnxruntime/core/providers/dml/OperatorAuthorHelper/OperatorHelper.cpp b/onnxruntime/core/providers/dml/OperatorAuthorHelper/OperatorHelper.cpp index babd32a560..2a51d122cf 100644 --- a/onnxruntime/core/providers/dml/OperatorAuthorHelper/OperatorHelper.cpp +++ b/onnxruntime/core/providers/dml/OperatorAuthorHelper/OperatorHelper.cpp @@ -2518,14 +2518,43 @@ namespace OperatorHelper m_sliceEnd = std::max(static_cast(trueEnd), m_sliceStart); } - std::vector ShapeHelper::GetOutputShapes(const MLShapeInferenceContext & shapeInfo) const + std::vector ShapeHelper::GetOutputShapes(const MLShapeInferenceContext& shapeInfo) const { return { EdgeShapes({m_sliceEnd - m_sliceStart}) }; } - std::vector SizeHelper::GetOutputShapes(const MLShapeInferenceContext & shapeInfo) const + std::vector SizeHelper::GetOutputShapes(const MLShapeInferenceContext& shapeInfo) const { return { EdgeShapes({}) }; } + std::vector EmbedLayerNormalizationHelper::GetOutputShapes(const MLShapeInferenceContext& shapeInfo) const + { + ML_CHECK_VALID_ARGUMENT(shapeInfo.GetInputCount() >= 3); + + auto inputIdsShape = shapeInfo.GetInputTensorShape(0); + auto wordEmbeddingShape = shapeInfo.GetInputTensorShape(2); + + // input_ids and word_embedding are 2D tensors + ML_CHECK_VALID_ARGUMENT(inputIdsShape.size() == 2); + ML_CHECK_VALID_ARGUMENT(wordEmbeddingShape.size() == 2); + + uint32_t batchSize = inputIdsShape[0]; + uint32_t sequenceLength = inputIdsShape[1]; + uint32_t hiddenSize = wordEmbeddingShape[1]; + + std::vector outputShapes; + outputShapes.reserve(3); + + outputShapes.push_back(EdgeShapes({batchSize, sequenceLength, hiddenSize})); + outputShapes.push_back(EdgeShapes({batchSize})); + + if (shapeInfo.GetOutputCount() == 3) + { + outputShapes.push_back(EdgeShapes({batchSize, sequenceLength, hiddenSize})); + } + + return outputShapes; + } + } // namespace OperatorHelper diff --git a/onnxruntime/core/providers/dml/OperatorAuthorHelper/OperatorHelper.h b/onnxruntime/core/providers/dml/OperatorAuthorHelper/OperatorHelper.h index 1b14aa3271..c3799793da 100644 --- a/onnxruntime/core/providers/dml/OperatorAuthorHelper/OperatorHelper.h +++ b/onnxruntime/core/providers/dml/OperatorAuthorHelper/OperatorHelper.h @@ -1383,6 +1383,19 @@ public: std::vector GetOutputShapes(const MLShapeInferenceContext& shapeInfo) const; }; +class EmbedLayerNormalizationHelper +{ + void Initialize( + const IKernelInformationAdapter& kernelInformation, + const IShapeInformationAdapter& shapeInformation + ); + +public: + template + EmbedLayerNormalizationHelper(const Info_t& info, const Shape_t& shapeInfo) { } + std::vector GetOutputShapes(const MLShapeInferenceContext& shapeInfo) const; +}; + using ShapeInferenceHelper_Conv = ConvHelper; using ShapeInferenceHelper_ConvTranspose = ConvTransposeHelper; using ShapeInferenceHelper_ConvTransposeWithDynamicPads = ConvTransposeWithDynamicPadsHelper; @@ -1406,6 +1419,7 @@ using ShapeInferenceHelper_MeanVarianceNormalization = GetOutputShapeAsInputShap using ShapeInferenceHelper_LayerNormalization = GetOutputShapeAsInputShapeHelper; using ShapeInferenceHelper_LayerNormalization17 = GetOutputShapeAsInputShapeHelper; using ShapeInferenceHelper_SkipLayerNormalization = GetOutputShapeAsInputShapeHelper; +using ShapeInferenceHelper_EmbedLayerNormalization = EmbedLayerNormalizationHelper; 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 7dc6096c68..c1b8b918b3 100644 --- a/onnxruntime/core/providers/dml/OperatorAuthorHelper/OperatorVersions.h +++ b/onnxruntime/core/providers/dml/OperatorAuthorHelper/OperatorVersions.h @@ -396,6 +396,7 @@ namespace OperatorHelper static const int sc_sinceVer_QLinearSigmoid = 1; static const int sc_sinceVer_Attention = 1; static const int sc_sinceVer_SkipLayerNormalization = 1; + static const int sc_sinceVer_EmbedLayerNormalization = 1; } // namespace MsftOperatorSet1 } // namespace OperatorHelper diff --git a/onnxruntime/test/contrib_ops/embed_layer_norm_op_test.cc b/onnxruntime/test/contrib_ops/embed_layer_norm_op_test.cc index df253c06ee..be384a20a3 100644 --- a/onnxruntime/test/contrib_ops/embed_layer_norm_op_test.cc +++ b/onnxruntime/test/contrib_ops/embed_layer_norm_op_test.cc @@ -17,9 +17,10 @@ static void RunTest(const embedlayernorm::OpData& data, int min_cuda_architecture = use_float16 ? 530 : 0; bool enable_cuda = HasCudaEnvironment(min_cuda_architecture); + bool enable_dml = DefaultDmlExecutionProvider().get() != nullptr; bool enable_cpu = !use_float16; - if (enable_cpu || enable_cuda) { + if (enable_cpu || enable_cuda || enable_dml) { // Input and output shapes // Input 0 - input_ids : (batch_size, sequence_size) // Input 1 - segment_ids : (batch_size, sequence_size) @@ -142,6 +143,10 @@ static void RunTest(const embedlayernorm::OpData& data, std::vector> execution_providers; execution_providers.push_back(DefaultCudaExecutionProvider()); tester.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); + } else if (enable_dml) { + std::vector> execution_providers; + execution_providers.push_back(DefaultDmlExecutionProvider()); + tester.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); } else { tester.Run(); }