[STFT][CPU] Improve performance of STFT for CPU by reusage RDFT jit Executor

src/core/shape_inference/include/stft_shape_inference.hpp

@@ -46,15 +46,17 @@ std::vector<TRShape> shape_infer(const STFT* op,
     if (signal_shape.rank().is_dynamic()) {
         return {signal_shape};
     } else if (!frame_size || !frame_step) {
-        return {TRShape{signal_shape[0], -1, -1, 2}};
+        return {TRShape{signal_shape[0], TDim(ov::util::dim::inf_bound), TDim(ov::util::dim::inf_bound), 2}};
     }
 
     const auto& frame_size_val = (*frame_size)[0];
     const auto& frame_step_val = (*frame_step)[0];
 
     NODE_SHAPE_INFER_CHECK(op,
                            input_shapes,
-                           0 < frame_size_val && frame_size_val < signal_shape[1].get_interval().get_max_val(),
+                           0 < frame_size_val && (signal_shape[1].is_static()
+                                                      ? frame_size_val < signal_shape[1].get_length()
+                                                      : frame_size_val < signal_shape[1].get_interval().get_max_val()),
                            "Provided frame size is ",
                            frame_size_val,
                            " but must be in range [1, ",
@@ -77,7 +79,7 @@ std::vector<TRShape> shape_infer(const STFT* op,
                            "].");
 
     const auto& batch_dim = signal_shape[0];
-    const TDim frame_size_dim = TDim{frame_size_val};
+    const TDim frame_size_dim = static_cast<TDim>(frame_size_val);
     const TDim signal_frame_size_diff = signal_shape[1] - frame_size_dim;
     TDim fft_samples_dim = (frame_size_val / 2) + 1;
 

src/plugins/intel_cpu/src/cpu_types.cpp

@@ -181,6 +181,7 @@ static const TypeToNameMap& get_type_to_name_tbl() {
         {"IDFT", Type::DFT},
         {"RDFT", Type::RDFT},
         {"IRDFT", Type::RDFT},
+        {"STFT", Type::STFT},
         {"Abs", Type::Math},
         {"Acos", Type::Math},
         {"Acosh", Type::Math},
@@ -339,6 +340,7 @@ std::string NameFromType(const Type type) {
         CASE(ShuffleChannels);
         CASE(DFT);
         CASE(RDFT);
+        CASE(STFT);
         CASE(Math);
         CASE(CTCLoss);
         CASE(Bucketize);

src/plugins/intel_cpu/src/cpu_types.h

@@ -89,6 +89,7 @@ enum class Type {
     ShuffleChannels,
     DFT,
     RDFT,
+    STFT,
     Math,
     CTCLoss,
     Bucketize,

src/plugins/intel_cpu/src/nodes/rdft.cpp

@@ -842,126 +842,123 @@ struct RDFTJitExecutor : public RDFTExecutor {
 };
 #endif
 
-struct RDFTRefExecutor : public RDFTExecutor {
-    RDFTRefExecutor(bool inverse) : RDFTExecutor(inverse) {}
+RDFTRefExecutor::RDFTRefExecutor(bool inverse) : RDFTExecutor(inverse) {}
 
-    private:
-        std::vector<float> generateTwiddlesDFT(size_t inputSize, size_t outputSize, enum dft_type type) override {
-            std::vector<float> twiddles(inputSize * outputSize * 2);
-            parallel_for2d(outputSize, inputSize, [&] (size_t k, size_t n) {
-                double angle = 2 * PI * k * n / inputSize;
-                if (!isInverse)
-                    angle = -angle;
-                twiddles[(k * inputSize + n) * 2] = std::cos(angle);
-                twiddles[(k * inputSize + n) * 2 + 1] = std::sin(angle);
-            });
-            return twiddles;
-        }
-
-        void dftRealToComplex(float* inputPtr, const float* twiddlesPtr, float* outputPtr,
-                     size_t inputSize, size_t outputSize, bool parallelize) {
-            auto dftIteration = [&] (size_t k) {
-                float real = 0, imag = 0;
-                for (size_t n = 0; n < inputSize; n++) {
-                    float cos = twiddlesPtr[2 * (k * inputSize + n)];
-                    float sin = twiddlesPtr[2 * (k * inputSize + n) + 1];
-                    real += inputPtr[n] * cos;
-                    imag += inputPtr[n] * sin;
-                }
-                outputPtr[2 * k] = real;
-                outputPtr[2 * k + 1] = imag;
-            };
-            if (parallelize) {
-                parallel_for(outputSize, dftIteration);
-            } else {
-                for (size_t k = 0; k < outputSize; k++) {
-                    dftIteration(k);
-                }
-            }
-        }
+std::vector<float> RDFTRefExecutor::generateTwiddlesDFT(size_t inputSize, size_t outputSize, enum dft_type type) {
+    std::vector<float> twiddles(inputSize * outputSize * 2);
+    parallel_for2d(outputSize, inputSize, [&] (size_t k, size_t n) {
+        double angle = 2 * PI * k * n / inputSize;
+        if (!isInverse)
+            angle = -angle;
+        twiddles[(k * inputSize + n) * 2] = std::cos(angle);
+        twiddles[(k * inputSize + n) * 2 + 1] = std::sin(angle);
+    });
+    return twiddles;
+}
 
-        void dftComplexToComplex(float* inputPtr, const float* twiddlesPtr, float* outputPtr,
-                     size_t inputSize, size_t signalSize, size_t outputSize, bool parallelize) {
-            auto dftIteration = [&] (size_t k) {
-                float real = 0, imag = 0;
-                for (size_t n = 0; n < inputSize; n++) {
-                    float cos = twiddlesPtr[2 * (k * outputSize + n)];
-                    float sin = twiddlesPtr[2 * (k * outputSize + n) + 1];
-                    float inputReal = inputPtr[2 * n];
-                    float inputImag = inputPtr[2 * n + 1];
-                    real += inputReal * cos - inputImag * sin;
-                    imag += inputImag * cos + inputReal * sin;
-                }
-                if (isInverse) {
-                    float* inp = inputPtr + 2 * (inputSize - 2 + outputSize % 2);
-                    for (size_t n = inputSize; n < signalSize; n++, inp -= 2) {
-                        float cos = twiddlesPtr[2 * (k * outputSize + n)];
-                        float sin = twiddlesPtr[2 * (k * outputSize + n) + 1];
-                        float inputReal = inp[0];
-                        float inputImag = -inp[1];
-                        real += inputReal * cos - inputImag * sin;
-                        imag += inputImag * cos + inputReal * sin;
-                    }
-                    real /= outputSize;
-                    imag /= outputSize;
-                }
-                outputPtr[2 * k] = real;
-                outputPtr[2 * k + 1] = imag;
-            };
-            if (parallelize) {
-                parallel_for(outputSize, dftIteration);
-            } else {
-                for (size_t k = 0; k < outputSize; k++) {
-                    dftIteration(k);
-                }
-            }
+void RDFTRefExecutor::dftRealToComplex(float* inputPtr, const float* twiddlesPtr, float* outputPtr,
+                size_t inputSize, size_t outputSize, bool parallelize) {
+    auto dftIteration = [&] (size_t k) {
+        float real = 0, imag = 0;
+        for (size_t n = 0; n < inputSize; n++) {
+            float cos = twiddlesPtr[2 * (k * inputSize + n)];
+            float sin = twiddlesPtr[2 * (k * inputSize + n) + 1];
+            real += inputPtr[n] * cos;
+            imag += inputPtr[n] * sin;
         }
-
-        void dftComplexToReal(float* inputPtr, const float* twiddlesPtr, float* outputPtr,
-                     size_t inputSize, size_t signalSize, size_t outputSize, bool parallelize) {
-            auto dftIteration = [&] (size_t k) {
-                float real = 0;
-                for (size_t n = 0; n < inputSize; n++) {
-                    float cos = twiddlesPtr[2 * (k * outputSize + n)];
-                    float sin = twiddlesPtr[2 * (k * outputSize + n) + 1];
-                    float inputReal = inputPtr[2 * n];
-                    float inputImag = inputPtr[2 * n + 1];
-                    real += inputReal * cos - inputImag * sin;
-                }
-                if (isInverse) {
-                    float* inp = inputPtr + 2 * (inputSize - 2 + outputSize % 2);
-                    for (size_t n = inputSize; n < signalSize; n++, inp -= 2) {
-                        float cos = twiddlesPtr[2 * (k * outputSize + n)];
-                        float sin = twiddlesPtr[2 * (k * outputSize + n) + 1];
-                        float inputReal = inp[0];
-                        float inputImag = inp[1];
-                        real += inputReal * cos + inputImag * sin;
-                    }
-                    real /= outputSize;
-                }
-                outputPtr[k] = real;
-            };
-            if (parallelize) {
-                parallel_for(outputSize, dftIteration);
-            } else {
-                for (size_t k = 0; k < outputSize; k++) {
-                    dftIteration(k);
-                }
+        outputPtr[2 * k] = real;
+        outputPtr[2 * k + 1] = imag;
+    };
+    if (parallelize) {
+        parallel_for(outputSize, dftIteration);
+    } else {
+        for (size_t k = 0; k < outputSize; k++) {
+            dftIteration(k);
+        }
+    }
+}
+
+void RDFTRefExecutor::dftComplexToComplex(float* inputPtr, const float* twiddlesPtr, float* outputPtr,
+                size_t inputSize, size_t signalSize, size_t outputSize, bool parallelize) {
+    auto dftIteration = [&] (size_t k) {
+        float real = 0, imag = 0;
+        for (size_t n = 0; n < inputSize; n++) {
+            float cos = twiddlesPtr[2 * (k * outputSize + n)];
+            float sin = twiddlesPtr[2 * (k * outputSize + n) + 1];
+            float inputReal = inputPtr[2 * n];
+            float inputImag = inputPtr[2 * n + 1];
+            real += inputReal * cos - inputImag * sin;
+            imag += inputImag * cos + inputReal * sin;
+        }
+        if (isInverse) {
+            float* inp = inputPtr + 2 * (inputSize - 2 + outputSize % 2);
+            for (size_t n = inputSize; n < signalSize; n++, inp -= 2) {
+                float cos = twiddlesPtr[2 * (k * outputSize + n)];
+                float sin = twiddlesPtr[2 * (k * outputSize + n) + 1];
+                float inputReal = inp[0];
+                float inputImag = -inp[1];
+                real += inputReal * cos - inputImag * sin;
+                imag += inputImag * cos + inputReal * sin;
             }
+            real /= outputSize;
+            imag /= outputSize;
         }
-
-        void dft(float* inputPtr, const float* twiddlesPtr, float* outputPtr,
-                 size_t inputSize, size_t signalSize, size_t outputSize,
-                 enum dft_type type, bool parallelize) override {
-            if (type == real_to_complex) {
-                dftRealToComplex(inputPtr, twiddlesPtr, outputPtr, inputSize, outputSize, parallelize);
-            } else if (type == complex_to_complex) {
-                dftComplexToComplex(inputPtr, twiddlesPtr, outputPtr, inputSize, signalSize, outputSize, parallelize);
-            } else if (type == complex_to_real) {
-                dftComplexToReal(inputPtr, twiddlesPtr, outputPtr, inputSize, signalSize, outputSize, parallelize);
+        outputPtr[2 * k] = real;
+        outputPtr[2 * k + 1] = imag;
+    };
+    if (parallelize) {
+        parallel_for(outputSize, dftIteration);
+    } else {
+        for (size_t k = 0; k < outputSize; k++) {
+            dftIteration(k);
+        }
+    }
+}
+
+void RDFTRefExecutor::dftComplexToReal(float* inputPtr, const float* twiddlesPtr, float* outputPtr,
+                size_t inputSize, size_t signalSize, size_t outputSize, bool parallelize) {
+    auto dftIteration = [&] (size_t k) {
+        float real = 0;
+        for (size_t n = 0; n < inputSize; n++) {
+            float cos = twiddlesPtr[2 * (k * outputSize + n)];
+            float sin = twiddlesPtr[2 * (k * outputSize + n) + 1];
+            float inputReal = inputPtr[2 * n];
+            float inputImag = inputPtr[2 * n + 1];
+            real += inputReal * cos - inputImag * sin;
+        }
+        if (isInverse) {
+            float* inp = inputPtr + 2 * (inputSize - 2 + outputSize % 2);
+            for (size_t n = inputSize; n < signalSize; n++, inp -= 2) {
+                float cos = twiddlesPtr[2 * (k * outputSize + n)];
+                float sin = twiddlesPtr[2 * (k * outputSize + n) + 1];
+                float inputReal = inp[0];
+                float inputImag = inp[1];
+                real += inputReal * cos + inputImag * sin;
             }
+            real /= outputSize;
         }
-};
+        outputPtr[k] = real;
+    };
+    if (parallelize) {
+        parallel_for(outputSize, dftIteration);
+    } else {
+        for (size_t k = 0; k < outputSize; k++) {
+            dftIteration(k);
+        }
+    }
+}
+
+void RDFTRefExecutor::dft(float* inputPtr, const float* twiddlesPtr, float* outputPtr,
+            size_t inputSize, size_t signalSize, size_t outputSize,
+            enum dft_type type, bool parallelize) {
+    if (type == real_to_complex) {
+        dftRealToComplex(inputPtr, twiddlesPtr, outputPtr, inputSize, outputSize, parallelize);
+    } else if (type == complex_to_complex) {
+        dftComplexToComplex(inputPtr, twiddlesPtr, outputPtr, inputSize, signalSize, outputSize, parallelize);
+    } else if (type == complex_to_real) {
+        dftComplexToReal(inputPtr, twiddlesPtr, outputPtr, inputSize, signalSize, outputSize, parallelize);
+    }
+}
 
 struct RDFTKey {
     bool isInverse;

src/plugins/intel_cpu/src/nodes/rdft.h

@@ -71,6 +71,44 @@ struct RDFTExecutor {
                                                   enum dft_type type, bool useFFT);
 };
 
+struct RDFTRefExecutor : public RDFTExecutor {
+    RDFTRefExecutor(bool inverse);
+
+private:
+    std::vector<float> generateTwiddlesDFT(size_t inputSize, size_t outputSize, enum dft_type type) override;
+    void dftRealToComplex(float* inputPtr,
+                          const float* twiddlesPtr,
+                          float* outputPtr,
+                          size_t inputSize,
+                          size_t outputSize,
+                          bool parallelize);
+
+    void dftComplexToComplex(float* inputPtr,
+                             const float* twiddlesPtr,
+                             float* outputPtr,
+                             size_t inputSize,
+                             size_t signalSize,
+                             size_t outputSize,
+                             bool parallelize);
+
+    void dftComplexToReal(float* inputPtr,
+                          const float* twiddlesPtr,
+                          float* outputPtr,
+                          size_t inputSize,
+                          size_t signalSize,
+                          size_t outputSize,
+                          bool parallelize);
+
+    void dft(float* inputPtr,
+             const float* twiddlesPtr,
+             float* outputPtr,
+             size_t inputSize,
+             size_t signalSize,
+             size_t outputSize,
+             enum dft_type type,
+             bool parallelize) override;
+};
+
 class RDFT : public Node {
 public:
     RDFT(const std::shared_ptr<ov::Node>& op, const GraphContext::CPtr context);