From 31be65899841af53928451311658b8c5a7500831 Mon Sep 17 00:00:00 2001
From: Carl Pearson <cwpearson@users.noreply.github.com>
Date: Tue, 25 Jun 2024 08:09:07 -0600
Subject: [PATCH] spgemm: add profiling regions to native implementations
 (#2253)

* spgemm: add profiling regions to native implementations

* Add profiling region to KokkosSPGEMM::KokkosSPGEMM_symbolic
---
 sparse/impl/KokkosSparse_spgemm_impl_def.hpp   | 2 ++
 sparse/impl/KokkosSparse_spgemm_impl_kkmem.hpp | 2 ++
 sparse/impl/KokkosSparse_spgemm_impl_seq.hpp   | 5 +++++
 sparse/impl/KokkosSparse_spgemm_impl_speed.hpp | 2 ++
 4 files changed, 11 insertions(+)

diff --git a/sparse/impl/KokkosSparse_spgemm_impl_def.hpp b/sparse/impl/KokkosSparse_spgemm_impl_def.hpp
index a420a81c90..54e4e228c8 100644
--- a/sparse/impl/KokkosSparse_spgemm_impl_def.hpp
+++ b/sparse/impl/KokkosSparse_spgemm_impl_def.hpp
@@ -59,6 +59,7 @@ void KokkosSPGEMM<HandleType, a_row_view_t_, a_lno_nnz_view_t_,
                   a_scalar_nnz_view_t_, b_lno_row_view_t_, b_lno_nnz_view_t_,
                   b_scalar_nnz_view_t_>::KokkosSPGEMM_symbolic(c_row_view_t
                                                                    rowmapC_) {
+  Kokkos::Profiling::pushRegion("KokkosSparse::spgemm_symbolic[NATIVE]");
   {
     if (KOKKOSKERNELS_VERBOSE) {
       std::cout << "SYMBOLIC PHASE" << std::endl;
@@ -162,6 +163,7 @@ void KokkosSPGEMM<HandleType, a_row_view_t_, a_lno_nnz_view_t_,
     }
 #endif
   }
+  Kokkos::Profiling::popRegion();
 }
 
 template <typename HandleType, typename a_row_view_t_,
diff --git a/sparse/impl/KokkosSparse_spgemm_impl_kkmem.hpp b/sparse/impl/KokkosSparse_spgemm_impl_kkmem.hpp
index 0a3113e5e4..a7fee71278 100644
--- a/sparse/impl/KokkosSparse_spgemm_impl_kkmem.hpp
+++ b/sparse/impl/KokkosSparse_spgemm_impl_kkmem.hpp
@@ -1241,6 +1241,7 @@ void KokkosSPGEMM<HandleType, a_row_view_t_, a_lno_nnz_view_t_,
         c_row_view_t rowmapC_, c_lno_nnz_view_t entriesC_,
         c_scalar_nnz_view_t valuesC_,
         KokkosKernels::Impl::ExecSpaceType lcl_my_exec_space) {
+  Kokkos::Profiling::pushRegion("KokkosSparse::spgemm_numeric[NATIVE/HASH]");
   if (KOKKOSKERNELS_VERBOSE) {
     std::cout << "\tHASH MODE" << std::endl;
   }
@@ -1699,6 +1700,7 @@ void KokkosSPGEMM<HandleType, a_row_view_t_, a_lno_nnz_view_t_,
   if (KOKKOSKERNELS_VERBOSE) {
     std::cout << "\t\tNumeric TIME:" << timer1.seconds() << std::endl;
   }
+  Kokkos::Profiling::popRegion();
 }
 
 // 01/30/2020: this code seems to be unused within any of the kokkos-kernels
diff --git a/sparse/impl/KokkosSparse_spgemm_impl_seq.hpp b/sparse/impl/KokkosSparse_spgemm_impl_seq.hpp
index b37328aa70..3915bb7d93 100644
--- a/sparse/impl/KokkosSparse_spgemm_impl_seq.hpp
+++ b/sparse/impl/KokkosSparse_spgemm_impl_seq.hpp
@@ -32,6 +32,8 @@ void spgemm_debug_symbolic(KernelHandle *handle,
                            bool /* transposeA */, blno_row_view_t_ row_mapB,
                            blno_nnz_view_t_ entriesB, bool /* transposeB */,
                            clno_row_view_t_ row_mapC) {
+  Kokkos::Profiling::pushRegion("KokkosSparse::spgemm_symbolic[NATIVE/DEBUG]");
+
   typename alno_row_view_t_::HostMirror h_rma =
       Kokkos::create_mirror_view(row_mapA);
   Kokkos::deep_copy(h_rma, row_mapA);
@@ -100,6 +102,7 @@ void spgemm_debug_symbolic(KernelHandle *handle,
   handle->get_spgemm_handle()->set_c_nnz(result_index);
   Kokkos::deep_copy(row_mapC, h_rmc);
   Kokkos::fence();
+  Kokkos::Profiling::popRegion();
 }
 
 template <typename KernelHandle, typename alno_row_view_t_,
@@ -119,6 +122,7 @@ void spgemm_debug_numeric(KernelHandle * /* handle */,
                           bscalar_nnz_view_t_ valuesB, bool /* transposeB */,
                           clno_row_view_t_ row_mapC, clno_nnz_view_t_ entriesC,
                           cscalar_nnz_view_t_ valuesC) {
+  Kokkos::Profiling::pushRegion("KokkosSparse::spgemm_numeric[NATIVE/DEBUG]");
   typename alno_row_view_t_::HostMirror h_rma =
       Kokkos::create_mirror_view(row_mapA);
   Kokkos::deep_copy(h_rma, row_mapA);
@@ -199,6 +203,7 @@ void spgemm_debug_numeric(KernelHandle * /* handle */,
   Kokkos::deep_copy(entriesC, h_entc);
   Kokkos::deep_copy(valuesC, h_valc);
   Kokkos::fence();
+  Kokkos::Profiling::popRegion();
 }
 
 }  // namespace Impl
diff --git a/sparse/impl/KokkosSparse_spgemm_impl_speed.hpp b/sparse/impl/KokkosSparse_spgemm_impl_speed.hpp
index 942e1be7bf..954bfb3f3e 100644
--- a/sparse/impl/KokkosSparse_spgemm_impl_speed.hpp
+++ b/sparse/impl/KokkosSparse_spgemm_impl_speed.hpp
@@ -468,6 +468,7 @@ void KokkosSPGEMM<HandleType, a_row_view_t_, a_lno_nnz_view_t_,
         c_row_view_t rowmapC_, c_lno_nnz_view_t entriesC_,
         c_scalar_nnz_view_t valuesC_,
         KokkosKernels::Impl::ExecSpaceType my_exec_space_) {
+  Kokkos::Profiling::pushRegion("KokkosSparse::spgemm_numeric[NATIVE/SPEED]");
   if (KOKKOSKERNELS_VERBOSE) {
     std::cout << "\tSPEED MODE" << std::endl;
   }
@@ -604,6 +605,7 @@ void KokkosSPGEMM<HandleType, a_row_view_t_, a_lno_nnz_view_t_,
     std::cout << "\t\tNumeric SPEED TIME WITH FREE:"
               << numeric_speed_timer_with_free.seconds() << std::endl;
   }
+  Kokkos::Profiling::popRegion();
 }
 }  // namespace Impl
 }  // namespace KokkosSparse