diff --git a/onnxruntime/contrib_ops/cpu/signal/dft.cc b/onnxruntime/contrib_ops/cpu/signal/dft.cc index d08852b84a..70aa7e2deb 100644 --- a/onnxruntime/contrib_ops/cpu/signal/dft.cc +++ b/onnxruntime/contrib_ops/cpu/signal/dft.cc @@ -39,20 +39,16 @@ ONNX_OPERATOR_KERNEL_EX( kMSExperimentalDomain, 1, kCpuExecutionProvider, - KernelDefBuilder().MayInplace(0, 0).TypeConstraint("T", BuildKernelDefConstraints()), + KernelDefBuilder().MayInplace(0, 0).TypeConstraint("T1", BuildKernelDefConstraints()) + .TypeConstraint("T2", BuildKernelDefConstraints()), STFT); static bool is_real_valued_signal(const onnxruntime::TensorShape & shape) { - // The first dimention is the batch size - // The second dimention is the signal value - return shape.NumDimensions() == 2; + return shape.NumDimensions() == 2 || shape[shape.NumDimensions() - 1] == 1; } static bool is_complex_valued_signal(const onnxruntime::TensorShape& shape) { - // The first dimention is the batch size - // The second dimention is the signal length - // The third dimention is set to 2 and represents the real and imaginary parts of the complex sample - return shape.NumDimensions() == 3 && shape[2] == 2; + return shape.NumDimensions() > 2 && shape[shape.NumDimensions() - 1] == 2; } static bool is_power_of_2(size_t size) { @@ -143,24 +139,27 @@ static T compute_angular_velocity(size_t number_of_samples, bool inverse) { } template -static Status fft_radix2(OpKernelContext* /*ctx*/, size_t batch_idx, - const Tensor* X, Tensor* Y, const Tensor* window, bool is_onesided, bool inverse, +static Status fft_radix2(OpKernelContext* /*ctx*/, + const Tensor* X, Tensor* Y, + size_t X_offset, size_t X_stride, size_t Y_offset, size_t Y_stride, int64_t axis, + const Tensor* window, bool is_onesided, bool inverse, std::vector>& V, std::vector>& temp_output) { // Get shape and significant bits const auto& X_shape = X->Shape(); - size_t number_of_samples = static_cast(X_shape[1]); + size_t number_of_samples = static_cast(X_shape[axis]); unsigned significant_bits = static_cast(log2(number_of_samples)); // Get data - auto* X_data = const_cast(reinterpret_cast(X->DataRaw())) + (batch_idx * number_of_samples); + auto* X_data = const_cast(reinterpret_cast(X->DataRaw())) + X_offset; // Get window U* window_data = nullptr; if (window) { window_data = const_cast(reinterpret_cast(window->DataRaw())); } + size_t Y_data_stride = 1; std::complex* Y_data; if (is_onesided) { if (temp_output.size() != number_of_samples) { @@ -168,7 +167,8 @@ static Status fft_radix2(OpKernelContext* /*ctx*/, size_t batch_idx, } Y_data = temp_output.data(); } else { - Y_data = reinterpret_cast*>(Y->MutableDataRaw()) + (batch_idx * number_of_samples); + Y_data = reinterpret_cast*>(Y->MutableDataRaw()) + Y_offset; + Y_data_stride = Y_stride; } auto angular_velocity = compute_angular_velocity(number_of_samples, inverse); @@ -184,9 +184,9 @@ static Status fft_radix2(OpKernelContext* /*ctx*/, size_t batch_idx, for (size_t i = 0; i < number_of_samples; i++) { size_t bit_reversed_index = bit_reverse(i, significant_bits); - auto x = *(X_data + bit_reversed_index); + auto x = *(X_data + bit_reversed_index*X_stride); auto window_element = window_data ? *(window_data + bit_reversed_index) : 1; - *(Y_data + i) = std::complex(1, 0) * x * window_element; + *(Y_data + i*Y_data_stride) = std::complex(1, 0) * x * window_element; } // Run fft_radix2 @@ -199,8 +199,8 @@ static Status fft_radix2(OpKernelContext* /*ctx*/, size_t batch_idx, auto first_idx = bit_reverse(k, current_significant_bits); auto second_idx = bit_reverse(midpoint + k, current_significant_bits); for (size_t j = 0; j < number_of_samples; j += i) { - std::complex* even = (Y_data + j) + k; - std::complex* odd = (Y_data + j) + (midpoint + k); + std::complex* even = (Y_data + j*Y_data_stride) + k; + std::complex* odd = (Y_data + j*Y_data_stride) + (midpoint + k); std::complex first = *even + (V[first_idx] * *odd); std::complex second = *even + (V[second_idx] * *odd); *even = first; @@ -212,32 +212,34 @@ static Status fft_radix2(OpKernelContext* /*ctx*/, size_t batch_idx, // Scale the output if inverse if (inverse) { for (size_t i = 0; i < number_of_samples; i++) { - std::complex& val = *(Y_data + i); + std::complex& val = *(Y_data + i * Y_data_stride); val /= static_cast(number_of_samples); } } if (is_onesided) { - const auto& Y_shape = Y->Shape(); - size_t fft_output_size = static_cast(Y_shape[1]); - auto destination = reinterpret_cast*>(Y->MutableDataRaw()) + (batch_idx * fft_output_size); - memcpy(destination, Y_data, sizeof(std::complex) * fft_output_size); + auto destination = reinterpret_cast*>(Y->MutableDataRaw()) + Y_offset; + for (size_t i = 0; i < number_of_samples; i++) { + *(destination + Y_stride * i) = *(Y_data + i); + } } return Status::OK(); } template -static Status dft_naive(size_t batch_idx, const Tensor* X, Tensor* Y, const Tensor* window, bool inverse) { +static Status dft_naive(const Tensor* X, Tensor* Y, + size_t X_offset, size_t X_stride, size_t Y_offset, size_t Y_stride, int64_t axis, + const Tensor* window, bool inverse) { // Get shape and significant bits const auto& X_shape = X->Shape(); - size_t number_of_samples = static_cast(X_shape[1]); + size_t number_of_samples = static_cast(X_shape[axis]); const auto& Y_shape = Y->Shape(); - size_t dft_output_size = static_cast(Y_shape[1]); + size_t dft_output_size = static_cast(Y_shape[axis]); // Get data - auto* X_data = const_cast(reinterpret_cast(X->DataRaw())) + (batch_idx * number_of_samples); - auto* Y_data = reinterpret_cast*>(Y->MutableDataRaw()) + (batch_idx * dft_output_size); + auto* X_data = const_cast(reinterpret_cast(X->DataRaw())) + X_offset; + auto* Y_data = reinterpret_cast*>(Y->MutableDataRaw()) + Y_offset; U* window_data = nullptr; if (window) { @@ -247,14 +249,14 @@ static Status dft_naive(size_t batch_idx, const Tensor* X, Tensor* Y, const Tens auto angular_velocity = compute_angular_velocity(number_of_samples, inverse); for (size_t i = 0; i < dft_output_size; i++) { - std::complex& out = *(Y_data + i); + std::complex& out = *(Y_data + i*Y_stride); out.real(0); out.imag(0); for (size_t j = 0; j < number_of_samples; j++) { // vectorize over this loop auto exponential = std::complex(cos(i * j * angular_velocity), sin(i * j * angular_velocity)); auto window_element = window_data ? * (window_data + j) : 1; - auto element = *(X_data + j) * window_element; + auto element = *(X_data + j*X_stride) * window_element; out += exponential * element; } @@ -267,26 +269,65 @@ static Status dft_naive(size_t batch_idx, const Tensor* X, Tensor* Y, const Tens } template -static Status discrete_fourier_transform(OpKernelContext* ctx, const Tensor* X, Tensor* Y, const Tensor* window, bool is_onesided, bool inverse, +static Status discrete_fourier_transform(OpKernelContext* ctx, const Tensor* X, Tensor* Y, int64_t axis, const Tensor* window, bool is_onesided, bool inverse, std::vector>& V, std::vector>& temp_output) { // Get shape const auto& X_shape = X->Shape(); - size_t number_of_batches = static_cast(X_shape[0]); - size_t number_of_samples = static_cast(X_shape[1]); - - // radix 2 fft - for (size_t i = 0; i < number_of_batches; i++) { - if (is_power_of_2(number_of_samples)) { - ORT_RETURN_IF_ERROR((fft_radix2(ctx, i, X, Y, window, is_onesided, inverse, V, temp_output))); - } else { - ORT_RETURN_IF_ERROR((dft_naive(i, X, Y, window, inverse))); - } + const auto& Y_shape = Y->Shape(); + size_t number_of_samples = static_cast(X_shape[axis]); + + auto batch_and_signal_rank = X->Shape().NumDimensions(); + auto total_dfts = static_cast(X->Shape().Size() / X->Shape()[axis]); + if (X->Shape().NumDimensions() > 2) + { + total_dfts /= X->Shape()[X->Shape().NumDimensions() - 1]; + batch_and_signal_rank -= 1; } + // Calculate x/y offsets/strides + for (size_t i = 0; i < total_dfts; i++) + { + size_t X_offset = 0; + size_t X_stride = X_shape.SizeFromDimension(axis+1); + size_t cumulative_packed_stride = total_dfts; + size_t temp = i; + for (size_t r = 0; r < batch_and_signal_rank; r++) { + if (r == static_cast(axis)) + { + continue; + } + cumulative_packed_stride /= X_shape[r]; + auto index = temp / cumulative_packed_stride; + temp -= (index * cumulative_packed_stride); + X_offset += index * X_shape.SizeFromDimension(r + 1); + } + + size_t Y_offset = 0; + size_t Y_stride = Y_shape.SizeFromDimension(axis + 1) / 2; + cumulative_packed_stride = total_dfts; + temp = i; + for (size_t r = 0; r < batch_and_signal_rank; r++) { + if (r == static_cast(axis)) + { + continue; + } + cumulative_packed_stride /= X_shape[r]; + auto index = temp / cumulative_packed_stride; + temp -= (index * cumulative_packed_stride); + Y_offset += index * Y_shape.SizeFromDimension(r + 1) / 2; + } + + if (is_power_of_2(number_of_samples)) { + ORT_RETURN_IF_ERROR((fft_radix2(ctx, X, Y, X_offset, X_stride, Y_offset, Y_stride, axis, window, is_onesided, inverse, V, temp_output))); + } else { + ORT_RETURN_IF_ERROR((dft_naive(X, Y, X_offset, X_stride, Y_offset, Y_stride, axis, window, inverse))); + } + } + return Status::OK(); } -static Status discrete_fourier_transform(OpKernelContext* ctx, bool is_onesided, bool inverse) { +static Status discrete_fourier_transform(OpKernelContext* ctx, int64_t axis, bool is_onesided, bool inverse) { // Get input shape const auto* X = ctx->Input(0); const auto& X_shape = X->Shape(); @@ -295,13 +336,21 @@ static Status discrete_fourier_transform(OpKernelContext* ctx, bool is_onesided, // Get the DFT output size. Onesided will return only the unique values! // note: x >> 1 === std::floor(x / 2.f) - int64_t number_of_samples = static_cast(X_shape[1]); + int64_t number_of_samples = static_cast(X_shape[axis]); auto dft_output_size = is_onesided ? ((number_of_samples >> 1) + 1) : number_of_samples; // Get output shape - auto Y_shape = onnxruntime::TensorShape({X_shape[0], dft_output_size, 2}); + auto Y_shape = onnxruntime::TensorShape(X_shape); + if (X_shape.NumDimensions() == 2) + { + Y_shape = onnxruntime::TensorShape({X_shape[0], dft_output_size, 2}); + } else + { + Y_shape[Y_shape.NumDimensions() - 1] = 2; + } + Y_shape[axis] = dft_output_size; auto Y = ctx->Output(0, Y_shape); // Get data type @@ -312,9 +361,9 @@ static Status discrete_fourier_transform(OpKernelContext* ctx, bool is_onesided, std::vector> V; std::vector> temp_output; if (is_real_valued) { - ORT_RETURN_IF_ERROR((discrete_fourier_transform(ctx, X, Y, nullptr, is_onesided, inverse, V, temp_output))); + ORT_RETURN_IF_ERROR((discrete_fourier_transform(ctx, X, Y, axis, nullptr, is_onesided, inverse, V, temp_output))); } else if (is_complex_valued) { - ORT_RETURN_IF_ERROR((discrete_fourier_transform>(ctx, X, Y, nullptr, is_onesided, inverse, V, temp_output))); + ORT_RETURN_IF_ERROR((discrete_fourier_transform>(ctx, X, Y, axis, nullptr, is_onesided, inverse, V, temp_output))); } else { ORT_THROW("Unsupported input signal shape. The signal's first dimenstion must be the batch dimension and its second dimension must be the signal length dimension. It may optionally include a 3rd dimension of size 2 for complex inputs.", data_type); } @@ -322,9 +371,9 @@ static Status discrete_fourier_transform(OpKernelContext* ctx, bool is_onesided, std::vector> V; std::vector> temp_output; if (is_real_valued) { - ORT_RETURN_IF_ERROR((discrete_fourier_transform(ctx, X, Y, nullptr, is_onesided, inverse, V, temp_output))); + ORT_RETURN_IF_ERROR((discrete_fourier_transform(ctx, X, Y, axis, nullptr, is_onesided, inverse, V, temp_output))); } else if (is_complex_valued) { - ORT_RETURN_IF_ERROR((discrete_fourier_transform>(ctx, X, Y, nullptr, is_onesided, inverse, V, temp_output))); + ORT_RETURN_IF_ERROR((discrete_fourier_transform>(ctx, X, Y, axis, nullptr, is_onesided, inverse, V, temp_output))); } else { ORT_THROW("Unsupported input signal shape. The signal's first dimenstion must be the batch dimension and its second dimension must be the signal length dimension. It may optionally include a 3rd dimension of size 2 for complex inputs.", data_type); } @@ -336,12 +385,12 @@ static Status discrete_fourier_transform(OpKernelContext* ctx, bool is_onesided, } Status DFT::Compute(OpKernelContext* ctx) const { - ORT_RETURN_IF_ERROR(discrete_fourier_transform(ctx, is_onesided_, false)); + ORT_RETURN_IF_ERROR(discrete_fourier_transform(ctx, axis_ + 1, is_onesided_, false)); return Status::OK(); } Status IDFT::Compute(OpKernelContext* ctx) const { - ORT_RETURN_IF_ERROR(discrete_fourier_transform(ctx, false, true)); + ORT_RETURN_IF_ERROR(discrete_fourier_transform(ctx, axis_ + 1, false, true)); return Status::OK(); } @@ -376,9 +425,9 @@ static Status short_time_fourier_transform(OpKernelContext* ctx, bool is_oneside // Get signal const auto* signal = ctx->Input(0); - const auto* window = ctx->Input(1); - const auto* frame_length_tensor = ctx->Input(2); - const auto frame_step = get_scalar_value_from_tensor(ctx->Input(3)); + const auto frame_step = get_scalar_value_from_tensor(ctx->Input(1)); + const auto* window = ctx->Input(2); + const auto* frame_length_tensor = ctx->Input(3); // Get input signal shape const auto& signal_shape = signal->Shape(); @@ -468,7 +517,7 @@ static Status short_time_fourier_transform(OpKernelContext* ctx, bool is_oneside 0); // Run individual dft - ORT_RETURN_IF_ERROR((discrete_fourier_transform(ctx, &input, &output, window, is_onesided, false, V, temp_output))); + ORT_RETURN_IF_ERROR((discrete_fourier_transform(ctx, &input, &output, 1, window, is_onesided, false, V, temp_output))); } } diff --git a/onnxruntime/contrib_ops/cpu/signal/dft.h b/onnxruntime/contrib_ops/cpu/signal/dft.h index 2b04781c70..fc90d48fab 100644 --- a/onnxruntime/contrib_ops/cpu/signal/dft.h +++ b/onnxruntime/contrib_ops/cpu/signal/dft.h @@ -8,16 +8,20 @@ namespace contrib { class DFT final : public OpKernel { bool is_onesided_ = true; + int64_t axis_ = 0; public: explicit DFT(const OpKernelInfo& info) : OpKernel(info) { is_onesided_ = static_cast(info.GetAttrOrDefault("onesided", 0)); + axis_ = info.GetAttrOrDefault("axis", 0); } Status Compute(OpKernelContext* ctx) const override; }; class IDFT final : public OpKernel { + int64_t axis_ = 0; public: explicit IDFT(const OpKernelInfo& info) : OpKernel(info) { + axis_ = info.GetAttrOrDefault("axis", 0); } Status Compute(OpKernelContext* ctx) const override; }; diff --git a/onnxruntime/core/graph/signal_ops/signal_defs.cc b/onnxruntime/core/graph/signal_ops/signal_defs.cc index 6b78bf075a..70720b8f85 100644 --- a/onnxruntime/core/graph/signal_ops/signal_defs.cc +++ b/onnxruntime/core/graph/signal_ops/signal_defs.cc @@ -42,6 +42,24 @@ static T get_scalar_value_from_tensor(const ONNX_NAMESPACE::TensorProto* t) { } } +inline const ONNX_NAMESPACE::TensorShapeProto* getOptionalInputShape(ONNX_NAMESPACE::InferenceContext& ctx, size_t n) { + const auto* input_type = ctx.getInputType(n); + + if (input_type == nullptr) { + return nullptr; + } + + const auto value_case = input_type->value_case(); + if (value_case != ONNX_NAMESPACE::TypeProto::kTensorType && value_case != ONNX_NAMESPACE::TypeProto::kSparseTensorType) { + fail_type_inference("Attribute expected to have tensor or sparse tensor type"); + } + if (value_case == ONNX_NAMESPACE::TypeProto::kTensorType) { + return &input_type->tensor_type().shape(); + } else { + return &input_type->sparse_tensor_type().shape(); + } +} + void RegisterSignalSchemas() { MS_SIGNAL_OPERATOR_SCHEMA(DFT) .SetDomain(kMSExperimentalDomain) @@ -53,132 +71,242 @@ void RegisterSignalSchemas() { "Values can be 0 or 1.", AttributeProto::AttributeType::AttributeProto_AttributeType_INT, static_cast(0)) + .Attr("axis", + "The axis on which to perform the DFT. By default this value is set to 0, which corresponds to the first dimension after the batch index." + "This value must be less than signal_dimN, where signal_dimN is the number of dimensions in the signal.", + AttributeProto::AttributeType::AttributeProto_AttributeType_INT, + static_cast(0)) .Input(0, - "input", - "For complex input, the following shape is expected: [batch_idx][n_fft][2]" - "The final dimension represents the real and imaginary parts of the value." - "For real input, the following shape is expected: [batch_idx][n_fft]" - "The first dimension is the batch dimension.", - "T") + "input", + "For real input, the following shape is expected: [batch_idx][n_fft]." + "For complex input, the following shape is expected: [batch_idx][n_fft][2]." + "The final dimension represents the real and imaginary parts of the value." + "For real multi-dimensional input, the following shape is expected: [batch_idx][signal_dim1][signal_dim2]...[signal_dimN][1]." + "For complex multi-dimensional input, the following shape is expected: [batch_idx][signal_dim1][signal_dim2]...[signal_dimN][2]." + "The first dimension is the batch dimension.", + "T") .Output(0, "output", "The Fourier Transform of the input vector." - "If onesided is 1, [batch_idx][floor(n_fft/2)+1][2]" - "If onesided is 0, [batch_idx][n_fft][2]", + "If signal_dimN = 1, and onesided is 0, [batch_idx][n_fft][2]" + "If signal_dimN = 1, and onesided is 1, [batch_idx][floor(n_fft/2)+1][2]" + "If signal_dimN = 2, and onesided is 0 and axis = 0, [batch_idx][signal_dim1][signal_dim2][2]" + "If signal_dimN = 2, and onesided is 0 and axis = 1, [batch_idx][signal_dim1][signal_dim2][2]" + "If signal_dimN = 2, and onesided is 1 and axis = 0, [batch_idx][floor(signal_dim1/2)+1][signal_dim2][2]" + "If signal_dimN = 2, and onesided is 1 and axis = 1, [batch_idx][signal_dim1][floor(signal_dim2/2)+1][2]", "T") - .TypeConstraint("T", {"tensor(float16)", "tensor(float)", "tensor(double)"}, "") + .TypeConstraint( + "T", + {"tensor(float16)", "tensor(float)", "tensor(double)", "tensor(bfloat16)"}, + "Constrain input and output types to float tensors.") .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) { - propagateElemTypeFromInputToOutput(ctx, 0, 0); - int64_t ndim = 1; + propagateElemTypeFromInputToOutput(ctx, 0, 0); + const int64_t batch_ndim = 1; - bool is_onesided = true; - auto attr_proto = ctx.getAttribute("onesided"); - if (attr_proto && attr_proto->has_i()) { - is_onesided = static_cast(attr_proto->i()); - } - - if (ctx.getInputType(0)->tensor_type().has_shape()) { auto& input_shape = getInputShape(ctx, 0); - ONNX_NAMESPACE::TensorShapeProto result_shape = input_shape; - - if (is_onesided) { - auto n_fft = input_shape.dim(1).dim_value(); - result_shape.mutable_dim(1)->set_dim_value((n_fft >> 1) + 1); - } - auto dim_size = static_cast(input_shape.dim_size()); - if (dim_size == ndim + 1) { // real input - result_shape.add_dim()->set_dim_value(2); // output is same shape, but with extra dim for 2 values (real/imaginary) - } else if (dim_size == ndim + 2) { // complex input, do nothing - } else { - fail_shape_inference( - "the input_shape must [batch_idx][n_fft] for real values or [batch_idx][n_fft][2] for complex values.") + auto has_component_dimension = dim_size > 2; + + ONNX_NAMESPACE::TensorShapeProto result_shape_proto = input_shape; + + bool is_onesided = static_cast(getAttribute(ctx, "onesided", 0)); + if (is_onesided) { + // Since signal_ndim = 1, and multidimensional DFT is not supported, + // only the single signal dim (1) needs to be updated + auto n_fft = input_shape.dim(1).dim_value(); + result_shape_proto.mutable_dim(1)->set_dim_value((n_fft >> 1) + 1); } - updateOutputShape(ctx, 0, result_shape); - } + + if (has_component_dimension) { + result_shape_proto.mutable_dim(static_cast(dim_size - 1))->set_dim_value(2); + } else { + result_shape_proto.add_dim()->set_dim_value(2); + } + + updateOutputShape(ctx, 0, result_shape_proto); }); - ; MS_SIGNAL_OPERATOR_SCHEMA(IDFT) .SetDomain(kMSExperimentalDomain) .SinceVersion(1) .SetDoc(R"DOC(IDFT)DOC") + .Attr("axis", + "The axis on which to perform the DFT. By default this value is set to 0, which corresponds to the first dimension after the batch index." + "This value must be less than signal_dimN, where signal_dimN is the number of dimensions in the signal.", + AttributeProto::AttributeType::AttributeProto_AttributeType_INT, + static_cast(0)) .Input(0, "input", - "A complex signal of dimension signal_ndim." - "The last dimension of the tensor should be 2," - "representing the real and imaginary components of complex numbers," - "and should have at least signal_ndim + 2 dimensions." + "For real input, the following shape is expected: [batch_idx][n_fft]." + "For complex input, the following shape is expected: [batch_idx][n_fft][2]." + "The final dimension represents the real and imaginary parts of the value." + "For real multi-dimensional input, the following shape is expected: [batch_idx][signal_dim1][signal_dim2]...[signal_dimN][1]." + "For complex multi-dimensional input, the following shape is expected: [batch_idx][signal_dim1][signal_dim2]...[signal_dimN][2]." "The first dimension is the batch dimension.", "T") .Output(0, "output", - "The inverse fourier transform of the input vector," - "using the same format as the input.", - "T") - .TypeConstraint("T", {"tensor(float16)", "tensor(float)", "tensor(double)"}, "") + "The inverse discrete Fourier transform of the input. " + "If signal_dimN = 1, [batch_idx][n_fft][2]" + "If signal_dimN = 2 and axis = 0, [batch_idx][signal_dim1][signal_dim2][2]" + "If signal_dimN = 2 and axis = 1, [batch_idx][signal_dim1][signal_dim2][2]" + "For all types of input, the last dimension of the output represents the components of a complex number.", + "T", + OpSchema::Single, + true, + 1, + OpSchema::NonDifferentiable) + .TypeConstraint( + "T", + {"tensor(float16)", "tensor(float)", "tensor(double)", "tensor(bfloat16)"}, + "Constrain input and output types to float tensors.") .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) { - propagateElemTypeFromInputToOutput(ctx, 0, 0); - int64_t ndim = 1; - auto attr_proto = ctx.getAttribute("signal_ndim"); - if (attr_proto && attr_proto->has_i()) { - ndim = static_cast(attr_proto->i()); - } + propagateElemTypeFromInputToOutput(ctx, 0, 0); + const int64_t batch_ndim = 1; + + auto& input_shape = getInputShape(ctx, 0); + ONNX_NAMESPACE::TensorShapeProto result_shape = input_shape; + auto dim_size = static_cast(input_shape.dim_size()); + auto has_component_dimension = dim_size > 2; - auto& input_shape = getInputShape(ctx, 0); - ONNX_NAMESPACE::TensorShapeProto result_shape = input_shape; + if (has_component_dimension) { + result_shape.mutable_dim(static_cast(dim_size - 1))->set_dim_value(2); + } else { + result_shape.add_dim()->set_dim_value(2); + } - auto dim_size = static_cast(input_shape.dim_size()); - if (dim_size == ndim + 1) { // real input - result_shape.add_dim()->set_dim_value(2); // output is same shape, but with extra dim for 2 values (real/imaginary) - } else if (dim_size == ndim + 2) { // complex input, do nothing - } else { - fail_shape_inference( - "the input_shape must have 1 + signal_ndim dimensions for real inputs, or 2 + signal_ndim dimensions for complex input.") - } - - updateOutputShape(ctx, 0, result_shape); + updateOutputShape(ctx, 0, result_shape); }); MS_SIGNAL_OPERATOR_SCHEMA(STFT) .SetDomain(kMSExperimentalDomain) .SinceVersion(1) .SetDoc(R"DOC(STFT)DOC") - .Attr("onesided", - "If True (default), only values for half of the fft size are returned because the real-to-complex Fourier transform satisfies the conjugate symmetry." - "The output tensor will return the first floor(n_fft/2) + 1 values from the DFT." - "Values can be 0 or 1.", - AttributeProto::AttributeType::AttributeProto_AttributeType_INT, - static_cast(1)) + .Attr( + "onesided", + "If onesided is 1, only values for w in [0, 1, 2, ..., floor(n_fft/2) + 1] are returned because " + "the real-to-complex Fourier transform satisfies the conjugate symmetry, i.e., X[m, w] = X[m,w]=X[m,n_fft-w]*. " + "Note if the input or window tensors are complex, then onesided output is not possible. " + "Enabling onesided with real inputs performs a Real-valued fast Fourier transform (RFFT)." + "When invoked with real or complex valued input, the default value is 0. " + "Values can be 0 or 1.", + AttributeProto::INT, + static_cast(0)) .Input(0, "signal", - "A complex signal of dimension signal_ndim." - "The last dimension of the tensor should be 2," - "representing the real and imaginary components of complex numbers," - "and should have at least signal_ndim + 2 dimensions." - "The first dimension is the batch dimension.", - "T1") - .Input(1, - "window", - "A tensor representing the window that will be slid over the input signal.", + "Input tensor representing a real or complex valued signal. " + "For real input, the following shape is expected: [batch_size][signal_length]. " + "For complex input, the following shape is expected: [batch_size][signal_length][2], where " + "[batch_size][signal_length][0] represents the real component and [batch_size][signal_length][1] represents the imaginary component of the signal.", "T1", - OpSchema::FormalParameterOption::Optional) - .Input(2, - "frame_length", // frame_length, fft_length, pad_mode - "Size of the fft.", - "T2", - OpSchema::FormalParameterOption::Optional) - .Input(3, + OpSchema::Single, + true, + 1, + OpSchema::NonDifferentiable) + .Input(1, "frame_step", "The number of samples to step between successive DFTs.", - "T2") + "T2", + OpSchema::Single, + true, + 1, + OpSchema::NonDifferentiable) + .Input(2, + "window", + "A tensor representing the window that will be slid over the signal." + "The window must have rank 1 with shape: [window_shape]. " + "It's an optional value. ", + "T1", + OpSchema::Optional, + true, + 1, + OpSchema::NonDifferentiable) + .Input(3, + "frame_length", + "A scalar representing the size of the DFT. " + "It's an optional value.", + "T2", + OpSchema::Optional, + true, + 1, + OpSchema::NonDifferentiable) .Output(0, "output", "The inverse fourier transform of the input vector," "using the same format as the input.", "T1") - .TypeConstraint("T1", {"tensor(float16)", "tensor(float)", "tensor(double)"}, "") - .TypeConstraint("T2", {"tensor(int64)"}, ""); + .TypeConstraint( + "T1", + {"tensor(float)", + "tensor(float16)", + "tensor(double)", + "tensor(bfloat16)"}, + "Constrain signal and output to float tensors.") + .TypeConstraint( + "T2", + {"tensor(int64)"}, + "Constrain scalar length types to int64_t.") + .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) { + propagateElemTypeFromInputToOutput(ctx, 0, 0); + constexpr int64_t batch_ndim = 1; + constexpr int64_t component_ndim = 1; + + // Get inputs + auto& input_shape = getInputShape(ctx, 0); + auto frame_step = get_scalar_value_from_tensor(ctx.getInputData(1)); + const ONNX_NAMESPACE::TensorShapeProto* window_input = nullptr; + try { + window_input = getOptionalInputShape(ctx, 2); + } catch (...) { + window_input = nullptr; + } + + const ONNX_NAMESPACE::TensorShapeProto* frame_length_input = nullptr; + try { + frame_length_input = getOptionalInputShape(ctx, 3); + } catch (...) { + frame_length_input = nullptr; + } + + // Determine the size of the DFT based on the 2 optional inputs window and frame_length. One must be set. + int64_t dft_size = 0; + if (window_input == nullptr && frame_length_input == nullptr) { + fail_type_inference("STFT expects to have at least one of these inputs set: [window, frame_length]."); + } else if (window_input != nullptr && window_input->dim_size() > 0 && frame_length_input != nullptr) { + if (window_input->dim_size() != 1) { + fail_type_inference("STFT's window input, must have rank = 1."); + } + auto window_length = window_input->dim(0).dim_value(); + auto frame_length = get_scalar_value_from_tensor(ctx.getInputData(3)); + if (window_length != frame_length) { + fail_type_inference("If STFT has both a window input and frame_length specified, the dimension of the window must match the frame_length specified!"); + } + dft_size = window_length; + } else if (window_input != nullptr && window_input->dim_size() > 0) { + if (window_input->dim_size() != 1) { + fail_type_inference("STFT's window input, must have rank = 1."); + } + dft_size = window_input->dim(0).dim_value(); + } else if (frame_length_input != nullptr) { + dft_size = get_scalar_value_from_tensor(ctx.getInputData(3)); + } + + bool is_onesided = static_cast(getAttribute(ctx, "onesided", 0)); + if (is_onesided) { + dft_size = is_onesided ? ((dft_size >> 1) + 1) : dft_size; + } + + auto signal_size = input_shape.dim(1).dim_value(); + auto n_dfts = static_cast(std::floor((signal_size - dft_size) / static_cast(frame_step)) + 1); + + // The output has the following shape: [batch_size][frames][dft_unique_bins][2] + ONNX_NAMESPACE::TensorShapeProto result_shape_proto; + result_shape_proto.add_dim()->set_dim_value(input_shape.dim(0).dim_value()); // batch size + result_shape_proto.add_dim()->set_dim_value(n_dfts); + result_shape_proto.add_dim()->set_dim_value(dft_size); + result_shape_proto.add_dim()->set_dim_value(2); + updateOutputShape(ctx, 0, result_shape_proto); + }); // Window Functions MS_SIGNAL_OPERATOR_SCHEMA(HannWindow) diff --git a/winml/test/api/LearningModelSessionAPITest.cpp b/winml/test/api/LearningModelSessionAPITest.cpp index c2aec18038..219c74525f 100644 --- a/winml/test/api/LearningModelSessionAPITest.cpp +++ b/winml/test/api/LearningModelSessionAPITest.cpp @@ -476,10 +476,13 @@ static void WindowFunction(const wchar_t* window_operator_name, TensorKind kind) #endif #if !defined(BUILD_INBOX) && defined(BUILD_MS_EXPERIMENTAL_OPS) -static void DiscreteFourierTransform(bool is_onesided = false) { - std::vector shape = {1, 5}; - std::vector output_shape = {1, 5, 2}; - output_shape[1] = is_onesided ? (1 + (shape[1] >> 1)) : shape[1]; +static void DiscreteFourierTransform(size_t axis, bool is_onesided = false) { + auto axis_dim = axis + 1; + printf("\nDiscrete Fourier Transform [axis=%d, is_onesided=%s]\n", static_cast(axis_dim), is_onesided ? "true" : "false"); + + std::vector shape = {2, 5, 8, 1}; + std::vector output_shape = {2, 5, 8, 2}; + output_shape[axis_dim] = is_onesided ? (1 + (shape[axis_dim] >> 1)) : shape[axis_dim]; auto model = LearningModelBuilder::Create(13) @@ -487,6 +490,7 @@ static void DiscreteFourierTransform(bool is_onesided = false) { .Outputs().Add(LearningModelBuilder::CreateTensorFeatureDescriptor(L"Output.Spectra", TensorKind::Float, output_shape)) .Operators().Add(Operator(L"DFT", MS_EXPERIMENTAL_DOMAIN) .SetInput(L"input", L"Input.Signal") + .SetAttribute(L"axis", TensorInt64Bit::CreateFromArray({}, {INT64(axis)})) .SetAttribute(L"onesided", TensorInt64Bit::CreateFromArray({}, {is_onesided})) .SetOutput(L"output", L"Output.Spectra")) .CreateModel(); @@ -495,19 +499,38 @@ static void DiscreteFourierTransform(bool is_onesided = false) { LearningModelBinding binding(session); // Populate binding - binding.Bind(L"Input.Signal", TensorFloat::CreateFromArray(shape, {1, 2, 3, 4, 5})); + binding.Bind( + L"Input.Signal", + TensorFloat::CreateFromArray( + shape, + {1, 2, 3, 4, 5, 6, 7, 8, + 1, 2, 3, 4, 5, 6, 7, 8, + 1, 2, 3, 4, 5, 6, 7, 8, + 1, 2, 3, 4, 5, 6, 7, 8, + 1, 2, 3, 4, 5, 6, 7, 8, + + 2, 4, 6, 8, 10, 12, 14, 16, + 2, 4, 6, 8, 10, 12, 14, 16, + 2, 4, 6, 8, 10, 12, 14, 16, + 2, 4, 6, 8, 10, 12, 14, 16, + 2, 4, 6, 8, 10, 12, 14, 16, + })); // Evaluate auto result = session.Evaluate(binding, L""); - // Check results - printf("Output.Spectra\n"); - auto y_tensor = result.Outputs().Lookup(L"Output.Spectra").as(); - auto y_ivv = y_tensor.GetAsVectorView(); - for (int i = 0; i < output_shape[0] * output_shape[1] * 2; i += 2) { - printf("(%f + %fi), ", y_ivv.GetAt(i), y_ivv.GetAt(i + 1)); - } - printf("\n"); + // // Check results + // printf("Output.Spectra\n"); + // auto y_tensor = result.Outputs().Lookup(L"Output.Spectra").as(); + // auto y_ivv = y_tensor.GetAsVectorView(); + // for (uint32_t i = 0; i < y_ivv.Size(); i+=2) { + // auto format_size = 16 * (!is_onesided || axis == 0) + 10 * (is_onesided && axis == 1); + // if (i % format_size == 0 && i != 0) { + // printf("\n"); + // } + // printf("(%.2f + %.2fi), ", y_ivv.GetAt(i), y_ivv.GetAt(i + 1)); + // } + // printf("\n"); } #endif @@ -612,20 +635,20 @@ static void STFT(size_t batch_size, size_t signal_size, size_t dft_size, printf("%f, ", window_ivv.GetAt(i)); } printf("\n"); - printf("Output.STFT\n"); - // Check results - auto y_tensor = result.Outputs().Lookup(L"Output.STFT").as(); - auto y_ivv = y_tensor.GetAsVectorView(); - auto size = y_ivv.Size(); - WINML_EXPECT_EQUAL(size, n_dfts * output_shape[2] * 2); - for (size_t dft_idx = 0; dft_idx < n_dfts; dft_idx++) { - for (size_t i = 0; INT64(i) < output_shape[2]; i++) { - auto real_idx = static_cast((i * 2) + (2 * dft_idx * output_shape[2])); - printf("(%d, %f , %fi), ", static_cast(i), y_ivv.GetAt(real_idx), y_ivv.GetAt(real_idx + 1)); - } - } - - printf("\n"); + //printf("Output.STFT\n"); + //// Check results + //auto y_tensor = result.Outputs().Lookup(L"Output.STFT").as(); + //auto y_ivv = y_tensor.GetAsVectorView(); + //auto size = y_ivv.Size(); + //WINML_EXPECT_EQUAL(size, n_dfts * output_shape[2] * 2); + //for (size_t dft_idx = 0; dft_idx < n_dfts; dft_idx++) { + // for (size_t i = 0; INT64(i) < output_shape[2]; i++) { + // auto real_idx = static_cast((i * 2) + (2 * dft_idx * output_shape[2])); + // printf("(%d, %f , %fi), ", static_cast(i), y_ivv.GetAt(real_idx), y_ivv.GetAt(real_idx + 1)); + // } + //} + // + //printf("\n"); } #endif @@ -913,8 +936,11 @@ static void ModelBuilding_ConstantMatmul() { static void ModelBuilding_DiscreteFourierTransform() { #if !defined(BUILD_INBOX) && defined(BUILD_MS_EXPERIMENTAL_OPS) - DiscreteFourierTransform(false /*onesided*/); - DiscreteFourierTransform(true /*onesided*/); + DiscreteFourierTransform(0, false /*onesided*/); + DiscreteFourierTransform(0, true /*onesided*/); + DiscreteFourierTransform(1, false /*onesided*/); + DiscreteFourierTransform(1, true /*onesided*/); + #endif }