saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-29 20:14:01 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

[DML EP] Register DFT-20 (#20341)

Browse source 
### Description
<!-- Describe your changes. -->

This PR registers DFT-20 to the DML EP.

### Motivation and Context
<!-- - Why is this change required? What problem does it solve?
- If it fixes an open issue, please link to the issue here. -->
This commit is contained in:
aamajumder 2024-05-02 11:08:39 -07:00 committed by GitHub
parent 0dda8b0c44
commit 589aeb7036
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 83 additions and 25 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

3
docs/OperatorKernels.md
View file

@ -984,7 +984,8 @@ Do not modify directly.*
|Crop|*in* input:**T**<br> *out* output:**T**|1+|**T** = tensor(float), tensor(float16)|
|CumSum|*in* x:**T**<br> *in* axis:**T2**<br> *out* y:**T**|14+|**T** = tensor(float), tensor(float16), tensor(int32), tensor(int64), tensor(uint32), tensor(uint64)|
|||11+|**T** = tensor(float), tensor(float16), tensor(int32), tensor(int64), tensor(uint32), tensor(uint64)|
|DFT|*in* input:**T1**<br> *in* dft_length:**T2**<br> *in* axis:**tensor(int64)**<br> *out* output:**T1**<br><br>or<br><br>*in* input:**T1**<br> *in* dft_length:**T2**<br> *out* output:**T1**|17+|**T1** = tensor(float), tensor(float16)<br/> **T2** = tensor(int64)|
|DFT|*in* input:**T1**<br> *in* dft_length:**T2**<br> *in* axis:**tensor(int64)**<br> *out* output:**T1**<br><br>or<br><br>*in* input:**T1**<br> *in* dft_length:**T2**<br> *out* output:**T1**|20+|**T1** = tensor(double), tensor(float), tensor(float16)<br/> **T2** = tensor(int32), tensor(int64)|
|||17+|**T1** = tensor(double), tensor(float), tensor(float16)<br/> **T2** = tensor(int32), tensor(int64)|
|DepthToSpace|*in* input:**T**<br> *out* output:**T**|13+|**T** = tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)|
|||11+|**T** = tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)|
|||1+|**T** = tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)|

90
onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlDFT.h
View file

