diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/DmlCommon.cpp b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/DmlCommon.cpp index 4132995c29..70a29a2d3d 100644 --- a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/DmlCommon.cpp +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/DmlCommon.cpp @@ -132,4 +132,16 @@ uint32_t GetSupportedDeviceDataTypeMask(IDMLDevice* dmlDevice) return deviceTypeMask; } +void GetDescendingPackedStrides(gsl::span sizes, /*out*/ gsl::span strides) +{ + assert(sizes.size() == strides.size()); + + uint32_t stride = 1; + for (size_t i = strides.size(); i-- > 0; ) + { + strides[i] = stride; + stride *= sizes[i]; + } +} + } // namespace Dml diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/DmlCommon.h b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/DmlCommon.h index 536fa9beeb..0f0c533558 100644 --- a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/DmlCommon.h +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/DmlCommon.h @@ -19,6 +19,7 @@ namespace Dml size_t ComputeByteSizeFromDimensions(gsl::span dimensions, MLOperatorTensorDataType tensorDataType); size_t ComputeByteSizeFromTensor(IMLOperatorTensor& tensor); uint32_t GetSupportedDeviceDataTypeMask(IDMLDevice* dmlDevice); + void GetDescendingPackedStrides(gsl::span sizes, /*out*/ gsl::span strides); bool IsSigned(DML_TENSOR_DATA_TYPE dataType); diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlOperatorEinSum.cpp b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlOperatorEinSum.cpp new file mode 100644 index 0000000000..84d8fc0e34 --- /dev/null +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlOperatorEinSum.cpp @@ -0,0 +1,176 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "precomp.h" + +namespace Dml +{ + +class DmlOperatorEinSum : public DmlOperator, public EinSumHelper +{ +public: + DmlOperatorEinSum(const MLOperatorKernelCreationContext& kernelCreationContext, uint32_t opsetVersion) + : DmlOperator(kernelCreationContext), + EinSumHelper(kernelCreationContext, kernelCreationContext.GetTensorShapeDescription(), opsetVersion) + { + ML_CHECK_VALID_ARGUMENT(kernelCreationContext.GetInputCount() + 1 == m_components.size(), "EinSum input tensor count is inconsistent with the equation component count."); + ML_CHECK_VALID_ARGUMENT(kernelCreationContext.GetOutputCount() == 1, "EinSum expects one output tensor."); + + DmlOperator::Initialize(kernelCreationContext); + + std::vector inputDescs = GetDmlInputDescs(); + std::vector outputDescs = GetDmlOutputDescs(); + + static_assert(RecognizedOperatorType::Total == static_cast(8), "Update this switch."); + switch (m_recognizedOperatorType) + { + case RecognizedOperatorType::Multiply: + { + DML_ELEMENT_WISE_MULTIPLY_OPERATOR_DESC operatorDesc = {}; + operatorDesc.ATensor = &inputDescs[0]; + operatorDesc.BTensor = &inputDescs[1]; + operatorDesc.OutputTensor = outputDescs.data(); + + SetDmlOperatorDesc({ DML_OPERATOR_ELEMENT_WISE_MULTIPLY, &operatorDesc}, kernelCreationContext); + } + break; + + case RecognizedOperatorType::MatMul: + case RecognizedOperatorType::MatMulTransposeA: + case RecognizedOperatorType::MatMulTransposeB: + { + DML_GEMM_OPERATOR_DESC operatorDesc = {}; + operatorDesc.ATensor = &inputDescs[0]; + operatorDesc.BTensor = &inputDescs[1]; + // No operatorDesc.CTensor + operatorDesc.OutputTensor = &outputDescs[0]; + operatorDesc.TransA = (m_recognizedOperatorType == RecognizedOperatorType::MatMulTransposeA) ? DML_MATRIX_TRANSFORM_TRANSPOSE : DML_MATRIX_TRANSFORM_NONE; + operatorDesc.TransB = (m_recognizedOperatorType == RecognizedOperatorType::MatMulTransposeB) ? DML_MATRIX_TRANSFORM_TRANSPOSE : DML_MATRIX_TRANSFORM_NONE; + operatorDesc.Alpha = 1.0; + operatorDesc.Beta = 0.0; + operatorDesc.FusedActivation = nullptr; + + SetDmlOperatorDesc({ DML_OPERATOR_GEMM, &operatorDesc }, kernelCreationContext); + } + break; + + case RecognizedOperatorType::ReduceSum: + { + // Get how many axes are kept in the final output, either 0 or 1 supported + // meaning full reduction or partial with one dimension left. *It could be + // generalized to support any number of output dimensions, but it would need + // to accomodate for Transposition too if the output labels are reordered. + auto keptAxes = m_components.back().GetLabels(m_labelIndices); + assert(keptAxes.size() <= 1); + + // DML expects output rank to match input rank (as if ONNX ReduceSum keepdims=1). + // So replace the existing tensor description with the input sizes, except that + // reduced dimensions have size 1. + std::vector reducedAxes; + std::vector inputSizes = kernelCreationContext.GetTensorShapeDescription().GetInputTensorShape(0); + std::vector outputSizes = inputSizes; + + // Determine which axes are being reduced by taking the opposite of those kept. + uint32_t keptAxesMask = 0; + for (auto axis : keptAxes) + { + keptAxesMask |= (1 << axis); + } + for (uint32_t axis = 0, axisCount = static_cast(outputSizes.size()); axis < axisCount; ++axis) + { + if (~keptAxesMask & (1<(reducedAxes.size()); + + SetDmlOperatorDesc({ DML_OPERATOR_REDUCE, &operatorDesc }, kernelCreationContext); + } + break; + + case RecognizedOperatorType::Transpose: + case RecognizedOperatorType::Identity: + { + if (m_recognizedOperatorType == RecognizedOperatorType::Transpose) + { + // Transpose via input strides. The output tensor is not strided. + assert(m_components.front().GetDimensionCount() == m_components.back().GetDimensionCount()); + auto originalStrides = m_inputTensorDescs.front().GetStrides(); + std::vector inputSizes = kernelCreationContext.GetTensorShapeDescription().GetInputTensorShape(0); + std::vector inputStrides(inputSizes.size()); + + // If there were no strides, compute them based in descending packed order + // based on the input sizes. + if (originalStrides.empty()) + { + Dml::GetDescendingPackedStrides(inputSizes, /*out*/ inputStrides); + } + else // Copy the original strides. + { + assert(originalStrides.size() >= inputStrides.size()); + size_t offset = originalStrides.size() - inputStrides.size(); + inputStrides.assign(originalStrides.begin() + offset, originalStrides.end()); + } + + // Remap transposed strides using the component labels from input to output. + auto labelIndices = m_components.back().GetLabels(m_labelIndices); + + std::vector newStrides(inputStrides.size()); + std::vector newSizes(inputStrides.size()); + for (size_t i = 0, dimensionCount = inputStrides.size(); i < dimensionCount; ++i) + { + uint32_t labelIndex = labelIndices[i]; + assert(labelIndex < inputStrides.size()); + newSizes[i] = inputSizes[labelIndex]; + newStrides[i] = inputStrides[labelIndex]; + } + + // Override the initial input tensor with the new strides. + m_inputTensorDescs.front() = TensorDesc(m_inputTensorDescs.front().GetDmlDataType(), newSizes, newStrides, 0); + m_outputTensorDescs.front() = TensorDesc(m_outputTensorDescs.front().GetDmlDataType(), newSizes, std::nullopt, 0); + m_inputTensorDescs.front().GetDmlDesc(); // Discard value, but keep side effect of refreshing the DML view. + m_outputTensorDescs.front().GetDmlDesc(); // Discard value, but keep side effect of refreshing the DML view. + } + + DML_ELEMENT_WISE_IDENTITY_OPERATOR_DESC operatorDesc = {}; + operatorDesc.InputTensor = inputDescs.data(); + operatorDesc.OutputTensor = outputDescs.data(); + + SetDmlOperatorDesc({ DML_OPERATOR_ELEMENT_WISE_IDENTITY, &operatorDesc}, kernelCreationContext); + } + break; + + default: + return; + } + } +}; + +void CALLBACK QueryEinSum(IMLOperatorSupportQueryContextPrivate* context, bool* isSupported) +{ + *isSupported = false; + + MLOperatorAttributes attributes(context); + EinSumHelper helper(attributes); + auto recognizedOperatorType = helper.GetRecognizedOperatorType(); + + static_assert(EinSumHelper::RecognizedOperatorType::Total == static_cast(8), "Verify this test still matches the switch above."); + *isSupported = (recognizedOperatorType != EinSumHelper::RecognizedOperatorType::None); +} + +DML_OP_DEFINE_CREATION_FUNCTION(Einsum12, VersionedKernel); + +} // namespace Dml diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlOperatorResize.cpp b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlOperatorResize.cpp index 5e554ec340..a3547188f0 100644 --- a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlOperatorResize.cpp +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlOperatorResize.cpp @@ -251,17 +251,6 @@ public: } }; -// A specific type of operation for registration. -template -struct DmlOperatorResizeTemplate : public DmlOperatorResize -{ -public: - DmlOperatorResizeTemplate(const MLOperatorKernelCreationContext& kernelInfo) - : DmlOperatorResize(kernelInfo, OpsetVersion) - { - } -}; - void CALLBACK QueryResize(IMLOperatorSupportQueryContextPrivate* context, bool* isSupported) { *isSupported = false; diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlOperatorTranspose.cpp b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlOperatorTranspose.cpp index 24258b3daf..851dc2ec9e 100644 --- a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlOperatorTranspose.cpp +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlOperatorTranspose.cpp @@ -19,41 +19,25 @@ public: ML_CHECK_VALID_ARGUMENT(kernelInfo.GetOutputCount() >= 1); DmlOperator::Initialize(kernelInfo); - const MLOperatorEdgeDescription inputEdgeDescription = kernelInfo.GetInputEdgeDescription(0); - const std::vector originalSizes = kernelInfo.GetTensorShapeDescription().GetInputTensorShape(0); ML_CHECK_VALID_ARGUMENT(m_permutations.size() == originalSizes.size()); // Calculate strides from original shape. ML_CHECK_VALID_ARGUMENT(!originalSizes.empty()); std::vector inputStrides(originalSizes.size()); - inputStrides.back() = 1; - for (int i = gsl::narrow_cast(inputStrides.size()) - 2; i >= 0; i--) - { - inputStrides[i] = inputStrides[i + 1] * gsl::narrow_cast(originalSizes[i + 1]); - } + Dml::GetDescendingPackedStrides(originalSizes, /*out*/ inputStrides); - const int leadingDims = gsl::narrow_cast(m_inputTensorDescs.front().GetDimensionCount() - originalSizes.size()); - - std::vector sizes(m_inputTensorDescs.front().GetDimensionCount()); - std::vector strides(m_inputTensorDescs.front().GetDimensionCount()); - - // Fill leading tensor desc sizes/strides with defaults. - for (int dimDML = 0; dimDML < leadingDims; ++dimDML) - { - sizes[dimDML] = 1; - strides[dimDML] = 0; - } + std::vector sizes(inputStrides.size()); + std::vector strides(inputStrides.size()); // Permute the shape and strides. for (int dimInput = 0, dimCount = gsl::narrow_cast(originalSizes.size()); dimInput < dimCount; ++dimInput) { - int dimDML = dimInput + leadingDims; int dimPermuted = m_permutations[dimInput]; ML_CHECK_VALID_ARGUMENT(gsl::narrow_cast(dimPermuted) < originalSizes.size()); - sizes[dimDML] = gsl::narrow_cast(originalSizes[dimPermuted]); - strides[dimDML] = inputStrides[dimPermuted]; + sizes[dimInput] = originalSizes[dimPermuted]; + strides[dimInput] = inputStrides[dimPermuted]; } // Override the initial tensor descs. The output tensor is not strided. diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/OperatorRegistration.cpp b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/OperatorRegistration.cpp index 442ecf4e89..3292e9d6fb 100644 --- a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/OperatorRegistration.cpp +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/OperatorRegistration.cpp @@ -143,6 +143,7 @@ DML_OP_EXTERN_CREATION_FUNCTION(Mean); DML_OP_EXTERN_CREATION_FUNCTION(Max); DML_OP_EXTERN_CREATION_FUNCTION(Min); DML_OP_EXTERN_CREATION_FUNCTION(ReduceSum); +DML_OP_EXTERN_CREATION_FUNCTION(Einsum12); DML_OP_EXTERN_CREATION_FUNCTION(ReduceMean); DML_OP_EXTERN_CREATION_FUNCTION(ReduceProd); DML_OP_EXTERN_CREATION_FUNCTION(ReduceLogSum); @@ -244,6 +245,7 @@ DML_OP_EXTERN_CREATION_FUNCTION(ConvInteger); DML_OP_EXTERN_QUERY_FUNCTION(MaxPool); DML_OP_EXTERN_QUERY_FUNCTION(Slice); DML_OP_EXTERN_QUERY_FUNCTION(Resize); +DML_OP_EXTERN_QUERY_FUNCTION(EinSum); constexpr static std::array typeNameListDefault = {"T"}; constexpr static std::array typeNameListTwo = { "T1", "T2" }; @@ -460,6 +462,7 @@ constexpr static OperatorRegistrationInformation operatorRegistrationInformation {REG_INFO( 9, Where, typeNameListWhere, supportedTypeListWhere, DmlGraphSupport::Supported)}, {REG_INFO( 7, ReduceSum, typeNameListDefault, supportedTypeListFloat16to32, DmlGraphSupport::Supported)}, {REG_INFO( 11, ReduceSum, typeNameListDefault, supportedTypeListFloat16to32, DmlGraphSupport::Supported)}, + {REG_INFO_VER( 12, Einsum, typeNameListDefault, supportedTypeListFloat16to32, DmlGraphSupport::Supported, requiredConstantCpuInputs(), std::nullopt, QueryEinSum )}, {REG_INFO( 7, ReduceMean, typeNameListDefault, supportedTypeListFloat16to32, DmlGraphSupport::Supported)}, {REG_INFO( 11, ReduceMean, typeNameListDefault, supportedTypeListFloat16to32, DmlGraphSupport::Supported)}, {REG_INFO( 7, ReduceProd, typeNameListDefault, supportedTypeListFloat16to32, DmlGraphSupport::Supported)}, @@ -600,7 +603,7 @@ void RegisterDmlOperators(IMLOperatorRegistry* registry) MLOperatorKernelDescription desc = {}; desc.domain = information.domain; desc.name = information.operatorName; - desc.executionType = MLOperatorExecutionType::D3D12; + desc.executionType = MLOperatorExecutionType::D3D12; // The graph must be configured with operators from only the legacy DML API, or only the new DML API bool kernelSupportsGraph = !bool(information.dmlGraphSupport & DmlGraphSupport::NotSupported); diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/TensorDesc.cpp b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/TensorDesc.cpp index e4fb01dfa9..2b352efb9a 100644 --- a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/TensorDesc.cpp +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/TensorDesc.cpp @@ -269,7 +269,7 @@ gsl::span TensorDesc::GetStrides() const { return {}; } - return { m_strides, m_strides + m_bufferTensorDesc.DimensionCount }; + return { m_strides, m_strides + m_bufferTensorDesc.DimensionCount }; } DML_TENSOR_DESC TensorDesc::GetDmlDesc() diff --git a/onnxruntime/core/providers/dml/OperatorAuthorHelper/Attributes.h b/onnxruntime/core/providers/dml/OperatorAuthorHelper/Attributes.h index deb194ff7e..95ce76a9bf 100644 --- a/onnxruntime/core/providers/dml/OperatorAuthorHelper/Attributes.h +++ b/onnxruntime/core/providers/dml/OperatorAuthorHelper/Attributes.h @@ -33,6 +33,7 @@ namespace AttrName static constexpr const char* Dtype = "dtype"; static constexpr const char* Ends = "ends"; static constexpr const char* Epsilon = "epsilon"; + static constexpr const char* Equation = "equation"; static constexpr const char* ExcludeOutside = "exclude_outside"; static constexpr const char* Exclusive = "exclusive"; static constexpr const char* Exponent = "exponent"; diff --git a/onnxruntime/core/providers/dml/OperatorAuthorHelper/OperatorHelper.cpp b/onnxruntime/core/providers/dml/OperatorAuthorHelper/OperatorHelper.cpp index c635407ee2..9698bbcbf0 100644 --- a/onnxruntime/core/providers/dml/OperatorAuthorHelper/OperatorHelper.cpp +++ b/onnxruntime/core/providers/dml/OperatorAuthorHelper/OperatorHelper.cpp @@ -858,7 +858,264 @@ namespace OperatorHelper std::iota(m_axes.begin(), m_axes.end(), 0); } } - + + void EinSumHelper::Initialize() + { + ParseEquationComponents(); + m_recognizedOperatorType = DetermineRecognizedOperatorType(); + } + + void EinSumHelper::ParseEquationComponents() + { + // Parse an equation like 'ij,jk->ik' into components {ij, jk, ik} mapping letters to + // numeric indices {(0,1}, {1,2}, {0,2}}. The last component is the output. + + std::map labelMap; + std::set repeatedLabels; + + uint32_t currentLabelIndex = 0; + Component currentComponent = {}; + bool foundOutput = false; + bool reachedEnd = false; + + // Read first to last character in equation, looking for letters, commas, and one arrow. + for (char* token = m_equation.data(); !reachedEnd; ++token) + { + char ch = *token; + + // Only ASCII letters are valid subscript symbols in numpy.einsum(). + if ((ch >= 'a' && ch <= 'z') || (ch >= 'A' && ch <= 'Z')) + { + // Check whether label already has an index. + const auto [i, inserted] = labelMap.insert({ch, currentLabelIndex}); + if (inserted) + { + ML_CHECK_VALID_ARGUMENT(!foundOutput, "Found label in equation output not matching any label from inputs.") + ++currentLabelIndex; // New label found. + } + else if (!foundOutput) + { + // If label in input already found earlier, then keep track of this later + // to generate the default output in case one is not specified. + repeatedLabels.insert(ch); + } + m_labelIndices.push_back(i->second); + } + else if (ch == ' ') + { + // Ignore spaces. + } + else + { + currentComponent.labelIndexEnd = static_cast(m_labelIndices.size()); + m_components.push_back(currentComponent); + currentComponent.labelIndexBegin = currentComponent.labelIndexEnd; + + switch (ch) + { + case ',': + // Note it's valid for 2 commas be adjacent, which indicates a scalar and generates + // an empty component. + break; + + case '-': // Start of "->" (must be atomic, no space between them). + ++token; // Skip '-'. + ML_CHECK_VALID_ARGUMENT(*token == '>', "Expected '->' for output.") + ML_CHECK_VALID_ARGUMENT(foundOutput == false, "Only one output arrow '->' is valid.") + foundOutput = true; + break; + + case '.': + // Ellipsis is unsupported. Leave recognized operator as None, deferring to another EP. + m_components.clear(); + return; + + case '\0': + reachedEnd = true; + break; // End of string. + + default: + ML_INVALID_ARGUMENT("Unsupported character in equation string. Must be a-z, A-Z, ',', or '->'."); + } + } + } + + if (!foundOutput) + { + // If no explicit output was given, generate an implicit output by ordering all the + // labels in alphabetic order (by ASCII value consistent with numpy, so Z < a). + // Exclude any labels that occurred more than once, as these cancel out. + + for (auto i : labelMap) + { + if (repeatedLabels.count(i.first) == 0) + { + m_labelIndices.push_back(i.second); + } + } + + // Push the final component, which is the output. + currentComponent.labelIndexEnd = static_cast(m_labelIndices.size()); + m_components.push_back(currentComponent); + } + } + + EinSumHelper::RecognizedOperatorType EinSumHelper::DetermineRecognizedOperatorType() + { + if (m_components.empty()) + { + return RecognizedOperatorType::None; // Parsing may have found unsupported components - treating as unknown. + } + + // std::ranges::equal is not supported yet. + auto equals = [](gsl::span a, gsl::span b) + { + return std::equal(a.begin(), a.end(), b.begin(), b.end()); + }; + + std::array componentRanks; + if (m_components.size() > componentRanks.size()) + { + // No recognized operator takes more than 2 inputs and 1 output. + // EinSum itself is generic and can handle any variable number of inputs, + // but DML's operators expect fixed counts. + return RecognizedOperatorType::None; + } + else if (m_components.size() == 2) + { + auto& inputLabels = m_components[0].GetLabels(m_labelIndices); + auto& outputLabels = m_components[1].GetLabels(m_labelIndices); + if (inputLabels.size() == outputLabels.size()) + { + // Check identity. + if (equals(inputLabels, outputLabels)) + { + // Handles: "->", "i->i", "ij->ij", "ijk->ijk", "ijkl->ijkl" ... + return RecognizedOperatorType::Identity; + } + else // Transpose since a permutation exists. + { + // Handles: "ij->ji", "ijk->kji", "ijkl->lkji", "ijkl->ijkl" ... + return RecognizedOperatorType::Transpose; + } + } + else if (outputLabels.empty()) // Scalar output, with all inputs reduced. + { + // Handles: "i->", "ij->", "ijk->", "ijkl->" ... + return RecognizedOperatorType::ReduceSum; + } + } + else if (m_components.size() == 3) + { + // If all components have the same size and label order, then apply elementwise multiplication. + auto& inputALabels = m_components[0].GetLabels(m_labelIndices); + auto& inputBLabels = m_components[1].GetLabels(m_labelIndices); + auto& outputLabels = m_components[2].GetLabels(m_labelIndices); + if (equals(inputALabels, outputLabels) && equals(inputBLabels, outputLabels)) + { + // Handles: "i,i->i", "ij,ij->ij", "ijk,ijk->ijk", "ijkl,ijkl->ijkl" ... + return RecognizedOperatorType::Multiply; + } + } + + // Otherwise check for special cases of dedicated operators... + + struct RecognizedOperatorInfo + { + RecognizedOperatorType recognizedOperatorType; + std::initializer_list componentRanks; + std::initializer_list labelIndices; + }; + + const RecognizedOperatorInfo recognizedOperators[] = { + {RecognizedOperatorType::MatMul, {2,2,2},{0,1, 1,2, 0,2}}, // ij,jk->ik + {RecognizedOperatorType::MatMul, {3,3,3},{0,1,2, 0,2,3, 0,1,3}}, // bij,bjk->bik + {RecognizedOperatorType::MatMul, {4,4,4},{0,1,2,3, 0,1,3,4, 0,1,2,4}}, // abij,abjk->abik + {RecognizedOperatorType::MatMulTransposeA, {2,2,2},{0,1, 0,2, 1,2}}, // ji,jk->ik + {RecognizedOperatorType::MatMulTransposeA, {3,3,3},{0,1,2, 0,1,3, 0,2,3}}, // bji,bjk->bik + {RecognizedOperatorType::MatMulTransposeA, {4,4,4},{0,1,2,3, 0,1,2,4, 0,1,3,4}}, // abji,abjk->abik + {RecognizedOperatorType::MatMulTransposeB, {2,2,2},{0,1, 2,1, 0,2}}, // ij,kj->ik + {RecognizedOperatorType::MatMulTransposeB, {3,3,3},{0,1,2, 0,3,2, 0,1,3}}, // bij,bkj->bik + {RecognizedOperatorType::MatMulTransposeB, {4,4,4},{0,1,2,3, 0,1,4,3, 0,1,2,4}}, // abij,abkj->abik + {RecognizedOperatorType::MatMulTransposeB, {1,1,0},{0,0,}}, // i,i-> (1D inner_prod) + {RecognizedOperatorType::ReduceSum, {2,1 },{0,1, 0}}, // ij->i + {RecognizedOperatorType::ReduceSum, {2,1 },{0,1, 1}}, // ij->j + }; + + // For each recognized operator, compare the labels-per-component and label indices. + for (auto& recognizedOperator : recognizedOperators) + { + if (equals(m_labelIndices, recognizedOperator.labelIndices) + && m_components.size() == recognizedOperator.componentRanks.size()) + { + for (size_t i = 0; i < m_components.size(); ++i) + { + componentRanks[i] = m_components[i].GetDimensionCount(); + } + + if (equals(gsl::make_span(componentRanks.data(), m_components.size()), recognizedOperator.componentRanks)) + { + return recognizedOperator.recognizedOperatorType; + } + } + } + + return RecognizedOperatorType::None; + } + + std::vector EinSumHelper::GetOutputShapes(const MLShapeInferenceContext& shapeInfo) const + { + assert(!m_components.empty()); // Should have already parsed components. + + uint32_t inputCount = shapeInfo.GetInputCount(); + uint32_t outputCount = shapeInfo.GetOutputCount(); + ML_CHECK_VALID_ARGUMENT(inputCount + 1 == m_components.size(), "Mismatch between input tensor count and string equation component count."); + ML_CHECK_VALID_ARGUMENT(outputCount == 1, "EinSum expects exactly 1 output tensor."); + + std::vector labelSizes(m_labelIndices.size(), INT_MIN); + + // Read every input tensor, comparing labels to ensure consistent sizes from the equation parsed earlier. + for (uint32_t i = 0; i < inputCount; ++i) + { + auto inputShape = shapeInfo.GetInputTensorShape(i); + auto& component = m_components[i]; + auto labelIndices = component.GetLabels(m_labelIndices); + uint32_t dimensionCount = component.GetDimensionCount(); + + ML_CHECK_VALID_ARGUMENT(inputShape.size() == dimensionCount, "Mismatch between input tensor shape and string equation label count."); + + for (uint32_t i = 0; i < dimensionCount; ++i) + { + // If this is the first time seeing this label, then record the size. + // Otherwise any following occurrences of the label must match sizes. + // e.g. Given "ij,ji", both i's and both j's must match dimension sizes. + uint32_t dimensionSize = inputShape[i]; + uint32_t labelIndex = labelIndices[i]; + assert(labelIndex < labelSizes.size()); + + if (labelSizes[labelIndex] == INT_MIN) + { + labelSizes[labelIndex] = dimensionSize; + } + else + { + ML_CHECK_VALID_ARGUMENT(labelSizes[labelIndex] == dimensionSize, "All labels must have the same dimension sizes."); + } + } + } + + // Generate output dimensions from corresponding input tensor labels. + // e.g. Given ij,jk->ij with [2,3] and [3,5], the output is [2,5]. + std::vector outputDimensions; + auto outputLabelIndices = m_components.back().GetLabels(m_labelIndices); + for (auto labelIndex : outputLabelIndices) + { + outputDimensions.push_back(labelSizes[labelIndex]); + } + + return { std::move(EdgeShapes(outputDimensions)) }; + } + std::vector MatMulHelperBase::GetOutputShapes(const MLShapeInferenceContext& shapeInfo) const { ML_CHECK_VALID_ARGUMENT(shapeInfo.GetInputCount() >= 2); diff --git a/onnxruntime/core/providers/dml/OperatorAuthorHelper/OperatorHelper.h b/onnxruntime/core/providers/dml/OperatorAuthorHelper/OperatorHelper.h index a32da47dbc..3bd8064dc6 100644 --- a/onnxruntime/core/providers/dml/OperatorAuthorHelper/OperatorHelper.h +++ b/onnxruntime/core/providers/dml/OperatorAuthorHelper/OperatorHelper.h @@ -771,6 +771,70 @@ class ReduceHelper : public ReduceHelperBase { ReduceHelper(const Info_t& info, const Shape_t& shape) : ReduceHelperBase(info, shape, true) {} }; +class EinSumHelper +{ +public: + void Initialize(); + + // Info_t is used to obtain attributes which will be used for calculating the output shape later. + // Shape_t is used to obtain input shape which will be used for adjusting attribute value. + template + EinSumHelper(const Info_t& info, const Shape_t& shape, uint32_t opsetVersion) + { + m_equation = info.GetAttribute(AttrName::Equation); + Initialize(); + } + + EinSumHelper(const MLOperatorAttributes& info) + { + m_equation = info.GetAttribute(AttrName::Equation); + Initialize(); + } + + std::vector GetOutputShapes(const MLShapeInferenceContext& shapeInfo) const; + + enum class RecognizedOperatorType + { + None, + Identity, + Multiply, + MatMul, + MatMulTransposeA, + MatMulTransposeB, + ReduceSum, + Transpose, + Total, + }; + + RecognizedOperatorType GetRecognizedOperatorType() const noexcept { return m_recognizedOperatorType; } + +protected: + void ParseEquationComponents(); + RecognizedOperatorType DetermineRecognizedOperatorType(); + +protected: + struct Component + { + uint32_t labelIndexBegin; + uint32_t labelIndexEnd; + + uint32_t GetDimensionCount() const noexcept + { + return labelIndexEnd - labelIndexBegin; + } + gsl::span GetLabels(gsl::span labels) const + { + return labels.subspan(labelIndexBegin, labelIndexEnd - labelIndexBegin); + }; + }; + + std::string m_equation; + std::vector m_labelIndices; // Concatenation of all labels as rebased indices ("ij,ai" -> 0,1,2,0). + std::vector m_components; // All components in order, including inputs and output. + std::vector m_outputDimensions; + RecognizedOperatorType m_recognizedOperatorType = RecognizedOperatorType::None; +}; + class MatMulHelperBase { public: // Info_t is used to obtain attributes which will be used for calculating the output shape later. @@ -1465,6 +1529,7 @@ using ShapeInferenceHelper_ReduceL1 = ReduceHelper; using ShapeInferenceHelper_ReduceL2 = ReduceHelper; using ShapeInferenceHelper_ReduceMax = ReduceHelper; using ShapeInferenceHelper_ReduceMin = ReduceHelper; +using ShapeInferenceHelper_Einsum12 = VersionedOpsetHelper; using ShapeInferenceHelper_ArgMax = ArgMinArgMaxHelper; using ShapeInferenceHelper_ArgMin = ArgMinArgMaxHelper; using ShapeInferenceHelper_Gemm = GemmHelper; diff --git a/onnxruntime/core/providers/dml/OperatorAuthorHelper/OperatorRegistration.h b/onnxruntime/core/providers/dml/OperatorAuthorHelper/OperatorRegistration.h index c8784a79d8..9d5c248559 100644 --- a/onnxruntime/core/providers/dml/OperatorAuthorHelper/OperatorRegistration.h +++ b/onnxruntime/core/providers/dml/OperatorAuthorHelper/OperatorRegistration.h @@ -247,19 +247,20 @@ namespace OperatorHelper namespace OnnxOperatorSet12 { - static const int sc_sinceVer_GreaterOrEqual = 12; - static const int sc_sinceVer_LessOrEqual = 12; + static const int sc_sinceVer_ArgMin = 12; + static const int sc_sinceVer_ArgMax = 12; static const int sc_sinceVer_Celu = 12; static const int sc_sinceVer_Clip = 12; + static const int sc_sinceVer_Einsum = 12; static const int sc_sinceVer_GatherND = 12; + static const int sc_sinceVer_GreaterOrEqual = 12; + static const int sc_sinceVer_LessOrEqual = 12; + static const int sc_sinceVer_MaxPool = 12; static const int sc_sinceVer_Min = 12; static const int sc_sinceVer_Max = 12; static const int sc_sinceVer_Pow = 12; - static const int sc_sinceVer_MaxPool = 12; static const int sc_sinceVer_ReduceMax = 12; static const int sc_sinceVer_ReduceMin = 12; - static const int sc_sinceVer_ArgMin = 12; - static const int sc_sinceVer_ArgMax = 12; } // namespace OnnxOperatorSet12 namespace MsftOperatorSet1 diff --git a/onnxruntime/core/providers/dml/OperatorAuthorHelper/precomp.h b/onnxruntime/core/providers/dml/OperatorAuthorHelper/precomp.h index 79e2cdba4e..6c47e60e63 100644 --- a/onnxruntime/core/providers/dml/OperatorAuthorHelper/precomp.h +++ b/onnxruntime/core/providers/dml/OperatorAuthorHelper/precomp.h @@ -7,6 +7,8 @@ #include #include #include +#include +#include #include #include