From a323455f042c6fa26a0e608b6ce404c605c21447 Mon Sep 17 00:00:00 2001 From: yasahi-hpc <57478230+yasahi-hpc@users.noreply.github.com> Date: Thu, 5 Dec 2024 15:02:01 +0100 Subject: [PATCH] Remove unnecessary cufftCreate and hipfftCreate to avoid creating plans twice (#212) Co-authored-by: Yuuichi Asahi --- fft/src/KokkosFFT_Cuda_plans.hpp | 58 ++++++++++++++------------------ fft/src/KokkosFFT_HIP_plans.hpp | 58 ++++++++++++++------------------ 2 files changed, 50 insertions(+), 66 deletions(-) diff --git a/fft/src/KokkosFFT_Cuda_plans.hpp b/fft/src/KokkosFFT_Cuda_plans.hpp index 9abd8451..7f4547f9 100644 --- a/fft/src/KokkosFFT_Cuda_plans.hpp +++ b/fft/src/KokkosFFT_Cuda_plans.hpp @@ -35,13 +35,6 @@ auto create_plan(const ExecutionSpace& exec_space, using in_value_type = typename InViewType::non_const_value_type; using out_value_type = typename OutViewType::non_const_value_type; - plan = std::make_unique(); - cufftResult cufft_rt = cufftCreate(&(*plan)); - KOKKOSFFT_THROW_IF(cufft_rt != CUFFT_SUCCESS, "cufftCreate failed"); - - cudaStream_t stream = exec_space.cuda_stream(); - cufftSetStream((*plan), stream); - auto type = KokkosFFT::Impl::transform_type::type(); auto [in_extents, out_extents, fft_extents, howmany] = @@ -50,9 +43,14 @@ auto create_plan(const ExecutionSpace& exec_space, int fft_size = std::accumulate(fft_extents.begin(), fft_extents.end(), 1, std::multiplies<>()); - cufft_rt = cufftPlan1d(&(*plan), nx, type, howmany); + plan = std::make_unique(); + cufftResult cufft_rt = cufftPlan1d(&(*plan), nx, type, howmany); KOKKOSFFT_THROW_IF(cufft_rt != CUFFT_SUCCESS, "cufftPlan1d failed"); + cudaStream_t stream = exec_space.cuda_stream(); + cufft_rt = cufftSetStream((*plan), stream); + KOKKOSFFT_THROW_IF(cufft_rt != CUFFT_SUCCESS, "cufftSetStream failed"); + return fft_size; } @@ -78,13 +76,6 @@ auto create_plan(const ExecutionSpace& exec_space, using in_value_type = typename InViewType::non_const_value_type; using out_value_type = typename OutViewType::non_const_value_type; - plan = std::make_unique(); - cufftResult cufft_rt = cufftCreate(&(*plan)); - KOKKOSFFT_THROW_IF(cufft_rt != CUFFT_SUCCESS, "cufftCreate failed"); - - cudaStream_t stream = exec_space.cuda_stream(); - cufftSetStream((*plan), stream); - auto type = KokkosFFT::Impl::transform_type::type(); [[maybe_unused]] auto [in_extents, out_extents, fft_extents, howmany] = @@ -93,9 +84,14 @@ auto create_plan(const ExecutionSpace& exec_space, int fft_size = std::accumulate(fft_extents.begin(), fft_extents.end(), 1, std::multiplies<>()); - cufft_rt = cufftPlan2d(&(*plan), nx, ny, type); + plan = std::make_unique(); + cufftResult cufft_rt = cufftPlan2d(&(*plan), nx, ny, type); KOKKOSFFT_THROW_IF(cufft_rt != CUFFT_SUCCESS, "cufftPlan2d failed"); + cudaStream_t stream = exec_space.cuda_stream(); + cufft_rt = cufftSetStream((*plan), stream); + KOKKOSFFT_THROW_IF(cufft_rt != CUFFT_SUCCESS, "cufftSetStream failed"); + return fft_size; } @@ -121,13 +117,6 @@ auto create_plan(const ExecutionSpace& exec_space, using in_value_type = typename InViewType::non_const_value_type; using out_value_type = typename OutViewType::non_const_value_type; - plan = std::make_unique(); - cufftResult cufft_rt = cufftCreate(&(*plan)); - KOKKOSFFT_THROW_IF(cufft_rt != CUFFT_SUCCESS, "cufftCreate failed"); - - cudaStream_t stream = exec_space.cuda_stream(); - cufftSetStream((*plan), stream); - auto type = KokkosFFT::Impl::transform_type::type(); [[maybe_unused]] auto [in_extents, out_extents, fft_extents, howmany] = @@ -138,9 +127,14 @@ auto create_plan(const ExecutionSpace& exec_space, int fft_size = std::accumulate(fft_extents.begin(), fft_extents.end(), 1, std::multiplies<>()); - cufft_rt = cufftPlan3d(&(*plan), nx, ny, nz, type); + plan = std::make_unique(); + cufftResult cufft_rt = cufftPlan3d(&(*plan), nx, ny, nz, type); KOKKOSFFT_THROW_IF(cufft_rt != CUFFT_SUCCESS, "cufftPlan3d failed"); + cudaStream_t stream = exec_space.cuda_stream(); + cufft_rt = cufftSetStream((*plan), stream); + KOKKOSFFT_THROW_IF(cufft_rt != CUFFT_SUCCESS, "cufftSetStream failed"); + return fft_size; } @@ -187,18 +181,16 @@ auto create_plan(const ExecutionSpace& exec_space, int istride = 1, ostride = 1; plan = std::make_unique(); - cufftResult cufft_rt = cufftCreate(&(*plan)); - KOKKOSFFT_THROW_IF(cufft_rt != CUFFT_SUCCESS, "cufftCreate failed"); - - cudaStream_t stream = exec_space.cuda_stream(); - cufftSetStream((*plan), stream); - - cufft_rt = cufftPlanMany(&(*plan), rank, fft_extents.data(), - in_extents.data(), istride, idist, - out_extents.data(), ostride, odist, type, howmany); + cufftResult cufft_rt = cufftPlanMany( + &(*plan), rank, fft_extents.data(), in_extents.data(), istride, idist, + out_extents.data(), ostride, odist, type, howmany); KOKKOSFFT_THROW_IF(cufft_rt != CUFFT_SUCCESS, "cufftPlanMany failed"); + cudaStream_t stream = exec_space.cuda_stream(); + cufft_rt = cufftSetStream((*plan), stream); + KOKKOSFFT_THROW_IF(cufft_rt != CUFFT_SUCCESS, "cufftSetStream failed"); + return fft_size; } diff --git a/fft/src/KokkosFFT_HIP_plans.hpp b/fft/src/KokkosFFT_HIP_plans.hpp index c94ed23e..1e0ba043 100644 --- a/fft/src/KokkosFFT_HIP_plans.hpp +++ b/fft/src/KokkosFFT_HIP_plans.hpp @@ -35,13 +35,6 @@ auto create_plan(const ExecutionSpace& exec_space, using in_value_type = typename InViewType::non_const_value_type; using out_value_type = typename OutViewType::non_const_value_type; - plan = std::make_unique(); - hipfftResult hipfft_rt = hipfftCreate(&(*plan)); - KOKKOSFFT_THROW_IF(hipfft_rt != HIPFFT_SUCCESS, "hipfftCreate failed"); - - hipStream_t stream = exec_space.hip_stream(); - hipfftSetStream((*plan), stream); - auto type = KokkosFFT::Impl::transform_type::type(); auto [in_extents, out_extents, fft_extents, howmany] = @@ -50,9 +43,14 @@ auto create_plan(const ExecutionSpace& exec_space, int fft_size = std::accumulate(fft_extents.begin(), fft_extents.end(), 1, std::multiplies<>()); - hipfft_rt = hipfftPlan1d(&(*plan), nx, type, howmany); + plan = std::make_unique(); + hipfftResult hipfft_rt = hipfftPlan1d(&(*plan), nx, type, howmany); KOKKOSFFT_THROW_IF(hipfft_rt != HIPFFT_SUCCESS, "hipfftPlan1d failed"); + hipStream_t stream = exec_space.hip_stream(); + hipfft_rt = hipfftSetStream((*plan), stream); + KOKKOSFFT_THROW_IF(hipfft_rt != HIPFFT_SUCCESS, "hipfftSetStream failed"); + return fft_size; } @@ -78,13 +76,6 @@ auto create_plan(const ExecutionSpace& exec_space, using in_value_type = typename InViewType::non_const_value_type; using out_value_type = typename OutViewType::non_const_value_type; - plan = std::make_unique(); - hipfftResult hipfft_rt = hipfftCreate(&(*plan)); - KOKKOSFFT_THROW_IF(hipfft_rt != HIPFFT_SUCCESS, "hipfftCreate failed"); - - hipStream_t stream = exec_space.hip_stream(); - hipfftSetStream((*plan), stream); - auto type = KokkosFFT::Impl::transform_type::type(); [[maybe_unused]] auto [in_extents, out_extents, fft_extents, howmany] = @@ -93,9 +84,14 @@ auto create_plan(const ExecutionSpace& exec_space, int fft_size = std::accumulate(fft_extents.begin(), fft_extents.end(), 1, std::multiplies<>()); - hipfft_rt = hipfftPlan2d(&(*plan), nx, ny, type); + plan = std::make_unique(); + hipfftResult hipfft_rt = hipfftPlan2d(&(*plan), nx, ny, type); KOKKOSFFT_THROW_IF(hipfft_rt != HIPFFT_SUCCESS, "hipfftPlan2d failed"); + hipStream_t stream = exec_space.hip_stream(); + hipfft_rt = hipfftSetStream((*plan), stream); + KOKKOSFFT_THROW_IF(hipfft_rt != HIPFFT_SUCCESS, "hipfftSetStream failed"); + return fft_size; } @@ -121,13 +117,6 @@ auto create_plan(const ExecutionSpace& exec_space, using in_value_type = typename InViewType::non_const_value_type; using out_value_type = typename OutViewType::non_const_value_type; - plan = std::make_unique(); - hipfftResult hipfft_rt = hipfftCreate(&(*plan)); - KOKKOSFFT_THROW_IF(hipfft_rt != HIPFFT_SUCCESS, "hipfftCreate failed"); - - hipStream_t stream = exec_space.hip_stream(); - hipfftSetStream((*plan), stream); - auto type = KokkosFFT::Impl::transform_type::type(); [[maybe_unused]] auto [in_extents, out_extents, fft_extents, howmany] = @@ -138,9 +127,14 @@ auto create_plan(const ExecutionSpace& exec_space, int fft_size = std::accumulate(fft_extents.begin(), fft_extents.end(), 1, std::multiplies<>()); - hipfft_rt = hipfftPlan3d(&(*plan), nx, ny, nz, type); + plan = std::make_unique(); + hipfftResult hipfft_rt = hipfftPlan3d(&(*plan), nx, ny, nz, type); KOKKOSFFT_THROW_IF(hipfft_rt != HIPFFT_SUCCESS, "hipfftPlan3d failed"); + hipStream_t stream = exec_space.hip_stream(); + hipfft_rt = hipfftSetStream((*plan), stream); + KOKKOSFFT_THROW_IF(hipfft_rt != HIPFFT_SUCCESS, "hipfftSetStream failed"); + return fft_size; } @@ -187,18 +181,16 @@ auto create_plan(const ExecutionSpace& exec_space, int istride = 1, ostride = 1; plan = std::make_unique(); - hipfftResult hipfft_rt = hipfftCreate(&(*plan)); - KOKKOSFFT_THROW_IF(hipfft_rt != HIPFFT_SUCCESS, "hipfftCreate failed"); - - hipStream_t stream = exec_space.hip_stream(); - hipfftSetStream((*plan), stream); - - hipfft_rt = hipfftPlanMany(&(*plan), rank, fft_extents.data(), - in_extents.data(), istride, idist, - out_extents.data(), ostride, odist, type, howmany); + hipfftResult hipfft_rt = hipfftPlanMany( + &(*plan), rank, fft_extents.data(), in_extents.data(), istride, idist, + out_extents.data(), ostride, odist, type, howmany); KOKKOSFFT_THROW_IF(hipfft_rt != HIPFFT_SUCCESS, "hipfftPlanMany failed"); + hipStream_t stream = exec_space.hip_stream(); + hipfft_rt = hipfftSetStream((*plan), stream); + KOKKOSFFT_THROW_IF(hipfft_rt != HIPFFT_SUCCESS, "hipfftSetStream failed"); + return fft_size; }