@ -208,7 +208,7 @@ public:
PrepareBluesteinZChirp(dataType);
}
GpuDFTOperator(IMLOperatorKernelCreationContext* context)
GpuDFTOperator(IMLOperatorKernelCreationContext* context, int32_t version)
{
ComPtr<IUnknown> executionObject;
context->GetExecutionInterface(executionObject.GetAddressOf());
@ -218,8 +218,6 @@ public:
ORT_THROW_IF_FAILED(commandList->GetDevice(IID_ID3D12Device, &m_device));
ORT_THROW_IF_FAILED(context->GetAttribute("axis", MLOperatorAttributeType::Int, 1, sizeof(int64_t), reinterpret_cast<void*>(&m_axis)));
int64_t isInverseInt;
ORT_THROW_IF_FAILED(context->GetAttribute("inverse", MLOperatorAttributeType::Int, 1, sizeof(int64_t), reinterpret_cast<void*>(&isInverseInt)));
m_isInverse = static_cast<bool>(isInverseInt);
@ -232,6 +230,15 @@ public:
ORT_THROW_IF_FAILED(context->GetInputEdgeDescription(0, &edgeDesc));
assert(edgeDesc.edgeType == MLOperatorEdgeType::Tensor);
if (version == 17)
{
ORT_THROW_IF_FAILED(context->GetAttribute("axis", MLOperatorAttributeType::Int, 1, sizeof(int64_t), reinterpret_cast<void*>(&m_axis)));
}
else
{
m_axis = -2; //-2 is the default axis value for DFT-20 if the optional axis input is not provided
}
PrepareStockhamFFT(edgeDesc.tensorDataType);
PrepareBluesteinZChirp(edgeDesc.tensorDataType);
}
@ -398,7 +405,15 @@ public:
// Get the input and output shape sizes
auto inputDims = GetTensorDimensions(inputTensor.Get());
ML_CHECK_VALID_ARGUMENT(static_cast<size_t>(m_axis) < inputDims.size())
auto rank = static_cast<int32_t>(inputDims.size());
ComPtr<IMLOperatorTensor> axisTensor;
if (SUCCEEDED(context->GetInputTensor(2, &axisTensor)) && axisTensor != nullptr)
{
MLOperatorTensor tensor(axisTensor.Get());
m_axis = OperatorHelper::ReadScalarTensorCastToInt64(tensor);
}
m_axis = OperatorHelper::HandleNegativeAxis(static_cast<int32_t>(m_axis), rank);
ML_CHECK_VALID_ARGUMENT(m_axis >= 0 && m_axis < rank);
auto outputDims = GetTensorDimensions(outputTensor.Get());
ORT_THROW_HR_IF(E_FAIL, inputDims.size() != outputDims.size());
@ -977,8 +992,6 @@ struct DFTShapeInferrer : public WRL::Base<IMLOperatorShapeInferrer>
{
try
{
int64_t axis;
ORT_THROW_IF_FAILED(context->GetAttribute("axis", MLOperatorAttributeType::Int, 1, sizeof(int64_t), reinterpret_cast<void*>(&axis)));
int64_t isInverseInt;
ORT_THROW_IF_FAILED(context->GetAttribute("inverse", MLOperatorAttributeType::Int, 1, sizeof(int64_t), reinterpret_cast<void*>(&isInverseInt)));
int64_t isOnesidedInt;
@ -999,7 +1012,25 @@ struct DFTShapeInferrer : public WRL::Base<IMLOperatorShapeInferrer>
throw;
}
auto axisIdx = OperatorHelper::HandleNegativeAxis(static_cast<int32_t>(axis), rank);
int64_t axisValue = -2; //Set default axis value to -2 for DFT-20 (last signal axis)
bool isDft17 = !context->IsInputValid(2); //Check if axis is provided as an input - if yes then it is DFT-20
if (isDft17)
{
axisValue = 1; //Default axis value for DFT-17 should be 1 if axis attribute is not provided
ORT_THROW_IF_FAILED(context->GetAttribute("axis", MLOperatorAttributeType::Int, 1, sizeof(int64_t), reinterpret_cast<void*>(&axisValue)));
}
else
{
ComPtr<IMLOperatorShapeInferenceContextPrivate> contextPrivate;
ORT_THROW_IF_FAILED(context->QueryInterface(IID_PPV_ARGS(&contextPrivate)));
ComPtr<IMLOperatorTensor> axisTensor;
ORT_THROW_IF_FAILED(contextPrivate->GetConstantInputTensor(2, &axisTensor));
MLOperatorTensor tensor(axisTensor.Get());
axisValue = OperatorHelper::ReadScalarTensorCastToInt64(tensor);
}
auto axisIdx = OperatorHelper::HandleNegativeAxis(static_cast<int32_t>(axisValue), rank);
// In general the output shape will match the input shape exactly
// So initialize the output shape with the input shape
@ -1010,6 +1041,8 @@ struct DFTShapeInferrer : public WRL::Base<IMLOperatorShapeInferrer>
// It corresponds to the real and imaginary parts of the DFT output.
outputDims.back() = 2;
auto dftLength = inputDims[axisIdx];
if (context->IsInputValid(1))
{
// If dft_length is specified, then we should honor the shape.
@ -1019,10 +1052,11 @@ struct DFTShapeInferrer : public WRL::Base<IMLOperatorShapeInferrer>
ComPtr<IMLOperatorTensor> dftLengthTensor;
ORT_THROW_IF_FAILED(contextPrivate->GetConstantInputTensor(1, &dftLengthTensor));
MLOperatorTensor tensor(dftLengthTensor.Get());
auto dftLength = onnxruntime::narrow<uint32_t>(OperatorHelper::ReadScalarTensorCastToInt64(tensor));
outputDims[axisIdx] = dftLength;
dftLength = onnxruntime::narrow<uint32_t>(OperatorHelper::ReadScalarTensorCastToInt64(tensor));
}
outputDims[axisIdx] = dftLength;
// When DFT is onesided, the output shape is half the size of the input shape
// along the specified axis.
if (isOnesided)
@ -1056,7 +1090,14 @@ public:
{
try
{
auto dftOperator = wil::MakeOrThrow<GpuDFTOperator>(context);
//If the axis value is provided as an input use DFT-20, otherwise use DFT-17
int32_t version = 17;
if (context->IsInputValid(2))
{
version = 20;
}
auto dftOperator = wil::MakeOrThrow<GpuDFTOperator>(context, version);
dftOperator.CopyTo(kernel);
return S_OK;
}
@ -1066,12 +1107,12 @@ public:
}
}
static void RegisterDFTKernel(IMLOperatorRegistry* registry)
static void RegisterDFTKernel(IMLOperatorRegistry* registry, int32_t version)
{
MLOperatorKernelDescription kernelDescription = {};
kernelDescription.domain = "";
kernelDescription.name = "DFT";
kernelDescription.minimumOperatorSetVersion = 17;
kernelDescription.minimumOperatorSetVersion = version;
kernelDescription.executionType = MLOperatorExecutionType::D3D12;
// T1: tensor(float16), tensor(float), tensor(double), tensor(bfloat16)
@ -1081,7 +1122,7 @@ public:
{
MLOperatorEdgeDescription { MLOperatorEdgeType::Tensor, (uint64_t)MLOperatorTensorDataType::Float16 },
MLOperatorEdgeDescription { MLOperatorEdgeType::Tensor, (uint64_t)MLOperatorTensorDataType::Float },
//MLOperatorEdgeDescription { MLOperatorEdgeType::Tensor, (uint64_t)MLOperatorTensorDataType::Double },
MLOperatorEdgeDescription { MLOperatorEdgeType::Tensor, (uint64_t)MLOperatorTensorDataType::Double },
};
t1Constraint.allowedTypes = t1AllowedEdges.data();
t1Constraint.allowedTypeCount = static_cast<uint32_t>(t1AllowedEdges.size());
@ -1091,7 +1132,7 @@ public:
t2Constraint.typeLabel = "T2";
std::vector<MLOperatorEdgeDescription> t2AllowedEdges
{
// MLOperatorEdgeDescription { MLOperatorEdgeType::Tensor, (uint64_t)MLOperatorTensorDataType::Int32 },
MLOperatorEdgeDescription { MLOperatorEdgeType::Tensor, (uint64_t)MLOperatorTensorDataType::Int32 },
MLOperatorEdgeDescription { MLOperatorEdgeType::Tensor, (uint64_t)MLOperatorTensorDataType::Int64 },
};
t2Constraint.allowedTypes = t2AllowedEdges.data();
@ -1101,13 +1142,6 @@ public:
kernelDescription.typeConstraints = typeConstraints.data();
kernelDescription.typeConstraintCount = static_cast<uint32_t>(typeConstraints.size());
MLOperatorAttributeNameValue axisAttributeValue;
axisAttributeValue.name = "axis";
axisAttributeValue.type = MLOperatorAttributeType::Int;
axisAttributeValue.valueCount = 1;
static const int64_t axis[] = { 1 };
axisAttributeValue.ints = axis;
MLOperatorAttributeNameValue inverseAttributeValue;
inverseAttributeValue.name = "inverse";
inverseAttributeValue.type = MLOperatorAttributeType::Int;
@ -1123,11 +1157,21 @@ public:
onesidedAttributeValue.ints = onesided;
std::vector<MLOperatorAttributeNameValue> attributeDefaultValues{
axisAttributeValue,
inverseAttributeValue,
onesidedAttributeValue
};
if (version == 17)
{
MLOperatorAttributeNameValue axisAttributeValue;
axisAttributeValue.name = "axis";
axisAttributeValue.type = MLOperatorAttributeType::Int;
axisAttributeValue.valueCount = 1;
static const int64_t axis[] = { 1 };
axisAttributeValue.ints = axis;
attributeDefaultValues.push_back(axisAttributeValue);
}
kernelDescription.defaultAttributes = attributeDefaultValues.data();
kernelDescription.defaultAttributeCount = static_cast<uint32_t>(attributeDefaultValues.size());
kernelDescription.options = MLOperatorKernelOptions::None;
@ -1136,7 +1180,7 @@ public:
auto shareInferrer = wil::MakeOrThrow<DFTShapeInferrer>();
auto factory = wil::MakeOrThrow<GpuDFTOperatorFactory>();
std::array<uint32_t, 1> requiredConstantCpuInputs = { 1 };
std::array<uint32_t, 2> requiredConstantCpuInputs = { 1, 2 };
ComPtr<IMLOperatorRegistryPrivate> registryPrivate;
ORT_THROW_IF_FAILED(registry->QueryInterface(IID_PPV_ARGS(&registryPrivate)));

3
onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/OperatorRegistration.cpp
View file

@ -1314,7 +1314,8 @@ void RegisterDmlOperators(IMLOperatorRegistry* registry)
));
}
GpuDFTOperatorFactory::RegisterDFTKernel(registry);
GpuDFTOperatorFactory::RegisterDFTKernel(registry, 17);
GpuDFTOperatorFactory::RegisterDFTKernel(registry, 20);
DmlSTFTOperatorFactory::RegisterSTFTKernel(registry);
DmlGridSampleOperatorFactory::RegisterGridSampleKernel(registry);
}

