From fe7b8b80ae35c62e5bf91eed068be4279b76173d Mon Sep 17 00:00:00 2001 From: Dwayne Robinson Date: Wed, 15 Jun 2022 21:49:18 -0700 Subject: [PATCH] Revert BatchNormalization change for now, falling back to CPU on mixed types until a more advanced solution is written --- .../DmlOperatorBatchNormalization.cpp | 98 ++++++++++++++++++- .../src/Operators/OperatorRegistration.cpp | 8 +- 2 files changed, 99 insertions(+), 7 deletions(-) diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlOperatorBatchNormalization.cpp b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlOperatorBatchNormalization.cpp index e7198824bd..c69e76731a 100644 --- a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlOperatorBatchNormalization.cpp +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlOperatorBatchNormalization.cpp @@ -21,6 +21,67 @@ class DmlOperatorBatchNormalization : public DmlOperator, BatchNormalizationHelp public: DmlOperatorBatchNormalization(const MLOperatorKernelCreationContext& kernelCreationContext) + : DmlOperator(kernelCreationContext), + BatchNormalizationHelper(kernelCreationContext, kernelCreationContext.GetTensorShapeDescription()) + { + std::vector> kernelInputIndices = {X, Mean, Variance, Scale, Bias}; + DmlOperator::Initialize(kernelCreationContext, kernelInputIndices); + + ML_CHECK_VALID_ARGUMENT(m_inputTensorDescs.size() == 5); + ML_CHECK_VALID_ARGUMENT(m_outputTensorDescs.size() >= 1); + + const float epsilon = kernelCreationContext.GetOptionalAttribute(AttrName::Epsilon, 0.0f); + const int spatial = kernelCreationContext.GetOptionalAttribute(AttrName::Spatial, 1); + const std::optional fusedActivation = FusionHelpers::TryGetFusedActivationDesc(kernelCreationContext); + DML_OPERATOR_DESC fusedActivationDmlDesc = fusedActivation ? fusedActivation->GetDmlDesc() : DML_OPERATOR_DESC(); + + m_inputTensorDescs[0] = CreateTensorDescFromInput(kernelCreationContext, 0, TensorAxis::DoNotCoerce, TensorAxis::N, TensorAxis::LeftAligned); + + // Massage each of these 1D tensors (of length C) into ND tensors of the form [1,C,1,1,...]. + for (uint32_t i = Scale; i < OnnxInputIndex::Count; ++i) + { + m_inputTensorDescs[i] = CreateTensorDescFromInput(kernelCreationContext, i, TensorAxis::DoNotCoerce, TensorAxis::C, TensorAxis::LeftAligned, std::nullopt, m_inputTensorDescs[0].GetDimensionCount()); + } + + m_outputTensorDescs[0] = CreateTensorDescFromOutput(kernelCreationContext, 0, TensorAxis::DoNotCoerce, TensorAxis::N, TensorAxis::LeftAligned, std::nullopt, m_inputTensorDescs[0].GetDimensionCount()); + + ML_CHECK_VALID_ARGUMENT(m_inputTensorDescs.size() == 5); + ML_CHECK_VALID_ARGUMENT(m_outputTensorDescs.size() >= 1); + + std::vector inputDescs = GetDmlInputDescs(); + std::vector outputDescs = GetDmlOutputDescs(); + + DML_BATCH_NORMALIZATION_OPERATOR_DESC operatorDesc = {}; + operatorDesc.InputTensor = &inputDescs[X]; + operatorDesc.MeanTensor = &inputDescs[Mean]; + operatorDesc.VarianceTensor = &inputDescs[Variance]; + operatorDesc.ScaleTensor = &inputDescs[Scale]; + operatorDesc.BiasTensor = &inputDescs[Bias]; + operatorDesc.OutputTensor = &outputDescs[0]; + operatorDesc.Spatial = static_cast(spatial); + operatorDesc.Epsilon = epsilon; + operatorDesc.FusedActivation = fusedActivation ? &fusedActivationDmlDesc : nullptr; + + DML_OPERATOR_DESC opDesc = { DML_OPERATOR_BATCH_NORMALIZATION, &operatorDesc }; + SetDmlOperatorDesc(opDesc, kernelCreationContext); + } +}; + +class DmlOperatorBatchNormalization15 : public DmlOperator, BatchNormalizationHelper +{ + // This order matches the ONNX schema. + enum OnnxInputIndex + { + X, // Input + Scale, + Bias, + Mean, + Variance, + Count, + }; + +public: + DmlOperatorBatchNormalization15(const MLOperatorKernelCreationContext& kernelCreationContext) : DmlOperator(kernelCreationContext), BatchNormalizationHelper(kernelCreationContext, kernelCreationContext.GetTensorShapeDescription()) { @@ -101,15 +162,46 @@ public: void CALLBACK QueryBatchNormalization(IMLOperatorSupportQueryContextPrivate* context, /*out*/ bool* isSupported) { - // training_mode=1 is unsupported as it isn't needed for inference (https://github.com/onnx/onnx/pull/3333). + *isSupported = false; + // training_mode=1 is unsupported as it isn't needed for inference (https://github.com/onnx/onnx/pull/3333). MLOperatorAttributes attributes(context); int32_t trainingMode = attributes.GetOptionalAttribute(AttrName::TrainingMode, 0); - *isSupported = (trainingMode == 0); + if (trainingMode != 0) + { + return; + } + + if (context->GetInputCount() < 5) + { + return; + } + + // Get the data type of each tensor. + MLOperatorEdgeDescription operatorEdgeDescription[5]; + for (uint32_t i = 0; i < 5; ++i) + { + if (FAILED(context->GetInputEdgeDescription(i, &operatorEdgeDescription[i])) + || operatorEdgeDescription[i].edgeType != MLOperatorEdgeType::Tensor) + { + return; + } + } + + // Fall back if the data types of the mean/variance or scale/bias differ from the input. + MLOperatorTensorDataType inputTensorDataType = operatorEdgeDescription[0].tensorDataType; + for (uint32_t i = 1; i < 5; ++i) + { + if (operatorEdgeDescription[i].tensorDataType != inputTensorDataType) + { + return; + } + } + + *isSupported = true; } DML_OP_DEFINE_CREATION_FUNCTION(BatchNormalization, DmlOperatorBatchNormalization); -DML_OP_DEFINE_CREATION_FUNCTION(BatchNormalization15, DmlOperatorBatchNormalization); DML_OP_DEFINE_CREATION_FUNCTION(FusedBatchNormalization, DmlOperatorBatchNormalization); } // 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 b8249f7c04..3ad7c4606b 100644 --- a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/OperatorRegistration.cpp +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/OperatorRegistration.cpp @@ -397,10 +397,10 @@ constexpr static OperatorRegistrationInformation operatorRegistrationInformation {REG_INFO( 7, MaxRoiPool, typeNameListDefault, supportedTypeListFloat16to32, DmlGraphSupport::Supported)}, {REG_INFO_VER( 10, RoiAlign, typeNameListTwo, supportedTypeListRoiAlign, DmlGraphSupport::Supported)}, {REG_INFO( 7, InstanceNormalization, typeNameListDefault, supportedTypeListFloat16to32, DmlGraphSupport::Supported)}, - {REG_INFO( 7, BatchNormalization, typeNameListDefault, supportedTypeListFloat16to32, DmlGraphSupport::NotSupported)}, - {REG_INFO( 9, BatchNormalization, typeNameListDefault, supportedTypeListFloat16to32, DmlGraphSupport::NotSupported)}, // v9 just removes 'spatial' attribute. - {REG_INFO( 14, BatchNormalization, typeNameListDefault, supportedTypeListFloat16to32, DmlGraphSupport::NotSupported, requiredConstantCpuInputs(), std::nullopt, QueryBatchNormalization)}, // v14 adds training_mode attribute - {REG_INFO_VER( 15, BatchNormalization, typeNameListDefault, supportedTypeListFloat16to32, DmlGraphSupport::NotSupported, requiredConstantCpuInputs(), std::nullopt, QueryBatchNormalization)}, // v15 adds differing types for scale and bias vs input. + {REG_INFO( 7, BatchNormalization, typeNameListDefault, supportedTypeListFloat16to32, DmlGraphSupport::Supported)}, + {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, LRN, typeNameListDefault, supportedTypeListFloat16to32, DmlGraphSupport::Supported)}, {REG_INFO( 13, LRN, typeNameListDefault, supportedTypeListFloat16to32, DmlGraphSupport::Supported)}, {REG_INFO( 7, MeanVarianceNormalization, typeNameListDefault, supportedTypeListFloat16to32, DmlGraphSupport::Supported)},