DirectML EP add CastLike-15 and Reshape-14 (#11568)

* Add CastLike15
* Add Reshape14
* Fix allowzero comment
* Rename REG_INFO_ID to REG_INFO_COPY to be clearer to readers
This commit is contained in:
Dwayne Robinson 2022-05-19 14:34:11 -07:00 committed by GitHub
parent 3a867d83d5
commit d2519ec0c2
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
6 changed files with 68 additions and 41 deletions

View file

@ -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<std::optional<uint32_t>> inputIndices = { 0 }; // For CastLike, the second tensor ('target_type') is not bound.
std::vector<std::optional<uint32_t>> outputIndices = { 0 };
DmlOperator::Initialize(kernelInfo, inputIndices, outputIndices);
std::vector<DML_TENSOR_DESC> inputDescs = GetDmlInputDescs();
std::vector<DML_TENSOR_DESC> 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

View file

@ -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<ShapeInferenceHelper_##operatorName##version>, 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)},

View file

@ -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";

View file

@ -1971,32 +1971,44 @@ namespace OperatorHelper
int inferDim = -1;
DimensionType inElementCount = ComputeElementCountFromDimensions(inputDimensions);
bool allowZero = shapeInfo.template GetOptionalAttribute<int32_t>(AttrName::AllowZero, 0);
for (int i = 0, ci = gsl::narrow_cast<int>(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<int>(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<int>(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) };

View file

@ -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<ResizeHelper, 10>;
@ -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;

View file

@ -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;