diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlOperator.cpp b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlOperator.cpp index 83b7498282..2308766d71 100644 --- a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlOperator.cpp +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlOperator.cpp @@ -133,12 +133,21 @@ namespace Dml )); } + void DmlOperator::Initialize( + const MLOperatorKernelCreationContext& kernelInfo, + uint32_t minDimensionCount + ) + { + Initialize(kernelInfo, std::nullopt, std::nullopt, std::nullopt, std::nullopt, minDimensionCount); + } + void DmlOperator::Initialize( const MLOperatorKernelCreationContext& kernelInfo, const std::optional>>& kernelInputIndices, const std::optional>>& kernelOutputIndices, const std::optional> inputShape, - const std::optional> outputShape + const std::optional> outputShape, + uint32_t minDimensionCount ) { if (kernelInputIndices) @@ -179,7 +188,7 @@ namespace Dml TensorAxis::W, TensorAxis::RightAligned, inputShape, - NchwDimensionCount)); + minDimensionCount)); } } @@ -200,7 +209,8 @@ namespace Dml TensorAxis::DoNotCoerce, TensorAxis::W, TensorAxis::RightAligned, - outputShape)); + outputShape, + minDimensionCount)); } } } diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlOperator.h b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlOperator.h index 3ef12aed91..dedfb41e71 100644 --- a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlOperator.h +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlOperator.h @@ -29,12 +29,18 @@ namespace Dml ComPtr m_persistentResourcePoolingUnk; // Controls when the persistent resource is returned to the pool std::optional m_persistentResourceBinding; + void Initialize( + const MLOperatorKernelCreationContext& kernelInfo, + uint32_t minDimensionCount + ); + void Initialize( const MLOperatorKernelCreationContext& kernelInfo, const std::optional>>& kernelInputIndices = std::nullopt, const std::optional>>& kernelOutputIndices = std::nullopt, const std::optional> inputShape = std::nullopt, - const std::optional> outputShape = std::nullopt + const std::optional> outputShape = std::nullopt, + uint32_t minDimensionCount = NchwDimensionCount ); bool AllowHalfPrecisionComputation() const; diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlOperatorGather.cpp b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlOperatorGather.cpp index dea317686f..6dfdf0cec6 100644 --- a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlOperatorGather.cpp +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlOperatorGather.cpp @@ -16,7 +16,14 @@ public: ML_CHECK_VALID_ARGUMENT(kernelCreationContext.GetInputCount() == 2, "Gather expects 2 inputs."); ML_CHECK_VALID_ARGUMENT(kernelCreationContext.GetOutputCount() == 1, "Gather expects 1 output."); - DmlOperator::Initialize(kernelCreationContext); + auto tensorShapeDescription = kernelCreationContext.GetTensorShapeDescription(); + std::vector dataDimensions = tensorShapeDescription.GetInputTensorShape(0); + std::vector indicesDimensions = tensorShapeDescription.GetInputTensorShape(1); + std::vector outputDimensions = tensorShapeDescription.GetOutputTensorShape(0); + + size_t dimensionCountMax = std::max({dataDimensions.size(), indicesDimensions.size(), outputDimensions.size()}); + DmlOperator::Initialize(kernelCreationContext, gsl::narrow_cast(dimensionCountMax)); + DmlOperator::Remap64bitDmlDataTypesTo32bitIfNeeded(); std::vector inputDescs = GetDmlInputDescs(); @@ -24,9 +31,6 @@ public: assert(inputDescs.size() == 2); assert(outputDescs.size() == 1); - auto outputTensorShapeDescription = kernelCreationContext.GetTensorShapeDescription(); - std::vector dataDimensions = outputTensorShapeDescription.GetInputTensorShape(0); - std::vector indicesDimensions = outputTensorShapeDescription.GetInputTensorShape(1); uint32_t dmlAxis = GetDmlAdjustedAxis(m_axis, kernelCreationContext, m_inputTensorDescs.front().GetDimensionCount()); DML_GATHER_OPERATOR_DESC operatorDesc = {}; @@ -50,7 +54,14 @@ public: ML_CHECK_VALID_ARGUMENT(kernelCreationContext.GetInputCount() == 2, "GatherElements expects 2 inputs."); ML_CHECK_VALID_ARGUMENT(kernelCreationContext.GetOutputCount() == 1, "GatherElements expects 1 output."); - DmlOperator::Initialize(kernelCreationContext); + auto tensorShapeDescription = kernelCreationContext.GetTensorShapeDescription(); + std::vector dataDimensions = tensorShapeDescription.GetInputTensorShape(0); + std::vector indicesDimensions = tensorShapeDescription.GetInputTensorShape(1); + std::vector outputDimensions = tensorShapeDescription.GetOutputTensorShape(0); + + size_t dimensionCountMax = std::max({dataDimensions.size(), indicesDimensions.size(), outputDimensions.size()}); + DmlOperator::Initialize(kernelCreationContext, gsl::narrow_cast(dimensionCountMax)); + DmlOperator::Remap64bitDmlDataTypesTo32bitIfNeeded(); std::vector inputDescs = GetDmlInputDescs(); @@ -59,10 +70,6 @@ public: assert(outputDescs.size() == 1); int32_t signedOnnxAxis = kernelCreationContext.GetOptionalAttribute(AttrName::Axis, 0); - auto outputTensorShapeDescription = kernelCreationContext.GetTensorShapeDescription(); - std::vector dataDimensions = outputTensorShapeDescription.GetInputTensorShape(0); - std::vector indicesDimensions = outputTensorShapeDescription.GetInputTensorShape(1); - ML_CHECK_VALID_ARGUMENT(dataDimensions.size() <= OperatorHelper::NchwDimensionCount); uint32_t dmlAxis = GetDmlAdjustedAxis(signedOnnxAxis, kernelCreationContext, m_inputTensorDescs.front().GetDimensionCount()); DML_GATHER_ELEMENTS_OPERATOR_DESC operatorDesc = {}; @@ -86,32 +93,21 @@ public: ML_CHECK_VALID_ARGUMENT(kernelCreationContext.GetInputCount() == 2, "GatherND expects 2 inputs."); ML_CHECK_VALID_ARGUMENT(kernelCreationContext.GetOutputCount() == 1, "GatherND expects 1 output."); - DmlOperator::Initialize(kernelCreationContext); + auto tensorShapeDescription = kernelCreationContext.GetTensorShapeDescription(); + std::vector dataDimensions = tensorShapeDescription.GetInputTensorShape(0); + std::vector indicesDimensions = tensorShapeDescription.GetInputTensorShape(1); + std::vector outputDimensions = tensorShapeDescription.GetOutputTensorShape(0); + + size_t dimensionCountMax = std::max({dataDimensions.size(), indicesDimensions.size(), outputDimensions.size()}); + DmlOperator::Initialize(kernelCreationContext, gsl::narrow_cast(dimensionCountMax)); + DmlOperator::Remap64bitDmlDataTypesTo32bitIfNeeded(); - uint32_t maxDimensionCount = std::max({ - m_inputTensorDescs[0].GetDimensionCount(), - m_inputTensorDescs[1].GetDimensionCount(), - m_outputTensorDescs[0].GetDimensionCount() - }); - - // DML expects all tensors to have the same dimension count. - // Update the tensor descriptions with new sizes. - m_inputTensorDescs[0].SetDimensionCount(maxDimensionCount, TensorAxis::RightAligned); - m_inputTensorDescs[1].SetDimensionCount(maxDimensionCount, TensorAxis::RightAligned); - m_outputTensorDescs[0].SetDimensionCount(maxDimensionCount, TensorAxis::RightAligned); - std::vector inputDescs = GetDmlInputDescs(); std::vector outputDescs = GetDmlOutputDescs(); assert(inputDescs.size() == 2); assert(outputDescs.size() == 1); - auto outputTensorShapeDescription = kernelCreationContext.GetTensorShapeDescription(); - std::vector dataDimensions = outputTensorShapeDescription.GetInputTensorShape(0); - std::vector indicesDimensions = outputTensorShapeDescription.GetInputTensorShape(1); - ML_CHECK_VALID_ARGUMENT(dataDimensions.size() > m_batchCount); - ML_CHECK_VALID_ARGUMENT(indicesDimensions.size() > m_batchCount); - DML_GATHER_ND1_OPERATOR_DESC operatorDesc = {}; operatorDesc.InputTensor = &inputDescs[0]; operatorDesc.IndicesTensor = &inputDescs[1]; diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlOperatorScatter.cpp b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlOperatorScatter.cpp index 94269b3efc..9a7f7de526 100644 --- a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlOperatorScatter.cpp +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlOperatorScatter.cpp @@ -23,7 +23,6 @@ public: ML_CHECK_VALID_ARGUMENT(dataDimensions == outputDimensions); ML_CHECK_VALID_ARGUMENT(indicesDimensions == updatesDimensions); ML_CHECK_VALID_ARGUMENT(dataDimensions.size() == indicesDimensions.size()); - ML_CHECK_VALID_ARGUMENT(dataDimensions.size() <= OperatorHelper::NchwDimensionCount); // When the indices tensor is empty, Scatter is basically Identity. But since DML doesn't support empty or null // tensors, we have to special-case it outside of DML. @@ -89,12 +88,9 @@ public: std::vector updatesDimensions = tensorShapeDescription.GetInputTensorShape(2); std::vector outputDimensions = tensorShapeDescription.GetOutputTensorShape(0); ML_CHECK_VALID_ARGUMENT(dataDimensions == outputDimensions); - ML_CHECK_VALID_ARGUMENT(dataDimensions.size() <= OperatorHelper::NchwDimensionCount); - ML_CHECK_VALID_ARGUMENT(indicesDimensions.size() <= OperatorHelper::NchwDimensionCount); - ML_CHECK_VALID_ARGUMENT(updatesDimensions.size() <= OperatorHelper::NchwDimensionCount); - ML_CHECK_VALID_ARGUMENT(outputDimensions.size() <= OperatorHelper::NchwDimensionCount); - DmlOperator::Initialize(kernelCreationContext); + size_t dimensionCountMax = std::max({dataDimensions.size(), updatesDimensions.size(), indicesDimensions.size(), outputDimensions.size()}); + DmlOperator::Initialize(kernelCreationContext, gsl::narrow_cast(dimensionCountMax)); DmlOperator::Remap64bitDmlDataTypesTo32bitIfNeeded(); std::vector inputDescs = GetDmlInputDescs();