Introduce Kokkos::Profiling::Region to profile tpl calls (#204)

* Introduce Kokkos::Profiling::Region to profile tpl calls * use ScopedRegion instead of pushRegion/popRegion * Include Kokkos_Profiling_ScopedRegion.hpp * fix: conflicts in KokkosFFT_ROCM_plans.hpp * Add an empty line before introducing the scoped region --------- Co-authored-by: Yuuichi Asahi <[email protected]>
kokkos · Jan 27, 2025 · d42b1f2 · d42b1f2
1 parent d991b0c
commit d42b1f2
Show file tree

Hide file tree

Showing 15 changed files with 139 additions and 7 deletions.
diff --git a/fft/src/KokkosFFT_Cuda_plans.hpp b/fft/src/KokkosFFT_Cuda_plans.hpp
@@ -6,6 +6,7 @@
 #define KOKKOSFFT_CUDA_PLANS_HPP
 
 #include <numeric>
+#include <Kokkos_Profiling_ScopedRegion.hpp>
 #include "KokkosFFT_Cuda_types.hpp"
 #include "KokkosFFT_Extents.hpp"
 #include "KokkosFFT_traits.hpp"
@@ -33,6 +34,8 @@ auto create_plan(const ExecutionSpace& exec_space,
       "InViewType and OutViewType.");
   using in_value_type  = typename InViewType::non_const_value_type;
   using out_value_type = typename OutViewType::non_const_value_type;
+
+  Kokkos::Profiling::ScopedRegion region("KokkosFFT::create_plan[TPL_cufft]");
   auto type = KokkosFFT::Impl::transform_type<ExecutionSpace, in_value_type,
                                               out_value_type>::type();
   auto [in_extents, out_extents, fft_extents, howmany] =
@@ -66,6 +69,8 @@ auto create_plan(const ExecutionSpace& exec_space,
       "InViewType and OutViewType.");
   using in_value_type  = typename InViewType::non_const_value_type;
   using out_value_type = typename OutViewType::non_const_value_type;
+
+  Kokkos::Profiling::ScopedRegion region("KokkosFFT::create_plan[TPL_cufft]");
   auto type = KokkosFFT::Impl::transform_type<ExecutionSpace, in_value_type,
                                               out_value_type>::type();
   [[maybe_unused]] auto [in_extents, out_extents, fft_extents, howmany] =
@@ -99,6 +104,8 @@ auto create_plan(const ExecutionSpace& exec_space,
       "InViewType and OutViewType.");
   using in_value_type  = typename InViewType::non_const_value_type;
   using out_value_type = typename OutViewType::non_const_value_type;
+
+  Kokkos::Profiling::ScopedRegion region("KokkosFFT::create_plan[TPL_cufft]");
   auto type = KokkosFFT::Impl::transform_type<ExecutionSpace, in_value_type,
                                               out_value_type>::type();
   [[maybe_unused]] auto [in_extents, out_extents, fft_extents, howmany] =
@@ -139,7 +146,9 @@ 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;
-  const int rank       = fft_rank;
+
+  Kokkos::Profiling::ScopedRegion region("KokkosFFT::create_plan[TPL_cufft]");
+  const int rank = fft_rank;
   constexpr auto type =
       KokkosFFT::Impl::transform_type<ExecutionSpace, in_value_type,
                                       out_value_type>::type();

diff --git a/fft/src/KokkosFFT_Cuda_transform.hpp b/fft/src/KokkosFFT_Cuda_transform.hpp
@@ -6,6 +6,7 @@
 #define KOKKOSFFT_CUDA_TRANSFORM_HPP
 
 #include <cufft.h>
+#include <Kokkos_Profiling_ScopedRegion.hpp>
 #include "KokkosFFT_asserts.hpp"
 #include "KokkosFFT_Cuda_types.hpp"
 
@@ -14,32 +15,42 @@ namespace Impl {
 
 inline void exec_plan(const ScopedCufftPlan& scoped_plan, cufftReal* idata,
                       cufftComplex* odata, int /*direction*/) {
+  Kokkos::Profiling::ScopedRegion region(
+      "KokkosFFT::exec_plan[TPL_cufftExecR2C]");
   cufftResult cufft_rt = cufftExecR2C(scoped_plan.plan(), idata, odata);
   KOKKOSFFT_THROW_IF(cufft_rt != CUFFT_SUCCESS, "cufftExecR2C failed");
 }
 
 inline void exec_plan(const ScopedCufftPlan& scoped_plan,
                       cufftDoubleReal* idata, cufftDoubleComplex* odata,
                       int /*direction*/) {
+  Kokkos::Profiling::ScopedRegion region(
+      "KokkosFFT::exec_plan[TPL_cufftExecD2Z]");
   cufftResult cufft_rt = cufftExecD2Z(scoped_plan.plan(), idata, odata);
   KOKKOSFFT_THROW_IF(cufft_rt != CUFFT_SUCCESS, "cufftExecD2Z failed");
 }
 
 inline void exec_plan(const ScopedCufftPlan& scoped_plan, cufftComplex* idata,
                       cufftReal* odata, int /*direction*/) {
+  Kokkos::Profiling::ScopedRegion region(
+      "KokkosFFT::exec_plan[TPL_cufftExecC2R]");
   cufftResult cufft_rt = cufftExecC2R(scoped_plan.plan(), idata, odata);
   KOKKOSFFT_THROW_IF(cufft_rt != CUFFT_SUCCESS, "cufftExecC2R failed");
 }
 
 inline void exec_plan(const ScopedCufftPlan& scoped_plan,
                       cufftDoubleComplex* idata, cufftDoubleReal* odata,
                       int /*direction*/) {
+  Kokkos::Profiling::ScopedRegion region(
+      "KokkosFFT::exec_plan[TPL_cufftExecZ2D]");
   cufftResult cufft_rt = cufftExecZ2D(scoped_plan.plan(), idata, odata);
   KOKKOSFFT_THROW_IF(cufft_rt != CUFFT_SUCCESS, "cufftExecZ2D failed");
 }
 
 inline void exec_plan(const ScopedCufftPlan& scoped_plan, cufftComplex* idata,
                       cufftComplex* odata, int direction) {
+  Kokkos::Profiling::ScopedRegion region(
+      "KokkosFFT::exec_plan[TPL_cufftExecC2C]");
   cufftResult cufft_rt =
       cufftExecC2C(scoped_plan.plan(), idata, odata, direction);
   KOKKOSFFT_THROW_IF(cufft_rt != CUFFT_SUCCESS, "cufftExecC2C failed");
@@ -48,6 +59,8 @@ inline void exec_plan(const ScopedCufftPlan& scoped_plan, cufftComplex* idata,
 inline void exec_plan(const ScopedCufftPlan& scoped_plan,
                       cufftDoubleComplex* idata, cufftDoubleComplex* odata,
                       int direction) {
+  Kokkos::Profiling::ScopedRegion region(
+      "KokkosFFT::exec_plan[TPL_cufftExecZ2Z]");
   cufftResult cufft_rt =
       cufftExecZ2Z(scoped_plan.plan(), idata, odata, direction);
   KOKKOSFFT_THROW_IF(cufft_rt != CUFFT_SUCCESS, "cufftExecZ2Z failed");

diff --git a/fft/src/KokkosFFT_Cuda_types.hpp b/fft/src/KokkosFFT_Cuda_types.hpp
@@ -7,6 +7,7 @@
 
 #include <cufft.h>
 #include <Kokkos_Abort.hpp>
+#include <Kokkos_Profiling_ScopedRegion.hpp>
 #include "KokkosFFT_common_types.hpp"
 #include "KokkosFFT_asserts.hpp"
 
@@ -56,6 +57,8 @@ struct ScopedCufftPlan {
   }
 
   ~ScopedCufftPlan() noexcept {
+    Kokkos::Profiling::ScopedRegion region(
+        "KokkosFFT::cleanup_plan[TPL_cufft]");
     cufftResult cufft_rt = cufftDestroy(m_plan);
     if (cufft_rt != CUFFT_SUCCESS) Kokkos::abort("cufftDestroy failed");
   }

diff --git a/fft/src/KokkosFFT_FFTW_Types.hpp b/fft/src/KokkosFFT_FFTW_Types.hpp
@@ -7,6 +7,7 @@
 
 #include <fftw3.h>
 #include <Kokkos_Core.hpp>
+#include <Kokkos_Profiling_ScopedRegion.hpp>
 #include "KokkosFFT_common_types.hpp"
 #include "KokkosFFT_utils.hpp"
 
@@ -104,6 +105,7 @@ struct ScopedFFTWPlan {
   }
 
   ~ScopedFFTWPlan() noexcept {
+    Kokkos::Profiling::ScopedRegion region("KokkosFFT::cleanup_plan[TPL_fftw]");
     if constexpr (std::is_same_v<plan_type, fftwf_plan>) {
       fftwf_destroy_plan(m_plan);
     } else {

diff --git a/fft/src/KokkosFFT_HIP_plans.hpp b/fft/src/KokkosFFT_HIP_plans.hpp
@@ -6,6 +6,7 @@
 #define KOKKOSFFT_HIP_PLANS_HPP
 
 #include <numeric>
+#include <Kokkos_Profiling_ScopedRegion.hpp>
 #include "KokkosFFT_HIP_types.hpp"
 #include "KokkosFFT_Extents.hpp"
 #include "KokkosFFT_traits.hpp"
@@ -33,6 +34,8 @@ auto create_plan(const ExecutionSpace& exec_space,
       "InViewType and OutViewType.");
   using in_value_type  = typename InViewType::non_const_value_type;
   using out_value_type = typename OutViewType::non_const_value_type;
+
+  Kokkos::Profiling::ScopedRegion region("KokkosFFT::create_plan[TPL_hipfft]");
   auto type = KokkosFFT::Impl::transform_type<ExecutionSpace, in_value_type,
                                               out_value_type>::type();
   auto [in_extents, out_extents, fft_extents, howmany] =
@@ -66,6 +69,8 @@ auto create_plan(const ExecutionSpace& exec_space,
       "InViewType and OutViewType.");
   using in_value_type  = typename InViewType::non_const_value_type;
   using out_value_type = typename OutViewType::non_const_value_type;
+
+  Kokkos::Profiling::ScopedRegion region("KokkosFFT::create_plan[TPL_hipfft]");
   auto type = KokkosFFT::Impl::transform_type<ExecutionSpace, in_value_type,
                                               out_value_type>::type();
   [[maybe_unused]] auto [in_extents, out_extents, fft_extents, howmany] =
@@ -99,6 +104,8 @@ auto create_plan(const ExecutionSpace& exec_space,
       "InViewType and OutViewType.");
   using in_value_type  = typename InViewType::non_const_value_type;
   using out_value_type = typename OutViewType::non_const_value_type;
+
+  Kokkos::Profiling::ScopedRegion region("KokkosFFT::create_plan[TPL_hipfft]");
   auto type = KokkosFFT::Impl::transform_type<ExecutionSpace, in_value_type,
                                               out_value_type>::type();
   [[maybe_unused]] auto [in_extents, out_extents, fft_extents, howmany] =
@@ -139,7 +146,9 @@ 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;
-  const int rank       = fft_rank;
+
+  Kokkos::Profiling::ScopedRegion region("KokkosFFT::create_plan[TPL_hipfft]");
+  const int rank = fft_rank;
   constexpr auto type =
       KokkosFFT::Impl::transform_type<ExecutionSpace, in_value_type,
                                       out_value_type>::type();

diff --git a/fft/src/KokkosFFT_HIP_transform.hpp b/fft/src/KokkosFFT_HIP_transform.hpp
@@ -6,6 +6,7 @@
 #define KOKKOSFFT_HIP_TRANSFORM_HPP
 
 #include <hipfft/hipfft.h>
+#include <Kokkos_Profiling_ScopedRegion.hpp>
 #include "KokkosFFT_asserts.hpp"
 #include "KokkosFFT_HIP_types.hpp"
 
@@ -14,32 +15,42 @@ namespace Impl {
 
 inline void exec_plan(const ScopedHIPfftPlan& scoped_plan, hipfftReal* idata,
                       hipfftComplex* odata, int /*direction*/) {
+  Kokkos::Profiling::ScopedRegion region(
+      "KokkosFFT::exec_plan[TPL_hipfftExecR2C]");
   hipfftResult hipfft_rt = hipfftExecR2C(scoped_plan.plan(), idata, odata);
   KOKKOSFFT_THROW_IF(hipfft_rt != HIPFFT_SUCCESS, "hipfftExecR2C failed");
 }
 
 inline void exec_plan(const ScopedHIPfftPlan& scoped_plan,
                       hipfftDoubleReal* idata, hipfftDoubleComplex* odata,
                       int /*direction*/) {
+  Kokkos::Profiling::ScopedRegion region(
+      "KokkosFFT::exec_plan[TPL_hipfftExecD2Z]");
   hipfftResult hipfft_rt = hipfftExecD2Z(scoped_plan.plan(), idata, odata);
   KOKKOSFFT_THROW_IF(hipfft_rt != HIPFFT_SUCCESS, "hipfftExecD2Z failed");
 }
 
 inline void exec_plan(const ScopedHIPfftPlan& scoped_plan, hipfftComplex* idata,
                       hipfftReal* odata, int /*direction*/) {
+  Kokkos::Profiling::ScopedRegion region(
+      "KokkosFFT::exec_plan[TPL_hipfftExecC2R]");
   hipfftResult hipfft_rt = hipfftExecC2R(scoped_plan.plan(), idata, odata);
   KOKKOSFFT_THROW_IF(hipfft_rt != HIPFFT_SUCCESS, "hipfftExecC2R failed");
 }
 
 inline void exec_plan(const ScopedHIPfftPlan& scoped_plan,
                       hipfftDoubleComplex* idata, hipfftDoubleReal* odata,
                       int /*direction*/) {
+  Kokkos::Profiling::ScopedRegion region(
+      "KokkosFFT::exec_plan[TPL_hipfftExecZ2D]");
   hipfftResult hipfft_rt = hipfftExecZ2D(scoped_plan.plan(), idata, odata);
   KOKKOSFFT_THROW_IF(hipfft_rt != HIPFFT_SUCCESS, "hipfftExecZ2D failed");
 }
 
 inline void exec_plan(const ScopedHIPfftPlan& scoped_plan, hipfftComplex* idata,
                       hipfftComplex* odata, int direction) {
+  Kokkos::Profiling::ScopedRegion region(
+      "KokkosFFT::exec_plan[TPL_hipfftExecC2C]");
   hipfftResult hipfft_rt =
       hipfftExecC2C(scoped_plan.plan(), idata, odata, direction);
   KOKKOSFFT_THROW_IF(hipfft_rt != HIPFFT_SUCCESS, "hipfftExecC2C failed");
@@ -48,6 +59,8 @@ inline void exec_plan(const ScopedHIPfftPlan& scoped_plan, hipfftComplex* idata,
 inline void exec_plan(const ScopedHIPfftPlan& scoped_plan,
                       hipfftDoubleComplex* idata, hipfftDoubleComplex* odata,
                       int direction) {
+  Kokkos::Profiling::ScopedRegion region(
+      "KokkosFFT::exec_plan[TPL_hipfftExecZ2Z]");
   hipfftResult hipfft_rt =
       hipfftExecZ2Z(scoped_plan.plan(), idata, odata, direction);
   KOKKOSFFT_THROW_IF(hipfft_rt != HIPFFT_SUCCESS, "hipfftExecZ2Z failed");

diff --git a/fft/src/KokkosFFT_HIP_types.hpp b/fft/src/KokkosFFT_HIP_types.hpp
@@ -7,6 +7,7 @@
 
 #include <hipfft/hipfft.h>
 #include <Kokkos_Abort.hpp>
+#include <Kokkos_Profiling_ScopedRegion.hpp>
 #include "KokkosFFT_common_types.hpp"
 #include "KokkosFFT_asserts.hpp"
 
@@ -56,6 +57,8 @@ struct ScopedHIPfftPlan {
   }
 
   ~ScopedHIPfftPlan() noexcept {
+    Kokkos::Profiling::ScopedRegion region(
+        "KokkosFFT::cleanup_plan[TPL_hipfft]");
     hipfftResult hipfft_rt = hipfftDestroy(m_plan);
     if (hipfft_rt != HIPFFT_SUCCESS) Kokkos::abort("hipfftDestroy failed");
   }

diff --git a/fft/src/KokkosFFT_Host_plans.hpp b/fft/src/KokkosFFT_Host_plans.hpp
@@ -6,6 +6,7 @@
 #define KOKKOSFFT_HOST_PLANS_HPP
 
 #include <numeric>
+#include <Kokkos_Profiling_ScopedRegion.hpp>
 #include "KokkosFFT_default_types.hpp"
 #include "KokkosFFT_Extents.hpp"
 #include "KokkosFFT_traits.hpp"
@@ -37,8 +38,9 @@ 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;
-  const int rank       = fft_rank;
 
+  Kokkos::Profiling::ScopedRegion region("KokkosFFT::create_plan[TPL_fftw]");
+  const int rank = fft_rank;
   auto [in_extents, out_extents, fft_extents, howmany] =
       KokkosFFT::Impl::get_extents(in, out, axes, s, is_inplace);
   int idist    = std::accumulate(in_extents.begin(), in_extents.end(), 1,

diff --git a/fft/src/KokkosFFT_Host_transform.hpp b/fft/src/KokkosFFT_Host_transform.hpp
@@ -6,43 +6,56 @@
 #define KOKKOSFFT_HOST_TRANSFORM_HPP
 
 #include <fftw3.h>
+#include <Kokkos_Profiling_ScopedRegion.hpp>
 
 namespace KokkosFFT {
 namespace Impl {
 
 template <typename ScopedPlanType>
 void exec_plan(const ScopedPlanType& scoped_plan, float* idata,
                fftwf_complex* odata, int /*direction*/) {
+  Kokkos::Profiling::ScopedRegion region(
+      "KokkosFFT::exec_plan[TPL_fftwExecR2C]");
   fftwf_execute_dft_r2c(scoped_plan.plan(), idata, odata);
 }
 
 template <typename ScopedPlanType>
 void exec_plan(const ScopedPlanType& scoped_plan, double* idata,
                fftw_complex* odata, int /*direction*/) {
+  Kokkos::Profiling::ScopedRegion region(
+      "KokkosFFT::exec_plan[TPL_fftwExecD2Z]");
   fftw_execute_dft_r2c(scoped_plan.plan(), idata, odata);
 }
 
 template <typename ScopedPlanType>
 void exec_plan(const ScopedPlanType& scoped_plan, fftwf_complex* idata,
                float* odata, int /*direction*/) {
+  Kokkos::Profiling::ScopedRegion region(
+      "KokkosFFT::exec_plan[TPL_fftwExecC2R]");
   fftwf_execute_dft_c2r(scoped_plan.plan(), idata, odata);
 }
 
 template <typename ScopedPlanType>
 void exec_plan(const ScopedPlanType& scoped_plan, fftw_complex* idata,
                double* odata, int /*direction*/) {
+  Kokkos::Profiling::ScopedRegion region(
+      "KokkosFFT::exec_plan[TPL_fftwExecZ2D]");
   fftw_execute_dft_c2r(scoped_plan.plan(), idata, odata);
 }
 
 template <typename ScopedPlanType>
 void exec_plan(const ScopedPlanType& scoped_plan, fftwf_complex* idata,
                fftwf_complex* odata, int /*direction*/) {
+  Kokkos::Profiling::ScopedRegion region(
+      "KokkosFFT::exec_plan[TPL_fftwExecC2C]");
   fftwf_execute_dft(scoped_plan.plan(), idata, odata);
 }
 
 template <typename ScopedPlanType>
 void exec_plan(const ScopedPlanType& scoped_plan, fftw_complex* idata,
                fftw_complex* odata, int /*direction*/) {
+  Kokkos::Profiling::ScopedRegion region(
+      "KokkosFFT::exec_plan[TPL_fftwExecZ2Z]");
   fftw_execute_dft(scoped_plan.plan(), idata, odata);
 }
 }  // namespace Impl

diff --git a/fft/src/KokkosFFT_ROCM_plans.hpp b/fft/src/KokkosFFT_ROCM_plans.hpp
@@ -5,6 +5,9 @@
 #ifndef KOKKOSFFT_ROCM_PLANS_HPP
 #define KOKKOSFFT_ROCM_PLANS_HPP
 
+#include <numeric>
+#include <algorithm>
+#include <Kokkos_Profiling_ScopedRegion.hpp>
 #include "KokkosFFT_ROCM_types.hpp"
 #include "KokkosFFT_Extents.hpp"
 #include "KokkosFFT_traits.hpp"
@@ -33,12 +36,14 @@ auto create_plan(const ExecutionSpace& exec_space,
       "and the same rank. ExecutionSpace must be accessible to the data in "
       "InViewType and OutViewType.");
 
-  static_assert(
-      InViewType::rank() >= fft_rank,
-      "KokkosFFT::create_plan: Rank of View must be larger than Rank of FFT.");
+  static_assert(InViewType::rank() >= fft_rank,
+                "KokkosFFT::create_plan: Rank of View must be larger than "
+                "Rank of FFT.");
 
   using in_value_type  = typename InViewType::non_const_value_type;
   using out_value_type = typename OutViewType::non_const_value_type;
+
+  Kokkos::Profiling::ScopedRegion region("KokkosFFT::create_plan[TPL_rocfft]");
   constexpr auto type =
       KokkosFFT::Impl::transform_type<ExecutionSpace, in_value_type,
                                       out_value_type>::type();