12
onnxruntime/test/providers/cpu/signal/signal_ops_test.cc
View file

@ -33,6 +33,10 @@ static void TestNaiveDFTFloat(bool onesided, int since_version) {
expected_output.resize(6);
}
test.AddInput<float>("input", shape, input);
if (since_version == 20) {
test.AddInput<int64_t>("dft_length", {}, {5});
test.AddInput<int64_t>("axis", {}, {-2});
}
test.AddAttribute<int64_t>("onesided", static_cast<int64_t>(onesided));
test.AddOutput<float>("output", output_shape, expected_output);
test.Run();
@ -53,6 +57,10 @@ static void TestRadix2DFTFloat(bool onesided, int since_version) {
expected_output.resize(10);
}
test.AddInput<float>("input", shape, input);
if (since_version == 20) {
test.AddInput<int64_t>("dft_length", {}, {8});
test.AddInput<int64_t>("axis", {}, {1});
}
test.AddAttribute<int64_t>("onesided", static_cast<int64_t>(onesided));
test.AddOutput<float>("output", output_shape, expected_output);
test.Run();
@ -67,6 +75,10 @@ static void TestInverseFloat(int since_version) {
vector<float> expected_output = {1.000f, 0.000f, 2.000f, 0.000f, 3.000f, 0.000f, 4.000f, 0.000f, 5.000f, 0.000f};
test.AddInput<float>("input", shape, input);
if (since_version == 20) {
test.AddInput<int64_t>("dft_length", {}, {5});
test.AddInput<int64_t>("axis", {}, {1});
}
test.AddAttribute<int64_t>("inverse", static_cast<int64_t>(true));
test.AddOutput<float>("output", shape, expected_output);
test.Run();