[phi] move stft to phi - Step 1 (#53517)

* [phi]mv StftKernel to phi

* [phi] fix KernelSignature

* [phi]fix arr error

* [phi] Disable check_dygraph

* [phi]fix include

* [phi] rewrite mutable_data, add output register

* [phi] fix  Alloc

* [phi] fix Alloc again

* [phi] fix mutable_data

* [phi] fix onesided_out Resize

paddle/fluid/operators/stft_op.cc

@@ -162,10 +162,6 @@ REGISTER_OPERATOR(stft,
 
 REGISTER_OPERATOR(stft_grad, ops::StftGradOp);
 
-REGISTER_OP_CPU_KERNEL(stft,
-                       ops::StftKernel<phi::CPUContext, float>,
-                       ops::StftKernel<phi::CPUContext, double>);
-
 REGISTER_OP_CPU_KERNEL(stft_grad,
                        ops::StftGradKernel<phi::CPUContext, float>,
                        ops::StftGradKernel<phi::CPUContext, double>);

paddle/fluid/operators/stft_op.cu

@@ -16,10 +16,6 @@
 
 namespace ops = paddle::operators;
 
-REGISTER_OP_CUDA_KERNEL(stft,
-                        ops::StftKernel<phi::GPUContext, float>,
-                        ops::StftKernel<phi::GPUContext, double>);
-
 REGISTER_OP_CUDA_KERNEL(stft_grad,
                         ops::StftGradKernel<phi::GPUContext, float>,
                         ops::StftGradKernel<phi::GPUContext, double>);

paddle/fluid/operators/stft_op.h

@@ -27,79 +27,6 @@
 namespace paddle {
 namespace operators {
 
-template <typename DeviceContext, typename T>
-class StftKernel : public framework::OpKernel<T> {
- public:
-  /*
-    Batch Signals (N, T) -> Frames (N, n_fft, num_frames) -> FFTR2C -> (N,
-    n_fft/2 + 1, num_frames) or (N, n_fft, num_frames)
-  */
-  void Compute(const framework::ExecutionContext& ctx) const override {
-    using C = paddle::platform::complex<T>;
-    const phi::DenseTensor* x = ctx.Input<phi::DenseTensor>("X");
-    const phi::DenseTensor* window = ctx.Input<phi::DenseTensor>("Window");
-    phi::DenseTensor* out = ctx.Output<phi::DenseTensor>("Out");
-    out->mutable_data<C>(ctx.GetPlace());
-
-    const size_t x_rank = x->dims().size();
-    const size_t out_rank = out->dims().size();
-
-    const int n_fft = ctx.Attr<int>("n_fft");
-    const int hop_length = ctx.Attr<int>("hop_length");
-    const bool normalized = ctx.Attr<bool>("normalized");
-    const bool onesided = ctx.Attr<bool>("onesided");
-
-    const int n_frames = out->dims()[out_rank - 1];
-    const int seq_length = x->dims()[x_rank - 1];
-
-    auto& dev_ctx = ctx.device_context<DeviceContext>();
-
-    std::vector<int64_t> axes = {1};
-
-    // Frame
-    phi::DenseTensor frames;
-    framework::DDim frames_dims(out->dims());
-    frames_dims.at(axes.back()) = n_fft;
-    frames.mutable_data<T>(frames_dims, ctx.GetPlace());
-    phi::funcs::FrameFunctor<DeviceContext, T>()(dev_ctx,
-                                                 x,
-                                                 &frames,
-                                                 seq_length,
-                                                 n_fft,
-                                                 n_frames,
-                                                 hop_length,
-                                                 /*is_grad*/ false);
-
-    // Window
-    phi::DenseTensor frames_w;
-    frames_w.mutable_data<T>(frames_dims, ctx.GetPlace());
-    ElementwiseComputeEx<MulFunctor<T>, DeviceContext, T>(
-        ctx, &frames, window, axes.back(), MulFunctor<T>(), &frames_w);
-
-    // FFTR2C
-    phi::funcs::FFTNormMode normalization;
-    if (normalized) {
-      normalization = phi::funcs::get_norm_from_string("ortho", true);
-    } else {
-      normalization = phi::funcs::get_norm_from_string("backward", true);
-    }
-    phi::funcs::FFTR2CFunctor<DeviceContext, T, C> fft_r2c_func;
-
-    if (onesided) {
-      fft_r2c_func(dev_ctx, frames_w, out, axes, normalization, true);
-    } else {
-      framework::DDim onesided_dims(out->dims());
-      const int64_t onesided_axis_size = out->dims().at(axes.back()) / 2 + 1;
-      onesided_dims.at(axes.back()) = onesided_axis_size;
-      phi::DenseTensor onesided_out;
-      onesided_out.mutable_data<C>(onesided_dims, ctx.GetPlace());
-      fft_r2c_func(dev_ctx, frames_w, &onesided_out, axes, normalization, true);
-      phi::funcs::FFTFillConj<DeviceContext, C>(
-          dev_ctx, &onesided_out, out, axes);
-    }
-  }
-};
-
 template <typename DeviceContext, typename T>
 class StftGradKernel : public framework::OpKernel<T> {
  public:

paddle/phi/kernels/cpu/stft_kernel.cc

+// Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "paddle/phi/kernels/stft_kernel.h"
+#include "paddle/phi/backends/cpu/cpu_context.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/impl/stft_kernel_impl.h"
+
+PD_REGISTER_KERNEL(stft, CPU, ALL_LAYOUT, phi::StftKernel, float, double) {
+  if (kernel_key.dtype() == phi::DataType::FLOAT16 &&
+      kernel_key.dtype() == phi::DataType::FLOAT32 &&
+      kernel_key.dtype() == phi::DataType::FLOAT64) {
+    kernel->OutputAt(0).SetDataType(phi::DataType::COMPLEX64);
+  } else {
+    kernel->OutputAt(0).SetDataType(phi::DataType::COMPLEX128);
+  }
+}

paddle/phi/kernels/gpu/stft_kernel.cu

+// Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "paddle/phi/kernels/stft_kernel.h"
+#include "paddle/phi/backends/gpu/gpu_context.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/impl/stft_kernel_impl.h"
+
+PD_REGISTER_KERNEL(stft, GPU, ALL_LAYOUT, phi::StftKernel, float, double) {}

paddle/phi/kernels/impl/stft_kernel_impl.h

+// Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+
+#include <vector>
+#include "paddle/phi/common/complex.h"
+#include "paddle/phi/core/dense_tensor.h"
+#include "paddle/phi/kernels/cpu/elementwise.h"
+#include "paddle/phi/kernels/funcs/elementwise_base.h"
+#include "paddle/phi/kernels/funcs/elementwise_functor.h"
+#include "paddle/phi/kernels/funcs/fft.h"
+#include "paddle/phi/kernels/funcs/fft_fill_conj.h"
+#include "paddle/phi/kernels/funcs/frame_functor.h"
+
+namespace phi {
+
+template <typename T, typename Context>
+void StftKernel(const Context& ctx,
+                const DenseTensor& x,
+                const DenseTensor& window,
+                int n_fft,
+                int hop_length,
+                bool normalized,
+                bool onesided,
+                DenseTensor* out) {
+  using C = phi::dtype::complex<T>;
+
+  ctx.template Alloc<C>(out);
+
+  const size_t x_rank = x.dims().size();
+  const size_t out_rank = out->dims().size();
+
+  const int n_frames = out->dims()[out_rank - 1];
+  const int seq_length = x.dims()[x_rank - 1];
+
+  std::vector<int64_t> axes = {1};
+
+  // Frame
+  phi::DenseTensor frames;
+  phi::DDim frames_dims(out->dims());
+  frames_dims.at(axes.back()) = n_fft;
+  frames.Resize(frames_dims);
+  ctx.template Alloc<T>(&frames);
+  phi::funcs::FrameFunctor<Context, T>()(ctx,
+                                         &x,
+                                         &frames,
+                                         seq_length,
+                                         n_fft,
+                                         n_frames,
+                                         hop_length,
+                                         /*is_grad*/ false);
+
+  // Window
+  phi::DenseTensor frames_w;
+  frames_w.Resize(frames_dims);
+  ctx.template Alloc<T>(&frames_w);
+  phi::funcs::ElementwiseCompute<phi::funcs::MultiplyFunctor<T>, T, T>(
+      ctx,
+      frames,
+      window,
+      phi::funcs::MultiplyFunctor<T>(),
+      &frames_w,
+      axes.back());
+
+  // FFTR2C
+  phi::funcs::FFTNormMode normalization;
+  if (normalized) {
+    normalization = phi::funcs::get_norm_from_string("ortho", true);
+  } else {
+    normalization = phi::funcs::get_norm_from_string("backward", true);
+  }
+  phi::funcs::FFTR2CFunctor<Context, T, C> fft_r2c_func;
+
+  if (onesided) {
+    fft_r2c_func(ctx, frames_w, out, axes, normalization, true);
+  } else {
+    phi::DDim onesided_dims(out->dims());
+    const int64_t onesided_axis_size = out->dims().at(axes.back()) / 2 + 1;
+    onesided_dims.at(axes.back()) = onesided_axis_size;
+    phi::DenseTensor onesided_out;
+    onesided_out.Resize(onesided_dims);
+    ctx.template Alloc<T>(&onesided_out);
+    fft_r2c_func(ctx, frames_w, &onesided_out, axes, normalization, true);
+    phi::funcs::FFTFillConj<Context, C>(ctx, &onesided_out, out, axes);
+  }
+}
+}  // namespace phi

paddle/phi/kernels/stft_kernel.h

+/* Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+    http://www.apache.org/licenses/LICENSE-2.0
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License. */
+
+#pragma once
+
+#include "paddle/phi/core/dense_tensor.h"
+
+namespace phi {
+template <typename T, typename Context>
+void StftKernel(const Context& ctx,
+                const DenseTensor& x,
+                const DenseTensor& window,
+                int n_fft,
+                int hop_length,
+                bool normalized,
+                bool onesided,
+                DenseTensor* out);
+
+}  // namespace phi

paddle/phi/ops/compat/stft_sig.cc

+// Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "paddle/phi/core/compat/op_utils.h"
+
+namespace phi {
+
+KernelSignature StftOpArgumentMapping(const ArgumentMappingContext& ctx) {
+  return KernelSignature("stft",
+                         {"X", "Window"},
+                         {"n_fft", "hop_length", "normalized", "onesided"},
+                         {"Out"});
+}
+
+}  // namespace phi
+
+PD_REGISTER_ARG_MAPPING_FN(stft, phi::StftOpArgumentMapping);

python/paddle/fluid/tests/unittests/test_stft_op.py

@@ -80,7 +80,7 @@ class TestStftOp(OpTest):
 
     def test_check_output(self):
         paddle.enable_static()
-        self.check_output()
+        self.check_output(check_dygraph=False)
         paddle.disable_static()
 
     def test_check_grad_normal(self):