diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlOperatorCast.cpp b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlOperatorCast.cpp index 311da999cc..76b9b308fe 100644 --- a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlOperatorCast.cpp +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlOperatorCast.cpp @@ -15,7 +15,11 @@ public: const MLOperatorKernelCreationContext& kernelInfo ) : DmlOperator(kernelInfo) { - Initialize(kernelInfo); + ML_CHECK_VALID_ARGUMENT(kernelInfo.GetInputCount() >= 1); + ML_CHECK_VALID_ARGUMENT(kernelInfo.GetOutputCount() == 1); + std::vector> inputIndices = { 0 }; // For CastLike, the second tensor ('target_type') is not bound. + std::vector> outputIndices = { 0 }; + DmlOperator::Initialize(kernelInfo, inputIndices, outputIndices); std::vector inputDescs = GetDmlInputDescs(); std::vector outputDescs = GetDmlOutputDescs(); @@ -38,10 +42,12 @@ public: m_compiledOperator.Get(), m_persistentResourceBinding ? &*m_persistentResourceBinding : nullptr, gsl::make_span(inputTensors), - gsl::make_span(outputTensors))); + gsl::make_span(outputTensors) + )); } }; DML_OP_DEFINE_CREATION_FUNCTION(Cast, DmlOperatorCast); +DML_OP_DEFINE_CREATION_FUNCTION(CastLike15, DmlOperatorCast); } // 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 700b6398c5..d39f177d37 100644 --- a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/OperatorRegistration.cpp +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/OperatorRegistration.cpp @@ -188,6 +188,7 @@ DML_OP_EXTERN_CREATION_FUNCTION(Affine); DML_OP_EXTERN_CREATION_FUNCTION(Dropout); DML_OP_EXTERN_CREATION_FUNCTION(MatMul); DML_OP_EXTERN_CREATION_FUNCTION(Cast); +DML_OP_EXTERN_CREATION_FUNCTION(CastLike15); DML_OP_EXTERN_CREATION_FUNCTION(MemcpyFromHost); DML_OP_EXTERN_CREATION_FUNCTION(MemcpyToHost); DML_OP_EXTERN_CREATION_FUNCTION(TopK7); @@ -349,7 +350,7 @@ constexpr auto requiredConstantCpuInputs(Args... args) // Identity operators use Copy, alias their first input, and use elementwise identity operators // when needed for striding support, but issue actual copies outside the graph. -#define REG_INFO_ID(version, operatorName, ...) \ +#define REG_INFO_COPY(version, operatorName, ...) \ #operatorName, OnnxOperatorSet##version::sc_sinceVer_##operatorName, onnxruntime::kOnnxDomain, CreateCopy, ShapeInferenceFunction, true, ##__VA_ARGS__, // MS-domain operators @@ -457,23 +458,22 @@ constexpr static OperatorRegistrationInformation operatorRegistrationInformation // {REG_INFO( 14, Trilu, typeNameListTrilu, supportedTypeListScalars8to32, DmlGraphSupport::Supported)}, // Data reorganization that merely changes the dimensions while keeping the data identical. - {REG_INFO_ID( 7, Identity, typeNameListDefault, supportedTypeListAllScalars, DmlGraphSupport::Supported)}, - {REG_INFO_ID( 13, Identity, typeNameListDefault, supportedTypeListAllScalars, DmlGraphSupport::Supported)}, - {REG_INFO_ID( 14, Identity, typeNameListDefault, supportedTypeListAllScalars, DmlGraphSupport::Supported)}, - {REG_INFO_ID( 7, Flatten, typeNameListDefault, supportedTypeListAllScalars, DmlGraphSupport::Supported)}, - {REG_INFO_ID( 9, Flatten, typeNameListDefault, supportedTypeListAllScalars, DmlGraphSupport::Supported)}, - {REG_INFO_ID( 11, Flatten, typeNameListDefault, supportedTypeListAllScalars, DmlGraphSupport::Supported)}, - {REG_INFO_ID( 13, Flatten, typeNameListDefault, supportedTypeListAllScalars, DmlGraphSupport::Supported)}, - {REG_INFO_ID( 7, Squeeze, typeNameListDefault, supportedTypeListAllScalars, DmlGraphSupport::Supported)}, - {REG_INFO_ID( 11, Squeeze, typeNameListDefault, supportedTypeListAllScalars, DmlGraphSupport::Supported)}, - {REG_INFO_ID( 13, Squeeze, typeNameListDefault, supportedTypeListAllScalars, DmlGraphSupport::Supported, requiredConstantCpuInputs(1))}, - {REG_INFO_ID( 7, Unsqueeze, typeNameListDefault, supportedTypeListAllScalars, DmlGraphSupport::Supported)}, - {REG_INFO_ID( 11, Unsqueeze, typeNameListDefault, supportedTypeListAllScalars, DmlGraphSupport::Supported)}, - {REG_INFO_ID( 13, Unsqueeze, typeNameListDefault, supportedTypeListAllScalars, DmlGraphSupport::Supported, requiredConstantCpuInputs(1))}, - {REG_INFO_ID( 7, Reshape, typeNameListDefault, supportedTypeListAllScalars, DmlGraphSupport::Supported, requiredConstantCpuInputs(1))}, - {REG_INFO_ID( 13, Reshape, typeNameListDefault, supportedTypeListAllScalars, DmlGraphSupport::Supported, requiredConstantCpuInputs(1))}, - // TODO: Add allowzero attribute. - // {REG_INFO_ID(14, Reshape, typeNameListDefault, supportedTypeListAllScalars, DmlGraphSupport::Supported, requiredConstantCpuInputs(1))}, + {REG_INFO_COPY( 7, Identity, typeNameListDefault, supportedTypeListAllScalars, DmlGraphSupport::Supported)}, + {REG_INFO_COPY(13, Identity, typeNameListDefault, supportedTypeListAllScalars, DmlGraphSupport::Supported)}, + {REG_INFO_COPY(14, Identity, typeNameListDefault, supportedTypeListAllScalars, DmlGraphSupport::Supported)}, + {REG_INFO_COPY( 7, Flatten, typeNameListDefault, supportedTypeListAllScalars, DmlGraphSupport::Supported)}, + {REG_INFO_COPY( 9, Flatten, typeNameListDefault, supportedTypeListAllScalars, DmlGraphSupport::Supported)}, + {REG_INFO_COPY(11, Flatten, typeNameListDefault, supportedTypeListAllScalars, DmlGraphSupport::Supported)}, + {REG_INFO_COPY(13, Flatten, typeNameListDefault, supportedTypeListAllScalars, DmlGraphSupport::Supported)}, + {REG_INFO_COPY( 7, Squeeze, typeNameListDefault, supportedTypeListAllScalars, DmlGraphSupport::Supported)}, + {REG_INFO_COPY(11, Squeeze, typeNameListDefault, supportedTypeListAllScalars, DmlGraphSupport::Supported)}, + {REG_INFO_COPY(13, Squeeze, typeNameListDefault, supportedTypeListAllScalars, DmlGraphSupport::Supported, requiredConstantCpuInputs(1))}, + {REG_INFO_COPY( 7, Unsqueeze, typeNameListDefault, supportedTypeListAllScalars, DmlGraphSupport::Supported)}, + {REG_INFO_COPY(11, Unsqueeze, typeNameListDefault, supportedTypeListAllScalars, DmlGraphSupport::Supported)}, + {REG_INFO_COPY(13, Unsqueeze, typeNameListDefault, supportedTypeListAllScalars, DmlGraphSupport::Supported, requiredConstantCpuInputs(1))}, + {REG_INFO_COPY( 7, Reshape, typeNameListDefault, supportedTypeListAllScalars, DmlGraphSupport::Supported, requiredConstantCpuInputs(1))}, + {REG_INFO_COPY(13, Reshape, typeNameListDefault, supportedTypeListAllScalars, DmlGraphSupport::Supported, requiredConstantCpuInputs(1))}, + {REG_INFO_COPY(14, Reshape, typeNameListDefault, supportedTypeListAllScalars, DmlGraphSupport::Supported, requiredConstantCpuInputs(1))}, // Elementwise {REG_INFO( 7, Sqrt, typeNameListDefault, supportedTypeListFloat16to32, DmlGraphSupport::Supported)}, @@ -666,8 +666,7 @@ constexpr static OperatorRegistrationInformation operatorRegistrationInformation {REG_INFO( 7, Cast, typeNameListTwo, supportedTypeListCast, DmlGraphSupport::Supported)}, {REG_INFO( 9, Cast, typeNameListTwo, supportedTypeListCast, DmlGraphSupport::Supported)}, {REG_INFO( 13, Cast, typeNameListTwo, supportedTypeListCast, DmlGraphSupport::Supported)}, - // TODO: Add CastLike-15 https://github.com/onnx/onnx/pull/3558 - // {REG_INFO( 15, CastLike, typeNameListTwo, supportedTypeListCast, DmlGraphSupport::Supported)}, + {REG_INFO_VER( 15, CastLike, typeNameListTwo, supportedTypeListCast, DmlGraphSupport::Supported)}, {REG_INFO( 7, MemcpyFromHost, typeNameListDefault, supportedTypeListAll)}, {REG_INFO( 7, MemcpyToHost, typeNameListDefault, supportedTypeListAll)}, {REG_INFO_VER( 7, TopK, typeNameListTopK, supportedTypeListTopK, DmlGraphSupport::Supported)}, diff --git a/onnxruntime/core/providers/dml/OperatorAuthorHelper/Attributes.h b/onnxruntime/core/providers/dml/OperatorAuthorHelper/Attributes.h index 71392bf155..99e7d1d390 100644 --- a/onnxruntime/core/providers/dml/OperatorAuthorHelper/Attributes.h +++ b/onnxruntime/core/providers/dml/OperatorAuthorHelper/Attributes.h @@ -9,6 +9,7 @@ namespace AttrName static constexpr const char* ActivationAlpha = "activation_alpha"; static constexpr const char* ActivationBeta = "activation_beta"; static constexpr const char* Activations = "activations"; + static constexpr const char* AllowZero = "allowzero"; static constexpr const char* Alpha = "alpha"; static constexpr const char* AutoPad = "auto_pad"; static constexpr const char* Axes = "axes"; diff --git a/onnxruntime/core/providers/dml/OperatorAuthorHelper/OperatorHelper.cpp b/onnxruntime/core/providers/dml/OperatorAuthorHelper/OperatorHelper.cpp index efe5d3f4a1..268e12706d 100644 --- a/onnxruntime/core/providers/dml/OperatorAuthorHelper/OperatorHelper.cpp +++ b/onnxruntime/core/providers/dml/OperatorAuthorHelper/OperatorHelper.cpp @@ -1971,32 +1971,44 @@ namespace OperatorHelper int inferDim = -1; DimensionType inElementCount = ComputeElementCountFromDimensions(inputDimensions); + bool allowZero = shapeInfo.template GetOptionalAttribute(AttrName::AllowZero, 0); - for (int i = 0, ci = gsl::narrow_cast(m_shapeDims.size()); i < ci; ++i) + if (allowZero) { - switch (m_shapeDims[i]) + // Just take the shape directly (no special handling for 0). + for (int i = 0, ci = gsl::narrow_cast(m_shapeDims.size()); i < ci; ++i) { - case -1: - ML_CHECK_VALID_ARGUMENT(inferDim == -1, "Only one dimension can be inferred."); - inferDim = i; - break; - - case 0: - outputDimensions[i] = inputDimensions[i]; - outElementCount *= outputDimensions[i]; - break; - - default: outputDimensions[i] = m_shapeDims[i]; - outElementCount *= outputDimensions[i]; - break; } } - - if (inferDim != -1) + else { - outputDimensions[inferDim] = inElementCount / outElementCount; - outElementCount *= outputDimensions[inferDim]; + // Special handling where 0 size means to copy the corresponding input tensor dimension. + for (int i = 0, ci = gsl::narrow_cast(m_shapeDims.size()); i < ci; ++i) + { + switch (m_shapeDims[i]) + { + case -1: + ML_CHECK_VALID_ARGUMENT(inferDim == -1, "Only one dimension can be inferred."); + inferDim = i; + break; + + case 0: + outputDimensions[i] = inputDimensions[i]; + outElementCount *= outputDimensions[i]; + break; + + default: + outputDimensions[i] = m_shapeDims[i]; + outElementCount *= outputDimensions[i]; + break; + } + } + + if (inferDim != -1) + { + outputDimensions[inferDim] = inElementCount / outElementCount; + } } return { EdgeShapes(outputDimensions) }; diff --git a/onnxruntime/core/providers/dml/OperatorAuthorHelper/OperatorHelper.h b/onnxruntime/core/providers/dml/OperatorAuthorHelper/OperatorHelper.h index 609a0c33ec..188eb3dff0 100644 --- a/onnxruntime/core/providers/dml/OperatorAuthorHelper/OperatorHelper.h +++ b/onnxruntime/core/providers/dml/OperatorAuthorHelper/OperatorHelper.h @@ -1365,6 +1365,7 @@ using ShapeInferenceHelper_EyeLike = GetOutputShapeAsInputShapeHelper; using ShapeInferenceHelper_Expand = ExpandHelper; using ShapeInferenceHelper_Reshape7 = ReshapeHelper; using ShapeInferenceHelper_Reshape13 = ReshapeHelper; +using ShapeInferenceHelper_Reshape14 = ReshapeHelper; using ShapeInferenceHelper_ConstantOfShape = ConstantOfShapeHelper; using ShapeInferenceHelper_Tile = TileHelper; using ShapeInferenceHelper_Resize10 = VersionedOpsetHelper; @@ -1494,6 +1495,8 @@ using ShapeInferenceHelper_CumSum11 = GetOutputShapeAsInputShapeHelper; using ShapeInferenceHelper_CumSum14 = GetOutputShapeAsInputShapeHelper; using ShapeInferenceHelper_Range = RangeHelper; +using ShapeInferenceHelper_CastLike15 = GetOutputShapeAsInputShapeHelper; + using ShapeInferenceHelper_FusedConv = ConvHelper; using ShapeInferenceHelper_FusedConvTranspose = ConvTransposeHelper; using ShapeInferenceHelper_FusedInstanceNormalization = GetOutputShapeAsInputShapeHelper; diff --git a/onnxruntime/core/providers/dml/OperatorAuthorHelper/OperatorVersions.h b/onnxruntime/core/providers/dml/OperatorAuthorHelper/OperatorVersions.h index e01fbb1ea4..ec6e4dca67 100644 --- a/onnxruntime/core/providers/dml/OperatorAuthorHelper/OperatorVersions.h +++ b/onnxruntime/core/providers/dml/OperatorAuthorHelper/OperatorVersions.h @@ -346,9 +346,15 @@ namespace OperatorHelper static const int sc_sinceVer_Identity = 14; static const int sc_sinceVer_Mul = 14; static const int sc_sinceVer_Relu = 14; + static const int sc_sinceVer_Reshape = 14; static const int sc_sinceVer_Sub = 14; } // namespace OnnxOperatorSet14 + namespace OnnxOperatorSet15 + { + static const int sc_sinceVer_CastLike = 15; + } // namespace OnnxOperatorSet14 + namespace MsftOperatorSet1 { static const int sc_sinceVer_FusedConv = 1;