Skip to content

Commit

Permalink
Remove unnecessary cufftCreate and hipfftCreate to avoid creating pla…
Browse files Browse the repository at this point in the history 
…ns twice
  • Loading branch information
Yuuichi Asahi committed Dec 5, 2024
1 parent 4b82641 commit 20c9978
Show file tree
Hide file tree
Showing 2 changed files with 50 additions and 66 deletions.
58 changes: 25 additions & 33 deletions fft/src/KokkosFFT_Cuda_plans.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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<PlanType>();
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<ExecutionSpace, in_value_type,
out_value_type>::type();
auto [in_extents, out_extents, fft_extents, howmany] =
Expand All @@ -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<PlanType>();
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;
}

Expand All @@ -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<PlanType>();
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<ExecutionSpace, in_value_type,
out_value_type>::type();
[[maybe_unused]] auto [in_extents, out_extents, fft_extents, howmany] =
Expand All @@ -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<PlanType>();
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;
}

Expand All @@ -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<PlanType>();
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<ExecutionSpace, in_value_type,
out_value_type>::type();
[[maybe_unused]] auto [in_extents, out_extents, fft_extents, howmany] =
Expand All @@ -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<PlanType>();
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;
}

Expand Down Expand Up @@ -187,18 +181,16 @@ auto create_plan(const ExecutionSpace& exec_space,
int istride = 1, ostride = 1;

plan = std::make_unique<PlanType>();
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;
}

Expand Down
58 changes: 25 additions & 33 deletions fft/src/KokkosFFT_HIP_plans.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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<PlanType>();
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<ExecutionSpace, in_value_type,
out_value_type>::type();
auto [in_extents, out_extents, fft_extents, howmany] =
Expand All @@ -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<PlanType>();
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;
}

Expand All @@ -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<PlanType>();
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<ExecutionSpace, in_value_type,
out_value_type>::type();
[[maybe_unused]] auto [in_extents, out_extents, fft_extents, howmany] =
Expand All @@ -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<PlanType>();
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;
}

Expand All @@ -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<PlanType>();
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<ExecutionSpace, in_value_type,
out_value_type>::type();
[[maybe_unused]] auto [in_extents, out_extents, fft_extents, howmany] =
Expand All @@ -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<PlanType>();
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;
}

Expand Down Expand Up @@ -187,18 +181,16 @@ auto create_plan(const ExecutionSpace& exec_space,
int istride = 1, ostride = 1;

plan = std::make_unique<PlanType>();
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;
}

Expand Down

0 comments on commit 20c9978

Please sign in to comment.