Fix stft not found (#9924)

Fix #9922
Oneflow-Inc · Mar 2, 2023 · a843da1 · a843da1
1 parent 8ba8eb6
commit a843da1
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 16 deletions.
diff --git a/oneflow/user/kernels/cufftplancache.h → oneflow/user/kernels/cufft_plan_cache.h b/oneflow/user/kernels/cufftplancache.h → oneflow/user/kernels/cufft_plan_cache.h
@@ -14,6 +14,9 @@ See the License for the specific language governing permissions and
 limitations under the License.
 */
 
+#ifndef ONEFLOW_USER_KERNELS_CUFFT_PLAN_CACHE_H_
+#define ONEFLOW_USER_KERNELS_CUFFT_PLAN_CACHE_H_
+
 #include <cufft.h>
 #include <cufftXt.h>
 #include "oneflow/core/framework/framework.h"
@@ -22,10 +25,13 @@ limitations under the License.
 #include "oneflow/core/kernel/kernel.h"
 
 namespace oneflow {
+
 namespace {
 
 constexpr int max_rank = 3;
 
+}
+
 struct CuFFtParams {
   int32_t ndim;
   int32_t output_shape[max_rank + 1];
@@ -34,8 +40,8 @@ struct CuFFtParams {
   int32_t output_strides[max_rank + 1];
   int32_t* rank;
   int32_t batch;
-  CuFFtParams(int32_t dims, int32_t* r, const Stride& in_strides, const Stride& out_strides,
-              const Shape& in_shape, const Shape& out_shape, int32_t b)
+  CuFFtParams(int32_t dims, int32_t* r, const Stride& in_strides,  // NOLINT
+              const Stride& out_strides, const Shape& in_shape, const Shape& out_shape, int32_t b)
       : ndim(dims), rank(r), batch(b) {
     std::copy(in_strides.begin(), in_strides.end(), input_strides);
     std::copy(out_strides.begin(), out_strides.end(), output_strides);
@@ -49,8 +55,9 @@ class CuFFtConfig {
  public:
   CuFFtConfig(const CuFFtConfig&) = delete;
   CuFFtConfig& operator=(CuFFtConfig const&) = delete;
+  ~CuFFtConfig() = default;
 
-  explicit CuFFtConfig(CuFFtParams& params) {
+  explicit CuFFtConfig(CuFFtParams& params) {  // NOLINT
     infer_cufft_type_();
     cufftPlanMany(&plan_handle_, params.ndim, params.rank, params.input_shape,
                   params.input_strides[0], params.input_strides[1], params.output_shape,
@@ -83,6 +90,6 @@ class CuFFtConfig {
   cufftType exectype_;
 };
 
-}  // namespace
+}  // namespace oneflow
 
-}  // namespace oneflow
+#endif  // ONEFLOW_USER_KERNELS_CUFFT_PLAN_CACHE_H_
diff --git a/oneflow/user/kernels/stft_kernel.cu b/oneflow/user/kernels/stft_kernel.cu
@@ -14,21 +14,15 @@ See the License for the specific language governing permissions and
 limitations under the License.
 */
 
+#include <cuda.h>
+
 #if CUDA_VERSION >= 11000
-#define CUDA_SUPPORT_CUFFT
-#endif
 
-#ifdef CUDA_SUPPORT_CUFFT
+#include "cufft_plan_cache.h"
 
-#include "oneflow/core/framework/framework.h"
-#include "oneflow/core/kernel/new_kernel_util.h"
-#include "oneflow/core/ep/cuda/cuda_stream.h"
-#include <cufft.h>
-#include <cufftXt.h>
-#include "oneflow/core/kernel/kernel.h"
-#include "cufftplancache.h"
 namespace oneflow {
-namespace {}  // namespace
+
+namespace {
 
 template<typename IN, typename OUT>
 __global__ void convert_complex_to_real(IN* dst, const OUT* src, size_t n) {
@@ -73,6 +67,8 @@ __global__ void convert_doublesided(const FFTTYPE* src, FFTTYPE* dst, size_t len
   }
 }
 
+}  // namespace
+
 template<typename IN, typename OUT>
 class StftGpuKernel final : public user_op::OpKernel {
  public:
@@ -158,8 +154,10 @@ class StftGpuKernel final : public user_op::OpKernel {
         const int64_t output_bytes = GetCudaAlignedSize(output_elem_cnt * sizeof(outtype)); \
         return onesided ? output_bytes : 2 * output_bytes;                                  \
       });
+
 REGISTER_STFT_GPU_KERNEL(float, cufftComplex)
 REGISTER_STFT_GPU_KERNEL(double, cufftDoubleComplex)
+
 }  // namespace oneflow
 
 #endif