From 4a676b011a71f5a1ddacdac566f5579f0e132976 Mon Sep 17 00:00:00 2001 From: Patrice Vignola Date: Mon, 10 Apr 2023 14:46:33 -0700 Subject: [PATCH] [DML EP] Add BiasAdd (#15211) --- docs/OperatorKernels.md | 1 + .../src/Operators/DmlOperatorBiasAdd.cpp | 113 ++++++++++++++++++ .../src/Operators/OperatorRegistration.cpp | 2 + .../dml/OperatorAuthorHelper/OperatorHelper.h | 1 + .../OperatorAuthorHelper/OperatorVersions.h | 1 + .../test/contrib_ops/bias_add_op_test.cc | 8 +- 6 files changed, 124 insertions(+), 2 deletions(-) create mode 100644 onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlOperatorBiasAdd.cpp diff --git a/docs/OperatorKernels.md b/docs/OperatorKernels.md index 7ed747da0a..99c265729d 100644 --- a/docs/OperatorKernels.md +++ b/docs/OperatorKernels.md @@ -1173,6 +1173,7 @@ Do not modify directly.* | | |**Operator Domain:** *com.microsoft*|||| |Attention|*in* input:**T**
*in* weights:**T**
*in* bias:**T**
*in* mask_index:**M**
*in* past:**T**
*in* relative_position_bias:**T**
*in* past_sequence_length:**M**
*out* output:**T**
*out* present:**T**|1+|**M** = tensor(int32)
**T** = tensor(float), tensor(float16)| +|BiasAdd|*in* X:**T**
*in* bias:**T**
*in* skip:**T**
*out* Y:**T**|1+|**T** = tensor(float), tensor(float16)| |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)| diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlOperatorBiasAdd.cpp b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlOperatorBiasAdd.cpp new file mode 100644 index 0000000000..1c851c94c4 --- /dev/null +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlOperatorBiasAdd.cpp @@ -0,0 +1,113 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "precomp.h" + +namespace Dml +{ + +class DmlOperatorBiasAdd : public DmlOperator +{ +public: + DmlOperatorBiasAdd(const MLOperatorKernelCreationContext& kernelCreationContext) + : DmlOperator(kernelCreationContext) + { + ML_CHECK_VALID_ARGUMENT(kernelCreationContext.GetInputCount() == 3); + ML_CHECK_VALID_ARGUMENT(kernelCreationContext.GetOutputCount() == 1); + + // Broadcast bias to have the same dimensions as the input + std::vector inputTensorShape = kernelCreationContext.GetTensorShapeDescription().GetInputTensorShape(0); + DmlOperator::Initialize(kernelCreationContext, std::nullopt, std::nullopt, inputTensorShape); + + ML_CHECK_VALID_ARGUMENT(m_inputTensorDescs.size() == 3); + ML_CHECK_VALID_ARGUMENT(m_outputTensorDescs.size() == 1); + ML_CHECK_VALID_ARGUMENT(m_inputTensorDescs[0].GetDimensionCount() == 4); + ML_CHECK_VALID_ARGUMENT(m_inputTensorDescs[1].GetDimensionCount() == 4); + ML_CHECK_VALID_ARGUMENT(m_inputTensorDescs[2].GetDimensionCount() == 4); + ML_CHECK_VALID_ARGUMENT(m_inputTensorDescs[0].GetSizes() == m_inputTensorDescs[2].GetSizes()); + + std::vector inputDescs = GetDmlInputDescs(); + std::vector outputDescs = GetDmlOutputDescs(); + + // 1. Add the bias + DML_ELEMENT_WISE_ADD_OPERATOR_DESC addBiasDesc{}; + addBiasDesc.ATensor = &inputDescs[0]; + addBiasDesc.BTensor = &inputDescs[1]; + addBiasDesc.OutputTensor = &inputDescs[0]; + DML_OPERATOR_DESC dmlAddBiasDesc = { DML_OPERATOR_ELEMENT_WISE_ADD, &addBiasDesc }; + + // 2. Add the residual inputs + DML_ELEMENT_WISE_ADD_OPERATOR_DESC addSkipDesc{}; + addSkipDesc.ATensor = &inputDescs[0]; + addSkipDesc.BTensor = &inputDescs[2]; + addSkipDesc.OutputTensor = &inputDescs[0]; + DML_OPERATOR_DESC dmlAddSkipDesc = { DML_OPERATOR_ELEMENT_WISE_ADD, &addSkipDesc }; + + enum NodeIndex + { + addBiasNodeIndex, + addSkipNodeIndex, + nodeCount, + }; + + // Construct the graph + std::vector opDescs(nodeCount); + opDescs[addBiasNodeIndex] = &dmlAddBiasDesc; + opDescs[addSkipNodeIndex] = &dmlAddSkipDesc; + + std::vector inputEdges; + inputEdges.reserve(3); + + std::vector intermediateEdges; + intermediateEdges.reserve(1); + + std::vector outputEdges; + outputEdges.reserve(1); + + DML_INPUT_GRAPH_EDGE_DESC inputToAddBiasEdge{}; + inputToAddBiasEdge.GraphInputIndex = 0; + inputToAddBiasEdge.ToNodeIndex = addBiasNodeIndex; + inputToAddBiasEdge.ToNodeInputIndex = 0; + inputEdges.push_back(inputToAddBiasEdge); + + DML_INPUT_GRAPH_EDGE_DESC biasToAddBiasEdge{}; + biasToAddBiasEdge.GraphInputIndex = 1; + biasToAddBiasEdge.ToNodeIndex = addBiasNodeIndex; + biasToAddBiasEdge.ToNodeInputIndex = 1; + inputEdges.push_back(biasToAddBiasEdge); + + DML_INTERMEDIATE_GRAPH_EDGE_DESC addBiasToAddSkipEdge{}; + addBiasToAddSkipEdge.FromNodeIndex = addBiasNodeIndex; + addBiasToAddSkipEdge.FromNodeOutputIndex = 0; + addBiasToAddSkipEdge.ToNodeIndex = addSkipNodeIndex; + addBiasToAddSkipEdge.ToNodeInputIndex = 0; + intermediateEdges.push_back(addBiasToAddSkipEdge); + + DML_INPUT_GRAPH_EDGE_DESC skipToAddSkipEdge{}; + skipToAddSkipEdge.GraphInputIndex = 2; + skipToAddSkipEdge.ToNodeIndex = addSkipNodeIndex; + skipToAddSkipEdge.ToNodeInputIndex = 1; + inputEdges.push_back(skipToAddSkipEdge); + + DML_OUTPUT_GRAPH_EDGE_DESC addSkipToOutputEdge{}; + addSkipToOutputEdge.FromNodeIndex = addSkipNodeIndex; + addSkipToOutputEdge.FromNodeOutputIndex = 0; + addSkipToOutputEdge.GraphOutputIndex = 0; + outputEdges.push_back(addSkipToOutputEdge); + + 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(BiasAdd, DmlOperatorBiasAdd); + +} // 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 349efdab3e..e3067a4a22 100644 --- a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/OperatorRegistration.cpp +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/OperatorRegistration.cpp @@ -211,6 +211,7 @@ 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(BiasAdd); DML_OP_EXTERN_CREATION_FUNCTION(LRN); DML_OP_EXTERN_CREATION_FUNCTION(MeanVarianceNormalization); DML_OP_EXTERN_CREATION_FUNCTION(GroupNorm); @@ -879,6 +880,7 @@ constexpr static OperatorRegistrationInformation operatorRegistrationInformation {REG_INFO( 7, LayerNormalization, typeNameListLayerNormContrib, supportedTypeListLayerNormalizationContrib, DmlGraphSupport::Supported, requiredConstantCpuInputs(), std::nullopt, QueryLayerNormalization)}, {REG_INFO_MS( 1, SkipLayerNormalization, typeNameListDefault, supportedTypeListFloat16to32, DmlGraphSupport::Supported, requiredConstantCpuInputs(), std::nullopt, QuerySkipLayerNormalization)}, {REG_INFO_MS( 1, EmbedLayerNormalization, typeNameListDefault, supportedTypeListFloat16to32, DmlGraphSupport::Supported)}, + {REG_INFO_MS( 1, BiasAdd, typeNameListDefault, supportedTypeListFloat16to32, DmlGraphSupport::Supported)}, {REG_INFO_MS( 1, QuickGelu, typeNameListDefault, supportedTypeListFloat16to32, DmlGraphSupport::Supported)}, {REG_INFO_MS( 1, GroupNorm, typeNameListGroupNorm, supportedTypeListGroupNorm, DmlGraphSupport::Supported)}, }; diff --git a/onnxruntime/core/providers/dml/OperatorAuthorHelper/OperatorHelper.h b/onnxruntime/core/providers/dml/OperatorAuthorHelper/OperatorHelper.h index c0d1562f8a..7ce27bfebe 100644 --- a/onnxruntime/core/providers/dml/OperatorAuthorHelper/OperatorHelper.h +++ b/onnxruntime/core/providers/dml/OperatorAuthorHelper/OperatorHelper.h @@ -1444,6 +1444,7 @@ using ShapeInferenceHelper_LpNormalization = GetOutputShapeAsInputShapeHelper; using ShapeInferenceHelper_RNN = RecurrentHelper; using ShapeInferenceHelper_GRU = RecurrentHelper; using ShapeInferenceHelper_LSTM = RecurrentHelper; +using ShapeInferenceHelper_BiasAdd = GetOutputShapeAsInputShapeHelper; using ShapeInferenceHelper_Gather = GatherHelper; using ShapeInferenceHelper_GatherElements = GetOutputShapeAsSpecificInputShapeHelper<1>; diff --git a/onnxruntime/core/providers/dml/OperatorAuthorHelper/OperatorVersions.h b/onnxruntime/core/providers/dml/OperatorAuthorHelper/OperatorVersions.h index 57464786d7..d5e4f8f134 100644 --- a/onnxruntime/core/providers/dml/OperatorAuthorHelper/OperatorVersions.h +++ b/onnxruntime/core/providers/dml/OperatorAuthorHelper/OperatorVersions.h @@ -410,6 +410,7 @@ namespace OperatorHelper static const int sc_sinceVer_Attention = 1; static const int sc_sinceVer_SkipLayerNormalization = 1; static const int sc_sinceVer_EmbedLayerNormalization = 1; + static const int sc_sinceVer_BiasAdd = 1; static const int sc_sinceVer_QuickGelu = 1; static const int sc_sinceVer_GroupNorm = 1; } // namespace MsftOperatorSet1 diff --git a/onnxruntime/test/contrib_ops/bias_add_op_test.cc b/onnxruntime/test/contrib_ops/bias_add_op_test.cc index 04bc68ac71..7699f4479c 100644 --- a/onnxruntime/test/contrib_ops/bias_add_op_test.cc +++ b/onnxruntime/test/contrib_ops/bias_add_op_test.cc @@ -13,7 +13,7 @@ using namespace onnxruntime::test; namespace onnxruntime { namespace test { -#if defined(USE_CUDA) || defined(USE_ROCM) +#if defined(USE_CUDA) || defined(USE_ROCM) || defined(USE_DML) static std::vector GetExpectedResult(const std::vector& input_data, const std::vector& bias_data, const std::vector& skip_data) { @@ -39,8 +39,9 @@ static void RunSkipBiasGpuTest(const std::vector& input_data, int min_cuda_architecture = use_float16 ? 530 : 0; bool enable_cuda = HasCudaEnvironment(min_cuda_architecture); bool enable_rocm = (nullptr != DefaultRocmExecutionProvider().get()); + bool enable_dml = (nullptr != DefaultDmlExecutionProvider().get()); - if (!enable_cuda && !enable_rocm) { + if (!enable_cuda && !enable_rocm && !enable_dml) { return; } @@ -65,6 +66,9 @@ static void RunSkipBiasGpuTest(const std::vector& input_data, if (enable_rocm) { execution_providers.push_back(DefaultRocmExecutionProvider()); } + if (enable_dml) { + execution_providers.push_back(DefaultDmlExecutionProvider()); + } tester.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); }